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
* Add (&:), arnil and arec presudo-constructors

Data/Vinyl.hs

@@ -15,7 +15,7 @@ import Data.Vinyl.Core
 import Data.Vinyl.Class.Method (RecMapMethod(..), RecPointed(..))
 import Data.Vinyl.Class.Method (rmapMethodF, mapFields)
 import Data.Vinyl.Class.Method (rtraverseInMethod, rsequenceInFields)
-import Data.Vinyl.ARec (ARec, toARec, fromARec)
+import Data.Vinyl.ARec (ARec, toARec, fromARec, ARecBuilder, arec, (&:), arnil)
 import Data.Vinyl.Derived
 import Data.Vinyl.FromTuple (record, fieldRec, ruple, xrec, xrecX, xrecTuple)
 import Data.Vinyl.Functor (ElField(..))

Data/Vinyl/ARec.hs

@@ -1,12 +1,35 @@
 {-# LANGUAGE Trustworthy #-}
 
--- | 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.
+-- | An 'ARec' (array backed record) is an extensible record using an array to
+-- store its fields. It has constant time access to fields, but fields cannot be
+-- shared between records.
+--
+-- = Usage
+--
+-- ARecs can be constructed via 'ARecBuilder'
+--
+-- Example (requires -XOverloadedLabels):
+--
+-- > import Data.Vinyl.Derived ((=:))
+-- >
+-- > user :: ARec ElField '[ "name"   ::: String
+-- >                       , "age"    ::: Int
+-- >                       , "active" ::: Bool]
+-- > user = arec (  #name   =: "Peter"
+-- >             &: #age    =: 4
+-- >             &: #active =: True
+-- >             &: arnil
+-- >             )
+--
+-- You can also convert a 'Rec' to an 'ARec' via 'toARec'
+--
+--
 module Data.Vinyl.ARec
   ( ARec -- Exported abstractly
+  , ARecBuilder -- Exported abstractly
+  , arec
+  , (&:)
+  , arnil
   , IndexableField
   , toARec
   , fromARec

Data/Vinyl/ARec/Internal.hs

@@ -1,11 +1,13 @@
 {-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE ConstraintKinds #-}
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE MagicHash #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE PolyKinds #-}
 #if __GLASGOW_HASKELL__ >= 806
@@ -17,15 +19,32 @@
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UnboxedTuples #-}
 {-# LANGUAGE UndecidableInstances #-}
 -- | 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.
+
+-- 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 constructor
+--   * 1 for the pointer to the SmallArray#
+--   * The SmallArray# has 2 words as header (1 for GC, 1 number of elements)
+--   * 1 pointer per element to the actual data
+-- * Rec requires (1 + 2n) words + size of Fields
 module Data.Vinyl.ARec.Internal
   ( ARec (..)
   , IndexableField
+  , ARecBuilder (..)
+  , arec
+  , (&:)
+  , arnil
   , toARec
   , fromARec
   , aget
@@ -39,28 +58,41 @@ module Data.Vinyl.ARec.Internal
 import Data.Vinyl.Core
 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 !(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 +101,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 +110,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 &: 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 +143,74 @@ 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 '(&:)' 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"
+-- >             &: #age    =: 4
+-- >             &: #active =: True
+-- >             &: 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 Any -- Arrray to write to
+              -> ST s ()
+              )
+
+infixr 1 &:
+-- | Pseudo-constructor for an ARecBuilder
+--
+-- "Cons" a field to an ARec under construction
+--
+-- See 'ARecBuilder'
+(&:) :: f u -> ARecBuilder f us -> ARecBuilder f ( u ': us )
+(&:) !v (ARecBuilder fvs) = ARecBuilder $ \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 (&:) #-}
+
+-- | 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) =
+  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 #-}
 
 -- | Defines a constraint that lets us index into an 'ARec' in order
 -- to produce a 'Rec' using 'fromARec'.
@@ -107,22 +221,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 +277,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))