internal/lsp/source: fix crash on short file

The import formatting code tries to extend the range it's diffing to
the next line after the last import so that it's working with whole
lines. Make sure that the next line actually exists before trying to use
it.

Fixes golang/go#35604.

Change-Id: I18fe61843aa11e62ed311a9ddff62ff876888a15
Reviewed-on: https://go-review.googlesource.com/c/tools/+/208672
Run-TryBot: Heschi Kreinick <heschi@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Muir Manders <muir@mnd.rs>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
diff --git a/internal/lsp/source/format.go b/internal/lsp/source/format.go
index 51059f3..c1f078e 100644
--- a/internal/lsp/source/format.go
+++ b/internal/lsp/source/format.go
@@ -285,8 +285,13 @@
 	if firstImport == nil {
 		return nil, 0
 	}
+	tok := fset.File(f.Pos())
 	start := firstImport.Pos()
-	end := fset.File(f.Pos()).LineStart(fset.Position(lastImport.End()).Line + 1)
+	end := lastImport.End()
+	if tok.LineCount() > fset.Position(end).Line {
+		end = fset.File(f.Pos()).LineStart(fset.Position(lastImport.End()).Line + 1)
+	}
+
 	startLineOffset := fset.Position(start).Line - 1 // lines are 1-indexed.
 	return src[fset.Position(firstImport.Pos()).Offset:fset.Position(end).Offset], startLineOffset
 }