Add permute_hash_values utility

Cargo.toml

@@ -10,6 +10,7 @@ members = [
     "crates/circuit_prover",
     "crates/stark_verifier",
     "crates/stark_verifier_examples",
+    "crates/unpacker",
 ]
 resolver = "2"
 
@@ -61,6 +62,7 @@ circuit-multiverifier = { path = "crates/circuit_multiverifier"}
 circuit-prover = { path = "crates/circuit_prover"}
 circuits-stark-verifier = { path = "crates/stark_verifier"}
 circuits-stark-verifier-examples = { path = "crates/stark_verifier_examples"}
+circuits-unpacker = { path = "crates/unpacker" }
 
 # Compile the circuit-prover and stwo crates with optimization level 3 to speed up the blake related tests.
 [profile.test.package.circuit-prover]

crates/unpacker/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "circuits-unpacker"
+version.workspace = true
+edition.workspace = true
+
+[dependencies]
+circuits.workspace = true
+itertools.workspace = true
+
+[dev-dependencies]
+stwo.workspace = true

crates/unpacker/src/lib.rs

@@ -0,0 +1 @@
+pub mod permutation;

crates/unpacker/src/permutation.rs

@@ -0,0 +1,114 @@
+use std::collections::HashMap;
+
+use circuits::blake::{BLAKE2S_DIGEST_N_WORDS, HashValue};
+use circuits::circuit::Permutation;
+use circuits::context::{Context, Var};
+use circuits::eval;
+use circuits::ivalue::{IValue, qm31_from_u32s};
+use circuits::ops::{Guess, eq};
+use circuits::wrappers::M31Wrapper;
+use itertools::zip_eq;
+
+#[cfg(test)]
+#[path = "permutation_test.rs"]
+pub mod test;
+
+/// Constrains `outputs` to be a multiset permutation of `inputs`, treating each [`HashValue`] as an
+/// atomic unit.
+///
+/// Every word is *tagged* with the index of its source hash in the `u` coordinate — words are
+/// packed as `(low_u16, high_u16, 0, 0)`, so the `u` coordinate is free and a tagged word is
+/// `(low_u16, high_u16, index, 0)`:
+///
+/// * each input hash `j` is tagged with the constant `j`;
+/// * each output hash is tagged with a single base-field-guessed source index `s`, reused for all
+///   [`BLAKE2S_DIGEST_N_WORDS`] of its words.
+///
+/// A per-column [`Permutation`] gate then ties the tagged inputs to fresh permutation-output
+/// variables, and each such variable is constrained equal to the corresponding tagged output. The
+/// per-column multiset checks plus the shared-per-hash tag together guarantee each output hash
+/// equals exactly one input hash, i.e. `outputs` is a genuine permutation of `inputs`.
+///
+/// Assumes the words of `inputs` and `outputs` are valid `u32` packings (`(low, high, 0, 0)`), i.e.
+/// the [`U32Wrapper`](circuits::wrappers::U32Wrapper) invariant; the soundness of the tag depends
+/// on the `u`/`iu` coordinates being zero.
+///
+/// Panics if `inputs` and `outputs` have different lengths, or if `outputs` is not a multiset
+/// permutation of `inputs` (some output hash has no matching input hash).
+pub fn permute_hash_values<Value: IValue>(
+    ctx: &mut Context<Value>,
+    inputs: &[HashValue<Var>],
+    outputs: &[HashValue<Var>],
+) {
+    let n = inputs.len();
+    assert_eq!(n, outputs.len(), "inputs and outputs must have the same length");
+
+    // The eight u32 words of a hash, used as the lookup key for matching outputs to inputs.
+    let key_from_hash =
+        |ctx: &Context<Value>, hash: &HashValue<Var>| -> [u32; BLAKE2S_DIGEST_N_WORDS] {
+            std::array::from_fn(|c| ctx.get(*hash[c].get()).unpack_u32())
+        };
+
+    // Tag each input word with its hash index `j`: `(low, high, 0, 0) + (0, 0, j, 0)`, and map each
+    // distinct input-hash value (its eight u32 words) to the input indices that hold it. Duplicate
+    // input hashes map to several indices and are consumed one per matching output.
+    let mut tagged_in: Vec<Vec<Var>> =
+        (0..BLAKE2S_DIGEST_N_WORDS).map(|_| Vec::with_capacity(n)).collect();
+    let mut indices_by_value: HashMap<[u32; BLAKE2S_DIGEST_N_WORDS], Vec<usize>> =
+        HashMap::with_capacity(n);
+    for (j, hash) in inputs.iter().enumerate() {
+        let tag = ctx.constant(qm31_from_u32s(0, 0, j as u32, 0));
+        for (col, word) in zip_eq(tagged_in.iter_mut(), hash.iter()) {
+            col.push(eval!(ctx, (*word.get()) + (tag)));
+        }
+
+        indices_by_value.entry(key_from_hash(ctx, hash)).or_default().push(j);
+    }
+
+    // `u = (0, 0, 1, 0)`; multiplying a base-field index `s` by it lifts it to `(0, 0, s, 0)`.
+    let u = ctx.u();
+
+    // `tagged_out[c]` collects the fresh permutation-output variables of word-column `c`.
+    let mut tagged_out: Vec<Vec<Var>> =
+        (0..BLAKE2S_DIGEST_N_WORDS).map(|_| Vec::with_capacity(n)).collect();
+
+    for out_hash in outputs {
+        // Find and guess the source index: an as-yet-unused input whose hash equals this output,
+        // constrained to the base field and lifted into the `u` coordinate. One guess per output
+        // hash, shared across all of its words.
+        let index_in_inputs = indices_by_value
+            .get_mut(&key_from_hash(ctx, out_hash))
+            .and_then(|idxs| idxs.pop())
+            .expect("output hash is not a permutation of the inputs");
+        let index_var = M31Wrapper::new_unsafe(Value::from_qm31(qm31_from_u32s(
+            index_in_inputs as u32,
+            0,
+            0,
+            0,
+        )))
+        .guess(ctx);
+        let tag_u = eval!(ctx, (*index_var.get()) * (u));
+
+        // Tag each word of this hash and add it to its word-column as one permutation output.
+        for (col, out_word) in zip_eq(tagged_out.iter_mut(), out_hash.iter()) {
+            // Tagged output word: `(low, high, 0, 0) + (0, 0, index, 0)`.
+            let tagged_word = eval!(ctx, (*out_word.get()) + (tag_u));
+            // The per-column Permutation gate yields its own fresh variable, so allocate one and
+            // pin it to `tagged_word` with the `eq`. The gate (added per column below) then proves
+            // these tagged outputs are a multiset-permutation of the column's tagged inputs.
+            let perm_out = ctx.new_var(ctx.get(tagged_word));
+            eq(ctx, perm_out, tagged_word);
+            col.push(perm_out);
+        }
+    }
+
+    // One Permutation gate per word-column: its tagged inputs and permutation outputs must agree as
+    // multisets.
+    for (in_col, out_col) in zip_eq(tagged_in, tagged_out) {
+        ctx.stats.permutation_inputs += in_col.len();
+        ctx.circuit.permutation.push(Permutation {
+            inputs: in_col.iter().map(|v| v.idx).collect(),
+            outputs: out_col.iter().map(|v| v.idx).collect(),
+        });
+    }
+}

crates/unpacker/src/permutation_test.rs

@@ -0,0 +1,138 @@
+use circuits::blake::HashValue;
+use circuits::circuit::Circuit;
+use circuits::context::{Context, TraceContext, Var};
+use circuits::ivalue::{IValue, NoValue};
+use circuits::ops::Guess;
+use circuits::wrappers::U32Wrapper;
+use stwo::core::fields::qm31::QM31;
+
+use crate::permutation::permute_hash_values;
+
+fn guess_hashes(context: &mut TraceContext, words: &[[u32; 8]]) -> Vec<HashValue<Var>> {
+    words.iter().map(|w| HashValue::from(*w).guess(context)).collect()
+}
+
+#[test]
+fn test_permute_hash_values() {
+    let mut context = TraceContext::default();
+
+    let inputs_u32: [[u32; 8]; 3] = [
+        [1, 2, 3, 4, 5, 6, 7, 8],
+        [9, 10, 11, 12, 13, 14, 15, 16],
+        [17, 18, 19, 20, 21, 22, 23, 24],
+    ];
+    // A genuine reordering of the input hashes (reverse).
+    let outputs_u32: Vec<[u32; 8]> = inputs_u32.iter().rev().copied().collect();
+
+    let inputs = guess_hashes(&mut context, &inputs_u32);
+    let outputs = guess_hashes(&mut context, &outputs_u32);
+
+    permute_hash_values(&mut context, &inputs, &outputs);
+
+    let context = context.finalize(false);
+    context.circuit().check_yields();
+    assert!(context.is_circuit_valid());
+}
+
+/// `outputs` that is not a permutation of `inputs` (a hash absent from the inputs) has no source
+/// index, so witness generation panics.
+#[test]
+#[should_panic(expected = "output hash is not a permutation of the inputs")]
+fn test_permute_hash_values_panics_on_non_permutation() {
+    let mut context = TraceContext::default();
+
+    let inputs_u32: [[u32; 8]; 2] = [[1, 2, 3, 4, 5, 6, 7, 8], [9, 10, 11, 12, 13, 14, 15, 16]];
+    // Second hash does not appear among the inputs.
+    let outputs_u32: [[u32; 8]; 2] = [[9, 10, 11, 12, 13, 14, 15, 16], [99, 2, 3, 4, 5, 6, 7, 8]];
+
+    let inputs = guess_hashes(&mut context, &inputs_u32);
+    let outputs = guess_hashes(&mut context, &outputs_u32);
+
+    permute_hash_values(&mut context, &inputs, &outputs);
+}
+
+/// An output hash assembled from words of *different* input hashes (each word-column is
+/// individually a valid permutation, but the mixed hash is not one of the inputs) is not a
+/// permutation, so witness generation panics.
+#[test]
+#[should_panic(expected = "output hash is not a permutation of the inputs")]
+fn test_permute_hash_values_panics_on_column_mixing() {
+    let mut context = TraceContext::default();
+
+    let inputs_u32: [[u32; 8]; 2] = [[1, 2, 3, 4, 5, 6, 7, 8], [11, 12, 13, 14, 15, 16, 17, 18]];
+    // Each column below is a valid permutation of the inputs' column, but word 1 is swapped between
+    // the two hashes, so neither output hash equals a whole input hash.
+    let outputs_u32: [[u32; 8]; 2] = [[1, 12, 3, 4, 5, 6, 7, 8], [11, 2, 13, 14, 15, 16, 17, 18]];
+
+    let inputs = guess_hashes(&mut context, &inputs_u32);
+    let outputs = guess_hashes(&mut context, &outputs_u32);
+
+    permute_hash_values(&mut context, &inputs, &outputs);
+}
+
+/// Inputs containing duplicate hashes: matching must consume the per-value index list one at a time
+/// and still produce a valid bijection.
+#[test]
+fn test_permute_hash_values_with_duplicates() {
+    let mut context = TraceContext::default();
+
+    let inputs_u32: [[u32; 8]; 3] = [
+        [1, 2, 3, 4, 5, 6, 7, 8],
+        [1, 2, 3, 4, 5, 6, 7, 8], // duplicate of input 0
+        [9, 10, 11, 12, 13, 14, 15, 16],
+    ];
+    // Same multiset, reordered; the two equal hashes are mapped to the two equal-valued positions.
+    let outputs_u32: [[u32; 8]; 3] =
+        [[9, 10, 11, 12, 13, 14, 15, 16], [1, 2, 3, 4, 5, 6, 7, 8], [1, 2, 3, 4, 5, 6, 7, 8]];
+
+    let inputs = guess_hashes(&mut context, &inputs_u32);
+    let outputs = guess_hashes(&mut context, &outputs_u32);
+
+    permute_hash_values(&mut context, &inputs, &outputs);
+
+    let context = context.finalize(false);
+    context.circuit().check_yields();
+    assert!(context.is_circuit_valid());
+}
+
+/// Builds a `HashValue<Value>` from raw words for any `Value` (witness or witness-less), so the
+/// same circuit-building code can run under `QM31` and `NoValue`.
+fn hash_from<Value: IValue>(words: [u32; 8]) -> HashValue<Value> {
+    HashValue(words.map(|w| U32Wrapper::new_unsafe(Value::pack_u32(w))))
+}
+
+/// Runs `permute_hash_values` in a fresh `Context<Value>` and returns the finalized circuit
+/// structure (gates only — no witness values).
+fn build_permute_circuit<Value: IValue>(
+    inputs_u32: &[[u32; 8]],
+    outputs_u32: &[[u32; 8]],
+) -> Circuit {
+    let mut ctx = Context::<Value>::default();
+    let inputs: Vec<HashValue<Var>> =
+        inputs_u32.iter().map(|w| hash_from::<Value>(*w).guess(&mut ctx)).collect();
+    let outputs: Vec<HashValue<Var>> =
+        outputs_u32.iter().map(|w| hash_from::<Value>(*w).guess(&mut ctx)).collect();
+    permute_hash_values(&mut ctx, &inputs, &outputs);
+    ctx.finalize(false).context.circuit
+}
+
+/// The circuit topology must not depend on the witness: building with concrete `QM31` values (where
+/// the source indices are recovered by matching wire values) must yield exactly the same gates as
+/// building with `NoValue` (where no values exist and matching degenerates to sequential indices).
+#[test]
+fn test_permute_hash_values_structure_is_witness_independent() {
+    let inputs_u32: [[u32; 8]; 3] = [
+        [1, 2, 3, 4, 5, 6, 7, 8],
+        [9, 10, 11, 12, 13, 14, 15, 16],
+        [17, 18, 19, 20, 21, 22, 23, 24],
+    ];
+    // A non-trivial reordering, so the QM31 build's value-matching yields a non-identity
+    // permutation while the NoValue build necessarily falls back to the identity — yet the
+    // gates must match.
+    let outputs_u32: Vec<[u32; 8]> = inputs_u32.iter().rev().copied().collect();
+
+    let with_values = build_permute_circuit::<QM31>(&inputs_u32, &outputs_u32);
+    let without_values = build_permute_circuit::<NoValue>(&inputs_u32, &outputs_u32);
+
+    assert_eq!(with_values, without_values);
+}