Bug fixes, example, readme

README.md

@@ -5,6 +5,51 @@
 The `dst` package attempts to provide a work-arround for [go/ast: Free-floating comments are 
 single-biggest issue when manipulating the AST](https://github.com/golang/go/issues/20744).
 
+### Example:
+
+```go
+func Example_Decorations() {
+	code := `package main
+
+	func main() {
+		var a int
+		a++
+		print(a)
+	}`
+	fset := token.NewFileSet()
+	f, err := parser.ParseFile(fset, "a.go", code, parser.ParseComments)
+	if err != nil {
+		panic(err)
+	}
+	df := decorator.Decorate(f, fset)
+	df = dstutil.Apply(df, func(c *dstutil.Cursor) bool {
+		switch n := c.Node().(type) {
+		case *dst.DeclStmt:
+			n.Decs.End.Replace("// foo")
+		case *dst.IncDecStmt:
+			n.Decs.AfterX.Add("/* bar */")
+		case *dst.CallExpr:
+			n.Decs.AfterLparen.Add("\n")
+			n.Decs.AfterArgs.Add("\n")
+		}
+		return true
+	}, nil).(*dst.File)
+	f, fset = decorator.Restore(df)
+	format.Node(os.Stdout, fset, f)
+
+	//Output:
+	//package main
+	//
+	//func main() {
+	//	var a int // foo
+	//	a /* bar */ ++
+	//	print(
+	//		a,
+	//	)
+	//}
+}
+```
+
 ### Progress as of 2nd October
 
 I've just finished a massive reorganisation of the code generation package. Instead of scanning the 

decorations.go

@@ -2,8 +2,12 @@ package dst
 
 type Decorations []string
 
-func (d *Decorations) Add(text string) {
-	*d = append(*d, text)
+func (d *Decorations) Add(decs ...string) {
+	*d = append(*d, decs...)
+}
+
+func (d *Decorations) Replace(decs ...string) {
+	*d = decs
 }
 
 func (d *Decorations) Clear() {

decorations_test.go

@@ -0,0 +1,54 @@
+package dst_test
+
+import (
+	"go/parser"
+	"go/token"
+
+	"go/format"
+	"os"
+
+	"github.com/dave/dst"
+	"github.com/dave/dst/decorator"
+	"github.com/dave/dst/dstutil"
+)
+
+func Example_Decorations() {
+	code := `package main
+
+	func main() {
+		var a int
+		a++
+		print(a)
+	}`
+	fset := token.NewFileSet()
+	f, err := parser.ParseFile(fset, "a.go", code, parser.ParseComments)
+	if err != nil {
+		panic(err)
+	}
+	df := decorator.Decorate(f, fset)
+	df = dstutil.Apply(df, func(c *dstutil.Cursor) bool {
+		switch n := c.Node().(type) {
+		case *dst.DeclStmt:
+			n.Decs.End.Replace("// foo")
+		case *dst.IncDecStmt:
+			n.Decs.AfterX.Add("/* bar */")
+		case *dst.CallExpr:
+			n.Decs.AfterLparen.Add("\n")
+			n.Decs.AfterArgs.Add("\n")
+		}
+		return true
+	}, nil).(*dst.File)
+	f, fset = decorator.Restore(df)
+	format.Node(os.Stdout, fset, f)
+
+	//Output:
+	//package main
+	//
+	//func main() {
+	//	var a int // foo
+	//	a /* bar */ ++
+	//	print(
+	//		a,
+	//	)
+	//}
+}

decorator/decorator.go

@@ -27,7 +27,7 @@ func (d *Decorator) Decorate(f *ast.File, fset *token.FileSet) *dst.File {
 	fragger := &Fragger{}
 	fragger.Fragment(f, fset)
 
-	//p.debug(os.Stdout, fset)
+	//fragger.debug(os.Stdout, fset)
 
 	d.decorations = fragger.Link()
 	return d.DecorateNode(f).(*dst.File)

decorator/decorator_test.go

@@ -20,6 +20,17 @@ func TestDecorator(t *testing.T) {
 		code       string
 		expect     string
 	}{
+		{
+			name: "value spec",
+			code: `package main
+
+				func main() {
+					var foo int
+				}`,
+			expect: `File [AfterName "\n" "\n"]
+            BlockStmt [AfterLbrace "\n"]
+            DeclStmt [End "\n"]`,
+		},
 		{
 			name: "chan type",
 			code: `package main
@@ -247,6 +258,10 @@ func TestDecorator(t *testing.T) {
 			if normalize(string(b)) != normalize(test.code) {
 				t.Fatalf("code changed after gofmt. before: \n%s\nafter:\n%s", test.code, string(b))
 			}
+
+			// use the formatted version (correct indents etc.)
+			test.code = string(b)
+
 			fset := token.NewFileSet()
 			f, err := parser.ParseFile(fset, "main.go", test.code, parser.ParseComments)
 			if err != nil {

decorator/fragger-generated.go

@@ -1426,10 +1426,14 @@ func (f *Fragger) ProcessNode(n ast.Node) {
 		}
 
 		// Token: Assign
-		f.AddToken(n, token.ASSIGN, token.NoPos)
+		if n.Values != nil {
+			f.AddToken(n, token.ASSIGN, token.NoPos)
+		}
 
 		// Decoration: AfterAssign
-		f.AddDecoration(n, "AfterAssign")
+		if n.Values != nil {
+			f.AddDecoration(n, "AfterAssign")
+		}
 
 		// List: Values
 		for _, v := range n.Values {

decorator/restorer-generated.go

@@ -1566,7 +1566,9 @@ func (r *FileRestorer) RestoreNode(n dst.Node) ast.Node {
 		}
 
 		// Token: Assign
-		r.cursor += token.Pos(len(token.ASSIGN.String()))
+		if n.Values != nil {
+			r.cursor += token.Pos(len(token.ASSIGN.String()))
+		}
 
 		// Decoration: AfterAssign
 		r.applyDecorations(n.Decs.AfterAssign)

decorator/restorer.go

@@ -73,8 +73,14 @@ func (r *Restorer) Restore(fname string, dstFile *dst.File) *ast.File {
 func (f *FileRestorer) applyDecorations(decorations dst.Decorations) {
 	for _, d := range decorations {
 
+		isNewline := d == "\n"
+		isLineComment := strings.HasPrefix(d, "//")
+		isInlineComment := strings.HasPrefix(d, "/*")
+		isComment := isLineComment || isInlineComment
+		isMultiLineComment := isInlineComment && strings.Contains(d, "\n")
+
 		// for multi-line comments, add a newline for each \n
-		if strings.HasPrefix(d, "/*") && strings.Contains(d, "\n") {
+		if isMultiLineComment {
 			for i, c := range d {
 				if c == '\n' {
 					f.Lines = append(f.Lines, int(f.cursor)+i)
@@ -83,13 +89,13 @@ func (f *FileRestorer) applyDecorations(decorations dst.Decorations) {
 		}
 
 		// if the decoration is a comment, add it and advance the cursor
-		if d != "\n" {
+		if isComment {
 			f.Comments = append(f.Comments, &ast.CommentGroup{List: []*ast.Comment{{Slash: f.cursor, Text: d}}})
 			f.cursor += token.Pos(len(d))
 		}
 
 		// for newline decorations and also line-comments, add a newline
-		if strings.HasPrefix(d, "//") || d == "\n" {
+		if isLineComment || isNewline {
 			f.Lines = append(f.Lines, int(f.cursor))
 			f.cursor++
 		}

gendst/fragment/fragment.go

@@ -1887,11 +1887,13 @@ var Info = map[string][]Part{
 			Type:  Type{"Expr", false},
 		},
 		Token{
-			Name:  "Assign",
-			Token: Basic{jen.Qual("go/token", "ASSIGN")},
+			Name:   "Assign",
+			Token:  Basic{jen.Qual("go/token", "ASSIGN")},
+			Exists: Expr(func(n *jen.Statement) *jen.Statement { return n.Dot("Values").Op("!=").Nil() }),
 		},
 		Decoration{
 			Name: "AfterAssign",
+			Use:  Expr(func(n *jen.Statement) *jen.Statement { return n.Dot("Values").Op("!=").Nil() }),
 		},
 		List{
 			Name:      "Values",

print.go

@@ -8,7 +8,6 @@ package dst
 
 import (
 	"fmt"
-	"go/token"
 	"io"
 	"os"
 	"reflect"
@@ -36,15 +35,14 @@ func NotNilFilter(_ string, v reflect.Value) bool {
 // struct fields for which f(fieldname, fieldvalue) is true are
 // printed; all others are filtered from the output. Unexported
 // struct fields are never printed.
-func Fprint(w io.Writer, fset *token.FileSet, x interface{}, f FieldFilter) error {
-	return fprint(w, fset, x, f)
+func Fprint(w io.Writer, x interface{}, f FieldFilter) error {
+	return fprint(w, x, f)
 }
 
-func fprint(w io.Writer, fset *token.FileSet, x interface{}, f FieldFilter) (err error) {
+func fprint(w io.Writer, x interface{}, f FieldFilter) (err error) {
 	// setup printer
 	p := printer{
 		output: w,
-		fset:   fset,
 		filter: f,
 		ptrmap: make(map[interface{}]int),
 		last:   '\n', // force printing of line number on first line
@@ -70,13 +68,12 @@ func fprint(w io.Writer, fset *token.FileSet, x interface{}, f FieldFilter) (err
 
 // Print prints x to standard output, skipping nil fields.
 // Print(fset, x) is the same as Fprint(os.Stdout, fset, x, NotNilFilter).
-func Print(fset *token.FileSet, x interface{}) error {
-	return Fprint(os.Stdout, fset, x, NotNilFilter)
+func Print(x interface{}) error {
+	return Fprint(os.Stdout, x, NotNilFilter)
 }
 
 type printer struct {
 	output io.Writer
-	fset   *token.FileSet
 	filter FieldFilter
 	ptrmap map[interface{}]int // *T -> line number
 	indent int                 // current indentation level
@@ -241,12 +238,6 @@ func (p *printer) print(x reflect.Value) {
 			// print strings in quotes
 			p.printf("%q", v)
 			return
-		case token.Pos:
-			// position values can be printed nicely if we have a file set
-			if p.fset != nil {
-				p.printf("%s", p.fset.Position(v))
-				return
-			}
 		}
 		// default
 		p.printf("%v", v)

print_test.go

@@ -87,7 +87,7 @@ func TestPrint(t *testing.T) {
 	var buf bytes.Buffer
 	for _, test := range tests {
 		buf.Reset()
-		if err := Fprint(&buf, nil, test.x, nil); err != nil {
+		if err := Fprint(&buf, test.x, nil); err != nil {
 			t.Errorf("Fprint failed: %s", err)
 		}
 		if s, ts := trim(buf.String()), trim(test.s); s != ts {