From a1eaeaff9c66ce8209fb77c0a4bcd1aa7243262f Mon Sep 17 00:00:00 2001
From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com>
Date: Fri, 5 Apr 2024 17:11:00 -0700
Subject: [PATCH] feat: clear virtual phase0 columns in builder before next
 phase (#295)

* feat: clear virtual phase0 columns in builder before next phase

In prover stage, these columns should never be touched again once they
have been raw assigned. This should save memory before the real phase0
assigned columns are KZG commited (MSM leads to memory spike).

* chore: bump axiom-eth v0.4.3, axiom-query v2.0.17
---
 axiom-eth/Cargo.toml                             |  2 +-
 .../utils/component/circuit/comp_circuit_impl.rs | 16 +++++++++++-----
 axiom-query/Cargo.toml                           |  4 ++--
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/axiom-eth/Cargo.toml b/axiom-eth/Cargo.toml
index a99864e7..ccbb5622 100644
--- a/axiom-eth/Cargo.toml
+++ b/axiom-eth/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "axiom-eth"
-version = "0.4.2"
+version = "0.4.3"
 authors = ["Intrinsic Technologies"]
 license = "MIT"
 edition = "2021"
diff --git a/axiom-eth/src/utils/component/circuit/comp_circuit_impl.rs b/axiom-eth/src/utils/component/circuit/comp_circuit_impl.rs
index f5f487f0..0f185f98 100644
--- a/axiom-eth/src/utils/component/circuit/comp_circuit_impl.rs
+++ b/axiom-eth/src/utils/component/circuit/comp_circuit_impl.rs
@@ -360,16 +360,22 @@ impl<F: Field, C: CoreBuilder<F>, P: PromiseBuilder<F>> Circuit<F>
         {
             let mut core_builder = self.core_builder.borrow_mut();
             let mut promise_builder = self.promise_builder.borrow_mut();
-            let rlc_builder = self.rlc_builder.borrow_mut();
+            let mut rlc_builder = self.rlc_builder.borrow_mut();
 
             let mut phase0_layouter = layouter.namespace(|| "raw synthesize phase0");
             core_builder.borrow_mut().raw_synthesize_phase0(&config.0, &mut phase0_layouter);
             promise_builder.raw_synthesize_phase0(&config.1, &mut phase0_layouter);
             rlc_builder.raw_synthesize_phase0(&config.2, phase0_layouter);
-        }
-        #[cfg(feature = "halo2-axiom")]
-        {
-            layouter.next_phase();
+
+            #[cfg(feature = "halo2-axiom")]
+            {
+                if rlc_builder.witness_gen_only() {
+                    // To save memory, clear virtual columns in phase0 because they should never be used again
+                    rlc_builder.base.pool(0).threads.clear();
+                }
+                drop(rlc_builder);
+                layouter.next_phase();
+            }
         }
         self.rlc_builder
             .borrow_mut()
diff --git a/axiom-query/Cargo.toml b/axiom-query/Cargo.toml
index f25a3d30..5f72dbca 100644
--- a/axiom-query/Cargo.toml
+++ b/axiom-query/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "axiom-query"
-version = "2.0.16"
+version = "2.0.17"
 authors = ["Intrinsic Technologies"]
 license = "MIT"
 edition = "2021"
@@ -35,7 +35,7 @@ rand = "0.8"
 rand_core = { version = "0.6", default-features = false, features = ["getrandom"] }
 
 # halo2, features turned on by axiom-eth
-axiom-eth = { version = "=0.4.2", path = "../axiom-eth", default-features = false, features = ["providers", "aggregation", "evm"] }
+axiom-eth = { version = "0.4.3", path = "../axiom-eth", default-features = false, features = ["providers", "aggregation", "evm"] }
 axiom-codec = { version = "0.2.1", path = "../axiom-codec", default-features = false }
 
 # crypto