Xhat step (#112)

martyall · claude · web-flow · commit 8bb86d0bfdd3 · 2026-03-14T14:29:09.000-07:00
* match xhat_step_circuit to OCaml constraint system

PublicInputCommit changes:
- Add CorrectionMode (PureCorrections | InCircuitCorrections) to control
  how correction points are combined
- PureCorrections: match OCaml's multiscale_known map-then-reduce structure:
  Phase 1 runs all scaleFast2' (collecting results), Phase 2 reduces results
  pairwise with addComplete, Phase 3 adds correction sum as constant point
- Corrections summed as pure field arithmetic (addPurePt) matching OCaml
  where add_fast on constant corrections is optimized away
- Plumb correctionMode through IncrementallyVerifyProofParams, WrapParams,
  and all callers

Fixture fix:
- Fix xhat_step_circuit domain from 2^15 to 2^16 in dump_circuit_impl.ml
  (step circuit uses Nat.N16, wrap uses Nat.N15)
- Regenerate fixture with correct Lagrange commitments

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;

* pedantic build

---------

Co-authored-by: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/mina b/mina
@@ -1 +1 @@
-Subproject commit 2a5ad7f585f32b857cb4e4e3d6890b8cd194f91d
+Subproject commit b739ec38204c055a5d38b9919a68596b01331d1b
diff --git a/packages/pickles-circuit-diffs-types/spago.yaml b/packages/pickles-circuit-diffs-types/spago.yaml
@@ -2,5 +2,3 @@ package:
   name: pickles-circuit-diffs-types
   dependencies:
     - maybe
-    - prelude
-    - simple-json
diff --git a/packages/pickles-circuit-diffs/circuits/ocaml/xhat_step_circuit.json b/packages/pickles-circuit-diffs/circuits/ocaml/xhat_step_circuit.json
diff --git a/packages/pickles-circuit-diffs/circuits/ocaml/xhat_step_circuit_cached_constants.json b/packages/pickles-circuit-diffs/circuits/ocaml/xhat_step_circuit_cached_constants.json
diff --git a/packages/pickles-circuit-diffs/package-lock.json b/packages/pickles-circuit-diffs/package-lock.json
diff --git a/packages/pickles-circuit-diffs/spago.yaml b/packages/pickles-circuit-diffs/spago.yaml
@@ -29,6 +29,7 @@ package:
   test:
     main: Test.Pickles.CircuitDiffs.Main
     dependencies:
+      - aff
       - effect
       - exceptions
       - node-buffer
diff --git a/packages/pickles-circuit-diffs/test/Test/Pickles/CircuitDiffs/Main.purs b/packages/pickles-circuit-diffs/test/Test/Pickles/CircuitDiffs/Main.purs
@@ -410,8 +410,8 @@ spec =
             { lagrangeComms: (coerce $ vestaSrsLagrangeCommitments stepSrs 16 30) :: Array (AffinePoint (F Fp))
             , blindingH: (coerce $ vestaSrsBlindingGenerator stepSrs) :: AffinePoint (F Fp)
             }
-        pending "xhat_step_circuit"
-      -- exactMatch "xhat_step_circuit" (fromCompiledCircuit $ compileXhatStep stepSrsData)
+        -- pending "xhat_step_circuit"
+        exactMatch "xhat_step_circuit" (fromCompiledCircuit $ compileXhatStep stepSrsData)
       describe "Pickles Wrap sub-circuits" do
         exactMatch "finalize_other_proof_wrap_circuit" (fromCompiledCircuit compileFopWrap)
         exactMatch "group_map_wrap_circuit" (fromCompiledCircuit compileGroupMap)
diff --git a/packages/pickles-circuit-diffs/visualizer/src/CircuitDiffTable.purs b/packages/pickles-circuit-diffs/visualizer/src/CircuitDiffTable.purs
@@ -78,8 +78,11 @@ computeCachedConstantRows ps ocaml =
           let
             p = Array.index psConsts i
             o = Array.index ocamlConsts i
+            -- Compare by positional index and value only.
+            -- Variable numbers differ between PS (scattered) and OCaml (sequential)
+            -- because of different internal variable allocation strategies.
             m = case p, o of
-              Just a, Just b -> a.variable == b.variable && a.varType == b.varType && a.value == b.value
+              Just a, Just b -> a.value == b.value
               _, _ -> false
           in
             { ps: p, ocaml: o, match: m }
diff --git a/packages/pickles-circuit-diffs/visualizer/src/Main.purs b/packages/pickles-circuit-diffs/visualizer/src/Main.purs
@@ -83,8 +83,8 @@ mkApp circuitDiffTable = component "App" \_ -> React.do
     pure (pure unit)
 
   let
-    matches = Array.filter (\e -> e.status == "match") manifest
-    mismatches = Array.filter (\e -> e.status /= "match") manifest
+    matches = Array.sortWith _.name $ Array.filter (\e -> e.status == "match") manifest
+    mismatches = Array.sortWith _.name $ Array.filter (\e -> e.status /= "match") manifest
 
     statusDot :: String -> JSX
     statusDot status =
diff --git a/packages/pickles/spago.yaml b/packages/pickles/spago.yaml
@@ -5,6 +5,7 @@ package:
     - blake2
     - curves
     - effect
+    - either
     - exceptions
     - foldable-traversable
     - identity
diff --git a/packages/pickles/src/Pickles/PublicInputCommit.purs b/packages/pickles/src/Pickles/PublicInputCommit.purs
@@ -29,10 +29,12 @@ import Prelude
 
 import Data.Array as Array
 import Data.Array.NonEmpty as NEA
-import Data.Foldable (foldM)
+import Data.Either (Either(..), fromLeft)
+import Data.Foldable (foldM, foldl)
 import Data.Maybe (Maybe(..), fromJust)
 import Data.Reflectable (class Reflectable, reflectType)
 import Data.Symbol (class IsSymbol)
+import Data.Traversable (for)
 import Data.Tuple (Tuple(..))
 import Data.Vector (Vector)
 import Partial.Unsafe (unsafePartial)
@@ -55,8 +57,8 @@ import Snarky.Types.Shifted (SplitField(..), Type1(..), Type2(..))
 import Type.Proxy (Proxy(..))
 
 -- | Controls how correction points are combined during public input commitment.
--- | Currently unused (always InCircuitCorrections), but needed by
--- | pickles-circuit-diffs to prepare for Step verifier's PureCorrections mode.
+-- | PureCorrections: sum as pure field arithmetic (no circuit cost) — for Step verifier.
+-- | InCircuitCorrections: sum via in-circuit addComplete gates — for Wrap verifier.
 data CorrectionMode = PureCorrections | InCircuitCorrections
 
 -- | A deferred scalar multiplication that captures type-level parameters.
@@ -239,6 +241,7 @@ publicInputCommit
   => { curveParams :: CurveParams f
      , lagrangeComms :: Array (AffinePoint (F f))
      , blindingH :: AffinePoint (F f)
+     , correctionMode :: CorrectionMode
      | r
      }
   -> a
@@ -248,46 +251,81 @@ publicInputCommit params input = label "public-input-commit" do
   case NEA.fromArray results of
     Nothing -> pure (constPt params.blindingH)
     Just results' -> unsafePartial do
-      -- Separate scalar mul terms (with corrections) from cond_add terms
+      -- Separate correction points from terms
       let
-        correctionPts = Array.mapMaybe
+        rawCorrectionPts = Array.mapMaybe
           ( case _ of
-              AddWithCorrection r -> Just (constPt r.correction)
+              AddWithCorrection r -> Just r.correction
               CondAdd _ _ -> Nothing
           )
           (NEA.toArray results')
 
-      -- 1. Sum corrections in-circuit: init accumulator = sum([2^n] * B_i)
-      --    Matches OCaml: correction = reduce corrections ~f:add_fast
-      --    NOTE: These are all constant points so this could be done as pure math,
-      --    but OCaml generates CompleteAdd gates here and we match for equivalence.
-      let { head: corrHead, tail: corrTail } = fromJust $ Array.uncons correctionPts
-      init <- foldM (\acc c -> _.p <$> addComplete acc c) corrHead corrTail
-
-      -- 2. Fold over all terms, matching OCaml's fold order:
-      --    - AddWithCorrection: acc = add_fast(acc, scale_fast2'(...))
-      --      OCaml calls scale_fast2' INLINE during the fold (step_verifier.ml:474),
-      --      so VarBaseMul gates precede CompleteAdd for each term.
-      --    - CondAdd: acc = if_ b then add_fast(lagrange, acc) else acc
-      acc <- foldM
-        ( \acc term -> case term of
-            AddWithCorrection { scaleMul: DeferredScaleMul doScaleMul } -> do
-              point <- doScaleMul
-              _.p <$> addComplete acc point
-            CondAdd b lagrangePt -> do
-              added <- _.p <$> addComplete (constPt lagrangePt) acc
-              -- OCaml's Inner_curve.if_ uses tuple (F.if_ x, F.if_ y),
-              -- right-to-left evaluation means y is processed before x
-              y' <- if_ b added.y acc.y
-              x' <- if_ b added.x acc.x
-              pure { x: x', y: y' }
-        )
-        init
-        (NEA.toArray results')
-
-      -- 3. Negate and add blinding generator
-      negAcc <- Curves.negate acc
-      _.p <$> addComplete negAcc (constPt params.blindingH)
+      let { head: corrHead, tail: corrTail } = fromJust $ Array.uncons rawCorrectionPts
+      let allTerms = NEA.toArray results'
+
+      case params.correctionMode of
+        PureCorrections -> do
+          -- OCaml's step verifier: corrections are constant, add_fast on constants
+          -- generates no seal gates, and all fold CompleteAdd are deferred to the end.
+          -- Phase 1: Execute all scaleFast2' calls, collecting results.
+          --   This generates VarBaseMul + internal CompleteAdd gates.
+          evaluated <- for allTerms \term -> case term of
+            AddWithCorrection { scaleMul: DeferredScaleMul doScaleMul } ->
+              Left <$> doScaleMul
+            CondAdd b lagrangePt ->
+              pure (Right { b, lagrangePt })
+
+          -- Phase 2: Reduce results pairwise with addComplete.
+          --   Matches OCaml's List.reduce_exn ~f:(fun (_,b1) (_,b2) -> (_, add_fast b1 b2)).
+          --   Corrections are summed as pure constants in parallel.
+          let { head: first, tail: rest } = fromJust $ Array.uncons evaluated
+          acc <- foldM
+            ( \acc result -> case result of
+                Left point -> _.p <$> addComplete acc point
+                Right { b, lagrangePt } -> do
+                  added <- _.p <$> addComplete (constPt lagrangePt) acc
+                  y' <- if_ b added.y acc.y
+                  x' <- if_ b added.x acc.x
+                  pure { x: x', y: y' }
+            )
+            (fromLeft (constPt corrHead) first)
+            rest
+
+          -- Phase 3: Add total correction (constant) to accumulator.
+          --   Matches OCaml's: acc + constant(negate correction |> add_opt constant_part)
+          --   Raw corrections are already negated (negate(pow2pow(base, shift))),
+          --   so their sum IS the negated total correction.
+          --   The constPt triggers reduceToVariable constant caching → 1 Generic gate,
+          --   which double-packs with the queued R1CS from Phase 1's last if_.
+          let correctionPt = constPt $ foldl (addPurePt params.curveParams) corrHead corrTail
+          accWithCorr <- _.p <$> addComplete acc correctionPt
+
+          -- Negate and add blinding generator
+          negAcc <- Curves.negate accWithCorr
+          _.p <$> addComplete negAcc (constPt params.blindingH)
+
+        InCircuitCorrections -> do
+          -- Wrap verifier: corrections summed in-circuit, fold interleaved.
+          let corrPts = map constPt rawCorrectionPts
+          let { head: ch, tail: ct } = fromJust $ Array.uncons corrPts
+          init <- foldM (\acc c -> _.p <$> addComplete acc c) ch ct
+
+          acc <- foldM
+            ( \acc term -> case term of
+                AddWithCorrection { scaleMul: DeferredScaleMul doScaleMul } -> do
+                  point <- doScaleMul
+                  _.p <$> addComplete acc point
+                CondAdd b lagrangePt -> do
+                  added <- _.p <$> addComplete (constPt lagrangePt) acc
+                  y' <- if_ b added.y acc.y
+                  x' <- if_ b added.x acc.x
+                  pure { x: x', y: y' }
+            )
+            init
+            allTerms
+
+          negAcc <- Curves.negate acc
+          _.p <$> addComplete negAcc (constPt params.blindingH)
 
 -------------------------------------------------------------------------------
 -- | Helpers
@@ -338,6 +376,13 @@ unwrapPt { x: F x', y: F y' } = { x: x', y: y' }
 wrapPt :: forall f. AffinePoint f -> AffinePoint (F f)
 wrapPt { x, y } = { x: F x, y: F y }
 
+-- | Pure affine addition for summing constant correction points.
+-- | Handles the doubling case (same point) via EC.double.
+addPurePt :: forall f. PrimeField f => CurveParams f -> AffinePoint (F f) -> AffinePoint (F f) -> AffinePoint (F f)
+addPurePt params p1 p2
+  | unwrapPt p1 == unwrapPt p2 = EC.double params p1
+  | otherwise = wrapPt $ unsafePartial $ fromJust $ EC.toAffine $ EC.addAffine (unwrapPt p1) (unwrapPt p2)
+
 -- | Compute [2^k] * p by iterating pure doubling.
 pow2pow :: forall f. PrimeField f => CurveParams f -> AffinePoint (F f) -> Int -> AffinePoint (F f)
 pow2pow _ p k
diff --git a/packages/pickles/src/Pickles/Verify.purs b/packages/pickles/src/Pickles/Verify.purs
@@ -25,7 +25,7 @@ import Data.Vector (Vector, (:<))
 import Data.Vector as Vector
 import Pickles.FtComm (ftComm)
 import Pickles.IPA (CheckBulletproofInput, checkBulletproof)
-import Pickles.PublicInputCommit (class PublicInputCommit, publicInputCommit)
+import Pickles.PublicInputCommit (class PublicInputCommit, CorrectionMode, publicInputCommit)
 import Pickles.ShiftOps (IpaScalarOps)
 import Pickles.Sponge (SpongeM, initialSpongeCircuit, labelM, liftSnarky)
 import Pickles.Sponge as Sponge
@@ -60,6 +60,7 @@ type IncrementallyVerifyProofParams f r =
   , indexDigest :: f
   , endo :: f -- ^ EndoScalar constant for challenge expansion
   , groupMapParams :: GroupMapParams f
+  , correctionMode :: CorrectionMode
   | r
   }
 
diff --git a/packages/pickles/src/Pickles/Wrap/Circuit.purs b/packages/pickles/src/Pickles/Wrap/Circuit.purs
@@ -28,6 +28,7 @@ import Data.Vector (Vector, (:<))
 import Data.Vector as Vector
 import Pickles.Dummy (dummyWrapChallengesExpanded)
 import Pickles.Linearization.Types (LinearizationPoly)
+import Pickles.PublicInputCommit (CorrectionMode)
 import Pickles.Sponge (evalSpongeM, initialSpongeCircuit)
 import Pickles.Types (StepStatement, WrapIPARounds, WrapStatement)
 import Pickles.Verify (verify)
@@ -79,6 +80,7 @@ type WrapParams f =
   , linearizationPoly :: LinearizationPoly f
   -- Shared
   , endo :: f
+  , correctionMode :: CorrectionMode
   }
 
 -- | The Wrap circuit: finalizes deferred values and verifies IPA opening.
diff --git a/packages/pickles/test/Test/Pickles/PublicInputCommitment.purs b/packages/pickles/test/Test/Pickles/PublicInputCommitment.purs
@@ -12,7 +12,7 @@ import Data.Vector as Vector
 import Effect.Aff (Aff)
 import Effect.Class (liftEffect)
 import Partial.Unsafe (unsafePartial)
-import Pickles.PublicInputCommit (publicInputCommit)
+import Pickles.PublicInputCommit (CorrectionMode(..), publicInputCommit)
 import Pickles.PublicInputCommitment (publicInputCommitment)
 import Pickles.Step.Circuit (StepStatement)
 import Pickles.Types (StepIPARounds, WrapIPARounds)
@@ -177,7 +177,7 @@ spec cfg =
         => Vector nPublic (FVar StepCircuitField)
         -> Snarky (KimchiConstraint StepCircuitField) t Identity (AffinePoint (FVar StepCircuitField))
       circuit inputs =
-        publicInputCommit { curveParams: curveParams (Proxy @Vesta.G), lagrangeComms: ctx.lagrangeComms, blindingH: ctx.blindingH } inputs
+        publicInputCommit { curveParams: curveParams (Proxy @Vesta.G), lagrangeComms: ctx.lagrangeComms, blindingH: ctx.blindingH, correctionMode: InCircuitCorrections } inputs
 
       gen = Vector.generator (Proxy @nPublic) fpRangeGen
 
@@ -212,7 +212,7 @@ spec cfg =
         => StepFullXhatVar
         -> Snarky (KimchiConstraint StepCircuitField) t Identity (AffinePoint (FVar StepCircuitField))
       circuit inputs =
-        publicInputCommit { curveParams: curveParams (Proxy @Vesta.G), lagrangeComms: ctx.lagrangeComms, blindingH: ctx.blindingH } inputs
+        publicInputCommit { curveParams: curveParams (Proxy @Vesta.G), lagrangeComms: ctx.lagrangeComms, blindingH: ctx.blindingH, correctionMode: InCircuitCorrections } inputs
 
       rustFn :: StepFullXhat -> AffinePoint (F StepCircuitField)
       rustFn tup = unsafePartial $
@@ -266,7 +266,7 @@ spec cfg =
         => Vector nPublic (FVar WrapCircuitField)
         -> Snarky (KimchiConstraint WrapCircuitField) t Identity (AffinePoint (FVar WrapCircuitField))
       circuit inputs =
-        publicInputCommit { curveParams: curveParams (Proxy @Pallas.G), lagrangeComms: ctx.lagrangeComms, blindingH: ctx.blindingH } inputs
+        publicInputCommit { curveParams: curveParams (Proxy @Pallas.G), lagrangeComms: ctx.lagrangeComms, blindingH: ctx.blindingH, correctionMode: InCircuitCorrections } inputs
 
       gen = Vector.generator (Proxy @nPublic) (arbitrary :: Gen (F WrapCircuitField))
 
diff --git a/packages/pickles/test/Test/Pickles/TestContext.purs b/packages/pickles/test/Test/Pickles/TestContext.purs
@@ -103,6 +103,7 @@ import Pickles.PlonkChecks.GateConstraints (buildChallenges, buildEvalPoint)
 import Pickles.PlonkChecks.Permutation (permContribution, permScalar)
 import Pickles.PlonkChecks.XiCorrect (FrSpongeInput, emptyPrevChallengeDigest, frSpongeChallengesPure)
 import Pickles.ProofWitness (ProofWitness)
+import Pickles.PublicInputCommit (CorrectionMode(..))
 import Pickles.Sponge (initialSponge, runPureSpongeM)
 import Pickles.Sponge as Pickles.Sponge
 import Pickles.Step.Advice (class StepWitnessM, getStepInputFields)
@@ -749,6 +750,7 @@ buildWrapCircuitParams ctx =
     , srsLengthLog2: ctx.domainLog2
     , linearizationPoly: Linearization.vesta
     , groupMapParams: Kimchi.groupMapParams (Proxy @Vesta.G)
+    , correctionMode: InCircuitCorrections
     }
 
 -------------------------------------------------------------------------------
@@ -1797,6 +1799,7 @@ buildStepIVPParams ctx =
     , indexDigest: ProofFFI.vestaVerifierIndexDigest ctx.verifierIndex
     , endo: stepEndo
     , groupMapParams: Kimchi.groupMapParams (Proxy @Pallas.G)
+    , correctionMode: PureCorrections
     }
 
 -- | Build IVP circuit input for an Fp circuit verifying a Wrap (Pallas) proof.
diff --git a/packages/snarky-kimchi/spago.yaml b/packages/snarky-kimchi/spago.yaml
@@ -23,9 +23,7 @@ package:
     - prelude
     - quickcheck
     - record
-    - foreign
     - safe-coerce
-    - simple-json
     - snarky
     - snarky-curves
     - snarky-test-utils
@@ -38,17 +36,13 @@ package:
   test:
     main: Test.Snarky.Circuit.Kimchi.Main
     dependencies:
-      - aff
       - effect
       - gen
       - identity
-      - node-buffer
-      - node-fs
       - partial
       - quickcheck
       - snarky-test-utils
       - spec
       - spec-node
       - spec-quickcheck
       - strings
-      - unsafe-coerce
diff --git a/spago.lock b/spago.lock
