implement lookup tables

.gitignore

@@ -1,5 +1,6 @@
 .DS_Store
+.ghc.*
 _bck*
 json/
-.ghc.*
-*.json
+tmp/
+*.json

README.md

@@ -8,7 +8,7 @@ This is a (WIP) implementation of a Plonky2 verifier written in Haskell.
 system developed by Polygon Zero, optimized for recursive proofs.
 
 The goal here is to provide an executable specification (along a with less precise,
-but [still detailed](commentary/Overview.md) human language description) of 
+but still detailed [human language description](commentary/README.md)) of 
 the Plonky2 verification algorithm. 
 
 Another goal is to be a basis for further tooling (for example:
@@ -27,7 +27,7 @@ try to focus on simplicity.
 - [ ] Recursive circuit subtle details (like [this](https://github.com/0xPolygonZero/plonky2/blob/356aefb6863ac881fb71f9bf851582c915428458/plonky2/src/fri/challenges.rs#L55-L64]))
 - [x] Constraints check
 - [ ] FRI check
-- [ ] Support lookup tables
+- [x] Support lookup tables
 - [x] Documenting Plonky2 internals and the verifier algorithm (WIP)
 - [ ] Cabalize
 
@@ -39,8 +39,8 @@ Supported gates:
 - [x] CosetInterpolationGate
 - [x] ConstantGate
 - [x] ExponentiationGate
-- [ ] LookupGate
-- [ ] LookupTableGate
+- [x] LookupGate
+- [x] LookupTableGate
 - [x] MulExtensionGate
 - [x] NoopGate
 - [x] PublicInputGate
@@ -53,6 +53,7 @@ Supported gates:
 Optional features:
 
 - [ ] Supporting different hash functions
+- [ ] Handle non-standard configurations
 - [ ] Field extensions with degree higher than 2
 - [ ] Being parametric over the field choice
 

src/Algebra/Poly.hs

@@ -0,0 +1,25 @@
+
+-- | Polynomials
+
+module Algebra.Poly where
+
+--------------------------------------------------------------------------------
+
+import Algebra.Goldilocks
+import Algebra.GoldilocksExt
+
+--------------------------------------------------------------------------------
+
+-- | Evaluate the Lagrange polynomial @L_0(x)@ of a subgroup of size @N@
+evalLagrange0 :: Int -> FExt -> FExt
+evalLagrange0 nn zeta 
+  | zeta == 1  = 1
+  | otherwise  = (powExt_ zeta nn - 1) / (fromIntegral nn * (zeta - 1))
+
+--------------------------------------------------------------------------------
+
+-- | Evaluate the zero polynomial @Z_H(x)@ on a subgroup of size @N@
+evalZeroPoly :: Int -> FExt -> FExt
+evalZeroPoly nn zeta = powExt_ zeta nn - 1 
+
+--------------------------------------------------------------------------------

src/Gate/Constraints.hs

@@ -130,6 +130,7 @@ exponentiationGateConstraints num_power_bits =
 --------------------------------------------------------------------------------
 -- * Debugging
 
+{-
 testCompute :: Compute () -> [FExt]
 testCompute = runComputation testEvaluationVarsExt 
 
@@ -146,5 +147,6 @@ testPoseidonGate    = testCompute $ gateComputation (PoseidonGate    12)
 testPoseidonMdsGate = testCompute $ gateComputation (PoseidonMdsGate 12)
 testReducingGate    = testCompute $ gateComputation (ReducingGate          13)
 testReducingExtGate = testCompute $ gateComputation (ReducingExtensionGate 13)
+-}
 
 --------------------------------------------------------------------------------

src/Misc/Aux.hs

@@ -27,6 +27,14 @@ exp2 (Log2 k) = shiftL 1 k
 
 --------------------------------------------------------------------------------
 
+divCeil :: Int -> Int -> Int
+divCeil n k = div (n+k-1) k
+
+divFloor :: Int -> Int -> Int
+divFloor = div
+
+--------------------------------------------------------------------------------
+
 range :: Int -> [Int]
 range k = [0..k-1]
 
@@ -70,6 +78,9 @@ select1 zs = go zs where
   go [x]     = [(x,[])]
   go (x:xs)  = (x,xs) : map (\(y,ys) -> (y,x:ys)) (go xs)
 
+remove1 :: [a] -> [[a]]
+remove1 = map snd . select1
+
 --------------------------------------------------------------------------------
 
 listToArray :: [a] -> Array Int a

src/Plonk/Lookups.hs

@@ -0,0 +1,134 @@
+
+-- | Evaluation the constraints of the lookup argument
+
+{-# LANGUAGE StrictData, RecordWildCards #-}
+module Plonk.Lookups where
+
+--------------------------------------------------------------------------------
+
+import Data.Array
+import Data.List ( foldl' )
+
+import Algebra.Goldilocks
+import Algebra.GoldilocksExt
+import Algebra.Poly
+
+import Challenge.Verifier
+import Gate.Selector
+
+import Types
+import Misc.Aux
+
+-- import Debug.Trace
+-- debug msg x y = trace (msg ++ ": " ++ show x) y
+
+--------------------------------------------------------------------------------
+
+data LookupSelector
+  = TransSre          -- ^ @TransSre@ is for Sum and RE transition constraints.
+  | TransLdc          -- ^ @TransLdc@ is for LDC transition constraints.
+  | InitSre           -- ^ @InitSre@  is for the initial constraint of Sum and Re.
+  | LastLdc           -- ^ @LastLdc@  is for the final (S)LDC constraint.
+  | StartEnd Int      -- ^ @StartEnd@ indicates where lookup end selectors begin.     
+  deriving (Eq,Ord,Show)
+
+lookupSelectorIndex :: LookupSelector -> Int
+lookupSelectorIndex sel = case sel of
+  TransSre   -> 0
+  TransLdc   -> 1
+  InitSre    -> 2
+  LastLdc    -> 3
+  StartEnd k -> 4 + k
+
+--------------------------------------------------------------------------------
+
+evalLookupEquations :: CommonCircuitData -> [FExt] -> OpeningSet -> ProofChallenges -> [FExt]
+evalLookupEquations (MkCommonCircuitData{..}) lkpSels openings challenges = final where
+
+  MkCircuitConfig  {..} = circuit_config
+  MkOpeningSet     {..} = openings
+  MkProofChallenges{..} = challenges 
+
+  nluts          = length circuit_luts
+  num_challenges = config_num_challenges
+
+  selector :: LookupSelector -> FExt
+  selector sel = lkpSels !! (lookupSelectorIndex sel)
+
+  -- one for each challenge round
+  roundChunks :: [[(FExt,FExt)]]
+  roundChunks = partition circuit_num_lookup_polys $ zip opening_lookup_zs opening_lookup_zs_next
+  
+  final = concat $ safeZipWith roundWorker plonk_deltas roundChunks
+
+  num_lu_slots   = config_num_routed_wires `Prelude.div` 2     -- lookups per LookupGate
+  num_lut_slots  = config_num_routed_wires `Prelude.div` 3     -- entries per LookupTableGate
+  num_sldc_polys = circuit_num_lookup_polys - 1
+  lu_degree      = circuit_quotient_degree_factor - 1 
+  lut_degree     = num_lut_slots `divCeil` num_sldc_polys
+
+  -- computation for a single challenge round
+  roundWorker :: LookupDelta -> [(FExt,FExt)] -> [FExt]
+  roundWorker (MkLookupDelta{..}) columns = final where
+
+    (re_pair:sldc_pairs) = columns
+    (re  ,re_next  ) = re_pair
+    (sldc,sldc_next) = unzip sldc_pairs
+
+    lu_combos    = [ (inp + scaleExt lookup_A out) | [inp,out]      <- take num_lu_slots  (partition 2 opening_wires) ] :: [FExt]
+    lut_combos_A = [ (inp + scaleExt lookup_A out) | [inp,out,mult] <- take num_lut_slots (partition 3 opening_wires) ] :: [FExt]
+    lut_combos_B = [ (inp + scaleExt lookup_B out) | [inp,out,mult] <- take num_lut_slots (partition 3 opening_wires) ] :: [FExt]
+
+    mult i = opening_wires !! (3*i+2)
+    mults  = [ mult i | i<-[0..num_lut_slots-1] ]
+
+    chunks_lu_combo  = partition lu_degree  lu_combos
+    chunks_lut_combo = partition lut_degree lut_combos_A
+    chunks_mults     = partition lut_degree mults
+
+    final  = [ eq_last_sldc , eq_ini_sum , eq_ini_re ]
+          ++ eq_finals_re
+          ++ [ eq_re_trans ]
+          ++ eqs_sldc
+
+    eq_last_sldc = selector LastLdc * (last sldc)         -- SLDC sums to zero
+    eq_ini_sum   = selector InitSre * (head sldc)         -- SUM starts with zero
+    eq_ini_re    = selector InitSre * re                  -- RE starts with zero
+
+    -- RE ends are correct (for each LUT)
+    eq_finals_re = 
+      [ selector (StartEnd k) * (re - fromBase (evalFinalRE table)) 
+      | (k,table) <- zip [0..] circuit_luts 
+      ] 
+    evalFinalRE table = cur_eval where
+      lut         = fromLookupTable table
+      lut_nrows   = length lut `divCeil` num_lut_slots               -- number of rows in this table
+      padded_size = lut_nrows * num_lut_slots                        -- padded size of the table
+      lut_padded  = take padded_size $ lut ++ repeat (0,0)           -- NOTE: this padding causes a soundness bug in Plonky2
+      cur_eval    = foldl' (\acc x -> lookup_delta * acc + x) 0 
+                      [ (inp + lookup_B * out) | (inp,out) <- lut_padded ]
+
+    -- RE transition is correct
+    eq_re_trans = selector TransSre * (re - cur_sum) where
+      cur_sum = foldl' (\acc elt -> scaleExt lookup_delta acc + elt) re_next lut_combos_B 
+
+    prevThisPairs = pairs (last sldc_next : sldc)
+    eqs_sldc = concatMap evalTransSLDC $ zip prevThisPairs (zip3 chunks_lu_combo chunks_lut_combo chunks_mults)
+
+    -- LDC and SUM transitions are correct
+    evalTransSLDC ((prev,this),(lu_combos,lut_combos,mults)) = [ eq_sum_trans , eq_ldc_trans ] where
+
+      alpha = fromBase lookup_alpha :: FExt
+
+      lu_prod     = product [ alpha - combo | combo <- lu_combos  ] :: FExt
+      lut_prod    = product [ alpha - combo | combo <- lut_combos ] :: FExt
+      lu_prods_i  = [        product [ alpha - combo | combo <- one_less ] |       one_less  <-            remove1 lu_combos   ] :: [FExt]
+      lut_prods_i = [ mult * product [ alpha - combo | combo <- one_less ] | (mult,one_less) <- zip mults (remove1 lut_combos) ] :: [FExt]
+
+      lu_sum_prod  = sum lu_prods_i
+      lut_sum_prod = sum lut_prods_i
+
+      eq_ldc_trans = selector TransLdc * (  lu_prod * (this - prev) + lu_sum_prod  )
+      eq_sum_trans = selector TransSre * ( lut_prod * (this - prev) - lut_sum_prod )
+
+--------------------------------------------------------------------------------

src/Plonk/Vanishing.hs

@@ -14,7 +14,7 @@
 -- This is done as many times as there are rounds, with different alphas (chosen from the base field).
 --
 
-{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE StrictData, RecordWildCards #-}
 module Plonk.Vanishing where
 
 --------------------------------------------------------------------------------
@@ -24,6 +24,7 @@ import Data.List ( foldl' , zipWith4 )
 
 import Algebra.Goldilocks
 import Algebra.GoldilocksExt
+import Algebra.Poly
 
 import Challenge.Verifier
 
@@ -34,19 +35,13 @@ import Gate.Computation
 import Gate.Constraints
 import Gate.Selector
 
+import Plonk.Lookups
+
 import Types
 import Misc.Aux
 
 --------------------------------------------------------------------------------
 
--- | Evaluate the Lagrange polynomial @L_0(x)@ of a subgroup of size @N@
-evalLagrange0 :: Int -> FExt -> FExt
-evalLagrange0 nn zeta 
-  | zeta == 1  = 1
-  | otherwise  = (powExt_ zeta nn - 1) / (fromIntegral nn * (zeta - 1))
-
---------------------------------------------------------------------------------
-
 -- | We combine the same things with all the challenge round alphas... 
 -- (Plonky2 is full with very strange design decisions...)
 --
@@ -65,8 +60,8 @@ combineWithPowersOfAlpha alpha xs = foldl' f 0 (reverse xs) where
 evalAllPlonkConstraints :: CommonCircuitData -> ProofWithPublicInputs -> ProofChallenges -> [FExt]
 evalAllPlonkConstraints 
     common_data@(MkCommonCircuitData{..}) 
-    (MkProofWithPublicInputs{..}) 
-    (MkProofChallenges{..}) = finals 
+    prf_with_inp@(MkProofWithPublicInputs{..}) 
+    challenges@(MkProofChallenges{..}) = finals 
   where
 
     finals = concat
@@ -76,13 +71,15 @@ evalAllPlonkConstraints
       , gates
       ]
 
-    lookup_checks = []    -- TODO
+    lookup_checks = if null circuit_luts 
+      then [] 
+      else evalLookupEquations common_data lookupSelectors openings challenges
 
     MkProof{..}       = the_proof
     MkOpeningSet{..}  = openings
     
-    MkSelectorConfig{..}   = getSelectorConfig common_data
-    opening_gate_selectors = take numGateSelectors opening_constants
+    selcfg@(MkSelectorConfig{..}) = getSelectorConfig common_data
+    MkConstantColumns{..}         = splitConstantColumns selcfg opening_constants
 
     nn      = fri_nrows circuit_fri_params
     maxdeg  = circuit_quotient_degree_factor
@@ -91,7 +88,7 @@ evalAllPlonkConstraints
     -- gate constraints
     eval_vars  = toEvaluationVars common_data pi_hash openings
     gate_prgs  = map gateProgram circuit_gates 
-    sel_values = evalGateSelectors circuit_selectors_info opening_gate_selectors
+    sel_values = evalGateSelectors circuit_selectors_info gateSelectors 
     unfiltered = map (runStraightLine eval_vars) gate_prgs
     filtered   = zipWith (\s cons -> map (*s) cons) sel_values unfiltered
     gates      = combineFilteredGateConstraints filtered
@@ -112,7 +109,7 @@ evalAllPlonkConstraints
       current = [z] ++ pp_chunk ++ [znext]
       f :: (FExt,FExt) -> [FExt] -> [FExt] -> FExt
       f (prev,next) numer denom = prev * product numer - next * product denom
-      
+
 --------------------------------------------------------------------------------
 
 -- | Each gate has some number of constrains (depending on which gate it is).

src/Plonk/Verifier.hs

@@ -26,14 +26,6 @@ import Misc.Aux
 
 --------------------------------------------------------------------------------
 
-{-
--- | Evaluate the zero polynomial @Z_H(x)@
-evalZeroPoly :: Int -> FExt -> FExt
-evalZeroPoly nn zeta = powExt_ zeta nn - 1 
--}
-
---------------------------------------------------------------------------------
-
 -- | Assuming valid FRI openings, this function checks the combined Plonk equations
 checkCombinedPlonkEquations :: CommonCircuitData -> ProofWithPublicInputs -> ProofChallenges -> Bool
 checkCombinedPlonkEquations common proof challenges = 

src/Types.hs

@@ -28,6 +28,9 @@ newtype LookupTable
   = MkLookupTable [(Word64,Word64)]
   deriving (Eq,Show,Generic)
 
+fromLookupTable :: LookupTable -> [(F,F)] 
+fromLookupTable (MkLookupTable pairs) = [ (toF inp, toF out) | (inp,out) <- pairs ]
+
 instance ToJSON   LookupTable where toJSON (MkLookupTable x) = toJSON x
 instance FromJSON LookupTable where parseJSON o = MkLookupTable <$> parseJSON o
 

src/testmain.hs

@@ -20,9 +20,10 @@ import qualified Data.ByteString.Lazy.Char8 as L
 --------------------------------------------------------------------------------
 
 main = do
-  let prefix = "fibonacci"
+  -- let prefix = "fibonacci"
   -- let prefix = "recursion_outer"
   -- let prefix = "lookup"
+  let prefix = "multi_lookup"
 
   text_common <- L.readFile ("../json/" ++ prefix ++ "_common.json")
   text_vkey   <- L.readFile ("../json/" ++ prefix ++ "_vkey.json"  )