feat: add nexus_sdk extensions support

core/src/lib.rs

@@ -20,5 +20,8 @@ pub mod nvm {
 
 /// Stwo proving
 pub mod stwo {
-    pub use nexus_vm_prover::{prove, verify, Proof, ProvingError, VerificationError};
+    pub use nexus_vm_prover::{
+        prove, prove_with_extensions, verify, verify_with_extensions, ExtensionComponent, Proof,
+        ProvingError, VerificationError,
+    };
 }

prover/src/lib.rs

@@ -7,6 +7,8 @@ pub mod components;
 pub mod extensions;
 pub mod trace;
 
+pub use extensions::ExtensionComponent;
+
 pub mod column;
 pub mod traits;
 pub mod virtual_column;
@@ -46,3 +48,32 @@ pub fn verify(proof: Proof, view: &nexus_vm::emulator::View) -> Result<(), Verif
         view.get_public_output(),
     )
 }
+
+pub fn prove_with_extensions(
+    extensions: &[ExtensionComponent],
+    trace: &impl nexus_vm::trace::Trace,
+    view: &nexus_vm::emulator::View,
+) -> Result<Proof, ProvingError> {
+    machine::Machine::<machine::BaseComponent>::prove_with_extensions(extensions, trace, view)
+}
+
+pub fn verify_with_extensions(
+    extensions: &[ExtensionComponent],
+    proof: Proof,
+    view: &nexus_vm::emulator::View,
+) -> Result<(), VerificationError> {
+    machine::Machine::<machine::BaseComponent>::verify_with_extensions(
+        extensions,
+        proof,
+        view.get_program_memory(),
+        view.view_associated_data().as_deref().unwrap_or_default(),
+        &[
+            view.get_ro_initial_memory(),
+            view.get_rw_initial_memory(),
+            view.get_public_input(),
+        ]
+        .concat(),
+        view.get_exit_code(),
+        view.get_public_output(),
+    )
+}

sdk/src/stwo/seq.rs

@@ -7,6 +7,9 @@ use thiserror::Error;
 
 use crate::error::{BuildError, ConfigurationError, IOError, PathError};
 
+/// Re-export ExtensionComponent.
+pub use nexus_core::stwo::ExtensionComponent;
+
 /// Errors that occur while proving using Stwo.
 #[derive(Debug, Error)]
 pub enum Error {
@@ -53,6 +56,8 @@ pub struct Stwo<C: Compute = Local> {
     pub elf: nexus_core::nvm::ElfFile,
     /// The associated data to prove with.
     pub ad: Vec<u8>,
+    /// Extensions to use during proving.
+    extensions: Vec<ExtensionComponent>,
     _compute: PhantomData<C>,
 }
 
@@ -76,6 +81,14 @@ where
     }
 }
 
+impl Stwo<Local> {
+    /// Add extensions to use during proving.
+    pub fn with_extensions(mut self, extensions: &[ExtensionComponent]) -> Self {
+        self.extensions = extensions.to_vec();
+        self
+    }
+}
+
 impl Prover for Stwo<Local> {
     type Proof = Proof;
     type View = nexus_core::nvm::View;
@@ -86,6 +99,7 @@ impl Prover for Stwo<Local> {
         Ok(Self {
             elf: elf.clone(),
             ad: Vec::new(),
+            extensions: Vec::new(),
             _compute: PhantomData,
         })
     }
@@ -136,7 +150,7 @@ impl Prover for Stwo<Local> {
             private_encoded.as_slice(),
             1,
         )?;
-        let proof = nexus_core::stwo::prove(&trace, &view)?;
+        let proof = nexus_core::stwo::prove_with_extensions(&self.extensions, &trace, &view)?;
 
         Ok((
             view,
@@ -148,6 +162,18 @@ impl Prover for Stwo<Local> {
     }
 }
 
+impl Proof {
+    /// Verify the proof with the given extensions.
+    pub fn verify_with_extensions(
+        &self,
+        extensions: &[ExtensionComponent],
+        view: &nexus_core::nvm::View,
+    ) -> Result<(), Error> {
+        nexus_core::stwo::verify_with_extensions(extensions, self.proof.clone(), view)?;
+        Ok(())
+    }
+}
+
 impl Verifiable for Proof {
     type View = nexus_core::nvm::View;
     type Error = Error;