Implement SmallArray#-based Arec (VinylRecords#153)

Philonous · May 14, 2021 · 4dc1dde · 4dc1dde
1 parent 8637345
commit 4dc1dde
Show file tree

Hide file tree

Showing 9 changed files with 20,391 additions and 60 deletions.
diff --git a/Data/Vinyl/ARec/Internal.hs b/Data/Vinyl/ARec/Internal.hs
@@ -18,49 +18,70 @@
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE UndecidableInstances #-}
+
+
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnboxedTuples #-}
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE UnliftedNewtypes #-}
+
 -- | Constant-time field accessors for extensible records. The
 -- trade-off is the usual lists vs arrays one: it is fast to add an
 -- element to the head of a list, but element access is linear time;
 -- array access time is uniform, but extending the array is more
 -- slower.
 module Data.Vinyl.ARec.Internal
-  ( ARec (..)
-  , IndexableField
-  , toARec
-  , fromARec
-  , aget
-  , unsafeAput
-  , unsafeAlens
-  , arecGetSubset
-  , arecSetSubset
-  , arecRepsMatchCoercion
-  , arecConsMatchCoercion
-  ) where
+  -- ( ARec (..)
+  -- , IndexableField
+  -- , toARec
+  -- , fromARec
+  -- , unsafeAget
+  -- , unsafeAput
+  -- , alens
+  -- , arecGetSubset
+  -- , arecSetSubset
+  -- , arecRepsMatchCoercion
+  -- , arecConsMatchCoercion
+  -- )
+  where
 import Data.Vinyl.Core
-import Data.Vinyl.Lens (RecElem(..), RecSubset(..))
+import Data.Vinyl.Lens        (RecElem(..), RecSubset(..))
 import Data.Vinyl.TypeLevel
+import GHC.ST
 
-import qualified Data.Array as Array
-import qualified Data.Array.Base as BArray
-import GHC.Exts (Any)
+import GHC.Exts               (Any)
 import Unsafe.Coerce
 #if __GLASGOW_HASKELL__ < 806
 import Data.Constraint.Forall (Forall)
 #endif
-import Data.Coerce (Coercible)
-import Data.Type.Coercion (Coercion (..))
+import Data.Coerce            (Coercible)
+import Data.Type.Coercion     (Coercion (..))
+
+import GHC.Prim
+import GHC.Types
 
 -- | An array-backed extensible record with constant-time field
 -- access.
-newtype ARec (f :: k -> *) (ts :: [k]) = ARec (Array.Array Int Any)
+data ARec (f :: k -> *) (ts :: [k]) = ARec {-# UNPACK #-} !(SmallArray# Any)
 type role ARec representational nominal
 
+-- | Get the ith element from the ARec
+unsafeIxARec
+  :: forall a k (f :: k -> *) (ts :: [k]).
+     ARec f ts
+  -> Int
+  -> a
+unsafeIxARec (ARec arr) (I# ix#) =
+  case indexSmallArray# arr ix# of
+    (# v #) -> unsafeCoerce v
+{-# INLINE unsafeIxARec #-}
+
 -- | Given that @xs@ and @ys@ have the same length, and mapping
 -- @f@ over @xs@ and @g@ over @ys@ produces lists whose elements
 -- are pairwise 'Coercible', @ARec f xs@ and @ARec g ys@ are
 -- 'Coercible'.
 arecRepsMatchCoercion :: AllRepsMatch f xs g ys => Coercion (ARec f xs) (ARec g ys)
-arecRepsMatchCoercion = Coercion
+arecRepsMatchCoercion = unsafeCoerce (Coercion :: Coercion () ())
 
 -- | Given that @forall x. Coercible (f x) (g x)@, produce a coercion from
 -- @ARec f xs@ to @ARec g xs@. While the constraint looks a lot like
@@ -69,7 +90,7 @@ arecRepsMatchCoercion = Coercion
 #if __GLASGOW_HASKELL__ >= 806
 arecConsMatchCoercion ::
   (forall (x :: k). Coercible (f x) (g x)) => Coercion (ARec f xs) (ARec g xs)
-arecConsMatchCoercion = Coercion
+arecConsMatchCoercion = unsafeCoerce (Coercion :: Coercion () ())
 #else
 arecConsMatchCoercion :: forall k (f :: k -> *) (g :: k -> *) (xs :: [k]).
   Forall (Similar f g) => Coercion (Rec f xs) (Rec g xs)
@@ -91,13 +112,87 @@ arecConsMatchCoercible f = f
 
 -- | Convert a 'Rec' into an 'ARec' for constant-time field access.
 toARec :: forall f ts. (NatToInt (RLength ts)) => Rec f ts -> ARec f ts
-toARec = go id
-  where go :: ([Any] -> [Any]) -> Rec f ts' -> ARec f ts
-        go acc RNil = ARec $! Array.listArray (0, n - 1) (acc [])
-        go acc (x :& xs) = go (acc . (unsafeCoerce x :)) xs
-        n = natToInt @(RLength ts)
+toARec r =
+  let n = natToInt @(RLength ts)
+      !(I# n#) = n
+  in runST $
+    ST $ \st0 -> case newSmallArray# n# undefined st0 of
+      (# st1 , mArr #) ->
+        case go mArr st1 0 r of
+          st2 ->
+            case unsafeFreezeSmallArray# mArr st2 of
+             (# st3, filledArray #) -> (# st3, ARec filledArray #)
+  where
+    go :: forall s us. SmallMutableArray# s Any
+       -> State# s
+       -> Int
+       -> Rec f us
+       -> State# s
+    go _mArr stn1 n RNil = stn1
+    go mArr stn1 n@(I# n#) ( x :& xs) =
+      case writeSmallArray# mArr n# (unsafeCoerce x) stn1 of
+        stn2 -> go mArr stn2 (n+1) xs
 {-# INLINE toARec #-}
 
+-- Don't export constructor
+newtype MkARec f us =
+  -- Takes the current index and returns the maximum index and an action that
+  -- fills the array
+  MkARec (forall s. Int -> SmallMutableArray# s Any -> ST s ())
+
+infixr 1 &:
+(&:) :: f u -> MkARec f us -> MkARec f ( u ': us )
+(&:) !v (MkARec fvs) = MkARec $ \i@(I# ix) mArr ->
+  let setRemainingFields = fvs (i+1)
+  in do
+    ST $ \st0 -> case writeSmallArray# mArr ix (unsafeCoerce v) st0 of
+      st1 -> (# st1, () #)
+    setRemainingFields mArr
+{-# INLINE (&:) #-}
+
+arnil :: MkARec f '[]
+arnil = MkARec $ \_i _arr -> return ()
+{-# INLINE arnil #-}
+
+arec
+  :: forall k (us :: [k] ) f
+  . (NatToInt (RLength us)) =>
+      MkARec f us
+  -> ARec f us
+arec (MkARec fillArray) =
+  let !(I# len#) = natToInt @(RLength us)
+  in runST $
+     ST $ \st0 -> case newSmallArray# len# undefined st0 of
+      (# st1 , mArr #) ->
+        case fillArray 0 mArr of
+          ST s ->
+            case s st1 of
+              (# st2, () #) ->
+                case unsafeFreezeSmallArray# mArr st2 of
+                  (# st3, filledArray #) -> (# st3, ARec filledArray #)
+{-# INLINE arec #-}
+
+
+class ToARec (us :: [k]) where
+  aRecValues :: Rec f us -> MkARec f us
+
+instance ToARec '[] where
+  aRecValues RNil = arnil
+  {-# INLINE aRecValues #-}
+
+instance ToARec us => ToARec (u ': us) where
+  aRecValues (x :& xs) = x &: aRecValues xs
+  {-# INLINE aRecValues #-}
+
+-- | Convert a 'Rec' into an 'ARec' for constant-time field access.
+toARecFast
+  :: forall f ts.
+     (NatToInt (RLength ts), ToARec ts)
+  => Rec f ts
+  -> ARec f ts
+toARecFast rs = arec (aRecValues rs)
+{-# INLINE toARecFast #-}
+
 -- | Defines a constraint that lets us index into an 'ARec' in order
 -- to produce a 'Rec' using 'fromARec'.
 class (NatToInt (RIndex t ts)) => IndexableField ts t where
@@ -107,22 +202,28 @@ instance (NatToInt (RIndex t ts)) => IndexableField ts t where
 fromARec :: forall f ts.
             (RecApplicative ts, RPureConstrained (IndexableField ts) ts)
          => ARec f ts -> Rec f ts
-fromARec (ARec arr) = rpureConstrained @(IndexableField ts) aux
+fromARec arec = rpureConstrained @(IndexableField ts) aux
   where aux :: forall t. NatToInt (RIndex t ts) => f t
-        aux = unsafeCoerce (arr Array.! natToInt @(RIndex t ts))
+        aux = unsafeIxARec arec (natToInt @(RIndex t ts))
 {-# INLINE fromARec #-}
 
 -- | Get a field from an 'ARec'.
 aget :: forall t f ts. (NatToInt (RIndex t ts)) => ARec f ts -> f t
-aget (ARec arr) =
-  unsafeCoerce (BArray.unsafeAt arr (natToInt @(RIndex t ts)))
+aget arec = unsafeIxARec arec (natToInt @(RIndex t ts))
 {-# INLINE aget #-}
 
 -- | Set a field in an 'ARec'.
 unsafeAput :: forall t t' f ts ts'. (NatToInt (RIndex t ts))
       => f t' -> ARec f ts -> ARec f ts'
-unsafeAput x (ARec arr) = ARec (arr Array.// [(i, unsafeCoerce x)])
-  where i = natToInt @(RIndex t ts)
+unsafeAput x (ARec arr) =
+  let !(I# z#) = 0
+      !(I# n#) = natToInt @(RIndex t ts)
+  in runST $ ST $ \st0 ->
+    case thawSmallArray# arr z# (sizeofSmallArray# arr) st0 of
+     (# st1, mArr #) ->
+      case writeSmallArray# mArr n# (unsafeCoerce x) st1 of
+        st2 -> case unsafeFreezeSmallArray# mArr st2 of
+          (# st3, arr' #) -> (# st3, ARec arr' #)
 {-# INLINE unsafeAput #-}
 
 -- | Define a lens for a field of an 'ARec'.
@@ -157,19 +258,40 @@ instance (RIndex t (s ': ts) ~ 'S i, NatToInt i,  RecElem ARec t t' ts ts' i)
 arecGetSubset :: forall rs ss f.
                  (IndexWitnesses (RImage rs ss), NatToInt (RLength rs))
               => ARec f ss -> ARec f rs
-arecGetSubset (ARec arr) = ARec (Array.listArray (0, n-1) $
-                                 go (indexWitnesses @(RImage rs ss)))
-  where go :: [Int] -> [Any]
-        go = map (arr Array.!)
-        n = natToInt @(RLength rs)
+arecGetSubset (ARec arr) = runST $ ST $ \st0 ->
+  case newSmallArray# n# undefined st0 of
+    (# st1 , mArr #) ->
+      case go mArr 0 st1 (indexWitnesses @(RImage rs ss)) of
+        st2 ->
+          case unsafeFreezeSmallArray# mArr st2 of
+            (# st3, filledArray #) -> (# st3, ARec filledArray #)
+  where go :: forall s. SmallMutableArray# s Any -> Int -> State# s -> [Int] -> State# s
+        go _mArr _i stn [] = stn
+        go mArr to@(I# to#) stn (I# from# : is) =
+          case indexSmallArray# arr from# of
+            (# v #) -> case writeSmallArray# mArr to# v stn of
+              stn' -> go mArr (to+1) stn' is
+        !(I# n#) = natToInt @(RLength rs)
 {-# INLINE arecGetSubset #-}
 
 -- | Set a subset of a larger record's fields to all of the fields of
 -- a smaller record.
 arecSetSubset :: forall rs ss f. (IndexWitnesses (RImage rs ss))
               => ARec f ss -> ARec f rs -> ARec f ss
-arecSetSubset (ARec arrBig) (ARec arrSmall) = ARec (arrBig Array.// updates)
-  where updates = zip (indexWitnesses @(RImage rs ss)) (Array.elems arrSmall)
+arecSetSubset (ARec arrBig) arecSmall =
+  runST $ ST $ \st0 ->
+  case thawSmallArray# arrBig z# (sizeofSmallArray# arrBig) st0 of
+    (# st1, mArr #) -> case go mArr 0 (indexWitnesses @(RImage rs ss)) st1 of
+      st2 -> case unsafeFreezeSmallArray# mArr st2 of
+          (# st3, arr' #) -> (# st3, ARec arr' #)
+    where
+  !(I# z#) = 0
+  go :: SmallMutableArray# s Any -> Int -> [Int] -> State# s -> State# s
+  go _mArr _ [] stn = stn
+  go mArr n (I# i# : is) stn =
+    case writeSmallArray# mArr i# (unsafeIxARec arecSmall n) stn of
+      stn' -> go mArr (n+1) is stn'
+
 {-# INLINE arecSetSubset #-}
 
 instance (is ~ RImage rs ss, IndexWitnesses is, NatToInt (RLength rs))