feat(rust): implement with macro argument (#699)

* feat(rust): implement `with` macro argument * feat: support `with` absolute/relative paths Signed-off-by: Brian H <[email protected]> * feat: support `with` absolute/relative paths Signed-off-by: Brian H <[email protected]> --------- Signed-off-by: Brian H <[email protected]> Co-authored-by: Brian H <[email protected]>
bytecodealliance · Oct 18, 2023 · de66aa5 · de66aa5
1 parent 434cb76
commit de66aa5
Show file tree

Hide file tree

Showing 3 changed files with 153 additions and 25 deletions.
diff --git a/crates/rust-macro/src/lib.rs b/crates/rust-macro/src/lib.rs
@@ -1,5 +1,6 @@
 use proc_macro2::{Span, TokenStream};
 use quote::ToTokens;
+use std::collections::HashMap;
 use std::path::{Path, PathBuf};
 use syn::parse::{Error, Parse, ParseStream, Result};
 use syn::punctuated::Punctuated;
@@ -79,6 +80,7 @@ impl Parse for Config {
                             .map(|p| p.into_token_stream().to_string())
                             .collect()
                     }
+                    Opt::With(with) => opts.with.extend(with),
                 }
             }
         } else {
@@ -166,6 +168,7 @@ mod kw {
     syn::custom_keyword!(stubs);
     syn::custom_keyword!(export_prefix);
     syn::custom_keyword!(additional_derives);
+    syn::custom_keyword!(with);
 }
 
 #[derive(Clone)]
@@ -225,6 +228,7 @@ enum Opt {
     ExportPrefix(syn::LitStr),
     // Parse as paths so we can take the concrete types/macro names rather than raw strings
     AdditionalDerives(Vec<syn::Path>),
+    With(HashMap<String, String>),
 }
 
 impl Parse for Opt {
@@ -322,6 +326,14 @@ impl Parse for Opt {
             syn::bracketed!(contents in input);
             let list = Punctuated::<_, Token![,]>::parse_terminated(&contents)?;
             Ok(Opt::AdditionalDerives(list.iter().cloned().collect()))
+        } else if l.peek(kw::with) {
+            input.parse::<kw::with>()?;
+            input.parse::<Token![:]>()?;
+            let contents;
+            let _lbrace = braced!(contents in input);
+            let fields: Punctuated<_, Token![,]> =
+                contents.parse_terminated(with_field_parse, Token![,])?;
+            Ok(Opt::With(HashMap::from_iter(fields.into_iter())))
         } else {
             Err(l.error())
         }
@@ -347,3 +359,46 @@ fn serialize(path: syn::Path) -> String {
         serialized
     }
 }
+
+fn with_field_parse(input: ParseStream<'_>) -> Result<(String, String)> {
+    let interface = input.parse::<syn::LitStr>()?.value();
+    input.parse::<Token![:]>()?;
+    let start = input.span();
+    let path = input.parse::<syn::Path>()?;
+
+    // It's not possible for the segments of a path to be empty
+    let span = start
+        .join(path.segments.last().unwrap().ident.span())
+        .unwrap_or(start);
+
+    let mut buf = String::new();
+    let append = |buf: &mut String, segment: syn::PathSegment| -> Result<()> {
+        if !segment.arguments.is_none() {
+            return Err(Error::new(
+                span,
+                "Module path must not contain angles or parens",
+            ));
+        }
+
+        buf.push_str(&segment.ident.to_string());
+
+        Ok(())
+    };
+
+    if path.leading_colon.is_some() {
+        buf.push_str("::");
+    }
+
+    let mut segments = path.segments.into_iter();
+
+    if let Some(segment) = segments.next() {
+        append(&mut buf, segment)?;
+    }
+
+    for segment in segments {
+        buf.push_str("::");
+        append(&mut buf, segment)?;
+    }
+
+    Ok((interface, buf))
+}
diff --git a/crates/rust/src/interface.rs b/crates/rust/src/interface.rs
@@ -1,7 +1,7 @@
 use crate::bindgen::FunctionBindgen;
 use crate::{
     dealias, int_repr, to_rust_ident, to_upper_camel_case, wasm_type, Direction, ExportKey, FnSig,
-    Identifier, Ownership, RustFlagsRepr, RustWasm, TypeMode,
+    Identifier, InterfaceName, Ownership, RustFlagsRepr, RustWasm, TypeMode,
 };
 use anyhow::Result;
 use heck::*;
@@ -213,9 +213,6 @@ impl InterfaceGenerator<'_> {
             }
         };
         let module_path = crate::compute_module_path(name, &self.resolve, !self.in_import);
-        if let Identifier::Interface(id, _) = self.identifier {
-            self.gen.interface_names.insert(id, module_path.join("::"));
-        }
         (snake, module_path)
     }
 
@@ -1499,26 +1496,30 @@ impl InterfaceGenerator<'_> {
     // }
 
     fn path_to_interface(&self, interface: InterfaceId) -> Option<String> {
-        let mut path = String::new();
-        if let Identifier::Interface(cur, name) = self.identifier {
-            if cur == interface {
-                return None;
-            }
-            if !self.in_import {
-                path.push_str("super::");
-            }
-            match name {
-                WorldKey::Name(_) => {
-                    path.push_str("super::");
+        let InterfaceName { path, remapped } = &self.gen.interface_names[&interface];
+        if *remapped {
+            return Some(path.to_string());
+        } else {
+            let mut full_path = String::new();
+            if let Identifier::Interface(cur, name) = self.identifier {
+                if cur == interface {
+                    return None;
                 }
-                WorldKey::Interface(_) => {
-                    path.push_str("super::super::super::");
+                if !self.in_import {
+                    full_path.push_str("super::");
+                }
+                match name {
+                    WorldKey::Name(_) => {
+                        full_path.push_str("super::");
+                    }
+                    WorldKey::Interface(_) => {
+                        full_path.push_str("super::super::super::");
+                    }
                 }
             }
+            full_path.push_str(&path);
+            Some(full_path)
         }
-        let name = &self.gen.interface_names[&interface];
-        path.push_str(name);
-        Some(path)
     }
 
     fn push_vec_name(&mut self) {

diff --git a/crates/rust/src/lib.rs b/crates/rust/src/lib.rs
@@ -33,6 +33,15 @@ struct ResourceInfo {
     owned: bool,
 }
 
+struct InterfaceName {
+    /// True when this interface name has been remapped through the use of `with` in the `bindgen!`
+    /// macro invocation.
+    remapped: bool,
+
+    /// The string name for this interface.
+    path: String,
+}
+
 #[derive(Default)]
 struct RustWasm {
     types: Types,
@@ -41,21 +50,29 @@ struct RustWasm {
     import_modules: Vec<(String, Vec<String>)>,
     export_modules: Vec<(String, Vec<String>)>,
     skip: HashSet<String>,
-    interface_names: HashMap<InterfaceId, String>,
+    interface_names: HashMap<InterfaceId, InterfaceName>,
     resources: HashMap<TypeId, ResourceInfo>,
     import_funcs_called: bool,
+    with_name_counter: usize,
+}
+
+#[cfg(feature = "clap")]
+fn iterate_hashmap_string(s: &str) -> impl Iterator<Item = Result<(&str, &str), String>> {
+    s.split(',').map(move |entry| {
+        entry.split_once('=').ok_or_else(|| {
+            format!("expected string of form `<key>=<value>[,<key>=<value>...]`; got `{s}`")
+        })
+    })
 }
 
 #[cfg(feature = "clap")]
 fn parse_exports(s: &str) -> Result<HashMap<ExportKey, String>, String> {
     if s.is_empty() {
         Ok(HashMap::default())
     } else {
-        s.split(',')
+        iterate_hashmap_string(s)
             .map(|entry| {
-                let (key, value) = entry.split_once('=').ok_or_else(|| {
-                    format!("expected string of form `<key>=<value>[,<key>=<value>...]`; got `{s}`")
-                })?;
+                let (key, value) = entry?;
                 Ok((
                     match key {
                         "world" => ExportKey::World,
@@ -74,6 +91,20 @@ pub enum ExportKey {
     Name(String),
 }
 
+#[cfg(feature = "clap")]
+fn parse_with(s: &str) -> Result<HashMap<String, String>, String> {
+    if s.is_empty() {
+        Ok(HashMap::default())
+    } else {
+        iterate_hashmap_string(s)
+            .map(|entry| {
+                let (key, value) = entry?;
+                Ok((key.to_owned(), value.to_owned()))
+            })
+            .collect()
+    }
+}
+
 #[derive(Default, Debug, Clone)]
 #[cfg_attr(feature = "clap", derive(clap::Args))]
 pub struct Opts {
@@ -149,6 +180,10 @@ pub struct Opts {
     /// These derive attributes will be added to any generated structs or enums
     #[cfg_attr(feature = "clap", arg(long = "additional_derive_attribute", short = 'd', default_values_t = Vec::<String>::new()))]
     pub additional_derive_attributes: Vec<String>,
+
+    /// Remapping of interface names to rust module names.
+    #[cfg_attr(feature = "clap", arg(long, value_parser = parse_with, default_value = ""))]
+    pub with: HashMap<String, String>,
 }
 
 impl Opts {
@@ -248,6 +283,37 @@ impl RustWasm {
         }
         bail!("expected `exports` map to contain key `{key}`")
     }
+
+    fn name_interface(
+        &mut self,
+        resolve: &Resolve,
+        id: InterfaceId,
+        name: &WorldKey,
+        is_export: bool,
+    ) -> bool {
+        let with_name = resolve.name_world_key(name);
+        let entry = if let Some(remapped_path) = self.opts.with.get(&with_name) {
+            let name = format!("__with_name{}", self.with_name_counter);
+            self.with_name_counter += 1;
+            uwriteln!(self.src, "use {remapped_path} as {name};");
+            InterfaceName {
+                remapped: true,
+                path: name,
+            }
+        } else {
+            let path = compute_module_path(name, resolve, is_export).join("::");
+
+            InterfaceName {
+                remapped: false,
+                path,
+            }
+        };
+
+        let remapped = entry.remapped;
+        self.interface_names.insert(id, entry);
+
+        remapped
+    }
 }
 
 /// If the package `id` is the only package with its namespace/name combo
@@ -316,6 +382,9 @@ impl WorldGenerator for RustWasm {
             true,
         );
         let (snake, module_path) = gen.start_append_submodule(name);
+        if gen.gen.name_interface(resolve, id, name, false) {
+            return;
+        }
         gen.types(id);
 
         gen.generate_imports(resolve.interfaces[id].functions.values());
@@ -349,6 +418,9 @@ impl WorldGenerator for RustWasm {
     ) -> Result<()> {
         let mut gen = self.interface(Identifier::Interface(id, name), None, resolve, false);
         let (snake, module_path) = gen.start_append_submodule(name);
+        if gen.gen.name_interface(resolve, id, name, true) {
+            return Ok(());
+        }
         gen.types(id);
         gen.generate_exports(resolve.interfaces[id].functions.values())?;
         gen.finish_append_submodule(&snake, module_path);