-- Copyright 2019-2021 Google LLC
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--      http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

-- | An approximation of a dependent pair type.

module Data.Ten.Sigma
         ( (:**)(..), overFragment, lmapFragment, eqKey
         , OpCostar(..), caseFragment
         ) where

import Data.Functor.Contravariant (Contravariant(..))
import Data.Maybe (fromMaybe, isJust)
import Data.Type.Equality ((:~:)(Refl), TestEquality(..))

import Control.DeepSeq (NFData(..))
import Data.GADT.Compare (GEq(..))
import Data.Portray (Portray(..), Portrayal(..), infixr_)
import Data.Portray.Diff (Diff(..), diffVs)

import Data.Ten.Entails ((:!:), Entails, withEntailment)
import Data.Ten.Foldable (Foldable10(..))
import Data.Ten.Foldable.WithIndex (Foldable10WithIndex(..))
import Data.Ten.Functor (Functor10(..))
import Data.Ten.Functor.WithIndex (Index10, Functor10WithIndex(..))
import Data.Ten.Representable (Representable10(..))
import Data.Ten.Traversable (Traversable10(..))
import Data.Ten.Traversable.WithIndex (Traversable10WithIndex(..))
import Data.Ten.Update (Update10, overRep10)

infixr 5 :**
-- | A pair of @k a@ and @m a@ for any (existential) @a@.
--
-- This is a lot like a dependent pair, in that it contains a left-hand-side
-- value that's meant to identify a type, and a right-hand-side parameterized
-- by that type.  For example, the true dependent pair type (in e.g. Idris)
-- @(n :: Nat ** Vec n Bool)@ could be approximated in Haskell as
-- @SInt :** Ap10 Bool Vec@.
--
-- This can be used to represent one field of a 'Representable10', where @k@ is
-- set to @Rep10 f@.  The @k a@ identifies which field (and locks down its
-- type), and the @m a@ provides its value.
data k :** m = forall a. k a :** m a

instance (forall a. NFData (k a), Entails k (NFData :!: m))
      => NFData (k :** m) where
  rnf :: (k :** m) -> ()
rnf (k a
k :** m a
m) = k a -> ((:!:) NFData m a => ()) -> ()
forall k1 (c :: k1 -> Constraint) (k2 :: k1 -> *) (a :: k1) r.
Entails k2 c =>
k2 a -> (c a => r) -> r
withEntailment @(NFData :!: m) k a
k (((:!:) NFData m a => ()) -> ()) -> ((:!:) NFData m a => ()) -> ()
forall a b. (a -> b) -> a -> b
$ k a -> ()
forall a. NFData a => a -> ()
rnf k a
k () -> () -> ()
`seq` m a -> ()
forall a. NFData a => a -> ()
rnf m a
m

instance (GEq k, Entails k (Eq :!: m))
      => Eq (k :** m) where
  (k a
kl :** m a
ml) == :: (k :** m) -> (k :** m) -> Bool
== (k a
kr :** m a
mr) = case k a -> k a -> Maybe (a :~: a)
forall k (f :: k -> *) (a :: k) (b :: k).
GEq f =>
f a -> f b -> Maybe (a :~: b)
geq k a
kl k a
kr of
     Maybe (a :~: a)
Nothing -> Bool
False
     Just a :~: a
Refl -> k a -> ((:!:) Eq m a => Bool) -> Bool
forall k1 (c :: k1 -> Constraint) (k2 :: k1 -> *) (a :: k1) r.
Entails k2 c =>
k2 a -> (c a => r) -> r
withEntailment @(Eq :!: m) k a
kl (((:!:) Eq m a => Bool) -> Bool) -> ((:!:) Eq m a => Bool) -> Bool
forall a b. (a -> b) -> a -> b
$ m a
ml m a -> m a -> Bool
forall a. Eq a => a -> a -> Bool
== m a
m a
mr

{-
instance ( Arbitrary (Exists k)
         , Representable10 rec, Entails k (Arbitrary :!: m)
         )
      => Arbitrary (k :** m) where
  arbitrary = do
    Exists k <- arbitrary
    a <- withEntailment @(Arbitrary :!: m) k arbitrary
    return $ k :** a
-}

instance (forall a. Show (k a), Entails k (Show :!: m))
      => Show (k :** m) where
  -- We have to write this by hand because the derived version doesn't know how
  -- to call into 'Constrained10' to find a 'Show' instance for @a@ based on
  -- @ka@.
  showsPrec :: Int -> (k :** m) -> ShowS
showsPrec Int
p (k a
ka :** m a
ma) = Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
prec) (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$
    Int -> k a -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
prec) k a
ka ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    String -> ShowS
showString String
" :** " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    k a -> ((:!:) Show m a => ShowS) -> ShowS
forall k1 (c :: k1 -> Constraint) (k2 :: k1 -> *) (a :: k1) r.
Entails k2 c =>
k2 a -> (c a => r) -> r
withEntailment @(Show :!: m) k a
ka (Int -> m a -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
prec) m a
ma)
   where
    prec :: Int
prec = Int
5

instance (forall a. Portray (k a), Entails k (Portray :!: m))
      => Portray (k :** m) where
  portray :: (k :** m) -> Portrayal
portray (k a
ka :** m a
ma) = k a -> ((:!:) Portray m a => Portrayal) -> Portrayal
forall k1 (c :: k1 -> Constraint) (k2 :: k1 -> *) (a :: k1) r.
Entails k2 c =>
k2 a -> (c a => r) -> r
withEntailment @(Portray :!: m) k a
ka (((:!:) Portray m a => Portrayal) -> Portrayal)
-> ((:!:) Portray m a => Portrayal) -> Portrayal
forall a b. (a -> b) -> a -> b
$
    Ident -> Infixity -> Portrayal -> Portrayal -> Portrayal
Binop Ident
":**" (Rational -> Infixity
infixr_ Rational
5) (k a -> Portrayal
forall a. Portray a => a -> Portrayal
portray k a
ka) (m a -> Portrayal
forall a. Portray a => a -> Portrayal
portray m a
ma)

instance ( TestEquality k, forall a. Portray (k a), forall a. Diff (k a)
         , Entails k (Portray :!: m), Entails k (Diff :!: m)
         )
      => Diff (k :** m) where
  diff :: (k :** m) -> (k :** m) -> Maybe Portrayal
diff (k a
ka :** m a
ma) (k a
kb :** m a
mb) = case k a -> k a -> Maybe (a :~: a)
forall k (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality k a
ka k a
kb of
    Just a :~: a
Refl -> k a -> ((:!:) Diff m a => Maybe Portrayal) -> Maybe Portrayal
forall k1 (c :: k1 -> Constraint) (k2 :: k1 -> *) (a :: k1) r.
Entails k2 c =>
k2 a -> (c a => r) -> r
withEntailment @(Diff :!: m) k a
ka (((:!:) Diff m a => Maybe Portrayal) -> Maybe Portrayal)
-> ((:!:) Diff m a => Maybe Portrayal) -> Maybe Portrayal
forall a b. (a -> b) -> a -> b
$
      case (k a -> k a -> Maybe Portrayal
forall a. Diff a => a -> a -> Maybe Portrayal
diff k a
ka k a
k a
kb, m a -> m a -> Maybe Portrayal
forall a. Diff a => a -> a -> Maybe Portrayal
diff m a
ma m a
m a
mb) of
        (Maybe Portrayal
Nothing, Maybe Portrayal
Nothing) -> Maybe Portrayal
forall a. Maybe a
Nothing
        (Maybe Portrayal
dk, Maybe Portrayal
dm) ->
          Portrayal -> Maybe Portrayal
forall a. a -> Maybe a
Just (Portrayal -> Maybe Portrayal) -> Portrayal -> Maybe Portrayal
forall a b. (a -> b) -> a -> b
$ Ident -> Infixity -> Portrayal -> Portrayal -> Portrayal
Binop Ident
":**" (Rational -> Infixity
infixr_ Rational
5)
            (Portrayal -> Maybe Portrayal -> Portrayal
forall a. a -> Maybe a -> a
fromMaybe (k a -> Portrayal
forall a. Portray a => a -> Portrayal
portray k a
ka) Maybe Portrayal
dk)
            (Portrayal -> Maybe Portrayal -> Portrayal
forall a. a -> Maybe a -> a
fromMaybe (Text -> Portrayal
Opaque Text
"_") Maybe Portrayal
dm)
    Maybe (a :~: a)
Nothing   -> Portrayal -> Maybe Portrayal
forall a. a -> Maybe a
Just (Portrayal -> Maybe Portrayal) -> Portrayal -> Maybe Portrayal
forall a b. (a -> b) -> a -> b
$
      k a -> ((:!:) Portray m a => Portrayal) -> Portrayal
forall k1 (c :: k1 -> Constraint) (k2 :: k1 -> *) (a :: k1) r.
Entails k2 c =>
k2 a -> (c a => r) -> r
withEntailment @(Portray :!: m) k a
ka ((k :** m) -> Portrayal
forall a. Portray a => a -> Portrayal
portray (k a
ka k a -> m a -> k :** m
forall (k :: * -> *) (m :: * -> *) a. k a -> m a -> k :** m
:** m a
ma)) Portrayal -> Portrayal -> Portrayal
`diffVs`
      k a -> ((:!:) Portray m a => Portrayal) -> Portrayal
forall k1 (c :: k1 -> Constraint) (k2 :: k1 -> *) (a :: k1) r.
Entails k2 c =>
k2 a -> (c a => r) -> r
withEntailment @(Portray :!: m) k a
kb ((k :** m) -> Portrayal
forall a. Portray a => a -> Portrayal
portray (k a
kb k a -> m a -> k :** m
forall (k :: * -> *) (m :: * -> *) a. k a -> m a -> k :** m
:** m a
mb))

instance Functor10 ((:**) k) where fmap10 :: (forall a. m a -> n a) -> (k :** m) -> k :** n
fmap10 forall a. m a -> n a
f (k a
k :** m a
m) = k a
k k a -> n a -> k :** n
forall (k :: * -> *) (m :: * -> *) a. k a -> m a -> k :** m
:** m a -> n a
forall a. m a -> n a
f m a
m
instance Foldable10 ((:**) k) where foldMap10 :: (forall a. m a -> w) -> (k :** m) -> w
foldMap10 forall a. m a -> w
f (k a
_ :** m a
m) = m a -> w
forall a. m a -> w
f m a
m
instance Traversable10 ((:**) k) where
  mapTraverse10 :: ((k :** n) -> r) -> (forall a. m a -> f (n a)) -> (k :** m) -> f r
mapTraverse10 (k :** n) -> r
r forall a. m a -> f (n a)
f (k a
k :** m a
m) = (k :** n) -> r
r ((k :** n) -> r) -> (n a -> k :** n) -> n a -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (k a
k k a -> n a -> k :** n
forall (k :: * -> *) (m :: * -> *) a. k a -> m a -> k :** m
:**) (n a -> r) -> f (n a) -> f r
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m a -> f (n a)
forall a. m a -> f (n a)
f m a
m

type instance Index10 ((:**) k) = k
instance Functor10WithIndex ((:**) k) where imap10 :: (forall a. Index10 ((:**) k) a -> m a -> n a)
-> (k :** m) -> k :** n
imap10 forall a. Index10 ((:**) k) a -> m a -> n a
f (k a
k :** m a
m) = k a
k k a -> n a -> k :** n
forall (k :: * -> *) (m :: * -> *) a. k a -> m a -> k :** m
:** Index10 ((:**) k) a -> m a -> n a
forall a. Index10 ((:**) k) a -> m a -> n a
f k a
Index10 ((:**) k) a
k m a
m
instance Foldable10WithIndex ((:**) k) where ifoldMap10 :: (forall a. Index10 ((:**) k) a -> m a -> w) -> (k :** m) -> w
ifoldMap10 forall a. Index10 ((:**) k) a -> m a -> w
f (k a
k :** m a
m) = Index10 ((:**) k) a -> m a -> w
forall a. Index10 ((:**) k) a -> m a -> w
f k a
Index10 ((:**) k) a
k m a
m
instance Traversable10WithIndex ((:**) k) where
  imapTraverse10 :: ((k :** n) -> r)
-> (forall a. Index10 ((:**) k) a -> m a -> g (n a))
-> (k :** m)
-> g r
imapTraverse10 (k :** n) -> r
r forall a. Index10 ((:**) k) a -> m a -> g (n a)
f (k a
k :** m a
m) = (k :** n) -> r
r ((k :** n) -> r) -> (n a -> k :** n) -> n a -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (k a
k k a -> n a -> k :** n
forall (k :: * -> *) (m :: * -> *) a. k a -> m a -> k :** m
:**) (n a -> r) -> g (n a) -> g r
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Index10 ((:**) k) a -> m a -> g (n a)
forall a. Index10 ((:**) k) a -> m a -> g (n a)
f k a
Index10 ((:**) k) a
k m a
m

-- | Check if two pairs have the same key.
eqKey :: GEq k => k :** m -> k :** n -> Bool
eqKey :: (k :** m) -> (k :** n) -> Bool
eqKey (k a
kl :** m a
_) (k a
kr :** n a
_) = Maybe (a :~: a) -> Bool
forall a. Maybe a -> Bool
isJust (k a -> k a -> Maybe (a :~: a)
forall k (f :: k -> *) (a :: k) (b :: k).
GEq f =>
f a -> f b -> Maybe (a :~: b)
geq k a
kl k a
kr)

-- | "Zip" a single field of a record with a (':**').
--
-- Since we're only operating on a single field, the @n@ type can't vary like
-- in a traditional zip function.
overFragment
  :: Update10 rec
  => (forall a. m a -> n a -> n a) -> Rep10 rec :** m -> rec n -> rec n
overFragment :: (forall a. m a -> n a -> n a)
-> (Rep10 rec :** m) -> rec n -> rec n
overFragment forall a. m a -> n a -> n a
f (Rep10 rec a
k :** m a
x) = Rep10 rec a -> (n a -> n a) -> rec n -> rec n
forall k (f :: (k -> *) -> *) (a :: k) (m :: k -> *).
Update10 f =>
Rep10 f a -> (m a -> m a) -> f m -> f m
overRep10 Rep10 rec a
k (m a -> n a -> n a
forall a. m a -> n a -> n a
f m a
x)

-- | Newtype used in implementing contravariant conversion of Fragments.  See
-- 'lmapFragment'.  Only exported because it's used in the type of
-- 'lmapFragment', but it can be largely ignored, like the many "ALens" etc.
-- types in "lens".
newtype OpCostar f r a = OpCostar { OpCostar f r a -> f a -> r
getOpCostar :: f a -> r }

instance Functor f => Contravariant (OpCostar f r) where
  contramap :: (a -> b) -> OpCostar f r b -> OpCostar f r a
contramap a -> b
f (OpCostar f b -> r
g) = (f a -> r) -> OpCostar f r a
forall (f :: * -> *) r a. (f a -> r) -> OpCostar f r a
OpCostar (f b -> r
g (f b -> r) -> (f a -> f b) -> f a -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f)

-- | Simulate a case statement on a (':**') with a record of functions.
--
-- @
--     caseFragment (MyRecord1 (OpCostar isJust) (OpCostar isNothing)) x
-- @
--
-- Is analogous to (pseudo-code):
--
-- @
--     case x of { (_mr1A :** mx) -> isJust mx; (_mr1B :** mx) -> isNothing mx }
-- @
--
-- This is just the action of `Representable10` (whereby @f m@ is isomorphic to
-- @forall a. Rep10 f a -> m a@) plus some newtyping:
--
-- @
--     f (OpCostar m r)                      ~=  (by Representable10)
--     forall a. Rep10 f a -> OpCostar m r a ~=  (by newtype)
--     forall a. Rep10 f a -> f a -> r       ~=  (by GADT constructor)
--     Rep10 f :** m -> r
-- @
caseFragment
  :: Representable10 f
  => f (OpCostar m r)
  -> Rep10 f :** m -> r
caseFragment :: f (OpCostar m r) -> (Rep10 f :** m) -> r
caseFragment f (OpCostar m r)
fco (Rep10 f a
k :** m a
v) = OpCostar m r a -> m a -> r
forall (f :: * -> *) r a. OpCostar f r a -> f a -> r
getOpCostar (f (OpCostar m r)
fco f (OpCostar m r) -> Rep10 f a -> OpCostar m r a
forall k (f :: (k -> *) -> *) (m :: k -> *) (a :: k).
Representable10 f =>
f m -> Rep10 f a -> m a
`index10` Rep10 f a
k) m a
v

-- | Convert a (':**') to a different key type contravariantly.
--
-- Example usage:
--
--     data MyRecord1 m = MyRecord1 { _mr1A :: Ap10 Int m, _mr1B :: Ap10 Int m }
--     data MyRecord2 m = MyRecord2 { _mr2A :: Ap10 Int m }
--
--     -- Collapse both fields _mr1A and _mr1B onto _mr2A.
--     example
--       :: Rep10 MyRecord1 :** Identity
--       -> Rep10 MyRecord2 :** Identity
--     example = lmapFragment $ \MyRecord2{..} -> MyRecord1
--       { _mr1A = _mr2A
--       , _mr1B = _mr2A
--       }
--
-- It looks weird that the argument converts from @recB@ to @recA@ in order
-- to convert (':**') the other way, so it merits some explanation: first,
-- note that, by @'Representable10' recA@, we know that @recA m@ is
-- isomorphic to @forall a. 'Rep10' recA a -> m a@.  That is, @Rep10 recA@
-- effectively appears in negative position in @recA m@.  So, a function from
-- @recB@ to @recA@ hand-wavingly contains a function in the opposite
-- direction from @Rep10 recA@ to @Rep10 recB@.
--
-- With the intuition out of the way, here's how we actually accomplish the
-- conversion: start off with a record @recB@ where each field is a function
-- that trivially rebuilds the corresponding @(:**)@ in each field with
-- @k :: Rep10 recB@ we literally just put @(k :**)@ with the appropriate
-- newtype constructors.  Then, apply the user's contravariant conversion
-- function, to turn our @recB@ of @recB@-pair-builders into an
-- @recA@ of @recB@-pair-builders.  If the user-provided conversion
-- function involves changing any field types, it must have done so by
-- @contramap@ping the pair-builders: instead of a function that just
-- directly applies @(k :=)@ to its argument, they will now contain functions
-- equivalent to @\ma -> k := _f ma@.  Finally, unpack the @recA@ pair
-- and use its @k@ to fetch that field's @recB@-pair-builder (potentially
-- with a conversion inserted at the front), and apply it to the payload.
--
-- Usage will typically involve applying contramap to some number of fields and
-- leaving the rest unchanged.  If you have a type-changing
-- 'Control.Lens.Setter' at hand, it's probably easier to use
-- 'Data.Ten.Lens.fragmented'.
lmapFragment
  :: forall recA recB m f
   . ( Representable10 recA, Representable10 recB
     , f ~ OpCostar m (Rep10 recB :** m)
     )
  => (recB f -> recA f)
  -> Rep10 recA :** m -> Rep10 recB :** m
lmapFragment :: (recB f -> recA f) -> (Rep10 recA :** m) -> Rep10 recB :** m
lmapFragment recB f -> recA f
f = recA (OpCostar m (Rep10 recB :** m))
-> (Rep10 recA :** m) -> Rep10 recB :** m
forall (f :: (* -> *) -> *) (m :: * -> *) r.
Representable10 f =>
f (OpCostar m r) -> (Rep10 f :** m) -> r
caseFragment recA (OpCostar m (Rep10 recB :** m))
fragmentBuilders
 where
  fragmentBuilders :: recA (OpCostar m (Rep10 recB :** m))
  fragmentBuilders :: recA (OpCostar m (Rep10 recB :** m))
fragmentBuilders = recB f -> recA f
f ((forall a. Rep10 recB a -> f a) -> recB f
forall k (f :: (k -> *) -> *) (m :: k -> *).
Representable10 f =>
(forall (a :: k). Rep10 f a -> m a) -> f m
tabulate10 (\Rep10 recB a
k' -> (m a -> Rep10 recB :** m) -> OpCostar m (Rep10 recB :** m) a
forall (f :: * -> *) r a. (f a -> r) -> OpCostar f r a
OpCostar (Rep10 recB a
k' Rep10 recB a -> m a -> Rep10 recB :** m
forall (k :: * -> *) (m :: * -> *) a. k a -> m a -> k :** m
:**)))