Merge pull request #2920 from o1-labs/optional-preimage-oracle

Optional PreImageOracle
o1-labs · Jan 2, 2025 · eb80434 · eb80434
2 parents 7870114 + 343132b
commit eb80434
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 16 deletions.
diff --git a/o1vm/src/legacy/main.rs b/o1vm/src/legacy/main.rs
@@ -29,7 +29,7 @@ use o1vm::{
         BaseSponge, Fp, OpeningProof, ScalarSponge,
     },
     lookups::LookupTableIDs,
-    preimage_oracle::PreImageOracle,
+    preimage_oracle::{NullPreImageOracle, PreImageOracle, PreImageOracleT},
     test_preimage_read,
 };
 use poly_commitment::SRS as _;
@@ -59,9 +59,6 @@ pub fn cannon_main(args: cli::cannon::RunArgs) {
             .unwrap_or_else(|_| panic!("Error deserializing metadata file {}", f))
     });
 
-    let mut po = PreImageOracle::create(&configuration.host);
-    let _child = po.start();
-
     // Initialize some data used for statistical computations
     let start = Start::create(state.step as usize);
 
@@ -78,8 +75,27 @@ pub fn cannon_main(args: cli::cannon::RunArgs) {
 
     // Initialize the environments
     // The Keccak environment is extracted inside the loop
-    let mut mips_wit_env =
-        mips_witness::Env::<Fp, PreImageOracle>::create(cannon::PAGE_SIZE as usize, state, po);
+    let mut mips_wit_env = match configuration.host.clone() {
+        Some(host) => {
+            let mut po = PreImageOracle::create(host);
+            let _child = po.start();
+            mips_witness::Env::<Fp, Box<dyn PreImageOracleT>>::create(
+                cannon::PAGE_SIZE as usize,
+                state,
+                Box::new(po),
+            )
+        }
+        None => {
+            debug!("No preimage oracle provided 🤞");
+            // warning: the null preimage oracle has no data and will crash the program if used
+            mips_witness::Env::<Fp, Box<dyn PreImageOracleT>>::create(
+                cannon::PAGE_SIZE as usize,
+                state,
+                Box::new(NullPreImageOracle),
+            )
+        }
+    };
+
     let mut mips_con_env = mips_constraints::Env::<Fp>::default();
     // The keccak environment is extracted inside the loop
 

diff --git a/o1vm/src/pickles/main.rs b/o1vm/src/pickles/main.rs
@@ -17,7 +17,7 @@ use o1vm::{
         Instruction,
     },
     pickles::{proof::ProofInputs, prover, verifier},
-    preimage_oracle::PreImageOracle,
+    preimage_oracle::{NullPreImageOracle, PreImageOracle, PreImageOracleT},
     test_preimage_read,
 };
 use poly_commitment::{ipa::SRS, SRS as _};
@@ -44,9 +44,6 @@ pub fn cannon_main(args: cli::cannon::RunArgs) {
             .unwrap_or_else(|_| panic!("Error deserializing metadata file {}", f))
     });
 
-    let mut po = PreImageOracle::create(&configuration.host);
-    let _child = po.start();
-
     // Initialize some data used for statistical computations
     let start = Start::create(state.step as usize);
 
@@ -58,8 +55,26 @@ pub fn cannon_main(args: cli::cannon::RunArgs) {
     };
 
     // Initialize the environments
-    let mut mips_wit_env =
-        mips_witness::Env::<Fp, PreImageOracle>::create(cannon::PAGE_SIZE as usize, state, po);
+    let mut mips_wit_env = match configuration.host.clone() {
+        Some(host) => {
+            let mut po = PreImageOracle::create(host);
+            let _child = po.start();
+            mips_witness::Env::<Fp, Box<dyn PreImageOracleT>>::create(
+                cannon::PAGE_SIZE as usize,
+                state,
+                Box::new(po),
+            )
+        }
+        None => {
+            debug!("No preimage oracle provided 🤞");
+            // warning: the null preimage oracle has no data and will crash the program if used
+            mips_witness::Env::<Fp, Box<dyn PreImageOracleT>>::create(
+                cannon::PAGE_SIZE as usize,
+                state,
+                Box::new(NullPreImageOracle),
+            )
+        }
+    };
 
     let constraints = mips_constraints::get_all_constraints::<Fp>();
 

diff --git a/o1vm/src/preimage_oracle.rs b/o1vm/src/preimage_oracle.rs
@@ -106,9 +106,7 @@ pub fn create_bidirectional_channel() -> Option<(RW, RW)> {
 }
 
 impl PreImageOracle {
-    pub fn create(hp_opt: &Option<HostProgram>) -> PreImageOracle {
-        let host_program = hp_opt.as_ref().expect("No host program given");
-
+    pub fn create(host_program: HostProgram) -> PreImageOracle {
         let mut cmd = Command::new(&host_program.name);
         cmd.args(&host_program.arguments);
 
@@ -157,6 +155,18 @@ impl PreImageOracle {
     }
 }
 
+pub struct NullPreImageOracle;
+
+impl PreImageOracleT for NullPreImageOracle {
+    fn get_preimage(&mut self, _key: [u8; 32]) -> Preimage {
+        panic!("No preimage oracle specified for preimage retrieval");
+    }
+
+    fn hint(&mut self, _hint: Hint) {
+        panic!("No preimage oracle specified for hints");
+    }
+}
+
 impl PreImageOracleT for PreImageOracle {
     // The preimage protocol goes as follows
     // 1. Ask for data through a key
@@ -224,6 +234,16 @@ impl PreImageOracleT for PreImageOracle {
     }
 }
 
+impl PreImageOracleT for Box<dyn PreImageOracleT> {
+    fn get_preimage(&mut self, key: [u8; 32]) -> Preimage {
+        self.as_mut().get_preimage(key)
+    }
+
+    fn hint(&mut self, hint: Hint) {
+        self.as_mut().hint(hint)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

diff --git a/o1vm/src/test_preimage_read.rs b/o1vm/src/test_preimage_read.rs
@@ -18,7 +18,8 @@ pub fn main(args: cli::cannon::RunArgs) -> ExitCode {
     let preimage_db_dir = args.preimage_db_dir;
 
     if let Some(preimage_key_dir) = preimage_db_dir {
-        let mut po = PreImageOracle::create(&configuration.host);
+        let host_program = configuration.host.expect("No host program specified");
+        let mut po = PreImageOracle::create(host_program);
         let _child = po.start();
         debug!("Let server start");
         std::thread::sleep(std::time::Duration::from_secs(5));