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
16 changes: 8 additions & 8 deletions backend/src/estark/polygon_wrapper.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use std::{fs, path::Path};
use std::{fs, path::PathBuf};

use powdr_ast::analyzed::Analyzed;
use powdr_executor::witgen::WitgenCallback;
Expand All @@ -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>,
Expand All @@ -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.
Expand All @@ -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");

Expand All @@ -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()))?;

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
35 changes: 33 additions & 2 deletions backend/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand All @@ -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")]
Expand All @@ -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;
Expand All @@ -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))
}
}
}
}
Expand Down Expand Up @@ -102,7 +133,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
4 changes: 2 additions & 2 deletions backend/src/plonky3/mod.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use std::{io, path::Path};
use std::{io, path::PathBuf};

use powdr_ast::analyzed::Analyzed;
use powdr_executor::witgen::WitgenCallback;
Expand All @@ -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>,
Expand Down
Loading
Loading