wasmtime_environ/
key.rs

1//! Keys for identifying functions during compilation, in call graphs, and when
2//! resolving relocations.
3
4#[cfg(feature = "component-model")]
5use crate::component;
6use crate::{
7    BuiltinFunctionIndex, DefinedFuncIndex, HostCall, ModuleInternedTypeIndex, StaticModuleIndex,
8};
9use core::{cmp, fmt};
10use serde_derive::{Deserialize, Serialize};
11
12/// The kind of a function that is being compiled, linked, or otherwise
13/// referenced.
14///
15/// This is like a `FuncKey` but without any payload values.
16#[repr(u32)]
17#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
18#[cfg_attr(test, derive(arbitrary::Arbitrary))]
19pub enum FuncKeyKind {
20    /// A Wasm-defined function.
21    DefinedWasmFunction = FuncKey::new_kind(0b000),
22
23    /// A trampoline from an array-caller to the given Wasm-callee.
24    ArrayToWasmTrampoline = FuncKey::new_kind(0b001),
25
26    /// A trampoline from a Wasm-caller to an array-callee of the given type.
27    WasmToArrayTrampoline = FuncKey::new_kind(0b010),
28
29    /// A trampoline from a Wasm-caller to the given builtin.
30    WasmToBuiltinTrampoline = FuncKey::new_kind(0b011),
31
32    /// A Pulley-specific host call.
33    PulleyHostCall = FuncKey::new_kind(0b100),
34
35    /// A Wasm-caller to component builtin trampoline.
36    #[cfg(feature = "component-model")]
37    ComponentTrampoline = FuncKey::new_kind(0b101),
38
39    /// A Wasm-caller to array-callee `resource.drop` trampoline.
40    #[cfg(feature = "component-model")]
41    ResourceDropTrampoline = FuncKey::new_kind(0b110),
42}
43
44impl From<FuncKeyKind> for u32 {
45    fn from(kind: FuncKeyKind) -> Self {
46        kind as u32
47    }
48}
49
50impl FuncKeyKind {
51    /// Get this kind's raw representation.
52    pub fn into_raw(self) -> u32 {
53        self.into()
54    }
55
56    /// Construct a `FuncKind` from its raw representation.
57    ///
58    /// Panics when given invalid raw representations.
59    pub fn from_raw(raw: u32) -> Self {
60        match raw {
61            x if x == Self::DefinedWasmFunction.into() => Self::DefinedWasmFunction,
62            x if x == Self::ArrayToWasmTrampoline.into() => Self::ArrayToWasmTrampoline,
63            x if x == Self::WasmToArrayTrampoline.into() => Self::WasmToArrayTrampoline,
64            x if x == Self::WasmToBuiltinTrampoline.into() => Self::WasmToBuiltinTrampoline,
65            x if x == Self::PulleyHostCall.into() => Self::PulleyHostCall,
66
67            #[cfg(feature = "component-model")]
68            x if x == Self::ComponentTrampoline.into() => Self::ComponentTrampoline,
69            #[cfg(feature = "component-model")]
70            x if x == Self::ResourceDropTrampoline.into() => Self::ResourceDropTrampoline,
71
72            _ => panic!("invalid raw value passed to `FuncKind::from_raw`: {raw}"),
73        }
74    }
75}
76
77/// The namespace half of a `FuncKey`.
78///
79/// This is an opaque combination of the key's kind and module index, if any.
80#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
81pub struct FuncKeyNamespace(u32);
82
83impl fmt::Debug for FuncKeyNamespace {
84    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
85        struct Hex<T: fmt::LowerHex>(T);
86        impl<T: fmt::LowerHex> fmt::Debug for Hex<T> {
87            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
88                write!(f, "{:#x}", self.0)
89            }
90        }
91        f.debug_struct("FuncKeyNamespace")
92            .field("raw", &Hex(self.0))
93            .field("kind", &self.kind())
94            .field("module", &self.module())
95            .finish()
96    }
97}
98
99impl From<FuncKeyNamespace> for u32 {
100    fn from(ns: FuncKeyNamespace) -> Self {
101        ns.0
102    }
103}
104
105impl FuncKeyNamespace {
106    /// Get this `FuncNamespace`'s raw representation.
107    pub fn into_raw(self) -> u32 {
108        self.0
109    }
110
111    /// Construct a `FuncNamespace` from its raw representation.
112    ///
113    /// Panics when given invalid raw representations.
114    pub fn from_raw(raw: u32) -> Self {
115        match FuncKeyKind::from_raw(raw & FuncKey::KIND_MASK) {
116            FuncKeyKind::DefinedWasmFunction | FuncKeyKind::ArrayToWasmTrampoline => Self(raw),
117            FuncKeyKind::WasmToArrayTrampoline
118            | FuncKeyKind::WasmToBuiltinTrampoline
119            | FuncKeyKind::PulleyHostCall => {
120                assert_eq!(raw & FuncKey::MODULE_MASK, 0);
121                Self(raw)
122            }
123
124            #[cfg(feature = "component-model")]
125            FuncKeyKind::ComponentTrampoline => {
126                let _ = Abi::from_raw(raw & FuncKey::MODULE_MASK);
127                Self(raw)
128            }
129
130            #[cfg(feature = "component-model")]
131            FuncKeyKind::ResourceDropTrampoline => {
132                assert_eq!(raw & FuncKey::MODULE_MASK, 0);
133                Self(raw)
134            }
135        }
136    }
137
138    /// Get this `FuncNamespace`'s kind.
139    pub fn kind(&self) -> FuncKeyKind {
140        let raw = self.0 & FuncKey::KIND_MASK;
141        FuncKeyKind::from_raw(raw)
142    }
143
144    fn module(&self) -> u32 {
145        self.0 & FuncKey::MODULE_MASK
146    }
147}
148
149/// The index half of a `FuncKey`.
150#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
151pub struct FuncKeyIndex(u32);
152
153impl From<FuncKeyIndex> for u32 {
154    fn from(index: FuncKeyIndex) -> Self {
155        index.0
156    }
157}
158
159impl FuncKeyIndex {
160    /// Get this index's raw representation.
161    pub fn into_raw(self) -> u32 {
162        self.0
163    }
164
165    /// Construct a `FuncKeyIndex` from its raw representation.
166    ///
167    /// Invalid raw representations will not be caught eagerly, but will cause
168    /// panics when paired with a `FuncKeyNamespace` to create a whole
169    /// `FuncKey`.
170    pub fn from_raw(raw: u32) -> Self {
171        FuncKeyIndex(raw)
172    }
173}
174
175/// ABI signature of functions that are generated here.
176#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
177#[cfg_attr(test, derive(arbitrary::Arbitrary))]
178pub enum Abi {
179    /// The "wasm" ABI, or suitable to be a `wasm_call` field of a `VMFuncRef`.
180    Wasm = 0,
181    /// The "array" ABI, or suitable to be an `array_call` field.
182    Array = 1,
183}
184
185#[cfg(feature = "component-model")]
186impl Abi {
187    fn from_raw(raw: u32) -> Self {
188        match raw {
189            x if x == Self::Wasm.into_raw() => Self::Wasm,
190            x if x == Self::Array.into_raw() => Self::Array,
191            _ => panic!("invalid raw representation passed to `Abi::from_raw`: {raw}"),
192        }
193    }
194
195    fn into_raw(self) -> u32 {
196        (self as u8).into()
197    }
198}
199
200/// A sortable, comparable function key for compilation output, call graph
201/// edges, and relocations.
202#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
203pub enum FuncKey {
204    /// A Wasm-defined function.
205    DefinedWasmFunction(StaticModuleIndex, DefinedFuncIndex),
206
207    /// A trampoline from an array-caller to the given Wasm-callee.
208    ArrayToWasmTrampoline(StaticModuleIndex, DefinedFuncIndex),
209
210    /// A trampoline from a Wasm-caller to an array-callee of the given type.
211    WasmToArrayTrampoline(ModuleInternedTypeIndex),
212
213    /// A trampoline from a Wasm-caller to the given builtin.
214    WasmToBuiltinTrampoline(BuiltinFunctionIndex),
215
216    /// A Pulley-specific host call.
217    PulleyHostCall(HostCall),
218
219    /// A Wasm-caller to component builtin trampoline.
220    #[cfg(feature = "component-model")]
221    ComponentTrampoline(Abi, component::TrampolineIndex),
222
223    /// A Wasm-caller to array-callee `resource.drop` trampoline.
224    #[cfg(feature = "component-model")]
225    ResourceDropTrampoline,
226}
227
228impl Ord for FuncKey {
229    fn cmp(&self, other: &Self) -> cmp::Ordering {
230        // Make sure to sort by our raw parts, because `CompiledFunctionsTable`
231        // relies on this for its binary search tables.
232        let raw_self = self.into_raw_parts();
233        let raw_other = other.into_raw_parts();
234        raw_self.cmp(&raw_other)
235    }
236}
237
238impl PartialOrd for FuncKey {
239    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
240        Some(self.cmp(other))
241    }
242}
243
244impl FuncKey {
245    const KIND_BITS: u32 = 3;
246    const KIND_OFFSET: u32 = 32 - Self::KIND_BITS;
247    const KIND_MASK: u32 = ((1 << Self::KIND_BITS) - 1) << Self::KIND_OFFSET;
248    const MODULE_MASK: u32 = !Self::KIND_MASK;
249
250    const fn new_kind(kind: u32) -> u32 {
251        assert!(kind < (1 << Self::KIND_BITS));
252        kind << Self::KIND_OFFSET
253    }
254
255    /// Split this key into its namespace and index halves.
256    #[inline]
257    pub fn into_parts(self) -> (FuncKeyNamespace, FuncKeyIndex) {
258        let (namespace, index) = match self {
259            FuncKey::DefinedWasmFunction(module, def_func) => {
260                assert_eq!(module.as_u32() & Self::KIND_MASK, 0);
261                let namespace = FuncKeyKind::DefinedWasmFunction.into_raw() | module.as_u32();
262                let index = def_func.as_u32();
263                (namespace, index)
264            }
265            FuncKey::ArrayToWasmTrampoline(module, def_func) => {
266                assert_eq!(module.as_u32() & Self::KIND_MASK, 0);
267                let namespace = FuncKeyKind::ArrayToWasmTrampoline.into_raw() | module.as_u32();
268                let index = def_func.as_u32();
269                (namespace, index)
270            }
271            FuncKey::WasmToArrayTrampoline(ty) => {
272                let namespace = FuncKeyKind::WasmToArrayTrampoline.into_raw();
273                let index = ty.as_u32();
274                (namespace, index)
275            }
276            FuncKey::WasmToBuiltinTrampoline(builtin) => {
277                let namespace = FuncKeyKind::WasmToBuiltinTrampoline.into_raw();
278                let index = builtin.index();
279                (namespace, index)
280            }
281            FuncKey::PulleyHostCall(host_call) => {
282                let namespace = FuncKeyKind::PulleyHostCall.into_raw();
283                let index = host_call.index();
284                (namespace, index)
285            }
286
287            #[cfg(feature = "component-model")]
288            FuncKey::ComponentTrampoline(abi, trampoline) => {
289                let abi = abi.into_raw();
290                assert_eq!(abi & Self::KIND_MASK, 0);
291                let namespace = FuncKeyKind::ComponentTrampoline.into_raw() | abi;
292                let index = trampoline.as_u32();
293                (namespace, index)
294            }
295            #[cfg(feature = "component-model")]
296            FuncKey::ResourceDropTrampoline => {
297                let namespace = FuncKeyKind::ResourceDropTrampoline.into_raw();
298                let index = 0;
299                (namespace, index)
300            }
301        };
302        (FuncKeyNamespace(namespace), FuncKeyIndex(index))
303    }
304
305    /// Get this key's kind.
306    pub fn kind(self) -> FuncKeyKind {
307        self.namespace().kind()
308    }
309
310    /// Get this key's namespace.
311    pub fn namespace(self) -> FuncKeyNamespace {
312        self.into_parts().0
313    }
314
315    /// Get this key's index.
316    pub fn index(self) -> FuncKeyIndex {
317        self.into_parts().1
318    }
319
320    /// Get the raw, underlying `(namespace, index)` representation of this
321    /// compilation key.
322    ///
323    /// The resulting values should only be used for (eventually) calling
324    /// `FuncKey::from_raw_parts` or `FuncKey{Namespace,Index}::from_raw`.
325    //
326    // NB: We use two `u32`s to exactly match
327    // `cranelift_codegen::ir::UserExternalName` and ensure that we can map
328    // one-to-one between that and `FuncKey`.
329    pub fn into_raw_parts(self) -> (u32, u32) {
330        let (ns, index) = self.into_parts();
331        (ns.into_raw(), index.into_raw())
332    }
333
334    /// Create a key from its namespace and index parts.
335    ///
336    /// Should only be called with namespaces and indices that are ultimately
337    /// derived from the same key. For example, if you attempt to pair an index
338    /// and namespace that come from different keys, that may panic. If it
339    /// happens not to panic, you'll end up with a valid key that names an
340    /// arbitrary function in the given namespace, but that function probably
341    /// does not actually exist in the compilation artifact.
342    pub fn from_parts(namespace: FuncKeyNamespace, index: FuncKeyIndex) -> Self {
343        Self::from_raw_parts(namespace.into_raw(), index.into_raw())
344    }
345
346    /// Create a key from its raw, underlying representation.
347    ///
348    /// Should only be given the results of a previous call to
349    /// `FuncKey::into_raw_parts`.
350    ///
351    /// Panics when given invalid raw parts.
352    pub fn from_raw_parts(a: u32, b: u32) -> Self {
353        match FuncKeyKind::from_raw(a & Self::KIND_MASK) {
354            FuncKeyKind::DefinedWasmFunction => {
355                let module = StaticModuleIndex::from_u32(a & Self::MODULE_MASK);
356                let def_func = DefinedFuncIndex::from_u32(b);
357                Self::DefinedWasmFunction(module, def_func)
358            }
359            FuncKeyKind::ArrayToWasmTrampoline => {
360                let module = StaticModuleIndex::from_u32(a & Self::MODULE_MASK);
361                let def_func = DefinedFuncIndex::from_u32(b);
362                Self::ArrayToWasmTrampoline(module, def_func)
363            }
364            FuncKeyKind::WasmToArrayTrampoline => {
365                assert_eq!(a & Self::MODULE_MASK, 0);
366                let ty = ModuleInternedTypeIndex::from_u32(b);
367                Self::WasmToArrayTrampoline(ty)
368            }
369            FuncKeyKind::WasmToBuiltinTrampoline => {
370                assert_eq!(a & Self::MODULE_MASK, 0);
371                let builtin = BuiltinFunctionIndex::from_u32(b);
372                Self::WasmToBuiltinTrampoline(builtin)
373            }
374            FuncKeyKind::PulleyHostCall => {
375                assert_eq!(a & Self::MODULE_MASK, 0);
376                let host_call = HostCall::from_index(b);
377                Self::PulleyHostCall(host_call)
378            }
379
380            #[cfg(feature = "component-model")]
381            FuncKeyKind::ComponentTrampoline => {
382                let abi = Abi::from_raw(a & Self::MODULE_MASK);
383                let trampoline = component::TrampolineIndex::from_u32(b);
384                Self::ComponentTrampoline(abi, trampoline)
385            }
386            #[cfg(feature = "component-model")]
387            FuncKeyKind::ResourceDropTrampoline => {
388                assert_eq!(a & Self::MODULE_MASK, 0);
389                assert_eq!(b, 0);
390                Self::ResourceDropTrampoline
391            }
392        }
393    }
394
395    /// Unwrap a `FuncKey::DefinedWasmFunction` or else panic.
396    pub fn unwrap_defined_wasm_function(self) -> (StaticModuleIndex, DefinedFuncIndex) {
397        match self {
398            Self::DefinedWasmFunction(module, def_func) => (module, def_func),
399            _ => panic!("`FuncKey::unwrap_defined_wasm_function` called on {self:?}"),
400        }
401    }
402
403    /// Unwrap a `FuncKey::ArrayToWasmTrampoline` or else panic.
404    pub fn unwrap_array_to_wasm_trampoline(self) -> (StaticModuleIndex, DefinedFuncIndex) {
405        match self {
406            Self::ArrayToWasmTrampoline(module, def_func) => (module, def_func),
407            _ => panic!("`FuncKey::unwrap_array_to_wasm_trampoline` called on {self:?}"),
408        }
409    }
410
411    /// Unwrap a `FuncKey::WasmToArrayTrampoline` or else panic.
412    pub fn unwrap_wasm_to_array_trampoline(self) -> ModuleInternedTypeIndex {
413        match self {
414            Self::WasmToArrayTrampoline(ty) => ty,
415            _ => panic!("`FuncKey::unwrap_wasm_to_array_trampoline` called on {self:?}"),
416        }
417    }
418
419    /// Unwrap a `FuncKey::WasmToBuiltinTrampoline` or else panic.
420    pub fn unwrap_wasm_to_builtin_trampoline(self) -> BuiltinFunctionIndex {
421        match self {
422            Self::WasmToBuiltinTrampoline(builtin) => builtin,
423            _ => panic!("`FuncKey::unwrap_wasm_to_builtin_trampoline` called on {self:?}"),
424        }
425    }
426
427    /// Unwrap a `FuncKey::PulleyHostCall` or else panic.
428    pub fn unwrap_pulley_host_call(self) -> HostCall {
429        match self {
430            Self::PulleyHostCall(host_call) => host_call,
431            _ => panic!("`FuncKey::unwrap_pulley_host_call` called on {self:?}"),
432        }
433    }
434
435    /// Unwrap a `FuncKey::ComponentTrampoline` or else panic.
436    #[cfg(feature = "component-model")]
437    pub fn unwrap_component_trampoline(self) -> (crate::Abi, component::TrampolineIndex) {
438        match self {
439            Self::ComponentTrampoline(abi, trampoline) => (abi, trampoline),
440            _ => panic!("`FuncKey::unwrap_component_trampoline` called on {self:?}"),
441        }
442    }
443
444    /// Unwrap a `FuncKey::ResourceDropTrampoline` or else panic.
445    #[cfg(feature = "component-model")]
446    pub fn unwrap_resource_drop_trampoline(self) {
447        match self {
448            Self::ResourceDropTrampoline => {}
449            _ => panic!("`FuncKey::unwrap_resource_drop_trampoline` called on {self:?}"),
450        }
451    }
452}