From d60efc9d68ead45a1990bde734c457cecd3b64e5 Mon Sep 17 00:00:00 2001
From: Andrew Brown <andrew.brown@intel.com>
Date: Mon, 22 Feb 2021 17:06:20 -0800
Subject: [PATCH] Optionally compile wasmtime-bench-api with wasi-nn and
 wasi-crypto

This adds the ability to add feature flags (e.g. `--features wasi-nn`) when compiling `wasmtime-bench-api` to allow benchmarking Wasmtime with WASI proposals included. Note that due to https://github.com/rust-lang/cargo/issues/5364, these features are only available:
 - in the `crates/bench-api` directory, e.g. `pushd crates/bench-api; cargo build --features wasi-crypto`
 - or from the top-level project directory using `-Zpackage-features`, e.g. `OPENVINO_INSTALL_DIR=/opt/intel/openvino cargo +nightly build -p wasmtime-bench-api -Zpackage-features --features wasi-nn`
---
 Cargo.lock                  |  2 ++
 crates/bench-api/Cargo.toml |  5 +++++
 crates/bench-api/src/lib.rs | 31 ++++++++++++++++++++++++++++---
 3 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 5d52e2543454..29ac6da5ee52 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3148,6 +3148,8 @@ dependencies = [
  "wasi-cap-std-sync",
  "wasmtime",
  "wasmtime-wasi",
+ "wasmtime-wasi-crypto",
+ "wasmtime-wasi-nn",
  "wat",
 ]
 
diff --git a/crates/bench-api/Cargo.toml b/crates/bench-api/Cargo.toml
index 5dcb9240b8e3..11af4b7c7452 100644
--- a/crates/bench-api/Cargo.toml
+++ b/crates/bench-api/Cargo.toml
@@ -19,6 +19,8 @@ anyhow = "1.0"
 shuffling-allocator = { version = "1.1.1", optional = true }
 wasmtime = { path = "../wasmtime", default-features = false }
 wasmtime-wasi = { path = "../wasi" }
+wasmtime-wasi-crypto = { path = "../wasi-crypto", optional = true }
+wasmtime-wasi-nn = { path = "../wasi-nn", optional = true }
 wasi-cap-std-sync = { path = "../wasi-common/cap-std-sync" }
 cap-std = "0.13"
 
@@ -27,3 +29,6 @@ wat = "1.0"
 
 [features]
 default = ["shuffling-allocator"]
+wasi-crypto = ["wasmtime-wasi-crypto"]
+wasi-nn = ["wasmtime-wasi-nn"]
+
diff --git a/crates/bench-api/src/lib.rs b/crates/bench-api/src/lib.rs
index 0d99ab635d6a..0ac83dadde04 100644
--- a/crates/bench-api/src/lib.rs
+++ b/crates/bench-api/src/lib.rs
@@ -223,9 +223,34 @@ impl BenchState {
             cx = cx.env("WASM_BENCH_USE_SMALL_WORKLOAD", &val)?;
         }
 
-        let cx = cx.build()?;
-        let wasi = Wasi::new(&store, cx);
-        wasi.add_to_linker(&mut linker)?;
+        Wasi::new(linker.store(), cx.build()?).add_to_linker(&mut linker)?;
+
+        #[cfg(feature = "wasi-nn")]
+        {
+            use std::cell::RefCell;
+            use std::rc::Rc;
+            use wasmtime_wasi_nn::{WasiNn, WasiNnCtx};
+
+            let wasi_nn = WasiNn::new(linker.store(), Rc::new(RefCell::new(WasiNnCtx::new()?)));
+            wasi_nn.add_to_linker(&mut linker)?;
+        }
+
+        #[cfg(feature = "wasi-crypto")]
+        {
+            use std::cell::RefCell;
+            use std::rc::Rc;
+            use wasmtime_wasi_crypto::{
+                WasiCryptoAsymmetricCommon, WasiCryptoCommon, WasiCryptoCtx, WasiCryptoSignatures,
+                WasiCryptoSymmetric,
+            };
+
+            let cx_crypto = Rc::new(RefCell::new(WasiCryptoCtx::new()));
+            WasiCryptoCommon::new(linker.store(), cx_crypto.clone()).add_to_linker(linker)?;
+            WasiCryptoAsymmetricCommon::new(linker.store(), cx_crypto.clone())
+                .add_to_linker(linker)?;
+            WasiCryptoSignatures::new(linker.store(), cx_crypto.clone()).add_to_linker(linker)?;
+            WasiCryptoSymmetric::new(linker.store(), cx_crypto).add_to_linker(linker)?;
+        }
 
         Ok(Self {
             engine,