internal/span: support line directives

When //line directives are in play, the ast.File's Offset function will
return offsets in the generated file. We want offsets in the authored
file, so we need to pass a Converter for the authored file, in addition
to the ast.File for the generated file. For the same reason, we have to
start (Range).Span() by translating into positions in the authored file,
then calculate offsets from that.

A lot of call sites outside of the LSP don't pass the Converter, but
they probably don't matter much. I think everything inside does because
it ends up using mappedRange.

Updates golang/go#35720.

Change-Id: I7be09b3a50720b078e862d48cfdb02208f8187ae
Reviewed-on: https://go-review.googlesource.com/c/tools/+/208501
Run-TryBot: Heschi Kreinick <heschi@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Ian Cottrell <iancottrell@google.com>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
diff --git a/internal/lsp/source/completion.go b/internal/lsp/source/completion.go
index 537264a..88aeb73 100644
--- a/internal/lsp/source/completion.go
+++ b/internal/lsp/source/completion.go
@@ -18,7 +18,6 @@
 	"golang.org/x/tools/internal/lsp/fuzzy"
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/lsp/snippet"
-	"golang.org/x/tools/internal/span"
 	"golang.org/x/tools/internal/telemetry/trace"
 	errors "golang.org/x/xerrors"
 )
@@ -254,11 +253,8 @@
 	c.surrounding = &Selection{
 		content: ident.Name,
 		cursor:  c.pos,
-		mappedRange: mappedRange{
-			// Overwrite the prefix only.
-			spanRange: span.NewRange(c.snapshot.View().Session().Cache().FileSet(), ident.Pos(), c.pos),
-			m:         c.mapper,
-		},
+		// Overwrite the prefix only.
+		mappedRange: newMappedRange(c.snapshot.View().Session().Cache().FileSet(), c.mapper, ident.Pos(), ident.End()),
 	}
 
 	if c.opts.FuzzyMatching {
@@ -273,12 +269,9 @@
 func (c *completer) getSurrounding() *Selection {
 	if c.surrounding == nil {
 		c.surrounding = &Selection{
-			content: "",
-			cursor:  c.pos,
-			mappedRange: mappedRange{
-				spanRange: span.NewRange(c.snapshot.View().Session().Cache().FileSet(), c.pos, c.pos),
-				m:         c.mapper,
-			},
+			content:     "",
+			cursor:      c.pos,
+			mappedRange: newMappedRange(c.snapshot.View().Session().Cache().FileSet(), c.mapper, c.pos, c.pos),
 		}
 	}
 	return c.surrounding
diff --git a/internal/lsp/source/completion_literal.go b/internal/lsp/source/completion_literal.go
index 004950d..9528671 100644
--- a/internal/lsp/source/completion_literal.go
+++ b/internal/lsp/source/completion_literal.go
@@ -14,7 +14,6 @@
 	"golang.org/x/tools/internal/lsp/diff"
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/lsp/snippet"
-	"golang.org/x/tools/internal/span"
 	"golang.org/x/tools/internal/telemetry/log"
 )
 
@@ -159,11 +158,7 @@
 // referenceEdit produces text edits that prepend a "&" operator to the
 // specified node.
 func referenceEdit(fset *token.FileSet, m *protocol.ColumnMapper, node ast.Node) ([]protocol.TextEdit, error) {
-	rng := span.Range{
-		FileSet: fset,
-		Start:   node.Pos(),
-		End:     node.Pos(),
-	}
+	rng := newMappedRange(fset, m, node.Pos(), node.Pos())
 	spn, err := rng.Span()
 	if err != nil {
 		return nil, err
diff --git a/internal/lsp/source/folding_range.go b/internal/lsp/source/folding_range.go
index 2fcb310..ee0d3c1 100644
--- a/internal/lsp/source/folding_range.go
+++ b/internal/lsp/source/folding_range.go
@@ -7,7 +7,6 @@
 	"sort"
 
 	"golang.org/x/tools/internal/lsp/protocol"
-	"golang.org/x/tools/internal/span"
 )
 
 type FoldingRangeInfo struct {
@@ -83,11 +82,8 @@
 		return nil
 	}
 	return &FoldingRangeInfo{
-		mappedRange: mappedRange{
-			m:         m,
-			spanRange: span.NewRange(fset, start, end),
-		},
-		Kind: kind,
+		mappedRange: newMappedRange(fset, m, start, end),
+		Kind:        kind,
 	}
 }
 
@@ -171,11 +167,8 @@
 		return nil
 	}
 	return &FoldingRangeInfo{
-		mappedRange: mappedRange{
-			m:         m,
-			spanRange: span.NewRange(fset, start, end),
-		},
-		Kind: kind,
+		mappedRange: newMappedRange(fset, m, start, end),
+		Kind:        kind,
 	}
 }
 
@@ -189,12 +182,9 @@
 			continue
 		}
 		comments = append(comments, &FoldingRangeInfo{
-			mappedRange: mappedRange{
-				m: m,
-				// Fold from the end of the first line comment to the end of the comment block.
-				spanRange: span.NewRange(fset, commentGrp.List[0].End(), commentGrp.End()),
-			},
-			Kind: protocol.Comment,
+			// Fold from the end of the first line comment to the end of the comment block.
+			mappedRange: newMappedRange(fset, m, commentGrp.List[0].End(), commentGrp.End()),
+			Kind:        protocol.Comment,
 		})
 	}
 	return comments
diff --git a/internal/lsp/source/util.go b/internal/lsp/source/util.go
index 00470f6..1cf8841 100644
--- a/internal/lsp/source/util.go
+++ b/internal/lsp/source/util.go
@@ -30,6 +30,18 @@
 	protocolRange *protocol.Range
 }
 
+func newMappedRange(fset *token.FileSet, m *protocol.ColumnMapper, start, end token.Pos) mappedRange {
+	return mappedRange{
+		spanRange: span.Range{
+			FileSet:   fset,
+			Start:     start,
+			End:       end,
+			Converter: m.Converter,
+		},
+		m: m,
+	}
+}
+
 func (s mappedRange) Range() (protocol.Range, error) {
 	if s.protocolRange == nil {
 		spn, err := s.spanRange.Span()
@@ -171,10 +183,7 @@
 	if !end.IsValid() {
 		return mappedRange{}, errors.Errorf("invalid position for %v", end)
 	}
-	return mappedRange{
-		m:         m,
-		spanRange: span.NewRange(view.Session().Cache().FileSet(), pos, end),
-	}, nil
+	return newMappedRange(view.Session().Cache().FileSet(), m, pos, end), nil
 }
 
 // Matches cgo generated comment as well as the proposed standard:
diff --git a/internal/span/token.go b/internal/span/token.go
index 01b5ed2..4028eaf 100644
--- a/internal/span/token.go
+++ b/internal/span/token.go
@@ -13,9 +13,10 @@
 // It also carries the FileSet that produced the positions, so that it is
 // self contained.
 type Range struct {
-	FileSet *token.FileSet
-	Start   token.Pos
-	End     token.Pos
+	FileSet   *token.FileSet
+	Start     token.Pos
+	End       token.Pos
+	Converter Converter
 }
 
 // TokenConverter is a Converter backed by a token file set and file.
@@ -64,33 +65,56 @@
 	if f == nil {
 		return Span{}, fmt.Errorf("file not found in FileSet")
 	}
-	s := Span{}
+	var s Span
 	var err error
-	s.v.Start.Offset, err = offset(f, r.Start)
+	var startFilename string
+	startFilename, s.v.Start.Line, s.v.Start.Column, err = position(f, r.Start)
 	if err != nil {
 		return Span{}, err
 	}
+	s.v.URI = FileURI(startFilename)
 	if r.End.IsValid() {
-		s.v.End.Offset, err = offset(f, r.End)
+		var endFilename string
+		endFilename, s.v.End.Line, s.v.End.Column, err = position(f, r.End)
 		if err != nil {
 			return Span{}, err
 		}
-	}
-	// In the presence of line directives, a single File can have sections from
-	// multiple file names.
-	filename := f.Position(r.Start).Filename
-	if r.End.IsValid() {
-		if endFilename := f.Position(r.End).Filename; filename != endFilename {
-			return Span{}, fmt.Errorf("span begins in file %q but ends in %q", filename, endFilename)
+		// In the presence of line directives, a single File can have sections from
+		// multiple file names.
+		if endFilename != startFilename {
+			return Span{}, fmt.Errorf("span begins in file %q but ends in %q", startFilename, endFilename)
 		}
 	}
-	s.v.URI = FileURI(filename)
-
 	s.v.Start.clean()
 	s.v.End.clean()
 	s.v.clean()
-	converter := NewTokenConverter(r.FileSet, f)
-	return s.WithPosition(converter)
+	if r.Converter != nil {
+		return s.WithOffset(r.Converter)
+	}
+	if startFilename != f.Name() {
+		return Span{}, fmt.Errorf("must supply Converter for file %q containing lines from %q", f.Name(), startFilename)
+	}
+	return s.WithOffset(NewTokenConverter(r.FileSet, f))
+}
+
+func position(f *token.File, pos token.Pos) (string, int, int, error) {
+	off, err := offset(f, pos)
+	if err != nil {
+		return "", 0, 0, err
+	}
+	return positionFromOffset(f, off)
+}
+
+func positionFromOffset(f *token.File, offset int) (string, int, int, error) {
+	if offset > f.Size() {
+		return "", 0, 0, fmt.Errorf("offset %v is past the end of the file %v", offset, f.Size())
+	}
+	pos := f.Pos(offset)
+	p := f.Position(pos)
+	if offset == f.Size() {
+		return p.Filename, p.Line + 1, 1, nil
+	}
+	return p.Filename, p.Line, p.Column, nil
 }
 
 // offset is a copy of the Offset function in go/token, but with the adjustment
@@ -118,22 +142,16 @@
 		return Range{}, fmt.Errorf("end offset %v is past the end of the file %v", s.End(), converter.file.Size())
 	}
 	return Range{
-		FileSet: converter.fset,
-		Start:   converter.file.Pos(s.Start().Offset()),
-		End:     converter.file.Pos(s.End().Offset()),
+		FileSet:   converter.fset,
+		Start:     converter.file.Pos(s.Start().Offset()),
+		End:       converter.file.Pos(s.End().Offset()),
+		Converter: converter,
 	}, nil
 }
 
 func (l *TokenConverter) ToPosition(offset int) (int, int, error) {
-	if offset > l.file.Size() {
-		return 0, 0, fmt.Errorf("offset %v is past the end of the file %v", offset, l.file.Size())
-	}
-	pos := l.file.Pos(offset)
-	p := l.fset.Position(pos)
-	if offset == l.file.Size() {
-		return p.Line + 1, 1, nil
-	}
-	return p.Line, p.Column, nil
+	_, line, col, err := positionFromOffset(l.file, offset)
+	return line, col, err
 }
 
 func (l *TokenConverter) ToOffset(line, col int) (int, error) {
diff --git a/internal/span/token_test.go b/internal/span/token_test.go
index c9fce77..db11df1 100644
--- a/internal/span/token_test.go
+++ b/internal/span/token_test.go
@@ -7,6 +7,7 @@
 import (
 	"fmt"
 	"go/token"
+	"path"
 	"testing"
 
 	"golang.org/x/tools/internal/span"
@@ -48,16 +49,18 @@
 	for _, test := range tokenTests {
 		f := files[test.URI()]
 		c := span.NewTokenConverter(fset, f)
-		checkToken(t, c, span.New(
-			test.URI(),
-			span.NewPoint(test.Start().Line(), test.Start().Column(), 0),
-			span.NewPoint(test.End().Line(), test.End().Column(), 0),
-		), test)
-		checkToken(t, c, span.New(
-			test.URI(),
-			span.NewPoint(0, 0, test.Start().Offset()),
-			span.NewPoint(0, 0, test.End().Offset()),
-		), test)
+		t.Run(path.Base(f.Name()), func(t *testing.T) {
+			checkToken(t, c, span.New(
+				test.URI(),
+				span.NewPoint(test.Start().Line(), test.Start().Column(), 0),
+				span.NewPoint(test.End().Line(), test.End().Column(), 0),
+			), test)
+			checkToken(t, c, span.New(
+				test.URI(),
+				span.NewPoint(0, 0, test.Start().Offset()),
+				span.NewPoint(0, 0, test.End().Offset()),
+			), test)
+		})
 	}
 }