code review 142360043: go/format, cmd/gofmt: fix issues with partial Go code w...

src/cmd/gofmt/gofmt.go

Index: src/cmd/gofmt/gofmt.go
===================================================================
--- a/src/cmd/gofmt/gofmt.go
+++ b/src/cmd/gofmt/gofmt.go
@@ -87,13 +87,13 @@
 		return err
 	}
 
-	file, adjust, err := parse(fileSet, filename, src, stdin)
+	file, sourceAdj, indentAdj, err := parse(fileSet, filename, src, stdin)
 	if err != nil {
 		return err
 	}
 
 	if rewrite != nil {
-		if adjust == nil {
+		if sourceAdj == nil {
 			file = rewrite(file)
 		} else {
 			fmt.Fprintf(os.Stderr, "warning: rewrite ignored for incomplete programs\n")
@@ -106,15 +106,10 @@
 		simplify(file)
 	}
 
-	var buf bytes.Buffer
-	err = (&printer.Config{Mode: printerMode, Tabwidth: tabWidth}).Fprint(&buf, fileSet, file)
+	res, err := format(fileSet, file, sourceAdj, indentAdj, src)
 	if err != nil {
 		return err
 	}
-	res := buf.Bytes()
-	if adjust != nil {
-		res = adjust(src, res)
-	}
 
 	if !bytes.Equal(src, res) {
 		// formatting has changed
@@ -242,17 +237,19 @@
 
 // parse parses src, which was read from filename,
 // as a Go source file or statement list.
-func parse(fset *token.FileSet, filename string, src []byte, stdin bool) (*ast.File, func(orig, src []byte) []byte, error) {
+func parse(fset *token.FileSet, filename string, src []byte, fragmentOk bool) (
+	file *ast.File,
+	sourceAdj func(src []byte, indent int) []byte,
+	indentAdj int,
+	err error,
+) {
 	// Try as whole source file.
-	file, err := parser.ParseFile(fset, filename, src, parserMode)
-	if err == nil {
-		return file, nil, nil
-	}
-	// If the error is that the source file didn't begin with a
-	// package line and this is standard input, fall through to
+	file, err = parser.ParseFile(fset, filename, src, parserMode)
+	// If there's no error, return.  If the error is that the source file didn't begin with a
+	// package line and source fragments are ok, fall through to
 	// try as a source fragment.  Stop and return on any other error.
-	if !stdin || !strings.Contains(err.Error(), "expected 'package'") {
-		return nil, nil, err
+	if err == nil || !fragmentOk || !strings.Contains(err.Error(), "expected 'package'") {
+		return
 	}
 
 	// If this is a declaration list, make it a source file
@@ -262,19 +259,19 @@
 	psrc := append([]byte("package p;"), src...)
 	file, err = parser.ParseFile(fset, filename, psrc, parserMode)
 	if err == nil {
-		adjust := func(orig, src []byte) []byte {
+		sourceAdj = func(src []byte, indent int) []byte {
 			// Remove the package clause.
 			// Gofmt has turned the ; into a \n.
-			src = src[len("package p\n"):]
-			return matchSpace(orig, src)
+			src = src[indent+len("package p\n"):]
+			return bytes.TrimSpace(src)
 		}
-		return file, adjust, nil
+		return
 	}
 	// If the error is that the source file didn't begin with a
 	// declaration, fall through to try as a statement list.
 	// Stop and return on any other error.
 	if !strings.Contains(err.Error(), "expected declaration") {
-		return nil, nil, err
+		return
 	}
 
 	// If this is a statement list, make it a source file
@@ -285,65 +282,89 @@
 	fsrc := append(append([]byte("package p; func _() {"), src...), '\n', '}')
 	file, err = parser.ParseFile(fset, filename, fsrc, parserMode)
 	if err == nil {
-		adjust := func(orig, src []byte) []byte {
+		sourceAdj = func(src []byte, indent int) []byte {
+			// Cap adjusted indent to zero.
+			if indent < 0 {
+				indent = 0
+			}
 			// Remove the wrapping.
 			// Gofmt has turned the ; into a \n\n.
-			src = src[len("package p\n\nfunc _() {"):]
-			src = src[:len(src)-len("\n}\n")]
-			// Gofmt has also indented the function body one level.
-			// Remove that indent.
-			src = bytes.Replace(src, []byte("\n\t"), []byte("\n"), -1)
-			return matchSpace(orig, src)
+			// There will be two non-blank lines with indent, hence 2*indent.
+			src = src[2*indent+len("package p\n\nfunc _() {"):]
+			src = src[:len(src)-(indent+len("\n}\n"))]
+			return bytes.TrimSpace(src)
 		}
-		return file, adjust, nil
+		// Gofmt has also indented the function body one level.
+		// Adjust that with indentAdj.
+		indentAdj = -1
 	}
 
-	// Failed, and out of options.
-	return nil, nil, err
+	// Succeeded, or out of options.
+	return
 }
 
-func cutSpace(b []byte) (before, middle, after []byte) {
-	i := 0
-	for i < len(b) && (b[i] == ' ' || b[i] == '\t' || b[i] == '\n') {
-		i++
+func format(fset *token.FileSet, file *ast.File, sourceAdj func(src []byte, indent int) []byte, indentAdj int, src []byte) ([]byte, error) {
+	if sourceAdj == nil {
+		// Complete source file.
+		var buf bytes.Buffer
+		err := (&printer.Config{Mode: printerMode, Tabwidth: tabWidth}).Fprint(&buf, fset, file)
+		if err != nil {
+			return nil, err
+		}
+		return buf.Bytes(), nil
 	}
-	j := len(b)
-	for j > 0 && (b[j-1] == ' ' || b[j-1] == '\t' || b[j-1] == '\n') {
-		j--
+
+	// Partial source file.
+	// Determine and prepend leading space.
+	i, j := 0, 0
+	for j < len(src) && isSpace(src[j]) {
+		if src[j] == '\n' {
+			i = j + 1 // byte offset of last line in leading space
+		}
+		j++
 	}
-	if i <= j {
-		return b[:i], b[i:j], b[j:]
+	var res []byte
+	res = append(res, src[:i]...)
+
+	// Determine and prepend indentation of first code line.
+	// Spaces are ignored unless there are no tabs,
+	// in which case spaces count as one tab.
+	indent := 0
+	hasSpace := false
+	for _, b := range src[i:j] {
+		switch b {
+		case ' ':
+			hasSpace = true
+		case '\t':
+			indent++
+		}
 	}
-	return nil, nil, b[j:]
+	if indent == 0 && hasSpace {
+		indent = 1
+	}
+	for i := 0; i < indent; i++ {
+		res = append(res, '\t')
+	}
+
+	// Format the source.
+	// Write it without any leading and trailing space.
+	cfg := &printer.Config{Mode: printerMode, Tabwidth: tabWidth}
+	cfg.Indent = indent + indentAdj
+	var buf bytes.Buffer
+	err := cfg.Fprint(&buf, fset, file)
+	if err != nil {
+		return nil, err
+	}
+	res = append(res, sourceAdj(buf.Bytes(), cfg.Indent)...)
+
+	// Determine and append trailing space.
+	i = len(src)
+	for i > 0 && isSpace(src[i-1]) {
+		i--
+	}
+	return append(res, src[i:]...), nil
 }
 
-// matchSpace reformats src to use the same space context as orig.
-// 1) If orig begins with blank lines, matchSpace inserts them at the beginning of src.
-// 2) matchSpace copies the indentation of the first non-blank line in orig
-//    to every non-blank line in src.
-// 3) matchSpace copies the trailing space from orig and uses it in place
-//   of src's trailing space.
-func matchSpace(orig []byte, src []byte) []byte {
-	before, _, after := cutSpace(orig)
-	i := bytes.LastIndex(before, []byte{'\n'})
-	before, indent := before[:i+1], before[i+1:]
-
-	_, src, _ = cutSpace(src)
-
-	var b bytes.Buffer
-	b.Write(before)
-	for len(src) > 0 {
-		line := src
-		if i := bytes.IndexByte(line, '\n'); i >= 0 {
-			line, src = line[:i+1], line[i+1:]
-		} else {
-			src = nil
-		}
-		if len(line) > 0 && line[0] != '\n' { // not blank
-			b.Write(indent)
-		}
-		b.Write(line)
-	}
-	b.Write(after)
-	return b.Bytes()
+func isSpace(b byte) bool {
+	return b == ' ' || b == '\t' || b == '\n' || b == '\r'
 }