{-# LANGUAGE AllowAmbiguousTypes #-}
{-# 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.
module Data.Vinyl.ARec.Internal
  ( ARec (..)
  , IndexableField
  , 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 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 (..))

-- | An array-backed extensible record with constant-time field
-- access.
newtype ARec (f :: k -> *) (ts :: [k]) = ARec (Array.Array Int Any)
type role ARec representational nominal

-- | 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 (ARec f xs) (ARec g ys)
arecRepsMatchCoercion = Coercion (ARec f xs) (ARec g ys)
forall k (a :: k) (b :: k). Coercible a b => Coercion a b
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 :: Coercion (ARec f xs) (ARec g xs)
arecConsMatchCoercion = Coercion (ARec f xs) (ARec g xs)
forall k (a :: k) (b :: k). Coercible a b => Coercion a b
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

{-
-- 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
-}

-- | 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 :: Rec f ts -> ARec f ts
toARec = ([Any] -> [Any]) -> Rec f ts -> ARec f ts
forall (ts' :: [k]). ([Any] -> [Any]) -> Rec f ts' -> ARec f ts
go [Any] -> [Any]
forall a. a -> a
id
  where go :: ([Any] -> [Any]) -> Rec f ts' -> ARec f ts
        go :: ([Any] -> [Any]) -> Rec f ts' -> ARec f ts
go [Any] -> [Any]
acc Rec f ts'
RNil = Array Int Any -> ARec f ts
forall k (f :: k -> *) (ts :: [k]). Array Int Any -> ARec f ts
ARec (Array Int Any -> ARec f ts) -> Array Int Any -> ARec f ts
forall a b. (a -> b) -> a -> b
$! (Int, Int) -> [Any] -> Array Int Any
forall i e. Ix i => (i, i) -> [e] -> Array i e
Array.listArray (Int
0, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ([Any] -> [Any]
acc [])
        go [Any] -> [Any]
acc (f r
x :& Rec f rs
xs) = ([Any] -> [Any]) -> Rec f rs -> ARec f ts
forall (ts' :: [k]). ([Any] -> [Any]) -> Rec f ts' -> ARec f ts
go ([Any] -> [Any]
acc ([Any] -> [Any]) -> ([Any] -> [Any]) -> [Any] -> [Any]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (f r -> Any
forall a b. a -> b
unsafeCoerce f r
x Any -> [Any] -> [Any]
forall a. a -> [a] -> [a]
:)) Rec f rs
xs
        n :: Int
n = NatToInt (RLength ts) => Int
forall (n :: Nat). NatToInt n => Int
natToInt @(RLength ts)
{-# INLINE toARec #-}

-- | 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 :: ARec f ts -> Rec f ts
fromARec (ARec Array Int Any
arr) = (forall (a :: u). IndexableField ts a => f a) -> Rec f ts
forall u (c :: u -> Constraint) (ts :: [u]) (f :: u -> *).
RPureConstrained c ts =>
(forall (a :: u). c a => f a) -> Rec f ts
rpureConstrained @(IndexableField ts) forall (t :: u). NatToInt (RIndex t ts) => f t
forall (a :: u). IndexableField ts a => f a
aux
  where aux :: forall t. NatToInt (RIndex t ts) => f t
        aux :: f t
aux = Any -> f t
forall a b. a -> b
unsafeCoerce (Array Int Any
arr Array Int Any -> Int -> Any
forall i e. Ix i => Array i e -> i -> e
Array.! NatToInt (RIndex t ts) => Int
forall (n :: Nat). NatToInt n => Int
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 f ts -> f t
aget (ARec Array Int Any
arr) =
  Any -> f t
forall a b. a -> b
unsafeCoerce (Array Int Any -> Int -> Any
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> Int -> e
BArray.unsafeAt Array Int Any
arr (NatToInt (RIndex t ts) => Int
forall (n :: Nat). NatToInt n => Int
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 :: f t' -> ARec f ts -> ARec f ts'
unsafeAput f t'
x (ARec Array Int Any
arr) = Array Int Any -> ARec f ts'
forall k (f :: k -> *) (ts :: [k]). Array Int Any -> ARec f ts
ARec (Array Int Any
arr Array Int Any -> [(Int, Any)] -> Array Int Any
forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
Array.// [(Int
i, f t' -> Any
forall a b. a -> b
unsafeCoerce f t'
x)])
  where i :: Int
i = NatToInt (RIndex t ts) => Int
forall (n :: Nat). NatToInt n => Int
natToInt @(RIndex t ts)
{-# 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 t -> g (f t')) -> ARec f ts -> g (ARec f ts')
unsafeAlens f t -> g (f t')
f ARec f ts
ar = (f t' -> ARec f ts') -> g (f t') -> g (ARec f ts')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((f t' -> ARec f ts -> ARec f ts')
-> ARec f ts -> f t' -> ARec f ts'
forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall (t' :: k) (f :: k -> *) (ts :: [k]) (ts' :: [k]).
NatToInt (RIndex t ts) =>
f t' -> ARec f ts -> ARec f ts'
forall k (t :: k) (t' :: k) (f :: k -> *) (ts :: [k]) (ts' :: [k]).
NatToInt (RIndex t ts) =>
f t' -> ARec f ts -> ARec f ts'
unsafeAput @t) ARec f ts
ar) (f t -> g (f t')
f (ARec f ts -> f t
forall k (t :: k) (f :: k -> *) (ts :: [k]).
NatToInt (RIndex t ts) =>
ARec f ts -> f t
aget ARec f ts
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 :: (f t -> g (f t')) -> ARec f (t : ts) -> g (ARec f (t' : ts))
rlensC = (f t -> g (f t')) -> ARec f (t : ts) -> g (ARec f (t' : ts))
forall k (f :: k -> *) (g :: * -> *) (t :: k) (t' :: k) (ts :: [k])
       (ts' :: [k]).
(Functor g, NatToInt (RIndex t ts)) =>
(f t -> g (f t')) -> ARec f ts -> g (ARec f ts')
unsafeAlens
  {-# INLINE rlensC #-}
  rgetC :: ARec f (t : ts) -> f t
rgetC = ARec f (t : ts) -> f t
forall k (t :: k) (f :: k -> *) (ts :: [k]).
NatToInt (RIndex t ts) =>
ARec f ts -> f t
aget
  {-# INLINE rgetC #-}
  rputC :: f t' -> ARec f (t : ts) -> ARec f (t' : ts)
rputC = forall (t' :: a) (f :: a -> *) (ts :: [a]) (ts' :: [a]).
NatToInt (RIndex t ts) =>
f t' -> ARec f ts -> ARec f ts'
forall k (t :: k) (t' :: k) (f :: k -> *) (ts :: [k]) (ts' :: [k]).
NatToInt (RIndex t ts) =>
f t' -> ARec f ts -> ARec f ts'
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 :: (f t -> g (f t')) -> ARec f (s : ts) -> g (ARec f (s : ts'))
rlensC = (f t -> g (f t')) -> ARec f (s : ts) -> g (ARec f (s : ts'))
forall k (f :: k -> *) (g :: * -> *) (t :: k) (t' :: k) (ts :: [k])
       (ts' :: [k]).
(Functor g, NatToInt (RIndex t ts)) =>
(f t -> g (f t')) -> ARec f ts -> g (ARec f ts')
unsafeAlens
  {-# INLINE rlensC #-}
  rgetC :: ARec f (s : ts) -> f t
rgetC = ARec f (s : ts) -> f t
forall k (t :: k) (f :: k -> *) (ts :: [k]).
NatToInt (RIndex t ts) =>
ARec f ts -> f t
aget
  {-# INLINE rgetC #-}
  rputC :: f t' -> ARec f (s : ts) -> ARec f (s : ts')
rputC = forall (t' :: a) (f :: a -> *) (ts :: [a]) (ts' :: [a]).
NatToInt (RIndex t ts) =>
f t' -> ARec f ts -> ARec f ts'
forall k (t :: k) (t' :: k) (f :: k -> *) (ts :: [k]) (ts' :: [k]).
NatToInt (RIndex t ts) =>
f t' -> ARec f ts -> ARec f ts'
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 f ss -> ARec f rs
arecGetSubset (ARec Array Int Any
arr) = Array Int Any -> ARec f rs
forall k (f :: k -> *) (ts :: [k]). Array Int Any -> ARec f ts
ARec ((Int, Int) -> [Any] -> Array Int Any
forall i e. Ix i => (i, i) -> [e] -> Array i e
Array.listArray (Int
0, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) ([Any] -> Array Int Any) -> [Any] -> Array Int Any
forall a b. (a -> b) -> a -> b
$
                                 [Int] -> [Any]
go (IndexWitnesses (RImage rs ss) => [Int]
forall (is :: [Nat]). IndexWitnesses is => [Int]
indexWitnesses @(RImage rs ss)))
  where go :: [Int] -> [Any]
        go :: [Int] -> [Any]
go = (Int -> Any) -> [Int] -> [Any]
forall a b. (a -> b) -> [a] -> [b]
map (Array Int Any
arr Array Int Any -> Int -> Any
forall i e. Ix i => Array i e -> i -> e
Array.!)
        n :: Int
n = NatToInt (RLength rs) => Int
forall (n :: Nat). NatToInt n => Int
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 f ss -> ARec f rs -> ARec f ss
arecSetSubset (ARec Array Int Any
arrBig) (ARec Array Int Any
arrSmall) = Array Int Any -> ARec f ss
forall k (f :: k -> *) (ts :: [k]). Array Int Any -> ARec f ts
ARec (Array Int Any
arrBig Array Int Any -> [(Int, Any)] -> Array Int Any
forall i e. Ix i => Array i e -> [(i, e)] -> Array i e
Array.// [(Int, Any)]
updates)
  where updates :: [(Int, Any)]
updates = [Int] -> [Any] -> [(Int, Any)]
forall a b. [a] -> [b] -> [(a, b)]
zip (IndexWitnesses (RImage rs ss) => [Int]
forall (is :: [Nat]). IndexWitnesses is => [Int]
indexWitnesses @(RImage rs ss)) (Array Int Any -> [Any]
forall i e. Array i e -> [e]
Array.elems Array Int Any
arrSmall)
{-# INLINE arecSetSubset #-}

instance (is ~ RImage rs ss, IndexWitnesses is, NatToInt (RLength rs))
         => RecSubset ARec rs ss is where
  rsubsetC :: (ARec f rs -> g (ARec f rs)) -> ARec f ss -> g (ARec f ss)
rsubsetC ARec f rs -> g (ARec f rs)
f ARec f ss
big = (ARec f rs -> ARec f ss) -> g (ARec f rs) -> g (ARec f ss)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ARec f ss -> ARec f rs -> ARec f ss
forall k (rs :: [k]) (ss :: [k]) (f :: k -> *).
IndexWitnesses (RImage rs ss) =>
ARec f ss -> ARec f rs -> ARec f ss
arecSetSubset ARec f ss
big) (ARec f rs -> g (ARec f rs)
f (ARec f ss -> ARec f rs
forall k (rs :: [k]) (ss :: [k]) (f :: k -> *).
(IndexWitnesses (RImage rs ss), NatToInt (RLength rs)) =>
ARec f ss -> ARec f rs
arecGetSubset ARec f ss
big))
  {-# INLINE rsubsetC #-}

instance (RPureConstrained (IndexableField rs) rs,
          RecApplicative rs,
          Show (Rec f rs)) => Show (ARec f rs) where
  show :: ARec f rs -> String
show = Rec f rs -> String
forall a. Show a => a -> String
show (Rec f rs -> String)
-> (ARec f rs -> Rec f rs) -> ARec f rs -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ARec f rs -> Rec f rs
forall u (f :: u -> *) (ts :: [u]).
(RecApplicative ts, RPureConstrained (IndexableField ts) ts) =>
ARec f ts -> Rec f ts
fromARec

instance (RPureConstrained (IndexableField rs) rs,
          RecApplicative rs,
          Eq (Rec f rs)) => Eq (ARec f rs) where
  ARec f rs
x == :: ARec f rs -> ARec f rs -> Bool
== ARec f rs
y = ARec f rs -> Rec f rs
forall u (f :: u -> *) (ts :: [u]).
(RecApplicative ts, RPureConstrained (IndexableField ts) ts) =>
ARec f ts -> Rec f ts
fromARec ARec f rs
x Rec f rs -> Rec f rs -> Bool
forall a. Eq a => a -> a -> Bool
== ARec f rs -> Rec f rs
forall u (f :: u -> *) (ts :: [u]).
(RecApplicative ts, RPureConstrained (IndexableField ts) ts) =>
ARec f ts -> Rec f ts
fromARec ARec f rs
y

instance (RPureConstrained (IndexableField rs) rs,
          RecApplicative rs,
          Ord (Rec f rs)) => Ord (ARec f rs) where
  compare :: ARec f rs -> ARec f rs -> Ordering
compare ARec f rs
x ARec f rs
y = Rec f rs -> Rec f rs -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (ARec f rs -> Rec f rs
forall u (f :: u -> *) (ts :: [u]).
(RecApplicative ts, RPureConstrained (IndexableField ts) ts) =>
ARec f ts -> Rec f ts
fromARec ARec f rs
x) (ARec f rs -> Rec f rs
forall u (f :: u -> *) (ts :: [u]).
(RecApplicative ts, RPureConstrained (IndexableField ts) ts) =>
ARec f ts -> Rec f ts
fromARec ARec f rs
y)