cmd/gogrep,filters: add basic var text filters support

cmd/gogrep/filters.go

@@ -4,7 +4,6 @@ import (
 	"fmt"
 	"go/ast"
 	"go/token"
-	"os"
 	"path/filepath"
 	"strconv"
 
@@ -47,6 +46,7 @@ const (
 	opVarIsIntLit
 	opVarIsFloatLit
 	opVarIsComplexLit
+	opVarText
 	opVarIsHot
 )
 
@@ -55,6 +55,14 @@ type filterContext struct {
 	w *worker
 }
 
+func (ctx *filterContext) NodeText(varname string) []byte {
+	n, ok := capturedByName(ctx.m, varname)
+	if !ok {
+		return nil
+	}
+	return ctx.w.nodeText(n)
+}
+
 func applyFilter(ctx filterContext, f *filters.Expr, n ast.Node) bool {
 	switch f.Op {
 	case filters.OpNot:
@@ -123,11 +131,26 @@ func applyFilter(ctx filterContext, f *filters.Expr, n ast.Node) bool {
 		})
 		return isHot
 
+	case filters.OpEq:
+		return applyEqFilter(ctx, f, n)
+	case filters.OpNotEq:
+		return !applyEqFilter(ctx, f, n)
+
 	default:
-		fmt.Fprintf(os.Stderr, "can't handle %s\n", filters.Sprint(ctx.w.filterInfo, f))
+		panic(fmt.Sprintf("can't handle %s\n", filters.Sprint(ctx.w.filterInfo, f)))
 	}
+}
 
-	return true
+func applyEqFilter(ctx filterContext, f *filters.Expr, n ast.Node) bool {
+	x := f.Args[0]
+	y := f.Args[1]
+	switch x.Op {
+	case opVarText:
+		if y.Op == filters.OpString {
+			return string(ctx.NodeText(x.Str)) == y.Str
+		}
+	}
+	panic(fmt.Sprintf("can't handle %s\n", filters.Sprint(ctx.w.filterInfo, f)))
 }
 
 func isPureExpr(expr ast.Expr) bool {

cmd/gogrep/main.go

@@ -330,6 +330,7 @@ func (p *program) compileFilter() error {
 		"IsFloatLit":   opVarIsFloatLit,
 		"IsComplexLit": opVarIsComplexLit,
 		"IsHot":        opVarIsHot,
+		"Text":         opVarText,
 	}
 	optab := filters.NewOperationTable(varOps)
 	expr, info, err := filters.Parse(optab, p.args.filter)

cmd/gogrep/worker.go

@@ -1,9 +1,11 @@
 package main
 
 import (
+	"bytes"
 	"fmt"
 	"go/ast"
 	"go/parser"
+	"go/printer"
 	"go/token"
 	"os"
 	"path/filepath"
@@ -161,6 +163,31 @@ func (w *worker) initMatchText(m *match, startPos, endPos int) {
 	m.matchLength = endPos - startPos
 }
 
+func (w *worker) nodeText(n ast.Node) []byte {
+	if gogrep.IsEmptyNodeSlice(n) {
+		return nil
+	}
+
+	from := w.fset.Position(n.Pos()).Offset
+	to := w.fset.Position(n.End()).Offset
+	src := w.data
+	if (from >= 0 && from < len(src)) && (to >= 0 && to < len(src)) {
+		return src[from:to]
+	}
+
+	// Go printer would panic on comments.
+	if n, ok := n.(*ast.Comment); ok {
+		return []byte(n.Text)
+	}
+
+	// Fallback to the printer.
+	var buf bytes.Buffer
+	if err := printer.Fprint(&buf, w.fset, n); err != nil {
+		panic(err)
+	}
+	return buf.Bytes()
+}
+
 func isAutogenFile(f *ast.File) bool {
 	for _, comment := range f.Comments {
 		if isAutogenComment(comment) {

filters/filters.go

@@ -119,6 +119,9 @@ const (
 	// OpNop = do nothing (should be optimized-away, unless it's a top level op)
 	OpNop Operation = math.MaxUint32 - iota
 
+	// OpString is a string literal that holds the value inside $Str.
+	OpString
+
 	// OpNot = !$Args[0]
 	OpNot
 
@@ -128,6 +131,12 @@ const (
 	// OpOr = $Args[0] || $Args[1]
 	OpOr
 
+	// OpEq = $Args[0] == $Args[1]
+	OpEq
+
+	// OpNotEq = $Args[0] != $Args[1]
+	OpNotEq
+
 	// OpFunctionVarFunc = function.$Str()
 	OpFunctionVarFunc
 

filters/parse.go

@@ -5,6 +5,7 @@ import (
 	"go/ast"
 	"go/parser"
 	"go/token"
+	"strconv"
 	"strings"
 )
 
@@ -98,11 +99,23 @@ func (p *filterParser) convertExpr(root ast.Expr) (*Expr, error) {
 		return p.convertExpr(root.X)
 	case *ast.CallExpr:
 		return p.convertCallExpr(root)
+	case *ast.BasicLit:
+		return p.convertBasicLit(root)
 	default:
 		return nil, fmt.Errorf("convert expr: unsupported %T", root)
 	}
 }
 
+func (p *filterParser) convertBasicLit(root *ast.BasicLit) (*Expr, error) {
+	switch root.Kind {
+	case token.STRING:
+		val, err := strconv.Unquote(root.Value)
+		return &Expr{Op: OpString, Str: val}, err
+	default:
+		return nil, fmt.Errorf("convert basic lit: unsupported %s", root.Kind)
+	}
+}
+
 func (p *filterParser) convertCallExpr(root *ast.CallExpr) (*Expr, error) {
 	if selector, ok := root.Fun.(*ast.SelectorExpr); ok {
 		return p.convertMethodCallExpr(root, selector)
@@ -198,8 +211,7 @@ func (p *filterParser) convertBinaryExpr(root *ast.BinaryExpr) (*Expr, error) {
 }
 
 func (p *filterParser) convertBinaryExprXY(op token.Token, x, y ast.Expr) (*Expr, error) {
-	switch op {
-	case token.LOR:
+	if op == token.LOR {
 		insideOr := p.insideOr
 		p.insideOr = true
 		lhs, err := p.convertExpr(x)
@@ -212,17 +224,25 @@ func (p *filterParser) convertBinaryExprXY(op token.Token, x, y ast.Expr) (*Expr
 		}
 		p.insideOr = insideOr
 		return &Expr{Op: OpOr, Args: []*Expr{lhs, rhs}}, nil
+	}
+
+	lhs, err := p.convertExpr(x)
+	if err != nil {
+		return nil, err
+	}
+	rhs, err := p.convertExpr(y)
+	if err != nil {
+		return nil, err
+	}
 
+	switch op {
 	case token.LAND:
-		lhs, err := p.convertExpr(x)
-		if err != nil {
-			return nil, err
-		}
-		rhs, err := p.convertExpr(y)
-		if err != nil {
-			return nil, err
-		}
 		return &Expr{Op: OpAnd, Args: []*Expr{lhs, rhs}}, nil
+
+	case token.EQL:
+		return &Expr{Op: OpEq, Args: []*Expr{lhs, rhs}}, nil
+	case token.NEQ:
+		return &Expr{Op: OpNotEq, Args: []*Expr{lhs, rhs}}, nil
 	}
 
 	return nil, fmt.Errorf("convert binary expr: unsupported %s", op)

filters/parse_test.go

@@ -117,15 +117,38 @@ func TestParse(t *testing.T) {
 			expr:  `(%IsPure "_Dollar2_")`,
 			info:  `$_Dollar2_`,
 		},
+
+		{
+			input: `$x.Text() == "String"`,
+			expr:  `(Eq (%Text "x") (String "String"))`,
+			info:  `$x`,
+		},
+		{
+			input: `"String" == $x.Text()`,
+			expr:  `(Eq (%Text "x") (String "String"))`,
+			info:  `$x`,
+		},
+		{
+			input: `$x.Text() != "String"`,
+			expr:  `(NotEq (%Text "x") (String "String"))`,
+			info:  `$x`,
+		},
+		{
+			input: `"String" != $x.Text()`,
+			expr:  `(NotEq (%Text "x") (String "String"))`,
+			info:  `$x`,
+		},
 	}
 
 	const (
 		opVarIsConst = iota + 1
 		opVarIsPure
+		opVarText
 	)
 	varOps := map[string]Operation{
 		"IsConst": opVarIsConst,
 		"IsPure":  opVarIsPure,
+		"Text":    opVarText,
 	}
 	optab := NewOperationTable(varOps)