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

src/go/format/format.go

 // Copyright 2012 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 format implements standard formatting of Go source.
 package format
 
 import (
         "bytes"
         "fmt"
@@ skipping 64 matching lines @@
 // or an (I/O or syntax) error. src is expected to be a syntactically
 // correct Go source file, or a list of Go declarations or statements.
 //
 // If src is a partial source file, the leading and trailing space of src
 // is applied to the result (such that it has the same leading and trailing
 // space as src), and the result is indented by the same amount as the first
 // line of src containing code. Imports are not sorted for partial source files.
 //
 func Source(src []byte) ([]byte, error) {
         fset := token.NewFileSet()
-        file, adjust, err := parse(fset, "", src, true)
+        file, sourceAdj, indentAdj, err := parse(fset, "", src, true)
         if err != nil {
                 return nil, err
         }
 
-        var res []byte
-        if adjust == nil {
-                // Complete source file.
-                ast.SortImports(fset, file)
-                var buf bytes.Buffer
-                err := config.Fprint(&buf, fset, file)
-                if err != nil {
-                        return nil, err
-                }
-                res = buf.Bytes()
-
-        } else {
-                // Partial source file.
-                // Determine and prepend leading space.
-                i, j := 0, 0
-                for ; j < len(src) && isSpace(src[j]); j++ {
-                        if src[j] == '\n' {
-                                i = j + 1 // index of last line in leading space
-                        }
-                }
-                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++
-                        }
-                }
-                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 := config
-                cfg.Indent = indent
-                var buf bytes.Buffer
-                err := cfg.Fprint(&buf, fset, file)
-                if err != nil {
-                        return nil, err
-                }
-                res = append(res, adjust(buf.Bytes(), indent)...)
-
-                // Determine and append trailing space.
-                i = len(src)
-                for i > 0 && isSpace(src[i-1]) {
-                        i--
-                }
-                res = append(res, src[i:]...)
-        }
-
-        return res, nil
+        return format(fset, file, sourceAdj, indentAdj, src)
 }
 
 func hasUnsortedImports(file *ast.File) bool {
         for _, d := range file.Decls {
                 d, ok := d.(*ast.GenDecl)
                 if !ok || d.Tok != token.IMPORT {
                         // Not an import declaration, so we're done.
                         // Imports are always first.
                         return false
                 }
                 if d.Lparen.IsValid() {
                         // For now assume all grouped imports are unsorted.
                         // TODO(gri) Should check if they are sorted already.
                         return true
                 }
                 // Ungrouped imports are sorted by default.
         }
         return false
 }
 
 // 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(src []byte, indent int) []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
         // by inserting a package clause.
         // Insert using a ;, not a newline, so that the line numbers
         // in psrc match the ones in src.
         psrc := append([]byte("package p;"), src...)
         file, err = parser.ParseFile(fset, filename, psrc, parserMode)
         if err == nil {
-                adjust := func(src []byte, indent int) []byte {
+                sourceAdj = func(src []byte, indent int) []byte {
                         // Remove the package clause.
                         // Gofmt has turned the ; into a \n.
                         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
         // by inserting a package clause and turning the list
         // into a function body.  This handles expressions too.
         // Insert using a ;, not a newline, so that the line numbers
         // in fsrc match the ones in src.
         fsrc := append(append([]byte("package p; func _() {"), src...), '\n', '}')
         file, err = parser.ParseFile(fset, filename, fsrc, parserMode)
         if err == nil {
-                adjust := func(src []byte, indent int) []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.
+                        // 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"))]
-                        // Gofmt has also indented the function body one level.
-                        // Remove that indent.
-                        src = bytes.Replace(src, []byte("\n\t"), []byte("\n"), -1)
                         return bytes.TrimSpace(src)
                 }
-                return file, adjust, nil
-        }
-
-        // Failed, and out of options.
-        return nil, nil, err
+                // Gofmt has also indented the function body one level.
+                // Adjust that with indentAdj.
+                indentAdj = -1
+        }
+
+        // Succeeded, or out of options.
+        return
+}
+
+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.
+                ast.SortImports(fset, file)
+                var buf bytes.Buffer
+                err := config.Fprint(&buf, fset, file)
+                if err != nil {
+                        return nil, err
+                }
+                return buf.Bytes(), nil
+        }
+
+        // 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++
+        }
+        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++
+                }
+        }
+        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 := config
+        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
 }
 
 func isSpace(b byte) bool {
         return b == ' ' || b == '\t' || b == '\n' || b == '\r'
 }