Refactor

backend/src/composite/mod.rs

@@ -2,11 +2,11 @@ use std::{collections::BTreeMap, io, marker::PhantomData, path::PathBuf, sync::A
 
 use powdr_ast::analyzed::Analyzed;
 use powdr_executor::witgen::WitgenCallback;
-use powdr_number::{DegreeType, FieldElement};
+use powdr_number::{DegreeType, FieldElement, FixedColumns};
 use serde::{Deserialize, Serialize};
 use split::select_machine_columns;
 
-use crate::{get_only_size, Backend, BackendFactory, BackendOptions, Error, Proof};
+use crate::{Backend, BackendFactory, BackendOptions, Error, Proof};
 
 mod split;
 
@@ -35,7 +35,7 @@ impl<F: FieldElement, B: BackendFactory<F>> BackendFactory<F> for CompositeBacke
     fn create<'a>(
         &self,
         pil: Arc<Analyzed<F>>,
-        fixed: Arc<Vec<(String, BTreeMap<usize, Vec<F>>)>>,
+        fixed: Arc<FixedColumns<F>>,
         output_dir: Option<PathBuf>,
         setup: Option<&mut dyn std::io::Read>,
         verification_key: Option<&mut dyn std::io::Read>,
@@ -47,7 +47,11 @@ impl<F: FieldElement, B: BackendFactory<F>> BackendFactory<F> for CompositeBacke
         }
 
         // TODO: Handle multiple sizes.
-        let fixed = Arc::new(get_only_size(&fixed)?);
+        let fixed = Arc::new(
+            fixed
+                .get_only_size()
+                .map_err(|_| Error::NoVariableDegreeAvailable)?,
+        );
 
         let per_machine_data = split::split_pil((*pil).clone())
             .into_iter()
@@ -70,7 +74,7 @@ impl<F: FieldElement, B: BackendFactory<F>> BackendFactory<F> for CompositeBacke
                     .map(|(column_name, values)| {
                         (column_name, [(values.len(), values)].into_iter().collect())
                     })
-                    .collect(),
+                    .into(),
                 );
                 let backend = self.factory.create(
                     pil.clone(),

backend/src/estark/mod.rs

@@ -5,19 +5,18 @@ pub mod polygon_wrapper;
 pub mod starky_wrapper;
 
 use std::{
-    collections::BTreeMap,
     fs::File,
     io::{self, BufWriter, Write},
     iter::{once, repeat},
     path::{Path, PathBuf},
     sync::Arc,
 };
 
-use crate::{get_only_size, Backend, BackendFactory, BackendOptions, Error, Proof};
+use crate::{Backend, BackendFactory, BackendOptions, Error, Proof};
 use powdr_ast::analyzed::Analyzed;
 
 use powdr_executor::witgen::WitgenCallback;
-use powdr_number::{DegreeType, FieldElement};
+use powdr_number::{DegreeType, FieldElement, FixedColumns};
 use serde::Serialize;
 use starky::types::{StarkStruct, Step, PIL};
 
@@ -223,14 +222,18 @@ impl<F: FieldElement> BackendFactory<F> for DumpFactory {
     fn create<'a>(
         &self,
         analyzed: Arc<Analyzed<F>>,
-        fixed: Arc<Vec<(String, BTreeMap<usize, Vec<F>>)>>,
+        fixed: Arc<FixedColumns<F>>,
         output_dir: Option<PathBuf>,
         setup: Option<&mut dyn std::io::Read>,
         verification_key: Option<&mut dyn std::io::Read>,
         verification_app_key: Option<&mut dyn std::io::Read>,
         options: BackendOptions,
     ) -> Result<Box<dyn crate::Backend<'a, F> + 'a>, Error> {
-        let fixed = Arc::new(get_only_size(&fixed)?);
+        let fixed = Arc::new(
+            fixed
+                .get_only_size()
+                .map_err(|_| Error::NoVariableDegreeAvailable)?,
+        );
         Ok(Box::new(DumpBackend(EStarkFilesCommon::create(
             &analyzed,
             fixed,

backend/src/estark/starky_wrapper.rs

@@ -1,12 +1,12 @@
+use std::io;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::time::Instant;
-use std::{collections::BTreeMap, io};
 
-use crate::{get_only_size, Backend, BackendFactory, BackendOptions, Error};
+use crate::{Backend, BackendFactory, BackendOptions, Error};
 use powdr_ast::analyzed::Analyzed;
 use powdr_executor::witgen::WitgenCallback;
-use powdr_number::{FieldElement, GoldilocksField, LargeInt};
+use powdr_number::{FieldElement, FixedColumns, GoldilocksField, LargeInt};
 
 use starky::{
     merklehash::MerkleTreeGL,
@@ -27,7 +27,7 @@ impl<F: FieldElement> BackendFactory<F> for Factory {
     fn create<'a>(
         &self,
         pil: Arc<Analyzed<F>>,
-        fixed: Arc<Vec<(String, BTreeMap<usize, Vec<F>>)>>,
+        fixed: Arc<FixedColumns<F>>,
         _output_dir: Option<PathBuf>,
         setup: Option<&mut dyn std::io::Read>,
         verification_key: Option<&mut dyn std::io::Read>,
@@ -50,7 +50,11 @@ impl<F: FieldElement> BackendFactory<F> for Factory {
             return Err(Error::NoVariableDegreeAvailable);
         }
 
-        let fixed = Arc::new(get_only_size(&fixed)?);
+        let fixed = Arc::new(
+            fixed
+                .get_only_size()
+                .map_err(|_| Error::NoVariableDegreeAvailable)?,
+        );
 
         let proof_type: ProofType = ProofType::from(options);
 

backend/src/halo2/mod.rs

@@ -1,14 +1,13 @@
 #![deny(clippy::print_stdout)]
 
-use std::collections::BTreeMap;
 use std::io;
 use std::path::PathBuf;
 use std::sync::Arc;
 
-use crate::{get_only_size, Backend, BackendFactory, BackendOptions, Error, Proof};
+use crate::{Backend, BackendFactory, BackendOptions, Error, Proof};
 use powdr_ast::analyzed::Analyzed;
 use powdr_executor::witgen::WitgenCallback;
-use powdr_number::{DegreeType, FieldElement};
+use powdr_number::{DegreeType, FieldElement, FixedColumns};
 use prover::{generate_setup, Halo2Prover};
 
 use serde::de::{self, Deserializer};
@@ -77,7 +76,7 @@ impl<F: FieldElement> BackendFactory<F> for Halo2ProverFactory {
     fn create<'a>(
         &self,
         pil: Arc<Analyzed<F>>,
-        fixed: Arc<Vec<(String, BTreeMap<usize, Vec<F>>)>>,
+        fixed: Arc<FixedColumns<F>>,
         _output_dir: Option<PathBuf>,
         setup: Option<&mut dyn io::Read>,
         verification_key: Option<&mut dyn io::Read>,
@@ -88,12 +87,12 @@ impl<F: FieldElement> BackendFactory<F> for Halo2ProverFactory {
             return Err(Error::NoVariableDegreeAvailable);
         }
         let proof_type = ProofType::from(options);
-        let mut halo2 = Box::new(Halo2Prover::new(
-            pil,
-            Arc::new(get_only_size(&fixed)?),
-            setup,
-            proof_type,
-        )?);
+        let fixed = Arc::new(
+            fixed
+                .get_only_size()
+                .map_err(|_| Error::NoVariableDegreeAvailable)?,
+        );
+        let mut halo2 = Box::new(Halo2Prover::new(pil, fixed, setup, proof_type)?);
         if let Some(vk) = verification_key {
             halo2.add_verification_key(vk);
         }
@@ -191,7 +190,7 @@ impl<F: FieldElement> BackendFactory<F> for Halo2MockFactory {
     fn create<'a>(
         &self,
         pil: Arc<Analyzed<F>>,
-        fixed: Arc<Vec<(String, BTreeMap<usize, Vec<F>>)>>,
+        fixed: Arc<FixedColumns<F>>,
         _output_dir: Option<PathBuf>,
         setup: Option<&mut dyn io::Read>,
         verification_key: Option<&mut dyn io::Read>,
@@ -208,10 +207,13 @@ impl<F: FieldElement> BackendFactory<F> for Halo2MockFactory {
             return Err(Error::NoAggregationAvailable);
         }
 
-        Ok(Box::new(Halo2Mock {
-            pil,
-            fixed: Arc::new(get_only_size(&fixed)?),
-        }))
+        let fixed = Arc::new(
+            fixed
+                .get_only_size()
+                .map_err(|_| Error::NoVariableDegreeAvailable)?,
+        );
+
+        Ok(Box::new(Halo2Mock { pil, fixed }))
     }
 }
 

backend/src/lib.rs

@@ -10,8 +10,8 @@ mod composite;
 
 use powdr_ast::analyzed::Analyzed;
 use powdr_executor::witgen::WitgenCallback;
-use powdr_number::{DegreeType, FieldElement};
-use std::{collections::BTreeMap, io, path::PathBuf, sync::Arc};
+use powdr_number::{DegreeType, FieldElement, FixedColumns};
+use std::{io, path::PathBuf, sync::Arc};
 use strum::{Display, EnumString, EnumVariantNames};
 
 #[derive(Clone, EnumString, EnumVariantNames, Display, Copy)]
@@ -131,11 +131,10 @@ pub type Proof = Vec<u8>;
 pub trait BackendFactory<F: FieldElement> {
     /// Create a new backend object.
     #[allow(clippy::too_many_arguments)]
-    #[allow(clippy::type_complexity)]
     fn create<'a>(
         &self,
         pil: Arc<Analyzed<F>>,
-        fixed: Arc<Vec<(String, BTreeMap<usize, Vec<F>>)>>,
+        fixed: Arc<FixedColumns<F>>,
         output_dir: Option<PathBuf>,
         setup: Option<&mut dyn io::Read>,
         verification_key: Option<&mut dyn io::Read>,
@@ -186,19 +185,3 @@ pub trait Backend<'a, F: FieldElement> {
         Err(Error::NoEthereumVerifierAvailable)
     }
 }
-
-fn get_only_size<F: Clone>(
-    columns: &[(String, BTreeMap<usize, Vec<F>>)],
-) -> Result<Vec<(String, Vec<F>)>, Error> {
-    // TODO: This clones the values
-    columns
-        .iter()
-        .map(|(name, column_by_size)| {
-            if column_by_size.len() != 1 {
-                return Err(Error::NoVariableDegreeAvailable);
-            }
-            let values = column_by_size.values().next().unwrap().clone();
-            Ok((name.clone(), values))
-        })
-        .collect()
-}

executor/src/constant_evaluator/mod.rs

@@ -11,7 +11,7 @@ use powdr_ast::{
         IndexAccess,
     },
 };
-use powdr_number::{BigInt, BigUint, DegreeType, FieldElement};
+use powdr_number::{BigInt, BigUint, DegreeType, FieldElement, FixedColumns};
 use powdr_pil_analyzer::evaluator::{self, Definitions, SymbolLookup, Value};
 use rayon::prelude::{IntoParallelIterator, ParallelIterator};
 
@@ -20,7 +20,7 @@ use rayon::prelude::{IntoParallelIterator, ParallelIterator};
 /// @returns the names (in source order) and the values for the columns.
 /// Arrays of columns are flattened, the name of the `i`th array element
 /// is `name[i]`.
-pub fn generate<T: FieldElement>(analyzed: &Analyzed<T>) -> Vec<(String, BTreeMap<usize, Vec<T>>)> {
+pub fn generate<T: FieldElement>(analyzed: &Analyzed<T>) -> FixedColumns<T> {
     let mut fixed_cols = HashMap::new();
     for (poly, value) in analyzed.constant_polys_in_source_order() {
         if let Some(value) = value {
@@ -38,7 +38,7 @@ pub fn generate<T: FieldElement>(analyzed: &Analyzed<T>) -> Vec<(String, BTreeMa
         .into_iter()
         .sorted_by_key(|(_, (id, _))| *id)
         .map(|(name, (_, values))| (name, [(values.len(), values)].into_iter().collect()))
-        .collect::<Vec<_>>()
+        .into()
 }
 
 fn generate_values<T: FieldElement>(

executor/src/witgen/block_processor.rs

@@ -152,10 +152,7 @@ mod tests {
         f: impl Fn(BlockProcessor<T, Q>, BTreeMap<String, PolyID>, u64, usize) -> R,
     ) -> R {
         let analyzed = analyze_string(src);
-        let constants = generate(&analyzed)
-            .into_iter()
-            .map(|(n, c)| (n.to_string(), c.into_values().next().unwrap()))
-            .collect::<Vec<_>>();
+        let constants = generate(&analyzed).get_only_size().unwrap();
         let fixed_data = FixedData::new(&analyzed, &constants, &[], Default::default(), 0);
 
         // No submachines

executor/src/witgen/global_constraints.rs

@@ -436,15 +436,16 @@ namespace Global(2**20);
     [ D ] in [ SHIFTED ];
 ";
         let analyzed = powdr_pil_analyzer::analyze_string::<GoldilocksField>(pil_source);
-        let constants = crate::constant_evaluator::generate(&analyzed);
+        let constants = crate::constant_evaluator::generate(&analyzed)
+            .get_only_size()
+            .unwrap();
         let fixed_polys = (0..constants.len())
             .map(|i| constant_poly_id(i as u64))
             .collect::<Vec<_>>();
         let mut known_constraints = fixed_polys
             .iter()
             .zip(&constants)
             .filter_map(|(&poly_id, (_, values))| {
-                let values = values.values().next().unwrap();
                 process_fixed_column(values).map(|(constraint, _full)| (poly_id, constraint))
             })
             .collect::<BTreeMap<_, _>>();

executor/src/witgen/mod.rs

@@ -8,7 +8,7 @@ use powdr_ast::analyzed::{
 };
 use powdr_ast::parsed::visitor::ExpressionVisitable;
 use powdr_ast::parsed::{FunctionKind, LambdaExpression};
-use powdr_number::{DegreeType, FieldElement};
+use powdr_number::{DegreeType, FieldElement, WitnessColumns};
 
 use self::data_structures::column_map::{FixedColumnMap, WitnessColumnMap};
 pub use self::eval_result::{
@@ -87,6 +87,7 @@ impl<T: FieldElement> WitgenCallback<T> {
         .with_external_witness_values(current_witness)
         .with_challenges(stage, challenges)
         .generate()
+        .0
     }
 }
 
@@ -154,7 +155,7 @@ impl<'a, 'b, T: FieldElement> WitnessGenerator<'a, 'b, T> {
 
     /// Generates the committed polynomial values
     /// @returns the values (in source order) and the degree of the polynomials.
-    pub fn generate(self) -> Vec<(String, Vec<T>)> {
+    pub fn generate(self) -> WitnessColumns<T> {
         record_start(OUTER_CODE_NAME);
         let fixed = FixedData::new(
             self.analyzed,
@@ -250,7 +251,7 @@ impl<'a, 'b, T: FieldElement> WitnessGenerator<'a, 'b, T> {
                 assert!(!column.is_empty());
                 (name, column)
             })
-            .collect::<Vec<_>>();
+            .into();
 
         log::debug!("Publics:");
         for (name, value) in extract_publics(&witness_cols, self.analyzed) {
@@ -261,10 +262,11 @@ impl<'a, 'b, T: FieldElement> WitnessGenerator<'a, 'b, T> {
 }
 
 pub fn extract_publics<T: FieldElement>(
-    witness: &[(String, Vec<T>)],
+    witness: &WitnessColumns<T>,
     pil: &Analyzed<T>,
 ) -> Vec<(String, T)> {
     let witness = witness
+        .0
         .iter()
         .map(|(name, col)| (name.clone(), col))
         .collect::<BTreeMap<_, _>>();