From 629539558bd2a1e6faea4399db84a2b10fca8253 Mon Sep 17 00:00:00 2001
From: Fatih Arslan <ftharsln@gmail.com>
Date: Sat, 31 Oct 2015 01:15:36 +0300
Subject: [PATCH] printer: collect standalone comments for printing

---
 ast/walk.go        |  2 ++
 parser/parser.go   |  1 -
 printer/nodes.go   | 59 ++++++++++++++++++++++++++++++++++++++++++++++
 printer/printer.go | 10 ++++----
 4 files changed, 66 insertions(+), 6 deletions(-)

diff --git a/ast/walk.go b/ast/walk.go
index feae611..c6dc75a 100644
--- a/ast/walk.go
+++ b/ast/walk.go
@@ -11,6 +11,8 @@ func Walk(node Node, fn func(Node) bool) {
 	}
 
 	switch n := node.(type) {
+	case *File:
+		Walk(n.Node, fn)
 	case *ObjectList:
 		for _, item := range n.Items {
 			Walk(item, fn)
diff --git a/parser/parser.go b/parser/parser.go
index 8623a83..9b5e2c9 100644
--- a/parser/parser.go
+++ b/parser/parser.go
@@ -17,7 +17,6 @@ type Parser struct {
 	// Last read token
 	tok token.Token
 
-	// comments
 	comments    []*ast.CommentGroup
 	leadComment *ast.CommentGroup // last lead comment
 	lineComment *ast.CommentGroup // last line comment
diff --git a/printer/nodes.go b/printer/nodes.go
index ec36302..f4d9c3a 100644
--- a/printer/nodes.go
+++ b/printer/nodes.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 
 	"github.com/fatih/hcl/ast"
+	"github.com/fatih/hcl/token"
 )
 
 const (
@@ -13,6 +14,64 @@ const (
 	tab     = byte('\t')
 )
 
+type printer struct {
+	cfg                Config
+	comments           []*ast.CommentGroup // may be nil, contains all comments
+	standaloneComments []*ast.CommentGroup // contains all standalone comments (not assigned to any node)
+}
+
+func (p *printer) collectComments(node ast.Node) {
+	leadComments := make([]*ast.CommentGroup, 0)
+	lineComments := make([]*ast.CommentGroup, 0)
+
+	ast.Walk(node, func(nn ast.Node) bool {
+		switch t := nn.(type) {
+		case *ast.File:
+			// will happen only once
+			p.comments = t.Comments
+		case *ast.ObjectItem:
+			if t.LeadComment != nil {
+				leadComments = append(leadComments, t.LeadComment)
+			}
+
+			if t.LineComment != nil {
+				lineComments = append(lineComments, t.LineComment)
+			}
+		}
+
+		return true
+	})
+
+	standaloneComments := make(map[token.Pos]*ast.CommentGroup, 0)
+	for _, c := range p.comments {
+		standaloneComments[c.Pos()] = c
+	}
+	for _, lead := range leadComments {
+		for _, comment := range lead.List {
+			if _, ok := standaloneComments[comment.Pos()]; ok {
+				delete(standaloneComments, comment.Pos())
+			}
+		}
+	}
+
+	for _, line := range lineComments {
+		for _, comment := range line.List {
+			if _, ok := standaloneComments[comment.Pos()]; ok {
+				delete(standaloneComments, comment.Pos())
+			}
+		}
+	}
+
+	for _, c := range standaloneComments {
+		p.standaloneComments = append(p.standaloneComments, c)
+	}
+
+	fmt.Printf("All comments len = %+v\n", len(p.comments))
+	fmt.Printf("Lead commetns = %+v\n", len(leadComments))
+	fmt.Printf("len(lineComments) = %+v\n", len(lineComments))
+	fmt.Printf("StandAlone Comments = %+v\n", len(p.standaloneComments))
+}
+
 // output prints creates a printable HCL output and returns it.
 func (p *printer) output(n interface{}) []byte {
 	var buf bytes.Buffer
diff --git a/printer/printer.go b/printer/printer.go
index a9a0d78..1d15fc8 100644
--- a/printer/printer.go
+++ b/printer/printer.go
@@ -12,10 +12,6 @@ import (
 
 var DefaultConfig = Config{}
 
-type printer struct {
-	cfg Config
-}
-
 // A Config node controls the output of Fprint.
 type Config struct {
 	SpacesWidth int // if set, it will use spaces instead of tabs for alignment
@@ -23,9 +19,13 @@ type Config struct {
 
 func (c *Config) Fprint(output io.Writer, node ast.Node) error {
 	p := &printer{
-		cfg: *c,
+		cfg:                *c,
+		comments:           make([]*ast.CommentGroup, 0),
+		standaloneComments: make([]*ast.CommentGroup, 0),
 	}
 
+	p.collectComments(node)
+
 	if _, err := output.Write(p.output(node)); err != nil {
 		return err
 	}