Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Continue CompositeBackend #1469

Merged
merged 15 commits into from
Jun 27, 2024
71 changes: 49 additions & 22 deletions backend/src/composite/mod.rs
Original file line number Diff line number Diff line change
@@ -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};

Expand All @@ -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:
Expand All @@ -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!();
}
}
21 changes: 11 additions & 10 deletions backend/src/estark/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}

Expand All @@ -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>,
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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>,
Expand All @@ -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>)],
Expand All @@ -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)?;
Expand Down
3 changes: 2 additions & 1 deletion backend/src/estark/starky_wrapper.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use std::path::PathBuf;
use std::time::Instant;
use std::{borrow::Cow, io};

Expand Down Expand Up @@ -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>,
Expand Down
7 changes: 4 additions & 3 deletions backend/src/halo2/mod.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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>,
Expand Down Expand Up @@ -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>,
Expand Down
28 changes: 26 additions & 2 deletions backend/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,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)]
Expand All @@ -18,6 +18,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")]
Expand All @@ -26,10 +29,17 @@ 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-composite")]
EStarkStarkyComposite,
#[strum(serialize = "estark-dump")]
EStarkDump,
#[strum(serialize = "estark-dump-composite")]
EStarkDumpComposite,
}

pub type BackendOptions = String;
Expand All @@ -43,15 +53,29 @@ 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))
}
}
}
}
Expand Down Expand Up @@ -95,7 +119,7 @@ pub trait BackendFactory<F: FieldElement> {
&self,
pil: &'a Analyzed<F>,
fixed: &'a [(String, Vec<F>)],
output_dir: Option<&'a Path>,
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This made it hard to change the path in the CompositeBackend. I don't see a good reason to pass a reference with a lifetime here anyway.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sgtm

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>,
Expand Down
12 changes: 6 additions & 6 deletions pipeline/src/pipeline.rs
Original file line number Diff line number Diff line change
Expand Up @@ -914,7 +914,7 @@ impl<T: FieldElement> Pipeline<T> {
.create(
pil.borrow(),
&fixed_cols[..],
self.output_dir(),
self.output_dir.clone(),
setup.as_io_read(),
vkey.as_io_read(),
vkey_app.as_io_read(),
Expand Down Expand Up @@ -950,8 +950,8 @@ impl<T: FieldElement> Pipeline<T> {
Ok(self.artifact.proof.as_ref().unwrap())
}

pub fn output_dir(&self) -> Option<&Path> {
self.output_dir.as_ref().map(|p| p.as_ref())
pub fn output_dir(&self) -> &Option<PathBuf> {
&self.output_dir
}

pub fn is_force_overwrite(&self) -> bool {
Expand Down Expand Up @@ -996,7 +996,7 @@ impl<T: FieldElement> Pipeline<T> {
.create(
pil.borrow(),
&fixed_cols[..],
self.output_dir(),
self.output_dir.clone(),
setup_file
.as_mut()
.map(|file| file as &mut dyn std::io::Read),
Expand Down Expand Up @@ -1050,7 +1050,7 @@ impl<T: FieldElement> Pipeline<T> {
.create(
pil.borrow(),
&fixed_cols[..],
self.output_dir(),
self.output_dir.clone(),
setup_file
.as_mut()
.map(|file| file as &mut dyn std::io::Read),
Expand Down Expand Up @@ -1096,7 +1096,7 @@ impl<T: FieldElement> Pipeline<T> {
.create(
pil.borrow(),
&fixed_cols[..],
self.output_dir(),
self.output_dir.clone(),
setup_file
.as_mut()
.map(|file| file as &mut dyn std::io::Read),
Expand Down
Loading
Loading