perf: 减少protoc-gen-route import引入判断

TBXark · Nov 9, 2024 · 0300f58 · 0300f58
1 parent c164941
commit 0300f58
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 11 deletions.
diff --git a/api/bot/v1/counter_bot.pb.go b/api/bot/v1/counter_bot.pb.go
diff --git a/contrib/protoc-gen-route/main.go b/contrib/protoc-gen-route/main.go
@@ -53,7 +53,7 @@ func genConfig() (*Config, error) {
 		requestType:      NewGoIdent(*requestModel),
 		responseType:     NewGoIdent(*responseModel),
 		extraType:        NewGoIdent(*extraDataModel),
-		extraConstructor: NewGoIdent(*extraDataConstructor),
+		extraConstructor: NewGoIdent(*extraDataConstructor, identIsFunc()),
 	}
 	if cfg.requestType == nil {
 		return nil, fmt.Errorf("flag request_model must be set")
@@ -68,8 +68,9 @@ func genConfig() (*Config, error) {
 }
 
 type GoIdent struct {
-	pkg   protogen.GoImportPath
-	ident string
+	pkg    protogen.GoImportPath
+	ident  string
+	isFunc bool
 }
 
 func (g GoIdent) GoIdent() protogen.GoIdent {
@@ -87,7 +88,7 @@ type Config struct {
 	extraConstructor *GoIdent
 }
 
-func NewGoIdent(s string) *GoIdent {
+func NewGoIdent(s string, options ...func(*GoIdent)) *GoIdent {
 	parts := strings.Split(s, ";")
 	if len(parts) != 2 {
 		return nil
@@ -97,3 +98,9 @@ func NewGoIdent(s string) *GoIdent {
 		ident: parts[1],
 	}
 }
+
+func identIsFunc() func(*GoIdent) {
+	return func(g *GoIdent) {
+		g.isFunc = true
+	}
+}
diff --git a/contrib/protoc-gen-route/route.go b/contrib/protoc-gen-route/route.go
@@ -53,10 +53,24 @@ func generateFileContent(gen *protogen.Plugin, file *protogen.File, g *protogen.
 		return
 	}
 	g.P("var _ = new(", contextPackage.Ident("Context"), ")")
-	g.P("var _ = new(", conf.requestType.GoIdent(), ")")
-	g.P("var _ = new(", conf.responseType.GoIdent(), ")")
-	g.P("var _ = new(", conf.extraType.GoIdent(), ")")
-	g.P("var _ = ", conf.extraConstructor.GoIdent())
+	idents := []*GoIdent{
+		conf.requestType,
+		conf.responseType,
+		conf.extraType,
+		conf.extraConstructor,
+	}
+	imported := make(map[string]struct{}, len(idents))
+	for _, i := range idents {
+		_, exist := imported[string(i.pkg)]
+		if !exist {
+			if i.isFunc {
+				g.P("var _ = ", i.GoIdent())
+			} else {
+				g.P("var _ = new(", i.GoIdent(), ")")
+			}
+			imported[string(i.pkg)] = struct{}{}
+		}
+	}
 	g.P()
 	for _, service := range file.Services {
 		genService(gen, file, g, service, conf)