1
use wasm_encoder::{BlockType, Catch, Encode, Function, HeapType, Instruction, MemArg, ValType};
2

            
3
pub struct FunctionEmitter {
4
    bytes: Vec<u8>,
5
    block_depth: u32,
6
}
7

            
8
impl FunctionEmitter {
9
259192
    pub fn new() -> Self {
10
259192
        Self {
11
259192
            bytes: Vec::new(),
12
259192
            block_depth: 0,
13
259192
        }
14
259192
    }
15

            
16
    /// A fresh scratch emitter whose `block_depth` is pre-seeded to
17
    /// `depth`, as if it were already nested that many structured-control
18
    /// frames deep. Used by BLOCK emit-time exit discovery: the body is
19
    /// compiled into a scratch buffer seeded at the depth the block's
20
    /// frame will occupy, so every relative `br` (RETURN-FROM / GO,
21
    /// including exits to outer blocks) computes the same target it will
22
    /// need once the bytes are spliced into the parent after the parent
23
    /// opens the block. Pair with [`Self::take_bytes`] + [`Self::splice`].
24
22152
    pub fn new_seeded(depth: u32) -> Self {
25
22152
        Self {
26
22152
            bytes: Vec::new(),
27
22152
            block_depth: depth,
28
22152
        }
29
22152
    }
30

            
31
    /// Consume the emitter, yielding its raw instruction bytes for
32
    /// splicing into a parent via [`Self::splice`].
33
19596
    pub fn take_bytes(self) -> Vec<u8> {
34
19596
        self.bytes
35
19596
    }
36

            
37
    /// Borrow the raw instruction bytes for splicing without consuming the
38
    /// emitter (e.g. handler-case keeps several clause scratches and splices
39
    /// each in turn).
40
2272
    pub fn bytes_ref(&self) -> &[u8] {
41
2272
        &self.bytes
42
2272
    }
43

            
44
    /// Append a scratch emitter's raw bytes verbatim. The scratch body is
45
    /// internally depth-balanced (it ends at the depth it started), so the
46
    /// parent's `block_depth` is unchanged — the caller manages the
47
    /// enclosing frame with `block_start_typed` / `block_end` around this.
48
21868
    pub fn splice(&mut self, bytes: &[u8]) {
49
21868
        self.bytes.extend_from_slice(bytes);
50
21868
    }
51

            
52
    /// Current wasm structured-control depth — the number of open
53
    /// `block` / `loop` / `if` frames around the next instruction. Used
54
    /// by labelled exits (BLOCK / RETURN-FROM, TAGBODY / GO) to compute
55
    /// the relative `br` depth.
56
21584
    pub fn block_depth(&self) -> u32 {
57
21584
        self.block_depth
58
21584
    }
59

            
60
246598
    pub fn finish(self, locals: &[(u32, ValType)]) -> Function {
61
246598
        let mut func = Function::new(locals.iter().copied());
62
246598
        func.raw(self.bytes);
63
246598
        func
64
246598
    }
65

            
66
5893
    pub fn memory_init(&mut self, data_idx: u32, dst: u32, len: u32) {
67
5893
        Instruction::I32Const(dst as i32).encode(&mut self.bytes);
68
5893
        Instruction::I32Const(0).encode(&mut self.bytes);
69
5893
        Instruction::I32Const(len as i32).encode(&mut self.bytes);
70
5893
        Instruction::MemoryInit {
71
5893
            mem: 0,
72
5893
            data_index: data_idx,
73
5893
        }
74
5893
        .encode(&mut self.bytes);
75
5893
    }
76

            
77
1794752
    pub fn local_get(&mut self, idx: u32) {
78
1794752
        Instruction::LocalGet(idx).encode(&mut self.bytes);
79
1794752
    }
80

            
81
747039
    pub fn local_set(&mut self, idx: u32) {
82
747039
        Instruction::LocalSet(idx).encode(&mut self.bytes);
83
747039
    }
84

            
85
2343
    pub fn local_tee(&mut self, idx: u32) {
86
2343
        Instruction::LocalTee(idx).encode(&mut self.bytes);
87
2343
    }
88

            
89
973010
    pub fn i32_const(&mut self, value: i32) {
90
973010
        Instruction::I32Const(value).encode(&mut self.bytes);
91
973010
    }
92

            
93
38202
    pub fn i64_const(&mut self, value: i64) {
94
38202
        Instruction::I64Const(value).encode(&mut self.bytes);
95
38202
    }
96

            
97
764543
    pub fn i32_add(&mut self) {
98
764543
        Instruction::I32Add.encode(&mut self.bytes);
99
764543
    }
100

            
101
83081
    pub fn i32_sub(&mut self) {
102
83081
        Instruction::I32Sub.encode(&mut self.bytes);
103
83081
    }
104

            
105
6177
    pub fn i32_eqz(&mut self) {
106
6177
        Instruction::I32Eqz.encode(&mut self.bytes);
107
6177
    }
108

            
109
9585
    pub fn i32_eq(&mut self) {
110
9585
        Instruction::I32Eq.encode(&mut self.bytes);
111
9585
    }
112

            
113
1633
    pub fn i32_lt_s(&mut self) {
114
1633
        Instruction::I32LtS.encode(&mut self.bytes);
115
1633
    }
116

            
117
3479
    pub fn i32_ge_s(&mut self) {
118
3479
        Instruction::I32GeS.encode(&mut self.bytes);
119
3479
    }
120

            
121
781
    pub fn i32_gt_s(&mut self) {
122
781
        Instruction::I32GtS.encode(&mut self.bytes);
123
781
    }
124

            
125
    pub fn i32_le_s(&mut self) {
126
        Instruction::I32LeS.encode(&mut self.bytes);
127
    }
128

            
129
9088
    pub fn i32_ne(&mut self) {
130
9088
        Instruction::I32Ne.encode(&mut self.bytes);
131
9088
    }
132

            
133
2201
    pub fn i32_shl(&mut self) {
134
2201
        Instruction::I32Shl.encode(&mut self.bytes);
135
2201
    }
136

            
137
2201
    pub fn i32_or(&mut self) {
138
2201
        Instruction::I32Or.encode(&mut self.bytes);
139
2201
    }
140

            
141
    pub fn i64_ne(&mut self) {
142
        Instruction::I64Ne.encode(&mut self.bytes);
143
    }
144

            
145
23146
    pub fn i32_mul(&mut self) {
146
23146
        Instruction::I32Mul.encode(&mut self.bytes);
147
23146
    }
148

            
149
    pub fn i32_div_s(&mut self) {
150
        Instruction::I32DivS.encode(&mut self.bytes);
151
    }
152

            
153
    pub fn i32_rem_s(&mut self) {
154
        Instruction::I32RemS.encode(&mut self.bytes);
155
    }
156

            
157
288790
    pub fn i32_load(&mut self, offset: u64) {
158
288790
        Instruction::I32Load(MemArg {
159
288790
            offset,
160
288790
            align: 2,
161
288790
            memory_index: 0,
162
288790
        })
163
288790
        .encode(&mut self.bytes);
164
288790
    }
165

            
166
7881
    pub fn i32_load8_u(&mut self, offset: u64) {
167
7881
        Instruction::I32Load8U(MemArg {
168
7881
            offset,
169
7881
            align: 0,
170
7881
            memory_index: 0,
171
7881
        })
172
7881
        .encode(&mut self.bytes);
173
7881
    }
174

            
175
10934
    pub fn i64_load(&mut self, offset: u64) {
176
10934
        Instruction::I64Load(MemArg {
177
10934
            offset,
178
10934
            align: 3,
179
10934
            memory_index: 0,
180
10934
        })
181
10934
        .encode(&mut self.bytes);
182
10934
    }
183

            
184
40473
    pub fn i32_ge_u(&mut self) {
185
40473
        Instruction::I32GeU.encode(&mut self.bytes);
186
40473
    }
187

            
188
521242
    pub fn struct_get(&mut self, type_idx: u32, field_idx: u32) {
189
521242
        Instruction::StructGet {
190
521242
            struct_type_index: type_idx,
191
521242
            field_index: field_idx,
192
521242
        }
193
521242
        .encode(&mut self.bytes);
194
521242
    }
195

            
196
23643
    pub fn if_block(&mut self, block_type: BlockType) {
197
23643
        Instruction::If(block_type).encode(&mut self.bytes);
198
23643
        self.block_depth += 1;
199
23643
    }
200

            
201
17324
    pub fn else_block(&mut self) {
202
17324
        Instruction::Else.encode(&mut self.bytes);
203
17324
    }
204

            
205
255065
    pub fn i32_store_raw(&mut self) {
206
255065
        Instruction::I32Store(MemArg {
207
255065
            offset: 0,
208
255065
            align: 2,
209
255065
            memory_index: 0,
210
255065
        })
211
255065
        .encode(&mut self.bytes);
212
255065
    }
213

            
214
38127
    pub fn i64_store_raw(&mut self) {
215
38127
        Instruction::I64Store(MemArg {
216
38127
            offset: 0,
217
38127
            align: 3,
218
38127
            memory_index: 0,
219
38127
        })
220
38127
        .encode(&mut self.bytes);
221
38127
    }
222

            
223
12567
    pub fn i64_extend_i32_u(&mut self) {
224
12567
        Instruction::I64ExtendI32U.encode(&mut self.bytes);
225
12567
    }
226

            
227
15478
    pub fn i64_extend_i32_s(&mut self) {
228
15478
        Instruction::I64ExtendI32S.encode(&mut self.bytes);
229
15478
    }
230

            
231
142
    pub fn i32_wrap_i64(&mut self) {
232
142
        Instruction::I32WrapI64.encode(&mut self.bytes);
233
142
    }
234

            
235
71289
    pub fn array_new_data(&mut self, type_idx: u32, data_idx: u32) {
236
71289
        Instruction::ArrayNewData {
237
71289
            array_type_index: type_idx,
238
71289
            array_data_index: data_idx,
239
71289
        }
240
71289
        .encode(&mut self.bytes);
241
71289
    }
242

            
243
36710
    pub fn array_get_u(&mut self, type_idx: u32) {
244
36710
        Instruction::ArrayGetU(type_idx).encode(&mut self.bytes);
245
36710
    }
246

            
247
3763
    pub fn array_new_default(&mut self, type_idx: u32) {
248
3763
        Instruction::ArrayNewDefault(type_idx).encode(&mut self.bytes);
249
3763
    }
250

            
251
3763
    pub fn array_set(&mut self, type_idx: u32) {
252
3763
        Instruction::ArraySet(type_idx).encode(&mut self.bytes);
253
3763
    }
254

            
255
13774
    pub fn array_len(&mut self) {
256
13774
        Instruction::ArrayLen.encode(&mut self.bytes);
257
13774
    }
258

            
259
140454
    pub fn block_start(&mut self) {
260
140454
        Instruction::Block(BlockType::Empty).encode(&mut self.bytes);
261
140454
        self.block_depth += 1;
262
140454
    }
263

            
264
    /// Open a wasm `block` whose result type is known up-front. Used
265
    /// by BLOCK / RETURN-FROM where the lexically-named block carries
266
    /// a value-yielding result.
267
437525
    pub fn block_start_typed(&mut self, block_type: BlockType) {
268
437525
        Instruction::Block(block_type).encode(&mut self.bytes);
269
437525
        self.block_depth += 1;
270
437525
    }
271

            
272
52259
    pub fn loop_start(&mut self) {
273
52259
        Instruction::Loop(BlockType::Empty).encode(&mut self.bytes);
274
52259
        self.block_depth += 1;
275
52259
    }
276

            
277
873777
    pub fn block_end(&mut self) {
278
873777
        Instruction::End.encode(&mut self.bytes);
279
873777
        self.block_depth = self.block_depth.saturating_sub(1);
280
873777
    }
281

            
282
303827
    pub fn br(&mut self, depth: u32) {
283
303827
        Instruction::Br(depth).encode(&mut self.bytes);
284
303827
    }
285

            
286
51123
    pub fn br_if(&mut self, depth: u32) {
287
51123
        Instruction::BrIf(depth).encode(&mut self.bytes);
288
51123
    }
289

            
290
    /// Indexed branch — pops an i32 selector, jumps to `targets[i]` when the
291
    /// selector is in range, else `default`. Used by TAGBODY / GO to dispatch
292
    /// to the segment matching the current pc.
293
1136
    pub fn br_table(&mut self, targets: &[u32], default: u32) {
294
1136
        Instruction::BrTable(std::borrow::Cow::Borrowed(targets), default).encode(&mut self.bytes);
295
1136
    }
296

            
297
605235
    pub fn call(&mut self, func_idx: u32) {
298
605235
        Instruction::Call(func_idx).encode(&mut self.bytes);
299
605235
    }
300

            
301
    /// Indirect call through a typed funcref popped from the stack.
302
    /// `type_idx` is the wasm fn-type index; the engine validates the
303
    /// funcref's declared shape against it at compile time. Used by the
304
    /// Tier 1.5 closure call site after `struct.get $closure_<sig>
305
    /// $fn_ref` has produced the funcref.
306
3266
    pub fn call_ref(&mut self, type_idx: u32) {
307
3266
        Instruction::CallRef(type_idx).encode(&mut self.bytes);
308
3266
    }
309

            
310
283973
    pub fn store_u8_dynamic(&mut self, local_base: u32, offset: u32, value: u8) {
311
283973
        Instruction::LocalGet(local_base).encode(&mut self.bytes);
312
283973
        Instruction::I32Const(offset as i32).encode(&mut self.bytes);
313
283973
        Instruction::I32Add.encode(&mut self.bytes);
314
283973
        Instruction::I32Const(i32::from(value)).encode(&mut self.bytes);
315
283973
        Instruction::I32Store8(MemArg {
316
283973
            offset: 0,
317
283973
            align: 0,
318
283973
            memory_index: 0,
319
283973
        })
320
283973
        .encode(&mut self.bytes);
321
283973
    }
322

            
323
518519
    pub fn store_i32_dynamic(&mut self, local_base: u32, offset: u32, value: i32) {
324
518519
        Instruction::LocalGet(local_base).encode(&mut self.bytes);
325
518519
        Instruction::I32Const(offset as i32).encode(&mut self.bytes);
326
518519
        Instruction::I32Add.encode(&mut self.bytes);
327
518519
        Instruction::I32Const(value).encode(&mut self.bytes);
328
518519
        Instruction::I32Store(MemArg {
329
518519
            offset: 0,
330
518519
            align: 2,
331
518519
            memory_index: 0,
332
518519
        })
333
518519
        .encode(&mut self.bytes);
334
518519
    }
335

            
336
88409
    pub fn store_i64_dynamic(&mut self, local_base: u32, offset: u32, value: i64) {
337
88409
        Instruction::LocalGet(local_base).encode(&mut self.bytes);
338
88409
        Instruction::I32Const(offset as i32).encode(&mut self.bytes);
339
88409
        Instruction::I32Add.encode(&mut self.bytes);
340
88409
        Instruction::I64Const(value).encode(&mut self.bytes);
341
88409
        Instruction::I64Store(MemArg {
342
88409
            offset: 0,
343
88409
            align: 3,
344
88409
            memory_index: 0,
345
88409
        })
346
88409
        .encode(&mut self.bytes);
347
88409
    }
348

            
349
50484
    pub fn i32_store8_raw(&mut self) {
350
50484
        Instruction::I32Store8(MemArg {
351
50484
            offset: 0,
352
50484
            align: 0,
353
50484
            memory_index: 0,
354
50484
        })
355
50484
        .encode(&mut self.bytes);
356
50484
    }
357

            
358
3195
    pub fn i32_load16_u(&mut self, offset: u64) {
359
3195
        Instruction::I32Load16U(MemArg {
360
3195
            offset,
361
3195
            align: 1,
362
3195
            memory_index: 0,
363
3195
        })
364
3195
        .encode(&mut self.bytes);
365
3195
    }
366

            
367
12855
    pub fn ref_null(&mut self, type_idx: u32) {
368
12855
        Instruction::RefNull(HeapType::Concrete(type_idx)).encode(&mut self.bytes);
369
12855
    }
370

            
371
    /// Emits `ref.null any` — the polymorphic null that satisfies every
372
    /// `(ref null any)` slot. Used as nomi-eval's "no result" return value
373
    /// for empty programs / definition-only forms.
374
3978
    pub fn ref_null_any(&mut self) {
375
3978
        Instruction::RefNull(HeapType::Abstract {
376
3978
            shared: false,
377
3978
            ty: wasm_encoder::AbstractHeapType::Any,
378
3978
        })
379
3978
        .encode(&mut self.bytes);
380
3978
    }
381

            
382
13774
    pub fn ref_is_null(&mut self) {
383
13774
        Instruction::RefIsNull.encode(&mut self.bytes);
384
13774
    }
385

            
386
    /// Boxes the top-of-stack i32 into an `(ref i31)` so it can ride an
387
    /// `anyref` slot — the canonical way to embed small ints in a
388
    /// `$pair` car. Pairs with [`Self::i31_get_s`] for the reverse.
389
52756
    pub fn ref_i31(&mut self) {
390
52756
        Instruction::RefI31.encode(&mut self.bytes);
391
52756
    }
392

            
393
    /// Unpacks a `(ref i31)` back to its signed i32 value.
394
3479
    pub fn i31_get_s(&mut self) {
395
3479
        Instruction::I31GetS.encode(&mut self.bytes);
396
3479
    }
397

            
398
    /// Downcasts a `(ref null any)` (or any subtype) to a concrete struct
399
    /// type's non-null reference. Traps at runtime if the cast fails.
400
    /// Caller is responsible for ensuring the value's dynamic type
401
    /// matches; the compile-time `PairElement` carries that proof.
402
494546
    pub fn ref_cast(&mut self, type_idx: u32) {
403
494546
        Instruction::RefCastNonNull(HeapType::Concrete(type_idx)).encode(&mut self.bytes);
404
494546
    }
405

            
406
    /// Nullable counterpart to [`ref_cast`]. Use when the value's
407
    /// declared type is `(ref null …)` and a null pointer is a valid
408
    /// runtime outcome — e.g. host fns returning
409
    /// `Option<Rooted<…>>`, which the wasmtime ABI surfaces as a
410
    /// nullable reference. The non-null variant traps on null; this
411
    /// one preserves null and only fails on type mismatch.
412
121056
    pub fn ref_cast_nullable(&mut self, type_idx: u32) {
413
121056
        Instruction::RefCastNullable(HeapType::Concrete(type_idx)).encode(&mut self.bytes);
414
121056
    }
415

            
416
    /// Downcast to the abstract `i31` heap type — paired with
417
    /// [`Self::i31_get_s`] to retrieve the boxed i32.
418
3479
    pub fn ref_cast_i31(&mut self) {
419
3479
        Instruction::RefCastNonNull(HeapType::I31).encode(&mut self.bytes);
420
3479
    }
421

            
422
248799
    pub fn end(&mut self) {
423
248799
        Instruction::End.encode(&mut self.bytes);
424
248799
    }
425

            
426
4828
    pub fn drop_value(&mut self) {
427
4828
        Instruction::Drop.encode(&mut self.bytes);
428
4828
    }
429

            
430
    /// Emit `unreachable`. Used after a `br` out of a value-yielding
431
    /// block where the validator still expects a typed value on the
432
    /// stack but control will never reach the next instruction —
433
    /// wasm's stack-polymorphism after `unreachable` keeps the
434
    /// type-check happy without forcing us to materialize a real value.
435
254266
    pub fn unreachable(&mut self) {
436
254266
        Instruction::Unreachable.encode(&mut self.bytes);
437
254266
    }
438

            
439
    /// Open a `try_table` frame (exception-handling proposal). `catches`
440
    /// routes matching `throw`s to a relative `br` depth resolved against
441
    /// the open blocks. Opens one structured-control frame like `block`,
442
    /// so `block_depth` increments — keeping `(return-from name v)` and
443
    /// `(go tag)` relative-`br` math correct across a wrapped body. Pair
444
    /// with [`Self::block_end`]. Used by the Tier 3 boundary wrapper and
445
    /// by `(handler-case)` / `(unwind-protect)` (ADR-0026).
446
259736
    pub fn try_table(&mut self, block_type: BlockType, catches: &[Catch]) {
447
259736
        Instruction::TryTable(block_type, std::borrow::Cow::Borrowed(catches))
448
259736
            .encode(&mut self.bytes);
449
259736
        self.block_depth += 1;
450
259736
    }
451

            
452
    /// Throw the exception with the given tag index, consuming one stack
453
    /// value per tag parameter (the `$nomi_error` tag takes a single
454
    /// `(ref null $nomi_condition)`). Control never returns past a
455
    /// `throw`, so the stack goes polymorphic — value-position callers
456
    /// pair this with a trailing type the validator accepts.
457
7100
    pub fn throw(&mut self, tag_idx: u32) {
458
7100
        Instruction::Throw(tag_idx).encode(&mut self.bytes);
459
7100
    }
460

            
461
    /// Pushes a typed funcref pointing at the wasm function with the
462
    /// given index. Tier 1.5 closure construction pairs this with
463
    /// [`Self::struct_new`] to populate a `$closure_<sig>` struct.
464
3621
    pub fn ref_func(&mut self, func_idx: u32) {
465
3621
        Instruction::RefFunc(func_idx).encode(&mut self.bytes);
466
3621
    }
467

            
468
    /// Constructs a struct of the registered struct type. Consumes one
469
    /// stack value per field in declaration order; pushes a
470
    /// `(ref $type_idx)` of the resulting struct.
471
6674
    pub fn struct_new(&mut self, type_idx: u32) {
472
6674
        Instruction::StructNew(type_idx).encode(&mut self.bytes);
473
6674
    }
474
}
475

            
476
impl Default for FunctionEmitter {
477
    fn default() -> Self {
478
        Self::new()
479
    }
480
}