Skip to content

Commit cee820e

Browse files
chore: add specialized returns, optimize parser
Signed-off-by: Henry <mail@henrygressmann.de>
1 parent 165b198 commit cee820e

8 files changed

Lines changed: 340 additions & 193 deletions

File tree

crates/parser/src/macros.rs

Lines changed: 97 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,103 @@
11
pub(crate) mod visit {
22
macro_rules! validate_then_visit {
3-
($( @$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*))*) => {$(
3+
($( @$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*))*) => {
4+
$(validate_then_visit!(@@$proposal $op $({ $($arg: $argty),* })? => $visit ($($ann)*));)*
5+
};
6+
7+
// These special-case arms exist so we only clone wasmparser's non-Copy payloads
8+
(@@mvp BrTable { $arg:ident: $argty:ty } => $visit:ident ($($ann:tt)*)) => {
9+
fn $visit(&mut self, $arg: $argty) -> Self::Output {
10+
self.0.$visit($arg.clone());
11+
let validation = self.0.validator.visitor(self.0.position).$visit($arg);
12+
if let Err(e) = validation {
13+
cold_path();
14+
self.0.record_error(crate::ParseError::ParseError { message: e.to_string(), offset: self.0.position });
15+
}
16+
}
17+
};
18+
19+
(@@reference_types TypedSelectMulti { $arg:ident: $argty:ty } => $visit:ident ($($ann:tt)*)) => {
20+
fn $visit(&mut self, $arg: $argty) -> Self::Output {
21+
self.0.$visit($arg.clone());
22+
let validation = self.0.validator.visitor(self.0.position).$visit($arg);
23+
if let Err(e) = validation {
24+
cold_path();
25+
self.0.record_error(crate::ParseError::ParseError { message: e.to_string(), offset: self.0.position });
26+
}
27+
}
28+
};
29+
30+
(@@exceptions TryTable { $arg:ident: $argty:ty } => $visit:ident ($($ann:tt)*)) => {
31+
fn $visit(&mut self, $arg: $argty) -> Self::Output {
32+
self.0.$visit($arg.clone());
33+
let validation = self.0.validator.visitor(self.0.position).$visit($arg);
34+
if let Err(e) = validation {
35+
cold_path();
36+
self.0.record_error(crate::ParseError::ParseError { message: e.to_string(), offset: self.0.position });
37+
}
38+
}
39+
};
40+
41+
(@@stack_switching Resume { cont_type_index: $cont:ty, resume_table: $table:ty } => $visit:ident ($($ann:tt)*)) => {
42+
fn $visit(&mut self, cont_type_index: $cont, resume_table: $table) -> Self::Output {
43+
self.0.$visit(cont_type_index, resume_table.clone());
44+
let validation = self.0.validator.visitor(self.0.position).$visit(cont_type_index, resume_table);
45+
if let Err(e) = validation {
46+
cold_path();
47+
self.0.record_error(crate::ParseError::ParseError { message: e.to_string(), offset: self.0.position });
48+
}
49+
}
50+
};
51+
52+
(@@stack_switching ResumeThrow { cont_type_index: $cont:ty, tag_index: $tag:ty, resume_table: $table:ty } => $visit:ident ($($ann:tt)*)) => {
53+
fn $visit(&mut self, cont_type_index: $cont, tag_index: $tag, resume_table: $table) -> Self::Output {
54+
self.0.$visit(cont_type_index, tag_index, resume_table.clone());
55+
let validation = self.0.validator.visitor(self.0.position).$visit(cont_type_index, tag_index, resume_table);
56+
if let Err(e) = validation {
57+
cold_path();
58+
self.0.record_error(crate::ParseError::ParseError { message: e.to_string(), offset: self.0.position });
59+
}
60+
}
61+
};
62+
63+
(@@stack_switching ResumeThrowRef { cont_type_index: $cont:ty, resume_table: $table:ty } => $visit:ident ($($ann:tt)*)) => {
64+
fn $visit(&mut self, cont_type_index: $cont, resume_table: $table) -> Self::Output {
65+
self.0.$visit(cont_type_index, resume_table.clone());
66+
let validation = self.0.validator.visitor(self.0.position).$visit(cont_type_index, resume_table);
67+
if let Err(e) = validation {
68+
cold_path();
69+
self.0.record_error(crate::ParseError::ParseError { message: e.to_string(), offset: self.0.position });
70+
}
71+
}
72+
};
73+
74+
(@@$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*)) => {
475
fn $visit(&mut self $($(,$arg: $argty)*)?) -> Self::Output {
5-
self.1.$visit($($($arg.clone()),*)?);
6-
self.1.validator_visitor(self.0).$visit($($($arg),*)?)?;
7-
Ok(())
76+
self.0.$visit($($($arg),*)?);
77+
let validation = self.0.validator.visitor(self.0.position).$visit($($($arg),*)?);
78+
if let Err(e) = validation {
79+
cold_path();
80+
self.0.record_error(crate::ParseError::ParseError { message: e.to_string(), offset: self.0.position });
81+
}
882
}
9-
)*};
83+
};
84+
}
85+
86+
macro_rules! validate_then_visit_simd {
87+
($( @$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*))*) => {
88+
$(validate_then_visit_simd!(@@$proposal $op $({ $($arg: $argty),* })? => $visit ($($ann)*));)*
89+
};
90+
91+
(@@$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*)) => {
92+
fn $visit(&mut self $($(,$arg: $argty)*)?) -> Self::Output {
93+
self.0.$visit($($($arg),*)?);
94+
let validation = self.0.validator.simd_visitor(self.0.position).$visit($($($arg),*)?);
95+
if let Err(e) = validation {
96+
cold_path();
97+
self.0.record_error(crate::ParseError::ParseError { message: e.to_string(), offset: self.0.position });
98+
}
99+
}
100+
};
10101
}
11102

12103
macro_rules! define_operand {
@@ -83,7 +174,7 @@ pub(crate) mod visit {
83174

84175
pub(crate) use {
85176
define_mem_operands, define_mem_operands_simd, define_mem_operands_simd_lane, define_operand, define_operands,
86-
impl_visit_operator, validate_then_visit,
177+
impl_visit_operator, validate_then_visit, validate_then_visit_simd,
87178
};
88179
}
89180

crates/parser/src/module.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,7 @@ impl ModuleReader {
208208
instructions,
209209
&mut data,
210210
options,
211+
results,
211212
self_func,
212213
import_mem_count,
213214
local_mem_alloc,

crates/parser/src/optimize.rs

Lines changed: 29 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use crate::ParserOptions;
22
use crate::macros::optimize::*;
33
use alloc::vec::Vec;
4-
use tinywasm_types::{BinOp, BinOp128, CmpOp, ConstIdx, Instruction, WasmFunctionData};
4+
use tinywasm_types::{BinOp, BinOp128, CmpOp, ConstIdx, Instruction, ValueCounts, WasmFunctionData};
55

66
pub(crate) struct OptimizeResult {
77
pub(crate) instructions: Vec<Instruction>,
@@ -12,12 +12,13 @@ pub(crate) fn optimize_instructions(
1212
mut instructions: Vec<Instruction>,
1313
function_data: &mut WasmFunctionData,
1414
options: &ParserOptions,
15+
function_results: ValueCounts,
1516
self_func_addr: u32,
1617
imported_memory_count: u32,
1718
track_local_memory_usage: bool,
1819
) -> OptimizeResult {
1920
let uses_local_memory = if options.optimize_rewrite() {
20-
rewrite(&mut instructions, self_func_addr, imported_memory_count, track_local_memory_usage)
21+
rewrite(&mut instructions, function_results, self_func_addr, imported_memory_count, track_local_memory_usage)
2122
} else {
2223
track_local_memory_usage
2324
&& instructions.iter().any(|instr| instr.memory_addr().is_some_and(|mem| mem >= imported_memory_count))
@@ -31,12 +32,20 @@ pub(crate) fn optimize_instructions(
3132

3233
fn rewrite(
3334
instrs: &mut [Instruction],
35+
function_results: ValueCounts,
3436
self_func_addr: u32,
3537
imported_memory_count: u32,
3638
track_local_memory_usage: bool,
3739
) -> bool {
3840
use Instruction::*;
3941
let mut uses_local_memory = false;
42+
let return_instr = match function_results {
43+
ValueCounts { c32: 0, c64: 0, c128: 0 } => Some(ReturnVoid),
44+
ValueCounts { c32: 1, c64: 0, c128: 0 } => Some(Return32),
45+
ValueCounts { c32: 0, c64: 1, c128: 0 } => Some(Return64),
46+
ValueCounts { c32: 0, c64: 0, c128: 1 } => Some(Return128),
47+
_ => None,
48+
};
4049

4150
for i in 0..instrs.len() {
4251
match instrs[i] {
@@ -45,6 +54,7 @@ fn rewrite(
4554
LocalCopy128(a, b) if a == b => instrs[i] = Nop,
4655
Call(addr) if addr == self_func_addr => instrs[i] = CallSelf,
4756
ReturnCall(addr) if addr == self_func_addr => instrs[i] = ReturnCallSelf,
57+
Return if let Some(return_instr) = return_instr => instrs[i] = return_instr,
4858
instr @ (I32Add | I32Mul | I32And | I32Or | I32Xor) => {
4959
let Some(op) = int_bin_op_32(instr) else { unreachable!() };
5060
rewrite!(instrs, i, [LocalGet32(a), LocalGet32(b)] => BinOpLocalLocal32(op, a, b));
@@ -370,11 +380,7 @@ fn rewrite(
370380
),
371381
Jump(ip) => {
372382
let target = resolve_jump_target(instrs, ip);
373-
if target == next_non_nop(instrs, i + 1) as u32 {
374-
instrs[i] = Nop;
375-
} else if target != ip {
376-
instrs[i] = Jump(target);
377-
}
383+
canonicalize_jump_like_with_target(instrs, i, target);
378384
}
379385
JumpIfZero(ip) => {
380386
let target = resolve_jump_target(instrs, ip);
@@ -430,10 +436,7 @@ fn rewrite(
430436
(0, CmpOp::Ne) => JumpIfNonZero64(target),
431437
(imm, op) => JumpCmpStackConst64 { target_ip: target, imm, op },
432438
});
433-
canonicalize_jump_like(instrs, i);
434-
if let JumpIfZero(current) = &mut instrs[i] {
435-
*current = target;
436-
}
439+
canonicalize_jump_like_with_target(instrs, i, target);
437440
}
438441
JumpIfNonZero(ip) => {
439442
let target = resolve_jump_target(instrs, ip);
@@ -489,42 +492,27 @@ fn rewrite(
489492
(0, CmpOp::Ne) => JumpIfNonZero64(target),
490493
(imm, op) => JumpCmpStackConst64 { target_ip: target, imm, op },
491494
});
492-
canonicalize_jump_like(instrs, i);
493-
if let JumpIfNonZero(current) = &mut instrs[i] {
494-
*current = target;
495-
}
495+
canonicalize_jump_like_with_target(instrs, i, target);
496496
}
497497
JumpIfZero32(ip) => {
498498
let target = resolve_jump_target(instrs, ip);
499499
rewrite!(instrs, i, [LocalGet32(local)] => JumpIfLocalZero32 { target_ip: target, local });
500-
canonicalize_jump_like(instrs, i);
501-
if let JumpIfZero32(current) = &mut instrs[i] {
502-
*current = target;
503-
}
500+
canonicalize_jump_like_with_target(instrs, i, target);
504501
}
505502
JumpIfNonZero32(ip) => {
506503
let target = resolve_jump_target(instrs, ip);
507504
rewrite!(instrs, i, [LocalGet32(local)] => JumpIfLocalNonZero32 { target_ip: target, local });
508-
canonicalize_jump_like(instrs, i);
509-
if let JumpIfNonZero32(current) = &mut instrs[i] {
510-
*current = target;
511-
}
505+
canonicalize_jump_like_with_target(instrs, i, target);
512506
}
513507
JumpIfZero64(ip) => {
514508
let target = resolve_jump_target(instrs, ip);
515509
rewrite!(instrs, i, [LocalGet64(local)] => JumpIfLocalZero64 { target_ip: target, local });
516-
canonicalize_jump_like(instrs, i);
517-
if let JumpIfZero64(current) = &mut instrs[i] {
518-
*current = target;
519-
}
510+
canonicalize_jump_like_with_target(instrs, i, target);
520511
}
521512
JumpIfNonZero64(ip) => {
522513
let target = resolve_jump_target(instrs, ip);
523514
rewrite!(instrs, i, [LocalGet64(local)] => JumpIfLocalNonZero64 { target_ip: target, local });
524-
canonicalize_jump_like(instrs, i);
525-
if let JumpIfNonZero64(current) = &mut instrs[i] {
526-
*current = target;
527-
}
515+
canonicalize_jump_like_with_target(instrs, i, target);
528516
}
529517
JumpCmpStackConst32 { target_ip, imm: 0, op } => {
530518
match op {
@@ -573,8 +561,8 @@ fn rewrite(
573561
_ => {}
574562
}
575563

576-
if track_local_memory_usage {
577-
uses_local_memory |= instrs[i].memory_addr().is_some_and(|mem| mem >= imported_memory_count);
564+
if track_local_memory_usage && !uses_local_memory {
565+
uses_local_memory = instrs[i].memory_addr().is_some_and(|mem| mem >= imported_memory_count);
578566
}
579567
}
580568

@@ -766,11 +754,14 @@ fn inverse_cmp_op(op: CmpOp) -> CmpOp {
766754
}
767755
}
768756

757+
const PREVIOUS_NON_NOP_BACKTRACK_LIMIT: usize = 32;
758+
769759
fn previous_non_nop<const N: usize>(instrs: &[Instruction], read: usize) -> Option<[(usize, Instruction); N]> {
770760
let mut out = [(0usize, Instruction::Nop); N];
771761
let mut filled = 0usize;
762+
let start = read.saturating_sub(PREVIOUS_NON_NOP_BACKTRACK_LIMIT);
772763

773-
for idx in (0..read).rev() {
764+
for idx in (start..read).rev() {
774765
let instr = instrs[idx];
775766
if matches!(instr, Instruction::MergeBarrier) {
776767
return None;
@@ -864,7 +855,10 @@ fn canonicalize_jump_like(instrs: &mut [Instruction], idx: usize) {
864855
return;
865856
};
866857

867-
let target = resolve_jump_target(instrs, target);
858+
canonicalize_jump_like_with_target(instrs, idx, resolve_jump_target(instrs, target));
859+
}
860+
861+
fn canonicalize_jump_like_with_target(instrs: &mut [Instruction], idx: usize, target: u32) {
868862
if matches!(instrs[idx], Instruction::Jump(_)) && target == next_non_nop(instrs, idx + 1) as u32 {
869863
instrs[idx] = Instruction::Nop;
870864
} else {

0 commit comments

Comments
 (0)