cue/ast: add API to modify comments

Adding more methods to nodes makes the API rather
cluttered. We deprecate the interface methods
in favor of functions.

Change-Id: Ie791855720249e198082c9a76a87a26e21e82375
Reviewed-on: https://cue-review.googlesource.com/c/cue/+/3221
Reviewed-by: Marcel van Lohuizen <mpvl@golang.org>
diff --git a/cue/ast/ast.go b/cue/ast/ast.go
index da6b72c..01f16c4 100644
--- a/cue/ast/ast.go
+++ b/cue/ast/ast.go
@@ -44,8 +44,12 @@
 
 	// TODO: SetPos(p token.RelPos)
 
+	// Deprecated: use ast.Comments
 	Comments() []*CommentGroup
+
+	// Deprecated: use ast.AddComment
 	AddComment(*CommentGroup)
+	commentInfo() *comments
 }
 
 // An Expr is implemented by all expression nodes.
@@ -117,6 +121,8 @@
 	groups *[]*CommentGroup
 }
 
+func (c *comments) commentInfo() *comments { return c }
+
 func (c *comments) Comments() []*CommentGroup {
 	if c.groups == nil {
 		return []*CommentGroup{}
@@ -139,14 +145,24 @@
 	*c.groups = append(*c.groups, cg)
 }
 
+func (c *comments) SetComments(cgs []*CommentGroup) {
+	if c.groups == nil {
+		a := cgs
+		c.groups = &a
+		return
+	}
+	*c.groups = cgs
+}
+
 // A Comment node represents a single //-style or /*-style comment.
 type Comment struct {
 	Slash token.Pos // position of "/" starting the comment
 	Text  string    // comment text (excluding '\n' for //-style comments)
 }
 
-func (g *Comment) Comments() []*CommentGroup { return nil }
-func (g *Comment) AddComment(*CommentGroup)  {}
+func (c *Comment) Comments() []*CommentGroup { return nil }
+func (c *Comment) AddComment(*CommentGroup)  {}
+func (c *Comment) commentInfo() *comments    { return nil }
 
 func (c *Comment) Pos() token.Pos { return c.Slash }
 func (c *Comment) End() token.Pos { return c.Slash.Add(len(c.Text)) }
@@ -171,6 +187,7 @@
 
 func (g *CommentGroup) Comments() []*CommentGroup { return nil }
 func (g *CommentGroup) AddComment(*CommentGroup)  {}
+func (g *CommentGroup) commentInfo() *comments    { return nil }
 
 func isWhitespace(ch byte) bool { return ch == ' ' || ch == '\t' || ch == '\n' || ch == '\r' }
 
diff --git a/cue/ast/astutil/apply.go b/cue/ast/astutil/apply.go
index 4e747f5..22deb16 100644
--- a/cue/ast/astutil/apply.go
+++ b/cue/ast/astutil/apply.go
@@ -42,7 +42,9 @@
 	Index() int
 
 	// Replace replaces the current Node with n.
-	// The replacement node is not walked by Apply.
+	// The replacement node is not walked by Apply. Comments of the old node
+	// are copied to the new node if it has not yet an comments associated
+	// with it.
 	Replace(n ast.Node)
 
 	// Delete deletes the current Node from its containing struct.
@@ -84,6 +86,9 @@
 func (c *cursor) Replace(n ast.Node) {
 	// panic if the value cannot convert to the original type.
 	reflect.ValueOf(n).Convert(reflect.TypeOf(c.typ).Elem())
+	if ast.Comments(n) != nil {
+		CopyComments(n, c.node)
+	}
 	c.node = n
 	c.replaced = true
 }
diff --git a/cue/ast/astutil/apply_test.go b/cue/ast/astutil/apply_test.go
index 4a2eebf..7ed8f08 100644
--- a/cue/ast/astutil/apply_test.go
+++ b/cue/ast/astutil/apply_test.go
@@ -116,13 +116,15 @@
 	}, {
 		name: "replace",
 		in: `
-		a: "string"
+		// keep comment
+		a: "string" // and this one
 		b: 3
 		c: [ 1, 2, 8, 4 ]
 		d: "\(foo) is \(0)"
 		`,
 		out: `
-a: s
+// keep comment
+a: s // and this one
 b: 4
 c: [4, 4, 4, 4]
 d: "\(foo) is \(4)"
diff --git a/cue/ast/astutil/util.go b/cue/ast/astutil/util.go
new file mode 100644
index 0000000..2a237b3
--- /dev/null
+++ b/cue/ast/astutil/util.go
@@ -0,0 +1,26 @@
+// Copyright 2019 CUE Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package astutil
+
+import "cuelang.org/go/cue/ast"
+
+// CopyComments associates comments of one node with another.
+// It may change the relative position of comments.
+func CopyComments(to, from ast.Node) {
+	if from == nil {
+		return
+	}
+	ast.SetComments(to, from.Comments())
+}
diff --git a/cue/ast/comments.go b/cue/ast/comments.go
new file mode 100644
index 0000000..09d5402
--- /dev/null
+++ b/cue/ast/comments.go
@@ -0,0 +1,46 @@
+// Copyright 2019 CUE Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package ast
+
+// Comments returns all comments associated with a given node.
+func Comments(n Node) []*CommentGroup {
+	c := n.commentInfo()
+	if c == nil {
+		return nil
+	}
+	return c.Comments()
+}
+
+// AddComment adds the given comment to the node if it supports it.
+// If a node does not support comments, such as for CommentGroup or Comment,
+// this call has no effect.
+func AddComment(n Node, cg *CommentGroup) {
+	c := n.commentInfo()
+	if c == nil {
+		return
+	}
+	c.AddComment(cg)
+}
+
+// SetComments replaces all comments of n with the given set of comments.
+// If a node does not support comments, such as for CommentGroup or Comment,
+// this call has no effect.
+func SetComments(n Node, cgs []*CommentGroup) {
+	c := n.commentInfo()
+	if c == nil {
+		return
+	}
+	c.SetComments(cgs)
+}
diff --git a/cue/ast/walk.go b/cue/ast/walk.go
index 1dfd1a0..85a45fc 100644
--- a/cue/ast/walk.go
+++ b/cue/ast/walk.go
@@ -64,7 +64,7 @@
 
 	// TODO: record the comment groups and interleave with the values like for
 	// parsing and printing?
-	for _, c := range node.Comments() {
+	for _, c := range Comments(node) {
 		walk(v, c)
 	}
 
@@ -221,7 +221,7 @@
 func (f *inspector) Before(node Node) visitor {
 	if f.before == nil || f.before(node) {
 		f.commentStack = append(f.commentStack, f.current)
-		f.current = commentFrame{cg: node.Comments()}
+		f.current = commentFrame{cg: Comments(node)}
 		f.visitComments(f.current.pos)
 		return f
 	}