wasmtime_environ/compile/
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};
9
10/// A sortable, comparable function key for compilation output, call graph
11/// edges, and relocations.
12#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
13pub enum FuncKey {
14    /// A Wasm-defined function.
15    DefinedWasmFunction(StaticModuleIndex, DefinedFuncIndex),
16
17    /// A trampoline from an array-caller to the given Wasm-callee.
18    ArrayToWasmTrampoline(StaticModuleIndex, DefinedFuncIndex),
19
20    /// A trampoline from a Wasm-caller to an array-callee of the given type.
21    WasmToArrayTrampoline(ModuleInternedTypeIndex),
22
23    /// A trampoline from a Wasm-caller to the given builtin.
24    WasmToBuiltinTrampoline(BuiltinFunctionIndex),
25
26    /// A Pulley-specific host call.
27    PulleyHostCall(HostCall),
28
29    /// A Wasm-caller to component builtin trampoline.
30    #[cfg(feature = "component-model")]
31    ComponentTrampoline(component::TrampolineIndex),
32
33    /// A Wasm-caller to array-callee `resource.drop` trampoline.
34    #[cfg(feature = "component-model")]
35    ResourceDropTrampoline,
36}
37
38impl FuncKey {
39    const KIND_BITS: u32 = 3;
40    const KIND_OFFSET: u32 = 32 - Self::KIND_BITS;
41    const KIND_MASK: u32 = ((1 << Self::KIND_BITS) - 1) << Self::KIND_OFFSET;
42    const MODULE_MASK: u32 = !Self::KIND_MASK;
43
44    const fn new_kind(kind: u32) -> u32 {
45        assert!(kind < (1 << Self::KIND_BITS));
46        kind << Self::KIND_OFFSET
47    }
48
49    const DEFINED_WASM_FUNCTION_KIND: u32 = Self::new_kind(0);
50    const ARRAY_TO_WASM_TRAMPOLINE_KIND: u32 = Self::new_kind(1);
51    const WASM_TO_ARRAY_TRAMPOLINE_KIND: u32 = Self::new_kind(2);
52    const WASM_TO_BUILTIN_TRAMPOLINE_KIND: u32 = Self::new_kind(3);
53    const PULLEY_HOST_CALL_KIND: u32 = Self::new_kind(4);
54
55    #[cfg(feature = "component-model")]
56    const COMPONENT_TRAMPOLINE_KIND: u32 = Self::new_kind(5);
57    #[cfg(feature = "component-model")]
58    const RESOURCE_DROP_TRAMPOLINE_KIND: u32 = Self::new_kind(6);
59
60    /// Get the raw, underlying representation of this compilation key.
61    ///
62    /// The resulting values should only be used for (eventually) calling
63    /// `CompileKey::from_raw_parts`.
64    //
65    // NB: We use two `u32`s to exactly match
66    // `cranelift_codegen::ir::UserExternalName` and ensure that we can map
67    // one-to-one between that and `FuncKey`.
68    pub fn into_raw_parts(self) -> (u32, u32) {
69        match self {
70            FuncKey::DefinedWasmFunction(module, def_func) => {
71                assert_eq!(module.as_u32() & Self::KIND_MASK, 0);
72                let namespace = Self::DEFINED_WASM_FUNCTION_KIND | module.as_u32();
73                let index = def_func.as_u32();
74                (namespace, index)
75            }
76            FuncKey::ArrayToWasmTrampoline(module, def_func) => {
77                assert_eq!(module.as_u32() & Self::KIND_MASK, 0);
78                let namespace = Self::ARRAY_TO_WASM_TRAMPOLINE_KIND | module.as_u32();
79                let index = def_func.as_u32();
80                (namespace, index)
81            }
82            FuncKey::WasmToArrayTrampoline(ty) => {
83                let namespace = Self::WASM_TO_ARRAY_TRAMPOLINE_KIND;
84                let index = ty.as_u32();
85                (namespace, index)
86            }
87            FuncKey::WasmToBuiltinTrampoline(builtin) => {
88                let namespace = Self::WASM_TO_BUILTIN_TRAMPOLINE_KIND;
89                let index = builtin.index();
90                (namespace, index)
91            }
92            FuncKey::PulleyHostCall(host_call) => {
93                let namespace = Self::PULLEY_HOST_CALL_KIND;
94                let index = host_call.index();
95                (namespace, index)
96            }
97
98            #[cfg(feature = "component-model")]
99            FuncKey::ComponentTrampoline(trampoline) => {
100                let namespace = Self::COMPONENT_TRAMPOLINE_KIND;
101                let index = trampoline.as_u32();
102                (namespace, index)
103            }
104            #[cfg(feature = "component-model")]
105            FuncKey::ResourceDropTrampoline => {
106                let namespace = Self::RESOURCE_DROP_TRAMPOLINE_KIND;
107                let index = 0;
108                (namespace, index)
109            }
110        }
111    }
112
113    /// Create a compilation key from its raw, underlying representation.
114    ///
115    /// Should only be given the results of a previous call to
116    /// `CompileKey::into_raw_parts`.
117    pub fn from_raw_parts(a: u32, b: u32) -> Self {
118        match a & Self::KIND_MASK {
119            Self::DEFINED_WASM_FUNCTION_KIND => {
120                let module = StaticModuleIndex::from_u32(a & Self::MODULE_MASK);
121                let def_func = DefinedFuncIndex::from_u32(b);
122                Self::DefinedWasmFunction(module, def_func)
123            }
124            Self::ARRAY_TO_WASM_TRAMPOLINE_KIND => {
125                let module = StaticModuleIndex::from_u32(a & Self::MODULE_MASK);
126                let def_func = DefinedFuncIndex::from_u32(b);
127                Self::ArrayToWasmTrampoline(module, def_func)
128            }
129            Self::WASM_TO_ARRAY_TRAMPOLINE_KIND => {
130                assert_eq!(a & Self::MODULE_MASK, 0);
131                let ty = ModuleInternedTypeIndex::from_u32(b);
132                Self::WasmToArrayTrampoline(ty)
133            }
134            Self::WASM_TO_BUILTIN_TRAMPOLINE_KIND => {
135                assert_eq!(a & Self::MODULE_MASK, 0);
136                let builtin = BuiltinFunctionIndex::from_u32(b);
137                Self::WasmToBuiltinTrampoline(builtin)
138            }
139            Self::PULLEY_HOST_CALL_KIND => {
140                assert_eq!(a & Self::MODULE_MASK, 0);
141                let host_call = HostCall::from_index(b);
142                Self::PulleyHostCall(host_call)
143            }
144
145            #[cfg(feature = "component-model")]
146            Self::COMPONENT_TRAMPOLINE_KIND => {
147                assert_eq!(a & Self::MODULE_MASK, 0);
148                let trampoline = component::TrampolineIndex::from_u32(b);
149                Self::ComponentTrampoline(trampoline)
150            }
151            #[cfg(feature = "component-model")]
152            Self::RESOURCE_DROP_TRAMPOLINE_KIND => {
153                assert_eq!(a & Self::MODULE_MASK, 0);
154                assert_eq!(b, 0);
155                Self::ResourceDropTrampoline
156            }
157
158            k => panic!(
159                "bad raw parts given to `FuncKey::from_raw_parts` call: ({a}, {b}), kind would be {k}"
160            ),
161        }
162    }
163
164    /// Unwrap a `FuncKey::DefinedWasmFunction` or else panic.
165    pub fn unwrap_defined_wasm_function(self) -> (StaticModuleIndex, DefinedFuncIndex) {
166        match self {
167            Self::DefinedWasmFunction(module, def_func) => (module, def_func),
168            _ => panic!("`FuncKey::unwrap_defined_wasm_function` called on {self:?}"),
169        }
170    }
171
172    /// Unwrap a `FuncKey::ArrayToWasmTrampoline` or else panic.
173    pub fn unwrap_array_to_wasm_trampoline(self) -> (StaticModuleIndex, DefinedFuncIndex) {
174        match self {
175            Self::ArrayToWasmTrampoline(module, def_func) => (module, def_func),
176            _ => panic!("`FuncKey::unwrap_array_to_wasm_trampoline` called on {self:?}"),
177        }
178    }
179
180    /// Unwrap a `FuncKey::WasmToArrayTrampoline` or else panic.
181    pub fn unwrap_wasm_to_array_trampoline(self) -> ModuleInternedTypeIndex {
182        match self {
183            Self::WasmToArrayTrampoline(ty) => ty,
184            _ => panic!("`FuncKey::unwrap_wasm_to_array_trampoline` called on {self:?}"),
185        }
186    }
187
188    /// Unwrap a `FuncKey::WasmToBuiltinTrampoline` or else panic.
189    pub fn unwrap_wasm_to_builtin_trampoline(self) -> BuiltinFunctionIndex {
190        match self {
191            Self::WasmToBuiltinTrampoline(builtin) => builtin,
192            _ => panic!("`FuncKey::unwrap_wasm_to_builtin_trampoline` called on {self:?}"),
193        }
194    }
195
196    /// Unwrap a `FuncKey::PulleyHostCall` or else panic.
197    pub fn unwrap_pulley_host_call(self) -> HostCall {
198        match self {
199            Self::PulleyHostCall(host_call) => host_call,
200            _ => panic!("`FuncKey::unwrap_pulley_host_call` called on {self:?}"),
201        }
202    }
203
204    /// Unwrap a `FuncKey::ComponentTrampoline` or else panic.
205    #[cfg(feature = "component-model")]
206    pub fn unwrap_component_trampoline(self) -> component::TrampolineIndex {
207        match self {
208            Self::ComponentTrampoline(trampoline) => trampoline,
209            _ => panic!("`FuncKey::unwrap_component_trampoline` called on {self:?}"),
210        }
211    }
212
213    /// Unwrap a `FuncKey::ResourceDropTrampoline` or else panic.
214    #[cfg(feature = "component-model")]
215    pub fn unwrap_resource_drop_trampoline(self) {
216        match self {
217            Self::ResourceDropTrampoline => {}
218            _ => panic!("`FuncKey::unwrap_resource_drop_trampoline` called on {self:?}"),
219        }
220    }
221}