Merge pull request #220 from Bernardstanislas/customize-entrypoint

Add CLI option to specify entrypoint

bench/criterion/criterion_benchmark.rs

@@ -21,6 +21,7 @@ pub fn criterion_benchmarks(c: &mut Criterion) {
             b.iter(|| {
                 cairo_run::cairo_run(
                     black_box(Path::new(&benchmark_name.1)),
+                    "main",
                     false,
                     &HINT_EXECUTOR,
                 )

cairo_programs/not_main.cairo

@@ -0,0 +1,8 @@
+func not_main():
+    [ap] = 123
+    ret
+end
+
+func main():
+    ret
+end

src/cairo_run.rs

@@ -9,10 +9,11 @@ use std::path::Path;
 
 pub fn cairo_run(
     path: &Path,
+    entrypoint: &str,
     trace_enabled: bool,
     hint_executor: &'static dyn HintExecutor,
 ) -> Result<CairoRunner, CairoRunError> {
-    let program = match Program::new(path) {
+    let program = match Program::new(path, entrypoint) {
         Ok(program) => program,
         Err(error) => return Err(CairoRunError::Program(error)),
     };
@@ -122,13 +123,14 @@ fn encode_relocated_memory(memory_bytes: &mut Vec<u8>, addr: usize, memory_cell:
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::bigint;
     use crate::vm::hints::execute_hint::BuiltinHintExecutor;
     use std::io::Read;
 
     static HINT_EXECUTOR: BuiltinHintExecutor = BuiltinHintExecutor {};
 
     fn run_test_program(program_path: &Path) -> Result<CairoRunner, CairoRunError> {
-        let program = match Program::new(program_path) {
+        let program = match Program::new(program_path, "main") {
             Ok(program) => program,
             Err(e) => return Err(CairoRunError::Program(e)),
         };
@@ -149,6 +151,24 @@ mod tests {
         Ok(cairo_runner)
     }
 
+    #[test]
+    fn cairo_run_custom_entry_point() {
+        let program_path = Path::new("cairo_programs/not_main.json");
+        let program = Program::new(program_path, "not_main").unwrap();
+
+        let mut cairo_runner = CairoRunner::new(&program, false, &HINT_EXECUTOR);
+        cairo_runner.initialize_segments(None);
+
+        let end = cairo_runner.initialize_main_entrypoint().unwrap();
+
+        assert!(cairo_runner.initialize_vm().is_ok());
+        assert!(cairo_runner.run_until_pc(end).is_ok());
+        assert!(cairo_runner.relocate().is_ok());
+        // `main` returns without doing nothing, but `not_main` sets `[ap]` to `1`
+        // Memory location was found empirically and simply hardcoded
+        assert_eq!(cairo_runner.relocated_memory[2], Some(bigint!(123)));
+    }
+
     fn compare_files(file_path_1: &Path, file_path_2: &Path) -> io::Result<()> {
         let mut file_1 = File::open(file_path_1)?;
         let mut file_2 = File::open(file_path_2)?;
@@ -173,7 +193,7 @@ mod tests {
         // it should fail when the program is loaded.
         let no_data_program_path = Path::new("cairo_programs/no_data_program.json");
 
-        assert!(cairo_run(no_data_program_path, false, &HINT_EXECUTOR).is_err());
+        assert!(cairo_run(no_data_program_path, "main", false, &HINT_EXECUTOR).is_err());
     }
 
     #[test]
@@ -182,7 +202,7 @@ mod tests {
         // it should fail when trying to run initialize_main_entrypoint.
         let no_main_program_path = Path::new("cairo_programs/no_main_program.json");
 
-        assert!(cairo_run(no_main_program_path, false, &HINT_EXECUTOR).is_err());
+        assert!(cairo_run(no_main_program_path, "main", false, &HINT_EXECUTOR).is_err());
     }
 
     #[test]
@@ -191,7 +211,7 @@ mod tests {
         // decode the instruction.
         let invalid_memory = Path::new("cairo_programs/invalid_memory.json");
 
-        assert!(cairo_run(invalid_memory, false, &HINT_EXECUTOR).is_err());
+        assert!(cairo_run(invalid_memory, "main", false, &HINT_EXECUTOR).is_err());
     }
 
     #[test]
@@ -242,7 +262,7 @@ mod tests {
     #[test]
     fn run_with_no_trace() {
         let program_path = Path::new("cairo_programs/struct.json");
-        let program = Program::new(program_path).unwrap();
+        let program = Program::new(program_path, "main").unwrap();
         let mut cairo_runner = CairoRunner::new(&program, false, &HINT_EXECUTOR);
 
         cairo_runner.initialize_segments(None);

src/main.rs

@@ -23,6 +23,8 @@ struct Args {
     trace_file: Option<PathBuf>,
     #[structopt(long = "--print_output")]
     print_output: bool,
+    #[structopt(long = "--entrypoint", default_value = "main")]
+    entrypoint: String,
     trace: Option<PathBuf>,
     #[structopt(long = "--memory_file")]
     memory_file: Option<PathBuf>,
@@ -33,8 +35,12 @@ fn main() -> Result<(), CairoRunError> {
 
     let args = Args::parse();
     let trace_enabled = args.trace_file.is_some();
-    let mut cairo_runner = match cairo_run::cairo_run(&args.filename, trace_enabled, &HINT_EXECUTOR)
-    {
+    let mut cairo_runner = match cairo_run::cairo_run(
+        &args.filename,
+        &args.entrypoint,
+        trace_enabled,
+        &HINT_EXECUTOR,
+    ) {
         Ok(runner) => runner,
         Err(error) => return Err(error),
     };

src/serde/deserialize_program.rs

@@ -255,13 +255,21 @@ pub fn deserialize_program_json(path: &Path) -> Result<ProgramJson, ProgramError
     Ok(program_json)
 }
 
-pub fn deserialize_program(path: &Path) -> Result<Program, ProgramError> {
+pub fn deserialize_program(path: &Path, entrypoint: &str) -> Result<Program, ProgramError> {
     let program_json: ProgramJson = deserialize_program_json(path)?;
+
+    let entrypoint_pc = match program_json
+        .identifiers
+        .get(&format!("__main__.{}", entrypoint))
+    {
+        Some(entrypoint_identifier) => entrypoint_identifier.pc,
+        None => return Err(ProgramError::EntrypointNotFound(entrypoint.to_string())),
+    };
     Ok(Program {
         builtins: program_json.builtins,
         prime: program_json.prime,
         data: program_json.data,
-        main: program_json.identifiers["__main__.main"].pc,
+        main: entrypoint_pc,
         hints: program_json.hints,
         reference_manager: program_json.reference_manager,
     })
@@ -579,11 +587,25 @@ mod tests {
         assert!(odd_result.is_err());
     }
 
+    #[test]
+    fn deserialize_missing_entrypoint_gives_error() {
+        let deserialization_result = deserialize_program(
+            Path::new("cairo_programs/manually_compiled/valid_program_a.json"),
+            "missing_function",
+        );
+        assert!(deserialization_result.is_err());
+        assert!(matches!(
+            deserialization_result,
+            Err(ProgramError::EntrypointNotFound(_))
+        ));
+    }
+
     #[test]
     fn deserialize_program_test() {
-        let program: Program = deserialize_program(Path::new(
-            "cairo_programs/manually_compiled/valid_program_a.json",
-        ))
+        let program: Program = deserialize_program(
+            Path::new("cairo_programs/manually_compiled/valid_program_a.json"),
+            "main",
+        )
         .expect("Failed to deserialize program");
 
         let builtins: Vec<String> = Vec::new();

src/types/errors/program_errors.rs

@@ -5,6 +5,7 @@ use std::io;
 pub enum ProgramError {
     IO(io::Error),
     Parse(serde_json::Error),
+    EntrypointNotFound(String),
 }
 
 impl From<serde_json::Error> for ProgramError {
@@ -30,6 +31,21 @@ impl fmt::Display for ProgramError {
                 write!(f, "Parsing error: ")?;
                 error.fmt(f)
             }
+            ProgramError::EntrypointNotFound(entrypoint) => {
+                f.write_fmt(format_args!("Entrypoint {} not found", entrypoint))
+            }
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn format_entrypoint_not_found_error() {
+        let error = ProgramError::EntrypointNotFound(String::from("my_function"));
+        let formatted_error = format!("{}", error);
+        assert_eq!(formatted_error, "Entrypoint my_function not found");
+    }
+}

src/types/program.rs

@@ -15,8 +15,8 @@ pub struct Program {
 }
 
 impl Program {
-    pub fn new(path: &Path) -> Result<Program, ProgramError> {
-        deserialize_program(path)
+    pub fn new(path: &Path, entrypoint: &str) -> Result<Program, ProgramError> {
+        deserialize_program(path, entrypoint)
     }
 }
 
@@ -27,9 +27,10 @@ mod tests {
 
     #[test]
     fn deserialize_program_test() {
-        let program: Program = Program::new(Path::new(
-            "cairo_programs/manually_compiled/valid_program_a.json",
-        ))
+        let program: Program = Program::new(
+            Path::new("cairo_programs/manually_compiled/valid_program_a.json"),
+            "main",
+        )
         .expect("Failed to deserialize program");
 
         let builtins: Vec<String> = Vec::new();

tests/bitwise_test.rs

@@ -11,8 +11,11 @@ use cleopatra_cairo::{
 static HINT_EXECUTOR: BuiltinHintExecutor = BuiltinHintExecutor {};
 #[test]
 fn bitwise_integration_test() {
-    let program = Program::new(Path::new("cairo_programs/bitwise_builtin_test.json"))
-        .expect("Failed to deserialize program");
+    let program = Program::new(
+        Path::new("cairo_programs/bitwise_builtin_test.json"),
+        "main",
+    )
+    .expect("Failed to deserialize program");
     let mut cairo_runner = CairoRunner::new(&program, true, &HINT_EXECUTOR);
     cairo_runner.initialize_segments(None);
     let end = cairo_runner.initialize_main_entrypoint().unwrap();