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
485602
    pub fn new() -> Self {
10
485602
        Self {
11
485602
            bytes: Vec::new(),
12
485602
            block_depth: 0,
13
485602
        }
14
485602
    }
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
33366
    pub fn new_seeded(depth: u32) -> Self {
25
33366
        Self {
26
33366
            bytes: Vec::new(),
27
33366
            block_depth: depth,
28
33366
        }
29
33366
    }
30

            
31
    /// Consume the emitter, yielding its raw instruction bytes for
32
    /// splicing into a parent via [`Self::splice`].
33
28801
    pub fn take_bytes(self) -> Vec<u8> {
34
28801
        self.bytes
35
28801
    }
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
3901
    pub fn bytes_ref(&self) -> &[u8] {
41
3901
        &self.bytes
42
3901
    }
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
32702
    pub fn splice(&mut self, bytes: &[u8]) {
49
32702
        self.bytes.extend_from_slice(bytes);
50
32702
    }
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
28967
    pub fn block_depth(&self) -> u32 {
57
28967
        self.block_depth
58
28967
    }
59

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            
149
337926
    pub fn i32_load(&mut self, offset: u64) {
150
337926
        Instruction::I32Load(MemArg {
151
337926
            offset,
152
337926
            align: 2,
153
337926
            memory_index: 0,
154
337926
        })
155
337926
        .encode(&mut self.bytes);
156
337926
    }
157

            
158
9296
    pub fn i32_load8_u(&mut self, offset: u64) {
159
9296
        Instruction::I32Load8U(MemArg {
160
9296
            offset,
161
9296
            align: 0,
162
9296
            memory_index: 0,
163
9296
        })
164
9296
        .encode(&mut self.bytes);
165
9296
    }
166

            
167
12782
    pub fn i64_load(&mut self, offset: u64) {
168
12782
        Instruction::I64Load(MemArg {
169
12782
            offset,
170
12782
            align: 3,
171
12782
            memory_index: 0,
172
12782
        })
173
12782
        .encode(&mut self.bytes);
174
12782
    }
175

            
176
47313
    pub fn i32_ge_u(&mut self) {
177
47313
        Instruction::I32GeU.encode(&mut self.bytes);
178
47313
    }
179

            
180
983332
    pub fn struct_get(&mut self, type_idx: u32, field_idx: u32) {
181
983332
        Instruction::StructGet {
182
983332
            struct_type_index: type_idx,
183
983332
            field_index: field_idx,
184
983332
        }
185
983332
        .encode(&mut self.bytes);
186
983332
    }
187

            
188
36853
    pub fn if_block(&mut self, block_type: BlockType) {
189
36853
        Instruction::If(block_type).encode(&mut self.bytes);
190
36853
        self.block_depth += 1;
191
36853
    }
192

            
193
28304
    pub fn else_block(&mut self) {
194
28304
        Instruction::Else.encode(&mut self.bytes);
195
28304
    }
196

            
197
298667
    pub fn i32_store_raw(&mut self) {
198
298667
        Instruction::I32Store(MemArg {
199
298667
            offset: 0,
200
298667
            align: 2,
201
298667
            memory_index: 0,
202
298667
        })
203
298667
        .encode(&mut self.bytes);
204
298667
    }
205

            
206
44488
    pub fn i64_store_raw(&mut self) {
207
44488
        Instruction::I64Store(MemArg {
208
44488
            offset: 0,
209
44488
            align: 3,
210
44488
            memory_index: 0,
211
44488
        })
212
44488
        .encode(&mut self.bytes);
213
44488
    }
214

            
215
15023
    pub fn i64_extend_i32_u(&mut self) {
216
15023
        Instruction::I64ExtendI32U.encode(&mut self.bytes);
217
15023
    }
218

            
219
120685
    pub fn i64_extend_i32_s(&mut self) {
220
120685
        Instruction::I64ExtendI32S.encode(&mut self.bytes);
221
120685
    }
222

            
223
166
    pub fn i32_wrap_i64(&mut self) {
224
166
        Instruction::I32WrapI64.encode(&mut self.bytes);
225
166
    }
226

            
227
87577
    pub fn array_new_data(&mut self, type_idx: u32, data_idx: u32) {
228
87577
        Instruction::ArrayNewData {
229
87577
            array_type_index: type_idx,
230
87577
            array_data_index: data_idx,
231
87577
        }
232
87577
        .encode(&mut self.bytes);
233
87577
    }
234

            
235
42831
    pub fn array_get_u(&mut self, type_idx: u32) {
236
42831
        Instruction::ArrayGetU(type_idx).encode(&mut self.bytes);
237
42831
    }
238

            
239
4482
    pub fn array_new_default(&mut self, type_idx: u32) {
240
4482
        Instruction::ArrayNewDefault(type_idx).encode(&mut self.bytes);
241
4482
    }
242

            
243
4482
    pub fn array_set(&mut self, type_idx: u32) {
244
4482
        Instruction::ArraySet(type_idx).encode(&mut self.bytes);
245
4482
    }
246

            
247
16102
    pub fn array_len(&mut self) {
248
16102
        Instruction::ArrayLen.encode(&mut self.bytes);
249
16102
    }
250

            
251
167593
    pub fn block_start(&mut self) {
252
167593
        Instruction::Block(BlockType::Empty).encode(&mut self.bytes);
253
167593
        self.block_depth += 1;
254
167593
    }
255

            
256
    /// Open a wasm `block` whose result type is known up-front. Used
257
    /// by BLOCK / RETURN-FROM where the lexically-named block carries
258
    /// a value-yielding result.
259
878246
    pub fn block_start_typed(&mut self, block_type: BlockType) {
260
878246
        Instruction::Block(block_type).encode(&mut self.bytes);
261
878246
        self.block_depth += 1;
262
878246
    }
263

            
264
64411
    pub fn loop_start(&mut self) {
265
64411
        Instruction::Loop(BlockType::Empty).encode(&mut self.bytes);
266
64411
        self.block_depth += 1;
267
64411
    }
268

            
269
1585518
    pub fn block_end(&mut self) {
270
1585518
        Instruction::End.encode(&mut self.bytes);
271
1585518
        self.block_depth = self.block_depth.saturating_sub(1);
272
1585518
    }
273

            
274
541759
    pub fn br(&mut self, depth: u32) {
275
541759
        Instruction::Br(depth).encode(&mut self.bytes);
276
541759
    }
277

            
278
63083
    pub fn br_if(&mut self, depth: u32) {
279
63083
        Instruction::BrIf(depth).encode(&mut self.bytes);
280
63083
    }
281

            
282
    /// Indexed branch — pops an i32 selector, jumps to `targets[i]` when the
283
    /// selector is in range, else `default`. Used by TAGBODY / GO to dispatch
284
    /// to the segment matching the current pc.
285
1328
    pub fn br_table(&mut self, targets: &[u32], default: u32) {
286
1328
        Instruction::BrTable(std::borrow::Cow::Borrowed(targets), default).encode(&mut self.bytes);
287
1328
    }
288

            
289
1218972
    pub fn call(&mut self, func_idx: u32) {
290
1218972
        Instruction::Call(func_idx).encode(&mut self.bytes);
291
1218972
    }
292

            
293
    /// Indirect call through a typed funcref popped from the stack.
294
    /// `type_idx` is the wasm fn-type index; the engine validates the
295
    /// funcref's declared shape against it at compile time. Used by the
296
    /// Tier 1.5 closure call site after `struct.get $closure_<sig>
297
    /// $fn_ref` has produced the funcref.
298
4316
    pub fn call_ref(&mut self, type_idx: u32) {
299
4316
        Instruction::CallRef(type_idx).encode(&mut self.bytes);
300
4316
    }
301

            
302
328392
    pub fn store_u8_dynamic(&mut self, local_base: u32, offset: u32, value: u8) {
303
328392
        Instruction::LocalGet(local_base).encode(&mut self.bytes);
304
328392
        Instruction::I32Const(offset as i32).encode(&mut self.bytes);
305
328392
        Instruction::I32Add.encode(&mut self.bytes);
306
328392
        Instruction::I32Const(i32::from(value)).encode(&mut self.bytes);
307
328392
        Instruction::I32Store8(MemArg {
308
328392
            offset: 0,
309
328392
            align: 0,
310
328392
            memory_index: 0,
311
328392
        })
312
328392
        .encode(&mut self.bytes);
313
328392
    }
314

            
315
606724
    pub fn store_i32_dynamic(&mut self, local_base: u32, offset: u32, value: i32) {
316
606724
        Instruction::LocalGet(local_base).encode(&mut self.bytes);
317
606724
        Instruction::I32Const(offset as i32).encode(&mut self.bytes);
318
606724
        Instruction::I32Add.encode(&mut self.bytes);
319
606724
        Instruction::I32Const(value).encode(&mut self.bytes);
320
606724
        Instruction::I32Store(MemArg {
321
606724
            offset: 0,
322
606724
            align: 2,
323
606724
            memory_index: 0,
324
606724
        })
325
606724
        .encode(&mut self.bytes);
326
606724
    }
327

            
328
100938
    pub fn store_i64_dynamic(&mut self, local_base: u32, offset: u32, value: i64) {
329
100938
        Instruction::LocalGet(local_base).encode(&mut self.bytes);
330
100938
        Instruction::I32Const(offset as i32).encode(&mut self.bytes);
331
100938
        Instruction::I32Add.encode(&mut self.bytes);
332
100938
        Instruction::I64Const(value).encode(&mut self.bytes);
333
100938
        Instruction::I64Store(MemArg {
334
100938
            offset: 0,
335
100938
            align: 3,
336
100938
            memory_index: 0,
337
100938
        })
338
100938
        .encode(&mut self.bytes);
339
100938
    }
340

            
341
63166
    pub fn i32_store8_raw(&mut self) {
342
63166
        Instruction::I32Store8(MemArg {
343
63166
            offset: 0,
344
63166
            align: 0,
345
63166
            memory_index: 0,
346
63166
        })
347
63166
        .encode(&mut self.bytes);
348
63166
    }
349

            
350
3818
    pub fn i32_load16_u(&mut self, offset: u64) {
351
3818
        Instruction::I32Load16U(MemArg {
352
3818
            offset,
353
3818
            align: 1,
354
3818
            memory_index: 0,
355
3818
        })
356
3818
        .encode(&mut self.bytes);
357
3818
    }
358

            
359
29636
    pub fn ref_null(&mut self, type_idx: u32) {
360
29636
        Instruction::RefNull(HeapType::Concrete(type_idx)).encode(&mut self.bytes);
361
29636
    }
362

            
363
    /// Emits `ref.null any` — the polymorphic null that satisfies every
364
    /// `(ref null any)` slot. Used as nomi-eval's "no result" return value
365
    /// for empty programs / definition-only forms.
366
8219
    pub fn ref_null_any(&mut self) {
367
8219
        Instruction::RefNull(HeapType::Abstract {
368
8219
            shared: false,
369
8219
            ty: wasm_encoder::AbstractHeapType::Any,
370
8219
        })
371
8219
        .encode(&mut self.bytes);
372
8219
    }
373

            
374
26395
    pub fn ref_is_null(&mut self) {
375
26395
        Instruction::RefIsNull.encode(&mut self.bytes);
376
26395
    }
377

            
378
    /// Boxes the top-of-stack i32 into an `(ref i31)` so it can ride an
379
    /// `anyref` slot — the canonical way to embed small ints in a
380
    /// `$pair` car. Pairs with [`Self::i31_get_s`] for the reverse.
381
2241
    pub fn ref_i31(&mut self) {
382
2241
        Instruction::RefI31.encode(&mut self.bytes);
383
2241
    }
384

            
385
    /// Unpacks a `(ref i31)` back to its signed i32 value.
386
6308
    pub fn i31_get_s(&mut self) {
387
6308
        Instruction::I31GetS.encode(&mut self.bytes);
388
6308
    }
389

            
390
    /// Downcasts a `(ref null any)` (or any subtype) to a concrete struct
391
    /// type's non-null reference. Traps at runtime if the cast fails.
392
    /// Caller is responsible for ensuring the value's dynamic type
393
    /// matches; the compile-time `PairElement` carries that proof.
394
943077
    pub fn ref_cast(&mut self, type_idx: u32) {
395
943077
        Instruction::RefCastNonNull(HeapType::Concrete(type_idx)).encode(&mut self.bytes);
396
943077
    }
397

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

            
408
    /// Downcast to the abstract `i31` heap type — paired with
409
    /// [`Self::i31_get_s`] to retrieve the boxed i32.
410
6308
    pub fn ref_cast_i31(&mut self) {
411
6308
        Instruction::RefCastNonNull(HeapType::I31).encode(&mut self.bytes);
412
6308
    }
413

            
414
473032
    pub fn end(&mut self) {
415
473032
        Instruction::End.encode(&mut self.bytes);
416
473032
    }
417

            
418
7304
    pub fn drop_value(&mut self) {
419
7304
        Instruction::Drop.encode(&mut self.bytes);
420
7304
    }
421

            
422
    /// Emit `unreachable`. Used after a `br` out of a value-yielding
423
    /// block where the validator still expects a typed value on the
424
    /// stack but control will never reach the next instruction —
425
    /// wasm's stack-polymorphism after `unreachable` keeps the
426
    /// type-check happy without forcing us to materialize a real value.
427
481000
    pub fn unreachable(&mut self) {
428
481000
        Instruction::Unreachable.encode(&mut self.bytes);
429
481000
    }
430

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

            
444
    /// Throw the exception with the given tag index, consuming one stack
445
    /// value per tag parameter (the `$nomi_error` tag takes a single
446
    /// `(ref null $nomi_condition)`). Control never returns past a
447
    /// `throw`, so the stack goes polymorphic — value-position callers
448
    /// pair this with a trailing type the validator accepts.
449
9047
    pub fn throw(&mut self, tag_idx: u32) {
450
9047
        Instruction::Throw(tag_idx).encode(&mut self.bytes);
451
9047
    }
452

            
453
    /// Pushes a typed funcref pointing at the wasm function with the
454
    /// given index. Tier 1.5 closure construction pairs this with
455
    /// [`Self::struct_new`] to populate a `$closure_<sig>` struct.
456
5063
    pub fn ref_func(&mut self, func_idx: u32) {
457
5063
        Instruction::RefFunc(func_idx).encode(&mut self.bytes);
458
5063
    }
459

            
460
    /// Constructs a struct of the registered struct type. Consumes one
461
    /// stack value per field in declaration order; pushes a
462
    /// `(ref $type_idx)` of the resulting struct.
463
8798
    pub fn struct_new(&mut self, type_idx: u32) {
464
8798
        Instruction::StructNew(type_idx).encode(&mut self.bytes);
465
8798
    }
466
}
467

            
468
impl Default for FunctionEmitter {
469
    fn default() -> Self {
470
        Self::new()
471
    }
472
}