numcodecs_wasm_host_reproducible/
engine.rs1use wasm_runtime_layer::{
2 ExportType, ExternType, FuncType, GlobalType, ImportType, MemoryType, TableType,
3 backend::{
4 AsContext, AsContextMut, Export, Extern, Imports, Value, WasmEngine, WasmExternRef,
5 WasmFunc, WasmGlobal, WasmInstance, WasmMemory, WasmModule, WasmStore, WasmStoreContext,
6 WasmStoreContextMut, WasmTable,
7 },
8};
9
10use crate::transform::{
11 instcnt::{InstructionCounterInjecter, PerfWitInterfaces},
12 nan::NaNCanonicaliser,
13};
14
15#[derive(Clone)]
16#[repr(transparent)]
17pub struct ReproducibleEngine<E: WasmEngine>(E);
18
19impl<E: WasmEngine> WasmEngine for ReproducibleEngine<E> {
20 type ExternRef = ReproducibleExternRef<E>;
21 type Func = ReproducibleFunc<E>;
22 type Global = ReproducibleGlobal<E>;
23 type Instance = ReproducibleInstance<E>;
24 type Memory = ReproducibleMemory<E>;
25 type Module = ReproducibleModule<E>;
26 type Store<T: 'static> = ReproducibleStore<T, E>;
27 type StoreContext<'a, T: 'static> = ReproducibleStoreContext<'a, T, E>;
28 type StoreContextMut<'a, T: 'static> = ReproducibleStoreContextMut<'a, T, E>;
29 type Table = ReproducibleTable<E>;
30}
31
32impl<E: WasmEngine> ReproducibleEngine<E> {
33 pub const fn new(engine: E) -> Self {
34 Self(engine)
35 }
36
37 const fn as_ref(&self) -> &E {
38 &self.0
39 }
40
41 const fn from_ref(engine: &E) -> &Self {
42 #[expect(unsafe_code)]
44 unsafe {
45 &*std::ptr::from_ref(engine).cast()
46 }
47 }
48}
49
50#[derive(Clone)]
51#[repr(transparent)]
52pub struct ReproducibleExternRef<E: WasmEngine>(E::ExternRef);
53
54impl<E: WasmEngine> WasmExternRef<ReproducibleEngine<E>> for ReproducibleExternRef<E> {
55 fn new<T: 'static + Send + Sync>(
56 mut ctx: impl AsContextMut<ReproducibleEngine<E>>,
57 object: T,
58 ) -> Self {
59 Self(<E::ExternRef as WasmExternRef<E>>::new(
60 ctx.as_context_mut().as_inner_context_mut(),
61 object,
62 ))
63 }
64
65 fn downcast<'a, 's: 'a, T: 'static, S: 'static>(
66 &'a self,
67 store: ReproducibleStoreContext<'s, S, E>,
68 ) -> anyhow::Result<&'a T> {
69 WasmExternRef::downcast(&self.0, store.0)
70 }
71}
72
73#[derive(Clone)]
74#[repr(transparent)]
75pub struct ReproducibleFunc<E: WasmEngine>(E::Func);
76
77impl<E: WasmEngine> WasmFunc<ReproducibleEngine<E>> for ReproducibleFunc<E> {
78 fn new<T: 'static>(
79 mut ctx: impl AsContextMut<ReproducibleEngine<E>, UserState = T>,
80 ty: FuncType,
81 func: impl 'static
82 + Send
83 + Sync
84 + Fn(
85 ReproducibleStoreContextMut<T, E>,
86 &[Value<ReproducibleEngine<E>>],
87 &mut [Value<ReproducibleEngine<E>>],
88 ) -> anyhow::Result<()>,
89 ) -> Self {
90 Self(<E::Func as WasmFunc<E>>::new(
91 ctx.as_context_mut().as_inner_context_mut(),
92 ty,
93 move |ctx, args, results| {
94 func(
95 ReproducibleStoreContextMut(ctx),
96 from_values(args),
97 from_values_mut(results),
98 )
99 },
100 ))
101 }
102
103 fn ty(&self, ctx: impl AsContext<ReproducibleEngine<E>>) -> FuncType {
104 WasmFunc::ty(&self.0, ctx.as_context().as_inner_context())
105 }
106
107 fn call<T>(
108 &self,
109 mut ctx: impl AsContextMut<ReproducibleEngine<E>>,
110 args: &[Value<ReproducibleEngine<E>>],
111 results: &mut [Value<ReproducibleEngine<E>>],
112 ) -> anyhow::Result<()> {
113 WasmFunc::call::<T>(
114 &self.0,
115 ctx.as_context_mut().as_inner_context_mut(),
116 as_values(args),
117 as_values_mut(results),
118 )
119 }
120}
121
122#[derive(Clone)]
123#[repr(transparent)]
124pub struct ReproducibleGlobal<E: WasmEngine>(E::Global);
125
126impl<E: WasmEngine> WasmGlobal<ReproducibleEngine<E>> for ReproducibleGlobal<E> {
127 fn new(
128 mut ctx: impl AsContextMut<ReproducibleEngine<E>>,
129 value: Value<ReproducibleEngine<E>>,
130 mutable: bool,
131 ) -> Self {
132 Self(<E::Global as WasmGlobal<E>>::new(
133 ctx.as_context_mut().as_inner_context_mut(),
134 into_value(value),
135 mutable,
136 ))
137 }
138
139 fn ty(&self, ctx: impl AsContext<ReproducibleEngine<E>>) -> GlobalType {
140 WasmGlobal::ty(&self.0, ctx.as_context().as_inner_context())
141 }
142
143 fn set(
144 &self,
145 mut ctx: impl AsContextMut<ReproducibleEngine<E>>,
146 new_value: Value<ReproducibleEngine<E>>,
147 ) -> anyhow::Result<()> {
148 WasmGlobal::set(
149 &self.0,
150 ctx.as_context_mut().as_inner_context_mut(),
151 into_value(new_value),
152 )
153 }
154
155 fn get(
156 &self,
157 mut ctx: impl AsContextMut<ReproducibleEngine<E>>,
158 ) -> Value<ReproducibleEngine<E>> {
159 from_value(WasmGlobal::get(
160 &self.0,
161 ctx.as_context_mut().as_inner_context_mut(),
162 ))
163 }
164}
165
166#[derive(Clone)]
167#[repr(transparent)]
168pub struct ReproducibleInstance<E: WasmEngine>(E::Instance);
169
170impl<E: WasmEngine> WasmInstance<ReproducibleEngine<E>> for ReproducibleInstance<E> {
171 fn new(
172 mut store: impl AsContextMut<ReproducibleEngine<E>>,
173 module: &ReproducibleModule<E>,
174 imports: &Imports<ReproducibleEngine<E>>,
175 ) -> anyhow::Result<Self> {
176 let mut new_imports = Imports::new();
177 new_imports.extend(
178 imports
179 .into_iter()
180 .map(|((module, name), value)| ((module, name), into_extern(value))),
181 );
182
183 let PerfWitInterfaces {
184 perf: perf_interface,
185 instruction_counter,
186 } = PerfWitInterfaces::get();
187 new_imports.define(
188 &format!("{perf_interface}"),
189 instruction_counter,
190 Extern::Global(
191 store
192 .as_context_mut()
193 .get_instruction_counter_global()
194 .0
195 .clone(),
196 ),
197 );
198
199 Ok(Self(<E::Instance as WasmInstance<E>>::new(
200 store.as_context_mut().as_inner_context_mut(),
201 &module.0,
202 &new_imports,
203 )?))
204 }
205
206 fn exports(
207 &self,
208 store: impl AsContext<ReproducibleEngine<E>>,
209 ) -> Box<dyn Iterator<Item = Export<ReproducibleEngine<E>>>> {
210 Box::new(
211 WasmInstance::exports(&self.0, store.as_context().as_inner_context()).map(
212 |Export { name, value }| Export {
213 name,
214 value: from_extern(value),
215 },
216 ),
217 )
218 }
219
220 fn get_export(
221 &self,
222 store: impl AsContext<ReproducibleEngine<E>>,
223 name: &str,
224 ) -> Option<Extern<ReproducibleEngine<E>>> {
225 WasmInstance::get_export(&self.0, store.as_context().as_inner_context(), name)
226 .map(from_extern)
227 }
228}
229
230#[derive(Clone)]
231#[repr(transparent)]
232pub struct ReproducibleMemory<E: WasmEngine>(E::Memory);
233
234impl<E: WasmEngine> WasmMemory<ReproducibleEngine<E>> for ReproducibleMemory<E> {
235 fn new(
236 mut ctx: impl AsContextMut<ReproducibleEngine<E>>,
237 ty: MemoryType,
238 ) -> anyhow::Result<Self> {
239 Ok(Self(<E::Memory as WasmMemory<E>>::new(
240 ctx.as_context_mut().as_inner_context_mut(),
241 ty,
242 )?))
243 }
244
245 fn ty(&self, ctx: impl AsContext<ReproducibleEngine<E>>) -> MemoryType {
246 WasmMemory::ty(&self.0, ctx.as_context().as_inner_context())
247 }
248
249 fn grow(
250 &self,
251 mut ctx: impl AsContextMut<ReproducibleEngine<E>>,
252 additional: u32,
253 ) -> anyhow::Result<u32> {
254 WasmMemory::grow(
255 &self.0,
256 ctx.as_context_mut().as_inner_context_mut(),
257 additional,
258 )
259 }
260
261 fn current_pages(&self, ctx: impl AsContext<ReproducibleEngine<E>>) -> u32 {
262 WasmMemory::current_pages(&self.0, ctx.as_context().as_inner_context())
263 }
264
265 fn read(
266 &self,
267 ctx: impl AsContext<ReproducibleEngine<E>>,
268 offset: usize,
269 buffer: &mut [u8],
270 ) -> anyhow::Result<()> {
271 WasmMemory::read(&self.0, ctx.as_context().as_inner_context(), offset, buffer)
272 }
273
274 fn write(
275 &self,
276 mut ctx: impl AsContextMut<ReproducibleEngine<E>>,
277 offset: usize,
278 buffer: &[u8],
279 ) -> anyhow::Result<()> {
280 WasmMemory::write(
281 &self.0,
282 ctx.as_context_mut().as_inner_context_mut(),
283 offset,
284 buffer,
285 )
286 }
287}
288
289pub const DETERMINISTIC_WASM_MODULE_FEATURES: wasmparser::WasmFeaturesInflated =
290 wasmparser::WasmFeaturesInflated {
291 mutable_global: true,
295 saturating_float_to_int: true,
297 sign_extension: true,
299 reference_types: false,
301 multi_value: true,
303 bulk_memory: true,
305 simd: true,
307 relaxed_simd: false,
309 threads: false,
311 shared_everything_threads: false,
313 tail_call: false,
317 floats: false,
319 multi_memory: true,
321 exceptions: false,
323 memory64: false,
326 extended_const: false,
330 component_model: false,
333 function_references: false,
335 memory_control: false,
337 gc: false,
339 custom_page_sizes: false,
343 legacy_exceptions: false,
345 gc_types: false,
348 stack_switching: false,
350 wide_arithmetic: true,
352 cm_values: false,
355 cm_nested_names: false,
358 cm_async: false,
361 cm_async_stackful: false,
364 cm_async_builtins: false,
367 cm_threading: false,
370 cm_error_context: false,
373 cm_fixed_size_list: false,
376 cm_gc: false,
379 call_indirect_overlong: false,
382 bulk_memory_opt: true,
385 };
386
387#[derive(Clone)]
388#[repr(transparent)]
389pub struct ReproducibleModule<E: WasmEngine>(E::Module);
390
391impl<E: WasmEngine> WasmModule<ReproducibleEngine<E>> for ReproducibleModule<E> {
392 fn new(engine: &ReproducibleEngine<E>, bytes: &[u8]) -> anyhow::Result<Self> {
393 let features = wasmparser::WasmFeatures::from(wasmparser::WasmFeaturesInflated {
394 floats: true,
397 ..DETERMINISTIC_WASM_MODULE_FEATURES
398 });
399
400 wasmparser::Validator::new_with_features(features).validate_all(bytes)?;
401
402 let bytes = InstructionCounterInjecter::apply_to_module(bytes, features)?;
404
405 let bytes = NaNCanonicaliser::apply_to_module(&bytes, features)?;
407
408 Ok(Self(<E::Module as WasmModule<E>>::new(
409 engine.as_ref(),
410 bytes.as_slice(),
411 )?))
412 }
413
414 fn exports(&self) -> Box<dyn '_ + Iterator<Item = ExportType<'_>>> {
415 WasmModule::exports(&self.0)
416 }
417
418 fn get_export(&self, name: &str) -> Option<ExternType> {
419 WasmModule::get_export(&self.0, name)
420 }
421
422 fn imports(&self) -> Box<dyn '_ + Iterator<Item = ImportType<'_>>> {
423 WasmModule::imports(&self.0)
424 }
425}
426
427struct StoreData<T, E: WasmEngine> {
428 data: T,
429 instruction_counter: Option<ReproducibleGlobal<E>>,
430}
431
432#[derive(Clone)]
433#[repr(transparent)]
434pub struct ReproducibleStore<T: 'static, E: WasmEngine>(E::Store<StoreData<T, E>>);
435
436impl<T: 'static, E: WasmEngine> WasmStore<T, ReproducibleEngine<E>> for ReproducibleStore<T, E> {
437 fn new(engine: &ReproducibleEngine<E>, data: T) -> Self {
438 Self(<E::Store<StoreData<T, E>> as WasmStore<
439 StoreData<T, E>,
440 E,
441 >>::new(
442 engine.as_ref(),
443 StoreData {
444 data,
445 instruction_counter: None,
446 },
447 ))
448 }
449
450 fn engine(&self) -> &ReproducibleEngine<E> {
451 ReproducibleEngine::from_ref(WasmStore::engine(&self.0))
452 }
453
454 fn data(&self) -> &T {
455 &WasmStore::data(&self.0).data
456 }
457
458 fn data_mut(&mut self) -> &mut T {
459 &mut WasmStore::data_mut(&mut self.0).data
460 }
461
462 fn into_data(self) -> T {
463 WasmStore::into_data(self.0).data
464 }
465}
466
467impl<T: 'static, E: WasmEngine> AsContext<ReproducibleEngine<E>> for ReproducibleStore<T, E> {
468 type UserState = T;
469
470 fn as_context(&self) -> ReproducibleStoreContext<'_, Self::UserState, E> {
471 ReproducibleStoreContext(AsContext::as_context(&self.0))
472 }
473}
474
475impl<T: 'static, E: WasmEngine> AsContextMut<ReproducibleEngine<E>> for ReproducibleStore<T, E> {
476 fn as_context_mut(&mut self) -> ReproducibleStoreContextMut<'_, Self::UserState, E> {
477 ReproducibleStoreContextMut(AsContextMut::as_context_mut(&mut self.0))
478 }
479}
480
481#[repr(transparent)]
482pub struct ReproducibleStoreContext<'a, T: 'static, E: WasmEngine>(
483 E::StoreContext<'a, StoreData<T, E>>,
484);
485
486impl<'a, T: 'static, E: WasmEngine> WasmStoreContext<'a, T, ReproducibleEngine<E>>
487 for ReproducibleStoreContext<'a, T, E>
488{
489 fn engine(&self) -> &ReproducibleEngine<E> {
490 ReproducibleEngine::from_ref(WasmStoreContext::engine(&self.0))
491 }
492
493 fn data(&self) -> &T {
494 &WasmStoreContext::data(&self.0).data
495 }
496}
497
498impl<T: 'static, E: WasmEngine> AsContext<ReproducibleEngine<E>>
499 for ReproducibleStoreContext<'_, T, E>
500{
501 type UserState = T;
502
503 fn as_context(&self) -> ReproducibleStoreContext<'_, Self::UserState, E> {
504 ReproducibleStoreContext(AsContext::as_context(&self.0))
505 }
506}
507
508impl<T: 'static, E: WasmEngine> ReproducibleStoreContext<'_, T, E> {
509 fn as_inner_context(&self) -> E::StoreContext<'_, StoreData<T, E>> {
510 self.0.as_context()
511 }
512}
513
514#[repr(transparent)]
515pub struct ReproducibleStoreContextMut<'a, T: 'static, E: WasmEngine>(
516 E::StoreContextMut<'a, StoreData<T, E>>,
517);
518
519impl<'a, T: 'static, E: WasmEngine> WasmStoreContext<'a, T, ReproducibleEngine<E>>
520 for ReproducibleStoreContextMut<'a, T, E>
521{
522 fn engine(&self) -> &ReproducibleEngine<E> {
523 ReproducibleEngine::from_ref(WasmStoreContext::engine(&self.0))
524 }
525
526 fn data(&self) -> &T {
527 &WasmStoreContext::data(&self.0).data
528 }
529}
530
531impl<'a, T: 'static, E: WasmEngine> WasmStoreContextMut<'a, T, ReproducibleEngine<E>>
532 for ReproducibleStoreContextMut<'a, T, E>
533{
534 fn data_mut(&mut self) -> &mut T {
535 &mut WasmStoreContextMut::data_mut(&mut self.0).data
536 }
537}
538
539impl<T: 'static, E: WasmEngine> AsContext<ReproducibleEngine<E>>
540 for ReproducibleStoreContextMut<'_, T, E>
541{
542 type UserState = T;
543
544 fn as_context(&self) -> ReproducibleStoreContext<'_, Self::UserState, E> {
545 ReproducibleStoreContext(AsContext::as_context(&self.0))
546 }
547}
548
549impl<T: 'static, E: WasmEngine> AsContextMut<ReproducibleEngine<E>>
550 for ReproducibleStoreContextMut<'_, T, E>
551{
552 fn as_context_mut(&mut self) -> ReproducibleStoreContextMut<'_, Self::UserState, E> {
553 ReproducibleStoreContextMut(AsContextMut::as_context_mut(&mut self.0))
554 }
555}
556
557impl<T: 'static, E: WasmEngine> ReproducibleStoreContextMut<'_, T, E> {
558 fn as_inner_context_mut(&mut self) -> E::StoreContextMut<'_, StoreData<T, E>> {
559 self.0.as_context_mut()
560 }
561
562 fn get_instruction_counter_global(&mut self) -> &ReproducibleGlobal<E> {
563 let mut this = self;
564
565 polonius_the_crab::polonius!(|this| -> &'polonius ReproducibleGlobal<E> {
567 let data: &mut StoreData<T, E> = WasmStoreContextMut::data_mut(&mut this.0);
568 if let Some(global) = &data.instruction_counter {
569 polonius_the_crab::polonius_return!(global);
570 }
571 });
572
573 let global = WasmGlobal::new(AsContextMut::as_context_mut(this), Value::I64(0), true);
574
575 let data: &mut StoreData<T, E> = WasmStoreContextMut::data_mut(&mut this.0);
576 data.instruction_counter.insert(global)
577 }
578}
579
580#[derive(Clone)]
581#[repr(transparent)]
582pub struct ReproducibleTable<E: WasmEngine>(E::Table);
583
584impl<E: WasmEngine> WasmTable<ReproducibleEngine<E>> for ReproducibleTable<E> {
585 fn new(
586 mut ctx: impl AsContextMut<ReproducibleEngine<E>>,
587 ty: TableType,
588 init: Value<ReproducibleEngine<E>>,
589 ) -> anyhow::Result<Self> {
590 Ok(Self(<E::Table as WasmTable<E>>::new(
591 ctx.as_context_mut().as_inner_context_mut(),
592 ty,
593 into_value(init),
594 )?))
595 }
596
597 fn ty(&self, ctx: impl AsContext<ReproducibleEngine<E>>) -> TableType {
598 WasmTable::ty(&self.0, ctx.as_context().as_inner_context())
599 }
600
601 fn size(&self, ctx: impl AsContext<ReproducibleEngine<E>>) -> u32 {
602 WasmTable::size(&self.0, ctx.as_context().as_inner_context())
603 }
604
605 fn grow(
606 &self,
607 mut ctx: impl AsContextMut<ReproducibleEngine<E>>,
608 delta: u32,
609 init: Value<ReproducibleEngine<E>>,
610 ) -> anyhow::Result<u32> {
611 WasmTable::grow(
612 &self.0,
613 ctx.as_context_mut().as_inner_context_mut(),
614 delta,
615 into_value(init),
616 )
617 }
618
619 fn get(
620 &self,
621 mut ctx: impl AsContextMut<ReproducibleEngine<E>>,
622 index: u32,
623 ) -> Option<Value<ReproducibleEngine<E>>> {
624 WasmTable::get(&self.0, ctx.as_context_mut().as_inner_context_mut(), index).map(from_value)
625 }
626
627 fn set(
628 &self,
629 mut ctx: impl AsContextMut<ReproducibleEngine<E>>,
630 index: u32,
631 value: Value<ReproducibleEngine<E>>,
632 ) -> anyhow::Result<()> {
633 WasmTable::set(
634 &self.0,
635 ctx.as_context_mut().as_inner_context_mut(),
636 index,
637 into_value(value),
638 )
639 }
640}
641
642const fn as_values<E: WasmEngine>(values: &[Value<ReproducibleEngine<E>>]) -> &[Value<E>] {
643 #[expect(unsafe_code)]
645 unsafe {
646 std::slice::from_raw_parts(values.as_ptr().cast(), values.len())
647 }
648}
649
650const fn as_values_mut<E: WasmEngine>(
651 values: &mut [Value<ReproducibleEngine<E>>],
652) -> &mut [Value<E>] {
653 #[expect(unsafe_code)]
655 unsafe {
656 std::slice::from_raw_parts_mut(values.as_mut_ptr().cast(), values.len())
657 }
658}
659
660const fn from_values<E: WasmEngine>(values: &[Value<E>]) -> &[Value<ReproducibleEngine<E>>] {
661 #[expect(unsafe_code)]
663 unsafe {
664 std::slice::from_raw_parts(values.as_ptr().cast(), values.len())
665 }
666}
667
668const fn from_values_mut<E: WasmEngine>(
669 values: &mut [Value<E>],
670) -> &mut [Value<ReproducibleEngine<E>>] {
671 #[expect(unsafe_code)]
673 unsafe {
674 std::slice::from_raw_parts_mut(values.as_mut_ptr().cast(), values.len())
675 }
676}
677
678fn into_value<E: WasmEngine>(value: Value<ReproducibleEngine<E>>) -> Value<E> {
679 match value {
680 Value::I32(v) => Value::I32(v),
681 Value::I64(v) => Value::I64(v),
682 Value::F32(v) => Value::F32(v),
683 Value::F64(v) => Value::F64(v),
684 Value::FuncRef(v) => Value::FuncRef(v.map(|v| v.0)),
685 Value::ExternRef(v) => Value::ExternRef(v.map(|v| v.0)),
686 }
687}
688
689fn from_value<E: WasmEngine>(value: Value<E>) -> Value<ReproducibleEngine<E>> {
690 match value {
691 Value::I32(v) => Value::I32(v),
692 Value::I64(v) => Value::I64(v),
693 Value::F32(v) => Value::F32(v),
694 Value::F64(v) => Value::F64(v),
695 Value::FuncRef(v) => Value::FuncRef(v.map(ReproducibleFunc)),
696 Value::ExternRef(v) => Value::ExternRef(v.map(ReproducibleExternRef)),
697 }
698}
699
700fn into_extern<E: WasmEngine>(value: Extern<ReproducibleEngine<E>>) -> Extern<E> {
701 match value {
702 Extern::Global(v) => Extern::Global(v.0),
703 Extern::Table(v) => Extern::Table(v.0),
704 Extern::Memory(v) => Extern::Memory(v.0),
705 Extern::Func(v) => Extern::Func(v.0),
706 }
707}
708
709fn from_extern<E: WasmEngine>(value: Extern<E>) -> Extern<ReproducibleEngine<E>> {
710 match value {
711 Extern::Global(v) => Extern::Global(ReproducibleGlobal(v)),
712 Extern::Table(v) => Extern::Table(ReproducibleTable(v)),
713 Extern::Memory(v) => Extern::Memory(ReproducibleMemory(v)),
714 Extern::Func(v) => Extern::Func(ReproducibleFunc(v)),
715 }
716}