From 03d90f30c96bca44fd57f7975d5be01cf25d0942 Mon Sep 17 00:00:00 2001
From: onsen-egg <work-account-root@protonmail.com>
Date: Tue, 23 May 2023 10:12:38 -0400
Subject: [PATCH] Faster multiplicity counting for lookup tables

---
 plonky2/src/plonk/prover.rs | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/plonky2/src/plonk/prover.rs b/plonky2/src/plonk/prover.rs
index df2249df..3748e68b 100644
--- a/plonky2/src/plonk/prover.rs
+++ b/plonky2/src/plonk/prover.rs
@@ -4,6 +4,7 @@ use core::cmp::min;
 use core::mem::swap;
 
 use anyhow::{ensure, Result};
+use hashbrown::HashMap;
 use plonky2_maybe_rayon::*;
 
 use super::circuit_builder::{LookupChallenges, LookupWire};
@@ -58,13 +59,18 @@ pub fn set_lookup_wires<
 
         // Compute multiplicities.
         let mut multiplicities = vec![0; lut_len];
+
+        let table_value_to_idx: HashMap<u16, usize> = common_data.luts[lut_index]
+            .iter()
+            .enumerate()
+            .map(|(i, (inp_target, _))| (inp_target.clone(), i))
+            .collect();
+
         for (inp_target, _) in prover_data.lut_to_lookups[lut_index].iter() {
             let inp_value = pw.get_target(*inp_target);
-            let mut idx = 0;
-            while F::from_canonical_u16(common_data.luts[lut_index][idx].0) != inp_value {
-                idx += 1;
-            }
-            multiplicities[idx] += 1;
+            let idx = table_value_to_idx.get(&u16::try_from(inp_value.to_canonical_u64()).unwrap()).unwrap();
+
+            multiplicities[*idx] += 1;
         }
 
         // Pad the last `LookupGate` with the first entry from the LUT.