BreakInCondition.hs

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

module Tests.Control.BreakInCondition where

-- fir
import FIR

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

-- Early exits from the condition of a loop, which refer to the enclosing loops.

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 @Word32 0

    -- 'break' and 'continue' in the condition of an inner loop.
    loop do
      modify @"x" ( + 1 )
      while
        ( do
            x <- get @"x"
            when ( x > 100 ) ( break @1 )
            when ( x > 50  ) continue
            pure ( x < n )
        ) do
        modify @"x" ( + 2 )

    -- A loop in the condition of a loop, breaking out of both the inner loop
    -- and the loop enclosing the outer loop's condition.
    loop do
      modify @"y" ( + 1 )
      while
        ( do
            loop do
              y <- get @"y"
              when ( y > 100 ) ( break @2 )
              when ( y > 10  ) ( break @1 )
              put @"y" ( y * 2 )
            y <- get @"y"
            pure ( y < n )
        ) do
        modify @"y" ( + 3 )

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