Skip to content

Commit 80621a2

Browse files
yordisyordis
committed
Add WitMap trait and configurable map_type for Rust bindings
Instead of hardcoding HashMap/BTreeMap, introduce a WitMap trait that generated code delegates to for map construction, insertion, and length. This lets users swap in their own map type via the new `map_type` bindgen option. The default is BTreeMap (always, regardless of std).
1 parent a9a8dd6 commit 80621a2

5 files changed

Lines changed: 90 additions & 8 deletions

File tree

‎crates/guest-rust/macro/src/lib.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,7 @@ impl Parse for Config {
111111
Opt::Ownership(ownership) => opts.ownership = ownership,
112112
Opt::Skip(list) => opts.skip.extend(list.iter().map(|i| i.value())),
113113
Opt::RuntimePath(path) => opts.runtime_path = Some(path.value()),
114+
Opt::MapType(path) => opts.map_type = Some(path.value()),
114115
Opt::BitflagsPath(path) => opts.bitflags_path = Some(path.value()),
115116
Opt::Stubs => {
116117
opts.stubs = true;
@@ -304,6 +305,7 @@ mod kw {
304305
syn::custom_keyword!(inline);
305306
syn::custom_keyword!(ownership);
306307
syn::custom_keyword!(runtime_path);
308+
syn::custom_keyword!(map_type);
307309
syn::custom_keyword!(bitflags_path);
308310
syn::custom_keyword!(exports);
309311
syn::custom_keyword!(stubs);
@@ -383,6 +385,7 @@ enum Opt {
383385
Skip(Vec<syn::LitStr>),
384386
Ownership(Ownership),
385387
RuntimePath(syn::LitStr),
388+
MapType(syn::LitStr),
386389
BitflagsPath(syn::LitStr),
387390
Stubs,
388391
ExportPrefix(syn::LitStr),
@@ -486,6 +489,10 @@ impl Parse for Opt {
486489
input.parse::<kw::runtime_path>()?;
487490
input.parse::<Token![:]>()?;
488491
Ok(Opt::RuntimePath(input.parse()?))
492+
} else if l.peek(kw::map_type) {
493+
input.parse::<kw::map_type>()?;
494+
input.parse::<Token![:]>()?;
495+
Ok(Opt::MapType(input.parse()?))
489496
} else if l.peek(kw::bitflags_path) {
490497
input.parse::<kw::bitflags_path>()?;
491498
input.parse::<Token![:]>()?;

‎crates/guest-rust/src/rt/mod.rs‎

Lines changed: 51 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -67,11 +67,59 @@ pub mod bitflags {
6767
pub use crate::bitflags;
6868
}
6969

70-
#[cfg(feature = "std")]
71-
pub type Map<K, V> = std::collections::HashMap<K, V>;
72-
#[cfg(not(feature = "std"))]
7370
pub type Map<K, V> = alloc::collections::BTreeMap<K, V>;
7471

72+
/// Trait abstracting the map operations needed by generated WIT bindings.
73+
///
74+
/// Generated code delegates to these methods rather than calling inherent
75+
/// methods on a concrete type, so users can swap in their own map
76+
/// implementation via the `map_type` bindgen option.
77+
///
78+
/// The concrete map type must also implement `IntoIterator` (both owned
79+
/// and by-ref) yielding key/value pairs.
80+
pub trait WitMap<K, V> {
81+
fn wit_map_new(capacity: usize) -> Self;
82+
fn wit_map_push(&mut self, key: K, value: V);
83+
fn wit_map_len(&self) -> usize;
84+
}
85+
86+
impl<'a, K, V, T: WitMap<K, V>> WitMap<K, V> for &'a T {
87+
fn wit_map_new(_capacity: usize) -> Self {
88+
unreachable!()
89+
}
90+
fn wit_map_push(&mut self, _key: K, _value: V) {
91+
unreachable!()
92+
}
93+
fn wit_map_len(&self) -> usize {
94+
T::wit_map_len(self)
95+
}
96+
}
97+
98+
impl<K: Ord, V> WitMap<K, V> for alloc::collections::BTreeMap<K, V> {
99+
fn wit_map_new(_capacity: usize) -> Self {
100+
alloc::collections::BTreeMap::new()
101+
}
102+
fn wit_map_push(&mut self, key: K, value: V) {
103+
self.insert(key, value);
104+
}
105+
fn wit_map_len(&self) -> usize {
106+
self.len()
107+
}
108+
}
109+
110+
#[cfg(feature = "std")]
111+
impl<K: core::hash::Hash + Eq, V> WitMap<K, V> for std::collections::HashMap<K, V> {
112+
fn wit_map_new(capacity: usize) -> Self {
113+
std::collections::HashMap::with_capacity(capacity)
114+
}
115+
fn wit_map_push(&mut self, key: K, value: V) {
116+
self.insert(key, value);
117+
}
118+
fn wit_map_len(&self) -> usize {
119+
self.len()
120+
}
121+
}
122+
75123
/// For more information about this see `./ci/rebuild-libwit-bindgen-cabi.sh`.
76124
#[cfg(not(target_env = "p2"))]
77125
mod wit_bindgen_cabi_realloc;

‎crates/rust/src/bindgen.rs‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -770,6 +770,7 @@ impl Bindgen for FunctionBindgen<'_, '_> {
770770
} => {
771771
let alloc = self.r#gen.path_to_std_alloc_module();
772772
let rt = self.r#gen.r#gen.runtime_path().to_string();
773+
let wit_map = self.r#gen.r#gen.wit_map_path();
773774
let body = self.blocks.pop().unwrap();
774775
let tmp = self.tmp();
775776
let map = format!("map{tmp}");
@@ -781,7 +782,7 @@ impl Bindgen for FunctionBindgen<'_, '_> {
781782
"let {map} = {operand0};\n",
782783
operand0 = operands[0]
783784
));
784-
self.push_str(&format!("let {len} = {map}.len();\n"));
785+
uwriteln!(self.src, "let {len} = {wit_map}::wit_map_len(&{map});");
785786
let entry = self.map_entry_layout(key, value);
786787
self.push_str(&format!(
787788
"let {layout} = {alloc}::Layout::from_size_align({len} * {}, {}).unwrap();\n",
@@ -876,7 +877,7 @@ impl Bindgen for FunctionBindgen<'_, '_> {
876877
let len = format!("len{tmp}");
877878
let base = format!("base{tmp}");
878879
let result = format!("result{tmp}");
879-
let map = format!("{}::Map", self.r#gen.r#gen.runtime_path());
880+
let wit_map = self.r#gen.r#gen.wit_map_path();
880881
self.push_str(&format!(
881882
"let {base} = {operand0};\n",
882883
operand0 = operands[0]
@@ -885,15 +886,15 @@ impl Bindgen for FunctionBindgen<'_, '_> {
885886
"let {len} = {operand1};\n",
886887
operand1 = operands[1]
887888
));
888-
self.push_str(&format!("let mut {result} = {map}::new();\n"));
889+
uwriteln!(self.src, "let mut {result} = {wit_map}::wit_map_new({len});");
889890
uwriteln!(self.src, "for i in 0..{len} {{");
890891
uwriteln!(
891892
self.src,
892893
"let base = {base}.add(i * {size});",
893894
size = entry.size.format(POINTER_SIZE_EXPRESSION),
894895
);
895896
uwriteln!(self.src, "let (map_key, map_value) = {body};");
896-
uwriteln!(self.src, "{result}.insert(map_key, map_value);");
897+
uwriteln!(self.src, "{wit_map}::wit_map_push(&mut {result}, map_key, map_value);");
897898
uwriteln!(self.src, "}}");
898899
results.push(result);
899900
let dealloc = self.r#gen.path_to_cabi_dealloc();

‎crates/rust/src/interface.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1949,7 +1949,7 @@ unsafe fn call_import(&mut self, _params: Self::ParamsLower, _results: *mut u8)
19491949
self.push_str(" ");
19501950
}
19511951
}
1952-
let path = format!("{}::Map", self.r#gen.runtime_path());
1952+
let path = self.r#gen.map_type_path();
19531953
self.push_str(&path);
19541954
self.push_str("::<");
19551955
self.print_ty(key, key_mode);

‎crates/rust/src/lib.rs‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,18 @@ pub struct Opts {
199199
#[cfg_attr(feature = "clap", arg(long, value_name = "PATH"))]
200200
pub runtime_path: Option<String>,
201201

202+
/// The optional path to the map type to use for WIT `map<K, V>`.
203+
///
204+
/// The specified type must accept two type parameters `<K, V>` and
205+
/// implement the `WitMap<K, V>` trait from the wit-bindgen runtime.
206+
/// It must also implement `IntoIterator<Item = (K, V)>` (owned) and
207+
/// its reference must implement `IntoIterator` yielding key/value
208+
/// pairs.
209+
///
210+
/// Defaults to `{runtime_path}::Map` which is `BTreeMap`.
211+
#[cfg_attr(feature = "clap", arg(long, value_name = "PATH"))]
212+
pub map_type: Option<String>,
213+
202214
/// The optional path to the bitflags crate to use.
203215
///
204216
/// This defaults to `wit_bindgen::bitflags`.
@@ -417,6 +429,17 @@ impl RustWasm {
417429
.unwrap_or("wit_bindgen::rt")
418430
}
419431

432+
fn map_type_path(&self) -> String {
433+
self.opts
434+
.map_type
435+
.clone()
436+
.unwrap_or_else(|| format!("{}::Map", self.runtime_path()))
437+
}
438+
439+
fn wit_map_path(&self) -> String {
440+
format!("{}::WitMap", self.runtime_path())
441+
}
442+
420443
fn bitflags_path(&self) -> String {
421444
self.opts
422445
.bitflags_path
@@ -1057,6 +1080,9 @@ impl WorldGenerator for RustWasm {
10571080
if let Some(runtime_path) = &self.opts.runtime_path {
10581081
uwriteln!(self.src_preamble, "// * runtime_path: {:?}", runtime_path);
10591082
}
1083+
if let Some(map_type) = &self.opts.map_type {
1084+
uwriteln!(self.src_preamble, "// * map_type: {:?}", map_type);
1085+
}
10601086
if let Some(bitflags_path) = &self.opts.bitflags_path {
10611087
uwriteln!(
10621088
self.src_preamble,

0 commit comments

Comments
 (0)