From 4a7f0ad4fd54a989a3c665c559fc795a92e7a0dd Mon Sep 17 00:00:00 2001
From: Jeffrey Charles <jeff.charles@shopify.com>
Date: Mon, 21 Oct 2024 13:24:57 -0400
Subject: [PATCH] Derive the import namespace from the provider module (#786)

---
 crates/cli/src/codegen/dynamic.rs |  4 ++--
 crates/cli/src/providers.rs       | 30 ++++++++++++++++++++++--------
 crates/core/src/lib.rs            |  4 ++++
 crates/core/src/namespace.rs      | 20 ++++++++++++++++++++
 4 files changed, 48 insertions(+), 10 deletions(-)
 create mode 100644 crates/core/src/namespace.rs

diff --git a/crates/cli/src/codegen/dynamic.rs b/crates/cli/src/codegen/dynamic.rs
index eced1084..9661f662 100644
--- a/crates/cli/src/codegen/dynamic.rs
+++ b/crates/cli/src/codegen/dynamic.rs
@@ -68,7 +68,7 @@ impl DynamicGenerator {
 
     /// Generate function imports.
     pub fn generate_imports(&self, module: &mut Module) -> Result<Imports> {
-        let import_namespace = self.provider.import_namespace();
+        let import_namespace = self.provider.import_namespace()?;
         let canonical_abi_realloc_type = module.types.add(
             &[ValType::I32, ValType::I32, ValType::I32, ValType::I32],
             &[ValType::I32],
@@ -146,7 +146,7 @@ impl DynamicGenerator {
                 &[],
             );
             let (invoke_fn, _) =
-                module.add_import_func(&self.provider.import_namespace(), "invoke", invoke_type);
+                module.add_import_func(&self.provider.import_namespace()?, "invoke", invoke_type);
 
             let fn_name_ptr_local = module.locals.add(ValType::I32);
             for export in &self.function_exports {
diff --git a/crates/cli/src/providers.rs b/crates/cli/src/providers.rs
index 756ad295..d680c474 100644
--- a/crates/cli/src/providers.rs
+++ b/crates/cli/src/providers.rs
@@ -1,5 +1,6 @@
 use crate::bytecode;
-use anyhow::Result;
+use anyhow::{anyhow, Result};
+use std::str;
 
 const QUICKJS_PROVIDER_MODULE: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/provider.wasm"));
 
@@ -30,12 +31,25 @@ impl Provider {
     }
 
     /// The import namespace to use for this provider.
-    pub fn import_namespace(&self) -> String {
-        let prefix = "javy_quickjs_provider_v";
-        let version = match self {
-            Self::Default => 3,
-            Self::V2 => 2,
-        };
-        format!("{prefix}{version}")
+    pub fn import_namespace(&self) -> Result<String> {
+        match self {
+            Self::V2 => Ok("javy_quickjs_provider_v2".to_string()),
+            Self::Default => {
+                let module = walrus::Module::from_buffer(self.as_bytes())?;
+                let import_namespace = module
+                    .customs
+                    .iter()
+                    .find_map(|(_, section)| {
+                        if section.name() == "import_namespace" {
+                            Some(section)
+                        } else {
+                            None
+                        }
+                    })
+                    .ok_or_else(|| anyhow!("Provider is missing import_namespace custom section"))?
+                    .data(&Default::default()); // Argument is required but not actually used for anything.
+                Ok(str::from_utf8(&import_namespace)?.to_string())
+            }
+        }
     }
 }
diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs
index a451a714..3d058c39 100644
--- a/crates/core/src/lib.rs
+++ b/crates/core/src/lib.rs
@@ -1,11 +1,13 @@
 use anyhow::anyhow;
 use javy::Runtime;
 use javy_config::Config;
+use namespace::import_namespace;
 use once_cell::sync::OnceCell;
 use std::slice;
 use std::str;
 
 mod execution;
+mod namespace;
 mod runtime;
 
 const FUNCTION_MODULE_NAME: &str = "function.mjs";
@@ -14,6 +16,8 @@ static mut COMPILE_SRC_RET_AREA: [u32; 2] = [0; 2];
 
 static mut RUNTIME: OnceCell<Runtime> = OnceCell::new();
 
+import_namespace!("javy_quickjs_provider_v3");
+
 /// Used by Wizer to preinitialize the module.
 #[export_name = "initialize_runtime"]
 pub extern "C" fn initialize_runtime() {
diff --git a/crates/core/src/namespace.rs b/crates/core/src/namespace.rs
new file mode 100644
index 00000000..d882d54a
--- /dev/null
+++ b/crates/core/src/namespace.rs
@@ -0,0 +1,20 @@
+// Create a custom section named `import_namespace` with the contents of the
+// string argument.
+macro_rules! import_namespace {
+    ($str:literal) => {
+        const IMPORT_NAMESPACE_BYTES: &[u8] = $str.as_bytes();
+
+        #[link_section = "import_namespace"]
+        pub static IMPORT_NAMESPACE: [u8; IMPORT_NAMESPACE_BYTES.len()] = {
+            let mut arr = [0u8; IMPORT_NAMESPACE_BYTES.len()];
+            let mut i = 0;
+            while i < IMPORT_NAMESPACE_BYTES.len() {
+                arr[i] = IMPORT_NAMESPACE_BYTES[i];
+                i += 1;
+            }
+            arr
+        };
+    };
+}
+
+pub(crate) use import_namespace;