LoopCondition.hs

{-# LANGUAGE BlockArguments   #-}
{-# LANGUAGE DataKinds        #-}
{-# LANGUAGE PolyKinds        #-}
{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators    #-}

module Tests.Control.LoopCondition where

-- fir
import FIR

------------------------------------------------
-- program

-- A loop condition containing control flow, which modifies a binding.

type Defs =
  '[ "ubo"  ':-> Uniform '[ DescriptorSet 0, Binding 0 ]
                   ( Struct '[ "n" ':-> Word32 ] )
   , "out"  ':-> StorageBuffer '[ DescriptorSet 0, Binding 1 ]
                   ( Struct '[ "res" ':-> Float ] )
   , "main" ':-> EntryPoint '[ LocalSize 1 1 1 ] Compute
   ]

program :: Module Defs
program = Module do

  entryPoint @"main" @Compute do

    n <- use @( Name "ubo" :.: Name "n" )
    _ <- def @"x" @RW @Word32 0
    _ <- def @"y" @RW @Float  0
    while
      ( do
          x <- get @"x"
          if x > 3
          then do
            modify @"y" ( + 1 )
            pure ( x < n )
          else pure ( x < 2 * n )
      ) do
      modify @"x" ( + 1 )

    y <- get @"y"
    assign @( Name "out" :.: Name "res" ) y