Optsponge

packages/pickles/src/Pickles/OptSponge.purs

@@ -0,0 +1,227 @@
+-- | Conditional sponge for optional absorption.
+-- |
+-- | Unlike the regular sponge which tracks position at compile time,
+-- | OptSponge tracks position as a circuit boolean because conditional
+-- | absorption makes the position data-dependent at runtime.
+-- |
+-- | Reference: mina/src/lib/crypto/pickles/opt_sponge.ml
+module Pickles.OptSponge
+  ( OptSponge
+  , create
+  , squeeze
+  ) where
+
+import Prelude
+
+import Data.Array as Array
+import Data.Fin (unsafeFinite)
+import Data.Foldable (foldM)
+import Data.Tuple (Tuple(..), fst, snd)
+import Data.Vector (Vector)
+import Data.Vector as Vector
+import Partial.Unsafe (unsafePartial)
+import Poseidon (class PoseidonField)
+import Safe.Coerce (coerce)
+import Snarky.Circuit.CVar (sub_)
+import Snarky.Circuit.DSL (class CircuitM, Bool(..), BoolVar, FVar, Snarky, addConstraint, all_, and_, any_, const_, exists, false_, if_, mul_, not_, or_, readCVar, xor_)
+import Snarky.Circuit.Kimchi.Poseidon (poseidon)
+import Snarky.Constraint.Basic (r1cs)
+import Snarky.Constraint.Kimchi (KimchiConstraint)
+import Snarky.Curves.Class (class PrimeField)
+
+type OptSponge f =
+  { state :: Vector 3 (FVar f)
+  , pos :: BoolVar f
+  , needsFinalPermuteIfEmpty :: Boolean
+  }
+
+-- | Create a fresh OptSponge with zero initial state.
+create :: forall f. PrimeField f => OptSponge f
+create =
+  { state: Vector.replicate (const_ zero)
+  , pos: false_
+  , needsFinalPermuteIfEmpty: true
+  }
+
+-- | Squeeze the sponge after consuming all pending absorptions.
+-- | Returns state[0] after the final permutation.
+squeeze
+  :: forall f t m
+   . PoseidonField f
+  => CircuitM f (KimchiConstraint f) t m
+  => OptSponge f
+  -> Array (Tuple (BoolVar f) (FVar f))
+  -> Snarky (KimchiConstraint f) t m (FVar f)
+squeeze sponge pending = do
+  finalState <- consume sponge pending
+  pure $ Vector.index finalState (unsafeFinite 0)
+
+-------------------------------------------------------------------------------
+-- Internal
+-------------------------------------------------------------------------------
+
+-- | Conditionally add x to rate position based on pos.
+-- | When pos = false (0), adds to state[0]; when pos = true (1), adds to state[1].
+-- |
+-- | For each rate position j, witnesses s_j' and constrains:
+-- |   x * flag_j = s_j' - s_j
+-- | where flag_0 = NOT pos, flag_1 = pos.
+addIn
+  :: forall f t m
+   . CircuitM f (KimchiConstraint f) t m
+  => Vector 3 (FVar f)
+  -> BoolVar f
+  -> FVar f
+  -> Snarky (KimchiConstraint f) t m (Vector 3 (FVar f))
+addIn state pos x = do
+  let
+    iEquals0 = not_ pos
+    iEquals1 = pos
+    s0 = Vector.index state (unsafeFinite 0)
+    s1 = Vector.index state (unsafeFinite 1)
+
+  -- Update position 0: s0' = s0 + (NOT pos) * x
+  s0' <- exists do
+    s0Val <- readCVar s0
+    flagVal <- readCVar (coerce iEquals0 :: FVar f)
+    xVal <- readCVar x
+    pure $ if flagVal /= zero then s0Val + xVal else s0Val
+  addConstraint $ r1cs
+    { left: x
+    , right: coerce iEquals0
+    , output: s0' `sub_` s0
+    }
+
+  -- Update position 1: s1' = s1 + pos * x
+  s1' <- exists do
+    s1Val <- readCVar s1
+    flagVal <- readCVar (coerce iEquals1 :: FVar f)
+    xVal <- readCVar x
+    pure $ if flagVal /= zero then s1Val + xVal else s1Val
+  addConstraint $ r1cs
+    { left: x
+    , right: coerce iEquals1
+    , output: s1' `sub_` s1
+    }
+
+  pure $ Vector.modifyAt (unsafeFinite 0) (const s0')
+    $ Vector.modifyAt (unsafeFinite 1) (const s1') state
+
+-- | Conditional poseidon permutation.
+-- | Runs the permutation, then selects between permuted and original state.
+condPermute
+  :: forall f t m
+   . PoseidonField f
+  => CircuitM f (KimchiConstraint f) t m
+  => BoolVar f
+  -> Vector 3 (FVar f)
+  -> Snarky (KimchiConstraint f) t m (Vector 3 (FVar f))
+condPermute permute state = do
+  permuted <- poseidon state
+  if_ permute permuted state
+
+-- | Process one pair of conditional absorptions.
+-- | Matches OCaml's consume_pairs fold body exactly.
+consumePair
+  :: forall f t m
+   . PoseidonField f
+  => CircuitM f (KimchiConstraint f) t m
+  => { state :: Vector 3 (FVar f), pos :: BoolVar f }
+  -> Tuple { b :: BoolVar f, x :: FVar f } { b :: BoolVar f, x :: FVar f }
+  -> Snarky (KimchiConstraint f) t m { state :: Vector 3 (FVar f), pos :: BoolVar f }
+consumePair { state, pos: p } (Tuple first second) = do
+  let { b, x } = first
+  let { b: b', x: y } = second
+
+  -- Position tracking
+  p' <- xor_ p b
+  posAfter <- xor_ p' b'
+
+  -- Mask y by b'
+  yMasked <- mul_ y (coerce b')
+
+  -- Only add y after permutation when (b=1, b'=1, p=1)
+  addInYAfter <- all_ [ b, b', p ]
+  let addInYBefore = not_ addInYAfter
+
+  -- Add x*b to state at position p
+  xb <- mul_ x (coerce b)
+  state1 <- addIn state p xb
+
+  -- Add yMasked*before_flag to state at position p'
+  yBefore <- mul_ yMasked (coerce addInYBefore)
+  state2 <- addIn state1 p' yBefore
+
+  -- Compute permute flag: (b && b') || (p && (b || b'))
+  -- OCaml evaluates right-to-left for list construction, so:
+  -- 1. all [p, b ||| b'] is computed first (rightmost list element)
+  -- 2. all [b, b'] is computed second (leftmost)
+  -- 3. any [left, right]
+  bOrB' <- or_ b b'
+  pAndBOrB' <- and_ p bOrB'
+  bAndB' <- and_ b b'
+  permute <- or_ bAndB' pAndBOrB'
+
+  -- Conditional permutation
+  state3 <- condPermute permute state2
+
+  -- Add yMasked*after_flag to state at position p'
+  yAfter <- mul_ yMasked (coerce addInYAfter)
+  state4 <- addIn state3 p' yAfter
+
+  pure { state: state4, pos: posAfter }
+
+-- | Consume all pending absorptions and perform final permutation.
+consume
+  :: forall f t m
+   . PoseidonField f
+  => CircuitM f (KimchiConstraint f) t m
+  => OptSponge f
+  -> Array (Tuple (BoolVar f) (FVar f))
+  -> Snarky (KimchiConstraint f) t m (Vector 3 (FVar f))
+consume { state: initState, pos: startPos, needsFinalPermuteIfEmpty } input = do
+  let
+    n = Array.length input
+    numPairs = n / 2
+    remaining = n - 2 * numPairs
+
+    -- Build pairs array matching OCaml's Array.init
+    pairs = Array.mapWithIndex
+      ( \i _ ->
+          let
+            first = unsafePartial $ Array.unsafeIndex input (2 * i)
+            second = unsafePartial $ Array.unsafeIndex input (2 * i + 1)
+          in
+            Tuple
+              { b: fst first, x: snd first }
+              { b: fst second, x: snd second }
+      )
+      (Array.replicate numPairs unit)
+
+  -- Process all pairs
+  { state, pos } <- foldM consumePair { state: initState, pos: startPos } pairs
+
+  -- Compute empty_input = not (any (map fst input))
+  emptyInput <- not $ any_ (map fst input)
+
+  -- Handle remainder and compute should_permute
+  case remaining of
+    0 -> do
+      shouldPermute <-
+        if needsFinalPermuteIfEmpty then or_ emptyInput pos
+        else pure pos
+      condPermute shouldPermute state
+    1 -> do
+      let
+        lastElem = unsafePartial $ Array.unsafeIndex input (n - 1)
+        b = fst lastElem
+        x = snd lastElem
+        p = pos
+      _posAfter <- xor_ p b
+      xb <- mul_ x (coerce b)
+      state' <- addIn state p xb
+      shouldPermute <-
+        if needsFinalPermuteIfEmpty then any_ [ p, b, emptyInput ]
+        else any_ [ p, b ]
+      condPermute shouldPermute state'
+    _ -> pure state -- unreachable

packages/pickles/src/Pickles/Step/ChallengeDigest.purs

@@ -4,113 +4,45 @@
 -- | challenge digest by conditionally absorbing old bulletproof challenges
 -- | based on which previous proofs are "real" vs "dummy".
 -- |
--- | The conditional absorption uses an "optional sponge" pattern: when the
--- | mask is true, absorb the value; when false, leave sponge state unchanged.
+-- | Uses OptSponge which tracks sponge position as a circuit boolean,
+-- | matching OCaml's `Opt_sponge` exactly.
 -- |
 -- | Reference: step_verifier.ml:880-909 (challenge digest computation)
 module Pickles.Step.ChallengeDigest
-  ( -- * Conditional Sponge Operations
-    absorbConditional
-  , absorbManyConditional
-  -- * Challenge Digest
-  , ChallengeDigestInput
+  ( ChallengeDigestInput
   , challengeDigestCircuit
   ) where
 
 import Prelude
 
-import Control.Monad.State.Trans (StateT(..))
-import Data.Newtype (wrap)
-import Data.Traversable (for_)
+import Data.Array as Array
+import Data.Foldable (foldMap)
 import Data.Tuple (Tuple(..))
 import Data.Vector (Vector)
 import Data.Vector as Vector
-import Pickles.Sponge (SpongeM, squeeze)
+import Pickles.OptSponge as OptSponge
 import Pickles.Verify.Types (BulletproofChallenges)
 import Poseidon (class PoseidonField)
-import Snarky.Circuit.DSL (class CircuitM, BoolVar, FVar, if_)
+import Snarky.Circuit.DSL (class CircuitM, BoolVar, FVar, Snarky)
 import Snarky.Circuit.DSL as SizedF
-import Snarky.Circuit.RandomOracle.Sponge as CircuitSponge
 import Snarky.Constraint.Kimchi (KimchiConstraint)
-import Snarky.Curves.Class (class FieldSizeInBits, class PrimeField)
-
--------------------------------------------------------------------------------
--- | Conditional Sponge Operations
--------------------------------------------------------------------------------
-
--- | Conditionally absorb a field element into the sponge.
--- |
--- | When `keep = true`, absorb the value. When `keep = false`, leave the
--- | sponge state unchanged.
--- |
--- | This implements the "optional sponge" pattern from OCaml's `Opt_sponge`.
--- |
--- | In circuit-land, both branches are computed but `if_` selects between:
--- | - New sponge state (after absorption)
--- | - Original sponge state (unchanged)
-absorbConditional
-  :: forall f t m
-   . PoseidonField f
-  => CircuitM f (KimchiConstraint f) t m
-  => BoolVar f
-  -> FVar f
-  -> SpongeM f (KimchiConstraint f) t m Unit
-absorbConditional keep x = wrap $ StateT \oldSponge -> do
-  -- Compute what the new state would be after absorption
-  newSponge <- CircuitSponge.absorb x oldSponge
-
-  -- Use if_ to conditionally select between new and old state vectors
-  -- The spongeState metadata tracks absorb/squeeze mode - we use the
-  -- newSponge's spongeState since we're always in "absorbing" mode here.
-  resultState <- if_ keep newSponge.state oldSponge.state
-
-  pure $ Tuple unit { state: resultState, spongeState: newSponge.spongeState }
-
--- | Conditionally absorb many field elements.
--- |
--- | All elements share the same `keep` flag, so either all are absorbed
--- | or none are.
-absorbManyConditional
-  :: forall f n t m
-   . PoseidonField f
-  => CircuitM f (KimchiConstraint f) t m
-  => BoolVar f
-  -> Vector n (FVar f)
-  -> SpongeM f (KimchiConstraint f) t m Unit
-absorbManyConditional keep xs = for_ xs (absorbConditional keep)
-
--------------------------------------------------------------------------------
--- | Challenge Digest
--------------------------------------------------------------------------------
+import Snarky.Curves.Class (class PrimeField)
 
 -- | Input for challenge digest computation.
 -- |
--- | This contains:
 -- | - `mask`: For each previous proof, whether it's "real" (should be absorbed)
 -- | - `oldChallenges`: The bulletproof challenges from each previous proof
 -- |
--- | The mask determines which proofs' challenges get absorbed into the sponge.
--- | Dummy proofs (mask = false) don't contribute to the digest.
--- |
 -- | Reference: step_verifier.ml:880-909
 type ChallengeDigestInput n d f b =
-  { -- | Which previous proofs are real (should have challenges absorbed)
-    mask :: Vector n b
-  , -- | Bulletproof challenges from each previous proof (d challenges each)
-    oldChallenges :: Vector n (BulletproofChallenges d f)
+  { mask :: Vector n b
+  , oldChallenges :: Vector n (BulletproofChallenges d f)
   }
 
 -- | Compute the challenge digest from old bulletproof challenges.
 -- |
--- | This circuit:
--- | 1. Takes a mask indicating which previous proofs are "real"
--- | 2. For each proof where mask = true, absorbs all bulletproof challenges
--- | 3. Squeezes the sponge to get the final digest
--- |
--- | The caller should initialize the sponge before calling this circuit.
--- | The sponge state after this call can be used for further operations.
--- |
--- | Reference: step_verifier.ml:880-909
+-- | Uses OptSponge to conditionally absorb challenges based on a mask,
+-- | matching OCaml's constraint structure exactly.
 -- |
 -- | ```ocaml
 -- | let challenge_digest =
@@ -126,17 +58,16 @@ type ChallengeDigestInput n d f b =
 challengeDigestCircuit
   :: forall n d f t m
    . PrimeField f
-  => FieldSizeInBits f 255
   => PoseidonField f
   => CircuitM f (KimchiConstraint f) t m
   => ChallengeDigestInput n d (FVar f) (BoolVar f)
-  -> SpongeM f (KimchiConstraint f) t m (FVar f)
-challengeDigestCircuit { mask, oldChallenges } = do
-  -- Absorb challenges from each previous proof based on mask
-  for_ (Vector.zip mask oldChallenges) \(Tuple keep chals) ->
-    -- For each proof, absorb all scalar challenges
-    for_ chals \chal ->
-      absorbConditional keep (SizedF.toField chal)
-
-  -- Squeeze to get the digest
-  squeeze
+  -> Snarky (KimchiConstraint f) t m (FVar f)
+challengeDigestCircuit { mask, oldChallenges } =
+  let
+    pending = Array.fromFoldable $ foldMap
+      ( \(Tuple keep chals) ->
+          map (Tuple keep <<< SizedF.toField) (Array.fromFoldable chals)
+      )
+      (Vector.zip mask oldChallenges)
+  in
+    OptSponge.squeeze OptSponge.create pending

packages/pickles/src/Pickles/Step/FinalizeOtherProof.purs

@@ -38,7 +38,7 @@ import Pickles.PlonkChecks.GateConstraints (GateConstraintInput)
 import Pickles.PlonkChecks.Permutation (PermutationInput)
 import Pickles.ProofWitness (ProofWitness)
 import Pickles.Sponge (SpongeM, absorb, liftSnarky, squeezeScalarChallenge)
-import Pickles.Step.ChallengeDigest (challengeDigestCircuit) as ChallengeDigest
+import Pickles.Step.ChallengeDigest (ChallengeDigestInput, challengeDigestCircuit) as ChallengeDigest
 import Pickles.Verify.Types (BulletproofChallenges, PlonkExpanded, UnfinalizedProof, expandPlonkMinimalCircuit, toPlonkMinimal)
 import Poseidon (class PoseidonField)
 import Prim.Int (class Add)