gogrep: improve FieldList matching inside func types (#225)

``` $*_ matches arbitrary field list (including a nil list) $_ matches a field list of exactly 1 element ($_ $_) matches a field list of 1 named element ``` This is useful for function type params and return type lists.
quasilyte · Apr 24, 2021 · 7184b66 · 7184b66
1 parent 0d77979
commit 7184b66
Show file tree

Hide file tree

Showing 4 changed files with 77 additions and 7 deletions.
diff --git a/internal/gogrep/compile.go b/internal/gogrep/compile.go
@@ -137,6 +137,21 @@ func (c *compiler) compileOptExpr(n ast.Expr) {
 	c.compileExpr(n)
 }
 
+func (c *compiler) compileOptFieldList(n *ast.FieldList) {
+	if len(n.List) == 1 {
+		if ident, ok := n.List[0].Type.(*ast.Ident); ok && isWildName(ident.Name) {
+			// `func (...) $*result` - result could be anything
+			// `func (...) $result`  - result is a field list of 1 element
+			info := decodeWildName(ident.Name)
+			if info.Seq {
+				c.compileWildIdent(ident, true)
+				return
+			}
+		}
+	}
+	c.compileFieldList(n)
+}
+
 func (c *compiler) compileFieldList(n *ast.FieldList) {
 	c.emitInstOp(opFieldList)
 	for _, x := range n.List {
@@ -452,9 +467,9 @@ func (c *compiler) compileFuncType(n *ast.FuncType) {
 	} else {
 		c.emitInstOp(opFuncType)
 	}
-	c.compileFieldList(n.Params)
+	c.compileOptFieldList(n.Params)
 	if !void {
-		c.compileFieldList(n.Results)
+		c.compileOptFieldList(n.Results)
 	}
 }
 

diff --git a/internal/gogrep/compile_test.go b/internal/gogrep/compile_test.go
@@ -223,6 +223,50 @@ func TestCompileWildcard(t *testing.T) {
 			`EmptyPackage`,
 			` • NamedNode p`,
 		},
+
+		// $*_ in a place of a field list implies a field list of 0 or more fields.
+		// It can also match a field list of 1 element and nil.
+		`func $_() $*_ { $*_ }`: {
+			`FuncDecl`,
+			` • Node`,
+			` • FuncType`,
+			` •  • FieldList`,
+			` •  •  • End`,
+			` •  • OptNode`,
+			` • BlockStmt`,
+			` •  • NodeSeq`,
+			` •  • End`,
+		},
+
+		// $y in a place of a field list implies a field list of exactly 1 field.
+		`func $_($x $y) $y { return $x }`: {
+			`FuncDecl`,
+			` • Node`,
+			` • FuncType`,
+			` •  • FieldList`,
+			` •  •  • Field`,
+			` •  •  •  • NamedNode x`,
+			` •  •  •  • NamedNode y`,
+			` •  •  • End`,
+			` •  • FieldList`,
+			` •  •  • UnnamedField`,
+			` •  •  •  • NamedNode y`,
+			` •  •  • End`,
+			` • BlockStmt`,
+			` •  • ReturnStmt`,
+			` •  •  • NamedNode x`,
+			` •  •  • End`,
+			` •  • End`,
+		},
+
+		`func _($*_) {}`: {
+			`FuncDecl`,
+			` • Ident _`,
+			` • VoidFuncType`,
+			` •  • OptNode`,
+			` • BlockStmt`,
+			` •  • End`,
+		},
 	})
 
 	for i := range tests {

diff --git a/internal/gogrep/match.go b/internal/gogrep/match.go
@@ -244,7 +244,7 @@ func (m *matcher) matchNodeWithInst(inst instruction, n ast.Node) bool {
 		return ok && n.Results == nil && m.matchNode(n.Params)
 	case opFuncType:
 		n, ok := n.(*ast.FuncType)
-		return ok && n.Results != nil && m.matchNode(n.Params) && m.matchNode(n.Results)
+		return ok && m.matchNode(n.Params) && m.matchNode(n.Results)
 
 	case opCompositeLit:
 		n, ok := n.(*ast.CompositeLit)
@@ -266,8 +266,9 @@ func (m *matcher) matchNodeWithInst(inst instruction, n ast.Node) bool {
 		n, ok := n.(*ast.Field)
 		return ok && len(n.Names) >= 2 && m.matchIdentSlice(n.Names) && m.matchNode(n.Type)
 	case opFieldList:
+		// FieldList could be nil in places like function return types.
 		n, ok := n.(*ast.FieldList)
-		return ok && m.matchFieldSlice(n.List)
+		return ok && n != nil && m.matchFieldSlice(n.List)
 
 	case opFuncLit:
 		n, ok := n.(*ast.FuncLit)

diff --git a/internal/gogrep/match_test.go b/internal/gogrep/match_test.go
@@ -692,9 +692,19 @@ func TestMatch(t *testing.T) {
 			`package p; func a(i int) int { return i }`,
 		},
 		{`func $x(i int)`, 1, `package p; func a(i int)`},
-		{`func $x(i int) {}`, 0, `package p; func a(i int)`},
-		// {`func $_() $*_ { $*_ }`, 1, `package p; func f() {}`},
-		// {`func $_() $*_ { $*_ }`, 1, `package p; func f() (int, error) { return 3, nil }`},
+		{`func $x(i int) {}`, 0, `package p; func b(i int)`},
+		{`func $_() $*_ { $*_ }`, 1, `package p; func f() { return 0 }`},
+		{`func $_() $*_ { $*_ }`, 1, `package p; func f() int { return 0 }`},
+		{`func $_() $*_ { $*_ }`, 1, `package p; func f() (int, string) { return 0, "" }`},
+		{`func $_() $*_ { $*_ }`, 1, `package p; func f() (a int, b string) { return }`},
+		{`func $_() $*_ { $*_ }`, 1, `package p; func f() (int, error) { return 3, nil }`},
+		{`func $_($*_) $_ { $*_ }`, 0, `package p; func f() {}`},
+		{`func $_($*_) $_ { $*_ }`, 0, `package p; func f() (int, string) { return }`},
+		{`func $_($*_) $_ { $*_ }`, 0, `package p; func f() (x int) { return 0 }`},
+		{`func $_($*_) ($_ $_) { $*_ }`, 1, `package p; func f() (x int) { return 0 }`},
+		{`func $_($*_) $_ { $*_ }`, 1, `package p; func f() int { return 0 }`},
+		{`func ($x $y) $f($*_) { $*_ }`, 1, `package p; func (t *MyType) methodName() {}`},
+		{`func ($x $y) $f($*_) $*_ { $*_ }`, 1, `package p; func (t *MyType) methodName() (int, int) { return 0, 0 }`},
 
 		// Gen decl.
 		{`const $x = $y`, 1, `const a = b`},