{-# LANGUAGE OverloadedLists      #-}
{-# LANGUAGE UndecidableInstances #-}

module ZkFold.Base.Protocol.Plonkup.PlonkConstraint where

import           Control.Monad                                       (guard, replicateM, return)
import           Data.Binary                                         (Binary)
import           Data.Containers.ListUtils                           (nubOrd)
import           Data.Eq                                             (Eq (..))
import           Data.Function                                       (($), (.))
import           Data.Functor                                        ((<$>))
import           Data.Functor.Rep                                    (Rep)
import           Data.List                                           (find, head, map, permutations, sort, (!!), (++))
import           Data.Map                                            (Map)
import qualified Data.Map                                            as Map
import           Data.Maybe                                          (Maybe (..), fromMaybe, mapMaybe)
import           Data.Ord                                            (Ord)
import           GHC.IsList                                          (IsList (..))
import           Numeric.Natural                                     (Natural)
import           Test.QuickCheck                                     (Arbitrary (..))
import           Text.Show                                           (Show)

import           ZkFold.Base.Algebra.Basic.Class
import           ZkFold.Base.Algebra.Polynomials.Multivariate        (Poly, evalMonomial, evalPolynomial, polynomial,
                                                                      var, variables)
import           ZkFold.Base.Data.ByteString                         (toByteString)
import           ZkFold.Prelude                                      (length, take)
import           ZkFold.Symbolic.Compiler.ArithmeticCircuit.Internal

data PlonkConstraint i a = PlonkConstraint
    { forall (i :: Type -> Type) a. PlonkConstraint i a -> a
qm :: a
    , forall (i :: Type -> Type) a. PlonkConstraint i a -> a
ql :: a
    , forall (i :: Type -> Type) a. PlonkConstraint i a -> a
qr :: a
    , forall (i :: Type -> Type) a. PlonkConstraint i a -> a
qo :: a
    , forall (i :: Type -> Type) a. PlonkConstraint i a -> a
qc :: a
    , forall (i :: Type -> Type) a. PlonkConstraint i a -> Var a i
x1 :: Var a i
    , forall (i :: Type -> Type) a. PlonkConstraint i a -> Var a i
x2 :: Var a i
    , forall (i :: Type -> Type) a. PlonkConstraint i a -> Var a i
x3 :: Var a i
    }

deriving instance (Show a, Show (Rep i)) => Show (PlonkConstraint i a)
deriving instance (Eq a, Eq (Rep i)) => Eq (PlonkConstraint i a)

instance (Ord a, Arbitrary a, Binary a, Ord (Rep i)) => Arbitrary (PlonkConstraint i a) where
    arbitrary :: Gen (PlonkConstraint i a)
arbitrary = do
        a
qm <- Gen a
forall a. Arbitrary a => Gen a
arbitrary
        a
ql <- Gen a
forall a. Arbitrary a => Gen a
arbitrary
        a
qr <- Gen a
forall a. Arbitrary a => Gen a
arbitrary
        a
qo <- Gen a
forall a. Arbitrary a => Gen a
arbitrary
        a
qc <- Gen a
forall a. Arbitrary a => Gen a
arbitrary
        let arbitraryNewVar :: Gen (Var a i)
arbitraryNewVar = SysVar i -> Var a i
forall a (i :: Type -> Type). SysVar i -> Var a i
SysVar (SysVar i -> Var a i) -> (a -> SysVar i) -> a -> Var a i
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> SysVar i
forall (i :: Type -> Type). ByteString -> SysVar i
NewVar (ByteString -> SysVar i) -> (a -> ByteString) -> a -> SysVar i
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Binary a => a -> ByteString
toByteString @a (a -> Var a i) -> Gen a -> Gen (Var a i)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen a
forall a. Arbitrary a => Gen a
arbitrary
        [Var a i]
xs <- [Var a i] -> [Var a i]
forall a. Ord a => [a] -> [a]
sort ([Var a i] -> [Var a i]) -> Gen [Var a i] -> Gen [Var a i]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Gen (Var a i) -> Gen [Var a i]
forall (m :: Type -> Type) a. Applicative m => Int -> m a -> m [a]
replicateM Int
3 Gen (Var a i)
arbitraryNewVar
        let x1 :: Var a i
x1 = [Var a i] -> Var a i
forall a. HasCallStack => [a] -> a
head [Var a i]
xs; x2 :: Var a i
x2 = [Var a i]
xs [Var a i] -> Int -> Var a i
forall a. HasCallStack => [a] -> Int -> a
!! Int
1; x3 :: Var a i
x3 = [Var a i]
xs [Var a i] -> Int -> Var a i
forall a. HasCallStack => [a] -> Int -> a
!! Int
2
        PlonkConstraint i a -> Gen (PlonkConstraint i a)
forall a. a -> Gen a
forall (m :: Type -> Type) a. Monad m => a -> m a
return (PlonkConstraint i a -> Gen (PlonkConstraint i a))
-> PlonkConstraint i a -> Gen (PlonkConstraint i a)
forall a b. (a -> b) -> a -> b
$ a
-> a
-> a
-> a
-> a
-> Var a i
-> Var a i
-> Var a i
-> PlonkConstraint i a
forall (i :: Type -> Type) a.
a
-> a
-> a
-> a
-> a
-> Var a i
-> Var a i
-> Var a i
-> PlonkConstraint i a
PlonkConstraint a
qm a
ql a
qr a
qo a
qc Var a i
x1 Var a i
x2 Var a i
x3

toPlonkConstraint :: forall a i . (Ord a, FiniteField a, Ord (Rep i)) => Poly a (Var a i) Natural -> PlonkConstraint i a
toPlonkConstraint :: forall a (i :: Type -> Type).
(Ord a, FiniteField a, Ord (Rep i)) =>
Poly a (Var a i) Natural -> PlonkConstraint i a
toPlonkConstraint Poly a (Var a i) Natural
p =
    let xs :: [Maybe (Var a i)]
xs    = Var a i -> Maybe (Var a i)
forall a. a -> Maybe a
Just (Var a i -> Maybe (Var a i)) -> [Var a i] -> [Maybe (Var a i)]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Set (Var a i) -> [Item (Set (Var a i))]
forall l. IsList l => l -> [Item l]
toList (Poly a (Var a i) Natural -> Set (Var a i)
forall c v. Ord v => Poly c v Natural -> Set v
variables Poly a (Var a i) Natural
p)
        perms :: [[Maybe (Var a i)]]
perms = [[Maybe (Var a i)]] -> [[Maybe (Var a i)]]
forall a. Ord a => [a] -> [a]
nubOrd ([[Maybe (Var a i)]] -> [[Maybe (Var a i)]])
-> [[Maybe (Var a i)]] -> [[Maybe (Var a i)]]
forall a b. (a -> b) -> a -> b
$ ([Maybe (Var a i)] -> [Maybe (Var a i)])
-> [[Maybe (Var a i)]] -> [[Maybe (Var a i)]]
forall a b. (a -> b) -> [a] -> [b]
map (Natural -> [Maybe (Var a i)] -> [Maybe (Var a i)]
forall a. HasCallStack => Natural -> [a] -> [a]
take Natural
3) ([[Maybe (Var a i)]] -> [[Maybe (Var a i)]])
-> [[Maybe (Var a i)]] -> [[Maybe (Var a i)]]
forall a b. (a -> b) -> a -> b
$ [Maybe (Var a i)] -> [[Maybe (Var a i)]]
forall a. [a] -> [[a]]
permutations ([Maybe (Var a i)] -> [[Maybe (Var a i)]])
-> [Maybe (Var a i)] -> [[Maybe (Var a i)]]
forall a b. (a -> b) -> a -> b
$ case [Maybe (Var a i)] -> Natural
forall (t :: Type -> Type) a. Foldable t => t a -> Natural
length [Maybe (Var a i)]
xs of
            Natural
0 -> [Maybe (Var a i)
Item [Maybe (Var a i)]
forall a. Maybe a
Nothing, Maybe (Var a i)
Item [Maybe (Var a i)]
forall a. Maybe a
Nothing, Maybe (Var a i)
Item [Maybe (Var a i)]
forall a. Maybe a
Nothing]
            Natural
1 -> [Maybe (Var a i)
Item [Maybe (Var a i)]
forall a. Maybe a
Nothing, Maybe (Var a i)
Item [Maybe (Var a i)]
forall a. Maybe a
Nothing, [Maybe (Var a i)] -> Maybe (Var a i)
forall a. HasCallStack => [a] -> a
head [Maybe (Var a i)]
xs, [Maybe (Var a i)] -> Maybe (Var a i)
forall a. HasCallStack => [a] -> a
head [Maybe (Var a i)]
xs]
            Natural
2 -> [Maybe (Var a i)
Item [Maybe (Var a i)]
forall a. Maybe a
Nothing] [Maybe (Var a i)] -> [Maybe (Var a i)] -> [Maybe (Var a i)]
forall a. [a] -> [a] -> [a]
++ [Maybe (Var a i)]
xs [Maybe (Var a i)] -> [Maybe (Var a i)] -> [Maybe (Var a i)]
forall a. [a] -> [a] -> [a]
++ [Maybe (Var a i)]
xs
            Natural
_ -> [Maybe (Var a i)]
xs [Maybe (Var a i)] -> [Maybe (Var a i)] -> [Maybe (Var a i)]
forall a. [a] -> [a] -> [a]
++ [Maybe (Var a i)]
xs

        getCoef :: Map (Maybe (Var a i)) Natural -> a
        getCoef :: Map (Maybe (Var a i)) Natural -> a
getCoef Map (Maybe (Var a i)) Natural
m = case ((a, Map (Var a i) Natural) -> Bool)
-> [(a, Map (Var a i) Natural)] -> Maybe (a, Map (Var a i) Natural)
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Maybe a
find (\(a
_, Map (Var a i) Natural
as) -> Map (Maybe (Var a i)) Natural
m Map (Maybe (Var a i)) Natural
-> Map (Maybe (Var a i)) Natural -> Bool
forall a. Eq a => a -> a -> Bool
== (Var a i -> Maybe (Var a i))
-> Map (Var a i) Natural -> Map (Maybe (Var a i)) Natural
forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
Map.mapKeys Var a i -> Maybe (Var a i)
forall a. a -> Maybe a
Just Map (Var a i) Natural
as) (Poly a (Var a i) Natural -> [Item (Poly a (Var a i) Natural)]
forall l. IsList l => l -> [Item l]
toList Poly a (Var a i) Natural
p) of
            Just (a
c, Map (Var a i) Natural
_) -> a
c
            Maybe (a, Map (Var a i) Natural)
_           -> a
forall a. AdditiveMonoid a => a
zero

        getCoefs :: [Maybe (Var a i)] -> Maybe (PlonkConstraint i a)
        getCoefs :: [Maybe (Var a i)] -> Maybe (PlonkConstraint i a)
getCoefs [Item [Maybe (Var a i)]
a, Item [Maybe (Var a i)]
b, Item [Maybe (Var a i)]
c] = do
            let xa :: [(Maybe (Var a i), Natural)]
xa = [(Maybe (Var a i)
Item [Maybe (Var a i)]
a, Natural
1)]
                xb :: [(Maybe (Var a i), Natural)]
xb = [(Maybe (Var a i)
Item [Maybe (Var a i)]
b, Natural
1)]
                xc :: [(Maybe (Var a i), Natural)]
xc = [(Maybe (Var a i)
Item [Maybe (Var a i)]
c, Natural
1)]
                xaxb :: [(Maybe (Var a i), Natural)]
xaxb = [(Maybe (Var a i), Natural)]
xa [(Maybe (Var a i), Natural)]
-> [(Maybe (Var a i), Natural)] -> [(Maybe (Var a i), Natural)]
forall a. [a] -> [a] -> [a]
++ [(Maybe (Var a i), Natural)]
xb

                qm :: a
qm = Map (Maybe (Var a i)) Natural -> a
getCoef (Map (Maybe (Var a i)) Natural -> a)
-> Map (Maybe (Var a i)) Natural -> a
forall a b. (a -> b) -> a -> b
$ (Natural -> Natural -> Natural)
-> [(Maybe (Var a i), Natural)] -> Map (Maybe (Var a i)) Natural
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith Natural -> Natural -> Natural
forall a. AdditiveSemigroup a => a -> a -> a
(+) [(Maybe (Var a i), Natural)]
xaxb
                ql :: a
ql = Map (Maybe (Var a i)) Natural -> a
getCoef (Map (Maybe (Var a i)) Natural -> a)
-> Map (Maybe (Var a i)) Natural -> a
forall a b. (a -> b) -> a -> b
$ [Item (Map (Maybe (Var a i)) Natural)]
-> Map (Maybe (Var a i)) Natural
forall l. IsList l => [Item l] -> l
fromList [(Maybe (Var a i), Natural)]
[Item (Map (Maybe (Var a i)) Natural)]
xa
                qr :: a
qr = Map (Maybe (Var a i)) Natural -> a
getCoef (Map (Maybe (Var a i)) Natural -> a)
-> Map (Maybe (Var a i)) Natural -> a
forall a b. (a -> b) -> a -> b
$ [Item (Map (Maybe (Var a i)) Natural)]
-> Map (Maybe (Var a i)) Natural
forall l. IsList l => [Item l] -> l
fromList [(Maybe (Var a i), Natural)]
[Item (Map (Maybe (Var a i)) Natural)]
xb
                qo :: a
qo = Map (Maybe (Var a i)) Natural -> a
getCoef (Map (Maybe (Var a i)) Natural -> a)
-> Map (Maybe (Var a i)) Natural -> a
forall a b. (a -> b) -> a -> b
$ [Item (Map (Maybe (Var a i)) Natural)]
-> Map (Maybe (Var a i)) Natural
forall l. IsList l => [Item l] -> l
fromList [(Maybe (Var a i), Natural)]
[Item (Map (Maybe (Var a i)) Natural)]
xc
                qc :: a
qc = Map (Maybe (Var a i)) Natural -> a
getCoef Map (Maybe (Var a i)) Natural
forall k a. Map k a
Map.empty
            Bool -> Maybe ()
forall (f :: Type -> Type). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ ((Var a i -> Poly a (Maybe (Var a i)) Natural)
 -> Mono (Var a i) Natural -> Poly a (Maybe (Var a i)) Natural)
-> (Var a i -> Poly a (Maybe (Var a i)) Natural)
-> Poly a (Var a i) Natural
-> Poly a (Maybe (Var a i)) Natural
forall c i j b.
(AdditiveMonoid b, Scale c b) =>
((i -> b) -> Mono i j -> b) -> (i -> b) -> Poly c i j -> b
evalPolynomial (Var a i -> Poly a (Maybe (Var a i)) Natural)
-> Mono (Var a i) Natural -> Poly a (Maybe (Var a i)) Natural
forall i j b.
(MultiplicativeMonoid b, Exponent b j) =>
(i -> b) -> Mono i j -> b
evalMonomial (Maybe (Var a i) -> Poly a (Maybe (Var a i)) Natural
forall c i j. Polynomial c i j => i -> Poly c i j
var (Maybe (Var a i) -> Poly a (Maybe (Var a i)) Natural)
-> (Var a i -> Maybe (Var a i))
-> Var a i
-> Poly a (Maybe (Var a i)) Natural
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Var a i -> Maybe (Var a i)
forall a. a -> Maybe a
Just) Poly a (Var a i) Natural
p Poly a (Maybe (Var a i)) Natural
-> Poly a (Maybe (Var a i)) Natural
-> Poly a (Maybe (Var a i)) Natural
forall a. AdditiveGroup a => a -> a -> a
- [(a, Mono (Maybe (Var a i)) Natural)]
-> Poly a (Maybe (Var a i)) Natural
forall c i j. Polynomial c i j => [(c, Mono i j)] -> Poly c i j
polynomial [(a
qm, [Item (Mono (Maybe (Var a i)) Natural)]
-> Mono (Maybe (Var a i)) Natural
forall l. IsList l => [Item l] -> l
fromList [(Maybe (Var a i), Natural)]
[Item (Mono (Maybe (Var a i)) Natural)]
xaxb), (a
ql, [Item (Mono (Maybe (Var a i)) Natural)]
-> Mono (Maybe (Var a i)) Natural
forall l. IsList l => [Item l] -> l
fromList [(Maybe (Var a i), Natural)]
[Item (Mono (Maybe (Var a i)) Natural)]
xa), (a
qr, [Item (Mono (Maybe (Var a i)) Natural)]
-> Mono (Maybe (Var a i)) Natural
forall l. IsList l => [Item l] -> l
fromList [(Maybe (Var a i), Natural)]
[Item (Mono (Maybe (Var a i)) Natural)]
xb), (a
qo, [Item (Mono (Maybe (Var a i)) Natural)]
-> Mono (Maybe (Var a i)) Natural
forall l. IsList l => [Item l] -> l
fromList [(Maybe (Var a i), Natural)]
[Item (Mono (Maybe (Var a i)) Natural)]
xc), (a
qc, Mono (Maybe (Var a i)) Natural
forall a. MultiplicativeMonoid a => a
one)] Poly a (Maybe (Var a i)) Natural
-> Poly a (Maybe (Var a i)) Natural -> Bool
forall a. Eq a => a -> a -> Bool
== Poly a (Maybe (Var a i)) Natural
forall a. AdditiveMonoid a => a
zero
            let va :: Var a i
va = Var a i -> Maybe (Var a i) -> Var a i
forall a. a -> Maybe a -> a
fromMaybe (a -> Var a i
forall a (i :: Type -> Type). a -> Var a i
ConstVar a
forall a. MultiplicativeMonoid a => a
one) Maybe (Var a i)
Item [Maybe (Var a i)]
a
                vb :: Var a i
vb = Var a i -> Maybe (Var a i) -> Var a i
forall a. a -> Maybe a -> a
fromMaybe (a -> Var a i
forall a (i :: Type -> Type). a -> Var a i
ConstVar a
forall a. MultiplicativeMonoid a => a
one) Maybe (Var a i)
Item [Maybe (Var a i)]
b
                vc :: Var a i
vc = Var a i -> Maybe (Var a i) -> Var a i
forall a. a -> Maybe a -> a
fromMaybe (a -> Var a i
forall a (i :: Type -> Type). a -> Var a i
ConstVar a
forall a. MultiplicativeMonoid a => a
one) Maybe (Var a i)
Item [Maybe (Var a i)]
c
            PlonkConstraint i a -> Maybe (PlonkConstraint i a)
forall a. a -> Maybe a
forall (m :: Type -> Type) a. Monad m => a -> m a
return (PlonkConstraint i a -> Maybe (PlonkConstraint i a))
-> PlonkConstraint i a -> Maybe (PlonkConstraint i a)
forall a b. (a -> b) -> a -> b
$ a
-> a
-> a
-> a
-> a
-> Var a i
-> Var a i
-> Var a i
-> PlonkConstraint i a
forall (i :: Type -> Type) a.
a
-> a
-> a
-> a
-> a
-> Var a i
-> Var a i
-> Var a i
-> PlonkConstraint i a
PlonkConstraint a
qm a
ql a
qr a
qo a
qc Var a i
va Var a i
vb Var a i
vc
        getCoefs [Maybe (Var a i)]
_ = Maybe (PlonkConstraint i a)
forall a. Maybe a
Nothing

    in case ([Maybe (Var a i)] -> Maybe (PlonkConstraint i a))
-> [[Maybe (Var a i)]] -> [PlonkConstraint i a]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe [Maybe (Var a i)] -> Maybe (PlonkConstraint i a)
getCoefs [[Maybe (Var a i)]]
perms of
        [] -> Poly a (Var a i) Natural -> PlonkConstraint i a
forall a (i :: Type -> Type).
(Ord a, FiniteField a, Ord (Rep i)) =>
Poly a (Var a i) Natural -> PlonkConstraint i a
toPlonkConstraint Poly a (Var a i) Natural
forall a. AdditiveMonoid a => a
zero
        [PlonkConstraint i a]
_  -> [PlonkConstraint i a] -> PlonkConstraint i a
forall a. HasCallStack => [a] -> a
head ([PlonkConstraint i a] -> PlonkConstraint i a)
-> [PlonkConstraint i a] -> PlonkConstraint i a
forall a b. (a -> b) -> a -> b
$ ([Maybe (Var a i)] -> Maybe (PlonkConstraint i a))
-> [[Maybe (Var a i)]] -> [PlonkConstraint i a]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe [Maybe (Var a i)] -> Maybe (PlonkConstraint i a)
getCoefs [[Maybe (Var a i)]]
perms

fromPlonkConstraint :: (Ord a, Field a, Ord (Rep i)) => PlonkConstraint i a -> Poly a (Var a i) Natural
fromPlonkConstraint :: forall a (i :: Type -> Type).
(Ord a, Field a, Ord (Rep i)) =>
PlonkConstraint i a -> Poly a (Var a i) Natural
fromPlonkConstraint (PlonkConstraint a
qm a
ql a
qr a
qo a
qc Var a i
a Var a i
b Var a i
c) =
    let xa :: Poly a (Var a i) Natural
xa = Var a i -> Poly a (Var a i) Natural
forall c i j. Polynomial c i j => i -> Poly c i j
var Var a i
a
        xb :: Poly a (Var a i) Natural
xb = Var a i -> Poly a (Var a i) Natural
forall c i j. Polynomial c i j => i -> Poly c i j
var Var a i
b
        xc :: Poly a (Var a i) Natural
xc = Var a i -> Poly a (Var a i) Natural
forall c i j. Polynomial c i j => i -> Poly c i j
var Var a i
c
        xaxb :: Poly a (Var a i) Natural
xaxb = Poly a (Var a i) Natural
xa Poly a (Var a i) Natural
-> Poly a (Var a i) Natural -> Poly a (Var a i) Natural
forall a. MultiplicativeSemigroup a => a -> a -> a
* Poly a (Var a i) Natural
xb
    in
              a -> Poly a (Var a i) Natural -> Poly a (Var a i) Natural
forall b a. Scale b a => b -> a -> a
scale a
qm Poly a (Var a i) Natural
xaxb
            Poly a (Var a i) Natural
-> Poly a (Var a i) Natural -> Poly a (Var a i) Natural
forall a. AdditiveSemigroup a => a -> a -> a
+ a -> Poly a (Var a i) Natural -> Poly a (Var a i) Natural
forall b a. Scale b a => b -> a -> a
scale a
ql Poly a (Var a i) Natural
xa
            Poly a (Var a i) Natural
-> Poly a (Var a i) Natural -> Poly a (Var a i) Natural
forall a. AdditiveSemigroup a => a -> a -> a
+ a -> Poly a (Var a i) Natural -> Poly a (Var a i) Natural
forall b a. Scale b a => b -> a -> a
scale a
qr Poly a (Var a i) Natural
xb
            Poly a (Var a i) Natural
-> Poly a (Var a i) Natural -> Poly a (Var a i) Natural
forall a. AdditiveSemigroup a => a -> a -> a
+ a -> Poly a (Var a i) Natural -> Poly a (Var a i) Natural
forall b a. Scale b a => b -> a -> a
scale a
qo Poly a (Var a i) Natural
xc
            Poly a (Var a i) Natural
-> Poly a (Var a i) Natural -> Poly a (Var a i) Natural
forall a. AdditiveSemigroup a => a -> a -> a
+ a -> Poly a (Var a i) Natural
forall a b. FromConstant a b => a -> b
fromConstant a
qc