2 files changed, 183 insertions, 0 deletions

diff --git a/src/Data/Primitive/ByteArray/Util.hs b/src/Data/Primitive/ByteArray/Util.hs
new file mode 100644
index 0000000..83a2e7d
--- /dev/null
+++ b/src/Data/Primitive/ByteArray/Util.hs
@@ -0,0 +1,71 @@
+{-# LANGUAGE DataKinds             #-}
+{-# LANGUAGE FlexibleContexts      #-}
+{-# LANGUAGE FlexibleInstances     #-}
+{-# LANGUAGE KindSignatures        #-}
+{-# LANGUAGE MagicHash             #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE ScopedTypeVariables   #-}
+{-# LANGUAGE TypeFamilies          #-}
+{-# LANGUAGE TypeOperators         #-}
+{-# LANGUAGE CPP                   #-}
+module Data.Primitive.ByteArray.Util where
+
+import GHC.Exts (Ptr(..))
+import GHC.TypeLits
+import Control.Monad.Primitive
+import qualified Data.ByteString as B
+import qualified Data.ByteString.Unsafe as B
+import Data.Primitive.Addr
+import Data.Primitive.Types
+import Data.Primitive.ByteArray
+import Data.Word
+import Foreign.Ptr
+
+newtype Offset (n :: Nat) = Offset Int
+
+offset :: KnownNat n => Offset n
+offset = let k = Offset $ fromIntegral $ natVal k in k
+
+(+.) :: Offset j -> Offset k -> Offset (j + k)
+Offset j +. Offset k = Offset (j + k)
+
+
+type family SizeOf a :: Nat
+
+class IsMultipleOf (n::Nat) (k::Nat)
+
+instance n ~ (q * k) => IsMultipleOf n k
+
+writeAtByte :: ( Prim a
+               , PrimMonad m
+#if __GLASGOW_HASKELL__ >= 802
+               , IsMultipleOf n (SizeOf a)
+#endif
+               ) => MutableByteArray (PrimState m) -> Offset n -> a -> m ()
+writeAtByte buf (Offset offset) a = writeByteArray buf (div offset $ (sizeOf a)) a
+{-# INLINE writeAtByte #-}
+
+readAtByte :: forall a m n.
+               ( Prim a
+               , PrimMonad m
+#if __GLASGOW_HASKELL__ >= 802
+               , IsMultipleOf n (SizeOf a)
+#endif
+               ) => MutableByteArray (PrimState m) -> Offset n -> m a
+readAtByte buf (Offset offset) = readByteArray buf (div offset $ (sizeOf (undefined :: a)))
+{-# INLINE readAtByte #-}
+
+writeStringAt :: PrimMonad m => MutableByteArray s -> Int -> B.ByteString -> m (Ptr a)
+writeStringAt src o bsname = do
+    (p,cnt) <- unsafeIOToPrim $ B.unsafeUseAsCStringLen bsname return
+    let nptr = ptr (mutableByteArrayContents src) `plusPtr` o
+    copyAddr (adr nptr) (adr p) cnt
+    writeOffAddr (adr nptr) cnt (0 :: Word8)
+    return nptr
+
+ptr :: Addr -> Ptr a
+ptr (Addr a) = Ptr a
+
+adr :: Ptr a -> Addr
+adr (Ptr a) = Addr a
+
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