From ed2c178fd4b226dfcc5fabd86fe92ae2e483a510 Mon Sep 17 00:00:00 2001
From: Georg Wiese <georgwiese@gmail.com>
Date: Fri, 20 Dec 2024 12:37:31 +0100
Subject: [PATCH] Log JIT statistics

---
 executor/src/witgen/machines/block_machine.rs | 21 ++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/executor/src/witgen/machines/block_machine.rs b/executor/src/witgen/machines/block_machine.rs
index 17ae57c35..fafe29624 100644
--- a/executor/src/witgen/machines/block_machine.rs
+++ b/executor/src/witgen/machines/block_machine.rs
@@ -77,6 +77,10 @@ pub struct BlockMachine<'a, T: FieldElement> {
     function_cache: FunctionCache<'a, T>,
     name: String,
     multiplicity_counter: MultiplicityCounter,
+    /// Counts the number of blocks created using the JIT.
+    block_count_jit: usize,
+    /// Counts the number of blocks created using the runtime solver.
+    block_count_runtime: usize,
 }
 
 impl<'a, T: FieldElement> BlockMachine<'a, T> {
@@ -143,6 +147,8 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
                 latch_row,
                 layout,
             ),
+            block_count_jit: 0,
+            block_count_runtime: 0,
         })
     }
 }
@@ -208,6 +214,13 @@ impl<'a, T: FieldElement> Machine<'a, T> for BlockMachine<'a, T> {
                     .map(|(id, values)| (self.fixed_data.column_name(&id).to_string(), values))
                     .collect();
             }
+        } else {
+            let total_block_count = self.block_count_jit + self.block_count_runtime;
+            log::debug!(
+                "{}: {} / {total_block_count} blocks computed via JIT.",
+                self.name,
+                self.block_count_jit
+            );
         }
         self.degree = compute_size_and_log(
             &self.name,
@@ -389,7 +402,10 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
             .compile_cached(identity_id, &known_inputs)
             .is_some()
         {
-            return self.process_lookup_via_jit(mutable_state, identity_id, outer_query);
+            let updates = self.process_lookup_via_jit(mutable_state, identity_id, outer_query)?;
+            assert!(updates.is_complete());
+            self.block_count_jit += 1;
+            return Ok(updates);
         }
 
         // TODO this assumes we are always using the same lookup for this machine.
@@ -417,6 +433,9 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
 
         match process_result {
             ProcessResult::Success(updated_data, updates) => {
+                assert!(updates.is_complete());
+                self.block_count_runtime += 1;
+
                 log::trace!(
                     "End processing block machine '{}' (successfully)",
                     self.name()