Internal.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds #-}
#if __GLASGOW_HASKELL__ >= 806
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
#endif
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# 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 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 (..)
  , ToARec
  , IndexableField
  , arec
  , ARecBuilder (..)
  , arcons
  , arnil
  , toARec
  , fromARec
  , aget
  , unsafeAput
  , unsafeAlens
  , arecGetSubset
  , arecSetSubset
  , arecRepsMatchCoercion
  , arecConsMatchCoercion
  ) where
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 Unsafe.Coerce
#if __GLASGOW_HASKELL__ < 806
import Data.Constraint.Forall (Forall)
#endif
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 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 = 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
-- @Coercible f g@, it is actually weaker.

#if __GLASGOW_HASKELL__ >= 806
arecConsMatchCoercion ::
  (forall (x :: k). Coercible (f x) (g x)) => Coercion (ARec f xs) (ARec g xs)
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)
-- Why do we need this? No idea, really. I guess some change in
-- newtype handling for Coercible in 8.6?
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
-- quantified constraints. Is there a good way to do this?

arecConsMatchCoercible :: forall k f g rep (r :: TYPE rep).
     (forall (x :: k). Coercible (f x) (g x))
  => ((forall (xs :: [k]). Coercible (ARec f xs) (ARec g xs)) => r) -> r
arecConsMatchCoercible f = f
-}

-- | 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'.
class (NatToInt (RIndex t ts)) => IndexableField ts t where
instance (NatToInt (RIndex t ts)) => IndexableField ts t where

-- | Convert an 'ARec' into a 'Rec'.
fromARec :: forall f ts.
            (RecApplicative ts, RPureConstrained (IndexableField ts) ts)
         => ARec f ts -> Rec f ts
fromARec ar = rpureConstrained @(IndexableField ts) aux
  where aux :: forall t. NatToInt (RIndex t ts) => f t
        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 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 $ runST $
  withThawedSmallArray arr $ \mArr ->
    writeSmallArray mArr (natToInt @(RIndex t ts)) x
{-# INLINE unsafeAput #-}

-- | Define a lens for a field of an 'ARec'.
unsafeAlens :: forall f g t t' ts ts'. (Functor g, NatToInt (RIndex t ts))
      => (f t -> g (f t')) -> ARec f ts -> g (ARec f ts')
unsafeAlens f ar = fmap (flip (unsafeAput @t) ar) (f (aget ar))
{-# INLINE unsafeAlens #-}

-- instance (i ~ RIndex t ts, i ~ RIndex t' ts', NatToInt (RIndex t ts)) => RecElem ARec t t' ts ts' i where
--   rlens = alens
--   rget = aget
--   rput = aput

instance RecElem ARec t t' (t ': ts) (t' ': ts) 'Z where
  rlensC = unsafeAlens
  {-# INLINE rlensC #-}
  rgetC = aget
  {-# INLINE rgetC #-}
  rputC = unsafeAput @t
  {-# INLINE rputC #-}

instance (RIndex t (s ': ts) ~ 'S i, NatToInt i,  RecElem ARec t t' ts ts' i)
  => RecElem ARec t t' (s ': ts) (s ': ts') ('S i) where
  rlensC = unsafeAlens
  {-# INLINE rlensC #-}
  rgetC = aget
  {-# INLINE rgetC #-}
  rputC = unsafeAput @t
  {-# INLINE rputC #-}

-- | Get a subset of a record's fields.
arecGetSubset :: forall rs ss f.
                 (IndexWitnesses (RImage rs ss), NatToInt (RLength rs))
              => ARec f ss -> ARec f 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 $ 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))
         => RecSubset ARec rs ss is where
  rsubsetC f big = fmap (arecSetSubset big) (f (arecGetSubset big))
  {-# INLINE rsubsetC #-}

instance (RPureConstrained (IndexableField rs) rs,
          RecApplicative rs,
          Show (Rec f rs)) => Show (ARec f rs) where
  show = show . fromARec

instance (RPureConstrained (IndexableField rs) rs,
          RecApplicative rs,
          Eq (Rec f rs)) => Eq (ARec f rs) where
  x == y = fromARec x == fromARec y

instance (RPureConstrained (IndexableField rs) rs,
          RecApplicative rs,
          Ord (Rec f rs)) => Ord (ARec f rs) where
  compare x y = compare (fromARec x) (fromARec y)