Internal.chs

{-# LANGUAGE MagicHash #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Strict #-}

module Data.SpirV.Reflect.FFI.Internal where

import Data.Coerce (Coercible, coerce)
import Data.Ord (comparing)
import Data.SpirV.Enum qualified as SpirV
import Data.SpirV.Enum.StorageClass qualified as StorageClass
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Vector (Vector)
import Data.Vector qualified as Vector
import Data.Vector.Algorithms.Heap qualified as Heap
import Data.Word (Word32)
import Data.IORef (IORef, newIORef, readIORef, modifyIORef')
import Foreign.C.String (CString)
import Foreign.C.Types (CULong)
import Foreign.Marshal.Utils (maybePeek)
import Foreign.Ptr (Ptr, castPtr, plusPtr, nullPtr)
import Foreign.Storable (peek, peekElemOff)
import GHC.Ptr qualified as GHC

import Data.SpirV.Reflect.BlockVariable (BlockVariable)
import Data.SpirV.Reflect.BlockVariable qualified as BlockVariable
import Data.SpirV.Reflect.DescriptorBinding (DescriptorBinding)
import Data.SpirV.Reflect.DescriptorBinding qualified as DescriptorBinding
import Data.SpirV.Reflect.DescriptorSet (DescriptorSet)
import Data.SpirV.Reflect.DescriptorSet qualified as DescriptorSet
import Data.SpirV.Reflect.Enums qualified as Reflect
import Data.SpirV.Reflect.InterfaceVariable (InterfaceVariable)
import Data.SpirV.Reflect.InterfaceVariable qualified as InterfaceVariable
import Data.SpirV.Reflect.Module (Module)
import Data.SpirV.Reflect.Module qualified as Module
import Data.SpirV.Reflect.Traits qualified as Traits
import Data.SpirV.Reflect.TypeDescription (TypeDescription)
import Data.SpirV.Reflect.TypeDescription qualified as TypeDescription
import Data.SpirV.Reflect.SpecializationConstant (SpecializationConstant)
import Data.SpirV.Reflect.SpecializationConstant qualified as SpecializationConstant

#include "spirv_reflect.h"

-- * Loader interface

{#pointer *SpvReflectShaderModule as ShaderModulePtr #}

{#enum
  SpvReflectResult as Result
  {underscoreToCase}
  deriving (Eq, Ord, Show)
#}

{#fun unsafe
  spvReflectCreateShaderModule as createShaderModule
  { id `CULong'
  , `Ptr ()'
  , `ShaderModulePtr'
  } -> `Result'
#}

{#enum
  SpvReflectModuleFlagBits as ModuleFlags
  {underscoreToCase}
  deriving (Eq, Ord, Show)
#}

{#fun unsafe
  spvReflectCreateShaderModule2 as createShaderModule2
  { `ModuleFlags'
  , id `CULong'
  , `Ptr ()'
  , `ShaderModulePtr'
  } -> `Result'
#}

{#fun unsafe
  spvReflectDestroyShaderModule as destroyShaderModule
  { `ShaderModulePtr'
  } -> `()'
#}

-- TODO: inflateEntryPoints :: ShaderModulePtr -> IO [EntryPoint]

-- * Module

shaderModuleSize :: Int
shaderModuleSize = {#sizeof SpvReflectShaderModule #}

inflateModule :: ShaderModulePtr -> IO Module
inflateModule smp = do
  let sm = castPtr smp

  generator <- inflateEnum $
    {#get SpvReflectShaderModule->generator #} sm

  entry_point_name <- inflateText $
    {#get SpvReflectShaderModule->entry_point_name #} sm

  entry_point_id <- inflateIntegral $
    {#get SpvReflectShaderModule->entry_point_id #} sm

  -- TODO: enums
  source_language <- inflateIntegral $
    {#get SpvReflectShaderModule->source_language #} sm

  source_language_version <- inflateIntegral $
    {#get SpvReflectShaderModule->source_language_version #} sm

  -- XXX: Uses value(s) from first entry point

  -- TODO: flags
  spirv_execution_model <- inflateIntegral $
    {#get SpvReflectShaderModule->spirv_execution_model #} sm

  -- TODO: flags
  shader_stage <- inflateIntegral $
    {#get SpvReflectShaderModule->shader_stage #} sm

  descriptor_bindings <-
    inflateVector
      ({#get SpvReflectShaderModule->descriptor_binding_count #} sm)
      ({#get SpvReflectShaderModule->descriptor_bindings #} sm)
      {#sizeof SpvReflectDescriptorBinding #}
      inflateDescriptorBinding

  descriptor_sets <-
    inflateVector
      ({#get SpvReflectShaderModule->descriptor_set_count #} sm)
      ({#get SpvReflectShaderModule->descriptor_sets #} sm)
      {#sizeof SpvReflectDescriptorSet #}
      inflateDescriptorSet

  interface_variables <-
    inflateVector
      ({#get SpvReflectShaderModule->interface_variable_count #} sm)
      ({#get SpvReflectShaderModule->interface_variables #} sm)
      {#sizeof SpvReflectInterfaceVariable #}
      inflateInterfaceVariable

  let
    ivLocation InterfaceVariable.InterfaceVariable{location} =
      location

    pickIvs query =
      Vector.modify (Heap.sortBy $ comparing ivLocation) $ Vector.filter query interface_variables

    input_variables =
      pickIvs $
        (== StorageClass.Input) . InterfaceVariable.storage_class

    output_variables =
      pickIvs $
        (== StorageClass.Output) . InterfaceVariable.storage_class

  push_constants <-
    inflateVector
      ({#get SpvReflectShaderModule->push_constant_block_count #} sm)
      ({#get SpvReflectShaderModule->push_constant_blocks #} sm)
      {#sizeof SpvReflectBlockVariable #}
      inflateBlockVariable

  spec_constants <-
    inflateVector
      ({#get SpvReflectShaderModule->spec_constant_count #} sm)
      ({#get SpvReflectShaderModule->spec_constants #} sm)
      {#sizeof SpvReflectSpecializationConstant #}
      inflateSpecConstant

  pure Module.Module{..}

{#pointer *SpvReflectDescriptorBinding as DescriptorBindingPtr #}

inflateDescriptorBinding :: DescriptorBindingPtr -> IO DescriptorBinding
inflateDescriptorBinding db = do
  spirv_id <- fmap just . inflateIntegral $
    {#get SpvReflectDescriptorBinding->spirv_id #} db

  name <- inflateText $
    {#get SpvReflectDescriptorBinding->name #} db

  binding <- inflateIntegral $
    {#get SpvReflectDescriptorBinding->binding #} db

  input_attachment_index <- inflateIntegral $
    {#get SpvReflectDescriptorBinding->input_attachment_index #} db

  set <- inflateIntegral $
    {#get SpvReflectDescriptorBinding->set #} db

  descriptor_type <- inflateEnum $
    {#get SpvReflectDescriptorBinding->descriptor_type #} db

  resource_type <- inflateFlags32 $
    {#get SpvReflectDescriptorBinding->resource_type #} db

  image <-
    inflateImageTraits
      db
      {#offsetof SpvReflectDescriptorBinding->image #}

  block <-
    {#get SpvReflectDescriptorBinding->block #} db >>=
      maybePeek inflateBlockVariable

  array <-
    inflateArrayTraits
      db
      {#offsetof SpvReflectDescriptorBinding->array #}

  count <- fmap just . inflateIntegral $
    {#get SpvReflectDescriptorBinding->count #} db

  accessed <- inflateIntegral $
    {#get SpvReflectDescriptorBinding->accessed #} db

  uav_counter_id <- inflateIntegral $
    {#get SpvReflectDescriptorBinding->uav_counter_id #} db

  uav_counter_binding <-
    {#get SpvReflectDescriptorBinding->uav_counter_binding #} db >>=
      maybePeek inflateDescriptorBinding

  type_description <-
    {#get SpvReflectDescriptorBinding->type_description #} db >>=
      maybePeek inflateTypeDescription

  let word_offset = DescriptorBinding.WordOffset{..}

  decoration_flags <- inflateFlags32 $
    {#get SpvReflectDescriptorBinding->decoration_flags #} db

  byte_address_buffer_offsets <-
    inflateVector
      ({#get SpvReflectDescriptorBinding->byte_address_buffer_offset_count #} db)
      ({#get SpvReflectDescriptorBinding->byte_address_buffer_offsets #} db)
      {#sizeof uint32_t #}
      (fmap fromIntegral . peek)

  user_type <- inflateMaybe (Reflect.UserType 0) inflateEnum $
    {#get SpvReflectDescriptorBinding->user_type #} db

  pure DescriptorBinding.DescriptorBinding{..}

inflateSpecConstant :: Ptr () -> IO SpecializationConstant
inflateSpecConstant sc = do
  spirv_id <- fmap just . inflateIntegral $
    {#get SpvReflectSpecializationConstant->spirv_id #} sc

  constant_id <- inflateIntegral $
    {#get SpvReflectSpecializationConstant->constant_id #} sc

  name <- inflateMaybe "" inflateText $
    {#get SpvReflectSpecializationConstant->name #} sc

  type_description <-
    {#get SpvReflectSpecializationConstant->type_description #} sc >>=
      maybePeek inflateTypeDescription

  default_value <- do
    size <- inflateIntegral $
      {#get SpvReflectSpecializationConstant->default_value_size #} sc
    valuePtr <- {#get SpvReflectSpecializationConstant->default_value #} sc
    if valuePtr == nullPtr then
      pure Vector.empty
    else
      -- default_value_size is a byte count, always a multiple of 4; copy it out
      -- as 32-bit words (low-order word first) while the SPIR-V buffer is alive.
      Vector.generateM (size `div` 4) $
        peekElemOff (castPtr valuePtr)

  pure SpecializationConstant.SpecializationConstant {..}

{#pointer *SpvReflectBlockVariable as BlockVariablePtr #}

inflateBlockVariable :: BlockVariablePtr -> IO BlockVariable
inflateBlockVariable bv0 = do
  cache <- newIORef []
  inflateBlockVariableCached cache bv0

-- | Inflate a 'BlockVariable', breaking cycles the same way
-- 'inflateTypeDescriptionCached' does. A @buffer_reference@ /
-- @PhysicalStorageBuffer@ member points back into its own block, so a naive deep
-- copy of @members@ never terminates. We three-colour keyed by the
-- 'SpvReflectBlockVariable' pointer (as for 'TypeDescription'): the C side builds
-- the recursive block as a cyclic graph of shared pointers, so pointer identity
-- detects the back-edge, and — unlike @spirv_id@, which is 0 for every member —
-- it never conflates two distinct members (each carries its own @offset@/@name@).
-- A back-edge into a block still on the current path becomes a shallow node (no
-- members), so acyclic shaders are unchanged while recursive ones stay finite.
inflateBlockVariableCached
  :: IORef [(BlockVariablePtr, Maybe BlockVariable)]
  -> BlockVariablePtr
  -> IO BlockVariable
inflateBlockVariableCached cache bv = do
  colour <- fmap (lookup bv) (readIORef cache)
  case colour of
    Just (Just black) ->
      pure black
    _grey_or_white -> do
      spirv_id <- fmap just . inflateIntegral $
        {#get SpvReflectBlockVariable->spirv_id #} bv

      name <- inflateMaybe "" inflateText $
        {#get SpvReflectBlockVariable->name #} bv

      offset <- inflateIntegral $
        {#get SpvReflectBlockVariable->offset #} bv

      absolute_offset <- inflateIntegral $
        {#get SpvReflectBlockVariable->absolute_offset #} bv

      size <- inflateIntegral $
        {#get SpvReflectBlockVariable->size #} bv

      padded_size <- inflateIntegral $
        {#get SpvReflectBlockVariable->padded_size #} bv

      decorations <- inflateFlags32 $
        {#get SpvReflectBlockVariable->decoration_flags #} bv

      numeric <-
        inflateNumericTraits
          bv
          {#offsetof SpvReflectBlockVariable->numeric #}

      array <-
        inflateArrayTraits
          bv
          {#offsetof SpvReflectBlockVariable->array #}

      word_offset <-
        fmap (Just . BlockVariable.WordOffset) . inflateIntegral $
          {#get SpvReflectBlockVariable->word_offset.offset #} bv

      case colour of
        Just Nothing ->
          -- grey: a back-edge into a block still being inflated; cut the cycle.
          pure BlockVariable.BlockVariable
            {spirv_id, members=Vector.empty, type_description=Nothing, ..}
        _white -> do
          modifyIORef' cache ((bv, Nothing) :) -- mark grey before descending

          members <-
            inflateVector
              ({#get SpvReflectBlockVariable->member_count #} bv)
              ({#get SpvReflectBlockVariable->members #} bv)
              {#sizeof SpvReflectBlockVariable #}
              (inflateBlockVariableCached cache)

          type_description <-
            {#get SpvReflectBlockVariable->type_description #} bv >>=
              maybePeek inflateTypeDescription

          let result = BlockVariable.BlockVariable{spirv_id, ..}
          modifyIORef' cache ((bv, Just result) :) -- mark black, memoise
          pure result

{#pointer *SpvReflectTypeDescription as TypeDescriptionPtr #}

inflateTypeDescription :: TypeDescriptionPtr -> IO TypeDescription
inflateTypeDescription td0 = do
  cache <- newIORef []
  inflateTypeDescriptionCached cache td0

-- | Inflate a 'TypeDescription', breaking self-referential cycles.
--
-- Mutually recursive @buffer_reference@ / @PhysicalStorageBuffer@ pointer types
-- form a cyclic graph (a struct whose members point back at it through
-- 'struct_type_description' / 'members'). The C side keeps types in a flat list
-- linked by reference, but here we copy the graph into a tree, so a naive deep
-- copy never terminates.
--
-- We walk the graph as a DFS with three-colouring, keyed by the
-- 'SpvReflectTypeDescription' pointer (identity is by pointer, not @id@: the C
-- side reuses one @id@ across a canonical type and its copied references):
--
--   * /white/ (unseen): expand fully, recursing into children.
--   * /grey/ (on the current path): a genuine back-edge — emit a shallow node
--     (scalar fields only, no children) to cut the cycle.
--   * /black/ (already finished): return the cached value, so a type referenced
--     more than once is built only once and shared.
--
-- Acyclic shaders never hit a grey node, and black hits return the same fully
-- expanded value, so their result is byte-for-byte unchanged while recursive
-- shaders stay finite (each distinct node is inflated at most once).
inflateTypeDescriptionCached
  :: IORef [(TypeDescriptionPtr, Maybe TypeDescription)]
  -> TypeDescriptionPtr
  -> IO TypeDescription
inflateTypeDescriptionCached cache td = do
  colour <- fmap (lookup td) (readIORef cache)
  case colour of
    Just (Just black) ->
      pure black
    _grey_or_white -> do
      id_ <- fmap just . inflateIntegral $
        {#get SpvReflectTypeDescription->id #} td

      op <- fmap (just . SpirV.Op) . inflateIntegral $
        {#get SpvReflectTypeDescription->op #} td

      type_name <- inflateMaybe "" inflateText $
        {#get SpvReflectTypeDescription->type_name #} td

      struct_member_name <- inflateMaybe "" inflateText $
        {#get SpvReflectTypeDescription->struct_member_name #} td

      storage_class <- fmap SpirV.StorageClass . inflateIntegral $
        {#get SpvReflectTypeDescription->storage_class #} td

      type_flags <- inflateFlags32 $
        {#get SpvReflectTypeDescription->type_flags #} td

      numeric <-
        inflateNumericTraits
          td
          {#offsetof SpvReflectTypeDescription->traits.numeric #}

      image <-
        inflateImageTraits
          td
          {#offsetof SpvReflectTypeDescription->traits.image #}

      array <-
        inflateArrayTraits
          td
          {#offsetof SpvReflectTypeDescription->traits.array #}

      let
        traits' = Just TypeDescription.Traits{..}
        traits = if traits' == Just TypeDescription.emptyTraits then Nothing else traits'

      copied <-
        fmap (\n -> if n > 0 then Just n else Nothing) . inflateIntegral $
          {#get SpvReflectTypeDescription->copied #} td

      case colour of
        Just Nothing -> do
          -- grey: a back-edge into a type still being inflated; cut the cycle.
          pure TypeDescription.TypeDescription
            {id=id_, struct_type_description=Nothing, members=Vector.empty, ..}
        _white -> do
          modifyIORef' cache ((td, Nothing) :) -- mark grey before descending

          struct_type_description <-
            {#get SpvReflectTypeDescription->struct_type_description #} td >>=
              maybePeek (inflateTypeDescriptionCached cache)

          -- XXX: deprecated, but struct_type_description isn't always available for structs
          members <-
            inflateVector
              ({#get SpvReflectTypeDescription->member_count #} td)
              ({#get SpvReflectTypeDescription->members #} td)
              {#sizeof SpvReflectTypeDescription #}
              (inflateTypeDescriptionCached cache)

          let result = TypeDescription.TypeDescription{id=id_, ..}
          modifyIORef' cache ((td, Just result) :) -- mark black, memoise
          pure result

{#pointer *SpvReflectDescriptorSet as DescriptorSetPtr #}

inflateDescriptorSet :: DescriptorSetPtr -> IO DescriptorSet
inflateDescriptorSet ds = do
  set <- inflateIntegral $
    {#get SpvReflectDescriptorSet->set #} ds

  bindingsPtr <- {#get SpvReflectDescriptorSet->bindings #} ds
  bindings <-
    inflateVector
      ({#get SpvReflectDescriptorSet->binding_count #} ds)
      (peek bindingsPtr)
      {#sizeof SpvReflectDescriptorBinding #}
      inflateDescriptorBinding

  pure DescriptorSet.DescriptorSet{..}

{#pointer *SpvReflectInterfaceVariable as InterfaceVariablePtr #}

inflateInterfaceVariable :: InterfaceVariablePtr -> IO InterfaceVariable
inflateInterfaceVariable iv = do
  spirv_id <- fmap just . inflateIntegral $
    {#get SpvReflectInterfaceVariable->spirv_id #} iv

  name <- inflateMaybe "" inflateText $
    {#get SpvReflectInterfaceVariable->name #} iv

  location <- inflateIntegral $
    {#get SpvReflectInterfaceVariable->location #} iv

  component <- inflateMaybe 4294967295 inflateIntegral $
    {#get SpvReflectInterfaceVariable->component #} iv

  storage_class <- fmap SpirV.StorageClass . inflateIntegral $
    {#get SpvReflectInterfaceVariable->storage_class #} iv

  semantic <- inflateMaybe "" inflateText $
    {#get SpvReflectInterfaceVariable->semantic #} iv

  decoration_flags <- inflateFlags32 $
    {#get SpvReflectInterfaceVariable->decoration_flags #} iv

  -- SPIRV-Reflect's invalid sentinel is 0xFFFFFFFF, which decodes to BuiltIn (-1)
  -- a non-built-in variable becomes Nothing.
  built_in <- inflateMaybe (SpirV.BuiltIn (-1)) (fmap SpirV.BuiltIn . inflateIntegral) $
    {#get SpvReflectInterfaceVariable->built_in #} iv

  numeric <-
    inflateNumericTraits
      iv
      {#offsetof SpvReflectInterfaceVariable->numeric #}

  array <-
    inflateArrayTraits
      iv
      {#offsetof SpvReflectInterfaceVariable->array #}

  members <-
    inflateVector
      ({#get SpvReflectInterfaceVariable->member_count #} iv)
      ({#get SpvReflectInterfaceVariable->members #} iv)
      {#sizeof SpvReflectInterfaceVariable #}
      inflateInterfaceVariable

  format <- inflateEnum $
    {#get SpvReflectInterfaceVariable->format #} iv

  type_description <-
    {#get SpvReflectInterfaceVariable->type_description #} iv >>=
      maybePeek inflateTypeDescription

  wo_location <- inflateIntegral $
    {#get SpvReflectInterfaceVariable->word_offset.location #} iv
  let word_offset = InterfaceVariable.WordOffset{location=wo_location}

  pure InterfaceVariable.InterfaceVariable{..}

-- * Traits

{#pointer *SpvReflectImageTraits as ImageTraitsPtr #}

inflateImageTraits :: Ptr struct -> Int -> IO Traits.Image
inflateImageTraits src offset = do
  let it = castPtr src `plusPtr` offset

  dim <- fmap SpirV.Dim . inflateIntegral $
    {#get SpvReflectImageTraits->dim #} it

  depth <- inflateIntegral $
    {#get SpvReflectImageTraits->depth #} it

  arrayed <- inflateIntegral $
    {#get SpvReflectImageTraits->arrayed #} it

  ms <- inflateIntegral $
    {#get SpvReflectImageTraits->ms #} it

  sampled <- inflateIntegral $
    {#get SpvReflectImageTraits->sampled #} it

  image_format <- fmap SpirV.ImageFormat . inflateEnum $
    {#get SpvReflectImageTraits->image_format #} it

  pure Traits.Image{..}

{#pointer *SpvReflectNumericTraits as NumericTraitsPtr #}

inflateNumericTraits :: Ptr struct -> Int -> IO Traits.Numeric
inflateNumericTraits src offset = do
  let nt = castPtr src `plusPtr` offset

  width <- inflateIntegral $
    {#get SpvReflectNumericTraits.scalar.width #} nt
  signedness <- inflateIntegral $
    {#get SpvReflectNumericTraits.scalar.signedness #} nt
  let signed = signedness /= (0 :: C2HSImp.CUInt)
  let scalar = Traits.Scalar{..}

  component_count <- inflateIntegral $
    {#get SpvReflectNumericTraits->vector.component_count #} nt
  let vector = Traits.Vector{..}

  column_count <- inflateIntegral $
    {#get SpvReflectNumericTraits->matrix.column_count #} nt
  row_count <- inflateIntegral $
    {#get SpvReflectNumericTraits->matrix.row_count #} nt
  stride <- inflateIntegral $
    {#get SpvReflectNumericTraits->matrix.stride #} nt
  let matrix = Traits.Matrix{..}

  pure Traits.Numeric{..}

{#pointer *SpvReflectArrayTraits as ArrayTraitsPtr #}

inflateArrayTraits :: Ptr struct -> Int -> IO Traits.Array
inflateArrayTraits src offset = do
  let at = castPtr src `plusPtr` offset

  dims_count <- inflateIntegral $
    {#get SpvReflectArrayTraits->dims_count #} at

  dims <- fmap Vector.convert $ inflateVector
    ({#get SpvReflectArrayTraits->dims_count #} at)
    ({#get SpvReflectArrayTraits->dims #} at)
    {#sizeof uint32_t #}
    (fmap fromIntegral . peek)

  stride <- inflateMaybe 0 inflateIntegral $
    {#get SpvReflectArrayTraits->stride #} at

  pure Traits.Array{..}

-- * Primitive types

inflateVector :: Integral i => IO i -> IO (Ptr p) -> Int -> (Ptr p -> IO a) -> IO (Vector a)
inflateVector getCount getItems itemSize inflate = do
  count <- getCount
  itemsPtr <- getItems
  Vector.generateM (fromIntegral count) \pos -> do
    x <- inflate $ itemsPtr `plusPtr` (itemSize * pos)
    pure $! x

inflateIntegral :: (Integral a, Num b) => IO a -> IO b
inflateIntegral getIntegral =
  getIntegral >>= \i ->
    pure $! fromIntegral i

inflateEnum :: (Integral a, Enum b) => IO a -> IO b
inflateEnum getEnum =
  getEnum >>= \i ->
    pure $! toEnum (fromIntegral i)

inflateFlags32 :: forall a b . (Integral a, Coercible Word32 b) => IO a -> IO b
inflateFlags32 gitBits =
  gitBits >>= \i ->
    pure $! coerce @Word32 @b (fromIntegral i)

inflateText :: IO CString -> IO Text
inflateText getPtr =
  getPtr >>= \ptr ->
    if nullPtr == ptr then
      pure mempty
    else
      case ptr of
        GHC.Ptr addr ->
          pure $! Text.unpackCString# addr

inflateMaybe :: Eq b => b -> (IO a -> IO b) -> IO a -> IO (Maybe b)
inflateMaybe nothingBurger inflate = fmap (\x -> if x == nothingBurger then Nothing else Just x) . inflate

just :: a -> Maybe a
just x = x `seq` Just x