Vector Unbox instances for Word256, Word128, UTxOKey, UTxOValue

IntersectMBO · May 7, 2024 · d024613 · d024613
1 parent 66cb182
commit d024613
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 8 deletions.
diff --git a/src-extras/Database/LSMTree/Extras/Orphans.hs b/src-extras/Database/LSMTree/Extras/Orphans.hs
@@ -1,11 +1,6 @@
-{-# LANGUAGE DeriveAnyClass             #-}
-{-# LANGUAGE DeriveGeneric              #-}
-{-# LANGUAGE DerivingStrategies         #-}
-{-# LANGUAGE FlexibleInstances          #-}
-{-# LANGUAGE GeneralisedNewtypeDeriving #-}
-{-# LANGUAGE MagicHash                  #-}
-{-# LANGUAGE NamedFieldPuns             #-}
-{-# LANGUAGE StandaloneDeriving         #-}
+{-# LANGUAGE DeriveAnyClass #-}
+{-# LANGUAGE DerivingVia    #-}
+{-# LANGUAGE TypeFamilies   #-}
 
 {-# OPTIONS_GHC -Wno-orphans #-}
 
@@ -17,7 +12,11 @@ module Database.LSMTree.Extras.Orphans (
 
 import           Control.DeepSeq
 import qualified Data.Primitive as P
+import qualified Data.Vector.Generic as VG
+import qualified Data.Vector.Generic.Mutable as VGM
 import qualified Data.Vector.Primitive as PV
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
 import           Data.WideWord.Word128 (Word128 (..), byteSwapWord128)
 import           Data.WideWord.Word256 (Word256 (..))
 import           Database.LSMTree.Internal.Primitive (indexWord8ArrayAsWord64)
@@ -84,6 +83,14 @@ indexWord8ArrayAsWord256 !ba !off =
             (indexWord8ArrayAsWord64 ba (off + 8))
             (indexWord8ArrayAsWord64 ba off)
 
+newtype instance VUM.MVector s Word256 = MV_Word256 (PV.MVector s Word256)
+newtype instance VU.Vector     Word256 = V_Word256  (PV.Vector    Word256)
+
+deriving via VU.UnboxViaPrim Word256 instance VGM.MVector VU.MVector Word256
+deriving via VU.UnboxViaPrim Word256 instance VG.Vector   VU.Vector  Word256
+
+instance VUM.Unbox Word256
+
 {-------------------------------------------------------------------------------
   Word128
 -------------------------------------------------------------------------------}
@@ -127,6 +134,14 @@ indexWord8ArrayAsWord128 !ba !off =
     Word128 (indexWord8ArrayAsWord64 ba (off + 8))
             (indexWord8ArrayAsWord64 ba off)
 
+newtype instance VUM.MVector s Word128 = MV_Word128 (PV.MVector s Word128)
+newtype instance VU.Vector     Word128 = V_Word128  (PV.Vector    Word128)
+
+deriving via VU.UnboxViaPrim Word128 instance VGM.MVector VU.MVector Word128
+deriving via VU.UnboxViaPrim Word128 instance VG.Vector   VU.Vector  Word128
+
+instance VUM.Unbox Word128
+
 {-------------------------------------------------------------------------------
   NFData
 -------------------------------------------------------------------------------}

diff --git a/src-extras/Database/LSMTree/Extras/UTxO.hs b/src-extras/Database/LSMTree/Extras/UTxO.hs
@@ -1,20 +1,27 @@
 {-# LANGUAGE DeriveAnyClass             #-}
 {-# LANGUAGE DeriveGeneric              #-}
 {-# LANGUAGE DerivingStrategies         #-}
+{-# LANGUAGE DerivingVia                #-}
 {-# LANGUAGE GeneralisedNewtypeDeriving #-}
 {-# LANGUAGE NamedFieldPuns             #-}
+{-# LANGUAGE TypeFamilies               #-}
 {- HLINT ignore "Redundant <$>" -}
 
 module Database.LSMTree.Extras.UTxO (
     UTxOKey (..)
   , UTxOValue (..)
+  , zeroUTxOValue
   , UTxOBlob (..)
   ) where
 
 import           Control.DeepSeq
 import qualified Data.ByteString as BS
 import qualified Data.Primitive as P
+import qualified Data.Vector.Generic as VG
+import qualified Data.Vector.Generic.Mutable as VGM
 import qualified Data.Vector.Primitive as PV
+import qualified Data.Vector.Unboxed as VU
+import qualified Data.Vector.Unboxed.Mutable as VUM
 import           Data.WideWord.Word128
 import           Data.WideWord.Word256
 import           Data.Word
@@ -56,6 +63,20 @@ instance Arbitrary UTxOKey where
   arbitrary = UTxOKey <$> arbitrary <*> arbitrary
   shrink (UTxOKey a b) = [ UTxOKey a' b' | (a', b') <- shrink (a, b) ]
 
+newtype instance VUM.MVector s UTxOKey = MV_UTxOKey (VU.MVector s (Word256, Word16))
+newtype instance VU.Vector     UTxOKey = V_UTxOKey  (VU.Vector    (Word256, Word16))
+
+instance VU.IsoUnbox UTxOKey (Word256, Word16) where
+  toURepr (UTxOKey a b) = (a, b)
+  fromURepr (a, b) = UTxOKey a b
+  {-# INLINE toURepr #-}
+  {-# INLINE fromURepr #-}
+
+deriving via VU.As UTxOKey (Word256, Word16) instance VGM.MVector VU.MVector UTxOKey
+deriving via VU.As UTxOKey (Word256, Word16) instance VG.Vector   VU.Vector  UTxOKey
+
+instance VUM.Unbox UTxOKey
+
 {-------------------------------------------------------------------------------
   UTxO values
 -------------------------------------------------------------------------------}
@@ -92,6 +113,23 @@ instance Arbitrary UTxOValue where
   shrink (UTxOValue a b c d) = [ UTxOValue a' b' c' d'
                                | (a', b', c', d') <- shrink (a, b, c, d) ]
 
+newtype instance VUM.MVector s UTxOValue = MV_UTxOValue (VU.MVector s (Word256, Word128, Word64, Word32))
+newtype instance VU.Vector     UTxOValue = V_UTxOValue  (VU.Vector    (Word256, Word128, Word64, Word32))
+
+instance VU.IsoUnbox UTxOValue (Word256, Word128, Word64, Word32) where
+  toURepr (UTxOValue a b c d) = (a, b, c, d)
+  fromURepr (a, b, c, d) = UTxOValue a b c d
+  {-# INLINE toURepr #-}
+  {-# INLINE fromURepr #-}
+
+deriving via VU.As UTxOValue (Word256, Word128, Word64, Word32) instance VGM.MVector VU.MVector UTxOValue
+deriving via VU.As UTxOValue (Word256, Word128, Word64, Word32) instance VG.Vector   VU.Vector  UTxOValue
+
+instance VUM.Unbox UTxOValue
+
+zeroUTxOValue :: UTxOValue
+zeroUTxOValue = UTxOValue 0 0 0 0
+
 {-------------------------------------------------------------------------------
   UTxO blobs
 -------------------------------------------------------------------------------}