-- | Vertex and index writers for the draw buffers. Each vertex is written with
-- two 128-bit GHC SIMD stores (FloatX4#) instead of eight scalar stores.
module NanoUI.SIMD
  ( pokeVertexSIMD
  , pokeQuadSIMD
  , pokeQuadGradientSIMD
  , concentricOffsetsSIMD
  ) where

import GHC.Ptr (Ptr (..))
import Foreign.Storable (pokeByteOff)
import GHC.Exts
  ( Float (F#)
  , Int (I#)
  , packFloatX4#
  , packWord32X4#
  , plusAddr#
  , writeFloatOffAddrAsFloatX4#
  , writeWord32OffAddrAsWord32X4#
  )
import GHC.Word (Word32 (W32#))
import GHC.IO (IO (..))
import Data.Word (Word8)

-- | Writes one 32-byte Vertex (8 floats) into memory using two 128-bit SIMD stores
-- instead of 8 scalar stores.
{-# INLINE pokeVertexSIMD #-}
pokeVertexSIMD ::
  Ptr Word8 ->
  Int ->
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  IO ()
pokeVertexSIMD :: Ptr Word8
-> Int
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> IO ()
pokeVertexSIMD (Ptr Addr#
addr#) (I# Int#
byteOff#) (F# Float#
px#) (F# Float#
py#) (F# Float#
r#) (F# Float#
g#) (F# Float#
b#) (F# Float#
a#) (F# Float#
u#) (F# Float#
v#) = (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, () #)) -> IO ())
-> (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s0 ->
  -- Offsets are recomputed inline (the address add is a single lea) so the
  -- simplified body stays free of let bindings; the inspection test guards
  -- this with a NoAllocation obligation.
  case (# Float#, Float#, Float#, Float# #) -> FloatX4#
packFloatX4# (# Float#
px#, Float#
py#, Float#
r#, Float#
g# #) of
    FloatX4#
v0# ->
      case (# Float#, Float#, Float#, Float# #) -> FloatX4#
packFloatX4# (# Float#
b#, Float#
a#, Float#
u#, Float#
v# #) of
        FloatX4#
v1# ->
          case Addr# -> Int# -> FloatX4# -> State# RealWorld -> State# RealWorld
forall d. Addr# -> Int# -> FloatX4# -> State# d -> State# d
writeFloatOffAddrAsFloatX4# (Addr# -> Int# -> Addr#
plusAddr# Addr#
addr# Int#
byteOff#) Int#
0# FloatX4#
v0# State# RealWorld
s0 of
            State# RealWorld
s1 -> case Addr# -> Int# -> FloatX4# -> State# RealWorld -> State# RealWorld
forall d. Addr# -> Int# -> FloatX4# -> State# d -> State# d
writeFloatOffAddrAsFloatX4# (Addr# -> Int# -> Addr#
plusAddr# (Addr# -> Int# -> Addr#
plusAddr# Addr#
addr# Int#
byteOff#) Int#
16#) Int#
0# FloatX4#
v1# State# RealWorld
s1 of
              State# RealWorld
s2 -> (# State# RealWorld
s2, () #)

-- | Vectorized Quad Poking: writes 4 vertices (128 bytes total) and 6 indices (24 bytes total)
-- with SIMD vector stores.
{-# INLINE pokeQuadSIMD #-}
pokeQuadSIMD ::
  Ptr Word8 ->
  Int ->
  Ptr Word8 ->
  Int ->
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  Word32 ->
  IO ()
pokeQuadSIMD :: Ptr Word8
-> Int
-> Ptr Word8
-> Int
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Word32
-> IO ()
pokeQuadSIMD Ptr Word8
vertices Int
vOffset Ptr Word8
indices Int
iOffset Float
x Float
y Float
w Float
h Float
u0 Float
v0 Float
u1 Float
v1 Float
r Float
g Float
b Float
a Word32
baseIdx = do
  let x1 :: Float
x1 = Float
x Float -> Float -> Float
forall a. Num a => a -> a -> a
+ Float
w
      y1 :: Float
y1 = Float
y Float -> Float -> Float
forall a. Num a => a -> a -> a
+ Float
h
  Ptr Word8
-> Int
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> IO ()
pokeVertexSIMD Ptr Word8
vertices Int
vOffset Float
x Float
y Float
r Float
g Float
b Float
a Float
u0 Float
v0
  Ptr Word8
-> Int
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> IO ()
pokeVertexSIMD Ptr Word8
vertices (Int
vOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
32) Float
x1 Float
y Float
r Float
g Float
b Float
a Float
u1 Float
v0
  Ptr Word8
-> Int
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> IO ()
pokeVertexSIMD Ptr Word8
vertices (Int
vOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
64) Float
x1 Float
y1 Float
r Float
g Float
b Float
a Float
u1 Float
v1
  Ptr Word8
-> Int
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> IO ()
pokeVertexSIMD Ptr Word8
vertices (Int
vOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
96) Float
x Float
y1 Float
r Float
g Float
b Float
a Float
u0 Float
v1
  Ptr Word8 -> Int -> Word32 -> IO ()
pokeQuadIndicesSIMD Ptr Word8
indices Int
iOffset Word32
baseIdx

-- Six indices form the same two triangles for both solid and gradient quads.
{-# INLINE pokeQuadIndicesSIMD #-}
pokeQuadIndicesSIMD :: Ptr Word8 -> Int -> Word32 -> IO ()
pokeQuadIndicesSIMD :: Ptr Word8 -> Int -> Word32 -> IO ()
pokeQuadIndicesSIMD (Ptr Addr#
addr#) offset :: Int
offset@(I# Int#
offset#) Word32
baseIdx = do
  let !(W32# Word32#
b0#) = Word32
baseIdx
      !(W32# Word32#
b1#) = Word32
baseIdx Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
1
      !(W32# Word32#
b2#) = Word32
baseIdx Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
2
      !idxVec# :: Word32X4#
idxVec# = (# Word32#, Word32#, Word32#, Word32# #) -> Word32X4#
packWord32X4# (# Word32#
b0#, Word32#
b1#, Word32#
b2#, Word32#
b0# #)
  (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, () #)) -> IO ())
-> (State# RealWorld -> (# State# RealWorld, () #)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s0 ->
    case Addr# -> Int# -> Word32X4# -> State# RealWorld -> State# RealWorld
forall d. Addr# -> Int# -> Word32X4# -> State# d -> State# d
writeWord32OffAddrAsWord32X4# (Addr# -> Int# -> Addr#
plusAddr# Addr#
addr# Int#
offset#) Int#
0# Word32X4#
idxVec# State# RealWorld
s0 of
      State# RealWorld
s1 -> (# State# RealWorld
s1, () #)
  Ptr (ZonkAny 0) -> Int -> Word32 -> IO ()
forall b. Ptr b -> Int -> Word32 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff (Addr# -> Ptr (ZonkAny 0)
forall a. Addr# -> Ptr a
Ptr Addr#
addr#) (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
16) (Word32
baseIdx Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
2)
  Ptr (ZonkAny 1) -> Int -> Word32 -> IO ()
forall b. Ptr b -> Int -> Word32 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff (Addr# -> Ptr (ZonkAny 1)
forall a. Addr# -> Ptr a
Ptr Addr#
addr#) (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
20) (Word32
baseIdx Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
3)

-- | Vectorized Quad with 4 distinct corner colors (top-left, top-right, bottom-right, bottom-left)
{-# INLINE pokeQuadGradientSIMD #-}
pokeQuadGradientSIMD ::
  Ptr Word8 ->
  Int ->
  Ptr Word8 ->
  Int ->
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  (Float, Float, Float, Float) ->
  (Float, Float, Float, Float) ->
  (Float, Float, Float, Float) ->
  (Float, Float, Float, Float) ->
  Word32 ->
  IO ()
pokeQuadGradientSIMD :: Ptr Word8
-> Int
-> Ptr Word8
-> Int
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> (Float, Float, Float, Float)
-> (Float, Float, Float, Float)
-> (Float, Float, Float, Float)
-> (Float, Float, Float, Float)
-> Word32
-> IO ()
pokeQuadGradientSIMD
  Ptr Word8
vertices Int
vOffset Ptr Word8
indices Int
iOffset Float
x Float
y Float
w Float
h Float
u Float
v
  (Float
r0, Float
g0, Float
b0, Float
a0) (Float
r1, Float
g1, Float
b1, Float
a1)
  (Float
r2, Float
g2, Float
b2, Float
a2) (Float
r3, Float
g3, Float
b3, Float
a3) Word32
baseIdx = do
  let x1 :: Float
x1 = Float
x Float -> Float -> Float
forall a. Num a => a -> a -> a
+ Float
w
      y1 :: Float
y1 = Float
y Float -> Float -> Float
forall a. Num a => a -> a -> a
+ Float
h
  Ptr Word8
-> Int
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> IO ()
pokeVertexSIMD Ptr Word8
vertices Int
vOffset Float
x Float
y Float
r0 Float
g0 Float
b0 Float
a0 Float
u Float
v
  Ptr Word8
-> Int
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> IO ()
pokeVertexSIMD Ptr Word8
vertices (Int
vOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
32) Float
x1 Float
y Float
r1 Float
g1 Float
b1 Float
a1 Float
u Float
v
  Ptr Word8
-> Int
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> IO ()
pokeVertexSIMD Ptr Word8
vertices (Int
vOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
64) Float
x1 Float
y1 Float
r2 Float
g2 Float
b2 Float
a2 Float
u Float
v
  Ptr Word8
-> Int
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> IO ()
pokeVertexSIMD Ptr Word8
vertices (Int
vOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
96) Float
x Float
y1 Float
r3 Float
g3 Float
b3 Float
a3 Float
u Float
v
  Ptr Word8 -> Int -> Word32 -> IO ()
pokeQuadIndicesSIMD Ptr Word8
indices Int
iOffset Word32
baseIdx

-- | Evaluates 4 concentric arc positions:
-- xs = cx + radii * ct
-- ys = cy + radii * st
--
-- Scalar on purpose: GHC 9.14.1 miscompiles the broadcast/pack/unpack FloatX4#
-- version at -O2 once it is inlined into a loop (liberate-case computed the y
-- lane from cx), corrupting anti-aliased border vertices. The results are
-- bit-identical to the vector version, which also multiplied and added separately.
{-# INLINE concentricOffsetsSIMD #-}
concentricOffsetsSIMD ::
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  Float ->
  ((Float, Float), (Float, Float), (Float, Float), (Float, Float))
concentricOffsetsSIMD :: Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> Float
-> ((Float, Float), (Float, Float), (Float, Float), (Float, Float))
concentricOffsetsSIMD Float
cx Float
cy Float
ct Float
st Float
r0 Float
r1 Float
r2 Float
r3 =
  ( (Float
cx Float -> Float -> Float
forall a. Num a => a -> a -> a
+ Float
r0 Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
ct, Float
cy Float -> Float -> Float
forall a. Num a => a -> a -> a
+ Float
r0 Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
st)
  , (Float
cx Float -> Float -> Float
forall a. Num a => a -> a -> a
+ Float
r1 Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
ct, Float
cy Float -> Float -> Float
forall a. Num a => a -> a -> a
+ Float
r1 Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
st)
  , (Float
cx Float -> Float -> Float
forall a. Num a => a -> a -> a
+ Float
r2 Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
ct, Float
cy Float -> Float -> Float
forall a. Num a => a -> a -> a
+ Float
r2 Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
st)
  , (Float
cx Float -> Float -> Float
forall a. Num a => a -> a -> a
+ Float
r3 Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
ct, Float
cy Float -> Float -> Float
forall a. Num a => a -> a -> a
+ Float
r3 Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
st)
  )