From 3615d40410dad83ab7ef58f621c775e04125b5c4 Mon Sep 17 00:00:00 2001
From: Evgeniy Fedotov <evspare@yandex.ru>
Date: Tue, 21 Nov 2023 01:34:25 +0300
Subject: [PATCH] feature: pool paths with wildcard

---
 cmd/protoc-gen-go-vtproto/main.go |  4 +-
 generator/generatedfile.go        |  7 ++-
 generator/generator.go            | 71 ++++++++++++++++++++++++++++---
 3 files changed, 73 insertions(+), 9 deletions(-)

diff --git a/cmd/protoc-gen-go-vtproto/main.go b/cmd/protoc-gen-go-vtproto/main.go
index 6f9209e..eec0fc6 100644
--- a/cmd/protoc-gen-go-vtproto/main.go
+++ b/cmd/protoc-gen-go-vtproto/main.go
@@ -22,8 +22,10 @@ func main() {
 	var f flag.FlagSet
 
 	f.BoolVar(&cfg.AllowEmpty, "allow-empty", false, "allow generation of empty files")
-	cfg.Poolable = make(generator.ObjectSet)
+	cfg.Poolable = generator.NewObjectSet()
+	cfg.PoolableExclude = generator.NewObjectSet()
 	f.Var(&cfg.Poolable, "pool", "use memory pooling for this object")
+	f.Var(&cfg.PoolableExclude, "pool-exclude", "do not use memory pooling for this object")
 	f.BoolVar(&cfg.Wrap, "wrap", false, "generate wrapper types")
 	f.BoolVar(&cfg.WellKnownTypes, "wkt", true, "generate optimized code for well-known types")
 	f.StringVar(&features, "features", "all", "list of features to generate (separated by '+')")
diff --git a/generator/generatedfile.go b/generator/generatedfile.go
index 5c0b5aa..bfd55db 100644
--- a/generator/generatedfile.go
+++ b/generator/generatedfile.go
@@ -35,12 +35,15 @@ func (p *GeneratedFile) Ident(path, ident string) string {
 }
 
 func (b *GeneratedFile) ShouldPool(message *protogen.Message) bool {
-	if message == nil {
+	// Do not generate pool if message is nil or message excluded by external rules
+	if message == nil || b.Config.PoolableExclude.Contains(message.GoIdent) {
 		return false
 	}
-	if b.Config.Poolable[message.GoIdent] {
+
+	if b.Config.Poolable.Contains(message.GoIdent) {
 		return true
 	}
+
 	ext := proto.GetExtension(message.Desc.Options(), vtproto.E_Mempool)
 	if mempool, ok := ext.(bool); ok {
 		return mempool
diff --git a/generator/generator.go b/generator/generator.go
index d3afaa2..0fb9843 100644
--- a/generator/generator.go
+++ b/generator/generator.go
@@ -20,13 +20,38 @@ type featureHelpers struct {
 	feature int
 }
 
-type ObjectSet map[protogen.GoIdent]bool
+type ObjectSet struct {
+	wildcards map[string]struct{}
+	mp        map[protogen.GoIdent]bool
+}
+
+func NewObjectSet() ObjectSet {
+	return ObjectSet{
+		mp:        map[protogen.GoIdent]bool{},
+		wildcards: map[string]struct{}{},
+	}
+}
 
 func (o ObjectSet) String() string {
 	return fmt.Sprintf("%#v", o)
 }
 
+func (o ObjectSet) Contains(g protogen.GoIdent) bool {
+	for wildcard := range o.wildcards {
+		if Match(wildcard, fmt.Sprintf("%s.%s", string(g.GoImportPath), g.GoName)) {
+			return true
+		}
+	}
+
+	return o.mp[g]
+}
+
 func (o ObjectSet) Set(s string) error {
+	if strings.Contains(s, "*") {
+		o.wildcards[s] = struct{}{}
+		return nil
+	}
+
 	idx := strings.LastIndexByte(s, '.')
 	if idx < 0 {
 		return fmt.Errorf("invalid object name: %q", s)
@@ -37,15 +62,49 @@ func (o ObjectSet) Set(s string) error {
 		GoName:       s[idx+1:],
 	}
 
-	o[ident] = true
+	o.mp[ident] = true
 	return nil
 }
 
+func Match(pattern, name string) (matched bool) {
+	if pattern == "" {
+		return name == pattern
+	}
+
+	if pattern == "*" {
+		return true
+	}
+	return deepMatchRune([]rune(name), []rune(pattern))
+
+}
+
+func deepMatchRune(str, pattern []rune) bool {
+	for len(pattern) > 0 {
+		switch pattern[0] {
+		default:
+			if len(str) == 0 || str[0] != pattern[0] {
+				return false
+			}
+		case '*':
+			return deepMatchRune(str, pattern[1:]) ||
+				(len(str) > 0 && deepMatchRune(str[1:], pattern))
+		}
+
+		str = str[1:]
+		pattern = pattern[1:]
+	}
+
+	return len(str) == 0 && len(pattern) == 0
+}
+
 type Config struct {
-	Poolable       ObjectSet
-	Wrap           bool
-	WellKnownTypes bool
-	AllowEmpty     bool
+	// Poolable rules determines if pool feature generate for particular message
+	Poolable ObjectSet
+	// PoolableExclude rules determines if pool feature disabled for particular message
+	PoolableExclude ObjectSet
+	Wrap            bool
+	WellKnownTypes  bool
+	AllowEmpty      bool
 }
 
 type Generator struct {