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