BreakContinueNested.hs

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

module Tests.Control.BreakContinueNested where

-- fir
import FIR

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

-- Early exits out of several loops at once, from the innermost of three nested loops:
--   - 'break @3' exits all three loops,
--   - 'breakContinue @2' continues the outermost loop,
--   - 'breakContinue @1' continues the middle loop.
--
-- In structured control flow, these must exit the innermost loop first,
-- and continue the early exit from its merge block.

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

program :: Module Defs
program = Module do

  entryPoint @"main" @Compute do

    _ <- def @"x" @RW @Float  0
    _ <- def @"i" @RW @Word32 0
    while ( ( < 10 ) <<$>> get @"i" ) do
      modify @"i" ( + 1 )
      loop do
        modify @"x" ( * 2 )
        loop do
          modify @"x" ( + 1 )
          x <- get @"x"
          if x > 100
          then do
            put @"x" ( -1 )
            break @3
          else
            if x > 50
            then breakContinue @2
            else
              if x > 20
              then breakContinue @1
              else continue

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