fir-0.1.0.0: src/CodeGen/Optics.hs
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}
{-|
Module: CodeGen.Optics
Code generation for operations involving the type-level optics used by this library,
such as @view@, @set@, @use@, @assign@.
-}
module CodeGen.Optics ( ) where
-- base
import Control.Arrow
( first )
import Control.Monad
( unless )
import Data.Foldable
( traverse_ )
import Data.Kind
( Type )
import Data.Proxy
( Proxy(Proxy) )
import Data.Type.Equality
( (:~:)(Refl) )
import Data.Typeable
( eqT )
import Data.Word
( Word32 )
import Unsafe.Coerce
( unsafeCoerce )
-- lens
import Control.Lens
( assign )
-- mtl
import Control.Monad.Except
( throwError )
-- text-short
import Data.Text.Short
( ShortText )
import qualified Data.Text.Short as ShortText
( pack )
-- fir
import CodeGen.Application
( ASTs(NilAST, ConsAST), pattern SnocAST
, Codes, pattern ConsCode
, Application(..)
)
import {-# SOURCE #-} CodeGen.CodeGen
( CodeGen(codeGenArgs), codeGen )
import CodeGen.Composite
( compositeConstruct
, compositeExtract
, compositeInsert
, vectorSwizzle
, productConstruct
)
import CodeGen.Debug
( whenAsserting )
import CodeGen.IDs
( constID, bindingID, typeID )
import CodeGen.Images
( imageTexel, writeTexel )
import CodeGen.Instruction
( ID(ID), Pointer(..) )
import CodeGen.Monad
( CGMonad )
import CodeGen.Pointers
( Safeness(Safe, Unsafe)
, Indices(RTInds, CTInds)
, temporaryVariable, accessChain
, load, store
, loadInstruction, storeInstruction
)
import CodeGen.Atomics
( atomicInstruction )
import CodeGen.State
( PointerState(Fresh, Modified)
, _localBinding
, setTemporaryPointer
)
import Control.Type.Optic
( Optic
, ProductComponents(..) )
import Data.Constraint.All
( All(allDict)
, AllDict(NilDict, ConsDict)
)
import Data.Product
( HList(HNil, (:>))
, MapHList
, IsProduct(toHList)
, AreProducts(productsDict)
, AreProductsDict(ConsProducts)
, Distribute
, distributeZipConsLemma1
)
import Data.Type.Known
( knownValue )
import Data.Type.List
( type (:++:), ZipCons
, KnownLength(sLength)
, SLength(SZero, SSucc)
, SameLength(sSameLength)
, SSameLength(SSameZero, SSameSucc)
, Snoc
)
import FIR.AST
( AST, Code
, pattern (:$), pattern Lit
, pattern MkID
, pattern NilHList, pattern ConsHList
, OpticF(..)
)
import FIR.AST.Type
( AugType(Val), MapVal )
import FIR.Prim.Atomic
( SAtomicOp, atomicOp, SAtomicShape(..), atomicShape )
import FIR.Prim.Types
( PrimTy, primTy, primTys
, SPrimTy(..), sPrimTy
)
import FIR.Syntax.Optics
( SOptic(..), showSOptic
, SProductComponents(..)
)
import qualified SPIRV.PrimTy as SPIRV
----------------------------------------------------------------------------
-- Code generation for optics.
data OpticArgs where
OpticView :: forall is s a (optic :: Optic is s a). SOptic optic -> Codes is -> Code s -> OpticArgs
OpticSet :: forall is s a (optic :: Optic is s a). SOptic optic -> Codes is -> Code a -> Code s -> OpticArgs
OpticUse :: forall is s a (optic :: Optic is s a). SOptic optic -> Codes is -> OpticArgs
OpticAssign :: forall is s a (optic :: Optic is s a). SOptic optic -> Codes is -> Code a -> OpticArgs
OpticAtomic :: forall is s a (optic :: Optic is s a) op
. SAtomicOp op -> SOptic optic -> Codes is -> [ Code a ]
-> OpticArgs
opticArgs :: Application (OpticF AST) f r -> OpticArgs
opticArgs ( Applied ( ViewF _ ( sOptic :: SOptic ( optic :: Optic is s a) ) ) is_s ) =
case unsafeCoerce is_s :: ( ASTs (MapVal is `Snoc` Val s) ) of
( is `SnocAST` s ) -> OpticView sOptic is s
opticArgs ( Applied ( SetF _ ( sOptic :: SOptic ( optic :: Optic is s a) ) ) is_a_s ) =
case unsafeCoerce is_a_s :: ( ASTs ( MapVal is `Snoc` Val a `Snoc` Val s ) ) of
( is `SnocAST` a `SnocAST` s ) -> OpticSet sOptic is a s
opticArgs ( Applied ( UseF _ ( sOptic :: SOptic ( optic :: Optic is s a) ) ) is ) =
case unsafeCoerce is :: ASTs ( MapVal is ) of
js -> OpticUse sOptic js
opticArgs ( Applied ( AssignF _ ( sOptic :: SOptic ( optic :: Optic is s a) ) ) is_a ) =
case unsafeCoerce is_a :: ( ASTs ( MapVal is `Snoc` Val a ) ) of
( is `SnocAST` a ) -> OpticAssign sOptic is a
opticArgs ( Applied ( AtomicF sOp _ ( sOptic :: SOptic ( optic :: Optic is s a) ) ) is_vs ) =
case atomicShape sOp of
SReads ->
OpticAtomic sOp sOptic ( unsafeCoerce is_vs :: Codes is ) []
SWrites ->
case unsafeCoerce is_vs :: ( ASTs ( MapVal is `Snoc` Val a ) ) of
( is `SnocAST` v ) -> OpticAtomic sOp sOptic is [ v ]
SModifies ->
case unsafeCoerce is_vs :: ( ASTs ( MapVal is `Snoc` Val a ) ) of
( is `SnocAST` v ) -> OpticAtomic sOp sOptic is [ v ]
SCompares ->
case unsafeCoerce is_vs :: ( ASTs ( MapVal is `Snoc` Val a `Snoc` Val a ) ) of
( is `SnocAST` expected `SnocAST` new )
-> OpticAtomic sOp sOptic is [ expected, new ]
instance CodeGen AST => CodeGen (OpticF AST) where
codeGenArgs opticApplication = case opticArgs opticApplication of {
-- View.
OpticView sOptic is s -> do
base <- codeGen s
extractUsingGetter base sOptic is
;
-- Set.
OpticSet sOptic is a s -> do
setUsingSetter s a sOptic is
;
-- Use.
OpticUse sOptic is -> case sOptic of
SBinding (_ :: Proxy name) ->
do let varName = knownValue @name
bd@(bdID, bdTy) <- bindingID varName
case bdTy of
SPIRV.Pointer ptrTy
-> load varName (Pointer bdID ptrTy)
_ -> pure bd
-- image
SImageTexel ( _ :: Proxy name ) ( _ :: Proxy props ) texelTy
-> case is of
( ops `ConsAST` coords `ConsAST` NilAST )
-> do
let imgName = knownValue @name
imgTexelTy = sPrimTy texelTy
bd@(bdID, bdTy) <- bindingID imgName
resTyID <- typeID imgTexelTy
img <- case bdTy of
SPIRV.Pointer ptrTy
-> load imgName (Pointer bdID ptrTy)
_ -> pure bd
(cdID, _) <- codeGen coords
imageTexel img ops cdID (resTyID, imgTexelTy)
SComposeO _ (SBinding (_ :: Proxy name )) getter ->
do let varName = knownValue @name
bd@(bdID, bdTy) <- bindingID varName
case bdTy of
SPIRV.Pointer ptrTy
-> loadThroughAccessChain (Pointer bdID ptrTy) getter is
_ -> extractUsingGetter bd getter is
SComposeO
(SSucc (SSucc _))
(SImageTexel ( _ :: Proxy name ) ( _ :: Proxy props ) texelTy)
getter
-> case is of
( ops `ConsAST` coords `ConsAST` nextindices )
->
do
-- copy-pasted from above
let imgName = knownValue @name
imgTexelTy = sPrimTy texelTy
bd@(bdID, bdTy) <- bindingID imgName
resTyID <- typeID imgTexelTy
img <- case bdTy of
SPIRV.Pointer ptrTy
-> load imgName (Pointer bdID ptrTy)
_ -> pure bd
(cdID, _) <- codeGen coords
-- end of copy-paste
texel <- imageTexel img ops cdID (resTyID, imgTexelTy)
extractUsingGetter texel getter nextindices
_ -> throwError ( "codeGen: cannot 'use', unsupported optic:\n"
<> ShortText.pack (showSOptic sOptic) <> "\n"
<> "Optic does not start by accessing a binding (or image texel).\n"
)
;
-- Assign.
OpticAssign sOptic is a -> case sOptic of
SBinding (_ :: Proxy name) ->
do let varName = knownValue @name
a_IDTy@(a_ID, _) <- codeGen a
(bdID, bdTy) <- bindingID varName
case bdTy of
SPIRV.Pointer ptrTy
-> store varName (Pointer bdID ptrTy) a_ID
_ -> assign ( _localBinding varName ) (Just a_IDTy)
pure (ID 0, SPIRV.Unit) -- ID should never be used
-- image
SImageTexel ( _ :: Proxy name ) ( _ :: Proxy props ) _
-> case is of
( ops `ConsAST` coords `ConsAST` NilAST )
-> do
let imgName = knownValue @name
bd@(bdID, bdTy) <- bindingID imgName
img <- case bdTy of
SPIRV.Pointer ptrTy
-> load imgName (Pointer bdID ptrTy)
_ -> pure bd
(cdID , _) <- codeGen coords
(texID, _) <- codeGen a
writeTexel img ops cdID texID
pure (ID 0, SPIRV.Unit) -- ID should never be used
SComposeO _ (SBinding (_ :: Proxy name)) setter ->
do let varName = knownValue @name
bd@(bdID, bdTy) <- bindingID varName
case bdTy of
SPIRV.Pointer ptrTy
-> storeThroughAccessChain (Pointer bdID ptrTy) a setter is
_ -> insertUsingSetter varName bd a setter is
pure (ID 0, SPIRV.Unit) -- ID should never be used
SComposeO _ SImageTexel {} _
-> throwError ( "codeGen: writing to individual components \
\of an image texel currently unsupported"
)
_ -> throwError ( "codeGen: cannot 'assign', unsupported optic:\n"
<> ShortText.pack (showSOptic sOptic) <> "\n"
<> "Optic does not start by accessing a binding."
)
;
-- Atomic operation.
OpticAtomic sOp sOptic is operands -> do
let opName = ShortText.pack ( show ( atomicOp sOp ) )
ptr <- opticPointer opName sOptic is
operandIDs <- traverse ( fmap fst . codeGen ) operands
atomicInstruction sOp ptr operandIDs
}
----------------------------------------------------------------------------
-- Pointers denoted by optics.
-- | The pointer an optic focuses on: the optic must start with a binding of
-- pointer type, possibly followed by an access chain.
opticPointer
:: forall is s a (optic :: Optic is s a)
. CodeGen AST
=> ShortText -- ^ Name of the operation (for error messages).
-> SOptic optic
-> Codes is
-> CGMonad Pointer
opticPointer opName sOptic is = case sOptic of
SBinding ( _ :: Proxy name ) ->
bindingPointer ( knownValue @name )
SComposeO _ ( SBinding ( _ :: Proxy name ) ) accessor -> do
basePtr <- bindingPointer ( knownValue @name )
pointerThroughAccessChain' opName basePtr =<< operationTree is accessor
_ -> unsupported "Optic does not start by accessing a binding."
where
bindingPointer :: ShortText -> CGMonad Pointer
bindingPointer varName = do
(bdID, bdTy) <- bindingID varName
case bdTy of
SPIRV.Pointer ptrTy
-> pure ( Pointer bdID ptrTy )
_ -> unsupported
( "Binding \"" <> varName <> "\" is not of pointer type." )
unsupported :: ShortText -> CGMonad x
unsupported reason =
throwError ( "codeGen: cannot '" <> opName <> "', unsupported optic:\n"
<> ShortText.pack (showSOptic sOptic) <> "\n"
<> reason
)
----------------------------------------------------------------------------
-- Perform code-generation by computing associated optical trees.
-- | Load through a pointer with the supplied getter.
loadThroughAccessChain
:: forall is s a (optic :: Optic is s a)
. CodeGen AST
=> Pointer -> SOptic optic -> Codes is -> CGMonad (ID, SPIRV.PrimTy)
loadThroughAccessChain basePtr sOptic is
= loadThroughAccessChain' basePtr =<< operationTree is sOptic
-- | Use a getter to @view@ a part of the provided object.
extractUsingGetter
:: forall is s a (optic :: Optic is s a)
. CodeGen AST
=> (ID, SPIRV.PrimTy) -> SOptic optic -> Codes is -> CGMonad (ID, SPIRV.PrimTy)
extractUsingGetter base sOptic is
= extractUsingGetter' base =<< operationTree is sOptic
-- | Store into a pointer with the supplied setter.
storeThroughAccessChain
:: forall is s a (optic :: Optic is s a)
. CodeGen AST
=> Pointer -> Code a -> SOptic optic -> Codes is -> CGMonad ()
storeThroughAccessChain basePtr val sOptic is
= storeThroughAccessChain' basePtr val =<< operationTree is sOptic
-- | Set a component of an object, with focus given by the supplied setter.
--
-- Returns the updated object.
setUsingSetter
:: forall is s a (optic :: Optic is s a)
. CodeGen AST
=> Code s
-> Code a
-> SOptic optic
-> Codes is
-> CGMonad (ID, SPIRV.PrimTy)
setUsingSetter base val sOptic is = do
baseID <- codeGen base
tree <- operationTree is sOptic
setUsingSetter' baseID val tree
-- | Set a component of a binding, with focus given by the supplied setter.
--
-- Doesn't return anything: updates the state of the provided binding.
insertUsingSetter
:: forall is s a (optic :: Optic is s a)
. CodeGen AST
=> ShortText -> (ID, SPIRV.PrimTy) -> Code a -> SOptic optic -> Codes is -> CGMonad ()
insertUsingSetter varName base val sOptic is = do
tree <- operationTree is sOptic
insertUsingSetter' varName base val tree
----------------------------------------------------------------------------
-- Optical trees.
data OpticalNode where
Access :: Safeness -> Indices -> OpticalNode
OfTypeOp :: SPrimTy s -> SPrimTy a -> OpticalNode
Combine :: (PrimTy p, All PrimTy as, IsProduct p as)
=> Proxy p -> Proxy as -> [ OpticalOperationTree ] -> OpticalNode
instance Show OpticalNode where
show (Access safe is) = "Access " ++ show safe ++ " " ++ show is
show (OfTypeOp s a) = "OfTypeOp @" ++ show s ++ " @" ++ show a
show (Combine (_ :: Proxy p) ( _ :: Proxy as) trees)
= "Combine @" ++ show (primTy @p) ++ " @" ++ show (primTys @as)
++ " " ++ show trees
type OpticalOperationTree = [ OpticalNode ]
operationTree
:: forall k is (s :: k) a (optic :: Optic is s a)
. CodeGen AST
=> Codes is -> SOptic optic -> CGMonad OpticalOperationTree
operationTree _ (SOfType s a) = pure [OfTypeOp s a]
operationTree (i `ConsAST` _) SAnIndex {}
= (:[]) . Access Unsafe . RTInds . (:[]) . fst <$> codeGen i
operationTree _ (SIndex s _ n)
-- if accessing a runtime array, a compile-time index may be unsafe
| SRuntimeArray <- s
= pure [ Access Unsafe (CTInds [n]) ]
-- otherwise, a compile-time index is guaranteed to be in-bounds
| otherwise
= pure [ Access Safe (CTInds [n]) ]
operationTree is (SComposeO lg1 opt1 opt2)
= do let (is1, is2) = composedIndices lg1 is
ops1 <- operationTree is1 opt1
ops2 <- operationTree is2 opt2
ops1 `continue` ops2
operationTree is
( SProd same (comps :: SProductComponents (os :: ProductComponents iss s as))
:: SOptic (optic :: Optic is s a)
) = do
comp_indices <- combinedIndices @iss @is @as sSameLength sLength is
comp_trees <- componentsTrees same comps comp_indices
case comp_trees of
[] -> return []
[ tree ] -> return tree
trees -> return [ Combine ( Proxy @a ) ( Proxy @as ) trees ]
-- N.B.: only combine optical subtrees when there is combining taking place.
-- If there are fewer than 2 subtrees, we won't be creating a composite structure.
-- See bug #88.
operationTree _ (SBinding {})
= throwError "operationTree: trying to access a binding within a binding"
operationTree _ (SImageTexel {})
= throwError "operationTree: unexpected image optic"
continue :: OpticalOperationTree -> OpticalOperationTree -> CGMonad OpticalOperationTree
continue ops1 [] = pure ops1
continue [] ops2 = pure ops2
-- try to combine successive accesses into a single access chain as much as possible
-- (this helps bypass Vulkan implementation bugs in tessellation evaluation shaders)
continue [Access safe1 (CTInds is1)] (Access safe2 (CTInds is2) : ops2)
= pure $ Access (safe1 <> safe2) (CTInds (is1 ++ is2)) : ops2
continue [Access safe1 (RTInds is1)] (Access safe2 (RTInds is2) : ops2)
= pure $ Access (safe1 <> safe2) (RTInds (is1 ++ is2)) : ops2
continue [Access safe1 (CTInds is1)] (Access safe2 (RTInds is2) : ops2)
= do
js1 <- traverse constID is1
pure $ ( Access (safe1 <> safe2) (RTInds (js1 ++ is2)) : ops2 )
continue [Access safe1 (RTInds is1)] (Access safe2 (CTInds is2) : ops2)
= do
js2 <- traverse constID is2
pure $ Access (safe1 <> safe2) (RTInds (is1 ++ js2)) : ops2
continue (op1 : ops1) ops2
= do
cont <- ops1 `continue` ops2
pure $ op1 : cont
componentsTrees
:: forall
(k :: Type) (iss :: [[Type]]) (s :: k)
(as :: [Type]) (os :: ProductComponents iss s as)
. CodeGen AST
=> SSameLength (Distribute iss as) as
-> SProductComponents os
-> Codes (MapHList (Distribute iss as))
-> CGMonad [OpticalOperationTree]
componentsTrees _ SEndProd _
= pure []
componentsTrees sameSucc@(SSameSucc lg) (so `SProductO` sos) (is `ConsAST` iss) =
case sameSucc of
( _ :: SSameLength (ds ': dss) (b ': bs) ) ->
case so of
( _ :: SOptic ( o :: Optic es s b ) ) ->
case sos of
( _ :: SProductComponents ( tail_os :: ProductComponents ess s bs ) ) ->
case distributeZipConsLemma1 @b @bs @ds @dss @es @ess of
(Refl, Refl) ->
(:) <$> operationTree (astsFromHList is) so <*> componentsTrees lg sos iss
composedIndices :: SLength is -> Codes (is :++: js) -> (Codes is, Codes js)
composedIndices SZero js = ( NilAST, js )
composedIndices (SSucc tail_is) (k `ConsAST` ks)
= first ( k `ConsAST` ) ( composedIndices tail_is ks )
combinedIndices
:: forall (jss :: [[Type]]) (is :: [Type]) (as :: [Type])
. ( CodeGen AST, AreProducts is jss as )
=> SSameLength is jss
-> SLength as
-> Codes is
-> CGMonad (Codes (MapHList (Distribute jss as)))
combinedIndices SSameZero SZero NilAST
= pure $ NilAST
combinedIndices SSameZero (SSucc lg) NilAST
= (NilHList `ConsAST`) <$> combinedIndices SSameZero lg NilAST
combinedIndices sSame@(SSameSucc sm) lg (i `ConsAST` is) =
case sSame of
( _ :: SSameLength (p ': ps) (es ': ess) ) ->
case productsDict @is @jss @as of
ConsProducts ->
zipIndices @es @(Distribute ess as) allDict sSameLength
<$> deconstruct allDict i
<*> combinedIndices sm lg is
zipIndices
:: AllDict PrimTy is
-> SSameLength is jss
-> Codes is
-> Codes (MapHList jss)
-> Codes (MapHList (ZipCons is jss))
zipIndices NilDict SSameZero NilAST NilAST = NilAST
zipIndices ConsDict (SSameSucc lg) (i `ConsAST` is) (js `ConsAST` jss) =
(ConsHList :$ i :$ js) `ConsAST` zipIndices allDict lg is jss
deconstruct
:: forall (p :: Type) (as :: [Type])
. ( CodeGen AST, IsProduct p as )
=> AllDict PrimTy as -> Code p -> CGMonad (Codes as)
deconstruct NilDict _ = pure NilAST
deconstruct dict (Lit x) = pure $ hListToASTs dict (toHList @p @as x)
deconstruct dict p
| Just hlist <- recogniseHList dict p
= pure hlist
-- TODO: more cases to avoid spurious "compositeConstruct ---> compositeExtract"
deconstruct ConsDict p = do
composite <- codeGen p
compositeExtractAll composite allDict 0
recogniseHList :: AllDict PrimTy as -> Code p -> Maybe (Codes as)
recogniseHList NilDict _ = Just NilAST
recogniseHList consDict@ConsDict ast
= case consDict of
( _ :: AllDict PrimTy (b ': bs) ) ->
case ast of
( ConsHList :$ (c :: Code c) :$ cs )
| Just Refl <- eqT @b @c
-> (c `ConsAST`) <$> recogniseHList allDict cs
_ -> Nothing
astsFromHList :: Code (HList as) -> Codes as
astsFromHList NilHList
= NilAST
astsFromHList (ConsHList :$ a :$ as)
= a `ConsAST` astsFromHList as
astsFromHList _
= error "'operationTree': indexing heterogenous list AST not of the expected form (TODO?)"
hListToASTs :: AllDict PrimTy as -> HList as -> Codes as
hListToASTs _ HNil = NilAST
hListToASTs ConsDict (a :> as) = Lit a `ConsAST` hListToASTs allDict as
compositeExtractAll
:: forall (as :: [Type])
. CodeGen AST
=> (ID, SPIRV.PrimTy)
-> AllDict PrimTy as
-> Word32
-> CGMonad (Codes as)
compositeExtractAll _ NilDict _ = pure NilAST
-- special case for matrix entries:
-- product for matrices can use individual components (rather than just columns)
compositeExtractAll comp@(_,compTy) dict@ConsDict i
| SPIRV.Matrix {..} <- compTy
, SPIRV.Scalar entryTy == expectedTy
= do
extracted <- traverse (compositeExtract comp) $ [ [i,j] | j <- [0..pred rows] ]
makeIDsAndContinue (succ i) dict extracted
where
expectedTy :: SPIRV.PrimTy
expectedTy = case dict of
( _ :: AllDict PrimTy (b ': bs) )
-> primTy @b
makeIDsAndContinue
:: Word32
-> AllDict PrimTy ts
-> [(ID, SPIRV.PrimTy)]
-> CGMonad (Codes ts)
makeIDsAndContinue _ NilDict _ = pure NilAST
makeIDsAndContinue nxt dict' [] = compositeExtractAll comp dict' nxt
makeIDsAndContinue nxt dict'@ConsDict
( extract@(_,extractedTy) : extracts ) =
case dict' of
( _ :: AllDict PrimTy (c ': cs) )
-> do
whenAsserting . unless ( expectedTy == extractedTy ) $
throwError $
"'compositeExtractAll': ASSERT failed\n"
<> "expected type" <> ShortText.pack (show expectedTy)
<> ", but composite component " <> ShortText.pack (show i)
<> " has type " <> ShortText.pack (show extractedTy)
next <- makeIDsAndContinue nxt allDict extracts
pure $ MkID extract `ConsAST` next
-- usual case
compositeExtractAll comp ConsDict i =
do
extracted <- compositeExtract comp [i]
next <- compositeExtractAll comp allDict (succ i)
pure $ MkID extracted `ConsAST` next
----------------------------------------------------------------------------
-- Generate code for a tree of optical operations.
-- check whether a vector swizzle operation can be used
swizzleIndices :: SPIRV.PrimTy -> [OpticalOperationTree] -> Maybe [Word32]
swizzleIndices (SPIRV.Vector _ _) = traverse simpleIndex
where
simpleIndex :: OpticalOperationTree -> Maybe Word32
simpleIndex [Access Safe (CTInds [i])] = Just i
simpleIndex _ = Nothing
swizzleIndices _ = const Nothing
loadThroughAccessChain'
:: CodeGen AST
=> Pointer -> OpticalOperationTree -> CGMonad (ID, SPIRV.PrimTy)
loadThroughAccessChain' basePtr []
= loadInstruction basePtr
loadThroughAccessChain' basePtr ( Access safe is : ops )
= do
newBasePtr <- accessChain basePtr safe is
loadThroughAccessChain' newBasePtr ops
loadThroughAccessChain' basePtr@(Pointer _ (SPIRV.PointerTy _ _ eltTy)) ( Combine (_ :: Proxy p) _ trees : ops )
-- special case: can we use a vector swizzle operation?
| Just is <- swizzleIndices eltTy trees
= do
base <- loadThroughAccessChain' basePtr []
vec <- vectorSwizzle base is
extractUsingGetter' vec ops
| otherwise
= do
subtrees <- traverse (loadThroughAccessChain' basePtr) trees
combined <- productConstruct (primTy @p) subtrees
extractUsingGetter' combined ops
loadThroughAccessChain' _ ( OfTypeOp _ _ : _ )
= throwError "loadThroughAccessChain': unexpected 'OfType' optic used as a getter"
-- | Follow a pointer through a tree of optical operations, which must be an
-- access chain (no product or 'OfType' optics).
pointerThroughAccessChain'
:: ShortText -- ^ Name of the operation (for error messages).
-> Pointer -> OpticalOperationTree -> CGMonad Pointer
pointerThroughAccessChain' _ basePtr []
= pure basePtr
pointerThroughAccessChain' opName basePtr ( Access safe is : ops )
= do
newBasePtr <- accessChain basePtr safe is
pointerThroughAccessChain' opName newBasePtr ops
pointerThroughAccessChain' opName _ ( OfTypeOp {} : _ )
= throwError ( "codeGen: cannot '" <> opName <> "' through an 'OfType' \
\optic, which does not denote a pointer." )
pointerThroughAccessChain' opName _ ( Combine {} : _ )
= throwError ( "codeGen: cannot '" <> opName <> "' through a product \
\optic, which does not denote a pointer." )
extractUsingGetter'
:: CodeGen AST
=> (ID, SPIRV.PrimTy) -> OpticalOperationTree -> CGMonad (ID, SPIRV.PrimTy)
extractUsingGetter' base []
= pure base
extractUsingGetter' base tree@( Access _ (CTInds is) : ops )
= do
whenAsserting
case base of
( _, SPIRV.Scalar _ ) ->
throwError
( "extractUsingGetter': ASSERT failed, trying to extract from a scalar type\n\
\base = " <> ShortText.pack ( show base ) <> "\n\
\tree = " <> ShortText.pack ( show tree )
)
_ -> pure ()
newBase <- compositeExtract base is
extractUsingGetter' newBase ops
extractUsingGetter' base@(baseID, _) ops@( Access _ (RTInds _) : _ )
-- run-time indices: revert to loading through pointers
= do
(basePtr, basePtrState) <- temporaryVariable base
unless
( basePtrState == Fresh )
( do storeInstruction basePtr baseID
-- reset temporary pointer to "fresh" as we just loaded the baseID into it
setTemporaryPointer baseID (basePtr, Fresh)
)
loadThroughAccessChain' basePtr ops
extractUsingGetter' base@(_, eltTy) ( Combine (_ :: Proxy p) _ trees : ops )
-- special case: can we use a vector swizzle operation?
| Just is <- swizzleIndices eltTy trees
= do
vec <- vectorSwizzle base is
extractUsingGetter' vec ops
| otherwise
= do
subtrees <- traverse (extractUsingGetter' base) trees
combined <- productConstruct (primTy @p) subtrees
extractUsingGetter' combined ops
extractUsingGetter' _ ( OfTypeOp _ _ : _ )
= throwError "extractUsingGetter': unexpected 'OfType' optic used as a getter"
storeThroughAccessChain'
:: CodeGen AST
=> Pointer -> Code a -> OpticalOperationTree -> CGMonad ()
-- Base case: identity.
storeThroughAccessChain' basePtr val []
= storeInstruction basePtr =<< fst <$> codeGen val
-- Access.
storeThroughAccessChain' basePtr val ( Access safe is : ops )
= do
newBasePtr <- accessChain basePtr safe is
storeThroughAccessChain' newBasePtr val ops
-- OfType optic.
storeThroughAccessChain' basePtr val ( OfTypeOp _ a : ops )
= storeAtTypeThroughAccessChain basePtr val ops (sPrimTy a) [ wholeObject basePtr ]
-- Product optic.
-- This is a bit more tricky and requires a second level of recursion.
--
-- * Base case.
storeThroughAccessChain' basePtr val ( Combine (_ :: Proxy p) (_ :: Proxy as) trees : [] )
= successiveStores basePtr trees =<< deconstruct @p @as allDict (unsafeCoerce val)
-- * Access.
storeThroughAccessChain' _ _ ( Combine {} : Access {} : _ )
= throwError "storeThroughAccessChain': currently unsupported composition ( product optic - indexing )"
-- * OfType optic.
-- If "OfType ty" pertains to the individual components, then we can
-- perform a 'storeAtTypeThroughAccessChain' on all the components in turn.
-- We must first catch the other cases: focus onto a part that arose from the creation of the product type.
storeThroughAccessChain' basePtr val ( Combine (px_p :: Proxy p) (px_as :: Proxy as) trees : OfTypeOp _ a : ops )
-- Start off by catching situations where we focus on a part that arose from the product.
-- * 'Combine' followed by 'OfType combinedTy': can ignore 'OfType combinedTy' (identity lens).
| combinedTy == ofTypeTy
= storeThroughAccessChain' basePtr val ( Combine px_p px_as trees : ops )
-- * Finicky corner case: we are combining entries (not columns) to form a matrix, but followed by `OfType colTy`.
| SPIRV.Matrix m _ eltTy <- combinedTy
, ( SPIRV.Scalar eltTy' : _ ) <- compTys
, eltTy' == eltTy
, ( SPIRV.Vector m' (SPIRV.Scalar eltTy'') ) <- ofTypeTy
, eltTy'' == eltTy && m' == m
= throwError "storeThroughAccessChain': currently unsupported corner case \
\(matrix product of entries followed by 'OfType' optic focused on column type)"
-- General case: call 'storeAtTypeThroughAccessChain' on all the components in turn.
| otherwise
= storeAtTypeThroughAccessChain basePtr val ops ofTypeTy ( zipWith Component trees compTys )
where
combinedTy, ofTypeTy :: SPIRV.PrimTy
compTys :: [SPIRV.PrimTy]
combinedTy = primTy @p
ofTypeTy = sPrimTy a
compTys = primTys @as
-- * Second product optic.
storeThroughAccessChain' _ _ ( Combine {} : Combine {} : _ )
= throwError "storeThroughAccessChain': currently unsupported composition ( product optic - product optic )"
-- | Store into all the subcomponents of the given type of the given components,
-- applying the given operations to each of these subcomponents.
storeAtTypeThroughAccessChain
:: CodeGen AST
=> Pointer
-> Code a
-> OpticalOperationTree
-> SPIRV.PrimTy
-> [ Component ]
-> CGMonad ()
storeAtTypeThroughAccessChain basePtr val ops a = traverse_ go
where
go :: Component -> CGMonad ()
go comp
| componentTy comp == a
= storeInto basePtr comp val ops
| [] <- ops
, Just filled <- replicated ( componentTy comp )
= do
v <- filled
storeInto basePtr comp ( MkID v ) []
| otherwise
= traverse_ go =<< subcomponents comp
-- A vector or matrix whose entries (or columns) are all of the given type
-- can be stored into all at once.
replicated :: SPIRV.PrimTy -> Maybe ( CGMonad (ID, SPIRV.PrimTy) )
replicated ty = case ty of
SPIRV.Vector n elt
| elt == a
-> Just $ fill ty n =<< value
SPIRV.Matrix m n elt
| a == colTy
-> Just $ fill ty n =<< value
| a == SPIRV.Scalar elt
-> Just $ fill ty n . fst =<< fill colTy m =<< value
where
colTy :: SPIRV.PrimTy
colTy = SPIRV.Vector m (SPIRV.Scalar elt)
_ -> Nothing
value :: CGMonad ID
value = fst <$> codeGen val
fill :: SPIRV.PrimTy -> Word32 -> ID -> CGMonad (ID, SPIRV.PrimTy)
fill ty n x = compositeConstruct ty ( replicate (fromIntegral n) x )
-- | A component of the object a pointer points to.
data Component
= Component
{ componentAccess :: OpticalOperationTree -- ^ How to access the component.
, componentTy :: SPIRV.PrimTy
}
-- | The object a pointer points to, as a whole.
wholeObject :: Pointer -> Component
wholeObject ptr = Component [] ( SPIRV.pointee ( pointerTy ptr ) )
-- | The immediate subcomponents of a component.
subcomponents :: Component -> CGMonad [ Component ]
subcomponents ( Component access ty ) = case ty of
SPIRV.Vector n elt -> indexed n elt
SPIRV.Matrix m n elt -> indexed n ( SPIRV.Vector m (SPIRV.Scalar elt) )
-- TODO: should be a loop as opposed to being fully unrolled,
-- but creating loops is unfortunately quite cumbersome
SPIRV.Array l elt -> indexed l elt
SPIRV.Struct eltTys -> sequence $ zipWith sub [0..] ( map snd eltTys )
SPIRV.RuntimeArray {} -> throwError "codeGen: cannot store into all components of a run-time array"
_ -> pure []
where
sub :: Word32 -> SPIRV.PrimTy -> CGMonad Component
sub i subTy = ( `Component` subTy ) <$> access `continue` [ Access Safe (CTInds [i]) ]
indexed :: Word32 -> SPIRV.PrimTy -> CGMonad [ Component ]
indexed n eltTy = sequence $ zipWith sub [0..] ( replicate (fromIntegral n) eltTy )
-- | Store a value into a component, applying the given operations to it.
storeInto :: CodeGen AST => Pointer -> Component -> Code b -> OpticalOperationTree -> CGMonad ()
storeInto basePtr comp v ops = storeThroughAccessChain' basePtr v =<< ( componentAccess comp `continue` ops )
successiveStores
:: CodeGen AST
=> Pointer
-> [OpticalOperationTree]
-> Codes as
-> CGMonad ()
successiveStores basePtr ( tree : trees ) (ConsCode a as) = do
storeThroughAccessChain' basePtr a tree
successiveStores basePtr trees as
successiveStores _ _ _
= pure ()
setUsingSetter'
:: CodeGen AST
=> (ID, SPIRV.PrimTy)
-> Code a
-> OpticalOperationTree
-> CGMonad (ID, SPIRV.PrimTy)
setUsingSetter' _ val [] = codeGen val
setUsingSetter' base val [Access _ (CTInds is)]
= do
valID <- codeGen val
compositeInsert valID base is
-- in more complex situations, revert to using load/store
setUsingSetter' base@(baseID, _) val ops
= do
(basePtr, basePtrState) <- temporaryVariable base
unless
( basePtrState == Fresh )
( storeInstruction basePtr baseID )
-- set the temporary pointer state to "modified"
-- as we are about to store something into it
setTemporaryPointer baseID (basePtr, Modified)
storeThroughAccessChain' basePtr val ops
loadInstruction basePtr
insertUsingSetter'
:: CodeGen AST
=> ShortText
-> (ID, SPIRV.PrimTy)
-> Code a
-> OpticalOperationTree
-> CGMonad ()
insertUsingSetter' varName base val ops =
assign ( _localBinding varName ) . Just =<< setUsingSetter' base val ops