fix ArithmeticExtensionGate and MulExtensionGate

src/Algebra/GoldilocksExt.hs

@@ -21,33 +21,36 @@ import Misc.Pretty
 
 --------------------------------------------------------------------------------
 
-type FExt = GoldilocksExt
+type GoldilocksExt = Ext Goldilocks
+type FExt          = GoldilocksExt
 
-data GoldilocksExt 
-  = MkExt !Goldilocks !Goldilocks
+--------------------------------------------------------------------------------
+
+data Ext a 
+  = MkExt !a !a
   deriving Eq
 
-fromBase :: Goldilocks -> GoldilocksExt
+fromBase :: Num a => a -> Ext a
 fromBase x = MkExt x 0
 
-instance Show GoldilocksExt where
+instance Show a => Show (Ext a) where
   show (MkExt real imag) = "(" ++ show real ++ " + X*" ++ show imag ++ ")"
 
-instance Pretty GoldilocksExt where 
+instance (Eq a, Num a, Pretty a) => Pretty (Ext a) where 
   prettyPrec d (MkExt real imag) 
     | imag == 0   = prettyPrec 0 real
     | otherwise   = showParen (d > 5) 
                   $ prettyPrec 0 real . showString " + X*" . prettyPrec 0 imag
 
-instance ToJSON GoldilocksExt where
+instance ToJSON a => ToJSON (Ext a) where
   toJSON (MkExt a b) = toJSON (a,b)
 
-instance FromJSON GoldilocksExt where
+instance FromJSON a => FromJSON (Ext a) where
   parseJSON o = (\(a,b) -> MkExt a b) <$> parseJSON o
 
 --------------------------------------------------------------------------------
 
-instance Num FExt where
+instance Num a => Num (Ext a) where
   fromInteger k = MkExt (fromInteger k) 0
   negate (MkExt r i) = MkExt (negate r) (negate i)
   (+) (MkExt r1 i1) (MkExt r2 i2) = MkExt (r1+r2) (i1+i2)
@@ -56,34 +59,34 @@ instance Num FExt where
   signum = const 1
   abs    = id
 
-instance Fractional FExt where
+instance Fractional a => Fractional (Ext a) where
   fromRational q = fromInteger (numerator q) / fromInteger (denominator q)
   recip = invExt
   (/)   = divExt
 
 --------------------------------------------------------------------------------
 
-scaleExt :: F -> FExt -> FExt
+scaleExt :: Num a => a -> Ext a -> Ext a
 scaleExt s (MkExt a b) = MkExt (s*a) (s*b)
 
-sqrExt :: FExt -> FExt
+sqrExt :: Num a => Ext a -> Ext a
 sqrExt x = x*x
 
-invExt :: FExt -> FExt
+invExt :: Fractional a => Ext a -> Ext a
 invExt (MkExt a b) = MkExt c d where
-  denom = inv (a*a - 7*b*b)
+  denom = recip (a*a - 7*b*b)
   c =        a * denom
   d = negate b * denom
 
-divExt :: FExt -> FExt -> FExt
+divExt :: Fractional a => Ext a -> Ext a -> Ext a
 divExt u v = u * invExt v
 
 --------------------------------------------------------------------------------
 
-powExt_ :: GoldilocksExt -> Int -> GoldilocksExt
+powExt_ :: Fractional a => Ext a -> Int -> Ext a
 powExt_ x e = powExt x (fromIntegral e)
 
-powExt :: GoldilocksExt -> Integer -> GoldilocksExt
+powExt :: Fractional a => Ext a -> Integer -> Ext a
 powExt x e 
   | e == 0    = 1
   | e <  0    = powExt (invExt x) (negate e)

src/Gate/Constraints.hs

@@ -13,6 +13,7 @@ module Gate.Constraints where
 import Prelude    hiding ( (^) )
 import Data.Array hiding (range)
 import Data.Char
+import Control.Monad
 
 import Algebra.Goldilocks
 import Algebra.GoldilocksExt
@@ -44,7 +45,12 @@ gateComputation gate =
 
     -- same but consecutive witness variables make up an extension field element
     ArithmeticExtensionGate num_ops 
-      -> commitList [ wireExt (j+6) - cnst 0 * wireExt j * wireExt (j+2) - cnst 1 * wireExt (j+4) | i<-range num_ops, let j = 8*i ]
+      -> forM_ (range num_ops) $ \i -> do
+           let j  = 8*i
+           let c0 = fromBase (cnst 0)
+           let c1 = fromBase (cnst 1)
+           let MkExt u v = wireExt (j+6) - c0 * wireExt j * wireExt (j+2) - c1 * wireExt (j+4)
+           commitList [ u , v ]
 
     -- `sum b^i * limbs[i] - out = 0`, and `0 <= limb[i] < B` is enforced
     BaseSumGate num_limbs base
@@ -71,7 +77,10 @@ gateComputation gate =
 
     -- `z[i] - c0*x[i]*y[i] = 0`, and two witness cells make up an extension field element
     MulExtensionGate num_ops 
-      -> commitList [ wireExt (j+4) - cnst 0 * wireExt j * wireExt (j+2) | i<-range num_ops, let j = 6*i ]
+      -> forM_ (range num_ops) $ \i -> do
+           let j  = 6*i
+           let MkExt u v = wireExt (j+4) - fromBase (cnst 0) * wireExt j * wireExt (j+2) 
+           commitList [ u , v ]
 
     NoopGate -> return ()
 
@@ -124,7 +133,8 @@ exponentiationGateConstraints num_power_bits =
 
 --------------------------------------------------------------------------------
 
-testExpoGate = runComputation testEvaluationVarsExt (gateComputation (ExponentiationGate 13))
+testArtihExtGate = runComputation testEvaluationVarsExt (gateComputation (ArithmeticExtensionGate 10))
+testMulExtGate   = runComputation testEvaluationVarsExt (gateComputation (MulExtensionGate        13))
+testExpoGate     = runComputation testEvaluationVarsExt (gateComputation (ExponentiationGate      13))
 
 --------------------------------------------------------------------------------
-

src/Gate/Vars.hs

@@ -7,6 +7,7 @@ module Gate.Vars where
 
 import Text.Show
 
+import Algebra.GoldilocksExt
 import Algebra.Expr
 import Misc.Pretty
 
@@ -49,7 +50,8 @@ wire i = VarE $ ProofVar $ WireV  i       -- witness variable
 cnst i = VarE $ ProofVar $ ConstV i       -- constant variable
 hash i = VarE $ ProofVar $ PIV    i       -- public input hash component
 
-wireExt :: Int -> Expr (Var PlonkyVar)
-wireExt i = wire i + ImgE (wire (i+1))    -- use two consecutive variables as an extension field element
+-- use two consecutive variables as a _simulated_ extension field element
+wireExt :: Int -> Ext (Expr (Var PlonkyVar))
+wireExt i = MkExt (wire i) (wire (i+1))     
 
 --------------------------------------------------------------------------------