Advanced.hs

-- | Full record representation
--
-- Intended for qualified import.
--
-- > import Data.Record.Anon.Internal.Advanced (Record)
-- > import Data.Record.Anon.Internal.Advanced qualified as A
module Data.Record.Anon.Internal.Advanced (
    -- * Definition
    Record -- opaque
    -- * Main API
  , Field(..)
  , empty
  , insert
  , insertA
  , get
  , set
  , merge
  , lens
  , project
  , inject
  , applyPending
    -- * Combinators
    -- ** " Foldable "
  , map
  , cmap
    -- ** " Applicable "
  , pure
  , cpure
  , ap
    -- ** " Foldable "
  , collapse
  , toList
    -- ** " Traversable "
  , mapM
  , cmapM
  , sequenceA
  , sequenceA'
  , distribute
  , distribute'
    -- ** Zipping
  , zip
  , zipWith
  , zipWithM
  , czipWith
  , czipWithM
    -- * Reification and reflection
  , reifyKnownFields
  , reflectKnownFields
  , reifyAllFields
  , reflectAllFields
  , InRow(..)
  , reifySubRow
  , reflectSubRow
    -- * Existential records
  , Some(..)
  , SomeRecord(..)
  , someRecord
  ) where

import Prelude hiding (map, mapM, zip, zipWith, sequenceA, pure)
import Prelude qualified

import Control.DeepSeq (NFData (..))
import Data.Aeson (ToJSON(..), FromJSON(..))
import Data.Bifunctor
import Data.Coerce (coerce)
import Data.Functor.Product
import Data.Kind
import Data.Primitive.SmallArray
import Data.Proxy
import Data.Record.Generic hiding (FieldName)
import Data.SOP.Classes (fn_2)
import Data.SOP.Constraint
import Data.Tagged
import GHC.Exts (Any)
import GHC.OverloadedLabels
import GHC.Records qualified as Base
import GHC.Records.Compat qualified as RecordHasfield
import GHC.TypeLits
import Optics.Core qualified as Optics

import Data.Record.Generic.Eq     qualified as Generic
import Data.Record.Generic.JSON   qualified as Generic
import Data.Record.Generic.NFData qualified as Generic
import Data.Record.Generic.Show   qualified as Generic

import Data.Record.Anon.Internal.Core.Canonical (Canonical)
import Data.Record.Anon.Internal.Core.Canonical qualified as Canon
import Data.Record.Anon.Internal.Core.Diff (Diff)
import Data.Record.Anon.Internal.Core.Diff qualified as Diff
import Data.Record.Anon.Internal.Core.FieldName
import Data.Record.Anon.Internal.Reflection (Reflected(..))
import Data.Record.Anon.Internal.Reflection qualified as Unsafe

import Data.Record.Anon.Plugin.Internal.Runtime

{-------------------------------------------------------------------------------
  Definition
-------------------------------------------------------------------------------}

-- | Anonymous record
data Record (f :: k -> Type) (r :: Row k) =
    NoPending  {-# UNPACK #-} !(Canonical f)
  | HasPending {-# UNPACK #-} !(Canonical f) !(Diff f)

type role Record representational representational

{-------------------------------------------------------------------------------
  Conversion
-------------------------------------------------------------------------------}

-- | Construct canonical form of the record (i.e., apply the internal 'Diff')
--
-- This is @O(n)@, and should be done only for operations on records that are
-- @O(n)@ /anyway/, so that the cost can be absorbed.
toCanonical :: Record f r -> Canonical f
toCanonical (NoPending  c)   = c
toCanonical (HasPending c d) = Diff.apply d c

-- | Construct 'Record' from 'Canonical' representation (empty 'Diff')
--
-- This function is unsafe because we cannot verify whether the record matches
-- it's row specification @r@.
unsafeFromCanonical :: Canonical f -> Record f r
unsafeFromCanonical = NoPending

{-------------------------------------------------------------------------------
  HasField from @record-hasfield@
-------------------------------------------------------------------------------}

-- | Proxy for a field name, with 'IsLabel' instance
--
-- The 'IsLabel' instance makes it possible to write
--
-- > #foo
--
-- to mean
--
-- > Field (Proxy @"foo")
data Field n where
  Field :: (KnownSymbol n, KnownHash n) => Proxy n -> Field n

-- | 'Show' instance relies on the 'IsLabel' instance
instance Show (Field n) where
  show (Field p) = "#" ++ symbolVal p

instance (n ~ n', KnownSymbol n, KnownHash n) => IsLabel n' (Field n) where
  fromLabel = Field (Proxy @n)

instance forall k (n :: Symbol) (f :: k -> Type) (r :: Row k) (a :: k).
       (KnownSymbol n, KnownHash n, RowHasField n r a)
    => RecordHasfield.HasField n (Record f r) (f a) where

  -- INLINE pragma important: it makes the 'NoPendingCases' cases very close
  -- to the performance of using a 'SmallArray' directly.
  {-# INLINE hasField #-}
  hasField r = (
        \x -> unsafeSetField ix name x r
      , unsafeGetField ix name r
      )
    where
      name :: FieldName
      name = mkFieldName (Proxy @n)

      ix :: Int
      ix = proxy rowHasField (Proxy @'(n, r, a))

-- | Compile-time construction of a 'FieldName'
mkFieldName :: (KnownSymbol n, KnownHash n) => Proxy n -> FieldName
mkFieldName p = FieldName (hashVal p) (symbolVal p)

instance (RowHasField n r a, KnownSymbol n, KnownHash n)
      => Optics.LabelOptic n Optics.A_Lens (Record f r) (Record f r) (f a) (f a) where
  labelOptic = aux (fromLabel @n)
    where
      aux :: Field n -> Optics.Lens' (Record f r) (f a)
      aux n = Optics.lens (get n) (flip (set n))

-- | Low level field accessor
--
-- It is the responsibility of the plugin to ensure that the field index and
-- the field type match.
unsafeGetField :: forall k (f :: k -> Type) (r :: Row k) (a :: k).
    Int -> FieldName -> Record f r -> f a
unsafeGetField i n = co . \case
    NoPending  c   -> Canon.getAtIndex c i
    HasPending c d -> Diff.get (i, n) d c
  where
    co  :: f Any -> f a
    co = noInlineUnsafeCo

-- | Low level field update
--
-- See comments in 'getField'.
unsafeSetField :: forall k (f :: k -> Type) (r :: Row k) (a :: k).
    Int -> FieldName -> f a -> Record f r -> Record f r
unsafeSetField i n x = \case
    NoPending  c   -> HasPending c (go Diff.empty)
    HasPending c d -> HasPending c (go d)
  where
    go :: Diff f -> Diff f
    go = Diff.set (i, n) (co x)

    co :: f a -> f Any
    co = noInlineUnsafeCo

get :: forall n f r a.
     RowHasField n r a
  => Field n -> Record f r -> f a
get (Field _) = RecordHasfield.getField @n @(Record f r)

set :: forall n f r a.
     RowHasField n r a
  => Field n -> f a -> Record f r -> Record f r
set (Field _) = flip (RecordHasfield.setField @n @(Record f r))

{-------------------------------------------------------------------------------
  Compatibility with HasField from base
-------------------------------------------------------------------------------}

instance (KnownSymbol n, KnownHash n, RowHasField n r a)
      => Base.HasField n (Record f r) (f a) where
  getField = snd . RecordHasfield.hasField @n

{-------------------------------------------------------------------------------
  Main API
-------------------------------------------------------------------------------}

empty :: Record f '[]
empty = NoPending mempty

insert :: forall k (f :: k -> Type) (r :: Row k) (a :: k) (n :: Symbol).
    Field n -> f a -> Record f r -> Record f (n := a : r)
insert (Field n) x = \case
    NoPending  c   -> HasPending c (go Diff.empty)
    HasPending c d -> HasPending c (go d)
  where
    go :: Diff f -> Diff f
    go = Diff.insert (mkFieldName n) (co x)

    co :: f a -> f Any
    co = noInlineUnsafeCo

insertA ::
     Applicative m
  => Field n -> m (f a) -> m (Record f r) -> m (Record f (n := a : r))
insertA f x r = insert f <$> x <*> r

merge :: Record f r -> Record f r' -> Record f (Merge r r')
merge (toCanonical -> r) (toCanonical -> r') =
    unsafeFromCanonical $ r <> r'

lens :: forall f r r'.
     SubRow r r'
  => Record f r -> (Record f r', Record f r' -> Record f r)
lens = \(toCanonical -> r) ->
    bimap getter setter $
      Canon.lens (proxy projectIndices (Proxy @'(r, r'))) r
  where
    getter :: Canonical f -> Record f r'
    getter = unsafeFromCanonical

    setter :: (Canonical f -> Canonical f) -> Record f r' -> Record f r
    setter f (toCanonical -> r) = unsafeFromCanonical (f r)

-- | Project out subrecord
--
-- This is just the 'lens' getter.
project :: SubRow r r' => Record f r -> Record f r'
project = fst . lens

-- | Inject subrecord
--
-- This is just the 'lens' setter.
inject :: SubRow r r' => Record f r' -> Record f r -> Record f r
inject small = ($ small) . snd . lens

applyPending :: Record f r -> Record f r
applyPending (toCanonical -> r) = unsafeFromCanonical r

{-------------------------------------------------------------------------------
  Unconstrained combinators
-------------------------------------------------------------------------------}

map :: (forall x. f x -> g x) -> Record f r -> Record g r
map f (toCanonical -> r) = unsafeFromCanonical $
    Canon.map f r

mapM ::
     Applicative m
  => (forall x. f x -> m (g x))
  -> Record f r -> m (Record g r)
mapM f (toCanonical -> r) = fmap unsafeFromCanonical $
    Canon.mapM f r

zip :: Record f r -> Record g r -> Record (Product f g) r
zip = zipWith Pair

zipWith ::
     (forall x. f x -> g x -> h x)
  -> Record f r -> Record g r -> Record h r
zipWith f (toCanonical -> r) (toCanonical -> r') = unsafeFromCanonical $
    Canon.zipWith f r r'

zipWithM ::
     Applicative m
  => (forall x. f x -> g x -> m (h x))
  -> Record f r -> Record g r -> m (Record h r)
zipWithM f (toCanonical -> r) (toCanonical -> r') = fmap unsafeFromCanonical $
    Canon.zipWithM f r r'

collapse :: Record (K a) r -> [a]
collapse (toCanonical -> r) =
    Canon.collapse r

sequenceA :: Applicative m => Record (m :.: f) r -> m (Record f r)
sequenceA (toCanonical -> r) = fmap unsafeFromCanonical $
    Canon.sequenceA r

sequenceA' :: Applicative m => Record m r -> m (Record I r)
sequenceA' = sequenceA . co
  where
    co :: Record m r -> Record (m :.: I) r
    co = noInlineUnsafeCo

distribute :: forall r m f.
     (Functor m, KnownFields r)
  => m (Record f r) -> Record (m :.: f) r
distribute =
      unsafeFromCanonical
    . Canon.fromRowOrderList
    . (\cs -> fieldVec cs <$> indices)
    . fmap toCanonical
  where
    indices :: [Int]
    indices = Canon.arrayIndicesInRowOrder (length $ proxy fieldNames (Proxy @r))

    fieldVec :: m (Canonical f) -> Int -> (m :.: f) Any
    fieldVec cs idx = Comp $ flip Canon.getAtIndex idx <$> cs

distribute' :: (Functor m, KnownFields r) => m (Record I r) -> Record m r
distribute' = co . distribute
  where
    co :: Record (m :.: I) r -> Record m r
    co = noInlineUnsafeCo

pure :: forall f r. KnownFields r => (forall x. f x) -> Record f r
pure f = unsafeFromCanonical $
    Canon.fromRowOrderList $
      Prelude.map (const f) (proxy fieldNames (Proxy @r))

ap :: Record (f -.-> g) r -> Record f r -> Record g r
ap (toCanonical -> r) (toCanonical -> r') = unsafeFromCanonical $
    Canon.ap r r'

{-------------------------------------------------------------------------------
  Reification and reflection

  The @KnownFields@ constraint on @reifyProject@ is a little dissatisfying, as
  it feels like an orthogonal concern. Ultimately the reason is that in

  > Record f (r :: Row k)

  we have @f :: k -> Type@, as opposed to @f :: Symbol -> k -> Type@. That is
  a generalization we could at some point consider, but until we do, the

  > RowHasField n r a

  constraint introduced in the body 'InRow' involves an /existential/ @n@;
  a /separate/ record with 'KnownSymbol' evidence would therefore not give us
  any information about /this/ @n@.
-------------------------------------------------------------------------------}

reifyKnownFields :: forall k (r :: Row k) proxy.
     KnownFields r
  => proxy r -> Record (K String) r
reifyKnownFields _ =
    unsafeFromCanonical $
      Canon.fromRowOrderList $ co $ proxy fieldNames (Proxy @r)
  where
    co :: [String] -> [K String Any]
    co = coerce

reflectKnownFields :: forall k (r :: Row k).
     Record (K String) r
  -> Reflected (KnownFields r)
reflectKnownFields names =
    Unsafe.reflectKnownFields $ Tagged $ collapse names

reifyAllFields :: forall k (r :: Row k) (c :: k -> Constraint) proxy.
     AllFields r c
  => proxy c -> Record (Dict c) r
reifyAllFields _ = unsafeFromCanonical $
    Canon.fromRowOrderArray $
      fmap aux $ proxy fieldDicts (Proxy @r)
  where
    aux :: DictAny c -> Dict c Any
    aux DictAny = Dict

reflectAllFields :: forall k (c :: k -> Constraint) (r :: Row k).
     Record (Dict c) r
  -> Reflected (AllFields r c)
reflectAllFields dicts =
    Unsafe.reflectAllFields $ Tagged $
      fmap aux $ Canon.toRowOrderArray $ toCanonical dicts
  where
    aux :: Dict c Any -> DictAny c
    aux Dict = DictAny

-- | @InRow r a@ is evidence that there exists some @n@ s.t. @(n := a)@ in @r@.
data InRow (r :: Row k) (a :: k) where
  InRow :: forall k (n :: Symbol) (r :: Row k) (a :: k).
       ( KnownSymbol n
       , RowHasField n r a
       )
    => Proxy n -> InRow r a

reifySubRow :: forall k (r :: Row k) (r' :: Row k).
     (SubRow r r', KnownFields r')
  => Record (InRow r) r'
reifySubRow =
    zipWith aux ixs (reifyKnownFields (Proxy @r'))
  where
    ixs :: Record (K Int) r'
    ixs = unsafeFromCanonical $
           Canon.fromRowOrderList $ co $ proxy projectIndices (Proxy @'(r, r'))

    co :: [Int] -> [K Int Any]
    co = coerce

    aux :: forall x. K Int x -> K String x -> InRow r x
    aux (K i) (K name) =
        case someSymbolVal name of
          SomeSymbol p -> unsafeInRow i p

reflectSubRow :: forall k (r :: Row k) (r' :: Row k).
     Record (InRow r) r'
  -> Reflected (SubRow r r')
reflectSubRow (toCanonical -> ixs) =
    Unsafe.reflectSubRow $ Tagged $
      (\inRow@(InRow p) -> aux inRow p) <$> Canon.toRowOrderList ixs
  where
    aux :: forall x n. RowHasField n r x => InRow r x -> Proxy n -> Int
    aux _ _ = proxy rowHasField (Proxy @'(n, r, x))

unsafeInRow :: forall n r a. KnownSymbol n => Int -> Proxy n -> InRow r a
unsafeInRow i p =
    case reflected of
      Reflected -> InRow p
  where
    reflected :: Reflected (RowHasField n r a)
    reflected = Unsafe.reflectRowHasField $ Tagged i

{-------------------------------------------------------------------------------
  Existential records
-------------------------------------------------------------------------------}

-- | Existential type ("there exists an @x@ such that @f x@")
data Some (f :: k -> Type) where
  Some :: forall k (f :: k -> Type) (x :: k). f x -> Some f

-- | Discovered row variable
--
-- See 'Data.Record.Anon.Advanced.someRecord' for detailed discussion.
data SomeRecord (f :: k -> Type) where
  SomeRecord :: forall k (r :: Row k) (f :: k -> Type).
       KnownFields r
    => Record (Product (InRow r) f) r
    -> SomeRecord f

someRecord :: forall k (f :: k -> Type). [(String, Some f)] -> SomeRecord f
someRecord fields =
    mkSomeRecord $
      unsafeFromCanonical . Canon.fromRowOrderList $
        Prelude.zipWith
          aux
          (Canon.arrayIndicesInRowOrder (length fields))
          (Prelude.map (first someSymbolVal) fields)
  where
    aux :: Int -> (SomeSymbol, Some f) -> Product (InRow r) f Any
    aux i (SomeSymbol n, Some fx) = Pair (unsafeInRow i n) (co fx)

    co :: f x -> f Any
    co = noInlineUnsafeCo

    mkSomeRecord :: forall r. Record (Product (InRow r) f) r -> SomeRecord f
    mkSomeRecord r =
        case reflected of
          Reflected -> SomeRecord r
      where
        reflected :: Reflected (KnownFields r)
        reflected = reflectKnownFields $ map getName r

        getName :: Product (InRow r) f x -> K String x
        getName (Pair (InRow p) _) = K $ symbolVal p

{-------------------------------------------------------------------------------
  Conversion to/from generic 'Rep'
-------------------------------------------------------------------------------}

recordToRep :: Record f r -> Rep I (Record f r)
recordToRep (toCanonical -> r) =
    Rep $ co $ Canon.toRowOrderArray r
  where
    -- Second @Any@ is really (f (Any))
    co :: SmallArray (f Any) -> SmallArray (I Any)
    co = noInlineUnsafeCo

repToRecord :: Rep I (Record f r) -> Record f r
repToRecord (Rep r) =
    unsafeFromCanonical $ Canon.fromRowOrderArray . co $ r
  where
    -- First @Any@ is really (f Any)@
    co :: SmallArray (I Any) -> SmallArray (f Any)
    co = noInlineUnsafeCo

{-------------------------------------------------------------------------------
  Generics instance
-------------------------------------------------------------------------------}

class    (AllFields r (Compose c f), KnownFields r) => RecordConstraints f r c
instance (AllFields r (Compose c f), KnownFields r) => RecordConstraints f r c

recordConstraints :: forall f r c.
     RecordConstraints f r c
  => Proxy c -> Rep (Dict c) (Record f r)
recordConstraints _ = Rep $
    co . aux <$> proxy fieldDicts (Proxy @r)
  where
    aux :: DictAny (Compose c f) -> Dict (Compose c f) Any
    aux DictAny = Dict

    -- The second 'Any' is really (f Any)
    co :: Dict (Compose c f) Any -> Dict c Any
    co = noInlineUnsafeCo

recordMetadata :: forall k (f :: k -> Type) (r :: Row k).
     KnownFields r
  => Metadata (Record f r)
recordMetadata = Metadata {
      recordName          = "Record"
    , recordConstructor   = "ANON_F"
    , recordSize          = length fields
    , recordFieldMetadata = Rep $ smallArrayFromList fields
    }
  where
    fields :: [FieldMetadata Any]
    fields = fieldMetadata (Proxy @r)

instance KnownFields r => Generic (Record f r) where
  type Constraints (Record f r) = RecordConstraints f r
  type MetadataOf  (Record f r) = FieldTypes        f r

  from     = recordToRep
  to       = repToRecord
  dict     = recordConstraints
  metadata = const recordMetadata

{-------------------------------------------------------------------------------
  Instances for standard type classes

  These instances all depend on the generics integration.
-------------------------------------------------------------------------------}

instance RecordConstraints f r Show => Show (Record f r) where
  showsPrec = Generic.gshowsPrec

instance RecordConstraints f r Eq => Eq (Record f r) where
  (==) = Generic.geq

instance ( RecordConstraints f r Eq
         , RecordConstraints f r Ord
         ) => Ord (Record f r) where
  compare = Generic.gcompare

instance RecordConstraints f r NFData => NFData (Record f r) where
  rnf = Generic.grnf

instance RecordConstraints f r ToJSON => ToJSON (Record f r) where
  toJSON = Generic.gtoJSON

instance RecordConstraints f r FromJSON => FromJSON (Record f r) where
  parseJSON = Generic.gparseJSON

{-------------------------------------------------------------------------------
  Constrained combinators
-------------------------------------------------------------------------------}

cpure :: forall r f c.
     AllFields r c
  => Proxy c
  -> (forall x. c x => f x)
  -> Record f r
cpure p f = map (\Dict -> f) $ reifyAllFields p

cmap :: forall r c f g.
     AllFields r c
  => Proxy c
  -> (forall x. c x => f x -> g x)
  -> Record f r -> Record g r
cmap p f = zipWith (\Dict -> f) (reifyAllFields p)

cmapM ::
     (Applicative m, AllFields r c)
  => Proxy c
  -> (forall x. c x => f x -> m (g x))
  -> Record f r -> m (Record g r)
cmapM p f = sequenceA . cmap p (Comp . f)

toList :: forall r a. KnownFields r => Record (K a) r -> [(String, a)]
toList = Prelude.zipWith aux (fieldMetadata (Proxy @r)) . collapse
  where
    aux :: FieldMetadata b -> a -> (String, a)
    aux (FieldMetadata p _) a = (symbolVal p, a)

czipWithM :: forall m r c f g h.
     (Applicative m, AllFields r c)
  => Proxy c
  -> (forall x. c x => f x -> g x -> m (h x))
  -> Record f r -> Record g r -> m (Record h r)
czipWithM p f r r' =
    sequenceA $ map (fn_2 . f') (reifyAllFields p) `ap` r `ap` r'
  where
    f' :: Dict c x -> f x -> g x -> (m :.: h) x
    f' Dict fx gx = Comp $ f fx gx

czipWith ::
     AllFields r c
  => Proxy c
  -> (forall x. c x => f x -> g x -> h x)
  -> Record f r -> Record g r -> Record h r
czipWith p f a b = unI $ czipWithM p (\x y -> I (f x y)) a b