CFG.hs

{-# LANGUAGE BlockArguments        #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE NamedFieldPuns        #-}
{-# LANGUAGE NamedWildCards        #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE PatternSynonyms       #-}
{-# LANGUAGE RecursiveDo           #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TupleSections         #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}

{-# OPTIONS_GHC -Wno-partial-type-signatures #-}

{-|
Module: CodeGen.CFG

Code generation for control flow operations: branching, selections, loops.

See also "CodeGen.Phi" for ϕ-functions necessary for such constructs when using SSA form.
-}

module CodeGen.CFG
  ( newBlock, locally, embed )
  where

-- base
import Prelude
  hiding ( break )
import Control.Arrow
  ( second )
import Control.Monad
  ( unless, when )
import Data.Foldable
  ( for_ )
import Data.Maybe
  ( fromMaybe, isJust )
import Data.Proxy
  ( Proxy )
import Data.Traversable
  ( for )
import Data.Word
  ( Word32 )

-- containers
import Data.Map
  ( Map )
import qualified Data.Map.Strict as Map
import Data.Set
  ( Set )
import qualified Data.Set as Set
  ( fromList, member, toList, isSubsetOf )

-- lens
import Control.Lens
  ( view, use, assign )

-- mtl
import Control.Monad.Except
  ( throwError )

-- text-short
import Data.Text.Short
  ( ShortText )
import qualified Data.Text.Short as ShortText
  ( pack, unpack )

-- fir
import CodeGen.Application
  ( ASTs(NilAST, ConsAST), UCode(UCode)
  , Application(Applied)
  )
import CodeGen.Binary
  ( instruction, structuralInstruction )
import {-# SOURCE #-} CodeGen.CodeGen
  ( CodeGen(codeGenArgs), codeGen )
import CodeGen.Debug
  ( whenAsserting )
import CodeGen.IDs
  ( typeID, constID )
import CodeGen.Instruction
  ( Args(Arg, EndArgs), toArgs
  , ID(ID), Instruction(..)
  , Pair(Pair)
  )
import CodeGen.Monad
  ( MonadFresh(fresh)
  , CGMonad
  )
import CodeGen.Phi
  ( phiInstruction, phiInstructions )
import CodeGen.State
  ( LexicalEnv(localBindings)
  , EarlyExit(..), LoopExit(..), LoopBlockIDs(..)
  , _localBindings, _localBinding, _writableBindings
  , _spirvVersion
  , scoped
  , withLoopFrame, inLoopBodyScope, innermostLoop, loopDepth, addLoopExit
  , currentBlockID, enterBlock
  )
import Data.Containers.Traversals
  ( traverseWithKey_ )
import Data.Type.Known
  ( knownValue )
import Data.Type.Map
  ( Keys )
import FIR.AST
  ( AST, Code
  , SelectionF(..), LoopF(..), StateF(..)
  , pattern (:$), pattern Return
  )
import FIR.AST.Type
  ( AugType(Val), Eff, Nullary )
import FIR.Prim.Types
  ( IntegralTy, primTy )
import FIR.ProgramState
  ( Bindings )
import qualified SPIRV.Control   as SPIRV
  ( SelectionControl, pattern NoSelectionControl
  , LoopControl, pattern NoLoopControl
  )
import qualified SPIRV.Operation as SPIRV.Op
import qualified SPIRV.PrimTy    as SPIRV
  ( PrimTy(..) )
import qualified SPIRV.Version as SPIRV
  ( Version(..) )

----------------------------------------------------------------------------
-- blocks and branching

-- | Start a new block with the given ID.
block :: ID -> CGMonad ()
block blockID = do
  structuralInstruction
    Instruction
      { operation = SPIRV.Op.Label
      , resTy = Nothing
      , resID = Just blockID
      , args  = EndArgs
      }
  enterBlock blockID

newBlock :: CGMonad ()
newBlock = fresh >>= block

-- | Whether the code currently being generated is reachable.
inReachableCode :: CGMonad Bool
inReachableCode = isJust <$> currentBlockID

branch :: ID -> CGMonad ()
branch branchID
  = instruction
      Instruction
       { operation = SPIRV.Op.Branch
       , resID = Nothing
       , resTy = Nothing
       , args  = Arg branchID EndArgs
       }

branchConditional :: ID -> ID -> ID -> CGMonad ()
branchConditional b t f
  = instruction
      Instruction
        { operation = SPIRV.Op.BranchConditional
        , resTy = Nothing
        , resID = Nothing
        , args  = Arg b
                $ Arg t
                $ Arg f EndArgs
        }

unreachable :: CGMonad ()
unreachable
  = instruction
      Instruction
        { operation = SPIRV.Op.Unreachable
        , resTy = Nothing
        , resID = Nothing
        , args  = EndArgs
        }

select :: ID -> (ID, SPIRV.PrimTy) -> (ID, SPIRV.PrimTy) -> CGMonad (ID, SPIRV.PrimTy)
select c (o1, ty1) (o2, ty2) = do
  ( whenAsserting . when (ty1 /= ty2) )
    ( throwError
      ( "select: objects of different types: \n\
        \object 1: " <> ShortText.pack (show ty1) <> "\n\
        \object 2: " <> ShortText.pack (show ty2) <> "."
      )
    )
  v <- fresh
  resTyID <- typeID ty1
  instruction
    Instruction
      { operation = SPIRV.Op.Select
      , resTy     = Just resTyID
      , resID     = Just v
      , args      = Arg c
                  $ Arg o1
                  $ Arg o2
                  EndArgs
      }
  pure (v, ty1)

selectionMerge :: ID -> Maybe SPIRV.SelectionControl -> CGMonad ()
selectionMerge mergeBlockID mbControl
  = instruction
      Instruction
        { operation = SPIRV.Op.SelectionMerge
        , resTy     = Nothing
        , resID     = Nothing
        , args      = Arg mergeBlockID
                    $ Arg mbControl
                    EndArgs
        }
loopMerge :: ID -> ID -> SPIRV.LoopControl Word32 -> CGMonad ()
loopMerge mergeBlockID loopBlockID loopControl
  = instruction
      Instruction
         { operation = SPIRV.Op.LoopMerge
         , resTy = Nothing
         , resID = Nothing
         , args  = Arg mergeBlockID
                 $ Arg loopBlockID
                 $ Arg loopControl
                 EndArgs
         }

multiWaySwitch :: IntegralTy s => ID -> ID -> [Pair s ID] -> CGMonad ()
multiWaySwitch scrut defBlock ids
  = instruction
      Instruction
        { operation = SPIRV.Op.Switch
        , resTy     = Nothing
        , resID     = Nothing
        , args      = Arg scrut
                    $ Arg defBlock
                    $ toArgs ids
        }

-- | The values of the bindings (satisfying the predicate) at the start of a block,
-- given the bindings at the end of each of its predecessors.
-- Emits ϕ-instructions as necessary.
mergeBindings
  :: ( ShortText -> Bool )
  -> [ ( ID, Map ShortText (ID, SPIRV.PrimTy) ) ]
  -> CGMonad ( Map ShortText (ID, SPIRV.PrimTy) )
mergeBindings isRelevant preds = case preds of
  [] -> pure Map.empty
  [ ( _, bds ) ] -> pure ( Map.filterWithKey ( \ name _ -> isRelevant name ) bds )
  _ -> do
    phis <- phiInstructions isRelevant ( map fst preds ) ( map snd preds )
    let
      common :: Map ShortText (ID, SPIRV.PrimTy)
      common = Map.filterWithKey ( \ name _ -> isRelevant name )
             $ foldr1 Map.intersection ( map snd preds )
    pure ( Map.mapWithKey ( \ name ( ident, ty ) -> ( Map.findWithDefault ident name phis, ty ) ) common )

-- | Update the values of the given bindings.
updateBindings :: Map ShortText (ID, SPIRV.PrimTy) -> CGMonad ()
updateBindings = traverseWithKey_ \ name bd -> assign ( _localBinding name ) ( Just bd )

----------------------------------------------------------------------------
-- code generation for selections (pure or branching)

instance CodeGen AST => CodeGen (SelectionF AST) where

  -- if statement
  codeGenArgs (Applied IfF (c `ConsAST` (t :: Code a) `ConsAST ` f `ConsAST` NilAST)) = do
    ver <- view _spirvVersion
    if canUseSelection (primTy @a) ver
    then selection c        t         f
    else ifM       c (UCode t) (UCode f)
      where
        canUseSelection :: SPIRV.PrimTy -> SPIRV.Version -> Bool
        canUseSelection (SPIRV.Scalar    _) _ = True
        canUseSelection (SPIRV.Pointer _) _ = True
        canUseSelection ty ver
          | ver < SPIRV.Version 1 4 = False
          | SPIRV.Matrix {} <- ty   = True
          | SPIRV.Vector {} <- ty   = True
          | SPIRV.Array  {} <- ty   = True
          | SPIRV.Struct {} <- ty   = True
          | otherwise               = False

  -- monadic if statement
  codeGenArgs (Applied IfMF (c `ConsAST` t `ConsAST ` f `ConsAST` NilAST)) = ifM c (UCode t) (UCode f)

  -- switch statement
  codeGenArgs (Applied (SwitchF scrut def cases) NilAST) =
    codeGen
      ( SwitchMF
       ( Return :$ scrut )
       ( Return :$ def   )
       ( map (second (Return :$)) cases )
      )

  -- monadic switch statement
  codeGenArgs (Applied (SwitchMF scrut def cases) NilAST) = switch scrut def (map (second UCode) cases)


selection :: ( CodeGen AST, Nullary a ) => Code Bool -> AST a -> AST a -> CGMonad (ID, SPIRV.PrimTy)
selection cond x y = do
  (condID, condTy) <- codeGen cond
  ( whenAsserting . unless ( condTy == SPIRV.Boolean) )
    ( throwError
    $  "codeGen: ASSERT failed, 'select' expected boolean conditional, but got "
    <> ShortText.pack (show condTy)
    )
  o1 <- codeGen x
  o2 <- codeGen y
  select condID o1 o2

ifM :: ( CodeGen AST, Nullary b ) => AST b -> UCode -> UCode -> CGMonad (ID, SPIRV.PrimTy)
ifM cond bodyTrue bodyFalse
  = branchingSelection conditionalHeader [bodyTrue, bodyFalse]
      where
        conditionalHeader :: [ID] -> ID -> CGMonad ()
        conditionalHeader [trueBlockID, falseBlockID] mergeBlockID = do
          (condID, condTy) <- locally (codeGen cond)
          ( whenAsserting . unless ( condTy == SPIRV.Boolean ) )
            ( throwError
            $  "codeGen: ASSERT failed, 'if' expected boolean conditional, but got "
            <> ShortText.pack (show condTy)
            )
          selectionMerge mergeBlockID SPIRV.NoSelectionControl
          branchConditional condID trueBlockID falseBlockID
        conditionalHeader ids _ =
          throwError
            ( "codeGen: if statement with "
            <> ShortText.pack ( show ( length ids ))
            <> " branches"
            )

switch :: ( CodeGen AST, IntegralTy t, Nullary s, Nullary a)
       => AST s -> AST a -> [(t, UCode)] -> CGMonad (ID, SPIRV.PrimTy)
switch scrut def cases
  = branchingSelection switchHeader (UCode def : map snd cases)
      where
        switchHeader :: [ID] -> ID -> CGMonad ()
        switchHeader [] _ = whenAsserting (throwError "codeGen: ASSERT failed, 'switch' statement missing default case")
        switchHeader ( defaultBlockID : caseBlockIDs ) mergeBlockID = do
          (scrutID, _) <- locally (codeGen scrut)
          selectionMerge mergeBlockID SPIRV.NoSelectionControl
          multiWaySwitch scrutID defaultBlockID
            (zipWith ( \(t,_) i -> Pair (t,i) ) cases caseBlockIDs)

branchingSelection :: CodeGen AST => ( [ID] -> ID -> CGMonad () ) -> [ UCode ] -> CGMonad (ID, SPIRV.PrimTy)
branchingSelection mkHeader cases = do
  live           <- inReachableCode
  headerBlockID  <- fresh
  caseBlockIDs   <- traverse (\x -> (,x) <$> fresh) cases
  mergeBlockID   <- fresh
  bindingsBefore <- use _localBindings

  -- header block
  when live do
    branch headerBlockID
    block  headerBlockID
  mkHeader ( map fst caseBlockIDs ) mergeBlockID

  -- code generation for each of the cases
  results <-
    for caseBlockIDs \ (caseBlockID, UCode caseBody) -> do
      ( ( bodyRes, mbEndBlock ), env ) <- scoped do
        when live ( block caseBlockID )
        bodyRes    <- codeGen caseBody
        mbEndBlock <- currentBlockID
        branch mergeBlockID
        pure ( bodyRes, mbEndBlock )
      pure ( bodyRes, ( , localBindings env ) <$> mbEndBlock )

  ( firstRes, resType ) <- case results of
    [] -> throwError "codeGen: empty list of branches in conditional"
    ( ( firstRes@( _, ty ), _ ) : others )
      -> do ( whenAsserting . unless ( all ( ( == ty ) . snd . fst ) others ) )
              ( throwError "codeGen: ASSERT failed, branches of conditional return different types" )
            pure ( firstRes, ty )
  let
    -- the cases which branch to the merge block
    preds :: [ ( ID, ( (ID, SPIRV.PrimTy), Map ShortText (ID, SPIRV.PrimTy) ) ) ]
    preds = [ ( endBlock, ( bodyRes, bds ) ) | ( bodyRes, Just ( endBlock, bds ) ) <- results ]

  if not live
  then pure firstRes -- (unreachable code)
  else do
    -- merge block
    block mergeBlockID
    case preds of
      [] -> do
        -- No case reaches the merge block: the code that follows is unreachable.
        unreachable
        pure firstRes
      _  -> do
        -- emit the appropriate ϕ-instructions, and update the bindings
        updateBindings
          =<< mergeBindings ( `Map.member` bindingsBefore ) [ ( b, bds ) | ( b, ( _, bds ) ) <- preds ]
        -- return the result
        if resType == SPIRV.Unit
        then pure (ID 0, resType) -- ID should never be used
        else do
          res <- fresh
          phiInstruction
            ( res, resType )
            [ Pair ( resID, b ) | ( b, ( ( resID, _ ), _ ) ) <- preds ]
          pure (res, resType)

----------------------------------------------------------------------------
-- code generation for loops

instance CodeGen AST => CodeGen (LoopF AST) where
  codeGenArgs ( Applied WhileF ( c `ConsAST` body `ConsAST` NilAST ) ) = while ( Just c ) body
  codeGenArgs ( Applied LoopF              ( body `ConsAST` NilAST ) ) = while ( Nothing :: Maybe ( AST ( Val Bool ) ) ) body
  codeGenArgs ( Applied ( BreakF         ( _ :: Proxy n ) ) NilAST   ) = break ( knownValue @n )
  codeGenArgs ( Applied ( BreakContinueF ( _ :: Proxy n ) ) NilAST   ) = breakContinue ( knownValue @n )

-- | Code generation for a loop (with an optional condition in the loop header).
--
-- Every modifiable binding gets a ϕ-instruction in the loop header. Its operand for the
-- back edge is only known after generating code for the loop body, which refers to the
-- ϕ-instruction: we tie the knot lazily.
while :: ( CodeGen AST, Nullary c, Nullary b ) => Maybe ( AST c ) -> AST b -> CGMonad (ID, SPIRV.PrimTy)
while mbCond loopBody = do
  mbBefore <- currentBlockID
  case mbBefore of
    Nothing            -> unreachableLoop
    Just beforeBlockID -> reachableLoop beforeBlockID
  pure (ID 0, SPIRV.Unit) -- ID should never be used

  where

    -- (generates code only to obtain the results; nothing is emitted)
    unreachableLoop :: CGMonad ()
    unreachableLoop = do
      loopIDs <- LoopBlockIDs <$> fresh <*> fresh
      _ <- scoped $ withLoopFrame loopIDs do
        for_ mbCond ( locally . codeGen )
        inLoopBodyScope ( codeGen loopBody )
      pure ()

    reachableLoop :: ID -> CGMonad ()
    reachableLoop beforeBlockID = mdo
      headerBlockID   <- fresh
      loopBlockID     <- fresh
      mergeBlockID    <- fresh -- block where control flow merges back
      continueBlockID <- fresh -- block at the end of the loop body from which we can continue for further loops
      branch headerBlockID

      -- Header block.
      block headerBlockID
      bindingsBefore <- use _localBindings
      writable       <- use _writableBindings
      let
        -- Bindings that the loop might modify.
        loopCarried :: Map ShortText (ID, SPIRV.PrimTy)
        loopCarried = bindingsBefore `Map.restrictKeys` writable

        -- value of a binding at the end of the continue block (see below)
        continueEndID :: ShortText -> ID
        continueEndID name =
          maybe
            ( error $ "codeGen: binding " <> ShortText.unpack name <> " missing at the end of a loop" )
            fst
            ( Map.lookup name continueEnd )
      headerPhis <-
        Map.traverseWithKey
          ( \ name ( beforeID, ty ) -> do
            v <- fresh
            -- NB: the second operand must not be forced during code generation.
            phiInstruction ( v, ty )
              [ Pair ( beforeID          , beforeBlockID   )
              , Pair ( continueEndID name, continueBlockID )
              ]
            pure ( v, ty )
          )
          loopCarried
      updateBindings headerPhis

      loopMerge mergeBlockID continueBlockID SPIRV.NoLoopControl
      -- The loop frame collects the early exits from the condition and the body.
      ( ( headerEndBindings, mbHeaderEndBlockID, mbCondID, mbBodyEndBlock, bodyEndBindings ), loopExits ) <-
        withLoopFrame ( LoopBlockIDs { continueBlockID, mergeBlockID } ) do

          -- Condition.
          mbCondition <- case mbCond of
            Nothing -> do
              branch loopBlockID
              pure Nothing
            Just cond -> do
              -- The condition gets its own block, as it might contain control flow.
              condBlockID <- fresh
              branch condBlockID
              block condBlockID
              (condID, condTy) <- locally (codeGen cond)
              ( whenAsserting . when ( condTy /= SPIRV.Boolean ) )
                ( throwError
                $  "codeGen: ASSERT failed, while loop expected boolean conditional, but got "
                <> ShortText.pack (show condTy)
                )
              pure ( Just condID )
          condEndBindings <- use _localBindings
          mbCondEndBlock  <- currentBlockID
          for_ mbCondition \ condID -> branchConditional condID loopBlockID mergeBlockID

          -- Loop body.
          ( mbEndBlock, env ) <- inLoopBodyScope $ scoped do
            block loopBlockID
            _ <- codeGen loopBody
            mbEndBlock <- currentBlockID
            branch continueBlockID
            pure mbEndBlock
          pure ( condEndBindings, mbCondEndBlock, mbCondition, mbEndBlock, localBindings env )

      -- Continue block.
      block continueBlockID
      let
        continuePreds :: [ ( ID, Map ShortText (ID, SPIRV.PrimTy) ) ]
        continuePreds
          =  [ ( b, bodyEndBindings ) | Just b <- [ mbBodyEndBlock ] ]
          ++ [ ( b, bds ) | ( b, ( ContinueLoop, bds ) ) <- Map.toList loopExits ]
      for_ continuePreds \ ( _, bds ) ->
        unless ( Map.keysSet loopCarried `Set.isSubsetOf` Map.keysSet bds ) $
          throwError "codeGen: early exit from a scope in which bindings modified by the loop are not visible"
      continueEnd <- case continuePreds of
        -- unreachable continue block: use the ϕ-instructions themselves
        [] -> pure headerPhis
        _  -> mergeBindings ( `Map.member` loopCarried ) continuePreds
      branch headerBlockID

      -- Merge block (first block after the loop).
      block mergeBlockID
      let
        -- with the early exit that remains to be performed (if any)
        mergePreds :: [ ( ID, Maybe EarlyExit, Map ShortText (ID, SPIRV.PrimTy) ) ]
        mergePreds
          =  [ ( b, Nothing, headerEndBindings ) | isJust mbCondID, Just b <- [ mbHeaderEndBlockID ] ]
          ++ [ ( b, mbEarlyExit, bds ) | ( b, ( ExitLoop mbEarlyExit, bds ) ) <- Map.toList loopExits ]
      case mergePreds of
        -- Nothing exits the loop: the code that follows is unreachable.
        [] -> unreachable
        _  -> do
          updateBindings
            =<< mergeBindings ( `Map.member` bindingsBefore ) [ ( b, bds ) | ( b, _, bds ) <- mergePreds ]
          continueEarlyExits [ ( b, mbEarlyExit ) | ( b, mbEarlyExit, _ ) <- mergePreds ]

-- | Perform the remaining early exits from the merge block of a loop,
-- given the predecessors of the merge block and the early exit that remains for each.
--
-- Emits ϕ-instructions: must be called at the start of the merge block.
continueEarlyExits :: [ ( ID, Maybe EarlyExit ) ] -> CGMonad ()
continueEarlyExits preds = do
  let
    -- Selector value for each early exit to perform (0 for none).
    selectors :: Map EarlyExit Word32
    selectors = Map.fromList $ zip ( Set.toList . Set.fromList $ [ ee | ( _, Just ee ) <- preds ] ) [ 1 .. ]
  unless ( null selectors ) do
    selectorID <- fresh
    selectorValues <-
      traverse
        ( \ ( b, mbEarlyExit ) -> ( , b ) <$> constID ( maybe 0 ( selectors Map.! ) mbEarlyExit ) )
        preds
    phiInstruction ( selectorID, primTy @Word32 ) ( map Pair selectorValues )
    caseBlockIDs <- traverse ( const fresh ) selectors
    afterBlockID <- fresh
    selectionMerge afterBlockID SPIRV.NoSelectionControl
    multiWaySwitch selectorID afterBlockID
      ( Map.elems $ Map.intersectionWith ( \ i caseBlockID -> Pair ( i, caseBlockID ) ) selectors caseBlockIDs )
    traverseWithKey_ ( \ ee caseBlockID -> block caseBlockID *> earlyExit ee ) caseBlockIDs
    block afterBlockID

-- | Break out of the given number of loops.
break :: Word32 -> CGMonad (ID, SPIRV.PrimTy)
break 0 = pure (ID 0, SPIRV.Unit) -- ID should never be used
break nbLoopsToBreak = do
  depth <- loopDepth
  when ( depth < fromIntegral nbLoopsToBreak ) $
    throwError
      ( "codeGen: 'break' statement inside too few loops.\n\
        \breaking " <> ShortText.pack ( show nbLoopsToBreak ) <> " loops,\n\
        \but only inside " <> ShortText.pack ( show depth ) <> " loops."
      )
  earlyExit ( Break nbLoopsToBreak )
  pure (ID 0, SPIRV.Unit) -- ID should never be used

-- | Break out of the given number of loops, and then continue to the next iteration
-- of the enclosing loop.
breakContinue :: Word32 -> CGMonad (ID, SPIRV.PrimTy)
breakContinue nbLoopsToBreak = do
  depth <- loopDepth
  when ( depth <= fromIntegral nbLoopsToBreak ) $
    throwError
      ( "codeGen: 'breakContinue' statement inside too few loops.\n\
        \continuing from start of loop after breaking " <> ShortText.pack ( show nbLoopsToBreak ) <> " loops,\n\
        \but only inside " <> ShortText.pack ( show depth ) <> " loops."
      )
  earlyExit ( BreakContinue nbLoopsToBreak )
  pure (ID 0, SPIRV.Unit) -- ID should never be used

-- | Perform an early exit.
--
-- Structured control flow only allows branching to the continue or merge block of the
-- innermost loop: to exit several loops, we exit the innermost loop, and perform the
-- remaining early exit from its merge block. The same goes for an early exit from the
-- condition of a loop, which is lexically outside of the loop.
earlyExit :: EarlyExit -> CGMonad ()
earlyExit ee = do
  mbLoop <- innermostLoop
  case mbLoop of
    Nothing -> throwError "codeGen: early exit outside of a loop"
    Just ( LoopBlockIDs { continueBlockID, mergeBlockID }, inBody ) -> do
      let
        ( loopExit, destBlockID ) = case ee of
          _ | not inBody  -> ( ExitLoop ( Just ee )                         , mergeBlockID    )
          BreakContinue 0 -> ( ContinueLoop                                 , continueBlockID )
          BreakContinue n -> ( ExitLoop ( Just ( BreakContinue ( n - 1 ) ) ), mergeBlockID    )
          Break n
            | n <= 1      -> ( ExitLoop Nothing                             , mergeBlockID    )
            | otherwise   -> ( ExitLoop ( Just ( Break ( n - 1 ) ) )        , mergeBlockID    )
      mbBlock <- currentBlockID
      for_ mbBlock \ currentBlock -> do
        currentBindings <- use _localBindings
        addLoopExit currentBlock loopExit currentBindings
        branch destBlockID

----------------------------------------------------------------------------
-- code-generation for locally / embed

instance CodeGen AST => CodeGen (StateF AST) where
  codeGenArgs (Applied LocallyF (a `ConsAST` NilAST)) =
    locally (codeGen a)
  codeGenArgs (Applied EmbedF   ( ( a :: AST ( Eff i i a ) ) `ConsAST` NilAST)) =
    embed ( Set.fromList $ knownValue @(Keys (Bindings i)) ) (codeGen a)

locally :: CGMonad a -> CGMonad a
locally action = do
  bindingsBefore <- use _localBindings
  ( res, env ) <- scoped action
  let
    bindingsAfter :: Map ShortText (ID, SPIRV.PrimTy)
    bindingsAfter = localBindings env

  -- keep the updated values of local bindings,
  -- but only for bindings that were defined before the "locally" block
  assign _localBindings ( bindingsAfter `Map.intersection` bindingsBefore )

  whenAsserting do
    let
      weirdos :: Map ShortText ( (ID, SPIRV.PrimTy), (ID, SPIRV.PrimTy) )
      weirdos = Map.mapMaybe id
              $ Map.intersectionWith ( \ ( id1, ty1 ) ( id2, ty2 ) -> if ty1 == ty2 then Nothing else Just ( ( id1, ty1 ), ( id2, ty2 ) ) )
                bindingsBefore bindingsAfter
    unless ( null weirdos ) $
      throwError ( "'codeGen': ASSERT failed, error in code-generation for 'locally'\nweirdos = " <> ShortText.pack ( show weirdos ) )

  pure res

embed :: Set ShortText -> CGMonad a -> CGMonad a
embed innerStartNames action = do
  bindingsBefore <- use _localBindings
  ( res, env ) <- scoped do
    assign _localBindings ( Map.restrictKeys bindingsBefore innerStartNames )
    action
  let
    bindingsAfter :: Map ShortText (ID, SPIRV.PrimTy)
    bindingsAfter = localBindings env

    -- discard the updated values of local bindings that shadowed previous bindings
    newBindings :: Map ShortText (ID, SPIRV.PrimTy)
    newBindings =
      Map.mapWithKey
        ( \ name aft -> fromMaybe aft (shadows name innerStartNames bindingsBefore ) )
        (  bindingsAfter `Map.intersection` bindingsBefore
        <> Map.withoutKeys bindingsBefore innerStartNames
        )

  assign _localBindings newBindings
  pure res

shadows :: ShortText -> Set ShortText -> Map ShortText (ID, SPIRV.PrimTy) -> Maybe (ID, SPIRV.PrimTy)
shadows name innerStartNames prevs
  | not ( name `Set.member` innerStartNames )
  , Just prev <- Map.lookup name prevs
  = Just prev
  | otherwise
  = Nothing