Accept --export to pick which func to run

Cargo.toml

@@ -7,6 +7,7 @@ members = [
   "benchmark/exit_code_function_zero",
   "benchmark/exit_code_function_one",
   "benchmark/log_truncation_function",
+  "benchmark/exports",
 ]
 
 [package]

benchmark/exports/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "exports"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]

benchmark/exports/src/lib.rs

@@ -0,0 +1,7 @@
+use std::io::Write;
+
+#[no_mangle]
+pub extern "C" fn export1() {
+    std::io::stdout().write("export1".as_bytes()).unwrap();
+    std::io::stdout().flush().unwrap();
+}

src/engine.rs

@@ -54,6 +54,7 @@ fn import_modules(
 pub fn run(
     function_path: PathBuf,
     input: Vec<u8>,
+    export: &str,
     profile_opts: Option<&ProfileOpts>,
 ) -> Result<FunctionRunResult> {
     let engine = Engine::new(
@@ -90,9 +91,7 @@ pub fn run(
         linker.module(&mut store, "Function", &module)?;
         let instance = linker.instantiate(&mut store, &module)?;
 
-        let func = instance
-            .get_typed_func::<(), ()>(store.as_context_mut(), "_start")
-            .unwrap();
+        let func = instance.get_typed_func::<(), ()>(store.as_context_mut(), export)?;
 
         let module_result;
         (module_result, profile_data) = if let Some(profile_opts) = profile_opts {
@@ -193,13 +192,15 @@ mod tests {
     use std::path::Path;
 
     const LINEAR_MEMORY_USAGE: u64 = 159 * 64;
+    const DEFAULT_EXPORT: &str = "_start";
 
     #[test]
     fn test_js_function() {
         let input = include_bytes!("../benchmark/build/js_function_input.json").to_vec();
         let function_run_result = run(
             Path::new("benchmark/build/js_function.wasm").to_path_buf(),
             input,
+            DEFAULT_EXPORT,
             None,
         );
 
@@ -212,6 +213,7 @@ mod tests {
         let function_run_result = run(
             Path::new("benchmark/build/exit_code_function_zero.wasm").to_path_buf(),
             input,
+            DEFAULT_EXPORT,
             None,
         )
         .unwrap();
@@ -225,6 +227,7 @@ mod tests {
         let function_run_result = run(
             Path::new("benchmark/build/exit_code_function_one.wasm").to_path_buf(),
             input,
+            DEFAULT_EXPORT,
             None,
         )
         .unwrap();
@@ -238,6 +241,7 @@ mod tests {
         let function_run_result = run(
             Path::new("benchmark/build/linear_memory_function.wasm").to_path_buf(),
             input,
+            DEFAULT_EXPORT,
             None,
         )
         .unwrap();
@@ -251,6 +255,7 @@ mod tests {
         let function_run_result = run(
             Path::new("benchmark/build/log_truncation_function.wasm").to_path_buf(),
             input,
+            DEFAULT_EXPORT,
             None,
         )
         .unwrap();
@@ -266,6 +271,7 @@ mod tests {
         let function_run_result = run(
             Path::new("benchmark/build/size_function.wasm").to_path_buf(),
             input,
+            DEFAULT_EXPORT,
             None,
         )
         .unwrap();

src/main.rs

@@ -24,6 +24,10 @@ struct Opts {
     /// Path to json file containing Function input; if omitted, stdin is used
     input: Option<PathBuf>,
 
+    /// Name of the export to invoke.
+    #[clap(short, long, default_value = "_start")]
+    export: String,
+
     /// Log the run result as a JSON object
     #[clap(short, long)]
     json: bool,
@@ -94,7 +98,12 @@ fn main() -> Result<()> {
         .map_err(|e| anyhow!("Invalid input JSON: {}", e))?;
 
     let profile_opts = opts.profile_opts();
-    let function_run_result = run(opts.function, buffer, profile_opts.as_ref())?;
+    let function_run_result = run(
+        opts.function,
+        buffer,
+        opts.export.as_ref(),
+        profile_opts.as_ref(),
+    )?;
 
     if opts.json {
         println!("{}", function_run_result.to_json());

tests/integration_tests.rs

@@ -197,6 +197,31 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn exports() -> Result<(), Box<dyn std::error::Error>> {
+        let mut cmd = Command::cargo_bin("function-runner")?;
+        cmd.args(["--function", "benchmark/build/exports.wasm"])
+            .args(["--export", "export1"])
+            .arg("benchmark/build/product_discount.json");
+
+        cmd.assert().success().stdout(contains("export1"));
+
+        Ok(())
+    }
+
+    #[test]
+    fn missing_export() -> Result<(), Box<dyn std::error::Error>> {
+        let mut cmd = Command::cargo_bin("function-runner")?;
+        cmd.args(["--function", "benchmark/build/exports.wasm"])
+            .arg("benchmark/build/product_discount.json");
+
+        cmd.assert()
+            .failure()
+            .stderr(contains(" failed to find function export `_start`"));
+
+        Ok(())
+    }
+
     fn profile_base_cmd_in_temp_dir(
     ) -> Result<(Command, assert_fs::TempDir), Box<dyn std::error::Error>> {
         let mut cmd = Command::cargo_bin("function-runner")?;