11use crate :: ParserOptions ;
22use crate :: macros:: optimize:: * ;
33use 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
66pub ( 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
3233fn 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+
769759fn 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