Continue CompositeBackend (#1469)

A few incremental changes towards implementing VadCop (#424):
- `CompositeBackend` now explicitly panics for stuff that I'm planning
to defer the implementation of, like exporting verification keys.
- It internally has a backend for each machine. Currently, the entire
PIL is considered one machine called `"main"`.
- I added wrappers for all existing backends.
- I added tests for the composite wrapper. We need to check if the CI
time is still fine!

backend/src/composite/mod.rs

@@ -1,8 +1,8 @@
-use std::{io, marker::PhantomData};
+use std::{collections::BTreeMap, io, marker::PhantomData, path::PathBuf};
 
 use powdr_ast::analyzed::Analyzed;
 use powdr_executor::witgen::WitgenCallback;
-use powdr_number::FieldElement;
+use powdr_number::{DegreeType, FieldElement};
 
 use crate::{Backend, BackendFactory, BackendOptions, Error, Proof};
 
@@ -25,27 +25,48 @@ impl<F: FieldElement, B: BackendFactory<F>> BackendFactory<F> for CompositeBacke
         &self,
         pil: &'a Analyzed<F>,
         fixed: &'a [(String, Vec<F>)],
-        output_dir: Option<&'a std::path::Path>,
+        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>,
         backend_options: BackendOptions,
     ) -> Result<Box<dyn Backend<'a, F> + 'a>, Error> {
-        let backend: Box<dyn Backend<'a, F> + 'a> = self.factory.create(
-            pil,
-            fixed,
-            output_dir,
-            setup,
-            verification_key,
-            verification_app_key,
-            backend_options,
-        )?;
-        Ok(Box::new(CompositeBackend { backend }))
+        if setup.is_some() || verification_key.is_some() || verification_app_key.is_some() {
+            unimplemented!();
+        }
+
+        let backend_by_machine = ["main"]
+            .iter()
+            .map(|machine_name| {
+                let output_dir = output_dir
+                    .clone()
+                    .map(|output_dir| output_dir.join(machine_name));
+                if let Some(ref output_dir) = output_dir {
+                    std::fs::create_dir_all(output_dir)?;
+                }
+                let backend = self.factory.create(
+                    pil,
+                    fixed,
+                    output_dir,
+                    // TODO: Handle setup, verification_key, verification_app_key
+                    None,
+                    None,
+                    None,
+                    backend_options.clone(),
+                );
+                backend.map(|backend| (machine_name.to_string(), backend))
+            })
+            .collect::<Result<BTreeMap<_, _>, _>>()?;
+        Ok(Box::new(CompositeBackend { backend_by_machine }))
+    }
+
+    fn generate_setup(&self, _size: DegreeType, _output: &mut dyn io::Write) -> Result<(), Error> {
+        Err(Error::NoSetupAvailable)
     }
 }
 
 pub(crate) struct CompositeBackend<'a, F: FieldElement> {
-    backend: Box<dyn Backend<'a, F> + 'a>,
+    backend_by_machine: BTreeMap<String, Box<dyn Backend<'a, F> + 'a>>,
 }
 
 // TODO: This just forwards to the backend for now. In the future this should:
@@ -60,22 +81,28 @@ impl<'a, F: FieldElement> Backend<'a, F> for CompositeBackend<'a, F> {
         prev_proof: Option<Proof>,
         witgen_callback: WitgenCallback<F>,
     ) -> Result<Proof, Error> {
-        self.backend.prove(witness, prev_proof, witgen_callback)
+        self.backend_by_machine
+            .get("main")
+            .unwrap()
+            .prove(witness, prev_proof, witgen_callback)
     }
 
     fn verify(&self, _proof: &[u8], instances: &[Vec<F>]) -> Result<(), Error> {
-        self.backend.verify(_proof, instances)
+        self.backend_by_machine
+            .get("main")
+            .unwrap()
+            .verify(_proof, instances)
     }
 
-    fn export_setup(&self, output: &mut dyn io::Write) -> Result<(), Error> {
-        self.backend.export_setup(output)
+    fn export_setup(&self, _output: &mut dyn io::Write) -> Result<(), Error> {
+        unimplemented!()
     }
 
-    fn export_verification_key(&self, output: &mut dyn io::Write) -> Result<(), Error> {
-        self.backend.export_verification_key(output)
+    fn export_verification_key(&self, _output: &mut dyn io::Write) -> Result<(), Error> {
+        unimplemented!();
     }
 
-    fn export_ethereum_verifier(&self, output: &mut dyn io::Write) -> Result<(), Error> {
-        self.backend.export_ethereum_verifier(output)
+    fn export_ethereum_verifier(&self, _output: &mut dyn io::Write) -> Result<(), Error> {
+        unimplemented!();
     }
 }

backend/src/estark/mod.rs

@@ -116,13 +116,13 @@ fn first_step_fixup<'a, F: FieldElement>(
     (pil, patched_constants)
 }
 
-struct EStarkFilesCommon<'a, F: FieldElement> {
+struct EStarkFilesCommon<F: FieldElement> {
     degree: DegreeType,
     pil: PIL,
     /// If this field is present, it means the constants were patched with
     /// "main.first_step" column and must be written again to a file.
     patched_constants: Option<Vec<(String, Vec<F>)>>,
-    output_dir: Option<&'a Path>,
+    output_dir: Option<PathBuf>,
     proof_type: ProofType,
 }
 
@@ -135,11 +135,11 @@ fn write_json_file<T: ?Sized + Serialize>(path: &Path, data: &T) -> Result<(), E
     Ok(())
 }
 
-impl<'a, F: FieldElement> EStarkFilesCommon<'a, F> {
+impl<F: FieldElement> EStarkFilesCommon<F> {
     fn create(
-        analyzed: &'a Analyzed<F>,
-        fixed: &'a [(String, Vec<F>)],
-        output_dir: Option<&'a Path>,
+        analyzed: &Analyzed<F>,
+        fixed: &[(String, Vec<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>,
@@ -176,7 +176,7 @@ struct ProverInputFilePaths {
     contraints: PathBuf,
 }
 
-impl<'a, F: FieldElement> EStarkFilesCommon<'a, F> {
+impl<F: FieldElement> EStarkFilesCommon<F> {
     /// Write the files in the EStark Polygon format.
     fn write_files(&self, output_dir: &Path) -> Result<ProverInputFilePaths, Error> {
         let paths = ProverInputFilePaths {
@@ -216,7 +216,7 @@ impl<F: FieldElement> BackendFactory<F> for DumpFactory {
         &self,
         analyzed: &'a Analyzed<F>,
         fixed: &'a [(String, Vec<F>)],
-        output_dir: Option<&'a Path>,
+        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>,
@@ -235,9 +235,9 @@ impl<F: FieldElement> BackendFactory<F> for DumpFactory {
 }
 
 /// A backend that just dumps the files to the output directory.
-struct DumpBackend<'a, F: FieldElement>(EStarkFilesCommon<'a, F>);
+struct DumpBackend<F: FieldElement>(EStarkFilesCommon<F>);
 
-impl<'a, F: FieldElement> Backend<'a, F> for DumpBackend<'a, F> {
+impl<'a, F: FieldElement> Backend<'a, F> for DumpBackend<F> {
     fn prove(
         &self,
         _witness: &[(String, Vec<F>)],
@@ -252,6 +252,7 @@ impl<'a, F: FieldElement> Backend<'a, F> for DumpBackend<'a, F> {
         let output_dir = self
             .0
             .output_dir
+            .as_ref()
             .ok_or(Error::BackendError("output_dir is None".to_owned()))?;
 
         self.0.write_files(output_dir)?;

backend/src/estark/polygon_wrapper.rs

@@ -1,4 +1,4 @@
-use std::{fs, path::Path};
+use std::{fs, path::PathBuf};
 
 use powdr_ast::analyzed::Analyzed;
 use powdr_executor::witgen::WitgenCallback;
@@ -15,7 +15,7 @@ impl<F: FieldElement> BackendFactory<F> for Factory {
         &self,
         analyzed: &'a Analyzed<F>,
         fixed: &'a [(String, Vec<F>)],
-        output_dir: Option<&'a Path>,
+        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>,
@@ -33,11 +33,11 @@ impl<F: FieldElement> BackendFactory<F> for Factory {
     }
 }
 
-struct PolygonBackend<'a, F: FieldElement>(EStarkFilesCommon<'a, F>);
+struct PolygonBackend<F: FieldElement>(EStarkFilesCommon<F>);
 
 // TODO: make both eStark backends interchangeable, from user perspective.
 // TODO: implement the other Backend trait methods.
-impl<'a, F: FieldElement> Backend<'a, F> for PolygonBackend<'a, F> {
+impl<'a, F: FieldElement> Backend<'a, F> for PolygonBackend<F> {
     fn prove(
         &self,
         // Witness is taken from file written by the pipeline.
@@ -51,14 +51,14 @@ impl<'a, F: FieldElement> Backend<'a, F> for PolygonBackend<'a, F> {
         }
 
         let tmp_dir;
-        let output_dir = if let Some(output_dir) = self.0.output_dir {
+        let output_dir = if let Some(output_dir) = self.0.output_dir.clone() {
             output_dir
         } else {
             tmp_dir = mktemp::Temp::new_dir()?;
-            tmp_dir.as_path()
+            tmp_dir.to_path_buf()
         };
 
-        let input_paths = self.0.write_files(output_dir)?;
+        let input_paths = self.0.write_files(&output_dir)?;
 
         let commits_path = output_dir.join("commits.bin");
 
@@ -68,7 +68,7 @@ impl<'a, F: FieldElement> Backend<'a, F> for PolygonBackend<'a, F> {
             &input_paths.stark_struct,
             &input_paths.constants,
             &commits_path,
-            output_dir,
+            &output_dir,
         )
         .map_err(|e| Error::BackendError(e.to_string()))?;
 

backend/src/estark/starky_wrapper.rs

@@ -1,3 +1,4 @@
+use std::path::PathBuf;
 use std::time::Instant;
 use std::{borrow::Cow, io};
 
@@ -26,7 +27,7 @@ impl<F: FieldElement> BackendFactory<F> for Factory {
         &self,
         pil: &'a Analyzed<F>,
         fixed: &'a [(String, Vec<F>)],
-        _output_dir: Option<&std::path::Path>,
+        _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>,

backend/src/halo2/mod.rs

@@ -1,6 +1,7 @@
 #![deny(clippy::print_stdout)]
 
-use std::{io, path::Path};
+use std::io;
+use std::path::PathBuf;
 
 use crate::{Backend, BackendFactory, BackendOptions, Error, Proof};
 use powdr_ast::analyzed::Analyzed;
@@ -75,7 +76,7 @@ impl<F: FieldElement> BackendFactory<F> for Halo2ProverFactory {
         &self,
         pil: &'a Analyzed<F>,
         fixed: &'a [(String, Vec<F>)],
-        _output_dir: Option<&'a Path>,
+        _output_dir: Option<PathBuf>,
         setup: Option<&mut dyn io::Read>,
         verification_key: Option<&mut dyn io::Read>,
         verification_app_key: Option<&mut dyn io::Read>,
@@ -181,7 +182,7 @@ impl<F: FieldElement> BackendFactory<F> for Halo2MockFactory {
         &self,
         pil: &'a Analyzed<F>,
         fixed: &'a [(String, Vec<F>)],
-        _output_dir: Option<&'a Path>,
+        _output_dir: Option<PathBuf>,
         setup: Option<&mut dyn io::Read>,
         verification_key: Option<&mut dyn io::Read>,
         verification_app_key: Option<&mut dyn io::Read>,

backend/src/lib.rs

@@ -11,7 +11,7 @@ mod composite;
 use powdr_ast::analyzed::Analyzed;
 use powdr_executor::witgen::WitgenCallback;
 use powdr_number::{DegreeType, FieldElement};
-use std::{io, path::Path};
+use std::{io, path::PathBuf};
 use strum::{Display, EnumString, EnumVariantNames};
 
 #[derive(Clone, EnumString, EnumVariantNames, Display, Copy)]
@@ -20,6 +20,9 @@ pub enum BackendType {
     #[strum(serialize = "halo2")]
     Halo2,
     #[cfg(feature = "halo2")]
+    #[strum(serialize = "halo2-composite")]
+    Halo2Composite,
+    #[cfg(feature = "halo2")]
     #[strum(serialize = "halo2-mock")]
     Halo2Mock,
     #[cfg(feature = "halo2")]
@@ -28,13 +31,23 @@ pub enum BackendType {
     #[cfg(feature = "estark-polygon")]
     #[strum(serialize = "estark-polygon")]
     EStarkPolygon,
+    #[cfg(feature = "estark-polygon")]
+    #[strum(serialize = "estark-polygon-composite")]
+    EStarkPolygonComposite,
     #[strum(serialize = "estark-starky")]
     EStarkStarky,
+    #[strum(serialize = "estark-starky-composite")]
+    EStarkStarkyComposite,
     #[strum(serialize = "estark-dump")]
     EStarkDump,
+    #[strum(serialize = "estark-dump-composite")]
+    EStarkDumpComposite,
     #[cfg(feature = "plonky3")]
     #[strum(serialize = "plonky3")]
     Plonky3,
+    #[cfg(feature = "plonky3")]
+    #[strum(serialize = "plonky3-composite")]
+    Plonky3Composite,
 }
 
 pub type BackendOptions = String;
@@ -48,17 +61,35 @@ impl BackendType {
             #[cfg(feature = "halo2")]
             BackendType::Halo2 => Box::new(halo2::Halo2ProverFactory),
             #[cfg(feature = "halo2")]
+            BackendType::Halo2Composite => Box::new(composite::CompositeBackendFactory::new(
+                halo2::Halo2ProverFactory,
+            )),
+            #[cfg(feature = "halo2")]
             BackendType::Halo2Mock => Box::new(halo2::Halo2MockFactory),
             #[cfg(feature = "halo2")]
             BackendType::Halo2MockComposite => Box::new(composite::CompositeBackendFactory::new(
                 halo2::Halo2MockFactory,
             )),
             #[cfg(feature = "estark-polygon")]
             BackendType::EStarkPolygon => Box::new(estark::polygon_wrapper::Factory),
+            #[cfg(feature = "estark-polygon")]
+            BackendType::EStarkPolygonComposite => Box::new(
+                composite::CompositeBackendFactory::new(estark::polygon_wrapper::Factory),
+            ),
             BackendType::EStarkStarky => Box::new(estark::starky_wrapper::Factory),
+            BackendType::EStarkStarkyComposite => Box::new(
+                composite::CompositeBackendFactory::new(estark::starky_wrapper::Factory),
+            ),
             BackendType::EStarkDump => Box::new(estark::DumpFactory),
+            BackendType::EStarkDumpComposite => {
+                Box::new(composite::CompositeBackendFactory::new(estark::DumpFactory))
+            }
             #[cfg(feature = "plonky3")]
             BackendType::Plonky3 => Box::new(plonky3::Factory),
+            #[cfg(feature = "plonky3")]
+            BackendType::Plonky3Composite => {
+                Box::new(composite::CompositeBackendFactory::new(plonky3::Factory))
+            }
         }
     }
 }
@@ -102,7 +133,7 @@ pub trait BackendFactory<F: FieldElement> {
         &self,
         pil: &'a Analyzed<F>,
         fixed: &'a [(String, Vec<F>)],
-        output_dir: Option<&'a Path>,
+        output_dir: Option<PathBuf>,
         setup: Option<&mut dyn io::Read>,
         verification_key: Option<&mut dyn io::Read>,
         verification_app_key: Option<&mut dyn io::Read>,

backend/src/plonky3/mod.rs

@@ -1,4 +1,4 @@
-use std::{io, path::Path};
+use std::{io, path::PathBuf};
 
 use powdr_ast::analyzed::Analyzed;
 use powdr_executor::witgen::WitgenCallback;
@@ -14,7 +14,7 @@ impl<T: FieldElement> BackendFactory<T> for Factory {
         &self,
         pil: &'a Analyzed<T>,
         _fixed: &'a [(String, Vec<T>)],
-        _output_dir: Option<&'a Path>,
+        _output_dir: Option<PathBuf>,
         setup: Option<&mut dyn io::Read>,
         verification_key: Option<&mut dyn io::Read>,
         verification_app_key: Option<&mut dyn io::Read>,