Implement improvements to ARec (VinylRecords#150, VinylRecords#153)

* Move from Array to SmallArray# to avoid intermediate list during construction
* Use class instead of recursive function in toARec for improved inlining

Data/Vinyl/ARec/Internal.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE ConstraintKinds #-}
 {-# LANGUAGE DataKinds #-}
@@ -23,9 +24,26 @@
 -- 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.
+
+-- Tradeoffs:
+--
+-- * No sharing of the spine (i.e. when you change elements in the front of the
+--   record the tail can't be re-used)
+-- * ARec requires (4 + n) words + size of the fields
+--   * 1 for the ARec constructor
+--   * 1 for the pointer to the SmallArray#
+--   * The SmallArray# has 2 words as header (1 for GC, 1 for number of elements)
+--   * 1 pointer per element to the actual data
+-- * Rec requires (2n) words + size of Fields
+--   * 1 word per (:&) constructor
+--   * 1 word for the pointer to the element
 module Data.Vinyl.ARec.Internal
   ( ARec (..)
   , IndexableField
+  , arec
+  , ARecBuilder (..)
+  , arcons
+  , arnil
   , toARec
   , fromARec
   , aget
@@ -39,28 +57,36 @@ module Data.Vinyl.ARec.Internal
 import Data.Vinyl.Core
 import Data.Vinyl.Lens (RecElem(..), RecSubset(..))
 import Data.Vinyl.TypeLevel
+import Data.Vinyl.ARec.Internal.SmallArray
+import Control.Monad.ST
 
-import qualified Data.Array as Array
-import qualified Data.Array.Base as BArray
-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.Type.Coercion     (Coercion (..))
+import GHC.Types
 
 -- | An array-backed extensible record with constant-time field
 -- access.
-newtype ARec (f :: k -> *) (ts :: [k]) = ARec (Array.Array Int Any)
+newtype ARec (f :: k -> *) (ts :: [k]) = ARec SmallArray
 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 ar) ix = indexSmallArray ar ix
+{-# 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 +95,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)
@@ -78,6 +104,28 @@ arecConsMatchCoercion :: forall k (f :: k -> *) (g :: k -> *) (xs :: [k]).
 arecConsMatchCoercion = unsafeCoerce (Coercion :: Coercion (Rec f xs) (Rec f xs))
 #endif
 
+-- Using a class instead of a recursive function allows aRecValues to be
+-- completely inlined
+class ToARec (us :: [k]) where
+  aRecValues :: Rec f us -> ARecBuilder f us
+
+instance ToARec '[] where
+  aRecValues RNil = arnil
+  {-# INLINE aRecValues #-}
+
+instance ToARec us => ToARec (u ': us) where
+  aRecValues (x :& xs) = x `arcons` aRecValues xs
+  {-# INLINE aRecValues #-}
+
+-- | Convert a 'Rec' into an 'ARec' for constant-time field access.
+toARec
+  :: forall f ts.
+     (NatToInt (RLength ts), ToARec ts)
+  => Rec f ts
+  -> ARec f ts
+toARec rs = arec (aRecValues rs)
+{-# INLINE toARec #-}
+
 {-
 -- This is sensible, but the ergonomics are likely quite bad thanks to the
 -- interaction between Coercible resolution and resolution in the presence of
@@ -89,14 +137,63 @@ arecConsMatchCoercible :: forall k f g rep (r :: TYPE rep).
 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)
-{-# INLINE toARec #-}
+-- | An efficient builder for ARec values
+--
+-- Use the pseudo-constructors 'arcons' and 'arnil' to construct an
+-- 'ARecBuilder' and then turn it into an 'ARec' with 'arec'
+--
+-- Example: (requires -XOverloadedLabels and )
+--
+-- > user :: ARec ElField '[ "name"   ::: String
+-- >                       , "age"    ::: Int
+-- >                       , "active" ::: Bool]
+-- > user = arec (  #name   =: "Peter"
+-- >             `arcons` #age    =: 4
+-- >             `arcons` #active =: True
+-- >             `arcons` arnil
+-- >             )
+newtype ARecBuilder f us =
+  -- A function that writes values to the correct position in the underlying array
+  -- Takes the current index
+  ARecBuilder ( forall s.
+                Int -- Index to write to
+              -> SmallMutableArray s -- Arrray to write to
+              -> ST s ()
+              )
+
+infixr 1 `arcons`
+-- | Pseudo-constructor for an ARecBuilder
+--
+-- "Cons" a field to an ARec under construction
+--
+-- See 'ARecBuilder'
+arcons :: f u -> ARecBuilder f us -> ARecBuilder f (u ': us)
+arcons !v (ARecBuilder fvs) = ARecBuilder $ \i mArr -> do
+    writeSmallArray mArr i v
+    fvs (i+1) mArr
+{-# INLINE arcons #-}
+
+-- | Pseudo-constructor for 'ARecBuilder'
+--
+-- Build an ARec without fields
+--
+-- See 'ARecBuilder'
+arnil :: ARecBuilder f '[]
+arnil = ARecBuilder $ \_i _arr -> return ()
+{-# INLINE arnil #-}
+
+-- | Turn an ARecBuilder into an ARec
+--
+-- See 'ARecBuilder'
+arec
+  :: forall k (us :: [k] ) f
+  . (NatToInt (RLength us)) =>
+      ARecBuilder f us
+  -> ARec f us
+arec (ARecBuilder fillArray) = ARec $
+  runST $ withNewSmallArray (natToInt @(RLength us))
+          $ fillArray 0
+{-# INLINE arec #-}
 
 -- | Defines a constraint that lets us index into an 'ARec' in order
 -- to produce a 'Rec' using 'fromARec'.
@@ -107,22 +204,22 @@ 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 ar = rpureConstrained @(IndexableField ts) aux
   where aux :: forall t. NatToInt (RIndex t ts) => f t
-        aux = unsafeCoerce (arr Array.! natToInt @(RIndex t ts))
+        aux = unsafeIxARec ar (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 ar = unsafeIxARec ar (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) = ARec $ runST $
+  withThawedSmallArray arr $ \mArr ->
+    writeSmallArray mArr (natToInt @(RIndex t ts)) x
 {-# INLINE unsafeAput #-}
 
 -- | Define a lens for a field of an 'ARec'.
@@ -157,19 +254,31 @@ 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) =
+  ARec $ runST $
+    withNewSmallArray (natToInt @(RLength rs)) $ \mArr ->
+      go mArr 0 (indexWitnesses @(RImage rs ss))
+  where
+    go :: SmallMutableArray s -> Int -> [Int] -> ST s ()
+    go _mArr _to [] = return ()
+    go mArr to (from : froms) = do
+      writeSmallArray mArr to (indexSmallArray arr from :: Any)
+      go mArr (to + 1) froms
 {-# 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) (ARec arrSmall) = ARec $ runST $
+  withThawedSmallArray arrBig $ \mArr -> do
+    go mArr 0 (indexWitnesses @(RImage rs ss))
+  where
+    go :: SmallMutableArray s -> Int -> [Int] -> ST s ()
+    go _mArr _ [] = return ()
+    go mArr from (to : tos) = do
+      writeSmallArray mArr to (indexSmallArray arrSmall from)
+      go mArr (from + 1) tos
 {-# INLINE arecSetSubset #-}
 
 instance (is ~ RImage rs ss, IndexWitnesses is, NatToInt (RLength rs))

Data/Vinyl/ARec/Internal/SmallArray.hs

@@ -0,0 +1,56 @@
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnboxedTuples #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE BangPatterns #-}
+
+-- | Helper functions for SmallArray#
+--
+-- This module exposes _unsafe_ functions to work with SmallArrays.  That means
+-- that specifically neither index bounds nor element types are checked So this
+-- functionality should only be used in a context that enforces them by some
+-- other means, e.g. ARec's type index
+
+module Data.Vinyl.ARec.Internal.SmallArray where
+
+import GHC.Prim
+import GHC.Types
+import Unsafe.Coerce
+import GHC.ST
+
+data SmallArray = SmallArray !(SmallArray# Any)
+data SmallMutableArray s = SmallMutableArray !(SmallMutableArray# s Any)
+
+indexSmallArray :: SmallArray -> Int -> a
+indexSmallArray (SmallArray arr) (I# ix) =
+  case indexSmallArray# arr ix of
+    (# v #) -> unsafeCoerce v
+{-# INLINE indexSmallArray #-}
+
+withNewSmallArray :: Int -> (SmallMutableArray s -> ST s ()) -> ST s SmallArray
+withNewSmallArray (I# len#) f =
+  ST $ \s0 ->  case newSmallArray# len# (error "withNewSmallArray exploded") s0 of
+       (# s1, mArr #) ->
+         case f (SmallMutableArray mArr) of
+           ST st -> case st s1 of
+             (# s2, () #) -> case unsafeFreezeSmallArray# mArr s2 of
+               (# s3, ar #) -> (# s3, SmallArray ar #)
+{-# INLINE withNewSmallArray #-}
+
+writeSmallArray :: SmallMutableArray s -> Int -> a -> ST s ()
+writeSmallArray (SmallMutableArray mArr) (I# n#) x = ST $ \s ->
+  case writeSmallArray# mArr n# (unsafeCoerce x) s of
+    s' -> (# s', () #)
+{-# INLINE writeSmallArray #-}
+
+withThawedSmallArray :: SmallArray
+               -> (SmallMutableArray s -> ST s ())
+               -> ST s SmallArray
+withThawedSmallArray (SmallArray arr) f = ST $ \s0 ->
+  let !(I# z#) = 0
+  in case thawSmallArray# arr z# (sizeofSmallArray# arr) s0 of
+       (# s1, mArr #) ->
+         case f (SmallMutableArray mArr) of
+           ST st -> case st s1 of
+             (# s2, () #) -> case unsafeFreezeSmallArray# mArr s2 of
+               (# s3, ar #) -> (# s3, SmallArray ar #)
+{-# INLINE withThawedSmallArray #-}