Add data structures for variably-sized fixed columns (#1542)

This PR adds `number::VariablySizedColumns`, which can store several sizes of the same column. Currently, we always just have one size, but as part of #1496, we can relax that.
powdr-labs · Jul 22, 2024 · de2905d · de2905d
1 parent d8218a2
commit de2905d
Show file tree

Hide file tree

Showing 20 changed files with 215 additions and 90 deletions.
diff --git a/backend/src/composite/mod.rs b/backend/src/composite/mod.rs
@@ -8,7 +8,10 @@ use std::{
 
 use itertools::Itertools;
 use powdr_ast::analyzed::Analyzed;
-use powdr_executor::witgen::WitgenCallback;
+use powdr_executor::{
+    constant_evaluator::{get_uniquely_sized_cloned, VariablySizedColumn},
+    witgen::WitgenCallback,
+};
 use powdr_number::{DegreeType, FieldElement};
 use serde::{Deserialize, Serialize};
 use split::{machine_fixed_columns, machine_witness_columns};
@@ -49,7 +52,7 @@ impl<F: FieldElement, B: BackendFactory<F>> BackendFactory<F> for CompositeBacke
     fn create<'a>(
         &self,
         pil: Arc<Analyzed<F>>,
-        fixed: Arc<Vec<(String, Vec<F>)>>,
+        fixed: Arc<Vec<(String, VariablySizedColumn<F>)>>,
         output_dir: Option<PathBuf>,
         setup: Option<&mut dyn std::io::Read>,
         verification_key: Option<&mut dyn std::io::Read>,
@@ -60,6 +63,11 @@ impl<F: FieldElement, B: BackendFactory<F>> BackendFactory<F> for CompositeBacke
             unimplemented!();
         }
 
+        // TODO: Handle multiple sizes.
+        let fixed = Arc::new(
+            get_uniquely_sized_cloned(&fixed).map_err(|_| Error::NoVariableDegreeAvailable)?,
+        );
+
         let pils = split::split_pil((*pil).clone());
 
         // Read the setup once (if any) to pass to all backends.
@@ -105,7 +113,12 @@ impl<F: FieldElement, B: BackendFactory<F>> BackendFactory<F> for CompositeBacke
                 if let Some(ref output_dir) = output_dir {
                     std::fs::create_dir_all(output_dir)?;
                 }
-                let fixed = Arc::new(machine_fixed_columns(&fixed, &pil));
+                let fixed = Arc::new(
+                    machine_fixed_columns(&fixed, &pil)
+                        .into_iter()
+                        .map(|(column_name, values)| (column_name, values.into()))
+                        .collect(),
+                );
                 let backend = self.factory.create(
                     pil.clone(),
                     fixed,

diff --git a/backend/src/estark/mod.rs b/backend/src/estark/mod.rs
@@ -15,7 +15,10 @@ use std::{
 use crate::{Backend, BackendFactory, BackendOptions, Error, Proof};
 use powdr_ast::analyzed::Analyzed;
 
-use powdr_executor::witgen::WitgenCallback;
+use powdr_executor::{
+    constant_evaluator::{get_uniquely_sized_cloned, VariablySizedColumn},
+    witgen::WitgenCallback,
+};
 use powdr_number::{DegreeType, FieldElement};
 use serde::Serialize;
 use starky::types::{StarkStruct, Step, PIL};
@@ -222,13 +225,16 @@ impl<F: FieldElement> BackendFactory<F> for DumpFactory {
     fn create<'a>(
         &self,
         analyzed: Arc<Analyzed<F>>,
-        fixed: Arc<Vec<(String, Vec<F>)>>,
+        fixed: Arc<Vec<(String, VariablySizedColumn<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_uniquely_sized_cloned(&fixed).map_err(|_| Error::NoVariableDegreeAvailable)?,
+        );
         Ok(Box::new(DumpBackend(EStarkFilesCommon::create(
             &analyzed,
             fixed,

diff --git a/backend/src/estark/polygon_wrapper.rs b/backend/src/estark/polygon_wrapper.rs
@@ -1,7 +1,10 @@
 use std::{fs, path::PathBuf, sync::Arc};
 
 use powdr_ast::analyzed::Analyzed;
-use powdr_executor::witgen::WitgenCallback;
+use powdr_executor::{
+    constant_evaluator::{get_uniquely_sized_cloned, VariablySizedColumn},
+    witgen::WitgenCallback,
+};
 use powdr_number::FieldElement;
 
 use crate::{Backend, BackendFactory, BackendOptions, Error, Proof};
@@ -14,13 +17,16 @@ impl<F: FieldElement> BackendFactory<F> for Factory {
     fn create<'a>(
         &self,
         analyzed: Arc<Analyzed<F>>,
-        fixed: Arc<Vec<(String, Vec<F>)>>,
+        fixed: Arc<Vec<(String, VariablySizedColumn<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_uniquely_sized_cloned(&fixed).map_err(|_| Error::NoVariableDegreeAvailable)?,
+        );
         Ok(Box::new(PolygonBackend(EStarkFilesCommon::create(
             &analyzed,
             fixed,

diff --git a/backend/src/estark/starky_wrapper.rs b/backend/src/estark/starky_wrapper.rs
@@ -4,7 +4,10 @@ use std::time::Instant;
 
 use crate::{Backend, BackendFactory, BackendOptions, Error};
 use powdr_ast::analyzed::Analyzed;
-use powdr_executor::witgen::WitgenCallback;
+use powdr_executor::{
+    constant_evaluator::{get_uniquely_sized_cloned, VariablySizedColumn},
+    witgen::WitgenCallback,
+};
 use powdr_number::{FieldElement, GoldilocksField, LargeInt};
 
 use starky::{
@@ -26,7 +29,7 @@ impl<F: FieldElement> BackendFactory<F> for Factory {
     fn create<'a>(
         &self,
         pil: Arc<Analyzed<F>>,
-        fixed: Arc<Vec<(String, Vec<F>)>>,
+        fixed: Arc<Vec<(String, VariablySizedColumn<F>)>>,
         _output_dir: Option<PathBuf>,
         setup: Option<&mut dyn std::io::Read>,
         verification_key: Option<&mut dyn std::io::Read>,
@@ -49,6 +52,10 @@ impl<F: FieldElement> BackendFactory<F> for Factory {
             return Err(Error::NoVariableDegreeAvailable);
         }
 
+        let fixed = Arc::new(
+            get_uniquely_sized_cloned(&fixed).map_err(|_| Error::NoVariableDegreeAvailable)?,
+        );
+
         let proof_type: ProofType = ProofType::from(options);
 
         let params = create_stark_struct(pil.degree(), proof_type.hash_type());

diff --git a/backend/src/halo2/mod.rs b/backend/src/halo2/mod.rs
@@ -6,6 +6,7 @@ use std::sync::Arc;
 
 use crate::{Backend, BackendFactory, BackendOptions, Error, Proof};
 use powdr_ast::analyzed::Analyzed;
+use powdr_executor::constant_evaluator::{get_uniquely_sized_cloned, VariablySizedColumn};
 use powdr_executor::witgen::WitgenCallback;
 use powdr_number::{DegreeType, FieldElement};
 use prover::{generate_setup, Halo2Prover};
@@ -76,7 +77,7 @@ impl<F: FieldElement> BackendFactory<F> for Halo2ProverFactory {
     fn create<'a>(
         &self,
         pil: Arc<Analyzed<F>>,
-        fixed: Arc<Vec<(String, Vec<F>)>>,
+        fixed: Arc<Vec<(String, VariablySizedColumn<F>)>>,
         _output_dir: Option<PathBuf>,
         setup: Option<&mut dyn io::Read>,
         verification_key: Option<&mut dyn io::Read>,
@@ -87,6 +88,9 @@ impl<F: FieldElement> BackendFactory<F> for Halo2ProverFactory {
             return Err(Error::NoVariableDegreeAvailable);
         }
         let proof_type = ProofType::from(options);
+        let fixed = Arc::new(
+            get_uniquely_sized_cloned(&fixed).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);
@@ -183,7 +187,7 @@ impl<F: FieldElement> BackendFactory<F> for Halo2MockFactory {
     fn create<'a>(
         &self,
         pil: Arc<Analyzed<F>>,
-        fixed: Arc<Vec<(String, Vec<F>)>>,
+        fixed: Arc<Vec<(String, VariablySizedColumn<F>)>>,
         _output_dir: Option<PathBuf>,
         setup: Option<&mut dyn io::Read>,
         verification_key: Option<&mut dyn io::Read>,
@@ -200,6 +204,10 @@ impl<F: FieldElement> BackendFactory<F> for Halo2MockFactory {
             return Err(Error::NoAggregationAvailable);
         }
 
+        let fixed = Arc::new(
+            get_uniquely_sized_cloned(&fixed).map_err(|_| Error::NoVariableDegreeAvailable)?,
+        );
+
         Ok(Box::new(Halo2Mock { pil, fixed }))
     }
 }

diff --git a/backend/src/lib.rs b/backend/src/lib.rs
@@ -9,7 +9,7 @@ mod plonky3;
 mod composite;
 
 use powdr_ast::analyzed::Analyzed;
-use powdr_executor::witgen::WitgenCallback;
+use powdr_executor::{constant_evaluator::VariablySizedColumn, witgen::WitgenCallback};
 use powdr_number::{DegreeType, FieldElement};
 use std::{io, path::PathBuf, sync::Arc};
 use strum::{Display, EnumString, EnumVariantNames};
@@ -134,7 +134,7 @@ pub trait BackendFactory<F: FieldElement> {
     fn create<'a>(
         &self,
         pil: Arc<Analyzed<F>>,
-        fixed: Arc<Vec<(String, Vec<F>)>>,
+        fixed: Arc<Vec<(String, VariablySizedColumn<F>)>>,
         output_dir: Option<PathBuf>,
         setup: Option<&mut dyn io::Read>,
         verification_key: Option<&mut dyn io::Read>,

diff --git a/backend/src/plonky3/mod.rs b/backend/src/plonky3/mod.rs
@@ -1,7 +1,10 @@
 use std::{io, path::PathBuf, sync::Arc};
 
 use powdr_ast::analyzed::Analyzed;
-use powdr_executor::witgen::WitgenCallback;
+use powdr_executor::{
+    constant_evaluator::{get_uniquely_sized_cloned, VariablySizedColumn},
+    witgen::WitgenCallback,
+};
 use powdr_number::{FieldElement, GoldilocksField, LargeInt};
 use powdr_plonky3::Plonky3Prover;
 
@@ -13,7 +16,7 @@ impl<T: FieldElement> BackendFactory<T> for Factory {
     fn create<'a>(
         &self,
         pil: Arc<Analyzed<T>>,
-        fixed: Arc<Vec<(String, Vec<T>)>>,
+        fixed: Arc<Vec<(String, VariablySizedColumn<T>)>>,
         _output_dir: Option<PathBuf>,
         setup: Option<&mut dyn io::Read>,
         verification_key: Option<&mut dyn io::Read>,
@@ -34,6 +37,10 @@ impl<T: FieldElement> BackendFactory<T> for Factory {
             return Err(Error::NoVariableDegreeAvailable);
         }
 
+        let fixed = Arc::new(
+            get_uniquely_sized_cloned(&fixed).map_err(|_| Error::NoVariableDegreeAvailable)?,
+        );
+
         let mut p3 = Box::new(Plonky3Prover::new(pil, fixed));
 
         if let Some(verification_key) = verification_key {

diff --git a/executor/Cargo.toml b/executor/Cargo.toml
@@ -20,6 +20,7 @@ bit-vec = "0.6.3"
 num-traits = "0.2.15"
 lazy_static = "1.4.0"
 indicatif = "0.17.7"
+serde = { version = "1.0", default-features = false, features = ["alloc", "derive", "rc"] }
 
 [dev-dependencies]
 test-log = "0.2.12"

diff --git a/executor/src/constant_evaluator/data_structures.rs b/executor/src/constant_evaluator/data_structures.rs
@@ -0,0 +1,48 @@
+use serde::{Deserialize, Serialize};
+use std::collections::BTreeMap;
+
+#[derive(Serialize, Deserialize)]
+pub struct VariablySizedColumn<F> {
+    column_by_size: BTreeMap<usize, Vec<F>>,
+}
+
+#[derive(Debug)]
+pub struct HasMultipleSizesError;
+
+impl<F> VariablySizedColumn<F> {
+    /// Create a view where each column has a single size. Fails if any column has multiple sizes.
+    pub fn get_uniquely_sized(&self) -> Result<&Vec<F>, HasMultipleSizesError> {
+        if self.column_by_size.len() != 1 {
+            return Err(HasMultipleSizesError);
+        }
+        Ok(self.column_by_size.values().next().unwrap())
+    }
+}
+
+pub fn get_uniquely_sized<F>(
+    column: &[(String, VariablySizedColumn<F>)],
+) -> Result<Vec<(String, &Vec<F>)>, HasMultipleSizesError> {
+    column
+        .iter()
+        .map(|(name, column)| Ok((name.clone(), column.get_uniquely_sized()?)))
+        .collect()
+}
+
+pub fn get_uniquely_sized_cloned<F: Clone>(
+    column: &[(String, VariablySizedColumn<F>)],
+) -> Result<Vec<(String, Vec<F>)>, HasMultipleSizesError> {
+    get_uniquely_sized(column).map(|column| {
+        column
+            .into_iter()
+            .map(|(name, column)| (name, column.clone()))
+            .collect()
+    })
+}
+
+impl<F> From<Vec<F>> for VariablySizedColumn<F> {
+    fn from(column: Vec<F>) -> Self {
+        VariablySizedColumn {
+            column_by_size: [(column.len(), column)].into_iter().collect(),
+        }
+    }
+}
diff --git a/executor/src/constant_evaluator/mod.rs b/executor/src/constant_evaluator/mod.rs
@@ -3,6 +3,7 @@ use std::{
     sync::{Arc, RwLock},
 };
 
+pub use data_structures::{get_uniquely_sized, get_uniquely_sized_cloned, VariablySizedColumn};
 use itertools::Itertools;
 use powdr_ast::{
     analyzed::{Analyzed, FunctionValueDefinition, Symbol, TypedExpression},
@@ -15,12 +16,14 @@ use powdr_number::{BigInt, BigUint, DegreeType, FieldElement};
 use powdr_pil_analyzer::evaluator::{self, Definitions, SymbolLookup, Value};
 use rayon::prelude::{IntoParallelIterator, ParallelIterator};
 
+mod data_structures;
+
 /// Generates the fixed column values for all fixed columns that are defined
 /// (and not just declared).
 /// @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, Vec<T>)> {
+pub fn generate<T: FieldElement>(analyzed: &Analyzed<T>) -> Vec<(String, VariablySizedColumn<T>)> {
     let mut fixed_cols = HashMap::new();
     for (poly, value) in analyzed.constant_polys_in_source_order() {
         if let Some(value) = value {
@@ -37,8 +40,8 @@ pub fn generate<T: FieldElement>(analyzed: &Analyzed<T>) -> Vec<(String, Vec<T>)
     fixed_cols
         .into_iter()
         .sorted_by_key(|(_, (id, _))| *id)
-        .map(|(name, (_, values))| (name, values))
-        .collect::<Vec<_>>()
+        .map(|(name, (_, values))| (name, values.into()))
+        .collect()
 }
 
 fn generate_values<T: FieldElement>(
@@ -169,17 +172,28 @@ impl<'a, T: FieldElement> SymbolLookup<'a, T> for CachedSymbols<'a, T> {
 
 #[cfg(test)]
 mod test {
+    use powdr_ast::analyzed::Analyzed;
     use powdr_number::GoldilocksField;
     use powdr_pil_analyzer::analyze_string;
     use pretty_assertions::assert_eq;
     use test_log::test;
 
-    use super::*;
+    use crate::constant_evaluator::{
+        data_structures::get_uniquely_sized, generate as generate_variably_sized,
+    };
 
     fn convert(input: Vec<i32>) -> Vec<GoldilocksField> {
         input.into_iter().map(|x| x.into()).collect()
     }
 
+    fn generate(analyzed: &Analyzed<GoldilocksField>) -> Vec<(String, Vec<GoldilocksField>)> {
+        get_uniquely_sized(&generate_variably_sized(analyzed))
+            .unwrap()
+            .into_iter()
+            .map(|(name, values)| (name, values.clone()))
+            .collect()
+    }
+
     #[test]
     fn last() {
         let src = r#"

diff --git a/executor/src/witgen/block_processor.rs b/executor/src/witgen/block_processor.rs
@@ -121,7 +121,7 @@ mod tests {
     use powdr_pil_analyzer::analyze_string;
 
     use crate::{
-        constant_evaluator::generate,
+        constant_evaluator::{generate, get_uniquely_sized},
         witgen::{
             data_structures::finalizable_data::FinalizableData,
             identity_processor::Machines,
@@ -152,10 +152,8 @@ 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))
-            .collect::<Vec<_>>();
+        let constants = generate(&analyzed);
+        let constants = get_uniquely_sized(&constants).unwrap();
         let fixed_data = FixedData::new(&analyzed, &constants, &[], Default::default(), 0);
 
         // No submachines

diff --git a/executor/src/witgen/global_constraints.rs b/executor/src/witgen/global_constraints.rs
@@ -365,6 +365,8 @@ mod test {
     use pretty_assertions::assert_eq;
     use test_log::test;
 
+    use crate::constant_evaluator::get_uniquely_sized;
+
     use super::*;
 
     #[test]
@@ -437,6 +439,7 @@ namespace Global(2**20);
 ";
         let analyzed = powdr_pil_analyzer::analyze_string::<GoldilocksField>(pil_source);
         let constants = crate::constant_evaluator::generate(&analyzed);
+        let constants = get_uniquely_sized(&constants).unwrap();
         let fixed_polys = (0..constants.len())
             .map(|i| constant_poly_id(i as u64))
             .collect::<Vec<_>>();