improve method selection on contracts

xunilrj · xunilrj · commit 62f7022c1b9b · 2025-10-16T09:50:20.000-03:00
diff --git a/sway-core/src/semantic_analysis/ast_node/declaration/auto_impl/abi_encoding.rs b/sway-core/src/semantic_analysis/ast_node/declaration/auto_impl/abi_encoding.rs
@@ -9,6 +9,8 @@ use crate::{
     Engines, TypeInfo, TypeParameter,
 };
 use std::collections::BTreeMap;
+use hashbrown::HashMap;
+use itertools::{join, Itertools};
 use sway_error::{
     error::CompileError,
     handler::{ErrorEmitted, Handler},
@@ -312,15 +314,16 @@ where
         fallback_fn: Option<DeclId<TyFunctionDecl>>,
         handler: &Handler,
     ) -> Result<TyAstNode, ErrorEmitted> {
-        let mut code = String::new();
-
         let mut reads = false;
         let mut writes = false;
 
         // used to check for name collisions
         let mut contract_methods: BTreeMap<String, Vec<Span>> = <_>::default();
 
+        let mut arm_by_size = BTreeMap::<usize, String>::default();
+
         // generate code
+        let mut method_names = String::new();
         for r in contract_fns {
             let decl = engines.de().get(r);
 
@@ -382,8 +385,20 @@ where
             };
 
             let method_name = decl.name.as_str();
+            let offset = if let Some(offset) = method_names.find(method_name) {
+                offset
+            } else {
+                let offset = method_names.len();
+                method_names.push_str(method_name);
+                offset
+            };
+
+            let method_name_len = method_name.len();
+            let code = arm_by_size.entry(method_name.len()).or_default();
 
-            code.push_str(&format!("if _method_name == \"{method_name}\" {{\n"));
+            code.push_str(&format!("
+            let is_this_method = asm(r, ptr: _method_name_ptr, name: _method_names_ptr, len: {method_name_len}) {{ addi r name i{offset}; meq r ptr r len; r: bool }};
+            if is_this_method {{\n"));
 
             if args_types == "()" {
                 code.push_str(&format!(
@@ -454,10 +469,18 @@ where
             (false, false) => "",
         };
 
+        let code = arm_by_size.iter().map(|(len, code)| {
+            format!("if _method_len == {len} {{ {code} }}")
+        }).join("");
         let code = format!(
             "{att} pub fn __entry() {{
+            let _method_names = \"{method_names}\";
             let mut _buffer = BufferReader::from_second_parameter();
-            let _method_name = decode_first_param::<str>();
+            
+            let mut _first_param_buffer = BufferReader::from_first_parameter();
+            let _method_len = _first_param_buffer.read::<u64>();
+            let _method_name_ptr = _first_param_buffer.ptr();
+            let _method_names_ptr = _method_names.as_ptr();
             {code}
             {fallback}
         }}"
diff --git a/sway-core/src/semantic_analysis/ast_node/expression/match_expression/typed/typed_match_expression.rs b/sway-core/src/semantic_analysis/ast_node/expression/match_expression/typed/typed_match_expression.rs
@@ -112,7 +112,7 @@ impl ty::TyMatchExpression {
         let typed_if_exp =
             handler.scope(
                 |handler| match &*ctx.engines().te().get(self.value_type_id) {
-                    TypeInfo::StringSlice => self.desugar_to_radix_trie(ctx),
+                    //TypeInfo::StringSlice => self.desugar_to_radix_trie(ctx),
                     _ => self.desugar_to_typed_if_expression(instantiate, ctx, handler),
                 },
             )?;
diff --git a/sway-ir/src/parser.rs b/sway-ir/src/parser.rs
@@ -1261,7 +1261,7 @@ mod ir_builder {
                                 },
                             )
                             .collect();
-                        let md_idx = meta_idx.map(|mdi| self.md_map.get(&mdi).unwrap()).copied();
+                        let md_idx = meta_idx.and_then(|mdi| self.md_map.get(&mdi)).copied();
                         let return_type = return_type.to_ir_type(context);
                         block
                             .append(context)
diff --git a/sway-lib-std/src/codec.sw b/sway-lib-std/src/codec.sw
@@ -156,6 +156,10 @@ impl BufferReader {
     {
         T::abi_decode(self)
     }
+
+    pub fn ptr(self) -> raw_ptr {
+        self.ptr
+    }
 }
 
 // Encode
diff --git a/test/src/e2e_vm_tests/test_programs/should_pass/language/match_expressions_all/snapshot.toml b/test/src/e2e_vm_tests/test_programs/should_pass/language/match_expressions_all/snapshot.toml
@@ -0,0 +1,4 @@
+cmds = [
+    "forc build --path {root}",
+    "forc build --path {root} --ir final --asm final | filter-fn {name} return_match_on_str_slice"
+]
diff --git a/test/src/e2e_vm_tests/test_programs/should_pass/language/match_expressions_all/stdout.snap b/test/src/e2e_vm_tests/test_programs/should_pass/language/match_expressions_all/stdout.snap
@@ -0,0 +1,125 @@
+---
+source: test/src/snapshot/mod.rs
+---
+> forc build --path test/src/e2e_vm_tests/test_programs/should_pass/language/match_expressions_all
+exit status: 0
+output:
+    Building test/src/e2e_vm_tests/test_programs/should_pass/language/match_expressions_all
+   Compiling library std (test/src/e2e_vm_tests/reduced_std_libs/sway-lib-std-assert)
+   Compiling script match_expressions_all (test/src/e2e_vm_tests/test_programs/should_pass/language/match_expressions_all)
+let packed_string = "get_a_b"
+if str.len() == 5
+    if str[0..4] == "get_" at packed_string[0]
+        if str[4..5] == "b" at packed_string[6]
+            return branch 2
+        if str[4..5] == "a" at packed_string[4]
+            return branch 0
+        return wildcard branch
+    return wildcard branch
+if str.len() == 7
+    if str[0..7] == "get_a_b" at packed_string[0]
+        return branch 1
+    return wildcard branch
+return wildcard branch
+let packed_string = "get_a_b"
+if str.len() == 5
+    if str[0..4] == "get_" at packed_string[0]
+        if str[4..5] == "b" at packed_string[6]
+            return branch 2
+        if str[4..5] == "a" at packed_string[4]
+            return branch 0
+        return wildcard branch
+    return wildcard branch
+if str.len() == 7
+    if str[0..7] == "get_a_b" at packed_string[0]
+        return branch 1
+    return wildcard branch
+return wildcard branch
+    Finished debug [unoptimized + fuel] target(s) [2.736 KB] in ???
+
+> forc build --path test/src/e2e_vm_tests/test_programs/should_pass/language/match_expressions_all --ir final --asm final | filter-fn match_expressions_all return_match_on_str_slice
+
+pshl i15                      ; save registers 16..40
+pshh i524288                  ; save registers 40..64
+move $$locbase $sp            ; save locals base register for function return_match_on_str_slice_8
+cfei i144                     ; allocate 144 bytes for locals and 0 slots for call arguments
+move $r3 $$reta               ; save return address
+addi $r0 $$locbase i128       ; get offset to local __ptr slice
+mcpi $r0 $$arg0 i16           ; copy memory
+movi $r1 i7                   ; initialize constant into register
+lw   $r0 $$locbase i17
+eq   $r0 $r0 $r1
+jnzf $r0 $zero i55
+movi $r1 i5                   ; initialize constant into register
+lw   $r0 $$locbase i17
+eq   $r0 $r0 $r1
+movi $r1 i1000                ; initialize constant into register
+jnzf $r0 $zero i1
+jmpf $zero i48
+addr $r0 data_NonConfigurable_9; get __const_global0's address in data section
+addi $r1 $$locbase i32        ; get offset to local __ptr { ptr, u64 }
+sw   $$locbase $r0 i4         ; store word
+movi $r0 i7                   ; initialize constant into register
+sw   $$locbase $r0 i5         ; store word
+addi $r0 $$locbase i48        ; get offset to local __ptr slice
+mcpi $r0 $r1 i16              ; copy memory
+movi $r2 i4                   ; initialize constant into register
+lw   $r0 $$locbase i16
+addi $r0 $r0 i0
+lw   $r1 $$locbase i6
+addi $r1 $r1 i0
+meq  $r0 $r0 $r1 $r2
+movi $r1 i1000                ; initialize constant into register
+jnzf $r0 $zero i1
+jmpf $zero i31
+addr $r0 data_NonConfigurable_9; get __const_global1's address in data section
+addi $r1 $$locbase i64        ; get offset to local __ptr { ptr, u64 }
+sw   $$locbase $r0 i8         ; store word
+movi $r0 i7                   ; initialize constant into register
+sw   $$locbase $r0 i9         ; store word
+addi $r0 $$locbase i80        ; get offset to local __ptr slice
+mcpi $r0 $r1 i16              ; copy memory
+lw   $r0 $$locbase i16
+addi $r0 $r0 i4
+lw   $r1 $$locbase i10
+addi $r1 $r1 i4
+meq  $r0 $r0 $r1 $one
+move $r1 $one                 ; move parameter from branch to block argument
+jnzf $r0 $zero i16
+addr $r0 data_NonConfigurable_9; get __const_global2's address in data section
+addi $r1 $$locbase i96        ; get offset to local __ptr { ptr, u64 }
+sw   $$locbase $r0 i12        ; store word
+movi $r0 i7                   ; initialize constant into register
+sw   $$locbase $r0 i13        ; store word
+addi $r0 $$locbase i112       ; get offset to local __ptr slice
+mcpi $r0 $r1 i16              ; copy memory
+lw   $r0 $$locbase i16
+addi $r0 $r0 i4
+lw   $r1 $$locbase i14
+addi $r1 $r1 i6
+meq  $r0 $r0 $r1 $one
+movi $r1 i3                   ; initialize constant into register
+jnzf $r0 $zero i1
+movi $r1 i1000                ; initialize constant into register
+jmpf $zero i16
+addr $r0 data_NonConfigurable_9; get __const_global's address in data section
+sw   $$locbase $r0 i0         ; store word
+movi $r0 i7                   ; initialize constant into register
+sw   $$locbase $r0 i1         ; store word
+addi $r0 $$locbase i16        ; get offset to local __ptr slice
+mcpi $r0 $$locbase i16        ; copy memory
+movi $r2 i7                   ; initialize constant into register
+lw   $r0 $$locbase i16
+addi $r0 $r0 i0
+lw   $r1 $$locbase i2
+addi $r1 $r1 i0
+meq  $r0 $r0 $r1 $r2
+movi $r1 i2                   ; initialize constant into register
+jnzf $r0 $zero i1
+movi $r1 i1000                ; initialize constant into register
+move $$retv $r1               ; set return value
+cfsi i144                     ; free 144 bytes for locals and 0 slots for extra call arguments
+move $$reta $r3               ; restore return address
+poph i524288                  ; restore registers 40..64
+popl i15                      ; restore registers 16..40
+jal  $zero $$reta i0          ; return from call
diff --git a/test/src/e2e_vm_tests/test_programs/should_pass/test_contracts/const_of_contract_call/stdout.snap b/test/src/e2e_vm_tests/test_programs/should_pass/test_contracts/const_of_contract_call/stdout.snap

Original file line number	Diff line number	Diff line change
`@@ -1261,7 +1261,7 @@ mod ir_builder {`
`1261`	`1261`	`},`
`1262`	`1262`	`)`
`1263`	`1263`	`.collect();`
`1264`		`- let md_idx = meta_idx.map(\|mdi\| self.md_map.get(&mdi).unwrap()).copied();`
	`1264`	`+ let md_idx = meta_idx.and_then(\|mdi\| self.md_map.get(&mdi)).copied();`
`1265`	`1265`	`let return_type = return_type.to_ir_type(context);`
`1266`	`1266`	`block`
`1267`	`1267`	`.append(context)`
Original file line number	Diff line number	Diff line change
`@@ -156,6 +156,10 @@ impl BufferReader {`
`156`	`156`	`{`
`157`	`157`	`T::abi_decode(self)`
`158`	`158`	`}`
	`159`	`+`
	`160`	`+ pub fn ptr(self) -> raw_ptr {`
	`161`	`+ self.ptr`
	`162`	`+ }`
`159`	`163`	`}`
`160`	`164`
`161`	`165`	`// Encode`