gogrep: use slices instead of maps for capture data

internal/mvdan.cc/gogrep/gogrep.go

@@ -10,7 +10,11 @@ type ExprList = exprList
 
 // Parse creates a gogrep pattern out of a given string expression.
 func Parse(fset *token.FileSet, expr string, strict bool) (*Pattern, error) {
-	m := matcher{fset: fset, strict: strict}
+	m := matcher{
+		fset:    fset,
+		strict:  strict,
+		capture: make([]CapturedNode, 0, 8),
+	}
 	node, err := m.parseExpr(expr)
 	if err != nil {
 		return nil, err
@@ -26,26 +30,35 @@ type Pattern struct {
 
 // MatchData describes a successful pattern match.
 type MatchData struct {
-	Node   ast.Node
-	Values map[string]ast.Node
+	Node    ast.Node
+	Capture []CapturedNode
+}
+
+type CapturedNode struct {
+	Name string
+	Node ast.Node
+}
+
+func (data MatchData) CapturedByName(name string) (ast.Node, bool) {
+	return findNamed(data.Capture, name)
 }
 
 // Clone creates a pattern copy.
 func (p *Pattern) Clone() *Pattern {
 	clone := *p
 	clone.m = &matcher{}
 	*clone.m = *p.m
-	clone.m.values = make(map[string]ast.Node)
+	clone.m.capture = make([]CapturedNode, 0, 8)
 	return &clone
 }
 
 // MatchNode calls cb if n matches a pattern.
 func (p *Pattern) MatchNode(n ast.Node, cb func(MatchData)) {
-	p.m.values = map[string]ast.Node{}
+	p.m.capture = p.m.capture[:0]
 	if p.m.node(p.Expr, n) {
 		cb(MatchData{
-			Values: p.m.values,
-			Node:   n,
+			Capture: p.m.capture,
+			Node:    n,
 		})
 	}
 }
@@ -62,14 +75,14 @@ func (p *Pattern) matchNodeList(pattern, list nodeList, cb func(MatchData)) {
 	listLen := list.len()
 	from := 0
 	for {
-		p.m.values = map[string]ast.Node{}
+		p.m.capture = p.m.capture[:0]
 		matched, offset := p.m.nodes(pattern, list.slice(from, listLen), true)
 		if matched == nil {
 			break
 		}
 		cb(MatchData{
-			Values: p.m.values,
-			Node:   matched,
+			Capture: p.m.capture,
+			Node:    matched,
 		})
 		from += offset - 1
 		if from >= listLen {

internal/mvdan.cc/gogrep/gogrep_perf_test.go

@@ -0,0 +1,76 @@
+package gogrep
+
+import (
+	"go/token"
+	"testing"
+)
+
+func BenchmarkMatch(b *testing.B) {
+	tests := []struct {
+		name  string
+		pat   string
+		input string
+	}{
+		{
+			name:  `simpleLit`,
+			pat:   `true`,
+			input: `true`,
+		},
+		{
+			name:  `capture1`,
+			pat:   `+$x`,
+			input: `+50`,
+		},
+		{
+			name:  `capture2`,
+			pat:   `$x + $y`,
+			input: `x + 4`,
+		},
+		{
+			name:  `capture8`,
+			pat:   `f($x1, $x2, $x3, $x4, $x5, $x6, $x7, $x8)`,
+			input: `f(1, 2, 3, 4, 5, 6, 7, 8)`,
+		},
+		{
+			name:  `capture2same`,
+			pat:   `$x + $x`,
+			input: `a + a`,
+		},
+		{
+			name:  `capture8same`,
+			pat:   `f($x, $x, $x, $x, $x, $x, $x, $x)`,
+			input: `f(1, 1, 1, 1, 1, 1, 1, 1)`,
+		},
+		{
+			name:  `captureBacktrackLeft`,
+			pat:   `f($*xs, $y)`,
+			input: `f(1, 2, 3, 4, 5, 6)`,
+		},
+		{
+			name:  `captureBacktrackRight`,
+			pat:   `f($x, $*ys)`,
+			input: `f(1, 2, 3, 4, 5, 6)`,
+		},
+	}
+
+	for i := range tests {
+		test := tests[i]
+		b.Run(test.name, func(b *testing.B) {
+			fset := token.NewFileSet()
+			pat, err := Parse(fset, test.pat, true)
+			if err != nil {
+				b.Errorf("parse `%s`: %v", test.pat, err)
+				return
+			}
+			target := testParseNode(b, test.input)
+			if err != nil {
+				b.Errorf("parse target `%s`: %v", test.input, err)
+				return
+			}
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				testAllMatches(pat, target, func(m MatchData) {})
+			}
+		})
+	}
+}

internal/mvdan.cc/gogrep/gogrep_test.go

@@ -149,8 +149,8 @@ func TestCapture(t *testing.T) {
 			}
 			capture := vars{}
 			pat.MatchNode(target, func(m MatchData) {
-				for k, n := range m.Values {
-					capture[k] = sprintNode(n)
+				for _, c := range m.Capture {
+					capture[c.Name] = sprintNode(c.Node)
 				}
 			})
 			if diff := cmp.Diff(capture, test.capture); diff != "" {
@@ -623,7 +623,7 @@ func testAllMatches(p *Pattern, target ast.Node, cb func(MatchData)) {
 	})
 }
 
-func testParseNode(t *testing.T, s string) ast.Node {
+func testParseNode(t testing.TB, s string) ast.Node {
 	if strings.HasPrefix(s, "package ") {
 		file, err := parser.ParseFile(token.NewFileSet(), "string", s, 0)
 		if err != nil {

internal/mvdan.cc/gogrep/match.go

@@ -15,7 +15,7 @@ type matcher struct {
 
 	// node values recorded by name, excluding "_" (used only by the
 	// actual matching phase)
-	values map[string]ast.Node
+	capture []CapturedNode
 
 	strict bool
 }
@@ -25,12 +25,10 @@ type varInfo struct {
 	Any  bool
 }
 
-func valsCopy(values map[string]ast.Node) map[string]ast.Node {
-	v2 := make(map[string]ast.Node, len(values))
-	for k, v := range values {
-		v2[k] = v
-	}
-	return v2
+func captureCopy(capture []CapturedNode) []CapturedNode {
+	copied := make([]CapturedNode, len(capture))
+	copy(copied, capture)
+	return copied
 }
 
 // optNode is like node, but for those nodes that can be nil and are not
@@ -88,10 +86,10 @@ func (m *matcher) node(expr, node ast.Node) bool {
 			// values are discarded, matches anything
 			return true
 		}
-		prev, ok := m.values[info.Name]
+		prev, ok := findNamed(m.capture, info.Name)
 		if !ok {
 			// first occurrence, record value
-			m.values[info.Name] = node
+			m.capture = append(m.capture, CapturedNode{Name: info.Name, Node: node})
 			return true
 		}
 		// multiple uses must match
@@ -396,7 +394,7 @@ func (m *matcher) nodes(ns1, ns2 nodeList, partial bool) (ast.Node, int) {
 	// with a different "any of" match while discarding any matches
 	// we found while trying it.
 	type restart struct {
-		matches      map[string]ast.Node
+		matches      []CapturedNode
 		next1, next2 int
 	}
 	// We need to stack these because otherwise some edge cases
@@ -409,12 +407,12 @@ func (m *matcher) nodes(ns1, ns2 nodeList, partial bool) (ast.Node, int) {
 		if n2 > ns2len {
 			return // would be discarded anyway
 		}
-		stack = append(stack, restart{valsCopy(m.values), n1, n2})
+		stack = append(stack, restart{captureCopy(m.capture), n1, n2})
 		next1, next2 = n1, n2
 	}
 	pop := func() {
 		i1, i2 = next1, next2
-		m.values = stack[len(stack)-1].matches
+		m.capture = stack[len(stack)-1].matches
 		stack = stack[:len(stack)-1]
 		next1, next2 = 0, 0
 		if len(stack) > 0 {
@@ -434,11 +432,11 @@ func (m *matcher) nodes(ns1, ns2 nodeList, partial bool) (ast.Node, int) {
 		}
 		list := ns2.slice(wildStart, i2)
 		// check that it matches any nodes found elsewhere
-		prev, ok := m.values[wildName]
+		prev, ok := findNamed(m.capture, wildName)
 		if ok && !m.node(prev, list) {
 			return false
 		}
-		m.values[wildName] = list
+		m.capture = append(m.capture, CapturedNode{Name: wildName, Node: list})
 		return true
 	}
 	for i1 < ns1len || i2 < ns2len {
@@ -642,3 +640,12 @@ func literalValue(lit *ast.BasicLit) interface{} {
 	}
 	return nil
 }
+
+func findNamed(capture []CapturedNode, name string) (ast.Node, bool) {
+	for _, c := range capture {
+		if c.Name == name {
+			return c.Node, true
+		}
+	}
+	return nil, false
+}

ruleguard/filters.go

@@ -224,7 +224,8 @@ func makeTextFilter(src, varname string, op token.Token, rhsVarname string) filt
 	return func(params *filterParams) matchFilterResult {
 		s1 := params.nodeText(params.subExpr(varname))
 		lhsValue := constant.MakeString(string(s1))
-		s2 := params.nodeText(params.values[rhsVarname])
+		n, _ := params.match.CapturedByName(rhsVarname)
+		s2 := params.nodeText(n)
 		rhsValue := constant.MakeString(string(s2))
 		if constant.Compare(lhsValue, op, rhsValue) {
 			return filterSuccess

ruleguard/gorule.go

@@ -53,7 +53,7 @@ type filterParams struct {
 
 	importer *goImporter
 
-	values map[string]ast.Node
+	match gogrep.MatchData
 
 	nodeText func(n ast.Node) []byte
 
@@ -62,7 +62,8 @@ type filterParams struct {
 }
 
 func (params *filterParams) subExpr(name string) ast.Expr {
-	switch n := params.values[name].(type) {
+	n, _ := params.match.CapturedByName(name)
+	switch n := n.(type) {
 	case ast.Expr:
 		return n
 	case *ast.ExprStmt:

ruleguard/ruleguard_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/quasilyte/go-ruleguard/internal/mvdan.cc/gogrep"
 )
 
 func TestRenderMessage(t *testing.T) {
@@ -89,12 +90,19 @@ func TestRenderMessage(t *testing.T) {
 		Fset: token.NewFileSet(),
 	}
 	for _, test := range tests {
-		nodes := make(map[string]ast.Node, len(test.vars))
-		for _, v := range test.vars {
-			nodes[v] = &ast.Ident{Name: v + "var"}
+		capture := make([]gogrep.CapturedNode, len(test.vars))
+		for i, v := range test.vars {
+			capture[i] = gogrep.CapturedNode{
+				Name: v,
+				Node: &ast.Ident{Name: v + "var"},
+			}
 		}
 
-		have := rr.renderMessage(test.msg, &ast.Ident{Name: "dd"}, nodes, false)
+		m := gogrep.MatchData{
+			Node:    &ast.Ident{Name: "dd"},
+			Capture: capture,
+		}
+		have := rr.renderMessage(test.msg, m, false)
 		if diff := cmp.Diff(have, test.want); diff != "" {
 			t.Errorf("render %s %v:\n(+want -have)\n%s", test.msg, test.vars, diff)
 		}