go/expect: add marker support for go.mod files

This change adds some basic marker support for go.mod files inside of go/expect. It requires all markers to be of the form "//@mark()", where mark can be anything. It is the same format as .go files, only difference is that it needs to have "//" since that is the only comment marker that go.mod files recognize.

Updates golang/go#36091

Change-Id: Ib9e325e01020181b8cee1c1be6bb257726ce913d
Reviewed-on: https://go-review.googlesource.com/c/tools/+/216838
Run-TryBot: Rohan Challa <rohan@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>

go/expect/expect.go

@@ -56,6 +56,7 @@ import (
 	"bytes"
 	"fmt"
 	"go/token"
+	"path/filepath"
 	"regexp"
 )
 
@@ -89,12 +90,20 @@ func MatchBefore(fset *token.FileSet, readFile ReadFile, end token.Pos, pattern
 		return token.NoPos, token.NoPos, fmt.Errorf("invalid file: %v", err)
 	}
 	position := f.Position(end)
-	startOffset := f.Offset(lineStart(f, position.Line))
+	startOffset := f.Offset(f.LineStart(position.Line))
 	endOffset := f.Offset(end)
 	line := content[startOffset:endOffset]
 	matchStart, matchEnd := -1, -1
 	switch pattern := pattern.(type) {
 	case string:
+		// If the file is a go.mod and we are matching // indirect, then we
+		// need to look for it on the line after the current line.
+		// TODO(golang/go#36894): have a more intuitive approach for // indirect
+		if filepath.Ext(f.Name()) == ".mod" && pattern == "// indirect" {
+			startOffset = f.Offset(f.LineStart(position.Line + 1))
+			endOffset = f.Offset(lineEnd(f, position.Line+1))
+			line = content[startOffset:endOffset]
+		}
 		bytePattern := []byte(pattern)
 		matchStart = bytes.Index(line, bytePattern)
 		if matchStart >= 0 {
@@ -118,32 +127,9 @@ func MatchBefore(fset *token.FileSet, readFile ReadFile, end token.Pos, pattern
 	return f.Pos(startOffset + matchStart), f.Pos(startOffset + matchEnd), nil
 }
 
-// this functionality was borrowed from the analysisutil package
+func lineEnd(f *token.File, line int) token.Pos {
-func lineStart(f *token.File, line int) token.Pos {
+	if line >= f.LineCount() {
-	// Use binary search to find the start offset of this line.
+		return token.Pos(f.Size() + 1)
-	//
-	// TODO(adonovan): eventually replace this function with the
-	// simpler and more efficient (*go/token.File).LineStart, added
-	// in go1.12.
-
-	min := 0        // inclusive
-	max := f.Size() // exclusive
-	for {
-		offset := (min + max) / 2
-		pos := f.Pos(offset)
-		posn := f.Position(pos)
-		if posn.Line == line {
-			return pos - (token.Pos(posn.Column) - 1)
-		}
-
-		if min+1 >= max {
-			return token.NoPos
-		}
-
-		if posn.Line < line {
-			min = offset
-		} else {
-			max = offset
-		}
 	}
+	return f.LineStart(line + 1)
 }

go/expect/expect_test.go

@@ -14,14 +14,16 @@ import (
 )
 
 func TestMarker(t *testing.T) {
-	const filename = "testdata/test.go"
+	for _, tt := range []struct {
-	content, err := ioutil.ReadFile(filename)
+		filename      string
-	if err != nil {
+		expectNotes   int
-		t.Fatal(err)
+		expectMarkers map[string]string
-	}
+		expectChecks  map[string][]interface{}
-
+	}{
-	const expectNotes = 13
+		{
-	expectMarkers := map[string]string{
+			filename:    "testdata/test.go",
+			expectNotes: 13,
+			expectMarkers: map[string]string{
 				"αSimpleMarker": "α",
 				"OffsetMarker":  "β",
 				"RegexMarker":   "γ",
@@ -33,17 +35,38 @@ func TestMarker(t *testing.T) {
 				"LineComment":   "someFunc",
 				"NonIdentifier": "+",
 				"StringMarker":  "\"hello\"",
-	}
+			},
-	expectChecks := map[string][]interface{}{
+			expectChecks: map[string][]interface{}{
 				"αSimpleMarker": nil,
 				"StringAndInt":  []interface{}{"Number %d", int64(12)},
 				"Bool":          []interface{}{true},
+			},
+		},
+		{
+			filename:    "testdata/go.mod",
+			expectNotes: 3,
+			expectMarkers: map[string]string{
+				"αMarker":        "αfake1α",
+				"IndirectMarker": "// indirect",
+				"βMarker":        "require golang.org/modfile v0.0.0",
+			},
+		},
+	} {
+		t.Run(tt.filename, func(t *testing.T) {
+			content, err := ioutil.ReadFile(tt.filename)
+			if err != nil {
+				t.Fatal(err)
 			}
-
 			readFile := func(string) ([]byte, error) { return content, nil }
+
 			markers := make(map[string]token.Pos)
-	for name, tok := range expectMarkers {
+			for name, tok := range tt.expectMarkers {
 				offset := bytes.Index(content, []byte(tok))
+				// Handle special case where we look for // indirect and we
+				// need to search the next line.
+				if tok == "// indirect" {
+					offset = bytes.Index(content, []byte(" "+tok)) + 1
+				}
 				markers[name] = token.Pos(offset + 1)
 				end := bytes.Index(content[offset:], []byte(tok))
 				if end > 0 {
@@ -52,12 +75,12 @@ func TestMarker(t *testing.T) {
 			}
 
 			fset := token.NewFileSet()
-	notes, err := expect.Parse(fset, filename, nil)
+			notes, err := expect.Parse(fset, tt.filename, content)
 			if err != nil {
 				t.Fatalf("Failed to extract notes: %v", err)
 			}
-	if len(notes) != expectNotes {
+			if len(notes) != tt.expectNotes {
-		t.Errorf("Expected %v notes, got %v", expectNotes, len(notes))
+				t.Errorf("Expected %v notes, got %v", tt.expectNotes, len(notes))
 			}
 			for _, n := range notes {
 				switch {
@@ -93,7 +116,7 @@ func TestMarker(t *testing.T) {
 						t.Errorf("%v: identifier, got %T", fset.Position(n.Pos), n.Args[0])
 						continue
 					}
-			args, ok := expectChecks[string(ident)]
+					args, ok := tt.expectChecks[string(ident)]
 					if !ok {
 						t.Errorf("%v: unexpected check %v", fset.Position(n.Pos), ident)
 						continue
@@ -111,6 +134,8 @@ func TestMarker(t *testing.T) {
 					t.Errorf("Unexpected note %v at %v", n.Name, fset.Position(n.Pos))
 				}
 			}
+		})
+	}
 }
 
 func checkMarker(t *testing.T, fset *token.FileSet, readFile expect.ReadFile, markers map[string]token.Pos, pos token.Pos, name string, pattern interface{}) {

go/expect/extract.go

@@ -9,15 +9,17 @@ import (
 	"go/ast"
 	"go/parser"
 	"go/token"
+	"path/filepath"
 	"regexp"
 	"strconv"
 	"strings"
 	"text/scanner"
+
+	"golang.org/x/mod/modfile"
 )
 
-const (
+const commentStart = "@"
-	commentStart = "@"
+const commentStartLen = len(commentStart)
-)
 
 // Identifier is the type for an identifier in an Note argument list.
 type Identifier string
@@ -34,6 +36,8 @@ func Parse(fset *token.FileSet, filename string, content []byte) ([]*Note, error
 	if content != nil {
 		src = content
 	}
+	switch filepath.Ext(filename) {
+	case ".go":
 		// TODO: We should write this in terms of the scanner.
 		// there are ways you can break the parser such that it will not add all the
 		// comments to the ast, which may result in files where the tests are silently
@@ -42,19 +46,85 @@ func Parse(fset *token.FileSet, filename string, content []byte) ([]*Note, error
 		if file == nil {
 			return nil, err
 		}
-	return Extract(fset, file)
+		return ExtractGo(fset, file)
+	case ".mod":
+		file, err := modfile.Parse(filename, content, nil)
+		if err != nil {
+			return nil, err
+		}
+		fset.AddFile(filename, -1, len(content)).SetLinesForContent(content)
+		return extractMod(fset, file)
+	}
+	return nil, nil
 }
 
-// Extract collects all the notes present in an AST.
+// extractMod collects all the notes present in a go.mod file.
 // Each comment whose text starts with @ is parsed as a comma-separated
 // sequence of notes.
 // See the package documentation for details about the syntax of those
 // notes.
-func Extract(fset *token.FileSet, file *ast.File) ([]*Note, error) {
+// Only allow notes to appear with the following format: "//@mark()" or // @mark()
+func extractMod(fset *token.FileSet, file *modfile.File) ([]*Note, error) {
+	var notes []*Note
+	for _, stmt := range file.Syntax.Stmt {
+		comment := stmt.Comment()
+		if comment == nil {
+			continue
+		}
+		// Handle the case for markers of `// indirect` to be on the line before
+		// the require statement.
+		// TODO(golang/go#36894): have a more intuitive approach for // indirect
+		for _, cmt := range comment.Before {
+			text, adjust := getAdjustedNote(cmt.Token)
+			if text == "" {
+				continue
+			}
+			parsed, err := parse(fset, token.Pos(int(cmt.Start.Byte)+adjust), text)
+			if err != nil {
+				return nil, err
+			}
+			notes = append(notes, parsed...)
+		}
+		// Handle the normal case for markers on the same line.
+		for _, cmt := range comment.Suffix {
+			text, adjust := getAdjustedNote(cmt.Token)
+			if text == "" {
+				continue
+			}
+			parsed, err := parse(fset, token.Pos(int(cmt.Start.Byte)+adjust), text)
+			if err != nil {
+				return nil, err
+			}
+			notes = append(notes, parsed...)
+		}
+	}
+	return notes, nil
+}
+
+// ExtractGo collects all the notes present in an AST.
+// Each comment whose text starts with @ is parsed as a comma-separated
+// sequence of notes.
+// See the package documentation for details about the syntax of those
+// notes.
+func ExtractGo(fset *token.FileSet, file *ast.File) ([]*Note, error) {
 	var notes []*Note
 	for _, g := range file.Comments {
 		for _, c := range g.List {
-			text := c.Text
+			text, adjust := getAdjustedNote(c.Text)
+			if text == "" {
+				continue
+			}
+			parsed, err := parse(fset, token.Pos(int(c.Pos())+adjust), text)
+			if err != nil {
+				return nil, err
+			}
+			notes = append(notes, parsed...)
+		}
+	}
+	return notes, nil
+}
+
+func getAdjustedNote(text string) (string, int) {
 	if strings.HasPrefix(text, "/*") {
 		text = strings.TrimSuffix(text, "*/")
 	}
@@ -70,23 +140,16 @@ func Extract(fset *token.FileSet, file *ast.File) ([]*Note, error) {
 		// Get the text before the commentStart.
 		pre := text[i-2 : i]
 		if pre != "//" {
-					continue
+			return "", 0
 		}
 		text = text[i:]
 		adjust = i
 	}
 	if !strings.HasPrefix(text, commentStart) {
-				continue
+		return "", 0
 	}
-			text = text[len(commentStart):]
+	text = text[commentStartLen:]
-			parsed, err := parse(fset, token.Pos(int(c.Pos())+4+adjust), text)
+	return text, commentStartLen + adjust + 1
-			if err != nil {
-				return nil, err
-			}
-			notes = append(notes, parsed...)
-		}
-	}
-	return notes, nil
 }
 
 const invalidToken rune = 0

go/expect/testdata/go.mod

@@ -0,0 +1,7 @@
+module αfake1α //@mark(αMarker, "αfake1α")
+
+go 1.14
+
+require golang.org/modfile v0.0.0 //@mark(βMarker, "require golang.org/modfile v0.0.0")
+//@mark(IndirectMarker, "// indirect")
+require golang.org/x/tools v0.0.0-20191219192050-56b0b28a00f7 // indirect

go/ssa/source_test.go

@@ -50,7 +50,7 @@ func TestObjValueLookup(t *testing.T) {
 	// Each note of the form @ssa(x, "BinOp") in testdata/objlookup.go
 	// specifies an expectation that an object named x declared on the
 	// same line is associated with an an ssa.Value of type *ssa.BinOp.
-	notes, err := expect.Extract(conf.Fset, f)
+	notes, err := expect.ExtractGo(conf.Fset, f)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -271,7 +271,7 @@ func testValueForExpr(t *testing.T, testfile string) {
 		return true
 	})
 
-	notes, err := expect.Extract(prog.Fset, f)
+	notes, err := expect.ExtractGo(prog.Fset, f)
 	if err != nil {
 		t.Fatal(err)
 	}