{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}

module ZkFold.Symbolic.Data.Class (
        SomeData (..),
        SymbolicData (..),
    ) where

import           Control.Applicative              ((<*>))
import           Data.Function                    (flip, (.))
import           Data.Functor                     ((<$>))
import           Data.Functor.Rep                 (Representable (..))
import           Data.Kind                        (Type)
import           Data.Type.Equality               (type (~))
import           Data.Typeable                    (Proxy (..), Typeable)
import           GHC.Generics                     (U1 (..), (:*:) (..), (:.:) (..))

import           ZkFold.Base.Algebra.Basic.Number (KnownNat)
import           ZkFold.Base.Control.HApplicative (HApplicative, hliftA2, hpure)
import           ZkFold.Base.Data.HFunctor        (hmap)
import           ZkFold.Base.Data.Package         (Package, pack)
import           ZkFold.Base.Data.Product         (fstP, sndP)
import           ZkFold.Base.Data.Vector          (Vector)

-- | A class for Symbolic data types.
class SymbolicData x where
    type Context x :: (Type -> Type) -> Type
    type Support x :: Type
    type Layout x :: Type -> Type

    -- | Returns the circuit that makes up `x`.
    pieces :: x -> Support x -> Context x (Layout x)

    -- | Restores `x` from the circuit's outputs.
    restore :: (Support x -> Context x (Layout x)) -> x

-- A wrapper for `SymbolicData` types.
data SomeData c where
    SomeData :: (Typeable t, SymbolicData t, Context t ~ c) => t -> SomeData c

instance SymbolicData (c (f :: Type -> Type)) where
    type Context (c f) = c
    type Support (c f) = Proxy c
    type Layout (c f) = f

    pieces :: c f -> Support (c f) -> Context (c f) (Layout (c f))
pieces c f
x Support (c f)
_ = c f
Context (c f) (Layout (c f))
x
    restore :: (Support (c f) -> Context (c f) (Layout (c f))) -> c f
restore Support (c f) -> Context (c f) (Layout (c f))
f = Support (c f) -> Context (c f) (Layout (c f))
f Proxy c
Support (c f)
forall {k} (t :: k). Proxy t
Proxy

instance HApplicative c => SymbolicData (Proxy (c :: (Type -> Type) -> Type)) where
    type Context (Proxy c) = c
    type Support (Proxy c) = Proxy c
    type Layout (Proxy c) = U1

    pieces :: Proxy c
-> Support (Proxy c) -> Context (Proxy c) (Layout (Proxy c))
pieces Proxy c
_ Support (Proxy c)
_ = (forall a. U1 a) -> c U1
forall {k} (c :: (k -> Type) -> Type) (f :: k -> Type).
HApplicative c =>
(forall (a :: k). f a) -> c f
forall (f :: Type -> Type). (forall a. f a) -> c f
hpure U1 a
forall a. U1 a
forall k (p :: k). U1 p
U1
    restore :: (Support (Proxy c) -> Context (Proxy c) (Layout (Proxy c)))
-> Proxy c
restore Support (Proxy c) -> Context (Proxy c) (Layout (Proxy c))
_ = Proxy c
forall {k} (t :: k). Proxy t
Proxy

instance
    ( SymbolicData x
    , SymbolicData y
    , HApplicative (Context x)
    , Context x ~ Context y
    , Support x ~ Support y
    ) => SymbolicData (x, y) where

    type Context (x, y) = Context x
    type Support (x, y) = Support x
    type Layout (x, y) = Layout x :*: Layout y

    pieces :: (x, y) -> Support (x, y) -> Context (x, y) (Layout (x, y))
pieces (x
a, y
b) = (forall a.
 Layout x a -> Layout y a -> (:*:) (Layout x) (Layout y) a)
-> Context y (Layout x)
-> Context y (Layout y)
-> Context y (Layout x :*: Layout y)
forall {k} (c :: (k -> Type) -> Type) (f :: k -> Type)
       (g :: k -> Type) (h :: k -> Type).
HApplicative c =>
(forall (a :: k). f a -> g a -> h a) -> c f -> c g -> c h
forall (f :: Type -> Type) (g :: Type -> Type) (h :: Type -> Type).
(forall a. f a -> g a -> h a)
-> Context y f -> Context y g -> Context y h
hliftA2 Layout x a -> Layout y a -> (:*:) (Layout x) (Layout y) a
forall a. Layout x a -> Layout y a -> (:*:) (Layout x) (Layout y) a
forall k (f :: k -> Type) (g :: k -> Type) (p :: k).
f p -> g p -> (:*:) f g p
(:*:) (Context y (Layout x)
 -> Context y (Layout y) -> Context y (Layout x :*: Layout y))
-> (Support y -> Context y (Layout x))
-> Support y
-> Context y (Layout y)
-> Context y (Layout x :*: Layout y)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> x -> Support x -> Context x (Layout x)
forall x. SymbolicData x => x -> Support x -> Context x (Layout x)
pieces x
a (Support y
 -> Context y (Layout y) -> Context y (Layout x :*: Layout y))
-> (Support y -> Context y (Layout y))
-> Support y
-> Context y (Layout x :*: Layout y)
forall a b.
(Support y -> a -> b) -> (Support y -> a) -> Support y -> b
forall (f :: Type -> Type) a b.
Applicative f =>
f (a -> b) -> f a -> f b
<*> y -> Support y -> Context y (Layout y)
forall x. SymbolicData x => x -> Support x -> Context x (Layout x)
pieces y
b
    restore :: (Support (x, y) -> Context (x, y) (Layout (x, y))) -> (x, y)
restore Support (x, y) -> Context (x, y) (Layout (x, y))
f = ((Support x -> Context x (Layout x)) -> x
forall x.
SymbolicData x =>
(Support x -> Context x (Layout x)) -> x
restore ((forall a. (:*:) (Layout x) (Layout y) a -> Layout x a)
-> Context y (Layout x :*: Layout y) -> Context y (Layout x)
forall {k} (c :: (k -> Type) -> Type) (f :: k -> Type)
       (g :: k -> Type).
HFunctor c =>
(forall (a :: k). f a -> g a) -> c f -> c g
forall (f :: Type -> Type) (g :: Type -> Type).
(forall a. f a -> g a) -> Context y f -> Context y g
hmap (:*:) (Layout x) (Layout y) a -> Layout x a
forall a. (:*:) (Layout x) (Layout y) a -> Layout x a
forall {k} (f :: k -> Type) (g :: k -> Type) (a :: k).
(:*:) f g a -> f a
fstP (Context y (Layout x :*: Layout y) -> Context y (Layout x))
-> (Support y -> Context y (Layout x :*: Layout y))
-> Support y
-> Context y (Layout x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Support y -> Context y (Layout x :*: Layout y)
Support (x, y) -> Context (x, y) (Layout (x, y))
f), (Support y -> Context y (Layout y)) -> y
forall x.
SymbolicData x =>
(Support x -> Context x (Layout x)) -> x
restore ((forall a. (:*:) (Layout x) (Layout y) a -> Layout y a)
-> Context y (Layout x :*: Layout y) -> Context y (Layout y)
forall {k} (c :: (k -> Type) -> Type) (f :: k -> Type)
       (g :: k -> Type).
HFunctor c =>
(forall (a :: k). f a -> g a) -> c f -> c g
forall (f :: Type -> Type) (g :: Type -> Type).
(forall a. f a -> g a) -> Context y f -> Context y g
hmap (:*:) (Layout x) (Layout y) a -> Layout y a
forall a. (:*:) (Layout x) (Layout y) a -> Layout y a
forall {k} (f :: k -> Type) (g :: k -> Type) (a :: k).
(:*:) f g a -> g a
sndP (Context y (Layout x :*: Layout y) -> Context y (Layout y))
-> (Support y -> Context y (Layout x :*: Layout y))
-> Support y
-> Context y (Layout y)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Support y -> Context y (Layout x :*: Layout y)
Support (x, y) -> Context (x, y) (Layout (x, y))
f))

instance
    ( SymbolicData x
    , SymbolicData y
    , SymbolicData z
    , HApplicative (Context x)
    , Context x ~ Context y
    , Context y ~ Context z
    , Support x ~ Support y
    , Support y ~ Support z
    ) => SymbolicData (x, y, z) where

    type Context (x, y, z) = Context (x, (y, z))
    type Support (x, y, z) = Support (x, (y, z))
    type Layout (x, y, z) = Layout (x, (y, z))

    pieces :: (x, y, z)
-> Support (x, y, z) -> Context (x, y, z) (Layout (x, y, z))
pieces (x
a, y
b, z
c) = (x, (y, z))
-> Support (x, (y, z)) -> Context (x, (y, z)) (Layout (x, (y, z)))
forall x. SymbolicData x => x -> Support x -> Context x (Layout x)
pieces (x
a, (y
b, z
c))
    restore :: (Support (x, y, z) -> Context (x, y, z) (Layout (x, y, z)))
-> (x, y, z)
restore Support (x, y, z) -> Context (x, y, z) (Layout (x, y, z))
f = let (x
a, (y
b, z
c)) = (Support (x, (y, z)) -> Context (x, (y, z)) (Layout (x, (y, z))))
-> (x, (y, z))
forall x.
SymbolicData x =>
(Support x -> Context x (Layout x)) -> x
restore Support (x, (y, z)) -> Context (x, (y, z)) (Layout (x, (y, z)))
Support (x, y, z) -> Context (x, y, z) (Layout (x, y, z))
f in (x
a, y
b, z
c)

instance
    ( SymbolicData w
    , SymbolicData x
    , SymbolicData y
    , SymbolicData z
    , HApplicative (Context x)
    , Context w ~ Context x
    , Context x ~ Context y
    , Context y ~ Context z
    , Support w ~ Support x
    , Support x ~ Support y
    , Support y ~ Support z
    ) => SymbolicData (w, x, y, z) where

    type Context (w, x, y, z) = Context (w, (x, y, z))
    type Support (w, x, y, z) = Support (w, (x, y, z))
    type Layout (w, x, y, z) = Layout (w, (x, y, z))

    pieces :: (w, x, y, z)
-> Support (w, x, y, z)
-> Context (w, x, y, z) (Layout (w, x, y, z))
pieces (w
a, x
b, y
c, z
d) = (w, (x, y, z))
-> Support (w, (x, y, z))
-> Context (w, (x, y, z)) (Layout (w, (x, y, z)))
forall x. SymbolicData x => x -> Support x -> Context x (Layout x)
pieces (w
a, (x
b, y
c, z
d))
    restore :: (Support (w, x, y, z)
 -> Context (w, x, y, z) (Layout (w, x, y, z)))
-> (w, x, y, z)
restore Support (w, x, y, z) -> Context (w, x, y, z) (Layout (w, x, y, z))
f = let (w
a, (x
b, y
c, z
d)) = (Support (w, (x, y, z))
 -> Context (w, (x, y, z)) (Layout (w, (x, y, z))))
-> (w, (x, y, z))
forall x.
SymbolicData x =>
(Support x -> Context x (Layout x)) -> x
restore Support (w, (x, y, z))
-> Context (w, (x, y, z)) (Layout (w, (x, y, z)))
Support (w, x, y, z) -> Context (w, x, y, z) (Layout (w, x, y, z))
f in (w
a, x
b, y
c, z
d)

instance
    ( SymbolicData v
    , SymbolicData w
    , SymbolicData x
    , SymbolicData y
    , SymbolicData z
    , HApplicative (Context x)
    , Context v ~ Context w
    , Context w ~ Context x
    , Context x ~ Context y
    , Context y ~ Context z
    , Support v ~ Support w
    , Support w ~ Support x
    , Support x ~ Support y
    , Support y ~ Support z
    ) => SymbolicData (v, w, x, y, z) where

    type Context (v, w, x, y, z) = Context (v, (w, x, y, z))
    type Support (v, w, x, y, z) = Support (v, (w, x, y, z))
    type Layout (v, w, x, y, z) = Layout (v, (w, x, y, z))

    pieces :: (v, w, x, y, z)
-> Support (v, w, x, y, z)
-> Context (v, w, x, y, z) (Layout (v, w, x, y, z))
pieces (v
a, w
b, x
c, y
d, z
e) = (v, (w, x, y, z))
-> Support (v, (w, x, y, z))
-> Context (v, (w, x, y, z)) (Layout (v, (w, x, y, z)))
forall x. SymbolicData x => x -> Support x -> Context x (Layout x)
pieces (v
a, (w
b, x
c, y
d, z
e))
    restore :: (Support (v, w, x, y, z)
 -> Context (v, w, x, y, z) (Layout (v, w, x, y, z)))
-> (v, w, x, y, z)
restore Support (v, w, x, y, z)
-> Context (v, w, x, y, z) (Layout (v, w, x, y, z))
f = let (v
a, (w
b, x
c, y
d, z
e)) = (Support (v, (w, x, y, z))
 -> Context (v, (w, x, y, z)) (Layout (v, (w, x, y, z))))
-> (v, (w, x, y, z))
forall x.
SymbolicData x =>
(Support x -> Context x (Layout x)) -> x
restore Support (v, (w, x, y, z))
-> Context (v, (w, x, y, z)) (Layout (v, (w, x, y, z)))
Support (v, w, x, y, z)
-> Context (v, w, x, y, z) (Layout (v, w, x, y, z))
f in (v
a, w
b, x
c, y
d, z
e)

instance
    ( SymbolicData x
    , Package (Context x)
    , KnownNat n
    ) => SymbolicData (Vector n x) where

    type Context (Vector n x) = Context x
    type Support (Vector n x) = Support x
    type Layout (Vector n x) = Vector n :.: Layout x

    pieces :: Vector n x
-> Support (Vector n x)
-> Context (Vector n x) (Layout (Vector n x))
pieces Vector n x
xs Support (Vector n x)
i = Vector n (Context x (Layout x))
-> Context x (Vector n :.: Layout x)
forall {k1} (c :: (k1 -> Type) -> Type) (f :: Type -> Type)
       (g :: k1 -> Type).
(Package c, Foldable f, Functor f) =>
f (c g) -> c (f :.: g)
forall (f :: Type -> Type) (g :: Type -> Type).
(Foldable f, Functor f) =>
f (Context x g) -> Context x (f :.: g)
pack ((x -> Support x -> Context x (Layout x))
-> Support x -> x -> Context x (Layout x)
forall a b c. (a -> b -> c) -> b -> a -> c
flip x -> Support x -> Context x (Layout x)
forall x. SymbolicData x => x -> Support x -> Context x (Layout x)
pieces Support x
Support (Vector n x)
i (x -> Context x (Layout x))
-> Vector n x -> Vector n (Context x (Layout x))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Vector n x
xs)
    restore :: (Support (Vector n x)
 -> Context (Vector n x) (Layout (Vector n x)))
-> Vector n x
restore Support (Vector n x) -> Context (Vector n x) (Layout (Vector n x))
f = (Rep (Vector n) -> x) -> Vector n x
forall a. (Rep (Vector n) -> a) -> Vector n a
forall (f :: Type -> Type) a.
Representable f =>
(Rep f -> a) -> f a
tabulate (\Rep (Vector n)
i -> (Support x -> Context x (Layout x)) -> x
forall x.
SymbolicData x =>
(Support x -> Context x (Layout x)) -> x
restore ((forall a. (:.:) (Vector n) (Layout x) a -> Layout x a)
-> Context x (Vector n :.: Layout x) -> Context x (Layout x)
forall {k} (c :: (k -> Type) -> Type) (f :: k -> Type)
       (g :: k -> Type).
HFunctor c =>
(forall (a :: k). f a -> g a) -> c f -> c g
forall (f :: Type -> Type) (g :: Type -> Type).
(forall a. f a -> g a) -> Context x f -> Context x g
hmap ((Vector n (Layout x a) -> Rep (Vector n) -> Layout x a)
-> Rep (Vector n) -> Vector n (Layout x a) -> Layout x a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Vector n (Layout x a) -> Rep (Vector n) -> Layout x a
forall a. Vector n a -> Rep (Vector n) -> a
forall (f :: Type -> Type) a. Representable f => f a -> Rep f -> a
index Rep (Vector n)
i (Vector n (Layout x a) -> Layout x a)
-> ((:.:) (Vector n) (Layout x) a -> Vector n (Layout x a))
-> (:.:) (Vector n) (Layout x) a
-> Layout x a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (:.:) (Vector n) (Layout x) a -> Vector n (Layout x a)
forall k2 k1 (f :: k2 -> Type) (g :: k1 -> k2) (p :: k1).
(:.:) f g p -> f (g p)
unComp1) (Context x (Vector n :.: Layout x) -> Context x (Layout x))
-> (Support x -> Context x (Vector n :.: Layout x))
-> Support x
-> Context x (Layout x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Support x -> Context x (Vector n :.: Layout x)
Support (Vector n x) -> Context (Vector n x) (Layout (Vector n x))
f))

instance SymbolicData f => SymbolicData (x -> f) where
    type Context (x -> f) = Context f
    type Support (x -> f) = (x, Support f)
    type Layout (x -> f) = Layout f

    pieces :: (x -> f) -> Support (x -> f) -> Context (x -> f) (Layout (x -> f))
pieces x -> f
f (x
x, Support f
i) = f -> Support f -> Context f (Layout f)
forall x. SymbolicData x => x -> Support x -> Context x (Layout x)
pieces (x -> f
f x
x) Support f
i
    restore :: (Support (x -> f) -> Context (x -> f) (Layout (x -> f))) -> x -> f
restore Support (x -> f) -> Context (x -> f) (Layout (x -> f))
f x
x = (Support f -> Context f (Layout f)) -> f
forall x.
SymbolicData x =>
(Support x -> Context x (Layout x)) -> x
restore ((x, Support f) -> Context f (Layout f)
Support (x -> f) -> Context (x -> f) (Layout (x -> f))
f ((x, Support f) -> Context f (Layout f))
-> (Support f -> (x, Support f))
-> Support f
-> Context f (Layout f)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (x
x,))