From 3fcc0657760096f755456725b4c66585261be6c5 Mon Sep 17 00:00:00 2001
From: Rute Figueiredo <rfigueiredo.dev@gmail.com>
Date: Tue, 27 Aug 2024 09:52:41 +0100
Subject: [PATCH 1/3] transform_meta to be more generic with traits

---
 src/poly/mod.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/poly/mod.rs b/src/poly/mod.rs
index b7c026ff..0d69af78 100644
--- a/src/poly/mod.rs
+++ b/src/poly/mod.rs
@@ -91,14 +91,14 @@ impl<F: Clone, V: Clone, M> Expr<F, V, M> {
     }
 }
 
-impl<F: Field + Hash, V: Debug + Clone + Eq + Hash, M: Clone> Expr<F, V, M> {
+impl<F: Clone, V: Debug + Clone + Eq, M: Clone> Expr<F, V, M> {
     pub fn transform_meta<N: Clone, ApplyMetaFn>(&self, apply_meta: ApplyMetaFn) -> Expr<F, V, N>
     where
         ApplyMetaFn: Fn(&Expr<F, V, M>) -> N + Clone,
     {
         let new_meta = apply_meta(self);
         match self {
-            Expr::Const(v, _) => Expr::Const(*v, new_meta),
+            Expr::Const(v, _) => Expr::Const(v.clone(), new_meta),
             Expr::Sum(ses, _) => Expr::Sum(
                 ses.iter()
                     .map(|e| e.transform_meta(apply_meta.clone()))

From 5ea705f91b9db945a4dceba42367d322d2242b70 Mon Sep 17 00:00:00 2001
From: Rute Figueiredo <rfigueiredo.dev@gmail.com>
Date: Tue, 27 Aug 2024 12:28:52 +0100
Subject: [PATCH 2/3] added debug symbol to the setup interpreter

---
 src/compiler/setup_inter.rs | 49 ++++++++++++++++++++++---------------
 1 file changed, 29 insertions(+), 20 deletions(-)

diff --git a/src/compiler/setup_inter.rs b/src/compiler/setup_inter.rs
index 9b2b3e24..36e11735 100644
--- a/src/compiler/setup_inter.rs
+++ b/src/compiler/setup_inter.rs
@@ -32,7 +32,7 @@ pub(super) fn interpret(ast: &[TLDecl<BigInt, Identifier>], _symbols: &SymTable)
 pub(super) type Setup<F> = HashMap<String, MachineSetup<F>>;
 
 pub(super) struct MachineSetup<F> {
-    poly_constraints: HashMap<String, Vec<Expr<F, Identifier, ()>>>,
+    poly_constraints: HashMap<String, Vec<Expr<F, Identifier, DebugSymRef>>>,
 
     input_signals: Vec<TypedIdDecl<Identifier>>,
     output_signals: Vec<TypedIdDecl<Identifier>>,
@@ -49,10 +49,10 @@ impl<F> Default for MachineSetup<F> {
 }
 impl MachineSetup<BigInt> {
     pub(crate) fn map_consts<F: Field>(&self) -> MachineSetup<F> {
-        let poly_constraints: HashMap<String, Vec<Expr<F, Identifier, ()>>> = self
+        let poly_constraints: HashMap<String, Vec<Expr<F, Identifier, DebugSymRef>>> = self
             .poly_constraints_iter()
             .map(|(step_id, step)| {
-                let new_step: Vec<Expr<F, Identifier, ()>> = step
+                let new_step: Vec<Expr<F, Identifier, DebugSymRef>> = step
                     .iter()
                     .map(|pi| Self::convert_const_to_field(pi))
                     .collect();
@@ -66,28 +66,32 @@ impl MachineSetup<BigInt> {
     }
 
     fn convert_const_to_field<F: Field>(
-        expr: &Expr<BigInt, Identifier, ()>,
-    ) -> Expr<F, Identifier, ()> {
+        expr: &Expr<BigInt, Identifier, DebugSymRef>,
+    ) -> Expr<F, Identifier, DebugSymRef> {
         use Expr::*;
         match expr {
-            Const(v, _) => Const(F::from_big_int(v), ()),
-            Sum(ses, _) => Sum(
+            Const(v, dsym) => Const(F::from_big_int(v), dsym.clone()),
+            Sum(ses, dsym) => Sum(
                 ses.iter()
                     .map(|se| Self::convert_const_to_field(se))
                     .collect(),
-                (),
+                dsym.clone(),
             ),
-            Mul(ses, _) => Mul(
+            Mul(ses, dsym) => Mul(
                 ses.iter()
                     .map(|se| Self::convert_const_to_field(se))
                     .collect(),
-                (),
+                dsym.clone(),
             ),
-            Neg(se, _) => Neg(Box::new(Self::convert_const_to_field(se)), ()),
-            Pow(se, exp, _) => Pow(Box::new(Self::convert_const_to_field(se)), *exp, ()),
-            Query(q, _) => Query(q.clone(), ()),
+            Neg(se, dsym) => Neg(Box::new(Self::convert_const_to_field(se)), dsym.clone()),
+            Pow(se, exp, dsym) => Pow(
+                Box::new(Self::convert_const_to_field(se)),
+                *exp,
+                dsym.clone(),
+            ),
+            Query(q, dsym) => Query(q.clone(), dsym.clone()),
             Halo2Expr(_, _) => todo!(),
-            MI(se, _) => MI(Box::new(Self::convert_const_to_field(se)), ()),
+            MI(se, dsym) => MI(Box::new(Self::convert_const_to_field(se)), dsym.clone()),
         }
     }
 }
@@ -125,7 +129,7 @@ impl<F: Clone> MachineSetup<F> {
     fn add_poly_constraints<S: Into<String>>(
         &mut self,
         state: S,
-        poly_constraints: Vec<Expr<F, Identifier, ()>>,
+        poly_constraints: Vec<Expr<F, Identifier, DebugSymRef>>,
     ) {
         self.poly_constraints
             .get_mut(&state.into())
@@ -135,13 +139,13 @@ impl<F: Clone> MachineSetup<F> {
 
     pub(super) fn poly_constraints_iter(
         &self,
-    ) -> std::collections::hash_map::Iter<String, Vec<Expr<F, Identifier, ()>>> {
+    ) -> std::collections::hash_map::Iter<String, Vec<Expr<F, Identifier, DebugSymRef>>> {
         self.poly_constraints.iter()
     }
 
     pub(super) fn replace_poly_constraints<F2>(
         &self,
-        poly_constraints: HashMap<String, Vec<Expr<F2, Identifier, ()>>>,
+        poly_constraints: HashMap<String, Vec<Expr<F2, Identifier, DebugSymRef>>>,
     ) -> MachineSetup<F2> {
         MachineSetup {
             poly_constraints,
@@ -157,7 +161,7 @@ impl<F: Clone> MachineSetup<F> {
     pub(super) fn get_poly_constraints<S: Into<String>>(
         &self,
         state: S,
-    ) -> Option<&Vec<Expr<F, Identifier, ()>>> {
+    ) -> Option<&Vec<Expr<F, Identifier, DebugSymRef>>> {
         self.poly_constraints.get(&state.into())
     }
 }
@@ -252,10 +256,15 @@ impl SetupInterpreter {
             SignalDecl(_, _) | WGVarDecl(_, _) => vec![],
         };
 
-        self.add_poly_constraints(result.into_iter().map(|cr| cr.anti_booly).collect());
+        self.add_poly_constraints(
+            result
+                .into_iter()
+                .map(|cr| cr.anti_booly.transform_meta(|_| cr.dsym.clone()))
+                .collect(),
+        );
     }
 
-    fn add_poly_constraints(&mut self, pis: Vec<Expr<BigInt, Identifier, ()>>) {
+    fn add_poly_constraints(&mut self, pis: Vec<Expr<BigInt, Identifier, DebugSymRef>>) {
         self.setup
             .get_mut(&self.current_machine)
             .unwrap()

From dd7a3eb6d8681a05ab08810cef80e407949b375a Mon Sep 17 00:00:00 2001
From: Rute Figueiredo <rfigueiredo.dev@gmail.com>
Date: Thu, 29 Aug 2024 14:19:40 +0100
Subject: [PATCH 3/3] added the debug symbol for the Expr nodes on the SBPIR
 while we are compiling

---
 src/compiler/compiler.rs        | 47 ++++++++++----------
 src/compiler/compiler_legacy.rs | 78 +++++++++++++++++++--------------
 src/sbpir/mod.rs                | 39 +++++++++++++----
 src/sbpir/sbpir_machine.rs      | 42 +++++++++++++++++-
 4 files changed, 141 insertions(+), 65 deletions(-)

diff --git a/src/compiler/compiler.rs b/src/compiler/compiler.rs
index 9a3c5584..470ba80b 100644
--- a/src/compiler/compiler.rs
+++ b/src/compiler/compiler.rs
@@ -4,6 +4,7 @@ use num_bigint::BigInt;
 
 use crate::{
     field::Field,
+    frontend::dsl::trace::DSLTraceGenerator,
     interpreter::InterpreterTraceGenerator,
     parser::{
         ast::{
@@ -269,6 +270,8 @@ impl<F: Field + Hash> Compiler<F> {
             sbpir.machines.insert(machine_name.clone(), sbpir_machine);
         }
 
+        let sbpir = sbpir.transform_metadata(|_| ());
+
         sbpir.without_trace()
     }
 
@@ -279,7 +282,7 @@ impl<F: Field + Hash> Compiler<F> {
         setup: &MachineSetup<F>,
         machine_name: &str,
         state_id: &str,
-    ) -> Vec<Constraint<F, ()>> {
+    ) -> Vec<Constraint<F, DebugSymRef>> {
         let exprs = setup.get_poly_constraints(state_id).unwrap();
 
         exprs
@@ -300,40 +303,40 @@ impl<F: Field + Hash> Compiler<F> {
         symbols: &SymTable,
         machine_name: &str,
         state_id: &str,
-        expr: &Expr<F, Identifier, ()>,
-    ) -> Expr<F, Queriable<F>, ()> {
+        expr: &Expr<F, Identifier, DebugSymRef>,
+    ) -> Expr<F, Queriable<F>, DebugSymRef> {
         use Expr::*;
         match expr {
-            Const(v, _) => Const(*v, ()),
-            Sum(ses, _) => Sum(
+            Const(v, dsym) => Const(*v, dsym.clone()),
+            Sum(ses, dsym) => Sum(
                 ses.iter()
                     .map(|se| self.translate_queries_expr(symbols, machine_name, state_id, se))
                     .collect(),
-                (),
+                dsym.clone(),
             ),
-            Mul(ses, _) => Mul(
+            Mul(ses, dsym) => Mul(
                 ses.iter()
                     .map(|se| self.translate_queries_expr(symbols, machine_name, state_id, se))
                     .collect(),
-                (),
+                dsym.clone(),
             ),
-            Neg(se, _) => Neg(
+            Neg(se, dsym) => Neg(
                 Box::new(self.translate_queries_expr(symbols, machine_name, state_id, se.as_ref())),
-                (),
+                dsym.clone(),
             ),
-            Pow(se, exp, _) => Pow(
+            Pow(se, exp, dsym) => Pow(
                 Box::new(self.translate_queries_expr(symbols, machine_name, state_id, se.as_ref())),
                 *exp,
-                (),
+                dsym.clone(),
             ),
-            MI(se, _) => MI(
+            MI(se, dsym) => MI(
                 Box::new(self.translate_queries_expr(symbols, machine_name, state_id, se.as_ref())),
-                (),
+                dsym.clone(),
             ),
-            Halo2Expr(se, _) => Halo2Expr(se.clone(), ()),
-            Query(id, _) => Query(
+            Halo2Expr(se, dsym) => Halo2Expr(se.clone(), dsym.clone()),
+            Query(id, dsym) => Query(
                 self.translate_query(symbols, machine_name, state_id, id),
-                (),
+                dsym.clone(),
             ),
         }
     }
@@ -428,10 +431,10 @@ impl<F: Field + Hash> Compiler<F> {
         machine_name: &str,
         machine_setup: &MachineSetup<F>,
         state_id: &str,
-    ) -> StepType<F, ()> {
+    ) -> StepType<F, DebugSymRef> {
         let handler = self.mapping.get_step_type_handler(machine_name, state_id);
 
-        let mut step_type: StepType<F, ()> =
+        let mut step_type: StepType<F, DebugSymRef> =
             StepType::new(handler.uuid(), handler.annotation.to_string());
 
         self.add_internal_signals(symbols, machine_name, &mut step_type, state_id);
@@ -448,7 +451,7 @@ impl<F: Field + Hash> Compiler<F> {
         &mut self,
         symbols: &SymTable,
         machine_name: &str,
-        step_type: &mut StepType<F, ()>,
+        step_type: &mut StepType<F, DebugSymRef>,
         state_id: &str,
     ) {
         let internal_ids = self.get_all_internals(symbols, machine_name, state_id);
@@ -467,7 +470,7 @@ impl<F: Field + Hash> Compiler<F> {
 
     fn add_step_type_handlers(
         &mut self,
-        machine: &mut SBPIRMachine<F>,
+        machine: &mut SBPIRMachine<F, DSLTraceGenerator<F>, DebugSymRef>,
         symbols: &SymTable,
         machine_name: &str,
     ) {
@@ -497,7 +500,7 @@ impl<F: Field + Hash> Compiler<F> {
 
     fn add_forward_signals(
         &mut self,
-        machine: &mut SBPIRMachine<F>,
+        machine: &mut SBPIRMachine<F, DSLTraceGenerator<F>, DebugSymRef>,
         symbols: &SymTable,
         machine_name: &str,
     ) {
diff --git a/src/compiler/compiler_legacy.rs b/src/compiler/compiler_legacy.rs
index c7bed275..01ed7e21 100644
--- a/src/compiler/compiler_legacy.rs
+++ b/src/compiler/compiler_legacy.rs
@@ -259,15 +259,16 @@ impl<F: Field + Hash> CompilerLegacy<F> {
         setup
             .iter()
             .map(|(machine_id, machine)| {
-                let poly_constraints: HashMap<String, Vec<Expr<F, Identifier, ()>>> = machine
-                    .poly_constraints_iter()
-                    .map(|(step_id, step)| {
-                        let new_step: Vec<Expr<F, Identifier, ()>> =
-                            step.iter().map(|pi| Self::map_pi_consts(pi)).collect();
+                let poly_constraints: HashMap<String, Vec<Expr<F, Identifier, DebugSymRef>>> =
+                    machine
+                        .poly_constraints_iter()
+                        .map(|(step_id, step)| {
+                            let new_step: Vec<Expr<F, Identifier, DebugSymRef>> =
+                                step.iter().map(|pi| Self::map_pi_consts(pi)).collect();
 
-                        (step_id.clone(), new_step)
-                    })
-                    .collect();
+                            (step_id.clone(), new_step)
+                        })
+                        .collect();
 
                 let new_machine: MachineSetup<F> =
                     machine.replace_poly_constraints(poly_constraints);
@@ -276,17 +277,25 @@ impl<F: Field + Hash> CompilerLegacy<F> {
             .collect()
     }
 
-    fn map_pi_consts(expr: &Expr<BigInt, Identifier, ()>) -> Expr<F, Identifier, ()> {
+    fn map_pi_consts(
+        expr: &Expr<BigInt, Identifier, DebugSymRef>,
+    ) -> Expr<F, Identifier, DebugSymRef> {
         use Expr::*;
         match expr {
-            Const(v, _) => Const(F::from_big_int(v), ()),
-            Sum(ses, _) => Sum(ses.iter().map(|se| Self::map_pi_consts(se)).collect(), ()),
-            Mul(ses, _) => Mul(ses.iter().map(|se| Self::map_pi_consts(se)).collect(), ()),
-            Neg(se, _) => Neg(Box::new(Self::map_pi_consts(se)), ()),
-            Pow(se, exp, _) => Pow(Box::new(Self::map_pi_consts(se)), *exp, ()),
-            Query(q, _) => Query(q.clone(), ()),
+            Const(v, dsym) => Const(F::from_big_int(v), dsym.clone()),
+            Sum(ses, dsym) => Sum(
+                ses.iter().map(|se| Self::map_pi_consts(se)).collect(),
+                dsym.clone(),
+            ),
+            Mul(ses, dsym) => Mul(
+                ses.iter().map(|se| Self::map_pi_consts(se)).collect(),
+                dsym.clone(),
+            ),
+            Neg(se, dsym) => Neg(Box::new(Self::map_pi_consts(se)), dsym.clone()),
+            Pow(se, exp, dsym) => Pow(Box::new(Self::map_pi_consts(se)), *exp, dsym.clone()),
+            Query(q, dsym) => Query(q.clone(), dsym.clone()),
             Halo2Expr(_, _) => todo!(),
-            MI(se, _) => MI(Box::new(Self::map_pi_consts(se)), ()),
+            MI(se, dsym) => MI(Box::new(Self::map_pi_consts(se)), dsym.clone()),
         }
     }
 
@@ -316,7 +325,7 @@ impl<F: Field + Hash> CompilerLegacy<F> {
                                 poly_constraints.iter().for_each(|poly| {
                                     let constraint = Constraint {
                                         annotation: format!("{:?}", poly),
-                                        expr: poly.clone(),
+                                        expr: poly.transform_meta(|_| ()),
                                         typing: Typing::AntiBooly,
                                     };
                                     ctx.constr(constraint);
@@ -372,7 +381,7 @@ impl<F: Field + Hash> CompilerLegacy<F> {
         setup: &Setup<F>,
         machine_id: &str,
         state_id: &str,
-    ) -> Vec<Expr<F, Queriable<F>, ()>> {
+    ) -> Vec<Expr<F, Queriable<F>, DebugSymRef>> {
         let exprs = setup
             .get(machine_id)
             .unwrap()
@@ -390,38 +399,41 @@ impl<F: Field + Hash> CompilerLegacy<F> {
         symbols: &SymTable,
         machine_id: &str,
         state_id: &str,
-        expr: &Expr<F, Identifier, ()>,
-    ) -> Expr<F, Queriable<F>, ()> {
+        expr: &Expr<F, Identifier, DebugSymRef>,
+    ) -> Expr<F, Queriable<F>, DebugSymRef> {
         use Expr::*;
         match expr {
-            Const(v, _) => Const(*v, ()),
-            Sum(ses, _) => Sum(
+            Const(v, dsym) => Const(*v, dsym.clone()),
+            Sum(ses, dsym) => Sum(
                 ses.iter()
                     .map(|se| self.translate_queries_expr(symbols, machine_id, state_id, se))
                     .collect(),
-                (),
+                dsym.clone(),
             ),
-            Mul(ses, _) => Mul(
+            Mul(ses, dsym) => Mul(
                 ses.iter()
                     .map(|se| self.translate_queries_expr(symbols, machine_id, state_id, se))
                     .collect(),
-                (),
+                dsym.clone(),
             ),
-            Neg(se, _) => Neg(
+            Neg(se, dsym) => Neg(
                 Box::new(self.translate_queries_expr(symbols, machine_id, state_id, se.as_ref())),
-                (),
+                dsym.clone(),
             ),
-            Pow(se, exp, _) => Pow(
+            Pow(se, exp, dsym) => Pow(
                 Box::new(self.translate_queries_expr(symbols, machine_id, state_id, se.as_ref())),
                 *exp,
-                (),
+                dsym.clone(),
             ),
-            MI(se, _) => MI(
+            MI(se, dsym) => MI(
                 Box::new(self.translate_queries_expr(symbols, machine_id, state_id, se.as_ref())),
-                (),
+                dsym.clone(),
+            ),
+            Halo2Expr(se, dsym) => Halo2Expr(se.clone(), dsym.clone()),
+            Query(id, dsym) => Query(
+                self.translate_query(symbols, machine_id, state_id, id),
+                dsym.clone(),
             ),
-            Halo2Expr(se, _) => Halo2Expr(se.clone(), ()),
-            Query(id, _) => Query(self.translate_query(symbols, machine_id, state_id, id), ()),
         }
     }
 
diff --git a/src/sbpir/mod.rs b/src/sbpir/mod.rs
index e1f90b74..279cb217 100644
--- a/src/sbpir/mod.rs
+++ b/src/sbpir/mod.rs
@@ -275,8 +275,8 @@ pub struct SBPIR<F: Clone, TG: TraceGenerator<F> = DSLTraceGenerator<F>, M: Clon
     pub identifiers: HashMap<String, UUID>,
 }
 
-impl<F: Field + Hash, TG: TraceGenerator<F>> SBPIR<F, TG> {
-    pub(crate) fn default() -> SBPIR<F, TG> {
+impl<F: Field + Hash, TG: TraceGenerator<F>, M: Clone> SBPIR<F, TG, M> {
+    pub(crate) fn default() -> SBPIR<F, TG, M> {
         let machines = HashMap::new();
         let identifiers = HashMap::new();
         SBPIR {
@@ -289,7 +289,7 @@ impl<F: Field + Hash, TG: TraceGenerator<F>> SBPIR<F, TG> {
         &self,
         // TODO does it have to be the same trace across all the machines?
         trace: &TG2,
-    ) -> SBPIR<F, TG2> {
+    ) -> SBPIR<F, TG2, M> {
         let mut machines_with_trace = HashMap::new();
         for (name, machine) in self.machines.iter() {
             let machine_with_trace = machine.with_trace(trace.clone());
@@ -301,7 +301,7 @@ impl<F: Field + Hash, TG: TraceGenerator<F>> SBPIR<F, TG> {
         }
     }
 
-    pub(crate) fn without_trace(&self) -> SBPIR<F, NullTraceGenerator> {
+    pub(crate) fn without_trace(&self) -> SBPIR<F, NullTraceGenerator, M> {
         let mut machines_without_trace = HashMap::new();
         for (name, machine) in self.machines.iter() {
             let machine_without_trace = machine.without_trace();
@@ -312,7 +312,28 @@ impl<F: Field + Hash, TG: TraceGenerator<F>> SBPIR<F, TG> {
             identifiers: self.identifiers.clone(),
         }
     }
+}
 
+impl<F: Field + Hash, TG: TraceGenerator<F> + Clone, M: Clone> SBPIR<F, TG, M> {
+    pub fn transform_metadata<N: Clone, ApplyMetaFn>(
+        self,
+        apply_meta: ApplyMetaFn,
+    ) -> SBPIR<F, TG, N>
+    where
+        ApplyMetaFn: Fn(&Expr<F, Queriable<F>, M>) -> N + Clone,
+    {
+        SBPIR {
+            machines: self
+                .machines
+                .into_iter()
+                .map(|(name, machine)| (name, machine.transform_meta(apply_meta.clone())))
+                .collect(),
+            identifiers: self.identifiers,
+        }
+    }
+}
+
+impl<F: Field + Hash, TG: TraceGenerator<F>> SBPIR<F, TG> {
     /// Eliminate multiplicative inverses
     pub(crate) fn eliminate_mul_inv(mut self) -> SBPIR<F, TG> {
         for machine in self.machines.values_mut() {
@@ -443,7 +464,7 @@ impl<F, M> StepType<F, M> {
     }
 }
 
-impl<F: Field + Hash, M: Clone> StepType<F, M> {
+impl<F: Clone + PartialEq + Eq + Hash, M: Clone> StepType<F, M> {
     pub fn transform_meta<N: Clone, ApplyMetaFn>(&self, apply_meta: ApplyMetaFn) -> StepType<F, N>
     where
         ApplyMetaFn: Fn(&Expr<F, Queriable<F>, M>) -> N + Clone,
@@ -470,7 +491,7 @@ impl<F: Field + Hash, M: Clone> StepType<F, M> {
             auto_signals: self
                 .auto_signals
                 .iter()
-                .map(|(k, v)| (*k, v.transform_meta(apply_meta.clone())))
+                .map(|(k, v)| (k.clone(), v.transform_meta(apply_meta.clone())))
                 .collect(),
             annotations: self.annotations.clone(),
         }
@@ -560,7 +581,7 @@ pub struct Constraint<F, M> {
     pub expr: PIR<F, M>,
 }
 
-impl<F: Field + Hash, M: Clone> Constraint<F, M> {
+impl<F: Clone + PartialEq + Eq, M: Clone> Constraint<F, M> {
     pub fn transform_meta<N: Clone, ApplyMetaFn>(&self, apply_meta: ApplyMetaFn) -> Constraint<F, N>
     where
         ApplyMetaFn: Fn(&Expr<F, Queriable<F>, M>) -> N + Clone,
@@ -579,7 +600,7 @@ pub struct TransitionConstraint<F, M> {
     pub expr: PIR<F, M>,
 }
 
-impl<F: Field + Hash, M: Clone> TransitionConstraint<F, M> {
+impl<F: Clone + PartialEq + Eq, M: Clone> TransitionConstraint<F, M> {
     pub fn transform_meta<N: Clone, ApplyMetaFn>(
         &self,
         apply_meta: ApplyMetaFn,
@@ -676,7 +697,7 @@ impl<F: Debug + Clone, M: Clone + Default> Lookup<F, M> {
     }
 }
 
-impl<F: Field + Hash, M: Clone> Lookup<F, M> {
+impl<F: Clone + PartialEq + Eq, M: Clone> Lookup<F, M> {
     pub fn transform_meta<N: Clone, ApplyMetaFn>(&self, apply_meta: ApplyMetaFn) -> Lookup<F, N>
     where
         ApplyMetaFn: Fn(&Expr<F, Queriable<F>, M>) -> N + Clone,
diff --git a/src/sbpir/sbpir_machine.rs b/src/sbpir/sbpir_machine.rs
index 1b1cc941..cdcc3789 100644
--- a/src/sbpir/sbpir_machine.rs
+++ b/src/sbpir/sbpir_machine.rs
@@ -4,12 +4,13 @@ use crate::{
         trace::{DSLTraceGenerator, TraceContext},
         StepTypeHandler,
     },
+    poly::Expr,
     sbpir::Halo2Column,
     util::{uuid, UUID},
     wit_gen::{FixedAssignment, NullTraceGenerator, TraceGenerator},
 };
 use halo2_proofs::plonk::{Advice, Fixed};
-use std::{collections::HashMap, fmt::Debug, rc::Rc};
+use std::{collections::HashMap, fmt::Debug, hash::Hash, rc::Rc};
 
 use super::{
     query::Queriable, ExposeOffset, FixedSignal, ForwardSignal, ImportedHalo2Advice,
@@ -273,6 +274,45 @@ impl<F: Clone + Field, TG: TraceGenerator<F>> SBPIRMachine<F, TG> {
     }
 }
 
+impl<F: Clone + PartialEq + Eq + Hash, TG: TraceGenerator<F> + Clone, M: Clone>
+    SBPIRMachine<F, TG, M>
+{
+    pub fn transform_meta<N: Clone, ApplyMetaFn>(
+        &self,
+        apply_meta: ApplyMetaFn,
+    ) -> SBPIRMachine<F, TG, N>
+    where
+        ApplyMetaFn: Fn(&Expr<F, Queriable<F>, M>) -> N + Clone,
+    {
+        SBPIRMachine {
+            step_types: self
+                .step_types
+                .iter()
+                .map(|(uuid, step_type)| (*uuid, step_type.transform_meta(apply_meta.clone())))
+                .collect(),
+
+            forward_signals: self.forward_signals.clone(),
+            shared_signals: self.shared_signals.clone(),
+            fixed_signals: self.fixed_signals.clone(),
+            halo2_advice: self.halo2_advice.clone(),
+            halo2_fixed: self.halo2_fixed.clone(),
+            exposed: self.exposed.clone(),
+
+            annotations: self.annotations.clone(),
+
+            trace_generator: self.trace_generator.clone(),
+            fixed_assignments: self.fixed_assignments.clone(),
+
+            first_step: self.first_step,
+            last_step: self.last_step,
+            num_steps: self.num_steps,
+            q_enable: self.q_enable,
+
+            id: self.id,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::wit_gen::NullTraceGenerator;