From decf642fbac73a7d655f819cc8db2120a77d3ecb Mon Sep 17 00:00:00 2001
From: arnaudberger <arnaudberger84@gmail.com>
Date: Wed, 3 Jul 2024 16:13:11 -0400
Subject: [PATCH] fix sanitize func

---
 ethfull/abi.go                    | 13 ----------
 ethfull/helpers.go                | 20 +++++++++++++++
 ethfull/helpers_test.go           | 42 +++++++++++++++++++++++++++++++
 ethfull/templates/build.rs.gotmpl |  8 +++---
 4 files changed, 65 insertions(+), 18 deletions(-)
 create mode 100644 ethfull/helpers.go
 create mode 100644 ethfull/helpers_test.go

diff --git a/ethfull/abi.go b/ethfull/abi.go
index 9bcb4e9..1447b88 100644
--- a/ethfull/abi.go
+++ b/ethfull/abi.go
@@ -3,7 +3,6 @@ package ethfull
 import (
 	"encoding/hex"
 	"fmt"
-	"regexp"
 	"sort"
 	"strconv"
 	"strings"
@@ -123,18 +122,6 @@ func (a *ABI) BuildEventModels() (out []codegenEvent, err error) {
 	return
 }
 
-func sanitizeABIStructName(rustABIStructName string) string {
-	reg := regexp.MustCompile("(_+)")
-	rustABIStructName = reg.ReplaceAllStringFunc(rustABIStructName, func(s string) string {
-		if len(s) > 1 {
-			return "_u"
-		}
-		return s
-	})
-
-	return rustABIStructName
-}
-
 func (a *ABI) BuildCallModels() (out []codegenCall, err error) {
 	abi := a.abi
 	names := maps.Keys(abi.FunctionsByNameMap)
diff --git a/ethfull/helpers.go b/ethfull/helpers.go
new file mode 100644
index 0000000..e7193f1
--- /dev/null
+++ b/ethfull/helpers.go
@@ -0,0 +1,20 @@
+package ethfull
+
+import (
+	"regexp"
+	"strings"
+)
+
+func sanitizeABIStructName(rustABIStructName string) string {
+	reg := regexp.MustCompile("_+")
+
+	result := reg.ReplaceAllStringFunc(rustABIStructName, func(s string) string {
+		count := len(s)
+
+		replacement := strings.Repeat("_u", count-1) + "_"
+
+		return replacement
+	})
+
+	return result
+}
diff --git a/ethfull/helpers_test.go b/ethfull/helpers_test.go
new file mode 100644
index 0000000..f5f20c8
--- /dev/null
+++ b/ethfull/helpers_test.go
@@ -0,0 +1,42 @@
+package ethfull
+
+import (
+	"github.com/stretchr/testify/require"
+	"testing"
+)
+
+func TestSanitizeStructABI(t *testing.T) {
+	cases := []struct {
+		name               string
+		inputStructAbiName string
+		expectedName       string
+	}{
+		{
+			inputStructAbiName: "foo_bar",
+			expectedName:       "foo_bar",
+		},
+		{
+			inputStructAbiName: "foo___bar",
+			expectedName:       "foo_u_u_bar",
+		},
+		{
+			inputStructAbiName: "__foo_bar",
+			expectedName:       "_u_foo_bar",
+		},
+		{
+			inputStructAbiName: "foobar__",
+			expectedName:       "foobar_u_",
+		},
+		{
+			inputStructAbiName: "__foobar",
+			expectedName:       "_u_foobar",
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.inputStructAbiName, func(t *testing.T) {
+			result := sanitizeABIStructName(c.inputStructAbiName)
+			require.Equal(t, c.expectedName, result)
+		})
+	}
+}
diff --git a/ethfull/templates/build.rs.gotmpl b/ethfull/templates/build.rs.gotmpl
index 6f104c3..b289a44 100644
--- a/ethfull/templates/build.rs.gotmpl
+++ b/ethfull/templates/build.rs.gotmpl
@@ -34,11 +34,9 @@ fn main() -> Result<(), anyhow::Error> {
         let re = Regex::new(r"_+").unwrap();
 
         let re_sanitized_abi_file = re.replace_all(&sanitized_abi_file, |caps: &regex::Captures| {
-            if caps[0].len() > 1 {
-                "_u".to_string()
-            } else {
-                caps[0].to_string()
-            }
+                let count = caps[0].len();
+                let replacement = format!("{}_", "_u".repeat(count - 1));
+                replacement
         });
 
         Abigen::from_bytes("Contract", re_sanitized_abi_file.as_bytes())?