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

module ZkFold.Base.Protocol.IVC.Oracle where

import qualified Data.Vector                                    as V
import           GHC.Generics
import           Prelude                                        (map, (.))
import qualified Prelude                                        as P

import           ZkFold.Base.Algebra.Basic.Class
import           ZkFold.Symbolic.Algorithms.Hash.MiMC           (mimcHashN)
import           ZkFold.Symbolic.Algorithms.Hash.MiMC.Constants (mimcConstants)

-- TODO: add more specific instances for efficiency

class HashAlgorithm algo a where
    hash :: [a] -> a

data MiMCHash
instance Ring a => HashAlgorithm MiMCHash a where
    hash :: [a] -> a
hash = forall a x. (FromConstant a x, Ring x) => [a] -> a -> [x] -> x
mimcHashN @a [a]
forall a. FromConstant Integer a => [a]
mimcConstants a
forall a. AdditiveMonoid a => a
zero

class RandomOracle algo x a where
    oracle :: x -> a
    default oracle :: (Generic x, RandomOracle' algo (Rep x) a) => x -> a
    oracle = forall (algo :: k) (f :: Type -> Type) a x.
RandomOracle' algo f a =>
f x -> a
forall {k} {k} (algo :: k) (f :: k -> Type) a (x :: k).
RandomOracle' algo f a =>
f x -> a
oracle' @algo (Rep x Any -> a) -> (x -> Rep x Any) -> x -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. x -> Rep x Any
forall x. x -> Rep x x
forall a x. Generic a => a -> Rep a x
from

instance (FromConstant P.Integer a, HashAlgorithm algo a) => RandomOracle algo P.Integer a where
    oracle :: Integer -> a
oracle = forall (algo :: k) x a. RandomOracle algo x a => x -> a
forall {k} (algo :: k) x a. RandomOracle algo x a => x -> a
oracle @algo @a (a -> a) -> (Integer -> a) -> Integer -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> a
forall a b. FromConstant a b => a -> b
fromConstant

instance HashAlgorithm algo a => RandomOracle algo a a where
    oracle :: a -> a
oracle a
x = forall (algo :: k) a. HashAlgorithm algo a => [a] -> a
forall {k} (algo :: k) a. HashAlgorithm algo a => [a] -> a
hash @algo [a
x]

instance HashAlgorithm algo a => RandomOracle algo (a, a) a where
    oracle :: (a, a) -> a
oracle (a
x, a
y) = forall (algo :: k) a. HashAlgorithm algo a => [a] -> a
forall {k} (algo :: k) a. HashAlgorithm algo a => [a] -> a
hash @algo [a
x, a
y]

instance HashAlgorithm algo a => RandomOracle algo [a] a where
    oracle :: [a] -> a
oracle = forall (algo :: k) a. HashAlgorithm algo a => [a] -> a
forall {k} (algo :: k) a. HashAlgorithm algo a => [a] -> a
hash @algo

instance (HashAlgorithm algo b, RandomOracle algo a b) => RandomOracle algo [a] b where
    oracle :: [a] -> b
oracle = forall (algo :: k) a. HashAlgorithm algo a => [a] -> a
forall {k} (algo :: k) a. HashAlgorithm algo a => [a] -> a
hash @algo ([b] -> b) -> ([a] -> [b]) -> [a] -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> [a] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (forall (algo :: k) x a. RandomOracle algo x a => x -> a
forall {k} (algo :: k) x a. RandomOracle algo x a => x -> a
oracle @algo)

instance (HashAlgorithm algo b, RandomOracle algo a b) => RandomOracle algo (V.Vector a) b where
    oracle :: Vector a -> b
oracle = (forall (algo :: k) x a. RandomOracle algo x a => x -> a
forall {k} (algo :: k) x a. RandomOracle algo x a => x -> a
oracle @algo) ([a] -> b) -> (Vector a -> [a]) -> Vector a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> [a]
forall a. Vector a -> [a]
V.toList

instance {-# OVERLAPPABLE #-} (Generic x, RandomOracle' algo (Rep x) a) => RandomOracle algo x a

class RandomOracle' algo f a where
    oracle' :: f x -> a

-- TODO: fix this instance
instance (RandomOracle' algo f b, RandomOracle' algo g b, HashAlgorithm algo b, Ring b) => RandomOracle' algo (f :+: g) b where
    oracle' :: forall (x :: k). (:+:) f g x -> b
oracle' (L1 f x
x) = forall (algo :: k) a. HashAlgorithm algo a => [a] -> a
forall {k} (algo :: k) a. HashAlgorithm algo a => [a] -> a
hash @algo [b
forall a. AdditiveMonoid a => a
zero, forall (algo :: k) (f :: k -> Type) a (x :: k).
RandomOracle' algo f a =>
f x -> a
forall {k} {k} (algo :: k) (f :: k -> Type) a (x :: k).
RandomOracle' algo f a =>
f x -> a
oracle' @algo f x
x]
    oracle' (R1 g x
x) = forall (algo :: k) (f :: k -> Type) a (x :: k).
RandomOracle' algo f a =>
f x -> a
forall {k} {k} (algo :: k) (f :: k -> Type) a (x :: k).
RandomOracle' algo f a =>
f x -> a
oracle' @algo g x
x

instance (RandomOracle' algo f a, RandomOracle' algo g a, HashAlgorithm algo a) => RandomOracle' algo (f :*: g) a where
    oracle' :: forall (x :: k). (:*:) f g x -> a
oracle' (f x
x :*: g x
y) =
        let z1 :: a
z1 = forall (algo :: k) (f :: k -> Type) a (x :: k).
RandomOracle' algo f a =>
f x -> a
forall {k} {k} (algo :: k) (f :: k -> Type) a (x :: k).
RandomOracle' algo f a =>
f x -> a
oracle' @algo f x
x :: a
            z2 :: a
z2 = forall (algo :: k) (f :: k -> Type) a (x :: k).
RandomOracle' algo f a =>
f x -> a
forall {k} {k} (algo :: k) (f :: k -> Type) a (x :: k).
RandomOracle' algo f a =>
f x -> a
oracle' @algo g x
y :: a
        in forall (algo :: k) a. HashAlgorithm algo a => [a] -> a
forall {k} (algo :: k) a. HashAlgorithm algo a => [a] -> a
hash @algo [a
z1, a
z2]

instance RandomOracle algo c a => RandomOracle' algo (Rec0 c) a where
    oracle' :: forall (x :: k). Rec0 c x -> a
oracle' (K1 c
x) = forall (algo :: k) x a. RandomOracle algo x a => x -> a
forall {k} (algo :: k) x a. RandomOracle algo x a => x -> a
oracle @algo c
x

-- | Handling constructors with no fields.
instance {-# OVERLAPPING #-}
    Ring a => RandomOracle' algo (M1 C ('MetaCons conName fixity selectors) U1) a where
    oracle' :: forall (x :: k).
M1 C ('MetaCons conName fixity selectors) U1 x -> a
oracle' M1 C ('MetaCons conName fixity selectors) U1 x
_ = a
forall a. AdditiveMonoid a => a
zero

instance RandomOracle' algo f a => RandomOracle' algo (M1 c m f) a where
    oracle' :: forall (x :: k). M1 c m f x -> a
oracle' (M1 f x
x) = forall (algo :: k) (f :: k -> Type) a (x :: k).
RandomOracle' algo f a =>
f x -> a
forall {k} {k} (algo :: k) (f :: k -> Type) a (x :: k).
RandomOracle' algo f a =>
f x -> a
oracle' @algo f x
x