State.hs

{-# LANGUAGE BlockArguments     #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleContexts   #-}
{-# LANGUAGE NamedFieldPuns     #-}
{-# LANGUAGE OverloadedStrings  #-}
{-# LANGUAGE RankNTypes         #-}

{-|
Module: CodeGen.State

This module defines the structure of the state used by the code generation,
in two parts:

  - 'CGContext', consisting of reified type-level data provided at the start
  of code generation (with a 'Control.Monad.Reader.ReaderT' transformer),
  - 'CGState', consisting of state accumulated during code generation,
  recording which objects have been declared and their IDs.

This module manually provides lenses for these. Most of these could be derived
automatically, but a few of them are slightly trickier, requiring an affine traversal.

These custom lenses also allow terse state modifications.
For instance, usage of a built-in variable simultaneously obtains an ID for the variable,
adds itself to the relevant entry point interface, and sets the necessary decorations.

-}

module CodeGen.State
  ( -- * Code generation state
    CGState(moduleState), initialState
  , ModuleState(..)
  , FunctionState(functionContext, localVariables)
  , LexicalEnv(..)
  , LoopBlockIDs(..), EarlyExit(..), LoopExit(..)
  , BlockState
  , PointerState(..)

    -- * Code generation context
  , CGContext(..), emptyContext

    -- * Lenses
  , affineTraverse
    -- ** Module state
  , _currentID
  , _neededCapabilities, _neededCapability
  , _neededExtensions, _neededExtension
  , _knownExtInsts, _knownExtInst
  , _knownStringLits, _knownStringLit
  , _names
  , _usedGlobals, _usedGlobal
  , _builtinVariables, _builtinVariable
  , _entryPoints, _entryPoint
  , _interfaces, _interface, _interfaceBinding
  , _builtin
  , _decorations, _decorate
  , _memberDecorations, _memberDecorate
  , _knownTypes, _knownType
  , _knownConstants, _knownConstant
  , _knownUndefineds, _knownUndefined
  , _knownBindings, _knownBinding
    -- ** Lexical environment
  , _localBindings, _localBinding
  , _writableBindings
  , _rayQueries, _rayQuery
  , defineLocalBinding
  , scoped
    -- ** Function state
  , currentFunctionContext
  , withFunctionState
  , addLocalVariable
    -- ** Loops
  , withLoopFrame
  , inLoopBodyScope, innermostLoop, loopDepth
  , addLoopExit
    -- ** Block state
  , currentBlockID
  , enterBlock, leaveBlock
  , lookupTemporaryPointer, setTemporaryPointer
    -- ** Code generation context
  , _userGlobals, _userGlobal
  , _userFunctions, _userFunction
  , _userEntryPoints, _userEntryPoint
  , _userCapabilities, _userExtensions
  , _userImages, _userImage
  , _spirvVersion, _backend
  , _debugging, _asserting, _emittingCode

    -- * Helper state update functions
  , requireCapability, requireCapabilities
  , requireExtension, requireExtensions
  , addName, addMemberName
  , addDecoration, addDecorations
  , addMemberDecoration, addMemberDecorations
  )
  where

-- base
import Data.Maybe
  ( fromMaybe )
import Data.Word
  ( Word32 )

-- containers
import Data.Map
  ( Map )
import qualified Data.Map as Map
import Data.Set
  ( Set )
import qualified Data.Set as Set

-- lens
import Control.Lens
  ( Lens', lens
  , at
  , view, set
  , use, assign, modifying
  )

-- mtl
import Control.Monad.Except
  ( MonadError, throwError )
import Control.Monad.State
  ( MonadState, gets, modify' )

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

-- fir
import CodeGen.Instruction
  ( ID(ID), TyID(tyID)
  , Pointer, Variable, variableID
  , Instruction
  )
import Data.Containers.Traversals
  ( traverseSet_ )
import FIR.Builtin
  ( modelBuiltins, builtinDecorations )
import FIR.Prim.Types
  ( AConstant )
import FIR.ProgramState
  ( FunctionContext(TopLevel), VLFunctionContext ) -- value-level function context
import qualified SPIRV.Capability    as SPIRV
import qualified SPIRV.Control       as SPIRV
import qualified SPIRV.Decoration    as SPIRV
import qualified SPIRV.ExecutionMode as SPIRV
import qualified SPIRV.Extension     as SPIRV
  ( ExtInst, Extension )
import qualified SPIRV.Image         as SPIRV
import qualified SPIRV.Layout        as SPIRV
  ( Layout )
import qualified SPIRV.PrimTy        as SPIRV
import qualified SPIRV.Stage         as SPIRV
import qualified SPIRV.Version       as SPIRV

----------------------------------------------------------------------------
-- * Code generation state and context

-- | Code generation state.
--
-- The state is split up according to the lifetime of the information it contains:
--
--  - 'ModuleState': information about the whole SPIR-V module (types, constants, ...),
--  - 'FunctionState': information about the function currently being generated,
--  - 'LexicalEnv': the bindings in scope (restored at the end of each lexical scope),
--  - 'LoopFrame's: the loops we are currently inside,
--  - 'BlockState': information about the current basic block.
--
-- Each of these is replaced wholesale when entering/leaving the corresponding scope,
-- using the functions 'withFunctionState', 'scoped', 'withLoopFrame', 'enterBlock'
-- and 'leaveBlock' (the constructors of 'FunctionState' and 'BlockState' are not exported).
-- Information that refers to IDs that are only valid in part of a function body
-- should be stored in the narrowest such scope, so that it can't be used where
-- it isn't valid.
--
-- NB: the code generation for loops and functions ties a knot through this state
-- (see 'CodeGen.CFG.while' and 'CodeGen.Functions.inContext'),
-- so the fields of these records must remain lazy.
data CGState
  = CGState
      { moduleState   :: ModuleState
      , functionState :: Maybe FunctionState -- ^ @Nothing@ at the top level.
      , lexicalEnv    :: LexicalEnv
      , loopStack     :: [ LoopFrame ]       -- ^ Innermost loop first.
      , blockState    :: BlockState
      }

-- | Information about the whole SPIR-V module being generated.
data ModuleState
  = ModuleState
      {
      -- | Current ID number (increases by 1 each time a new ID is needed).
        currentID           :: ID

      -- | Capability requirements that have been declared.
      , neededCapabilities  :: Set                                   SPIRV.Capability

      -- | Needed extensions.
      , neededExtensions    :: Set                                   SPIRV.Extension

      -- | IDs of all used extended instruction sets.
      , knownExtInsts       :: Map SPIRV.ExtInst                     ID

      -- | IDs of constant text literals.
      , knownStringLits     :: Map ShortText                         ID

      -- | IDs which have been annotated with a name / member name.
      , names               :: Set                                   (ID, Either ShortText (Word32, ShortText))

      -- | Entry point IDs.
      , entryPoints         :: Map (ShortText, SPIRV.ExecutionModel) ID

      -- | Entry point interfaces, keeping track of which global variables are used.
      , interfaces          :: Map (ShortText, SPIRV.ExecutionModel) (Map ShortText ID)

      -- | Decorations for given types.
      , decorations         :: Map ID                                SPIRV.Decorations

      -- | Decorations for members of a given struct type.
      , memberDecorations   :: Map (TyID, Word32)                    SPIRV.Decorations

      -- | Map of all declared types, together with their layout.
      --
      -- We keep the full 'Instruction' rather than just its 'TyID',
      -- as the types have to be declared at the top of the SPIR-V module,
      -- rather than inline with the rest of the code.
      , knownTypes          :: Map (SPIRV.Layout, SPIRV.PrimTy)      Instruction

      -- | Map of all declared constants.
      --
      -- Same comment than for 'knownTypes' regarding 'Instruction' vs 'ID'.
      , knownConstants      :: Map AConstant                         Instruction

      -- | Map of all types who have a corresponding "Undefined" instruction.
      , knownUndefineds     :: Map SPIRV.PrimTy                      (ID, TyID)

      -- | Which user-defined global variables have been used?
      , usedGlobals         :: Map ShortText                         Variable

      -- | Built-in variables that have been used (each entry point has its own), with their names.
      , builtinVariables    :: Map ID                                (ShortText, Variable)

      -- | Top-level bindings available, such as top-level functions.
      , knownBindings       :: Map ShortText                         (ID, SPIRV.PrimTy)
      }

-- | Information about the function currently being generated.
data FunctionState
  = FunctionState
      { -- | Whether we are in a function or an entry point.
        functionContext :: VLFunctionContext

        -- | Local variables of the function.
        --
        -- These are declared at the start of the function body.
      , localVariables  :: Map ID Variable
      }

-- | The bindings in scope.
data LexicalEnv
  = LexicalEnv
      { -- | Variables declared by the user in the program.
        localBindings    :: Map ShortText (ID, SPIRV.PrimTy)

        -- | Which of the local bindings can be modified.
        --
        -- This doesn't include bindings that are pointers (e.g. a mutable array).
      , writableBindings :: Set ShortText

        -- | IDs of locally declared ray query variables.
      , rayQueries       :: Map ShortText ID
      }

-- | Information about a loop we are currently inside of.
data LoopFrame
  = LoopFrame
      { loopBlockIDs :: LoopBlockIDs

        -- | Whether we are in the body of this loop (as opposed to its condition).
      , inLoopBody   :: Bool

        -- | The early exits which branch to the continue or merge block of this
        -- loop, keyed by the ID of the block in which the early exit occurred,
        -- together with the values of the bindings at that point.
      , loopExits    :: Map ID ( LoopExit, Map ShortText (ID, SPIRV.PrimTy) )
      }

-- | Information about the current basic block.
data BlockState
  = BlockState
      { -- | ID of the current block, or @Nothing@ if there is no current block.
        currentBlock      :: Maybe ID

        -- | Temporary pointers holding copies of values, keyed by the ID of the value.
      , temporaryPointers :: Map ID (Pointer, PointerState)
      }

data PointerState
  = Fresh
  | Modified
  deriving stock ( Eq, Show )

data LoopBlockIDs
  = LoopBlockIDs
    { continueBlockID :: !ID
    , mergeBlockID    :: !ID
    }
    deriving stock ( Show, Eq, Ord )

-- | An early exit from loops ('break' or 'breakContinue').
data EarlyExit
  -- | Break out of the given number of loops (at least 1).
  = Break Word32
  -- | Break out of the given number of loops, and then continue
  -- to the next iteration of the enclosing loop.
  | BreakContinue Word32
  deriving stock ( Show, Eq, Ord )

-- | How an early exit leaves the innermost loop.
--
-- In structured control flow, it is only possible to branch to the continue or merge block
-- of the innermost loop. Exiting several loops at once thus amounts to exiting the innermost
-- loop, and then performing the remaining early exit from its merge block.
data LoopExit
  -- | Continue to the next iteration of the loop.
  = ContinueLoop
  -- | Exit the loop, and then perform the given early exit from the enclosing loops (if any).
  | ExitLoop ( Maybe EarlyExit )
  deriving stock ( Show, Eq, Ord )

initialState :: CGContext -> CGState
initialState CGContext { userCapabilities, userExtensions }
  = CGState
      { moduleState   =
          ModuleState
            { currentID           = ID 1
            , neededCapabilities  = userCapabilities
            , neededExtensions    = userExtensions
            , knownExtInsts       = Map.empty
            , knownStringLits     = Map.empty
            , names               = Set.empty
            , entryPoints         = Map.empty
            , interfaces          = Map.empty
            , decorations         = Map.empty
            , memberDecorations   = Map.empty
            , knownTypes          = Map.empty
            , knownConstants      = Map.empty
            , knownUndefineds     = Map.empty
            , usedGlobals         = Map.empty
            , builtinVariables    = Map.empty
            , knownBindings       = Map.empty
            }
      , functionState = Nothing
      , lexicalEnv    =
          LexicalEnv
            { localBindings    = Map.empty
            , writableBindings = Set.empty
            , rayQueries       = Map.empty
            }
      , loopStack     = []
      , blockState    = noBlock
      }

-- | Block state when there is no current block.
noBlock :: BlockState
noBlock = BlockState { currentBlock = Nothing, temporaryPointers = Map.empty }

-- | Code generation context.
--
-- Consists of information provided at the start of code generation,
-- as provided by the user-provided type annotations of
-- top-level functions, entry points and input/output variables.
data CGContext
  = CGContext
     { -- | User defined inputs/outputs (not builtins).
       userGlobals
          :: Map ShortText (SPIRV.PointerTy, SPIRV.Decorations)

       -- | User defined functions (not entry points).
     , userFunctions
          :: Map ShortText SPIRV.FunctionControl

       -- | User defined entry points.
     , userEntryPoints
          :: Map (ShortText, SPIRV.ExecutionModel) SPIRV.ExecutionModes

       -- | Capabilities that are required (computed from user definitions).
     , userCapabilities
          :: Set                                   SPIRV.Capability

       -- | Extensions that are required (computed from user definitions).
     , userExtensions
          :: Set                                   SPIRV.Extension

       -- | User defined images.
     , userImages
          :: Map ShortText SPIRV.Image

       -- | SPIR-V version to use.
     , spirvVersion :: SPIRV.Version

       -- | SPIR-V backend to use (Vulkan or OpenCL).
     , backend :: SPIRV.Backend

       -- | Whether to add extra source information in the generated SPIR-V assembly.
     , debugging :: Bool

       -- | Whether to enable assertions.
     , asserting :: Bool

       -- | Whether to emit assembly code.
     , emittingCode :: Bool
     }
  deriving stock Show

emptyContext :: CGContext
emptyContext
  = CGContext
      { userGlobals      = Map.empty
      , userFunctions    = Map.empty
      , userEntryPoints  = Map.empty
      , userCapabilities = Set.empty
      , userExtensions   = Set.empty
      , userImages       = Map.empty
      , spirvVersion     = SPIRV.Version 1 5
      , backend          = SPIRV.Vulkan
      , debugging        = False
      , asserting        = False
      , emittingCode     = True
      }

----------------------------------------------------------------------------
-- * Lenses

-- | Affine traversal relative to 'Maybe'.
--
-- Useful function to deal with nested data structures
-- such as @Map a (Map b c)@
affineTraverse :: (Monoid a, Functor f) => (a -> f b) -> (Maybe a -> f (Maybe b))
affineTraverse f Nothing  = fmap Just (f mempty)
affineTraverse f (Just a) = fmap Just (f a)

----------------------------------------------------------------------------
-- ** Module state

-- Internal
_moduleState :: Lens' CGState ModuleState
_moduleState = lens moduleState ( \s v -> s { moduleState = v } )

_currentID :: Lens' CGState ID
_currentID = _moduleState . lens currentID ( \s v -> s { currentID = v } )

_neededCapabilities :: Lens' CGState (Set SPIRV.Capability)
_neededCapabilities = _moduleState . lens neededCapabilities ( \s v -> s { neededCapabilities = v } )

_neededCapability :: SPIRV.Capability -> Lens' CGState (Maybe ())
_neededCapability capability = _neededCapabilities . at capability

_neededExtensions :: Lens' CGState (Set SPIRV.Extension)
_neededExtensions = _moduleState . lens neededExtensions ( \s v -> s { neededExtensions = v } )

_neededExtension :: SPIRV.Extension -> Lens' CGState (Maybe ())
_neededExtension extension = _neededExtensions . at extension

_knownExtInsts :: Lens' CGState (Map SPIRV.ExtInst ID)
_knownExtInsts = _moduleState . lens knownExtInsts ( \s v -> s { knownExtInsts = v } )

_knownExtInst :: SPIRV.ExtInst -> Lens' CGState (Maybe ID)
_knownExtInst ext = _knownExtInsts . at ext

_knownStringLits :: Lens' CGState (Map ShortText ID)
_knownStringLits = _moduleState . lens knownStringLits ( \s v -> s { knownStringLits = v } )

_knownStringLit :: ShortText -> Lens' CGState (Maybe ID)
_knownStringLit lit = _knownStringLits . at lit

_names :: Lens' CGState ( Set (ID, Either ShortText (Word32, ShortText)) )
_names = _moduleState . lens names ( \s v -> s { names = v } )

_usedGlobals :: Lens' CGState (Map ShortText Variable)
_usedGlobals = _moduleState . lens usedGlobals ( \s v -> s { usedGlobals = v } )

_usedGlobal :: ShortText -> Lens' CGState (Maybe Variable)
_usedGlobal name = _usedGlobals . at name

_builtinVariables :: Lens' CGState (Map ID (ShortText, Variable))
_builtinVariables = _moduleState . lens builtinVariables ( \s v -> s { builtinVariables = v } )

_builtinVariable :: ID -> Lens' CGState (Maybe (ShortText, Variable))
_builtinVariable v = _builtinVariables . at v

_entryPoints :: Lens' CGState (Map (ShortText, SPIRV.ExecutionModel) ID)
_entryPoints = _moduleState . lens entryPoints ( \s v -> s { entryPoints = v } )

_entryPoint :: ShortText -> SPIRV.ExecutionModel -> Lens' CGState (Maybe ID)
_entryPoint name stage = _entryPoints . at (name, stage)

_interfaces :: Lens' CGState (Map (ShortText, SPIRV.ExecutionModel) (Map ShortText ID))
_interfaces = _moduleState . lens interfaces ( \s v -> s { interfaces = v } )

_interface :: ShortText -> SPIRV.ExecutionModel -> Lens' CGState (Maybe (Map ShortText ID))
_interface stageName stage = _interfaces . at (stageName, stage)

_interfaceBinding :: ShortText -> SPIRV.ExecutionModel -> ShortText -> Lens' CGState (Maybe ID)
_interfaceBinding stageName stage varName
  = _interface stageName stage
  . affineTraverse
  . at varName

-- | The ID of a built-in variable, for a given entry point.
--
-- Setting this adds the built-in variable to the interface of the entry point,
-- and sets the necessary decorations.
_builtin :: ShortText -> SPIRV.ExecutionInfo Word32 stage -> ShortText -> Lens' CGState (Maybe ID)
_builtin stageName executionInfo builtinName
  = lens
      ( view _interfaceBuiltin )
      ( \s mb_i -> case mb_i of
         Nothing -> s
         Just i  -> set _interfaceBuiltin ( Just i )
                  . set ( _decorate i ) ( Just decs )
                  $ s
      )
  where executionModel :: SPIRV.ExecutionModel
        executionModel = SPIRV.modelOf executionInfo

        _interfaceBuiltin :: Lens' CGState (Maybe ID)
        _interfaceBuiltin = _interfaceBinding stageName executionModel builtinName

        decs :: SPIRV.Decorations
        decs = builtinDecorations executionModel builtinName builtinTy

        builtinTy :: SPIRV.PointerTy
        builtinTy =
          fromMaybe
            ( error
              ( "_builtin: builtin with name " ++ ShortText.unpack builtinName ++ " cannot be found,\n\
                  \among builtins for " ++ show executionModel ++ " named " ++ ShortText.unpack stageName
              )
            )
            ( lookup builtinName $ modelBuiltins executionInfo )


_decorations :: Lens' CGState (Map ID SPIRV.Decorations)
_decorations = _moduleState . lens decorations ( \s v -> s { decorations = v } )

_decorate :: ID -> Lens' CGState (Maybe SPIRV.Decorations)
_decorate bindingID = _decorations . at bindingID

_memberDecorations :: Lens' CGState ( Map (TyID, Word32) SPIRV.Decorations )
_memberDecorations = _moduleState . lens memberDecorations ( \s v -> s { memberDecorations = v } )

_memberDecorate :: TyID -> Word32 -> Lens' CGState ( Maybe SPIRV.Decorations )
_memberDecorate bindingTyID index
  = _memberDecorations
  . at (bindingTyID, index)

_knownTypes :: Lens' CGState (Map (SPIRV.Layout, SPIRV.PrimTy) Instruction)
_knownTypes = _moduleState . lens knownTypes ( \s v -> s { knownTypes = v } )

_knownType :: SPIRV.Layout -> SPIRV.PrimTy -> Lens' CGState (Maybe Instruction)
_knownType lay primTy = _knownTypes . at (lay, primTy)

_knownConstants :: Lens' CGState (Map AConstant Instruction)
_knownConstants = _moduleState . lens knownConstants ( \s v -> s { knownConstants = v } )

_knownConstant :: AConstant -> Lens' CGState (Maybe Instruction)
_knownConstant constant = _knownConstants . at constant

_knownUndefineds :: Lens' CGState (Map SPIRV.PrimTy (ID, TyID))
_knownUndefineds = _moduleState . lens knownUndefineds ( \s v -> s { knownUndefineds = v } )

_knownUndefined :: SPIRV.PrimTy -> Lens' CGState (Maybe (ID,TyID))
_knownUndefined primTy = _knownUndefineds . at primTy

_knownBindings :: Lens' CGState (Map ShortText (ID, SPIRV.PrimTy))
_knownBindings = _moduleState . lens knownBindings ( \s v -> s { knownBindings = v } )

_knownBinding :: ShortText -> Lens' CGState (Maybe (ID, SPIRV.PrimTy))
_knownBinding binding = _knownBindings . at binding

----------------------------------------------------------------------------
-- ** Lexical environment

-- Internal
_lexicalEnv :: Lens' CGState LexicalEnv
_lexicalEnv = lens lexicalEnv ( \s v -> s { lexicalEnv = v } )

_localBindings :: Lens' CGState (Map ShortText (ID, SPIRV.PrimTy))
_localBindings = _lexicalEnv . lens localBindings ( \s v -> s { localBindings = v } )

_localBinding :: ShortText -> Lens' CGState (Maybe (ID, SPIRV.PrimTy))
_localBinding binding = _localBindings . at binding

_writableBindings :: Lens' CGState (Set ShortText)
_writableBindings = _lexicalEnv . lens writableBindings ( \s v -> s { writableBindings = v } )

_rayQueries :: Lens' CGState (Map ShortText ID)
_rayQueries = _lexicalEnv . lens rayQueries ( \s v -> s { rayQueries = v } )

_rayQuery :: ShortText -> Lens' CGState (Maybe ID)
_rayQuery v = _rayQueries . at v

-- | Define a local binding, recording whether it can be modified.
defineLocalBinding :: MonadState CGState m => ShortText -> (ID, SPIRV.PrimTy) -> Bool -> m ()
defineLocalBinding name bd writable = do
  assign ( _localBinding name ) ( Just bd )
  modifying _writableBindings ( if writable then Set.insert name else Set.delete name )

-- | Run an action in a new lexical scope; the previous lexical environment is
-- restored afterwards.
--
-- Returns the lexical environment at the end of the scope
scoped :: MonadState CGState m => m a -> m ( a, LexicalEnv )
scoped action = do
  envBefore <- use _lexicalEnv
  a         <- action
  envAfter  <- use _lexicalEnv
  assign _lexicalEnv envBefore
  pure ( a, envAfter )

----------------------------------------------------------------------------
-- ** Function state

-- | The current function context.
currentFunctionContext :: MonadState CGState m => m VLFunctionContext
currentFunctionContext = gets ( maybe TopLevel functionContext . functionState )

-- | Run an action with a fresh function state, returning the function state
-- at the end.
withFunctionState :: MonadState CGState m => VLFunctionContext -> m a -> m ( a, FunctionState )
withFunctionState context action = do
  before <- gets functionState
  modify' \ s -> s { functionState = Just ( FunctionState { functionContext = context, localVariables = Map.empty } ) }
  a <- action
  after <- gets functionState
  modify' \ s -> s { functionState = before }
  pure ( a, fromMaybe ( error "withFunctionState: function state disappeared" ) after )

-- | Record a local variable of the current function.
addLocalVariable :: ( MonadState CGState m, MonadError ShortText m ) => Variable -> m ()
addLocalVariable var = do
  mbFunState <- gets functionState
  case mbFunState of
    Nothing ->
      throwError "codeGen: cannot declare a local variable outside of a function"
    Just funState ->
      modify' \ s -> s { functionState = Just funState { localVariables = Map.insert ( variableID var ) var ( localVariables funState ) } }

----------------------------------------------------------------------------
-- ** Loops

-- | Run an action (code generation for a loop) inside a new loop frame,
-- returning the early exits that branch to the continue or merge block of this loop.
withLoopFrame
  :: MonadState CGState m
  => LoopBlockIDs -> m a -> m ( a, Map ID ( LoopExit, Map ShortText (ID, SPIRV.PrimTy) ) )
withLoopFrame ids action = do
  modify' \ s -> s { loopStack = LoopFrame { loopBlockIDs = ids, inLoopBody = False, loopExits = Map.empty } : loopStack s }
  a <- action
  frames <- gets loopStack
  case frames of
    frame : outer -> do
      modify' \ s -> s { loopStack = outer }
      pure ( a, loopExits frame )
    [] -> error "withLoopFrame: loop frame disappeared"

-- | Run an action (code generation for the body of the innermost loop) in the body of the loop.
inLoopBodyScope :: MonadState CGState m => m a -> m a
inLoopBodyScope action = do
  setInLoopBody True
  a <- action
  setInLoopBody False
  pure a

setInLoopBody :: MonadState CGState m => Bool -> m ()
setInLoopBody b = modify' \ s -> case loopStack s of
  frame : outer -> s { loopStack = frame { inLoopBody = b } : outer }
  []            -> s

-- | The innermost loop, and whether we are in its body.
innermostLoop :: MonadState CGState m => m ( Maybe ( LoopBlockIDs, Bool ) )
innermostLoop = gets \ s -> case loopStack s of
  frame : _ -> Just ( loopBlockIDs frame, inLoopBody frame )
  []        -> Nothing

-- | Number of loops whose body we are in.
loopDepth :: MonadState CGState m => m Int
loopDepth = gets ( length . filter inLoopBody . loopStack )

-- | Record an early exit from the innermost loop.
addLoopExit
  :: ( MonadState CGState m, MonadError ShortText m )
  => ID -> LoopExit -> Map ShortText (ID, SPIRV.PrimTy) -> m ()
addLoopExit blockID loopExit bindings = do
  frames <- gets loopStack
  case frames of
    frame : outer ->
      modify' \ s -> s { loopStack = frame { loopExits = Map.insert blockID ( loopExit, bindings ) ( loopExits frame ) } : outer }
    [] -> throwError "codeGen: early exit outside of a loop"

----------------------------------------------------------------------------
-- ** Block state

-- | The ID of the current block, if any.
--
-- When there is no current block, the code being generated is unreachable.
currentBlockID :: MonadState CGState m => m ( Maybe ID )
currentBlockID = gets ( currentBlock . blockState )

-- | Start a new block, with an empty block state.
--
-- NB: this doesn't emit the label instruction (see 'CodeGen.CFG.block').
enterBlock :: MonadState CGState m => ID -> m ()
enterBlock blockID = modify' \ s -> s { blockState = noBlock { currentBlock = Just blockID } }

-- | End the current block (e.g. after a terminator instruction).
leaveBlock :: MonadState CGState m => m ()
leaveBlock = modify' \ s -> s { blockState = noBlock }

lookupTemporaryPointer :: MonadState CGState m => ID -> m ( Maybe (Pointer, PointerState) )
lookupTemporaryPointer v = gets ( Map.lookup v . temporaryPointers . blockState )

setTemporaryPointer :: MonadState CGState m => ID -> (Pointer, PointerState) -> m ()
setTemporaryPointer v ptr =
  modify' \ s -> s { blockState = ( blockState s ) { temporaryPointers = Map.insert v ptr ( temporaryPointers ( blockState s ) ) } }

----------------------------------------------------------------------------
-- ** Code generation context

_userGlobals
  :: Lens' CGContext
        ( Map ShortText
            ( SPIRV.PointerTy
            , SPIRV.Decorations
            )
        )
_userGlobals = lens userGlobals ( \c v -> c { userGlobals = v } )

_userGlobal
  :: ShortText
  -> Lens' CGContext
        ( Maybe
            ( SPIRV.PointerTy
            , SPIRV.Decorations
            )
        )
_userGlobal global = _userGlobals . at global

_userFunctions :: Lens' CGContext ( Map ShortText SPIRV.FunctionControl )
_userFunctions = lens userFunctions ( \c v -> c { userFunctions = v } )

_userFunction :: ShortText -> Lens' CGContext ( Maybe SPIRV.FunctionControl )
_userFunction function = _userFunctions . at function

_userEntryPoints :: Lens' CGContext ( Map (ShortText, SPIRV.ExecutionModel) SPIRV.ExecutionModes )
_userEntryPoints = lens userEntryPoints ( \c v -> c { userEntryPoints = v } )

_userEntryPoint :: ShortText -> SPIRV.ExecutionModel -> Lens' CGContext ( Maybe SPIRV.ExecutionModes )
_userEntryPoint name stage = _userEntryPoints . at (name, stage)

_userCapabilities :: Lens' CGContext (Set SPIRV.Capability)
_userCapabilities = lens userCapabilities ( \c v -> c { userCapabilities = v } )

_userExtensions :: Lens' CGContext (Set SPIRV.Extension)
_userExtensions = lens userExtensions ( \c v -> c { userExtensions = v } )

_userImages :: Lens' CGContext ( Map ShortText SPIRV.Image )
_userImages = lens userImages ( \c v -> c { userImages = v } )

_userImage :: ShortText -> Lens' CGContext ( Maybe SPIRV.Image )
_userImage image = _userImages . at image

_spirvVersion :: Lens' CGContext SPIRV.Version
_spirvVersion = lens spirvVersion ( \c v -> c { spirvVersion = v } )

_backend :: Lens' CGContext SPIRV.Backend
_backend = lens backend ( \c v -> c { backend = v } )

_debugging :: Lens' CGContext Bool
_debugging = lens debugging ( \c v -> c { debugging = v } )

_asserting :: Lens' CGContext Bool
_asserting = lens asserting ( \c v -> c { asserting = v } )

_emittingCode :: Lens' CGContext Bool
_emittingCode = lens emittingCode ( \c v -> c { emittingCode = v } )

-----------------------------------------------------------------------------
-- * Helper state update functions

requireCapability :: MonadState CGState m => SPIRV.Capability -> m ()
requireCapability cap = assign ( _neededCapability cap ) (Just ())

requireCapabilities :: MonadState CGState m => Set SPIRV.Capability -> m ()
requireCapabilities = traverseSet_ requireCapability

requireExtension :: MonadState CGState m => SPIRV.Extension -> m ()
requireExtension ext = assign ( _neededExtension ext ) (Just ())

requireExtensions :: MonadState CGState m => Set SPIRV.Extension -> m ()
requireExtensions = traverseSet_ requireExtension

addName :: MonadState CGState m
        => ID -> ShortText -> m ()
addName bdID name
  = modifying _names
      ( Set.insert (bdID, Left name) )

addMemberName :: MonadState CGState m 
              => TyID -> Word32 -> ShortText -> m ()
addMemberName structTyID index name
  = modifying _names
      ( Set.insert (tyID structTyID, Right (index,name)) )

addDecoration :: MonadState CGState m
              => ID -> SPIRV.Decoration Word32 -> m ()
addDecoration bdID dec
  = modifying ( _decorate bdID )
      ( Just . maybe (Set.singleton dec) (Set.insert dec) )

addDecorations :: MonadState CGState m
               => ID -> SPIRV.Decorations -> m ()
addDecorations bdID decs
  = modifying ( _decorate bdID )
      ( Just . maybe decs (Set.union decs) )

addMemberDecoration :: MonadState CGState m
                    => TyID -> Word32 -> SPIRV.Decoration Word32 -> m ()
addMemberDecoration structTyID index dec
  = modifying ( _memberDecorate structTyID index )
      ( Just . maybe (Set.singleton dec) (Set.insert dec) )

addMemberDecorations :: MonadState CGState m
                     => TyID -> Word32 -> SPIRV.Decorations -> m ()
addMemberDecorations structTyID index decs
  = modifying ( _memberDecorate structTyID index )
      ( Just . maybe decs (Set.union decs) )