Fix for #107

3 files changed, 51 insertions, 3 deletions

diff --git a/src/codegen/codegen_decl.c b/src/codegen/codegen_decl.c
index 18d81f5..b82e1af 100644
--- a/src/codegen/codegen_decl.c
+++ b/src/codegen/codegen_decl.c
@@ -953,7 +953,20 @@ void emit_impl_vtables(ParserContext *ctx, FILE *out)
             ASTNode *m = node->impl_trait.methods;
             while (m)
             {
-                const char *orig = parse_original_method_name(m->func.name);
+                // Calculate expected prefix: Struct__Trait_
+                char prefix[256];
+                sprintf(prefix, "%s__%s_", strct, trait);
+                const char *orig = m->func.name;
+                if (strncmp(orig, prefix, strlen(prefix)) == 0)
+                {
+                    orig += strlen(prefix);
+                }
+                else
+                {
+                    // Fallback if mangling schema differs (shouldn't happen)
+                    orig = parse_original_method_name(m->func.name);
+                }
+
                 fprintf(out, ".%s = (__typeof__(((%s_VTable*)0)->%s))%s__%s_%s", orig, trait, orig,
                         strct, trait, orig);
                 if (m->next)
diff --git a/src/codegen/codegen_utils.c b/src/codegen/codegen_utils.c
index a7e4925..3169eba 100644
--- a/src/codegen/codegen_utils.c
+++ b/src/codegen/codegen_utils.c
@@ -519,8 +519,22 @@ char *extract_call_args(const char *args)
 // Parse original method name from mangled name.
 const char *parse_original_method_name(const char *mangled)
 {
-    const char *last = strrchr(mangled, '_');
-    return last ? last + 1 : mangled;
+    const char *sep = strstr(mangled, "__");
+    if (!sep)
+    {
+        return mangled;
+    }
+
+    // Let's iterate to find the last `__`.
+    const char *last_double = NULL;
+    const char *p = mangled;
+    while ((p = strstr(p, "__")))
+    {
+        last_double = p;
+        p += 2;
+    }
+
+    return last_double ? last_double + 2 : mangled;
 }
 
 // Replace string type in arguments.
diff --git a/tests/features/test_traits_suite.zc b/tests/features/test_traits_suite.zc
index 7a2a262..8410de8 100644
--- a/tests/features/test_traits_suite.zc
+++ b/tests/features/test_traits_suite.zc
@@ -121,3 +121,24 @@ fn print_default_shape(s: Shape = &g_def_circle) {
 test "implicit_trait_cast_default" {
     print_default_shape();
 }
+
+trait UnderscoreTest {
+    fn method_with_underscores_123(self) -> int;
+}
+
+struct UnderscoreStruct {
+    val: int;
+}
+
+impl UnderscoreTest for UnderscoreStruct {
+    fn method_with_underscores_123(self) -> int {
+        return self.val;
+    }
+}
+
+test "trait_underscores" {
+    var u = UnderscoreStruct { val: 100 };
+    var t: UnderscoreTest = &u;
+    
+    assert(t.method_with_underscores_123() == 100, "Method with underscores call failed");
+}