[WIP] Retain comments in AST

runtime/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 {
@@ -199,9 +199,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(),
 	})
 }
 

runtime/ast/attachment_test.go

@@ -56,8 +56,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},
@@ -141,8 +145,12 @@ func TestAttachmentDeclaration_Doc(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},

runtime/ast/block.go

@@ -29,11 +29,12 @@ 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{

runtime/ast/comments.go

@@ -0,0 +1,93 @@
+package ast
+
+import (
+	"bytes"
+	"github.com/onflow/cadence/runtime/common"
+	"strings"
+)
+
+type Comments struct {
+	Leading  []*Comment `json:"-"`
+	Trailing []*Comment `json:"-"`
+}
+
+func (c Comments) PackToList() []*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.Doc() {
+			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 {
+	// TODO(preserve-comments): Track memory usage
+	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) Doc() bool {
+	if c.Multiline() {
+		return bytes.HasPrefix(c.source, blockCommentDocStringPrefix)
+	} else {
+		return bytes.HasPrefix(c.source, lineCommentDocStringPrefix)
+	}
+}
+
+// Text without opening/closing comment characters /*, /**, */, //
+func (c Comment) Text() []byte {
+	withoutPrefixes := cutOptionalPrefixes(c.source, [][]byte{
+		blockCommentDocStringPrefix, // must be before blockCommentStringPrefix
+		blockCommentStringPrefix,
+		lineCommentDocStringPrefix, // must be before lineCommentStringPrefix
+		lineCommentStringPrefix,
+	})
+	return cutOptionalSuffixes(withoutPrefixes, [][]byte{
+		blockCommentStringSuffix,
+	})
+}
+
+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
+}

runtime/ast/composite.go

@@ -51,12 +51,12 @@ func IsResourceDestructionDefaultEvent(identifier string) bool {
 
 type CompositeDeclaration struct {
 	Members      *Members
-	DocString    string
 	Conformances []*NominalType
 	Identifier   Identifier
 	Range
 	Access        Access
 	CompositeKind common.CompositeKind
+	Comments
 }
 
 var _ Element = &CompositeDeclaration{}
@@ -71,8 +71,8 @@ func NewCompositeDeclaration(
 	identifier Identifier,
 	conformances []*NominalType,
 	members *Members,
-	docString string,
 	declarationRange Range,
+	comments Comments,
 ) *CompositeDeclaration {
 	common.UseMemory(memoryGauge, common.CompositeDeclarationMemoryUsage)
 
@@ -82,8 +82,8 @@ func NewCompositeDeclaration(
 		Identifier:    identifier,
 		Conformances:  conformances,
 		Members:       members,
-		DocString:     docString,
 		Range:         declarationRange,
+		Comments:      comments,
 	}
 }
 
@@ -120,17 +120,19 @@ func (d *CompositeDeclaration) DeclarationMembers() *Members {
 }
 
 func (d *CompositeDeclaration) DeclarationDocString() string {
-	return d.DocString
+	return d.Comments.LeadingDocString()
 }
 
 func (d *CompositeDeclaration) MarshalJSON() ([]byte, error) {
 	type Alias CompositeDeclaration
 	return json.Marshal(&struct {
 		*Alias
-		Type string
+		Type      string
+		DocString string
 	}{
-		Type:  "CompositeDeclaration",
-		Alias: (*Alias)(d),
+		Type:      "CompositeDeclaration",
+		Alias:     (*Alias)(d),
+		DocString: d.DeclarationDocString(),
 	})
 }
 
@@ -304,12 +306,14 @@ const (
 
 type FieldDeclaration struct {
 	TypeAnnotation *TypeAnnotation
-	DocString      string
-	Identifier     Identifier
+	// TODO(preserve-comments): Remove
+	DocString  string
+	Identifier Identifier
 	Range
 	Access       Access
 	VariableKind VariableKind
 	Flags        FieldDeclarationFlags
+	Comments
 }
 
 var _ Element = &FieldDeclaration{}
@@ -323,8 +327,8 @@ func NewFieldDeclaration(
 	variableKind VariableKind,
 	identifier Identifier,
 	typeAnnotation *TypeAnnotation,
-	docString string,
 	declRange Range,
+	comments Comments,
 ) *FieldDeclaration {
 	common.UseMemory(memoryGauge, common.FieldDeclarationMemoryUsage)
 
@@ -342,8 +346,8 @@ func NewFieldDeclaration(
 		VariableKind:   variableKind,
 		Identifier:     identifier,
 		TypeAnnotation: typeAnnotation,
-		DocString:      docString,
 		Range:          declRange,
+		Comments:       comments,
 	}
 }
 
@@ -382,16 +386,18 @@ func (d *FieldDeclaration) MarshalJSON() ([]byte, error) {
 	type Alias FieldDeclaration
 	return json.Marshal(&struct {
 		*Alias
-		Type     string
-		Flags    FieldDeclarationFlags `json:",omitempty"`
-		IsStatic bool
-		IsNative bool
+		Type      string
+		Flags     FieldDeclarationFlags `json:",omitempty"`
+		IsStatic  bool
+		IsNative  bool
+		DocString string
 	}{
-		Type:     "FieldDeclaration",
-		Alias:    (*Alias)(d),
-		IsStatic: d.IsStatic(),
-		IsNative: d.IsNative(),
-		Flags:    0,
+		Type:      "FieldDeclaration",
+		Alias:     (*Alias)(d),
+		IsStatic:  d.IsStatic(),
+		IsNative:  d.IsNative(),
+		Flags:     0,
+		DocString: d.DeclarationDocString(),
 	})
 }
 

runtime/ast/composite_test.go

@@ -406,8 +406,12 @@ func TestCompositeDeclaration_MarshalJSON(t *testing.T) {
 				},
 			},
 		},
-		Members:   NewUnmeteredMembers([]Declaration{}),
-		DocString: "test",
+		Members: NewUnmeteredMembers([]Declaration{}),
+		Comments: Comments{
+			Leading: []*Comment{
+				NewComment(nil, []byte("///test")),
+			},
+		},
 		Range: Range{
 			StartPos: Position{Offset: 7, Line: 8, Column: 9},
 			EndPos:   Position{Offset: 10, Line: 11, Column: 12},