Make csv export work with multiple fixed column sizes (#1744)

Author: Schaeff

number/src/serialize.rs

@@ -29,14 +29,14 @@ const ROW_NAME: &str = "Row";
 pub fn write_polys_csv_file<T: FieldElement>(
     file: impl Write,
     render_mode: CsvRenderMode,
-    polys: &[&(String, Vec<T>)],
+    polys: &[(&String, &[T])],
 ) {
     let mut writer = Writer::from_writer(file);
 
     // Write headers, adding a "Row" column
     let mut headers = vec![ROW_NAME];
     headers.extend(polys.iter().map(|(name, _)| {
-        assert!(name != ROW_NAME);
+        assert!(*name != ROW_NAME);
         name.as_str()
     }));
     writer.write_record(&headers).unwrap();
@@ -172,7 +172,10 @@ mod tests {
             .into_iter()
             .map(|(name, values)| (name.to_string(), values))
             .collect::<Vec<_>>();
-        let polys_ref = polys.iter().collect::<Vec<_>>();
+        let polys_ref = polys
+            .iter()
+            .map(|(name, values)| (name, values.as_slice()))
+            .collect::<Vec<_>>();
 
         for render_mode in &[
             CsvRenderMode::SignedBase10,

pipeline/src/pipeline.rs

@@ -1,5 +1,6 @@
 use std::{
     borrow::Borrow,
+    collections::HashMap,
     fmt::Display,
     fs,
     io::{self, BufReader},
@@ -20,7 +21,7 @@ use powdr_ast::{
 };
 use powdr_backend::{Backend, BackendOptions, BackendType, Proof};
 use powdr_executor::{
-    constant_evaluator::{self, get_uniquely_sized_cloned, VariablySizedColumn},
+    constant_evaluator::{self, VariablySizedColumn},
     witgen::{
         chain_callbacks, extract_publics, unused_query_callback, QueryCallback, WitgenCallback,
         WitgenCallbackContext, WitnessGenerator,
@@ -549,9 +550,30 @@ impl<T: FieldElement> Pipeline<T> {
 
         if self.arguments.export_witness_csv {
             if let Some(path) = self.path_if_should_write(|name| format!("{name}_columns.csv"))? {
-                // TODO: Handle multiple sizes
-                let fixed = get_uniquely_sized_cloned(fixed).unwrap();
-                let columns = fixed.iter().chain(witness.iter()).collect::<Vec<_>>();
+                // get the column size for each namespace. This assumes all witness columns of the same namespace have the same size.
+                let witness_sizes: HashMap<&str, u64> = witness
+                    .iter()
+                    .map(|(name, values)| {
+                        let namespace = name.split("::").next().unwrap();
+                        (namespace, values.len() as u64)
+                    })
+                    .collect();
+
+                // choose the fixed column of the correct size. This assumes any namespace with no witness columns has a unique size
+                let fixed = fixed.iter().map(|(name, columns)| {
+                    let namespace = name.split("::").next().unwrap();
+                    let columns = witness_sizes
+                        .get(&namespace)
+                        // if we have witness columns, use their size
+                        .map(|size| columns.get_by_size(*size).unwrap())
+                        // otherwise, return the unique size
+                        .unwrap_or_else(|| columns.get_uniquely_sized().unwrap());
+                    (name, columns)
+                });
+
+                let columns = fixed
+                    .chain(witness.iter().map(|(name, values)| (name, values.as_ref())))
+                    .collect::<Vec<_>>();
 
                 let csv_file = fs::File::create(path).map_err(|e| vec![format!("{}", e)])?;
                 write_polys_csv_file(csv_file, self.arguments.csv_render_mode, &columns);