Fix addMod

j178 · Jan 29, 2023 · f3ebe00 · f3ebe00
1 parent 9b5cc9d
commit f3ebe00
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/config/config.go b/config/config.go
@@ -201,8 +201,8 @@ func Default() *Config {
 					Modifiers: []Modifier{
 						{Name: "removeUselessComments"},
 						{Name: "changeReceiverName"},
-						{Name: "addMod"},
 						{Name: "addNamedReturn"},
+						{Name: "addMod"},
 					},
 				},
 			},

diff --git a/lang/go.go b/lang/go.go
@@ -103,8 +103,13 @@ func addMod(code string, q *leetcode.QuestionData) string {
 	var returnType string
 	for _, line := range lines {
 		if strings.HasPrefix(line, "func ") {
-			rightBrace := strings.LastIndex(line, ")")
-			returnType = strings.TrimSpace(line[rightBrace+1 : strings.LastIndex(line, "{")])
+			if strings.Count(line, "(") == 1 {
+				rightBrace := strings.LastIndex(line, ")")
+				returnType = strings.TrimSpace(line[rightBrace+1 : strings.LastIndex(line, "{")])
+			} else {
+				s := line[strings.LastIndex(line, "(")+1 : strings.LastIndex(line, ")")]
+				returnType = s[strings.LastIndex(s, " ")+1:]
+			}
 			newLines = append(newLines, line)
 			newLines = append(newLines, "\tconst mod int = 1e9 + 7\n")
 		} else if strings.HasPrefix(line, "\treturn") {