Pointers.hs

{-# LANGUAGE BlockArguments      #-}
{-# LANGUAGE DerivingStrategies  #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections       #-}
{-# LANGUAGE TypeApplications    #-}

{-|
Module: CodeGen.Pointers

Code generation for load/store instruction using pointers.

Values always have logical types, whereas memory can impose a layout on the types it contains.
Loading and storing converts between the two.
-}

module CodeGen.Pointers
  ( Safeness(Safe, Unsafe)
  , Indices(RTInds, CTInds)
  , newVariable, temporaryVariable
  , accessChain
  , load, loadInstruction
  , store, storeInstruction
  , relayout
  ) where

-- base
import Data.Word
  ( Word32 )

-- containers
import qualified Data.Map.Strict as Map

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

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

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

-- fir
import CodeGen.Binary
  ( instruction )
import CodeGen.Composite
  ( accessedLaidOutTy )
import CodeGen.IDs
  ( typeID, laidOutTypeID, constID )
import CodeGen.Instruction
  ( Args(..), toArgs
  , ID, TyID, Instruction(..)
  , Pointer(..), Variable(..)
  )
import CodeGen.Monad
  ( CGMonad
  , MonadFresh(fresh)
  , note
  )
import CodeGen.State
  ( PointerState(Fresh)
  , _interfaceBinding
  , _knownConstants
  , _spirvVersion
  , currentFunctionContext
  , addLocalVariable
  , lookupTemporaryPointer, setTemporaryPointer
  )
import FIR.Prim.Types
  ( SPrimTy(..), SScalarTy(..)
  , PrimTy(primTySing), ScalarTy(scalarTySing)
  , AConstant(AConstant)
  )
import FIR.ProgramState
  ( FunctionContext(TopLevel, InEntryPoint) )
import qualified SPIRV.Layout    as SPIRV
import qualified SPIRV.Operation as SPIRV.Op
import qualified SPIRV.PrimTy    as SPIRV
import qualified SPIRV.Stage     as SPIRV
import qualified SPIRV.Storage   as Storage
import qualified SPIRV.Version   as SPIRV

----------------------------------------------------------------------------
-- safeness of pointer access (whether indices known to be within bounds)

data Safeness
  = Unsafe
  | Safe
  deriving stock ( Eq, Show )

instance Semigroup Safeness where
  Safe <> x = x
  _    <> _ = Unsafe

data Indices
  = CTInds [Word32] -- compile-time indices
  | RTInds [ID]     -- run-time indices
  deriving stock Show

----------------------------------------------------------------------------
-- creating pointers

-- | Create a new local variable of the current function, holding values of the given type.
newVariable :: SPIRV.PrimTy -> CGMonad Pointer
newVariable ty
  = do
      v <- fresh
      ptrTyID <- typeID ( SPIRV.Pointer ptrTy )
      let ptr = Pointer v ptrTy
      addLocalVariable ( Variable ptr ptrTyID )
      pure ptr
  where
    ptrTy :: SPIRV.PointerTy
    ptrTy = SPIRV.PointerTy Storage.Function SPIRV.logical ty

-- | Obtain a local variable holding a copy of the given value, in the current block.
temporaryVariable :: (ID, SPIRV.PrimTy) -> CGMonad (Pointer, PointerState)
temporaryVariable (baseID, baseTy)
  = do
      mbTemp <- lookupTemporaryPointer baseID
      case mbTemp of
        Just temp -> pure temp
        Nothing   -> do
          ptr <- newVariable baseTy
          storeInstruction ptr baseID
          setTemporaryPointer baseID (ptr, Fresh)
          pure (ptr, Fresh)

indicesIDs :: Indices -> CGMonad [ID]
indicesIDs (CTInds ws) = traverse constID ws
indicesIDs (RTInds is) = pure is

-- create a pointer into a composite object with a list of successive indices
accessChain :: Pointer -> Safeness -> Indices -> CGMonad Pointer
accessChain (Pointer basePtrID (SPIRV.PointerTy storage baseLayout baseTy)) safe indices
  = do
      is <- indicesIDs indices
      (eltLayout, eltTy) <-
        note ( "'accessChain': could not compute accessee type.\n\
               \base: " <> ShortText.pack (show baseTy) <> "\n\
               \indices: " <> ShortText.pack (show indices) <> "."
             )
          =<< ( flip accessedLaidOutTy (baseLayout, baseTy) <$> reverseLookupIndices indices )
      let opAccessChain
            = case safe of
                  Safe -> SPIRV.Op.InBoundsAccessChain
                  _    -> SPIRV.Op.AccessChain
          accessPtrTy = SPIRV.PointerTy storage eltLayout eltTy
      accessPtrTyID <- typeID (SPIRV.Pointer accessPtrTy)
      v <- fresh
      instruction
        Instruction
          { operation = opAccessChain
          , resTy     = Just accessPtrTyID
          , resID     = Just v
          , args      = Arg basePtrID
                      $ toArgs is
          }
      pure (Pointer v accessPtrTy)

reverseLookupIndices :: Indices -> CGMonad [ Maybe Word32 ]
reverseLookupIndices (CTInds is) = pure (map Just is)
reverseLookupIndices (RTInds is) = traverse reverseConstantLookup is

reverseConstantLookup :: ID -> CGMonad (Maybe Word32)
reverseConstantLookup c = do
  lookups <- filter ( (== Just c) . resID . snd ) . Map.assocs <$> use _knownConstants
  case lookups of
    ( (AConstant (a :: ty), _) : _ )
      -> case primTySing @ty of
            SScalar ->
              case scalarTySing @ty of
                SWord8  -> pure . Just $ fromIntegral a
                SWord16 -> pure . Just $ fromIntegral a
                SWord32 -> pure . Just $ fromIntegral a
                SWord64 -> pure . Just $ fromIntegral a
                SInt8   -> pure . Just $ fromIntegral a
                SInt16  -> pure . Just $ fromIntegral a
                SInt32  -> pure . Just $ fromIntegral a
                SInt64  -> pure . Just $ fromIntegral a
                _       -> pure Nothing
            _ -> pure Nothing
    _ -> pure Nothing

----------------------------------------------------------------------------
-- load/store through pointers

-- | Load a value through the pointer of the variable with the given name.
load :: ShortText -> Pointer -> CGMonad (ID, SPIRV.PrimTy)
load loadeeName ptr@(Pointer loadeeID (SPIRV.PointerTy storage _ _))
  = do
      context <- currentFunctionContext
      case context of
        TopLevel -> throwError "codeGen: load operation not allowed at top level"
        InEntryPoint entryPointName stageInfo _
          | storage == Storage.Input
          -- add this variable to the interface of the entry point
          -> assign
                ( _interfaceBinding entryPointName (SPIRV.modelOf stageInfo) loadeeName )
                ( Just loadeeID )
        _ -> pure ()
      loadInstruction ptr

-- | Load a value through a pointer.
loadInstruction :: Pointer -> CGMonad (ID, SPIRV.PrimTy)
loadInstruction (Pointer loadeeID (SPIRV.PointerTy _ lay ty))
  = do  tyID <- laidOutTypeID lay ty
        v <- fresh
        instruction
          Instruction
            { operation = SPIRV.Op.Load
            , resTy = Just tyID
            , resID = Just v
            , args  = Arg loadeeID EndArgs
            }
        (, ty) <$> relayout lay SPIRV.logical (v, ty)

-- | Store a value through the pointer of the variable with the given name.
store :: ShortText -> Pointer -> ID -> CGMonad ()
store storeeName ptr@(Pointer ptrID (SPIRV.PointerTy storage _ _)) storeeID
  = do
      context <- currentFunctionContext
      case context of
        TopLevel -> throwError "codeGen: store operation not allowed at top level"
        InEntryPoint entryPointName stageInfo _
          | storage == Storage.Output
          -- add this variable to the interface of the entry point
          -> assign
                ( _interfaceBinding entryPointName (SPIRV.modelOf stageInfo) storeeName )
                ( Just ptrID )
        _ -> pure ()
      storeInstruction ptr storeeID

-- | Store a value through a pointer.
storeInstruction :: Pointer -> ID -> CGMonad ()
storeInstruction (Pointer ptrID (SPIRV.PointerTy _ lay ty)) storeeID
  = do
      laidOutStoreeID <- relayout SPIRV.logical lay (storeeID, ty)
      instruction
        Instruction
          { operation = SPIRV.Op.Store
          , resTy = Nothing
          , resID = Nothing
          , args = Arg ptrID
                 $ Arg laidOutStoreeID EndArgs
          }

----------------------------------------------------------------------------
-- converting between layouts

-- | Convert a value of the given type from one layout to another.
relayout :: SPIRV.Layout -> SPIRV.Layout -> (ID, SPIRV.PrimTy) -> CGMonad ID
relayout from to (v, ty) = do
  fromTyID <- laidOutTypeID from ty
  toTyID   <- laidOutTypeID to   ty
  if fromTyID == toTyID
  then pure v
  else do
    ver <- view _spirvVersion
    if ver >= SPIRV.Version 1 4
    then
      newValue SPIRV.Op.CopyLogical toTyID ( Arg v EndArgs )
    else do
      components <- case ty of
        SPIRV.Array size eltTy
          -> pure $ replicate ( fromIntegral size ) eltTy
        SPIRV.Struct eltTys
          -> pure $ map snd eltTys
        _ -> throwError
               ( "'relayout': cannot convert a value of type " <> ShortText.pack ( show ty )
               <> " from layout " <> ShortText.pack ( show from )
               <> " to layout " <> ShortText.pack ( show to ) <> "."
               )
      let
        fromComp, toComp :: SPIRV.Layout
        fromComp = SPIRV.componentLayout ty from
        toComp   = SPIRV.componentLayout ty to
      convertedComps <-
        traverse
          ( \ ( i, compTy ) -> do
              compTyID <- laidOutTypeID fromComp compTy
              comp <- newValue SPIRV.Op.CompositeExtract compTyID ( Arg v $ Arg i EndArgs )
              relayout fromComp toComp ( comp, compTy )
          )
          ( zip [ ( 0 :: Word32 ) .. ] components )
      newValue SPIRV.Op.CompositeConstruct toTyID ( toArgs convertedComps )

  where
    newValue :: SPIRV.Op.Operation -> TyID -> Args -> CGMonad ID
    newValue op resTyID operands = do
      u <- fresh
      instruction
        Instruction
          { operation = op
          , resTy     = Just resTyID
          , resID     = Just u
          , args      = operands
          }
      pure u