From 2c9924d2b44dd5cd0452d2821b738ec16c4ccb04 Mon Sep 17 00:00:00 2001
From: chriseth <chris@ethereum.org>
Date: Wed, 13 Nov 2024 17:35:19 +0100
Subject: [PATCH] Extract some functions from split_out_machines. (#2085)

A pure refactoring PR that tries to simplify split_out_machines with the
goal of being able to JIT-compile at least some machines.
---
 .../src/witgen/machines/machine_extractor.rs  | 76 ++++++++++---------
 1 file changed, 41 insertions(+), 35 deletions(-)

diff --git a/executor/src/witgen/machines/machine_extractor.rs b/executor/src/witgen/machines/machine_extractor.rs
index 3efbd37d75..9a8a113c51 100644
--- a/executor/src/witgen/machines/machine_extractor.rs
+++ b/executor/src/witgen/machines/machine_extractor.rs
@@ -140,53 +140,27 @@ pub fn split_out_machines<'a, T: FieldElement>(
             })
             .collect::<Vec<(_, &analyzed::Expression)>>();
 
-        log::trace!(
-            "\nExtracted a machine with the following witnesses:\n{}\n identities:\n{}\n connecting identities:\n{}\n and prover functions:\n{}",
-            machine_witnesses
-                .iter()
-                .map(|s| fixed.column_name(s))
-                .sorted()
-                .format(", "),
-            machine_identities
-                .iter()
-                .format("\n"),
-            machine_connections
-                .values()
-                .map(|id| id.to_string())
-                .format("\n"),
-            prover_functions
-                .iter()
-                .map(|(_, pf)| format!("{pf}"))
-                .format("\n")
+        let machine_parts = MachineParts::new(
+            fixed,
+            machine_connections,
+            machine_identities,
+            machine_witnesses,
+            prover_functions.iter().map(|&(_, pf)| pf).collect(),
         );
 
+        log_extracted_machine(&machine_parts);
+
         for (i, pf) in &prover_functions {
             if !extracted_prover_functions.insert(*i) {
                 log::warn!("Prover function was assigned to multiple machines:\n{pf}");
             }
         }
 
-        let first_witness = machine_witnesses.iter().next().unwrap();
-        let first_witness_name = fixed.column_name(first_witness);
-        let namespace = first_witness_name
-            .rfind("::")
-            .map(|idx| &first_witness_name[..idx]);
-
-        // For machines compiled using Powdr ASM we'll always have a namespace, but as a last
-        // resort we'll use the first witness name.
-        let name = namespace.unwrap_or(first_witness_name);
+        let name = suggest_machine_name(&machine_parts);
         let id = id_counter;
         id_counter += 1;
         let name_with_type = |t: &str| format!("Secondary machine {id}: {name} ({t})");
 
-        let machine_parts = MachineParts::new(
-            fixed,
-            machine_connections,
-            machine_identities,
-            machine_witnesses,
-            prover_functions.iter().map(|&(_, pf)| pf).collect(),
-        );
-
         machines.push(build_machine(fixed, machine_parts, name_with_type));
     }
     publics.add_all(base_identities.as_slice()).unwrap();
@@ -235,6 +209,38 @@ pub fn split_out_machines<'a, T: FieldElement>(
     }
 }
 
+fn log_extracted_machine<T: FieldElement>(parts: &MachineParts<'_, T>) {
+    log::trace!(
+        "\nExtracted a machine with the following witnesses:\n{}\n identities:\n{}\n connecting identities:\n{}\n and prover functions:\n{}",
+        parts.witnesses
+            .iter()
+            .map(|s|parts.column_name(s))
+            .sorted()
+            .format(", "),
+        parts.identities
+            .iter()
+            .format("\n"),
+        parts.connections
+            .values()
+            .format("\n"),
+        parts.prover_functions
+            .iter()
+            .format("\n")
+    );
+}
+
+fn suggest_machine_name<T: FieldElement>(parts: &MachineParts<'_, T>) -> String {
+    let first_witness = parts.witnesses.iter().next().unwrap();
+    let first_witness_name = parts.column_name(first_witness);
+    let namespace = first_witness_name
+        .rfind("::")
+        .map(|idx| &first_witness_name[..idx]);
+
+    // For machines compiled using Powdr ASM we'll always have a namespace, but as a last
+    // resort we'll use the first witness name.
+    namespace.unwrap_or(first_witness_name).to_string()
+}
+
 #[derive(Default)]
 /// Keeps track of the global set of publics that are referenced by the machine's identities.
 struct PublicsTracker<'a>(BTreeSet<&'a String>);