Layout.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE ScopedTypeVariables #-}

{-| Cross-checks gl-block's std140/std430 layout against the reference
produced by glslangValidator + spirv-cross (see test/tools/gen-goldens.py).

For each corpus struct we poke the real 'Std140'/'Std430' 'Storable' into a
buffer twice: once over 0x00 fill, once over 0xFF fill. A byte was written
by gl-block iff the two results agree (untouched bytes keep their differing
fill), so we recover the exact set of occupied bytes without depending on
the field values -- which also handles 'Bool' serialising as 0x00000001.
That occupancy and the buffer size must match the reference reflection.
-}
module Spec.Layout where

import Control.Exception (SomeException, try)
import Data.List (nub, sort)
import Foreign
import Foreign.C.Types (CDouble, CFloat, CInt, CShort, CUInt, CUShort)
import GHC.Generics (Generic)
import Numeric.Half (Half)

import Graphics.Gl.Block (Block, Packed (..), Std140 (..), Std430 (..))
import Spec.Array (Array, unsafeFromList)
import Spec.Layout.Golden (Layout (..), goldens)
import Spec.Types (Mat4 (..), Vec2 (..), Vec3 (..), Vec3Packed (..), Vec4 (..))

import Test.Tasty
import Test.Tasty.HUnit

-- Corpus: every type mirrors a test/glsl/*.glslfrag struct, field for field.

data Scalars = Scalars Int32 Word32 Float
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

data Doubles = Doubles Float Double Float
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

data Wide64 = Wide64 Word32 Int64 Float Word64
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

{- | @those@ is a tuple, so gl-block lays it out as the @struct Pair@ in
Mixed.glslfrag (not a vec2). @wat@ is a 'Bool' against GLSL @uint@.
-}
data Mixed = Mixed Bool Double Float (Float, Float)
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

data Nested2 = Nested2 (Float, Double) Float
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

{- | A 4-aligned sub-struct followed by a scalar: isolates whether a nested
struct's std140 size is padded up to 16 (which sets @b@'s offset).
-}
data NestPad = NestPad (Float, Float) Float
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

-- | Foreign.C.Types scalars, to cover the @C*@ 'Block' instances.
data CTypes = CTypes CInt CUInt CFloat CDouble
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

{- | 16-bit integer scalars. @c@ (a 4-byte @uint@) follows two 2-byte fields, so
the reference packs @a\/b@ tight (alignment 2) and only then bumps @c@ to a
4-byte boundary -- this pins that a tail field of higher alignment doesn't
over-pad the 16-bit fields before it.
-}
data Shorts16 = Shorts16 Word16 Int16 Word32 Word16
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

-- | The @Foreign.C.Types@ 16-bit family (@CShort@/@CUShort@).
data CShorts16 = CShorts16 CShort CUShort Float
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

{- | A fixed-size 16-bit array: std140 rounds each element's stride up to 16
(64 bytes for 4 elements), while std430 packs them tight (8 bytes).
-}
data ShortArr = ShortArr (Array 4 Word16)
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

{- | Half-precision floats (GLSL @float16_t@). Two 2-byte halves pack tight, then
the 32-bit @c@ bumps to a 4-byte boundary -- same shape as 'CShorts16'.
-}
data Halves = Halves Half Half Float
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

-- | A fixed-size half-precision array (std140 element stride rounds up to 16).
data HalfArr = HalfArr (Array 4 Half)
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

{- | A record (named fields). Same 'Generic' 'Rep' as the positional product, so
it must lay out identically -- this is what locks that in.
-}
data Recordy = Recordy {rA :: Float, rB :: Double, rC :: Int32}
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

-- | A fixed-size array field (std140 rounds each element's stride up to 16).
data FloatArr = FloatArr (Array 4 Float)
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

-- | A scalar followed by a fixed-size array.
data MixArr = MixArr Float (Array 3 Float)
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

{- | A nullary constructor (size 0): exercises the @U1@ generic instance.
Not in the glslang corpus (an empty block isn't valid GLSL) -- round-trip only.
-}
data Unit0 = Unit0
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

-- | Names must match the @goldens@ keys (checked by 'unit_corpusMatchesGoldens').
data Some = forall a. (Block a, Eq a, Show a) => Some a

corpus :: [(String, Some)]
corpus =
  [ ("Scalars", Some (Scalars 1 2 3))
  , ("Doubles", Some (Doubles 1 2 3))
  , ("Wide64", Some (Wide64 1 2 3 4))
  , ("Mixed", Some (Mixed True 1 2 (3, 4)))
  , ("Nested2", Some (Nested2 (1, 2) 3))
  , ("NestPad", Some (NestPad (1, 2) 3))
  , ("CTypes", Some (CTypes 1 2 3 4))
  , ("Shorts16", Some (Shorts16 1 2 3 4))
  , ("CShorts16", Some (CShorts16 1 2 3))
  , ("ShortArr", Some (ShortArr (unsafeFromList [1, 2, 3, 4])))
  , ("Halves", Some (Halves 1 2 3))
  , ("HalfArr", Some (HalfArr (unsafeFromList [1, 2, 3, 4])))
  , ("Recordy", Some (Recordy 1 2 3))
  , ("FloatArr", Some (FloatArr (unsafeFromList [1, 2, 3, 4])))
  , ("MixArr", Some (MixArr 1 (unsafeFromList [2, 3, 4])))
  ]

-- | The set of bytes gl-block actually writes, via the two-fill trick.
occupiedMask :: (Storable w) => w -> IO [Bool]
occupiedMask w = do
  let n = sizeOf w
  allocaBytes n $ \p0 -> allocaBytes n $ \pF -> do
    fillBytes p0 0x00 n
    fillBytes pF 0xFF n
    poke (castPtr p0 :: Ptr w) w
    poke (castPtr pF :: Ptr w) w
    b0 <- peekArray n (p0 :: Ptr Word8)
    bF <- peekArray n (pF :: Ptr Word8)
    pure (zipWith (==) b0 bF)

-- | Bytes covered by the reference's scalar leaves, over a buffer of @stride@.
goldenMask :: Layout -> [Bool]
goldenMask lay =
  [ any (\(o, s) -> o <= i && i < o + s) (gLeaves lay)
  | i <- [0 .. gStride lay - 1]
  ]

{- | Collapse a byte mask to @(offset, length)@ runs -- a faithful, readable
encoding of the mask for assertion messages.
-}
runs :: [Bool] -> [(Int, Int)]
runs = go 0
  where
    go _ [] = []
    go i bs =
      let
        (gap, rest) = span not bs
        i' = i + length gap
        (run, rest') = span id rest
      in
        if null run then
          []
        else
          (i', length run) : go (i' + length run) rest'

checkLayout :: (Storable w) => Layout -> w -> Assertion
checkLayout lay w = do
  assertEqual "padded size (array stride)" (gStride lay) (sizeOf w)
  mask <- occupiedMask w
  assertEqual
    "occupied byte ranges (offset, length)"
    (runs (goldenMask lay))
    (runs mask)

{- | poke-then-peek must recover the value, for every layout. This is what
exercises the read side (read140\/read430\/readPacked and the generic struct
traversal), plus the packed size\/write paths.
-}
roundTrip :: (Block a, Eq a, Show a) => a -> Assertion
roundTrip v = do
  Std140 v140 <- with (Std140 v) peek
  v140 @?= v
  Std430 v430 <- with (Std430 v) peek
  v430 @?= v
  Packed vP <- with (Packed v) peek
  vP @?= v

test_layout :: TestTree
test_layout = testGroup "layout vs glslang/spirv-cross" (map perStruct corpus)
  where
    perStruct (name, Some a) =
      case lookup3 name goldens of
        Nothing -> testCase name $ assertFailure ("no golden for " ++ name)
        Just (g140, g430) ->
          testGroup
            name
            [ testCase "std140" $ checkLayout g140 (Std140 a)
            , testCase "std430" $ checkLayout g430 (Std430 a)
            , testCase "roundtrip" $ roundTrip a
            ]

lookup3 :: String -> [(String, b, c)] -> Maybe (b, c)
lookup3 k xs = case [(b, c) | (n, b, c) <- xs, n == k] of
  (x : _) -> Just x
  [] -> Nothing

-- | Guards against the corpus and the generated goldens drifting apart.
unit_corpusMatchesGoldens :: Assertion
unit_corpusMatchesGoldens =
  sort (nub (map fst corpus ++ [g | GCase _ g _ _ <- geomCases]))
    @?= sort [name | (name, _, _) <- goldens]

-- | A zero-field block round-trips (and its size is 0).
unit_emptyBlock :: Assertion
unit_emptyBlock = do
  sizeOf (Std140 Unit0) @?= 0
  roundTrip Unit0

-- | Covers the 3-tuple 'Block' instance (the corpus only uses 2-tuples).
unit_tuple3 :: Assertion
unit_tuple3 = roundTrip ((1, 2, 3) :: (Float, Float, Float))

-- Wrapper structs, mirroring the test/glsl/{Vecs,MatVec,Vec3Tail,Vec3Pack}.glslfrag.

data Vecs = Vecs Vec2 Vec3Packed Vec4
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

data MatVec = MatVec Mat4 Vec4
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

data Vec3TailP = Vec3TailP Float Vec3Packed
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

data Vec3TailD = Vec3TailD Float Vec3
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

data Vec3PackS = Vec3PackS Vec3Packed Float
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

{- | @float float vec2@: a low-alignment field (float, align 4) precedes a
higher-alignment sibling (vec2, align 8) in the same right-leaning product
subtree -- the same shape that over-padded 'CShorts16'\/'Halves'. The second
float must stay at offset 4, not be bumped to 8.
-}
data Vec2Tail = Vec2Tail Float Float Vec2
  deriving stock (Eq, Show, Generic)
  deriving anyclass (Block)

-- | Whether a specimen is expected to match glslang, or to diverge (and why).
data Mode = Conform | Diverges String

-- | (case label, golden key, mode, value)
data GCase = forall a. (Block a, Eq a, Show a) => GCase String String Mode a

geomCases :: [GCase]
geomCases =
  [ GCase
      "Vecs"
      "Vecs"
      Conform
      (Vecs (Vec2 1 2) (Vec3Packed (Vec3 3 4 5)) (Vec4 6 7 8 9))
  , GCase
      "MatVec"
      "MatVec"
      Conform
      (MatVec (Mat4 [1 .. 16]) (Vec4 1 2 3 4))
  , GCase
      "Vec3Tail (packed)"
      "Vec3Tail"
      Conform
      (Vec3TailP 1 (Vec3Packed (Vec3 2 3 4)))
  , GCase
      "Vec3Tail (default, writes 4th lane)"
      "Vec3Tail"
      (Diverges "default Vec3 writes its 4th lane (1.0): 16 occupied bytes vs the spec's 12")
      (Vec3TailD 1 (Vec3 2 3 4))
  , GCase
      "Vec3Pack (no vec3+scalar packing)"
      "Vec3Pack"
      (Diverges "sizeOf140(vec3)=16, so a vec3+scalar does not pack (glslang places the scalar at 12)")
      (Vec3PackS (Vec3Packed (Vec3 1 2 3)) 4)
  , GCase
      "Vec2Tail (float float vec2)"
      "Vec2Tail"
      Conform
      (Vec2Tail 1 2 (Vec2 3 4))
  ]

-- | Passes iff @act@ fails, i.e. gl-block diverges from glslang as expected.
expectMismatch :: String -> Assertion -> Assertion
expectMismatch why act = do
  res <- try act
  case res of
    Left (_ :: SomeException) -> pure ()
    Right () ->
      assertFailure
        ( "expected a known divergence ("
            ++ why
            ++ ") but gl-block now matches glslang -- promote this case"
        )

test_layoutGeom :: TestTree
test_layoutGeom =
  testGroup "geomancy vec/mat specimens vs glslang" (map perCase geomCases)
  where
    perCase (GCase name gname mode a) =
      case lookup3 gname goldens of
        Nothing -> testCase name $ assertFailure ("no golden for " ++ gname)
        Just (g140, g430) ->
          testGroup
            name
            [ testCase "std140" $ run mode g140 (Std140 a)
            , testCase "std430" $ run mode g430 (Std430 a)
            , testCase "roundtrip" $ roundTrip a
            ]
    run Conform lay w = checkLayout lay w
    run (Diverges why) lay w = expectMismatch why (checkLayout lay w)