ProgramState.hs

{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DerivingStrategies    #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}

{-|
Module: FIR.ProgramState

Defines the type-level data used for the state of the indexed monad used to construct programs:

> program :: Program (i :: ProgramState) (j :: ProgramState) (a :: Type)

See "FIR.Module.Program".

The program state consists of:

  * the bindings that are in scope (such as built-in variables or user-defined variables),
  * the function context (whether the current code is contained within a function/entry-point, or is at the top-level),
  * information associated to the functions and entry-points relevant to the program,
    such as whether they have been defined yet or only declared.
  * the data of which SPIR-V backend we are using: Vulkan or OpenCL.
-}

module FIR.ProgramState where

-- base
import Data.Kind
  ( Type )
import Data.Word
  ( Word32 )
import GHC.TypeLits
  ( Symbol, KnownSymbol
  , TypeError, ErrorMessage(..)
  )
import GHC.TypeNats
  ( Nat )

-- text-short
import Data.Text.Short
  ( ShortText )

-- fir
import Data.Type.Known
  ( Demotable(Demote)
  , Known(known), knownValue
  )
import Data.Type.Map
  ( Map, (:->) )
import FIR.Binding
  ( BindingsMap
  , Permissions
  )
import FIR.Prim.RayTracing
  ( RayQueryState )
import FIR.Prim.Types
  ( KnownVars(knownVars) )
import qualified SPIRV.Control    as SPIRV
import qualified SPIRV.Decoration as SPIRV
import qualified SPIRV.PrimTy     as SPIRV
  ( PrimTy )
import qualified SPIRV.Stage      as SPIRV
import qualified SPIRV.Storage    as SPIRV

-------------------------------------------------------------------
-- | Data (kind) used to keep track of function context at the type-level.
data FunctionContext lit nat bindings mbIface where
  TopLevel :: FunctionContext lit nat bindings iface
  InFunction :: lit -> bindings -> FunctionContext lit nat bindings mbIface
  InEntryPoint
    :: lit
    -> SPIRV.ExecutionInfo nat em
    -> mbIface
    -> FunctionContext lit nat bindings mbIface

deriving stock instance ( Show lit, Show bindings, Show nat, Show mbIface )
                     => Show ( FunctionContext lit nat bindings mbIface )


type InterfaceVariable nat ty = ([SPIRV.Decoration nat], ty)
type TLInterfaceVariable = InterfaceVariable Nat    Type
type VLInterfaceVariable = InterfaceVariable Word32 SPIRV.PrimTy

type TLInterface
  = [ SPIRV.StorageClass :-> [ Symbol    :-> TLInterfaceVariable ] ]
type VLInterface
  = [ SPIRV.StorageClass :-> [ ShortText :-> VLInterfaceVariable ] ]

-- | A type-level function context.
type TLFunctionContext
  = FunctionContext
      Symbol
      Nat
      BindingsMap
      ( Maybe TLInterface )
-- | A value-level function context.
type VLFunctionContext
  = FunctionContext
      ShortText
      Word32
      [ ( ShortText, (SPIRV.PrimTy, Permissions) ) ]
      ( Maybe VLInterface )

instance Demotable TLFunctionContext where
  type Demote TLFunctionContext = VLFunctionContext

instance Known TLFunctionContext 'TopLevel where
  known = TopLevel
instance ( KnownVars args
         , KnownSymbol fnName
         )
       => Known TLFunctionContext ('InFunction fnName args) where
  known = InFunction
            ( knownValue @fnName )
            ( knownVars  @args   )
instance ( Known (SPIRV.ExecutionInfo Nat stage) stageInfo
         , KnownSymbol stageName
         )
      => Known TLFunctionContext
            ( 'InEntryPoint
                  stageName
                  ( stageInfo :: SPIRV.ExecutionInfo Nat stage )
                  ('Just iface)
            ) where
  known = InEntryPoint
            ( knownValue     @stageName )
            ( knownValue     @stageInfo )
            Nothing

data Definedness
  = Declared
  | Defined

-- | Keeps track of a function at the type level.
data FunctionInfo where
  FunctionInfo
    :: BindingsMap
    -> Type
    -> SPIRV.FunctionControl
    -> Definedness
    -> FunctionInfo

-- | Keeps track of an entry point at the type level:
--
--   - a type level 'SPIRV.ExecutionInfo' recording the stage as well as additional
--     stage information,
--   - the entry point interface (input/output variables which it uses),
--     (or Nothing if the interface has not yet been computed).
data EntryPointInfo where
  EntryPointInfo
    :: SPIRV.ExecutionInfo Nat stage
    -> TLInterface
    -> Definedness
    -> EntryPointInfo

newtype CFGState = CFGState { loopDepth :: Nat }

-- | State that is used in the user-facing indexed monad (at the type level).
-- Consists of:
--
--    - a type-level list of which bindings are in scope,
--    such as user-defined functions, user-defined variables
--    and built-in variables,
--    - a context – whether code in the current indexed monadic state
--    occurs inside a function or entry-point body –
--    - the current control-flow state,
--    - which functions are to be declared,
--    whether they have been declared yet,
--    and some further information (arguments, function control),
--    - which entry points are to be declared,
--    whether they have been declared yet, together with
--    some additional info pertaining to their respective execution modes,
--    and their interfaces (user defined inputs/outputs),
--    - the state of all current ray-query operations,
--    - which SPIR-V backend is being used (Vulkan or OpenCL).
data ProgramState
  = ProgramState
      { bindings    :: BindingsMap
      , context     :: TLFunctionContext
      , cfgState    :: CFGState
      , functions   :: Map Symbol FunctionInfo
      , entryPoints :: Map Symbol EntryPointInfo
      , interface   :: TLInterface
      , rayQueries  :: Map Symbol RayQueryState
      , backend     :: SPIRV.Backend
      }

type family Bindings ( s :: ProgramState ) :: BindingsMap where
  Bindings ('ProgramState bds _ _ _ _ _ _ _) = bds

type family CFGInfo ( s :: ProgramState ) :: CFGState where
  CFGInfo ('ProgramState _ _ cfg _ _ _ _ _) = cfg

type family FunctionInfos ( s :: ProgramState ) :: Map Symbol FunctionInfo where
  FunctionInfos ('ProgramState _ _ _ fs _ _ _ _) = fs

type family EntryPointInfos ( s :: ProgramState ) :: Map Symbol EntryPointInfo where
  EntryPointInfos ('ProgramState _ _ _ _ eps _ _ _) = eps

type family ExecutionContext ( s :: ProgramState ) :: Maybe SPIRV.ExecutionModel where
  ExecutionContext ('ProgramState _ ('InEntryPoint _ (info :: SPIRV.ExecutionInfo Nat stage) _) _ _ _ _ _ _)
    = Just stage
  ExecutionContext _
    = Nothing

type family RayQueries ( s :: ProgramState ) :: Map Symbol RayQueryState where
  RayQueries ('ProgramState _ _ _ _ _ _ rayQueries _) = rayQueries

executionContext :: VLFunctionContext -> Maybe (ShortText, SPIRV.ExecutionModel)
executionContext (InEntryPoint stageName stageInfo _) = Just (stageName, SPIRV.modelOf stageInfo)
executionContext _ = Nothing

type family ExecutionContext' ( s :: ProgramState ) :: SPIRV.ExecutionModel where
  ExecutionContext' ('ProgramState _ ('InEntryPoint _ (info :: SPIRV.ExecutionInfo Nat stage) _) _ _ _ _ _ _)
    = stage
  ExecutionContext' _
    = TypeError
        ( Text "Cannot access stage context: not within a stage." )

type family ExecutionInfoContext
                ( s :: ProgramState )
              :: Maybe (SPIRV.ExecutionInfo Nat (ExecutionContext' s))
                where
  ExecutionInfoContext ('ProgramState _ ('InEntryPoint _ info _) _ _ _ _ _ _)
    = Just info
  ExecutionInfoContext _ = 'Nothing