The Cook–Mertz procedure for Tree Evaluation

The procedure keeps three registers, each holding ℓ elements of a finite field K, as global storage shared by all levels of the recursion. Called on a subtree, a coefficient c and a target register t, it adds c times the value of the subtree (read as 0/1 field elements) to register t and leaves the other two registers unchanged, whatever their contents.

At an internal node with gate f, the procedure runs one round for every unit u of K. A round scales the two other registers a and b by u, adds the children's values v₀ and v₁ to them by recursive calls, subtracts c times the multilinear extension of f at the masked values u • a + v₀ and u • b + v₁ from register t, and then undoes the recursive calls and the scaling. Summed over all units, the masks cancel (FiniteField.sum_units_eval_mul_add), because the extension has degree at most 2ℓ < |K| - 1, and register t has gained c • f v₀ v₁.

Only the recursion depth, the round counter and the stage within a round are local to a level; this is what makes the procedure space-efficient. It is the fan-in-d procedure Cslib.TreeEval.Ary.proc at d = 2, step for step, on the embedded tree: child 0 uses register t + 1 and child 1 register t + 2, and the extensions of a binary gate and of its fan-in-2 reading differ only by the naming of their variables. Its correctness is therefore the case d = 2 of Cslib.TreeEval.Ary.proc_eq. This file proves correctness only; space bounds are separate.

Main definitions

• Cslib.TreeEval.gatePoly: the multilinear extension of an output bit of a gate.
• Cslib.TreeEval.proc: the procedure, parametrized by an enumeration of the units of K.

Main results

• Cslib.TreeEval.Ary.proc_ofBinary: on an embedded binary tree, the fan-in-d procedure at d = 2 is the binary procedure.
• Cslib.TreeEval.proc_eq: the procedure adds c times the value of the tree to the target register and leaves the other registers unchanged.
• Cslib.TreeEval.proc_zero: started from zero registers, it leaves the value of the tree in the target register.
• Cslib.TreeEval.embed_eval_eq_proc: the same, for the canonical enumeration of the units.

References

• J. Cook and I. Mertz, Tree Evaluation is in Space O(log n · log log n), Lemma 13 and Theorem 10.
• O. Goldreich, On the Cook–Mertz Tree Evaluation procedure, Algorithm 1, whose signed recursive calls are followed here.

@[reducible, inline]

abbrev Cslib.TreeEval.State (K : Type u_2) (ℓ : ℕ) : Type u_2

The global storage: three registers of ℓ field elements.

noncomputable def Cslib.TreeEval.gatePoly {K : Type u_1} [Field K] {ℓ : ℕ} (f : Gate ℓ) (i : Fin ℓ) : MvPolynomial (Fin ℓ ⊕ Fin ℓ) K

The multilinear extension of output bit i of a gate, in the bits of both inputs.

def Cslib.TreeEval.scale {K : Type u_1} [Field K] {ℓ : ℕ} (r : Fin 3) (a : K) (S : State K ℓ) : State K ℓ

Multiply register r by a.

noncomputable def Cslib.TreeEval.gateUpdate {K : Type u_1} [Field K] {ℓ : ℕ} (f : Gate ℓ) (c : K) (t : Fin 3) (S : State K ℓ) : State K ℓ

The update of a round at a gate f: subtract c times the gate's extension at registers t + 1 and t + 2 from register t.

noncomputable def Cslib.TreeEval.gateRound {K : Type u_1} [Field K] {ℓ : ℕ} (pl pr : K → Fin 3 → State K ℓ → State K ℓ) (f : Gate ℓ) (c : K) (t : Fin 3) (u : Kˣ) (S : State K ℓ) : State K ℓ

One round of the procedure at a gate f, for the unit u, given procedures pl and pr for the two subtrees. The target is t, and the registers t + 1 and t + 2 hold the masks.

noncomputable def Cslib.TreeEval.proc {K : Type u_1} [Field K] {ℓ : ℕ} (us : List Kˣ) {h : ℕ} : Tree ℓ h → K → Fin 3 → State K ℓ → State K ℓ

The Cook–Mertz procedure: add c times the value of the tree to register t. At a gate it runs one gateRound for every unit in the enumeration us.

theorem Cslib.TreeEval.Ary.proc_ofBinary {K : Type u_1} [Field K] {ℓ : ℕ} (us : List Kˣ) {h : ℕ} (T : TreeEval.Tree ℓ h) : proc us (Tree.ofBinary T) = TreeEval.proc us T

The binary procedure is the fan-in-d procedure at d = 2. On a binary tree embedded with fan-in 2, the fan-in-d procedure is the binary procedure, step for step: it acts on the same three registers (Fin (2 + 1) is Fin 3), calls the children in the same order with the same coefficients and target registers, and performs the same updates and scalings.

theorem Cslib.TreeEval.proc_eq {K : Type u_1} [Field K] {ℓ : ℕ} [Fintype K] {us : List Kˣ} (hnd : us.Nodup) (hmem : ∀ (u : Kˣ), u ∈ us) (hK : 2 * ℓ < Fintype.card K - 1) {h : ℕ} (T : Tree ℓ h) (c : K) (t : Fin 3) (S : State K ℓ) : proc us T c t S = Function.update S t (S t + c • embed T.eval)

Correctness of the Cook–Mertz procedure. If us enumerates the units of K without repetition and 2ℓ < |K| - 1, then the procedure adds c times the value of the tree to the target register and leaves the other registers unchanged, whatever their initial contents.

theorem Cslib.TreeEval.proc_zero {K : Type u_1} [Field K] {ℓ : ℕ} [Fintype K] {us : List Kˣ} (hnd : us.Nodup) (hmem : ∀ (u : Kˣ), u ∈ us) (hK : 2 * ℓ < Fintype.card K - 1) {h : ℕ} (T : Tree ℓ h) (t : Fin 3) : proc us T 1 t 0 t = embed T.eval

Started from zero registers, the procedure leaves the value of the tree in the target register.

theorem Cslib.TreeEval.embed_eval_eq_proc {K : Type u_1} [Field K] {ℓ : ℕ} [Fintype K] [DecidableEq K] (hK : 2 * ℓ < Fintype.card K - 1) {h : ℕ} (T : Tree ℓ h) (t : Fin 3) : embed T.eval = proc Finset.univ.toList T 1 t 0 t

Tree Evaluation by the Cook–Mertz procedure. Run over the units of K in any order with 2ℓ < |K| - 1, from zero registers, the procedure computes the value of the tree into the target register.