internal/lsp: make format work on the ast not the source

This makes the format code use the AST that is already cached on the file to do
the formatting. It also moves the core format code into the source directory.

Change-Id: Iaa79169708e92525cce326ea094ab98144fe1011
Reviewed-on: https://go-review.googlesource.com/c/148198
Run-TryBot: Ian Cottrell <iancottrell@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
diff --git a/internal/lsp/format.go b/internal/lsp/format.go
deleted file mode 100644
index 8066a5d..0000000
--- a/internal/lsp/format.go
+++ /dev/null
@@ -1,88 +0,0 @@
-// Copyright 2018 The Go Authors. All rights reserved.
-// Use of this source code is governed by a BSD-style
-// license that can be found in the LICENSE file.
-
-package lsp
-
-import (
-	"bytes"
-	"fmt"
-	"go/format"
-
-	"golang.org/x/tools/internal/lsp/protocol"
-	"golang.org/x/tools/internal/lsp/source"
-)
-
-// formatRange formats a document with a given range.
-func formatRange(v *source.View, uri protocol.DocumentURI, rng *protocol.Range) ([]protocol.TextEdit, error) {
-	data, err := v.GetFile(source.URI(uri)).Read()
-	if err != nil {
-		return nil, err
-	}
-	if rng != nil {
-		start, err := positionToOffset(data, int(rng.Start.Line), int(rng.Start.Character))
-		if err != nil {
-			return nil, err
-		}
-		end, err := positionToOffset(data, int(rng.End.Line), int(rng.End.Character))
-		if err != nil {
-			return nil, err
-		}
-		data = data[start:end]
-		// format.Source will fail if the substring is not a balanced expression tree.
-		// TODO(rstambler): parse the file and use astutil.PathEnclosingInterval to
-		// find the largest ast.Node n contained within start:end, and format the
-		// region n.Pos-n.End instead.
-	}
-	// format.Source changes slightly from one release to another, so the version
-	// of Go used to build the LSP server will determine how it formats code.
-	// This should be acceptable for all users, who likely be prompted to rebuild
-	// the LSP server on each Go release.
-	fmted, err := format.Source([]byte(data))
-	if err != nil {
-		return nil, err
-	}
-	if rng == nil {
-		// Get the ending line and column numbers for the original file.
-		line := bytes.Count(data, []byte("\n"))
-		col := len(data) - bytes.LastIndex(data, []byte("\n")) - 1
-		if col < 0 {
-			col = 0
-		}
-		rng = &protocol.Range{
-			Start: protocol.Position{
-				Line:      0,
-				Character: 0,
-			},
-			End: protocol.Position{
-				Line:      float64(line),
-				Character: float64(col),
-			},
-		}
-	}
-	// TODO(rstambler): Compute text edits instead of replacing whole file.
-	return []protocol.TextEdit{
-		{
-			Range:   *rng,
-			NewText: string(fmted),
-		},
-	}, nil
-}
-
-// positionToOffset converts a 0-based line and column number in a file
-// to a byte offset value.
-func positionToOffset(contents []byte, line, col int) (int, error) {
-	start := 0
-	for i := 0; i < int(line); i++ {
-		if start >= len(contents) {
-			return 0, fmt.Errorf("file contains %v lines, not %v lines", i, line)
-		}
-		index := bytes.IndexByte(contents[start:], '\n')
-		if index == -1 {
-			return 0, fmt.Errorf("file contains %v lines, not %v lines", i, line)
-		}
-		start += index + 1
-	}
-	offset := start + int(col)
-	return offset, nil
-}
diff --git a/internal/lsp/position.go b/internal/lsp/position.go
index 2d0146a..de4ce67 100644
--- a/internal/lsp/position.go
+++ b/internal/lsp/position.go
@@ -28,11 +28,8 @@
 	tokFile := v.Config.Fset.File(r.Start)
 	file := v.GetFile(source.ToURI(tokFile.Name()))
 	return protocol.Location{
-		URI: protocol.DocumentURI(file.URI),
-		Range: protocol.Range{
-			Start: toProtocolPosition(tokFile, r.Start),
-			End:   toProtocolPosition(tokFile, r.End),
-		},
+		URI:   protocol.DocumentURI(file.URI),
+		Range: toProtocolRange(tokFile, r),
 	}
 }
 
@@ -56,6 +53,14 @@
 	}
 }
 
+// toProtocolRange converts from a source range back to a protocol range.
+func toProtocolRange(f *token.File, r source.Range) protocol.Range {
+	return protocol.Range{
+		Start: toProtocolPosition(f, r.Start),
+		End:   toProtocolPosition(f, r.End),
+	}
+}
+
 // fromProtocolPosition converts a protocol position (0-based line and column
 // number) to a token.Pos (byte offset value).
 // It requires the token file the pos belongs to in order to do this.
diff --git a/internal/lsp/server.go b/internal/lsp/server.go
index 0818195..f7cf697 100644
--- a/internal/lsp/server.go
+++ b/internal/lsp/server.go
@@ -6,6 +6,7 @@
 
 import (
 	"context"
+	"go/token"
 	"os"
 	"sync"
 
@@ -240,11 +241,46 @@
 }
 
 func (s *server) Formatting(ctx context.Context, params *protocol.DocumentFormattingParams) ([]protocol.TextEdit, error) {
-	return formatRange(s.view, params.TextDocument.URI, nil)
+	return formatRange(ctx, s.view, params.TextDocument.URI, nil)
 }
 
 func (s *server) RangeFormatting(ctx context.Context, params *protocol.DocumentRangeFormattingParams) ([]protocol.TextEdit, error) {
-	return formatRange(s.view, params.TextDocument.URI, &params.Range)
+	return formatRange(ctx, s.view, params.TextDocument.URI, &params.Range)
+}
+
+// formatRange formats a document with a given range.
+func formatRange(ctx context.Context, v *source.View, uri protocol.DocumentURI, rng *protocol.Range) ([]protocol.TextEdit, error) {
+	f := v.GetFile(source.URI(uri))
+	tok, err := f.GetToken()
+	if err != nil {
+		return nil, err
+	}
+	var r source.Range
+	if rng == nil {
+		r.Start = tok.Pos(0)
+		r.End = tok.Pos(tok.Size())
+	} else {
+		r = fromProtocolRange(tok, *rng)
+	}
+	edits, err := source.Format(ctx, f, r)
+	if err != nil {
+		return nil, err
+	}
+	return toProtocolEdits(tok, edits), nil
+}
+
+func toProtocolEdits(f *token.File, edits []source.TextEdit) []protocol.TextEdit {
+	if edits == nil {
+		return nil
+	}
+	result := make([]protocol.TextEdit, len(edits))
+	for i, edit := range edits {
+		result[i] = protocol.TextEdit{
+			Range:   toProtocolRange(f, edit.Range),
+			NewText: edit.NewText,
+		}
+	}
+	return result
 }
 
 func (s *server) OnTypeFormatting(context.Context, *protocol.DocumentOnTypeFormattingParams) ([]protocol.TextEdit, error) {
diff --git a/internal/lsp/source/file.go b/internal/lsp/source/file.go
index 8b9ddb2..dc7bd85 100644
--- a/internal/lsp/source/file.go
+++ b/internal/lsp/source/file.go
@@ -33,6 +33,13 @@
 	End   token.Pos
 }
 
+// TextEdit represents a change to a section of a document.
+// The text within the specified range should be replaced by the supplied new text.
+type TextEdit struct {
+	Range   Range
+	NewText string
+}
+
 // SetContent sets the overlay contents for a file.
 // Setting it to nil will revert it to the on disk contents, and remove it
 // from the active set.
diff --git a/internal/lsp/source/format.go b/internal/lsp/source/format.go
new file mode 100644
index 0000000..5e7c1e7
--- /dev/null
+++ b/internal/lsp/source/format.go
@@ -0,0 +1,39 @@
+// Copyright 2018 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package source
+
+import (
+	"bytes"
+	"context"
+	"go/format"
+)
+
+// Format formats a document with a given range.
+func Format(ctx context.Context, f *File, rng Range) ([]TextEdit, error) {
+	fAST, err := f.GetAST()
+	if err != nil {
+		return nil, err
+	}
+
+	// TODO(rstambler): use astutil.PathEnclosingInterval to
+	// find the largest ast.Node n contained within start:end, and format the
+	// region n.Pos-n.End instead.
+
+	// format.Node changes slightly from one release to another, so the version
+	// of Go used to build the LSP server will determine how it formats code.
+	// This should be acceptable for all users, who likely be prompted to rebuild
+	// the LSP server on each Go release.
+	buf := &bytes.Buffer{}
+	if err := format.Node(buf, f.view.Config.Fset, fAST); err != nil {
+		return nil, err
+	}
+	// TODO(rstambler): Compute text edits instead of replacing whole file.
+	return []TextEdit{
+		{
+			Range:   rng,
+			NewText: buf.String(),
+		},
+	}, nil
+}