[WIP] Retain comments in AST

ast/access_test.go

@@ -43,7 +43,7 @@ func TestMappedAccess_MarshalJSON(t *testing.T) {
 
 	t.Parallel()
 
-	e := NewNominalType(nil, NewIdentifier(nil, "E", Position{Offset: 1, Line: 2, Column: 3}), []Identifier{})
+	e := NewNominalType(nil, NewIdentifier(nil, "E", Position{Offset: 1, Line: 2, Column: 3}), []Identifier{}, EmptyComments)
 
 	access := NewMappedAccess(e, Position{Offset: 0, Line: 0, Column: 0})
 	actual, err := json.Marshal(access)
@@ -68,8 +68,8 @@ func TestEntitlementAccess_MarshalJSON(t *testing.T) {
 
 	t.Parallel()
 
-	e := NewNominalType(nil, NewIdentifier(nil, "E", Position{Offset: 0, Line: 0, Column: 0}), []Identifier{})
-	f := NewNominalType(nil, NewIdentifier(nil, "F", Position{Offset: 1, Line: 2, Column: 3}), []Identifier{})
+	e := NewNominalType(nil, NewIdentifier(nil, "E", Position{Offset: 0, Line: 0, Column: 0}), []Identifier{}, EmptyComments)
+	f := NewNominalType(nil, NewIdentifier(nil, "F", Position{Offset: 1, Line: 2, Column: 3}), []Identifier{}, EmptyComments)
 
 	t.Run("conjunction", func(t *testing.T) {
 		t.Parallel()

ast/attachment.go

@@ -34,8 +34,8 @@ type AttachmentDeclaration struct {
 	BaseType     *NominalType
 	Conformances []*NominalType
 	Members      *Members
-	DocString    string
 	Range
+	Comments
 }
 
 var _ Element = &AttachmentDeclaration{}
@@ -50,8 +50,8 @@ func NewAttachmentDeclaration(
 	baseType *NominalType,
 	conformances []*NominalType,
 	members *Members,
-	docString string,
 	declarationRange Range,
+	comments Comments,
 ) *AttachmentDeclaration {
 	common.UseMemory(memoryGauge, common.AttachmentDeclarationMemoryUsage)
 
@@ -61,8 +61,8 @@ func NewAttachmentDeclaration(
 		BaseType:     baseType,
 		Conformances: conformances,
 		Members:      members,
-		DocString:    docString,
 		Range:        declarationRange,
+		Comments:     comments,
 	}
 }
 
@@ -99,7 +99,7 @@ func (d *AttachmentDeclaration) DeclarationMembers() *Members {
 }
 
 func (d *AttachmentDeclaration) DeclarationDocString() string {
-	return d.DocString
+	return d.Comments.LeadingDocString()
 }
 
 func (*AttachmentDeclaration) Kind() common.CompositeKind {
@@ -193,9 +193,11 @@ func (d *AttachmentDeclaration) MarshalJSON() ([]byte, error) {
 	return json.Marshal(&struct {
 		Type string
 		*Alias
+		DocString string
 	}{
-		Type:  "AttachmentDeclaration",
-		Alias: (*Alias)(d),
+		Type:      "AttachmentDeclaration",
+		Alias:     (*Alias)(d),
+		DocString: d.DeclarationDocString(),
 	})
 }
 
@@ -290,6 +292,7 @@ type RemoveStatement struct {
 	Attachment *NominalType
 	Value      Expression
 	StartPos   Position `json:"-"`
+	Comments
 }
 
 var _ Element = &RemoveStatement{}
@@ -300,13 +303,15 @@ func NewRemoveStatement(
 	attachment *NominalType,
 	value Expression,
 	startPos Position,
+	comments Comments,
 ) *RemoveStatement {
 	common.UseMemory(gauge, common.RemoveStatementMemoryUsage)
 
 	return &RemoveStatement{
 		Attachment: attachment,
 		Value:      value,
 		StartPos:   startPos,
+		Comments:   comments,
 	}
 }
 

ast/attachment_test.go

@@ -46,6 +46,7 @@ func TestAttachmentDeclaration_MarshallJSON(t *testing.T) {
 				Position{Offset: 1, Line: 2, Column: 3},
 			),
 			[]Identifier{},
+			EmptyComments,
 		),
 		Conformances: []*NominalType{
 			{
@@ -56,8 +57,12 @@ func TestAttachmentDeclaration_MarshallJSON(t *testing.T) {
 				),
 			},
 		},
-		Members:   NewMembers(nil, []Declaration{}),
-		DocString: "test",
+		Members: NewMembers(nil, []Declaration{}),
+		Comments: Comments{
+			Leading: []*Comment{
+				NewComment(nil, []byte("///test")),
+			},
+		},
 		Range: Range{
 			StartPos: Position{Offset: 1, Line: 2, Column: 3},
 			EndPos:   Position{Offset: 4, Line: 5, Column: 6},
@@ -138,6 +143,11 @@ func TestAttachmentDeclaration_Doc(t *testing.T) {
 				},
 			},
 			Members: NewMembers(nil, []Declaration{}),
+			Comments: Comments{
+				Leading: []*Comment{
+					NewComment(nil, []byte("///test")),
+				},
+			},
 		}
 
 		require.Equal(
@@ -260,6 +270,7 @@ func TestAttachExpressionMarshallJSON(t *testing.T) {
 				"foo",
 				Position{Offset: 1, Line: 2, Column: 3},
 			),
+			Comments{},
 		),
 		Attachment: NewInvocationExpression(
 			nil,
@@ -270,6 +281,7 @@ func TestAttachExpressionMarshallJSON(t *testing.T) {
 					"bar",
 					Position{Offset: 1, Line: 2, Column: 3},
 				),
+				Comments{},
 			),
 			[]*TypeAnnotation{},
 			Arguments{},
@@ -291,6 +303,7 @@ func TestAttachExpressionMarshallJSON(t *testing.T) {
             "EndPos": {"Offset": 3, "Line": 2, "Column": 5},
             "Base":  {
                 "Type": "IdentifierExpression",
+                "Comments": {},
                 "Identifier": { 
                     "Identifier": "foo",
                     "StartPos": {"Offset": 1, "Line": 2, "Column": 3},
@@ -303,6 +316,7 @@ func TestAttachExpressionMarshallJSON(t *testing.T) {
                 "Type": "InvocationExpression",
                 "InvokedExpression": {
                     "Type": "IdentifierExpression",
+                    "Comments": {},
                     "Identifier": { 
                         "Identifier": "bar",
                         "StartPos": {"Offset": 1, "Line": 2, "Column": 3},
@@ -437,6 +451,7 @@ func TestRemoveStatement_MarshallJSON(t *testing.T) {
 				Position{Offset: 1, Line: 2, Column: 3},
 			),
 			[]Identifier{},
+			EmptyComments,
 		),
 		Value: NewIdentifierExpression(
 			nil,
@@ -445,6 +460,7 @@ func TestRemoveStatement_MarshallJSON(t *testing.T) {
 				"baz",
 				Position{Offset: 1, Line: 2, Column: 3},
 			),
+			Comments{},
 		),
 		StartPos: Position{Offset: 1, Line: 2, Column: 3},
 	}
@@ -461,6 +477,7 @@ func TestRemoveStatement_MarshallJSON(t *testing.T) {
             "EndPos": {"Offset": 3, "Line": 2, "Column": 5},
             "Value":  {
                 "Type": "IdentifierExpression",
+                "Comments": {},
                 "Identifier": { 
                     "Identifier": "baz",
                     "StartPos": {"Offset": 1, "Line": 2, "Column": 3},
@@ -503,6 +520,7 @@ func TestRemoveStatement_Doc(t *testing.T) {
 				Identifier{
 					Identifier: "baz",
 				},
+				Comments{},
 			),
 		}
 
@@ -563,6 +581,7 @@ func TestRemoveStatement_String(t *testing.T) {
 				Identifier{
 					Identifier: "baz",
 				},
+				Comments{},
 			),
 		}
 

ast/block.go

@@ -29,16 +29,23 @@ import (
 type Block struct {
 	Statements []Statement
 	Range
+	Comments
 }
 
 var _ Element = &Block{}
 
-func NewBlock(memoryGauge common.MemoryGauge, statements []Statement, astRange Range) *Block {
+func NewBlock(
+	memoryGauge common.MemoryGauge,
+	statements []Statement,
+	astRange Range,
+	comments Comments,
+) *Block {
 	common.UseMemory(memoryGauge, common.BlockMemoryUsage)
 
 	return &Block{
 		Statements: statements,
 		Range:      astRange,
+		Comments:   comments,
 	}
 }
 

ast/block_test.go

@@ -360,6 +360,7 @@ func TestFunctionBlock_MarshalJSON(t *testing.T) {
 											"Type": "InvocationExpression",
 											"InvokedExpression": {
 											   "Type": "IdentifierExpression",
+											   "Comments": {},
 											   "Identifier": {
 												   "Identifier": "foobar",
 												   "StartPos": {"Offset": 31, "Line": 32, "Column": 33},

ast/comments.go

@@ -0,0 +1,101 @@
+package ast
+
+import (
+	"bytes"
+	"strings"
+
+	"github.com/onflow/cadence/common"
+)
+
+type Comments struct {
+	Leading  []*Comment `json:"-"`
+	Trailing []*Comment `json:"-"`
+}
+
+var EmptyComments = Comments{}
+
+// All combines Leading and Trailing comments in a single array.
+func (c Comments) All() []*Comment {
+	var comments []*Comment
+	comments = append(comments, c.Leading...)
+	comments = append(comments, c.Trailing...)
+	return comments
+}
+
+// LeadingDocString prints the leading doc comments to string
+func (c Comments) LeadingDocString() string {
+	var s strings.Builder
+	for _, comment := range c.Leading {
+		if comment.IsDoc() {
+			if s.Len() > 0 {
+				s.WriteRune('\n')
+			}
+			s.Write(comment.Text())
+		}
+	}
+	return s.String()
+}
+
+type Comment struct {
+	source []byte
+}
+
+func NewComment(memoryGauge common.MemoryGauge, source []byte) *Comment {
+	common.UseMemory(memoryGauge, common.NewRawStringMemoryUsage(len(source)))
+	return &Comment{
+		source: source,
+	}
+}
+
+var blockCommentDocStringPrefix = []byte("/**")
+var blockCommentStringPrefix = []byte("/*")
+var lineCommentDocStringPrefix = []byte("///")
+var lineCommentStringPrefix = []byte("//")
+var blockCommentStringSuffix = []byte("*/")
+
+func (c Comment) Multiline() bool {
+	return bytes.HasPrefix(c.source, blockCommentStringPrefix)
+}
+
+func (c Comment) IsDoc() bool {
+	if c.Multiline() {
+		return bytes.HasPrefix(c.source, blockCommentDocStringPrefix)
+	} else {
+		return bytes.HasPrefix(c.source, lineCommentDocStringPrefix)
+	}
+}
+
+var commentPrefixes = [][]byte{
+	blockCommentDocStringPrefix, // must be before blockCommentStringPrefix
+	blockCommentStringPrefix,
+	lineCommentDocStringPrefix, // must be before lineCommentStringPrefix
+	lineCommentStringPrefix,
+}
+
+var commentSuffixes = [][]byte{
+	blockCommentStringSuffix,
+}
+
+// Text without opening/closing comment characters /*, /**, */, //
+func (c Comment) Text() []byte {
+	withoutPrefixes := cutOptionalPrefixes(c.source, commentPrefixes)
+	return cutOptionalSuffixes(withoutPrefixes, commentSuffixes)
+}
+
+func cutOptionalPrefixes(input []byte, prefixes [][]byte) (output []byte) {
+	output = input
+	for _, prefix := range prefixes {
+		cut, _ := bytes.CutPrefix(output, prefix)
+		output = cut
+	}
+	return
+}
+
+func cutOptionalSuffixes(input []byte, suffixes [][]byte) (output []byte) {
+	output = input
+	for _, suffix := range suffixes {
+		cut, _ := bytes.CutSuffix(output, suffix)
+		output = cut
+	}
+	return
+}