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

module ZkFold.Symbolic.Algorithms.Hash.MiMC where

import           Data.Foldable                                  (toList)
import           Data.Functor.Rep                               (fmapRep, liftR3, pureRep)
import           Data.List.NonEmpty                             (NonEmpty ((:|)), nonEmpty)
import           Data.Proxy                                     (Proxy (..))
import           GHC.Generics                                   (Par1 (..), (:*:) (..))
import           Numeric.Natural                                (Natural)
import           Prelude                                        hiding (Eq (..), Num (..), any, length, not, (!!), (/),
                                                                 (^), (||))

import           ZkFold.Base.Algebra.Basic.Class
import           ZkFold.Base.Control.HApplicative               (hpair)
import           ZkFold.Base.Data.HFunctor                      (hmap)
import           ZkFold.Base.Data.Package                       (unpacked)
import           ZkFold.Symbolic.Algorithms.Hash.MiMC.Constants (mimcConstants)
import           ZkFold.Symbolic.Class
import           ZkFold.Symbolic.Data.Class
import           ZkFold.Symbolic.Data.Combinators
import           ZkFold.Symbolic.Data.FieldElement
import           ZkFold.Symbolic.Interpreter
import           ZkFold.Symbolic.MonadCircuit                   (MonadCircuit (newAssigned))

-- | MiMC-2n/n (Feistel) hash function.
-- See https://eprint.iacr.org/2016/492.pdf, page 5
mimcHash2 :: forall c x a.(FromConstant a x, c ~ Context x, Symbolic c, SymbolicData x) => [a] -> a -> x -> x -> x
mimcHash2 :: forall (c :: (Type -> Type) -> Type) x a.
(FromConstant a x, c ~ Context x, Symbolic c, SymbolicData x) =>
[a] -> a -> x -> x -> x
mimcHash2 ((a -> x) -> [a] -> [x]
forall a b. (a -> b) -> [a] -> [b]
map a -> x
forall a b. FromConstant a b => a -> b
fromConstant -> [x]
xs) (forall a b. FromConstant a b => a -> b
fromConstant @a @x -> x
k) =
  case [x] -> Maybe (NonEmpty x)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty ([x] -> [x]
forall a. [a] -> [a]
reverse [x]
xs) of
    Just NonEmpty x
cs -> NonEmpty x -> x -> x -> x
go NonEmpty x
cs
    Maybe (NonEmpty x)
Nothing -> [Char] -> x -> x -> x
forall a. HasCallStack => [Char] -> a
error [Char]
"mimcHash: empty list"
  where
    go :: NonEmpty x -> x -> x -> x
    go :: NonEmpty x -> x -> x -> x
go (x
c :| [x]
cs) x
xL x
xR =
      let
        t5 :: x
t5 = (Support x -> (c (Layout x), Payload x (WitnessField c))) -> x
forall x (c :: (Type -> Type) -> Type).
(SymbolicData x, Context x ~ c) =>
(Support x -> (c (Layout x), Payload x (WitnessField c))) -> x
forall (c :: (Type -> Type) -> Type).
(Context x ~ c) =>
(Support x -> (c (Layout x), Payload x (WitnessField c))) -> x
restore ((Support x -> (c (Layout x), Payload x (WitnessField c))) -> x)
-> (Support x -> (c (Layout x), Payload x (WitnessField c))) -> x
forall a b. (a -> b) -> a -> b
$ \Support x
s ->
          (c (Layout x :*: Layout x)
-> c (Layout x)
-> c (Layout x)
-> CircuitFun
     '[Layout x :*: Layout x, Layout x, Layout x] (Layout x) c
-> c (Layout x)
forall (c :: (Type -> Type) -> Type) (f :: Type -> Type)
       (g :: Type -> Type) (h :: Type -> Type) (k :: Type -> Type).
Symbolic c =>
c f -> c g -> c h -> CircuitFun '[f, g, h] k c -> c k
fromCircuit3F
              (c (Layout x) -> c (Layout x) -> c (Layout x :*: Layout x)
forall {k} (c :: (k -> Type) -> Type) (f :: k -> Type)
       (g :: k -> Type).
HApplicative c =>
c f -> c g -> c (f :*: g)
forall (f :: Type -> Type) (g :: Type -> Type).
c f -> c g -> c (f :*: g)
hpair (x -> Support x -> Context x (Layout x)
forall x. SymbolicData x => x -> Support x -> Context x (Layout x)
arithmetize x
xL Support x
s) (x -> Support x -> Context x (Layout x)
forall x. SymbolicData x => x -> Support x -> Context x (Layout x)
arithmetize x
k Support x
s)) (x -> Support x -> Context x (Layout x)
forall x. SymbolicData x => x -> Support x -> Context x (Layout x)
arithmetize x
c Support x
s) (x -> Support x -> Context x (Layout x)
forall x. SymbolicData x => x -> Support x -> Context x (Layout x)
arithmetize x
xR Support x
s)
              (\ (Layout x i
a1 :*: Layout x i
a2) Layout x i
a3 Layout x i
r
                  -> do Layout x i
s3 <- Layout x (m i) -> m (Layout x i)
forall (t :: Type -> Type) (f :: Type -> Type) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
forall (f :: Type -> Type) a.
Applicative f =>
Layout x (f a) -> f (Layout x a)
sequenceA ((i -> i -> i -> m i)
-> Layout x i -> Layout x i -> Layout x i -> Layout x (m i)
forall (f :: Type -> Type) a b c d.
Representable f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftR3 (\i
p i
h i
t -> ClosedPoly i (BaseField c) -> m i
forall var a w (m :: Type -> Type).
MonadCircuit var a w m =>
ClosedPoly var a -> m var
newAssigned (((i -> x) -> i -> x
forall a b. (a -> b) -> a -> b
$ i
h) ((i -> x) -> x) -> ((i -> x) -> x) -> (i -> x) -> x
forall a. AdditiveSemigroup a => a -> a -> a
+ ((i -> x) -> i -> x
forall a b. (a -> b) -> a -> b
$ i
t) ((i -> x) -> x) -> ((i -> x) -> x) -> (i -> x) -> x
forall a. AdditiveSemigroup a => a -> a -> a
+ ((i -> x) -> i -> x
forall a b. (a -> b) -> a -> b
$ i
p))) Layout x i
a1 Layout x i
a2 Layout x i
a3)
                        Layout x i
p5 <- Layout x (m i) -> m (Layout x i)
forall (t :: Type -> Type) (f :: Type -> Type) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
forall (f :: Type -> Type) a.
Applicative f =>
Layout x (f a) -> f (Layout x a)
sequenceA ((i -> m i) -> Layout x i -> Layout x (m i)
forall (f :: Type -> Type) a b.
Representable f =>
(a -> b) -> f a -> f b
fmapRep (\i
ss -> ClosedPoly i (BaseField c) -> m i
forall var a w (m :: Type -> Type).
MonadCircuit var a w m =>
ClosedPoly var a -> m var
newAssigned (((i -> x) -> i -> x
forall a b. (a -> b) -> a -> b
$ i
ss) ((i -> x) -> x) -> Natural -> (i -> x) -> x
forall a b. Exponent a b => a -> b -> a
^ (Natural
5 :: Natural))) Layout x i
s3)
                        (i -> i -> m i) -> Layout x i -> Layout x i -> m (Layout x i)
forall (f :: Type -> Type) (m :: Type -> Type) a b c.
(Representable f, Traversable f, Applicative m) =>
(a -> b -> m c) -> f a -> f b -> m (f c)
mzipWithMRep (\i
h i
t -> ClosedPoly i (BaseField c) -> m i
forall var a w (m :: Type -> Type).
MonadCircuit var a w m =>
ClosedPoly var a -> m var
newAssigned (((i -> x) -> i -> x
forall a b. (a -> b) -> a -> b
$ i
h) ((i -> x) -> x) -> ((i -> x) -> x) -> (i -> x) -> x
forall a. AdditiveSemigroup a => a -> a -> a
+ ((i -> x) -> i -> x
forall a b. (a -> b) -> a -> b
$ i
t))) Layout x i
p5 Layout x i
r )
          , WitnessField c -> Payload x (WitnessField c)
forall (f :: Type -> Type) a. Representable f => a -> f a
pureRep WitnessField c
forall a. AdditiveMonoid a => a
zero)
      in case [x] -> Maybe (NonEmpty x)
forall a. [a] -> Maybe (NonEmpty a)
nonEmpty [x]
cs of
          Just NonEmpty x
cs' -> NonEmpty x -> x -> x -> x
go NonEmpty x
cs' x
t5 x
xL
          Maybe (NonEmpty x)
Nothing  -> x
t5

mimcHashN :: forall a x. (FromConstant a x, Ring x, SymbolicData x) => [a] -> a -> [x] -> x
mimcHashN :: forall a x.
(FromConstant a x, Ring x, SymbolicData x) =>
[a] -> a -> [x] -> x
mimcHashN [a]
xs a
k = [x] -> x
go
  where
    go :: [x] -> x
go [x]
zs = case [x]
zs of
      []          -> [a] -> a -> x -> x -> x
forall (c :: (Type -> Type) -> Type) x a.
(FromConstant a x, c ~ Context x, Symbolic c, SymbolicData x) =>
[a] -> a -> x -> x -> x
mimcHash2 [a]
xs a
k x
forall a. AdditiveMonoid a => a
zero x
forall a. AdditiveMonoid a => a
zero
      [x
z]         -> [a] -> a -> x -> x -> x
forall (c :: (Type -> Type) -> Type) x a.
(FromConstant a x, c ~ Context x, Symbolic c, SymbolicData x) =>
[a] -> a -> x -> x -> x
mimcHash2 [a]
xs a
k x
forall a. AdditiveMonoid a => a
zero x
z
      [x
zL, x
zR]    -> [a] -> a -> x -> x -> x
forall (c :: (Type -> Type) -> Type) x a.
(FromConstant a x, c ~ Context x, Symbolic c, SymbolicData x) =>
[a] -> a -> x -> x -> x
mimcHash2 [a]
xs a
k x
zL x
zR
      (x
zL:x
zR:[x]
zs') -> [x] -> x
go ([a] -> a -> x -> x -> x
forall (c :: (Type -> Type) -> Type) x a.
(FromConstant a x, c ~ Context x, Symbolic c, SymbolicData x) =>
[a] -> a -> x -> x -> x
mimcHash2 [a]
xs a
k x
zL x
zR x -> [x] -> [x]
forall a. a -> [a] -> [a]
: [x]
zs')

hash :: forall context x a .
    ( SymbolicOutput x
    , BaseField context ~ a
    , Context x ~ context
    ) => x -> FieldElement context
hash :: forall (context :: (Type -> Type) -> Type) x a.
(SymbolicOutput x, BaseField context ~ a, Context x ~ context) =>
x -> FieldElement context
hash = [a] -> a -> [FieldElement context] -> FieldElement context
forall a x.
(FromConstant a x, Ring x, SymbolicData x) =>
[a] -> a -> [x] -> x
mimcHashN [a]
forall a. FromConstant Integer a => [a]
mimcConstants (a
forall a. AdditiveMonoid a => a
zero :: a) ([FieldElement context] -> FieldElement context)
-> (x -> [FieldElement context]) -> x -> FieldElement context
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (context Par1 -> FieldElement context)
-> [context Par1] -> [FieldElement context]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap context Par1 -> FieldElement context
forall (c :: (Type -> Type) -> Type). c Par1 -> FieldElement c
FieldElement ([context Par1] -> [FieldElement context])
-> (x -> [context Par1]) -> x -> [FieldElement context]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. context [] -> [context Par1]
forall (c :: (Type -> Type) -> Type) (f :: Type -> Type).
(Package c, Functor f) =>
c f -> f (c Par1)
unpacked (context [] -> [context Par1])
-> (x -> context []) -> x -> [context Par1]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Layout x a -> [a]) -> context (Layout x) -> context []
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 f -> context g
hmap Layout x a -> [a]
forall a. Layout x a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
toList (context (Layout x) -> context [])
-> (x -> context (Layout x)) -> x -> context []
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (x -> Proxy context -> context (Layout x))
-> Proxy context -> x -> context (Layout x)
forall a b c. (a -> b -> c) -> b -> a -> c
flip x -> Proxy context -> context (Layout x)
x -> Support x -> Context x (Layout x)
forall x. SymbolicData x => x -> Support x -> Context x (Layout x)
arithmetize Proxy context
forall {k} (t :: k). Proxy t
Proxy

mimcHashN' :: forall a x. (FromConstant a x, Arithmetic x) => [a] -> a -> [x] -> x
mimcHashN' :: forall a x.
(FromConstant a x, Arithmetic x) =>
[a] -> a -> [x] -> x
mimcHashN' [a]
xs a
k [x]
as = x
ans
  where
    FieldElement (Interpreter (Par1 x
ans)) = [FieldElement (Interpreter x)] -> FieldElement (Interpreter x)
go ((x -> FieldElement (Interpreter x))
-> [x] -> [FieldElement (Interpreter x)]
forall a b. (a -> b) -> [a] -> [b]
map (Interpreter x Par1 -> FieldElement (Interpreter x)
forall (c :: (Type -> Type) -> Type). c Par1 -> FieldElement c
FieldElement (Interpreter x Par1 -> FieldElement (Interpreter x))
-> (x -> Interpreter x Par1) -> x -> FieldElement (Interpreter x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.Par1 x -> Interpreter x Par1
forall {k} (a :: k) (f :: k -> Type). f a -> Interpreter a f
Interpreter (Par1 x -> Interpreter x Par1)
-> (x -> Par1 x) -> x -> Interpreter x Par1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. x -> Par1 x
forall p. p -> Par1 p
Par1) [x]
as)

    go :: [FieldElement (Interpreter x)] -> FieldElement (Interpreter x)
    go :: [FieldElement (Interpreter x)] -> FieldElement (Interpreter x)
go [FieldElement (Interpreter x)]
zs = case  [FieldElement (Interpreter x)]
zs of
      []          -> [a]
-> a
-> FieldElement (Interpreter x)
-> FieldElement (Interpreter x)
-> FieldElement (Interpreter x)
forall (c :: (Type -> Type) -> Type) x a.
(FromConstant a x, c ~ Context x, Symbolic c, SymbolicData x) =>
[a] -> a -> x -> x -> x
mimcHash2 [a]
xs a
k FieldElement (Interpreter x)
forall a. AdditiveMonoid a => a
zero FieldElement (Interpreter x)
forall a. AdditiveMonoid a => a
zero
      [FieldElement (Interpreter x)
z]         -> [a]
-> a
-> FieldElement (Interpreter x)
-> FieldElement (Interpreter x)
-> FieldElement (Interpreter x)
forall (c :: (Type -> Type) -> Type) x a.
(FromConstant a x, c ~ Context x, Symbolic c, SymbolicData x) =>
[a] -> a -> x -> x -> x
mimcHash2 [a]
xs a
k FieldElement (Interpreter x)
forall a. AdditiveMonoid a => a
zero FieldElement (Interpreter x)
z
      [FieldElement (Interpreter x)
zL, FieldElement (Interpreter x)
zR]    -> [a]
-> a
-> FieldElement (Interpreter x)
-> FieldElement (Interpreter x)
-> FieldElement (Interpreter x)
forall (c :: (Type -> Type) -> Type) x a.
(FromConstant a x, c ~ Context x, Symbolic c, SymbolicData x) =>
[a] -> a -> x -> x -> x
mimcHash2 [a]
xs a
k FieldElement (Interpreter x)
zL FieldElement (Interpreter x)
zR
      (FieldElement (Interpreter x)
zL:FieldElement (Interpreter x)
zR:[FieldElement (Interpreter x)]
zs') -> [FieldElement (Interpreter x)] -> FieldElement (Interpreter x)
go ([a]
-> a
-> FieldElement (Interpreter x)
-> FieldElement (Interpreter x)
-> FieldElement (Interpreter x)
forall (c :: (Type -> Type) -> Type) x a.
(FromConstant a x, c ~ Context x, Symbolic c, SymbolicData x) =>
[a] -> a -> x -> x -> x
mimcHash2 [a]
xs a
k FieldElement (Interpreter x)
zL FieldElement (Interpreter x)
zR FieldElement (Interpreter x)
-> [FieldElement (Interpreter x)] -> [FieldElement (Interpreter x)]
forall a. a -> [a] -> [a]
: [FieldElement (Interpreter x)]
zs')