PrimTy.hs

{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GADTs              #-}
{-# LANGUAGE NamedFieldPuns     #-}

{-|
Module: SPIRV.PrimTy

This module enumerates the types that exist within the SPIR-V type system.

These are logical types: they don't include any memory layout information.
The layout of a type in memory is specified separately, by a 'Layout'.

See the SPIR-V specification, §3.22.6 __Type-Declaration instructions__.
-}

module SPIRV.PrimTy
  ( PrimTy(..)
  , PointerTy(..)
  , componentLayout
  , tyOp
  , scalars
  ) where

-- base
import Data.Word
  ( Word32 )
import Prelude
  hiding (Integer, Floating)

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

-- fir
import SPIRV.Image
  ( Image(texelComponent) )
import SPIRV.Layout
  ( Layout(blockKind) )
import SPIRV.Operation
  hiding ( Function, Image, SampledImage )
import SPIRV.ScalarTy
  ( ScalarTy(Integer, Floating) )
import SPIRV.Storage
  ( StorageClass )

--------------------------------------------------
-- SPIR-V types

data PrimTy where
  Unit    :: PrimTy -- known as Void in the SPIR-V specification
  Boolean :: PrimTy
  Scalar  :: ScalarTy -> PrimTy
  Vector  ::
    { size  :: Word32
    , eltTy :: PrimTy
    } -> PrimTy
  Matrix ::
    { rows    :: Word32
    , cols    :: Word32
    , entryTy :: ScalarTy
    } -> PrimTy
  Array ::
    { size  :: Word32
    , eltTy :: PrimTy
    } -> PrimTy
  RuntimeArray ::
    { eltTy :: PrimTy
    } -> PrimTy
  Struct ::
    { eltTys :: [ ( Maybe ShortText, PrimTy ) ]
    -- ^ the (optional) field names are only used to annotate the source code
    } -> PrimTy
  Pointer  :: PointerTy -> PrimTy
  Function ::
    { argumentTypes :: [PrimTy]
    , resultType    :: PrimTy
    } -> PrimTy
  Image        :: Image -> PrimTy
  Sampler      ::          PrimTy -- opaque
  SampledImage :: Image -> PrimTy
  AccelerationStructure :: PrimTy -- opaque
  RayQuery              :: PrimTy -- opaque
  deriving stock ( Show, Eq, Ord )

-- | A pointer type, to a type laid out in memory in the given way.
data PointerTy
  = PointerTy
  { pointerStorage :: StorageClass
  , pointerLayout  :: Layout
  , pointee        :: PrimTy
  }
  deriving stock ( Show, Eq, Ord )

-- | The layout of the components of a composite type with the given layout.
--
-- The elements of an array of interface blocks are interface blocks,
-- but the members of an interface block are not.
componentLayout :: PrimTy -> Layout -> Layout
componentLayout Array        {} lay = lay
componentLayout RuntimeArray {} lay = lay
componentLayout _               lay = lay { blockKind = Nothing }

tyOp :: PrimTy -> Operation
tyOp Unit                    = TypeVoid
tyOp Boolean                 = TypeBool
tyOp (Scalar Integer   {})   = TypeInt
tyOp (Scalar Floating  {})   = TypeFloat
tyOp (Vector _ (Vector _ _)) = TypeMatrix
  -- SPIR-V doesn't have vectors of vectors, so declare them as matrices instead.
  -- This is a bit of a workaround for issue #95.
tyOp Vector             {}   = TypeVector
tyOp Matrix             {}   = TypeMatrix
tyOp Array              {}   = TypeArray
tyOp RuntimeArray       {}   = TypeRuntimeArray
tyOp Struct             {}   = TypeStruct
tyOp Function           {}   = TypeFunction
tyOp Pointer            {}   = TypePointer
tyOp Image              {}   = TypeImage
tyOp Sampler                 = TypeSampler
tyOp SampledImage       {}   = TypeSampledImage
tyOp AccelerationStructure   = TypeAccelerationStructure
tyOp RayQuery                = TypeRayQuery

scalars :: PrimTy -> [ ScalarTy ]
scalars Unit                     = [ ]
scalars Boolean                  = [ ]
scalars (Scalar s)               = [ s ]
scalars (Vector _ a)             = scalars a
scalars (Matrix _ _ a)           = [ a ]
scalars (Array { eltTy })        = scalars eltTy
scalars (RuntimeArray { eltTy }) = scalars eltTy
scalars (Struct { eltTys })      = scalars . snd =<< eltTys
scalars (Pointer ptrTy)          = scalars ( pointee ptrTy )
scalars (Function as b)          = scalars b ++ ( scalars =<< as )
scalars (Image img)              = [ texelComponent img ]
scalars Sampler                  = [ ]
scalars (SampledImage img)       = [ texelComponent img ]
scalars AccelerationStructure    = []
scalars RayQuery                 = []