Attach comments to AST nodes. (#64)

* Attach comments to AST nodes.

This feature is off by default. More testing will be needed before
exposing it to the ParseFile function.

Logic is copied from Buildifier (https://github.com/bazelbuild/buildtools/tree/master/build).

bug #63

* Fix tests by saving the state of 'blank' in the scanner.

* - Rename flattenAST
- Add new argument to the Parser and the Scanner
- Update tests

* Remove global constant keepComments

* Update more tests (new argument to the parser)

* Add CommentsRef to allow allocating comments

Address a few other issues

* Remove the COMMENT tokens

Parser won't see COMMENT tokens anymore. The list is kept by the
scanner. This simplifies the code and reverts some of my previous
changes.

* Address review comments

* - Removed the .Suffix boolean
- Renamed CommentsRef to commentsRef
- Simplified assignComments function (reversing was not useful)

* assignComments leaves early if there is no comments
+ address other review comments

* Address review comments (for -> if, removed useless code about suffix comments)

1 files changed, 68 insertions, 0 deletions

diff --git a/syntax/syntax.go b/syntax/syntax.go
index 2b1c965..113c9c4 100644
--- a/syntax/syntax.go
+++ b/syntax/syntax.go
@@ -9,6 +9,47 @@ package syntax
 type Node interface {
 	// Span returns the start and end position of the expression.
 	Span() (start, end Position)
+
+	// Comments returns the comments associated with this node.
+	// It returns nil if RetainComments was not specified during parsing,
+	// or if AllocComments was not called.
+	Comments() *Comments
+
+	// AllocComments allocates a new Comments node if there was none.
+	// This makes possible to add new comments using Comments() method.
+	AllocComments()
+}
+
+// A Comment represents a single # comment.
+type Comment struct {
+	Start Position
+	Text  string // without trailing newline
+}
+
+// Comments collects the comments associated with an expression.
+type Comments struct {
+	Before []Comment // whole-line comments before this expression
+	Suffix []Comment // end-of-line comments after this expression (up to 1)
+
+	// For top-level expressions only, After lists whole-line
+	// comments following the expression.
+	After []Comment
+}
+
+// A commentsRef is a possibly-nil reference to a set of comments.
+// A commentsRef is embedded in each type of syntax node,
+// and provides its Comments and AllocComments methods.
+type commentsRef struct{ ref *Comments }
+
+// Comments returns the comments associated with a syntax node,
+// or nil if AllocComments has not yet been called.
+func (cr commentsRef) Comments() *Comments { return cr.ref }
+
+// AllocComments enables comments to be associated with a syntax node.
+func (cr *commentsRef) AllocComments() {
+	if cr.ref == nil {
+		cr.ref = new(Comments)
+	}
 }
 
 // Start returns the start position of the expression.
@@ -25,6 +66,7 @@ func End(n Node) Position {
 
 // A File represents a Skylark file.
 type File struct {
+	commentsRef
 	Path  string
 	Stmts []Stmt
 
@@ -61,6 +103,7 @@ func (*ReturnStmt) stmt() {}
 //	x, y = y, x
 // 	x += 1
 type AssignStmt struct {
+	commentsRef
 	OpPos Position
 	Op    Token // = EQ | {PLUS,MINUS,STAR,PERCENT}_EQ
 	LHS   Expr
@@ -75,6 +118,7 @@ func (x *AssignStmt) Span() (start, end Position) {
 
 // A Function represents the common parts of LambdaExpr and DefStmt.
 type Function struct {
+	commentsRef
 	StartPos Position // position of DEF or LAMBDA token
 	Params   []Expr   // param = ident | ident=expr | *ident | **ident
 	Body     []Stmt
@@ -93,6 +137,7 @@ func (x *Function) Span() (start, end Position) {
 
 // A DefStmt represents a function definition.
 type DefStmt struct {
+	commentsRef
 	Def  Position
 	Name *Ident
 	Function
@@ -105,6 +150,7 @@ func (x *DefStmt) Span() (start, end Position) {
 
 // An ExprStmt is an expression evaluated for side effects.
 type ExprStmt struct {
+	commentsRef
 	X Expr
 }
 
@@ -115,6 +161,7 @@ func (x *ExprStmt) Span() (start, end Position) {
 // An IfStmt is a conditional: If Cond: True; else: False.
 // 'elseif' is desugared into a chain of IfStmts.
 type IfStmt struct {
+	commentsRef
 	If      Position // IF or ELIF
 	Cond    Expr
 	True    []Stmt
@@ -140,6 +187,7 @@ func (x *IfStmt) Span() (start, end Position) {
 // without.  For consistency we create fake identifiers for all the
 // strings.
 type LoadStmt struct {
+	commentsRef
 	Load   Position
 	Module *Literal // a string
 	From   []*Ident // name defined in loading module
@@ -153,6 +201,7 @@ func (x *LoadStmt) Span() (start, end Position) {
 
 // A BranchStmt changes the flow of control: break, continue, pass.
 type BranchStmt struct {
+	commentsRef
 	Token    Token // = BREAK | CONTINUE | PASS
 	TokenPos Position
 }
@@ -163,6 +212,7 @@ func (x *BranchStmt) Span() (start, end Position) {
 
 // A ReturnStmt returns from a function.
 type ReturnStmt struct {
+	commentsRef
 	Return Position
 	Result Expr // may be nil
 }
@@ -199,6 +249,7 @@ func (*UnaryExpr) expr()     {}
 
 // An Ident represents an identifier.
 type Ident struct {
+	commentsRef
 	NamePos Position
 	Name    string
 
@@ -214,6 +265,7 @@ func (x *Ident) Span() (start, end Position) {
 
 // A Literal represents a literal string or number.
 type Literal struct {
+	commentsRef
 	Token    Token // = STRING | INT
 	TokenPos Position
 	Raw      string      // uninterpreted text
@@ -226,6 +278,7 @@ func (x *Literal) Span() (start, end Position) {
 
 // A CallExpr represents a function call expression: Fn(Args).
 type CallExpr struct {
+	commentsRef
 	Fn     Expr
 	Lparen Position
 	Args   []Expr
@@ -239,6 +292,7 @@ func (x *CallExpr) Span() (start, end Position) {
 
 // A DotExpr represents a field or method selector: X.Name.
 type DotExpr struct {
+	commentsRef
 	X       Expr
 	Dot     Position
 	NamePos Position
@@ -254,6 +308,7 @@ func (x *DotExpr) Span() (start, end Position) {
 // A Comprehension represents a list or dict comprehension:
 // [Body for ... if ...] or {Body for ... if ...}
 type Comprehension struct {
+	commentsRef
 	Curly   bool // {x:y for ...} or {x for ...}, not [x for ...]
 	Lbrack  Position
 	Body    Expr
@@ -267,6 +322,7 @@ func (x *Comprehension) Span() (start, end Position) {
 
 // A ForStmt represents a loop: for Vars in X: Body.
 type ForStmt struct {
+	commentsRef
 	For  Position
 	Vars Expr // name, or tuple of names
 	X    Expr
@@ -280,6 +336,7 @@ func (x *ForStmt) Span() (start, end Position) {
 
 // A ForClause represents a for clause in a list comprehension: for Vars in X.
 type ForClause struct {
+	commentsRef
 	For  Position
 	Vars Expr // name, or tuple of names
 	In   Position
@@ -293,6 +350,7 @@ func (x *ForClause) Span() (start, end Position) {
 
 // An IfClause represents an if clause in a list comprehension: if Cond.
 type IfClause struct {
+	commentsRef
 	If   Position
 	Cond Expr
 }
@@ -304,6 +362,7 @@ func (x *IfClause) Span() (start, end Position) {
 
 // A DictExpr represents a dictionary literal: { List }.
 type DictExpr struct {
+	commentsRef
 	Lbrace Position
 	List   []Expr // all *DictEntrys
 	Rbrace Position
@@ -316,6 +375,7 @@ func (x *DictExpr) Span() (start, end Position) {
 // A DictEntry represents a dictionary entry: Key: Value.
 // Used only within a DictExpr.
 type DictEntry struct {
+	commentsRef
 	Key   Expr
 	Colon Position
 	Value Expr
@@ -333,6 +393,7 @@ func (x *DictEntry) Span() (start, end Position) {
 // currently part of the Skylark spec, so their use is controlled by the
 // resolver.AllowLambda flag.
 type LambdaExpr struct {
+	commentsRef
 	Lambda Position
 	Function
 }
@@ -344,6 +405,7 @@ func (x *LambdaExpr) Span() (start, end Position) {
 
 // A ListExpr represents a list literal: [ List ].
 type ListExpr struct {
+	commentsRef
 	Lbrack Position
 	List   []Expr
 	Rbrack Position
@@ -355,6 +417,7 @@ func (x *ListExpr) Span() (start, end Position) {
 
 // CondExpr represents the conditional: X if COND else ELSE.
 type CondExpr struct {
+	commentsRef
 	If      Position
 	Cond    Expr
 	True    Expr
@@ -370,6 +433,7 @@ func (x *CondExpr) Span() (start, end Position) {
 
 // A TupleExpr represents a tuple literal: (List).
 type TupleExpr struct {
+	commentsRef
 	Lparen Position // optional (e.g. in x, y = 0, 1), but required if List is empty
 	List   []Expr
 	Rparen Position
@@ -385,6 +449,7 @@ func (x *TupleExpr) Span() (start, end Position) {
 
 // A UnaryExpr represents a unary expression: Op X.
 type UnaryExpr struct {
+	commentsRef
 	OpPos Position
 	Op    Token
 	X     Expr
@@ -397,6 +462,7 @@ func (x *UnaryExpr) Span() (start, end Position) {
 
 // A BinaryExpr represents a binary expression: X Op Y.
 type BinaryExpr struct {
+	commentsRef
 	X     Expr
 	OpPos Position
 	Op    Token
@@ -411,6 +477,7 @@ func (x *BinaryExpr) Span() (start, end Position) {
 
 // A SliceExpr represents a slice or substring expression: X[Lo:Hi:Step].
 type SliceExpr struct {
+	commentsRef
 	X            Expr
 	Lbrack       Position
 	Lo, Hi, Step Expr // all optional
@@ -424,6 +491,7 @@ func (x *SliceExpr) Span() (start, end Position) {
 
 // An IndexExpr represents an index expression: X[Y].
 type IndexExpr struct {
+	commentsRef
 	X      Expr
 	Lbrack Position
 	Y      Expr