Dynamic Lookups

"No, no!" said the Queen. "Sentence first—verdict afterwards."

Dynamic Lookups

The previous section introduced the concept of lookups. Namely static lookups, in which the values to look up are fixed at circuit generation time and the only freedom afforded to the prover is to choose the entries to look up. This is fine if you know all the values you want to look up at compile time and the set is small enough (e.g. the XOR table of 8-bit values). But what if you don't? Or maybe the table would be too large to encode explicitly?

In this section we will introduce the concept of dynamic lookups. It is a natural extension of the static lookups we have seen so far: rather than looking up rows in Constant columns we will look up rows in Advice columns. Such dynamic lookups are widely used in machine emulation, let us start by getting a grasp of why such lookups are useful.

Motivation

The instruction set architecture (ISA) of a CPU defines the set of potential instructions that the processor can execute. This set can be rather large, for example the x86-64 ISA has over 1000 instructions, for the EVM (Ethereum Virtual Machine) this number is a more modest 140 instructions but includes very complex operations like KECCAK256. A naive approach to implementing a CPU in a zkSNARK would be to create a circuit for each instruction, then run every possible instruction in parallel and multiplex the result. This requires executing every possible instruction at every cycle of the CPU. Without lookups, this is essentially what people did in the early days of zkSNARKs for machine emulation, which in turn restricted the size and complexity of the ISA that could be emulated.

Dynamic lookups allow us to avoid this naive approach as follows:

1. We can create a table for each instruction. The table simply contains inputs/outputs for the instruction, for instance a table for the ADD instruction would contain rows of the form:

$$(a, b, a + b)$$

2. When emulating the CPU, we can then do a lookup in each table to retrieve the result for every possible instruction and then multiplex the results.

The key insight is that we need only a single row in the table for the instruction we are actually executing: all the other lookups can be "nopped out" and the result of the lookup is ignored. If every instruction previously required a separate circuit with \( m \) gates and we have \( n \) instructions, the dynamic lookup approach requires only \( n \) tables with \( m \) rows each whereas the original approach would require \( n \cdot m \) gates.

Example: Conditional Hashing

The example that we are going to explore is a gate that conditionally hashes a value, i.e.

$$ \mathsf{ConditionalHash}(x, b) = \begin{cases} \mathsf{Hash}(x) & \text{if } b = 1 \\ 0 & \text{if } b = 0 \\ \end{cases} $$

The goal is to only "pay" for the hash operation if \( b = 1 \): there will be some a priori bound on the maximum number of hash operations that can be performed, but we don't know ahead of time where / how many of these operations will be performed. We want to incur a cost that is proportional to this upper bound.

A Naive Approach

Before we proceed, let us consider a naive approach to this problem, without dynamic lookups. The baseline solution to this problem would be to create a circuit that hashes the value unconditionally, then use a multiplexer to select the output of the hash operation if the condition is true:

$$ \mathsf{ConditionalHash}(x, b) = \mathsf{Hash}(x) \cdot b $$

As hinted above, the issue with this approach is that the hash operation is always performed, even if the condition is false: we need to generate a circuit for \( \mathsf{Hash}(x) \) and assign the witness, even when \( b = 0 \). So if, for instance, you have a CPU which may compute a hash at every cycle, the circuit of the CPU would have to include the hash operation at every cycle even if the hash is not computed (e.g. if a simple arithmetic operation is performed instead).

In the EVM, the hash might be keccak256 and the condition might be the result of a comparison between the current instruction and the KECCAK256 opcode. But in order to keep things simple, we will use a simplified and round-reduced variant of the "Poseidon" hash function instead. This variant is not secure for cryptographic use!

The Poseidon Hash Function

Our simplified Poseidon hash function has a state width of 3 field elements and 8 rounds, we split the state into two parts: the "RATE" and the "CAPACITY" part:

#![allow(unused)]
fn main() {
const ROUNDS: usize = 8;
const WIDTH: usize = 3;

const CAPACITY: usize = 1;
const RATE: usize = WIDTH - CAPACITY;
}

The CAPACITY part is the "internal state" of the hash function, while the message to be hashed is added to the "RATE" part. As mentioned above there are 8 rounds in the hash function, each round consists of the following operations:

1. AddRoundConstant: $$(x, y, z) \mapsto (x + a, y + b, z + c)$$ Where a, b, c are constants (and different each round).

2. SBox: $$ (x, y, z) \mapsto (x^5, y^5, z^5) $$ The exponent 5 is chosen such that the map is a permutation.

3. Mix: $$ (x, y, z) \mapsto M (x, y, z) $$ Where M is a 3x3 matrix (the same for all rounds).

To hash a message \( (x, y) \) we initialize the state with \( (x, y, 0) \) and then apply the 8 rounds.

The hash digest is the first element of the final state:

$$ \mathsf{Hash}(x, y) = \mathsf{output} \text{ where } (\mathsf{output}, ?, ?) = \mathsf{PoseidonPermutation}(x, y, 0) $$

The Poseidon Table

We are going to have a table which contains every invocation of the Poseidon hash function and all their intermediate rounds. The table will contain the following columns:

// a simplified version of Poseidon permutation
#[derive(Debug, Clone)]
struct PoseidonTable<F: Field + Clone> {
    matrix: [[F; WIDTH]; WIDTH],
    round_constants: [[F; WIDTH]; ROUNDS],
    flag_start: Column<Fixed>, // start of permutation
    flag_round: Column<Fixed>, // apply round
    flag_final: Column<Fixed>, // end of permutation
    inp1: Column<Advice>,
    inp2: Column<Advice>,
    rndc: [Column<Fixed>; WIDTH],
    cols: [Column<Advice>; WIDTH],
    _ph: PhantomData<F>,
}

Some explanation is in order:

• matrix is the fixed matrix M from the Mix operation.
• round_constants are the constants a, b, c from the AddRoundConstant operation.
• flag_start is a flag that indicates the start of a new hash invocation.
• flag_round is a flag that indicates that the Poseidon round should be applied.
• flag_final is a flag that indicates that the result is ready.
• inp1 is the first input to the Poseidon hash function.
• inp2 is the second input to the Poseidon hash function.
• rndc a set of fixed columns containing the round constants.
• cols a set of advice columns storing the state of the Poseidon hash function.

As a table we are going to fill it out like this:

Constraining the Poseidon Table

There are two types of constraints that we need to add to the Poseidon table:

1. The flag_start constraint: This resets the state to the input values of the hash function.
2. The flag_round constraint: This applies the Poseidon round to the state.

Let's take a look at each in turn:

The Start Constraint

The start constraint in its totality is as follows:

        meta.create_gate("start", |meta| {
            let flag_start = meta.query_fixed(flag_start, Rotation::cur());
            let inp1 = meta.query_advice(inp1, Rotation::cur());
            let inp2 = meta.query_advice(inp2, Rotation::cur());
            let col1 = meta.query_advice(cols[0], Rotation::cur());
            let col2 = meta.query_advice(cols[1], Rotation::cur());
            let col3 = meta.query_advice(cols[2], Rotation::cur());
            vec![
                flag_start.clone() * (inp1 - col1), // col1 = inp1
                flag_start.clone() * (inp2 - col2), // col2 = inp2
                flag_start.clone() * col3,          // col3 = 0
            ]
        });

What this enforces is:

• inp1.cur() = col[0].cur() : load the first input into the first element of the state.
• inp2.cur() = col[1].cur() : load the second input into the second element of the state.
• col[2].cur() = 0 : set the third element of the state to 0.

These constraints are enforced when flag_start is set to 1. It corresponds to a row of this form:

The Round Constraint

The round constraint is substantially more complex than the start constraint. It is likely the most complex gate you have encountered so far. It applies an entire round of the Poseidon hash function to the state, including the addition of the round constants, the SBox, and the full matrix operation.

So let's break it down into parts.

We start by "reading" the cells of the current row and the next:

        meta.create_gate("round", |meta| {
            let flag_round = meta.query_fixed(flag_round, Rotation::cur());

            let rndc = [
                meta.query_fixed(rndc[0], Rotation::cur()),
                meta.query_fixed(rndc[1], Rotation::cur()),
                meta.query_fixed(rndc[2], Rotation::cur()),
            ];

            let cols_cur = [
                meta.query_advice(cols[0], Rotation::cur()),
                meta.query_advice(cols[1], Rotation::cur()),
                meta.query_advice(cols[2], Rotation::cur()),
            ];

            let cols_nxt = [
                meta.query_advice(cols[0], Rotation::next()),
                meta.query_advice(cols[1], Rotation::next()),
                meta.query_advice(cols[2], Rotation::next()),
            ];

            let inp_cur = [
                meta.query_advice(inp1, Rotation::cur()),
                meta.query_advice(inp2, Rotation::cur()),
            ];

            let inp_nxt = [
                meta.query_advice(inp1, Rotation::next()),
                meta.query_advice(inp2, Rotation::next()),
            ];

We then add the round constants to the state:

            // add round constants
            let cols_arc = [
                cols_cur[0].clone() + rndc[0].clone(), //
                cols_cur[1].clone() + rndc[1].clone(), //
                cols_cur[2].clone() + rndc[2].clone(), //
            ];

Note that this results in an array of Expressions: in what follows we are essentially composing constraints as if we were applying functions to concrete values. We now apply the SBox to the elements of the state:

            // apply sbox: this is pretty inefficient, the degree of the gate is high
            assert_eq!(POWER, 5);

            fn sbox<F: Field>(x: Expression<F>) -> Expression<F> {
                x.clone() * x.clone() * x.clone() * x.clone() * x.clone()
            }

            let cols_sbox = [
                sbox(cols_arc[0].clone()),
                sbox(cols_arc[1].clone()),
                sbox(cols_arc[2].clone()),
            ];

Finally, we apply the matrix operation to the state (consisting of Expressions):

            // apply matrix
            let cols_mat: [Expression<F>; WIDTH] = [
                Expression::Constant(F::ZERO)
                    + cols_sbox[0].clone() * matrix[0][0]
                    + cols_sbox[1].clone() * matrix[0][1]
                    + cols_sbox[2].clone() * matrix[0][2],
                Expression::Constant(F::ZERO)
                    + cols_sbox[0].clone() * matrix[1][0]
                    + cols_sbox[1].clone() * matrix[1][1]
                    + cols_sbox[2].clone() * matrix[1][2],
                Expression::Constant(F::ZERO)
                    + cols_sbox[0].clone() * matrix[2][0]
                    + cols_sbox[1].clone() * matrix[2][1]
                    + cols_sbox[2].clone() * matrix[2][2],
            ];

Finally we enforce that the next row is the result of applying this complex transformation to the current row:

            // enforce that the next state is the round applied to the current state.
            vec![
                // round application
                flag_round.clone() * (cols_mat[0].clone() - cols_nxt[0].clone()), // inp1 = col1
                flag_round.clone() * (cols_mat[1].clone() - cols_nxt[1].clone()), // inp2 = col2
                flag_round.clone() * (cols_mat[2].clone() - cols_nxt[2].clone()), // 0 = col3
                // maintain input
                flag_round.clone() * (inp_cur[0].clone() - inp_nxt[0].clone()), // inp1 = inp1
                flag_round.clone() * (inp_cur[1].clone() - inp_nxt[1].clone()), // inp2 = inp2
            ]
        });

Overall, this corresponds to a row of the form:

Where:

$$(x', y', z') = \text{PoseidonRound}(\mathsf{st} = (x, y, z), \mathsf{rndc} = (a1, b1, c1)) $$

Filling in the Poseidon Table

To aid in the construction of the Poseidon table, we can define a simple helper function that fills in a single row:

    fn assign_row(
        &self,
        idx: usize,
        reg: &mut Region<'_, F>,
        flag_start: bool,
        flag_round: bool,
        flag_final: bool,
        rndc: [F; 3],
        cols: [F; 3],
        inp: [F; 2],
    ) -> Result<(), plonk::Error> {
        reg.assign_fixed(
            || "flag_start",
            self.flag_start,
            idx,
            || Value::known(if flag_start { F::ONE } else { F::ZERO }),
        )?;
        reg.assign_fixed(
            || "flag_round",
            self.flag_round,
            idx,
            || Value::known(if flag_round { F::ONE } else { F::ZERO }),
        )?;
        reg.assign_fixed(
            || "flag_final",
            self.flag_final,
            idx,
            || Value::known(if flag_final { F::ONE } else { F::ZERO }),
        )?;
        reg.assign_fixed(|| "rndc0", self.rndc[0], idx, || Value::known(rndc[0]))?;
        reg.assign_fixed(|| "rndc1", self.rndc[1], idx, || Value::known(rndc[1]))?;
        reg.assign_fixed(|| "rndc2", self.rndc[2], idx, || Value::known(rndc[2]))?;
        reg.assign_advice(|| "cols", self.cols[0], idx, || Value::known(cols[0]))?;
        reg.assign_advice(|| "cols", self.cols[1], idx, || Value::known(cols[1]))?;
        reg.assign_advice(|| "cols", self.cols[2], idx, || Value::known(cols[2]))?;
        reg.assign_advice(|| "inp1", self.inp1, idx, || Value::known(inp[0]))?;
        reg.assign_advice(|| "inp2", self.inp2, idx, || Value::known(inp[1]))?;
        Ok(())
    }

To assign the entire table we create a function which takes a list of pairs of elements to be hashed and fills the table accordingly. All it does is hash each pair, round-by-round and fill in the rows of the table sequentially. There is one special row however: the first row is the all-zero row (all the flags are set to 0 as well), this is to enable the lookup of a "dummy" value whenever the Poseidon gate is disabled.

    fn populate(
        &self,
        layouter: &mut impl Layouter<F>,
        inputs: Vec<(F, F)>,
    ) -> Result<(), plonk::Error> {
        // ensure padded
        assert_eq!(inputs.len(), MAX_OPS_POSEIDON);

        // assign poseidon table
        layouter.assign_region(
            || "poseidon",
            |mut reg| {
                let mut st = [F::ZERO; WIDTH];
                let mut inp = [F::ZERO; 2];
                let mut nxt = 0;

                // zero row
                {
                    self.assign_row(
                        nxt,
                        &mut reg,
                        false,
                        false,
                        false,
                        [F::ZERO, F::ZERO, F::ZERO],
                        [F::ZERO, F::ZERO, F::ZERO],
                        [F::ZERO, F::ZERO],
                    )?;
                    nxt += 1;
                }

                for op in 0..MAX_OPS_POSEIDON {
                    // apply rounds
                    for r in 0..ROUNDS {
                        // load input
                        if r == 0 {
                            inp = [inputs[op].0, inputs[op].1];
                            st[0] = inp[0];
                            st[1] = inp[1];
                            st[2] = F::ZERO;
                        }

                        self.assign_row(
                            nxt,
                            &mut reg,
                            r == 0,
                            r > 0,
                            false,
                            self.round_constants[r],
                            st,
                            inp,
                        )?;

                        // apply poseidon round (out of circuit)
                        st = poseidon_round(&self.matrix, &self.round_constants[r], st);

                        // next row
                        nxt += 1;
                    }

                    // output
                    self.assign_row(
                        nxt,
                        &mut reg,
                        false,
                        false,
                        true,
                        [F::ZERO, F::ZERO, F::ZERO],
                        st,
                        inp,
                    )?;
                    nxt += 1;
                }
                Ok(())
            },
        )?;

        Ok(())
    }

Note that when generating the verification key, this function will be run on junk values: the inputs are completely arbitrary and the state of the Poseidon hash function does not matter. The flags and round constants are fixed however.

The Poseidon Chip

At this point we have a table, guaranteed to contain correct invocations of the Poseidon hash function. Now we need to create a chip that can look up the entries (input/output pairs) in the table dynamically so we can "use the hash function" in a circuit.

Configuration

Towards this end, we define a chip responsible for looking up the Poseidon table:

#[derive(Clone, Debug)]
pub struct PoseidonChip<F: Field> {
    inputs: RefCell<Vec<(F, F)>>,
    sel: Selector,
    tbl: PoseidonTable<F>,
    in1: Column<Advice>,
    in2: Column<Advice>,
    out: Column<Advice>,
    on: Column<Advice>,
}

The fields are mostly self-explanatory, but here is a brief overview:

• inputs is a simple way for us to collect all the inputs to the Poseidon hash function we encounter during witness generation. Whenever we are asked to hash a pair of values \( (x, y) \), we simply hash them out-of-circuit \( \mathsf{Hash}(x, y) \) then we add them to this list.

• sel is just a selector to turn on this chip at the current offset.

• tbl is the Poseidon table we constructed earlier.

• in1 and in2 are the inputs to the Poseidon hash function.

• out is the output of the Poseidon hash function.

• on is the flag that determines whether the Poseidon chip is enabled or not. Unlike a selector, which is constant, this can be toggled on and off dynamically (it's an Advice column).

The gates of the Poseidon chip are where the magic happens:

impl<F: Field> PoseidonChip<F> {
    fn configure(meta: &mut ConstraintSystem<F>) -> Self {
        let sel = meta.complex_selector();
        let in1 = meta.advice_column();
        let in2 = meta.advice_column();
        let out = meta.advice_column();
        let on = meta.advice_column();
        let tbl = PoseidonTable::new(meta);

        meta.enable_equality(in1);
        meta.enable_equality(in2);
        meta.enable_equality(out);
        meta.enable_equality(on);

        meta.create_gate("bit", |meta| {
            let on = meta.query_advice(on, Rotation::cur());
            let sel = meta.query_selector(sel);
            vec![
                sel * on.clone() * (on.clone() - Expression::Constant(F::ONE)), //
            ]
        });

        meta.lookup_any("poseidon_lookup", |cells| {
            let on = cells.query_advice(on, Rotation::cur());
            let sel = cells.query_selector(sel);
            let in1 = cells.query_advice(in1, Rotation::cur());
            let in2 = cells.query_advice(in2, Rotation::cur());
            let out = cells.query_advice(out, Rotation::cur());

            let do_lookup = on * sel;

            let table = tbl.table_expr(cells);

            // (1, in1, in2, out) in PoseidonTable
            vec![
                (do_lookup.clone() * Expression::Constant(F::ONE), table.flag),
                (do_lookup.clone() * in1.clone(), table.inp1),
                (do_lookup.clone() * in2.clone(), table.inp2),
                (do_lookup.clone() * out.clone(), table.out),
            ]
        });

        Self {
            sel,
            tbl,
            in1,
            in2,
            out,
            on,
            inputs: RefCell::new(Vec::new()),
        }
    }

The first constraint is pretty simple: it simply enforces that on is a boolean value whenever sel is set to 1. The poseidon_lookup "lookup_any" constraint is where the actual lookup happens:

1. We read on at the current offset.
2. We read sel at the current offset.
3. We read in1 at the current offset.
4. We read in2 at the current offset.
5. We read out at the current offset.

We then define do_lookup = sel * on, which means that:

$$ \text{do_lookup} = 1 \iff \text{sel} = 1 \land \text{on} = 1 $$

This dynamic value will be used to turn the lookup "on" and "off".

In order for us to access the Poseidon table, we need to gain access to its columns, to do this we follow the approach of PSE and create a little helper which reads the column expressions. This is done by let table = tbl.table_expr(cells); and inside table_expr we simply query the columns at the current offset and return the resulting expressions:

    fn table_expr(&self, meta: &mut VirtualCells<F>) -> PoseidonExprs<F> {
        PoseidonExprs {
            flag: meta.query_any(self.flag_final, Rotation::cur()),
            inp1: meta.query_any(self.inp1, Rotation::cur()),
            inp2: meta.query_any(self.inp2, Rotation::cur()),
            out: meta.query_any(self.cols[0], Rotation::cur()),
        }
    }

With that out of the way, we can now define the actual lookup: when do_lookup = 1 the lookup requires that the row (1, inp1, inp2, out) is in the columns (flag_final, inp1, inp2, col1) of the Poseidon table. In other words: the inputs/output pair is from a row which has been marked by flag_final = 1 i.e. for which the Poseidon hash function has been completely evaluated.

Synthesizing

What remains is to use our new Poseidon chip in a circuit. To do this we define a function which creates a new region and assigns the output appropriately:

    fn hash(
        &self,
        layouter: &mut impl Layouter<F>,
        on: AssignedCell<F, F>,
        in1: AssignedCell<F, F>,
        in2: AssignedCell<F, F>,
    ) -> Result<AssignedCell<F, F>, plonk::Error> {
        // store inputs
        in1.value().and_then(|in1| {
            in2.value()
                .map(|in2| self.inputs.borrow_mut().push((*in1, *in2)))
        });

        layouter.assign_region(
            || "poseidon",
            |mut reg| {
                self.sel.enable(&mut reg, 0)?;

                on.copy_advice(|| "on", &mut reg, self.on, 0)?;
                in1.copy_advice(|| "in1", &mut reg, self.in1, 0)?;
                in2.copy_advice(|| "in2", &mut reg, self.in2, 0)?;

                let hsh = in1
                    .value()
                    .and_then(|in1| in2.value().map(|in2| self.tbl.hash(*in1, *in2)));

                // if on = 0, hsh = 0
                let hsh = on.value().and_then(|on| hsh.map(|hsh| hsh * on));

                let out = reg.assign_advice(|| "out", self.out, 0, || hsh)?;
                Ok(out)
            },
        )
    }

The hash function stores the inputs into inputs, computes the Poseidon hash function out-of-circuit, and assigns the inputs/output pairs of the Poseidon chip. Finally, we turn on the selector sel for the Poseidon chip (lookup). This constrains the inputs/output pair to be in the Poseidon table, which will only be the case if the Poseidon hash function was correctly evaluated out-of-circuit.

After the last hash invocation, we need to finalize the Poseidon chip:

    fn finalize(self, layouter: &mut impl Layouter<F>) -> Result<(), plonk::Error> {
        let mut inputs = self.inputs.borrow().clone();
        while inputs.len() < MAX_OPS_POSEIDON {
            inputs.push((F::ZERO, F::ZERO));
        }
        self.tbl.populate(layouter, inputs)
    }

Which simply pads the inputs vector with dummy values to hash, before populating the Poseidon table. The overall use of the Poseidon chip is as follows:

impl<F: Field> Circuit<F> for TestCircuit<F> {
    type Config = TestConfig<F>;
    type FloorPlanner = SimpleFloorPlanner;

    fn without_witnesses(&self) -> Self {
        TestCircuit { _ph: PhantomData }
    }

    #[allow(unused_variables)]
    fn configure(meta: &mut ConstraintSystem<F>) -> Self::Config {
        let poseidon = PoseidonChip::configure(meta);
        let free = meta.advice_column();
        meta.enable_equality(free);
        TestConfig { poseidon, free }
    }

    #[allow(unused_variables)]
    fn synthesize(
        &self,
        config: Self::Config,
        mut layouter: impl Layouter<F>,
    ) -> Result<(), plonk::Error> {
        let hashes = vec![(F::ZERO, F::ZERO); MAX_OPS_POSEIDON];

        let in1 = layouter.assign_region(
            || "free1",
            |mut region| {
                region.assign_advice(
                    || "free", //
                    config.free,
                    0,
                    || Value::known(F::ONE),
                )
            },
        )?;

        let in2 = layouter.assign_region(
            || "free2",
            |mut region| {
                region.assign_advice(
                    || "free", //
                    config.free,
                    0,
                    || Value::known(F::ONE),
                )
            },
        )?;

        let on = layouter.assign_region(
            || "free3",
            |mut region| {
                region.assign_advice(
                    || "free", //
                    config.free,
                    0,
                    || Value::known(F::ONE),
                )
            },
        )?;

        // populate poseidon
        let out = config.poseidon.hash(&mut layouter, on, in1, in2)?;
        println!("hash done: {:?}", out);
        config.poseidon.finalize(&mut layouter)?;
        Ok(())
    }
}

Which should be familiar by now:

• We create three free variables in1, in2, and on.
• We hash the inputs using the Poseidon chip.
• We finalize the Poseidon chip.

Of course in a real application, the variables in1, in2, and on would likely be constrained by other parts of the circuit.

Exercises