String -> Arc<str>

Author: gabriel-barrett

src/aiur/bytecode.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use rustc_hash::FxHashMap;
 
 use crate::FxIndexMap;
@@ -7,7 +9,7 @@ use super::G;
 pub struct Toplevel {
   pub(crate) functions: Vec<Function>,
   #[allow(dead_code)]
-  pub(crate) filtered_functions: Vec<FxHashMap<String, Function>>,
+  pub(crate) filtered_functions: Vec<FxHashMap<Arc<str>, Function>>,
   pub(crate) memory_sizes: Vec<usize>,
 }
 
@@ -70,7 +72,7 @@ pub enum Op {
 #[derive(Clone)]
 pub enum Ctrl {
   Match(ValIdx, FxIndexMap<G, Block>, Option<Box<Block>>),
-  Return(SelIdx, String, Vec<ValIdx>),
+  Return(SelIdx, Arc<str>, Vec<ValIdx>),
   Yield(SelIdx, Vec<ValIdx>),
   MatchContinue(
     ValIdx,

src/aiur/execute.rs

@@ -1,6 +1,7 @@
 use multi_stark::p3_field::{PrimeCharacteristicRing, PrimeField64};
 use rustc_hash::FxHashMap;
 use std::collections::hash_map::Entry;
+use std::sync::Arc;
 
 use crate::{
   FxIndexMap,
@@ -18,7 +19,7 @@ use crate::{
 pub struct QueryResult {
   pub(crate) output: Vec<G>,
   pub(crate) multiplicity: G,
-  pub(crate) return_group: String,
+  pub(crate) return_group: Arc<str>,
 }
 
 pub type QueryMap = FxIndexMap<Vec<G>, QueryResult>;
@@ -257,7 +258,7 @@ impl Function {
             let result = QueryResult {
               output: vec![ptr],
               multiplicity: G::from_bool(!unconstrained),
-              return_group: String::new(),
+              return_group: Arc::from(""),
             };
             memory_queries.insert(values, result);
             map.push(ptr);

src/aiur/split.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use rustc_hash::{FxHashMap, FxHashSet};
 
 use crate::FxIndexMap;
@@ -6,7 +8,7 @@ use super::bytecode::{Block, Ctrl, Function};
 use super::layout::compute_layout;
 
 impl Function {
-  pub fn return_groups(&self) -> FxHashSet<String> {
+  pub fn return_groups(&self) -> FxHashSet<Arc<str>> {
     let mut groups = FxHashSet::default();
     collect_block(&self.body, &mut groups);
     groups
@@ -36,7 +38,7 @@ impl Function {
     self
   }
 
-  pub fn split(&self) -> FxHashMap<String, Function> {
+  pub fn split(&self) -> FxHashMap<Arc<str>, Function> {
     self
       .return_groups()
       .into_iter()
@@ -58,7 +60,7 @@ fn filter_block(block: &Block, target: &str) -> Option<Block> {
 fn filter_ctrl(ctrl: &Ctrl, target: &str) -> Option<Ctrl> {
   match ctrl {
     Ctrl::Return(sel, group, vs) => {
-      if group == target {
+      if group.as_ref() == target {
         Some(Ctrl::Return(*sel, group.clone(), vs.clone()))
       } else {
         None
@@ -148,11 +150,11 @@ fn filter_cases(
   new_cases
 }
 
-fn collect_block(block: &Block, groups: &mut FxHashSet<String>) {
+fn collect_block(block: &Block, groups: &mut FxHashSet<Arc<str>>) {
   collect_ctrl(&block.ctrl, groups);
 }
 
-fn collect_ctrl(ctrl: &Ctrl, groups: &mut FxHashSet<String>) {
+fn collect_ctrl(ctrl: &Ctrl, groups: &mut FxHashSet<Arc<str>>) {
   match ctrl {
     Ctrl::Return(_, group, _) => {
       groups.insert(group.clone());

src/aiur/synthesis.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use multi_stark::{
   lookup::LookupAir,
   p3_air::{Air, AirBuilder, BaseAir},
@@ -71,7 +73,7 @@ where
 }
 
 enum CircuitType {
-  Function { idx: usize, group: String },
+  Function { idx: usize, group: Arc<str> },
   Memory { width: usize },
   Bytes1,
   Bytes2,
@@ -85,8 +87,8 @@ impl AiurSystem {
     let toplevel_ref = &toplevel;
     let function_circuits =
       (0..toplevel_ref.functions.len()).flat_map(move |i| {
-        let groups: Vec<String> = if toplevel_ref.functions[i].constrained {
-          let mut gs: Vec<String> =
+        let groups: Vec<Arc<str>> = if toplevel_ref.functions[i].constrained {
+          let mut gs: Vec<Arc<str>> =
             toplevel_ref.filtered_functions[i].keys().cloned().collect();
           gs.sort();
           gs
@@ -142,8 +144,9 @@ impl AiurSystem {
     let _g = tracing::info_span!("aiur/witness").entered();
     let functions: Vec<CircuitType> = (0..self.toplevel.functions.len())
       .flat_map(|idx| {
-        let groups: Vec<String> = if self.toplevel.functions[idx].constrained {
-          let mut gs: Vec<String> =
+        let groups: Vec<Arc<str>> = if self.toplevel.functions[idx].constrained
+        {
+          let mut gs: Vec<Arc<str>> =
             self.toplevel.filtered_functions[idx].keys().cloned().collect();
           gs.sort();
           gs

src/aiur/trace.rs

@@ -89,7 +89,7 @@ impl Toplevel {
     let queries = unfiltered_queries
       .iter()
       .filter(|(_, res)| {
-        !res.multiplicity.is_zero() && res.return_group == group
+        !res.multiplicity.is_zero() && res.return_group.as_ref() == group
       })
       .collect::<Vec<_>>();
     let height_no_padding = queries.len();

src/ffi/aiur/protocol.rs

@@ -136,28 +136,25 @@ extern "C" fn rs_aiur_toplevel_execute(
 
   // Per-function array of `(group, width, unique_rows, total_hits)` quadruples,
   // one per split. Queries within a function are partitioned by return group.
-  let function_stats: Vec<Vec<(String, usize, usize, usize)>> = (0..toplevel
-    .functions
-    .len())
+  let function_stats: Vec<Vec<(std::sync::Arc<str>, usize, usize, usize)>> = (0
+    ..toplevel.functions.len())
     .map(|i| {
       let queries = &query_record.function_queries[i];
-      let mut stats: Vec<(String, usize, usize, usize)> = toplevel
+      let mut stats: Vec<(std::sync::Arc<str>, usize, usize, usize)> = toplevel
         .filtered_functions[i]
         .iter()
         .map(|(group, func)| {
           let (rows, hits) = queries
             .iter()
-            .filter(|(_, res)| res.return_group == *group)
+            .filter(|(_, res)| res.return_group.as_ref() == group.as_ref())
             .fold((0usize, 0usize), |(r, h), (_, res)| {
               let m = usize::try_from(res.multiplicity.as_canonical_u64())
                 .expect("multiplicity exceeds usize");
               if m != 0 { (r + 1, h + m) } else { (r, h) }
             });
           let l = func.layout;
-          let width = l.input_size
-            + l.selectors
-            + l.auxiliaries
-            + 4 * (1 + l.lookups);
+          let width =
+            l.input_size + l.selectors + l.auxiliaries + 4 * (1 + l.lookups);
           (group.clone(), width, rows, hits)
         })
         .collect();
@@ -181,7 +178,7 @@ extern "C" fn rs_aiur_toplevel_execute(
       for (j, (group, width, rows, hits)) in per_fn.iter().enumerate() {
         // (String × Nat × Nat × Nat) — right-nested pair encoding
         let quad = LeanProd::new(
-          LeanString::new(group),
+          LeanString::new(group.as_ref()),
           LeanProd::new(
             LeanOwned::box_usize(*width),
             LeanProd::new(

src/ffi/aiur/toplevel.rs

@@ -171,7 +171,9 @@ fn decode_ctrl(ctor: LeanCtor<LeanBorrowed<'_>>) -> Ctrl {
     1 => {
       let [sel_idx_obj, group_obj, val_idxs_obj] = ctor.objs::<3>();
       let sel_idx = lean_unbox_nat_as_usize(&sel_idx_obj);
-      let group = group_obj.as_string().to_string();
+      let group_lean = group_obj.as_string();
+      let group: std::sync::Arc<str> =
+        std::sync::Arc::from(group_lean.as_str());
       let val_idxs = decode_vec_val_idx(val_idxs_obj);
       Ctrl::Return(sel_idx, group, val_idxs)
     },