1 files changed, 112 insertions, 0 deletions

diff --git a/src/Data/Primitive/Struct.hs b/src/Data/Primitive/Struct.hs
new file mode 100644
index 0000000..154b750
--- /dev/null
+++ b/src/Data/Primitive/Struct.hs
@@ -0,0 +1,112 @@
+{-# LANGUAGE AllowAmbiguousTypes   #-}
+{-# LANGUAGE CPP                   #-}
+{-# LANGUAGE DataKinds             #-}
+{-# LANGUAGE FlexibleContexts      #-}
+{-# LANGUAGE FlexibleInstances     #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE ScopedTypeVariables   #-}
+{-# LANGUAGE TypeApplications      #-}
+{-# LANGUAGE TypeFamilies          #-}
+{-# LANGUAGE TypeOperators         #-}
+module Data.Primitive.Struct where
+
+import Control.Monad.Primitive
+import Data.Primitive.ByteArray
+import Data.Primitive.ByteArray.Util
+import Data.Primitive.Types
+import Data.Tagged
+import Data.Typeable
+import Foreign.Ptr
+import Foreign.Storable
+import GHC.TypeLits
+
+newtype Field tag typ n = Field (Offset n)
+
+data Struct m base tag = Struct
+    { structOffset :: !(Offset base)
+    , structArray  :: !(MutableByteArray (PrimState m))
+    }
+
+newStruct :: forall tag m. (KnownNat (SizeOf tag), PrimMonad m) => m (Struct m 0 tag)
+newStruct = Struct (Offset 0) <$> newPinnedByteArray (fromIntegral sz)
+ where
+    sz = natVal (Proxy :: Proxy (SizeOf tag))
+
+newtype Nested tag subtag n = Nested (Offset n)
+
+class IsStruct m p where
+    type BaseOffset p :: Nat
+    type NestedStruct m p (offset::Nat) subtag
+
+    setField :: ( Prim a
+#if __GLASGOW_HASKELL__ >= 802
+           , IsMultipleOf ((BaseOffset p) + k) (SizeOf a)
+#endif
+           ) => p tag -> Field tag a k -> a -> m ()
+
+    getField :: ( Prim a
+#if __GLASGOW_HASKELL__ >= 802
+           , IsMultipleOf ((BaseOffset p) + k) (SizeOf a)
+#endif
+           ) => p tag -> Field tag a k -> m a
+
+    nestedField :: p tag -> Field tag subtag k -> proxy m -> NestedStruct m p k subtag
+
+class IsField (lbl::Symbol) tag where
+    type FieldOffset lbl tag :: Nat
+    type FieldType lbl tag
+    field :: p tag -> Field tag (FieldType lbl tag) (FieldOffset lbl tag)
+
+set :: forall lbl tag m p. (IsField lbl tag, IsStruct m p, Prim (FieldType lbl tag),
+#if __GLASGOW_HASKELL__ >= 802
+              IsMultipleOf (BaseOffset p + FieldOffset lbl tag) (SizeOf (FieldType lbl tag))
+#endif
+              ) =>
+             p tag -> FieldType lbl tag -> m ()
+set p a = setField p (field @lbl p) a
+{-# INLINE set #-}
+
+get :: forall lbl tag m p. (IsField lbl tag, IsStruct m p, Prim (FieldType lbl tag),
+#if __GLASGOW_HASKELL__ >= 802
+              IsMultipleOf (BaseOffset p + FieldOffset lbl tag) (SizeOf (FieldType lbl tag))
+#endif
+              ) =>
+             p tag -> m (FieldType lbl tag)
+get p = getField p (field @lbl p)
+{-# INLINE get #-}
+
+
+modify :: forall lbl tag m p.
+    ( Monad m
+    , IsField lbl tag
+    , IsStruct m p
+    , Prim (FieldType lbl tag)
+#if __GLASGOW_HASKELL__ >= 802
+    , IsMultipleOf (BaseOffset p + FieldOffset lbl tag) (SizeOf (FieldType lbl tag))
+#endif
+    ) => p tag -> (FieldType lbl tag -> FieldType lbl tag) -> m ()
+modify p f = get @lbl p >>= set @lbl p . f
+
+nested :: forall lbl m p tag. (IsField lbl tag, IsStruct m p) =>
+                p tag
+                -> NestedStruct m p (FieldOffset lbl tag) (FieldType lbl tag)
+nested p = nestedField p (field @lbl p) (Proxy @m)
+
+instance PrimMonad m => IsStruct m (Struct m base) where
+    type BaseOffset (Struct m base) = base
+    type NestedStruct m (Struct m base) j t = Struct m (base + j) t
+    setField (Struct o c) (Field field) value = writeAtByte c (o +. field) value
+    getField (Struct o c) (Field field) = readAtByte c (o +. field)
+    nestedField (Struct base ary) (Field offset) _ = Struct (base +. offset) ary
+
+instance IsStruct IO Ptr where
+    type BaseOffset Ptr = 0
+    type NestedStruct IO Ptr j t = Ptr t
+    setField ptr (Field (Offset o)) value = poke (ptr `plusPtr` o) $ PrimStorable value
+    getField ptr (Field (Offset o)) = getPrimStorable <$> peek (ptr `plusPtr` o)
+    nestedField ptr (Field (Offset offset)) _ = castPtr (plusPtr ptr offset)
+
+withPointer :: Struct IO base tag -> (Ptr tag -> IO x) -> IO x
+withPointer (Struct (Offset off) ary) f = do
+    x <- f (ptr (mutableByteArrayContents ary) `plusPtr` off)
+    seq ary $ return x