Provide unique name from closures with same name (#488)

In YUL it's possible to declare closure with a same name, but in
different scope, which previously resulted to func name conflict during
LLVM IR construction.
To solve this, track YUL scope when building LLVM IR and a special
counter that increments every time function with the same name exists in
LLVM module.
There's still assert if somehow within the same scope there 2 functions
with the same name, which should be forbidden by Yul

Fixes: #474

---------

Co-authored-by: LJ <81748770+elle-j@users.noreply.github.com>

crates/llvm-context/src/polkavm/context/mod.rs

@@ -1,6 +1,7 @@
 //! The LLVM IR generator context.
 
 use std::cell::RefCell;
+use std::collections::BTreeMap;
 use std::collections::HashMap;
 use std::rc::Rc;
 
@@ -84,6 +85,10 @@ pub struct Context<'ctx> {
     current_function: Option<Rc<RefCell<Function<'ctx>>>>,
     /// The loop context stack.
     loop_stack: Vec<Loop<'ctx>>,
+    /// Monotonic counter for unique frontend function name mangling.
+    function_counter: usize,
+    /// Scope stack mapping Yul names to mangled LLVM names.
+    function_scope: Vec<BTreeMap<String, String>>,
     /// The PVM memory configuration.
     memory_config: SolcStandardJsonInputSettingsPolkaVMMemory,
 
@@ -246,6 +251,8 @@ impl<'ctx> Context<'ctx> {
             functions: HashMap::with_capacity(Self::FUNCTIONS_HASHMAP_INITIAL_CAPACITY),
             current_function: None,
             loop_stack: Vec::with_capacity(Self::LOOP_STACK_INITIAL_CAPACITY),
+            function_counter: 0,
+            function_scope: Vec::new(),
             memory_config,
 
             debug_info,
@@ -446,6 +453,12 @@ impl<'ctx> Context<'ctx> {
         &self.llvm_runtime
     }
 
+    pub fn get_current_scope(&mut self) -> &mut BTreeMap<String, String> {
+        self.function_scope
+            .last_mut()
+            .expect("ICE: function scope must be pushed before declaring frontend functions")
+    }
+
     /// Appends a function to the current module.
     pub fn add_function(
         &mut self,
@@ -456,12 +469,20 @@ impl<'ctx> Context<'ctx> {
         location: Option<(u32, u32)>,
         is_frontend: bool,
     ) -> anyhow::Result<Rc<RefCell<Function<'ctx>>>> {
-        assert!(
-            self.get_function(name, is_frontend).is_none(),
-            "ICE: function '{name}' declared subsequentally"
-        );
-
-        let name = self.internal_function_name(name, is_frontend);
+        let name = if is_frontend {
+            assert!(
+                !self.get_current_scope().contains_key(name),
+                "ICE: function '{name}' declared subsequently in the same scope"
+            );
+            let counter = self.function_counter;
+            self.function_counter += 1;
+            let mangled = format!("{name}_{}__{counter}", self.code_type().unwrap());
+            self.get_current_scope()
+                .insert(name.to_string(), mangled.clone());
+            mangled
+        } else {
+            name.to_string()
+        };
         let value = self.module().add_function(&name, r#type, linkage);
 
         if self.debug_info().is_some() {
@@ -524,9 +545,16 @@ impl<'ctx> Context<'ctx> {
         name: &str,
         is_frontend: bool,
     ) -> Option<Rc<RefCell<Function<'ctx>>>> {
-        self.functions
-            .get(&self.internal_function_name(name, is_frontend))
-            .cloned()
+        if is_frontend {
+            let mangled = self
+                .function_scope
+                .iter()
+                .rev()
+                .find_map(|scope| scope.get(name))?;
+            self.functions.get(mangled).cloned()
+        } else {
+            self.functions.get(name).cloned()
+        }
     }
 
     /// Returns a shared reference to the current active function.
@@ -654,6 +682,18 @@ impl<'ctx> Context<'ctx> {
             .expect("The current context is not in a loop")
     }
 
+    /// Pushes a new function scope.
+    pub fn push_function_scope(&mut self) {
+        self.function_scope.push(BTreeMap::new());
+    }
+
+    /// Pops the current function scope.
+    pub fn pop_function_scope(&mut self) {
+        self.function_scope
+            .pop()
+            .expect("ICE: tried to pop an empty function scope stack");
+    }
+
     /// Returns the debug info.
     pub fn debug_info(&self) -> Option<&DebugInfo<'ctx>> {
         self.debug_info.as_ref()
@@ -1429,15 +1469,6 @@ impl<'ctx> Context<'ctx> {
         )
     }
 
-    /// Returns the internal function name.
-    fn internal_function_name(&self, name: &str, is_frontend: bool) -> String {
-        if is_frontend {
-            format!("{name}_{}", self.code_type().unwrap())
-        } else {
-            name.to_string()
-        }
-    }
-
     /// Scans all functions in the module and removes the `MinSize` attribute
     /// if the function contains any large sdiv, udiv, srem, urem instructions with either unknown
     /// NOTE: The check here could be relaxed by checking denominator: if the denominator is

crates/resolc/src/cli_utils.rs

@@ -24,6 +24,12 @@ pub const YUL_MEMSET_CONTRACT_PATH: &str = "src/tests/data/yul/memset.yul";
 pub const YUL_RETURN_CONTRACT_PATH: &str = "src/tests/data/yul/return.yul";
 /// The invalid YUL contract test fixture path (hex literal with odd number of nibbles).
 pub const YUL_INVALID_HEX_NIBBLES_PATH: &str = "src/tests/data/yul/invalid_hex_nibbles.yul";
+/// Yul contract with duplicate function names in switch cases.
+pub const YUL_DUPLICATE_FUNCTIONS_SWITCH_PATH: &str =
+    "src/tests/data/yul/duplicate_functions_switch.yul";
+/// Yul contract with duplicate function names in deeply nested switch cases.
+pub const YUL_DUPLICATE_FUNCTIONS_DEEP_NESTING_PATH: &str =
+    "src/tests/data/yul/duplicate_functions_deep_nesting.yul";
 
 /// The standard JSON contracts test fixture path.
 pub const STANDARD_JSON_CONTRACTS_PATH: &str =

crates/resolc/src/tests/cli/yul.rs

@@ -2,7 +2,9 @@
 
 use crate::cli_utils::{
     assert_command_failure, assert_command_success, assert_equal_exit_codes, execute_resolc,
-    execute_solc, RESOLC_YUL_FLAG, SOLC_YUL_FLAG, YUL_CONTRACT_PATH, YUL_INVALID_HEX_NIBBLES_PATH,
+    execute_solc, RESOLC_YUL_FLAG, SOLC_YUL_FLAG, YUL_CONTRACT_PATH,
+    YUL_DUPLICATE_FUNCTIONS_DEEP_NESTING_PATH, YUL_DUPLICATE_FUNCTIONS_SWITCH_PATH,
+    YUL_INVALID_HEX_NIBBLES_PATH,
 };
 
 #[test]
@@ -35,3 +37,29 @@ fn bails_with_invalid_input_file() {
     let solc_result = execute_solc(&[YUL_INVALID_HEX_NIBBLES_PATH, SOLC_YUL_FLAG]);
     assert_equal_exit_codes(&solc_result, &resolc_result);
 }
+
+#[test]
+fn duplicate_functions_in_switch_cases() {
+    let resolc_result = execute_resolc(&[
+        YUL_DUPLICATE_FUNCTIONS_SWITCH_PATH,
+        RESOLC_YUL_FLAG,
+        "--bin",
+    ]);
+    assert_command_success(
+        &resolc_result,
+        "Duplicate function names in different switch cases",
+    );
+}
+
+#[test]
+fn duplicate_functions_deep_nesting() {
+    let resolc_result = execute_resolc(&[
+        YUL_DUPLICATE_FUNCTIONS_DEEP_NESTING_PATH,
+        RESOLC_YUL_FLAG,
+        "--bin",
+    ]);
+    assert_command_success(
+        &resolc_result,
+        "Duplicate function names in deeply nested switch cases",
+    );
+}

crates/resolc/src/tests/data/yul/duplicate_functions_deep_nesting.yul

@@ -0,0 +1,58 @@
+object "Test" {
+    code {
+        {
+            let size := datasize("Test_deployed")
+            codecopy(0, dataoffset("Test_deployed"), size)
+            return(0, size)
+        }
+    }
+    object "Test_deployed" {
+        code {
+            {
+                switch calldataload(0)
+                case 0 {
+                    function f() -> ret {
+                        ret := 1
+                    }
+
+                    switch calldataload(32)
+                    case 0 {
+                        function g() -> ret {
+                            ret := 10
+                        }
+                        mstore(0, add(f(), g()))
+                        return(0, 32)
+                    }
+                    case 1 {
+                        function g() -> ret {
+                            ret := 20
+                        }
+                        mstore(0, add(f(), g()))
+                        return(0, 32)
+                    }
+                }
+                case 1 {
+                    function f() -> ret {
+                        ret := 2
+                    }
+
+                    switch calldataload(32)
+                    case 0 {
+                        function g() -> ret {
+                            ret := 30
+                        }
+                        mstore(0, add(f(), g()))
+                        return(0, 32)
+                    }
+                    case 1 {
+                        function g() -> ret {
+                            ret := 40
+                        }
+                        mstore(0, add(f(), g()))
+                        return(0, 32)
+                    }
+                }
+            }
+        }
+    }
+}

crates/resolc/src/tests/data/yul/duplicate_functions_switch.yul

@@ -0,0 +1,30 @@
+object "Test" {
+    code {
+        {
+            let size := datasize("Test_deployed")
+            codecopy(0, dataoffset("Test_deployed"), size)
+            return(0, size)
+        }
+    }
+    object "Test_deployed" {
+        code {
+            {
+                switch calldataload(0)
+                case 0 {
+                    function f() -> ret {
+                        ret := 1
+                    }
+                    mstore(0, f())
+                    return(0, 32)
+                }
+                case 1 {
+                    function f() -> ret {
+                        ret := 2
+                    }
+                    mstore(0, f())
+                    return(0, 32)
+                }
+            }
+        }
+    }
+}

crates/resolc/src/tests/unit/mod.rs

@@ -9,3 +9,4 @@ mod remappings;
 mod runtime_code;
 mod standard_json;
 mod unsupported_opcodes;
+mod yul_function_scoping;

crates/resolc/src/tests/unit/yul_function_scoping.rs

@@ -0,0 +1,43 @@
+//! Tests for Yul function scoping and duplicate function name handling.
+
+use crate::test_utils::build_yul;
+
+/// Reproducer from GH-474: duplicate `f` across switch cases.
+#[test]
+fn duplicate_function_names_in_switch_cases() {
+    let code = r#"
+object "Test" {
+    code {
+        {
+            let size := datasize("Test_deployed")
+            codecopy(0, dataoffset("Test_deployed"), size)
+            return(0, size)
+        }
+    }
+    object "Test_deployed" {
+        code {
+            {
+                switch calldataload(0)
+                case 0 {
+                    function f() -> ret {
+                        ret := 1
+                    }
+                    mstore(0, f())
+                    return(0, 32)
+                }
+                case 1 {
+                    function f() -> ret {
+                        ret := 2
+                    }
+                    mstore(0, f())
+                    return(0, 32)
+                }
+            }
+        }
+    }
+}
+"#;
+
+    build_yul(&[("test.yul", code)])
+        .expect("should compile duplicate function names in switch cases");
+}

crates/yul/src/parser/statement/block.rs

@@ -136,6 +136,7 @@ impl Block {
 
 impl PolkaVMWriteLLVM for Block {
     fn into_llvm(self, context: &mut PolkaVMContext) -> anyhow::Result<()> {
+        context.push_function_scope();
         let current_function = context.current_function().borrow().name().to_owned();
         let current_block = context.basic_block();
 
@@ -223,6 +224,7 @@ impl PolkaVMWriteLLVM for Block {
         }
 
         context.pop_debug_scope();
+        context.pop_function_scope();
 
         Ok(())
     }