Preserve comments in AST

runtime/ast/block.go

@@ -29,10 +29,18 @@ import (
 type Block struct {
 	Statements []Statement
 	Range
+	Comments
 }
 
 var _ Element = &Block{}
 
+// TODO(preserve-comments): Migrate and remove
+func NewBlockWithComments(memoryGauge common.MemoryGauge, statements []Statement, astRange Range, comments Comments) *Block {
+	block := NewBlock(memoryGauge, statements, astRange)
+	block.Comments = comments
+	return block
+}
+
 func NewBlock(memoryGauge common.MemoryGauge, statements []Statement, astRange Range) *Block {
 	common.UseMemory(memoryGauge, common.BlockMemoryUsage)
 

runtime/ast/comments.go

@@ -0,0 +1,71 @@
+package ast
+
+import (
+	"bytes"
+	"github.com/onflow/cadence/runtime/common"
+)
+
+type Comments struct {
+	Leading  []Comment `json:"-"`
+	Trailing []Comment `json:"-"`
+}
+
+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/function_declaration.go

@@ -20,6 +20,7 @@ package ast
 
 import (
 	"encoding/json"
+	"strings"
 
 	"github.com/turbolent/prettier"
 
@@ -69,17 +70,51 @@ type FunctionDeclaration struct {
 	ParameterList        *ParameterList
 	ReturnTypeAnnotation *TypeAnnotation
 	FunctionBlock        *FunctionBlock
-	DocString            string
 	Identifier           Identifier
 	StartPos             Position `json:"-"`
 	Access               Access
 	Flags                FunctionDeclarationFlags
+	Comments
 }
 
 var _ Element = &FunctionDeclaration{}
 var _ Declaration = &FunctionDeclaration{}
 var _ Statement = &FunctionDeclaration{}
 
+// TODO(preserve-comments): Temporary, add `comments` param to NewFunctionDeclaration in the future
+func NewFunctionDeclarationWithComments(
+	gauge common.MemoryGauge,
+	access Access,
+	purity FunctionPurity,
+	isStatic bool,
+	isNative bool,
+	identifier Identifier,
+	typeParameterList *TypeParameterList,
+	parameterList *ParameterList,
+	returnTypeAnnotation *TypeAnnotation,
+	functionBlock *FunctionBlock,
+	startPos Position,
+	docString string,
+	comments Comments,
+) *FunctionDeclaration {
+	decl := NewFunctionDeclaration(
+		gauge,
+		access,
+		purity,
+		isStatic,
+		isNative,
+		identifier,
+		typeParameterList,
+		parameterList,
+		returnTypeAnnotation,
+		functionBlock,
+		startPos,
+		docString,
+	)
+	decl.Comments = comments
+	return decl
+}
+
 func NewFunctionDeclaration(
 	gauge common.MemoryGauge,
 	access Access,
@@ -114,7 +149,6 @@ func NewFunctionDeclaration(
 		ReturnTypeAnnotation: returnTypeAnnotation,
 		FunctionBlock:        functionBlock,
 		StartPos:             startPos,
-		DocString:            docString,
 	}
 }
 
@@ -176,7 +210,16 @@ func (d *FunctionDeclaration) DeclarationMembers() *Members {
 }
 
 func (d *FunctionDeclaration) DeclarationDocString() string {
-	return d.DocString
+	var s strings.Builder
+	for _, comment := range d.Comments.Leading {
+		if comment.Doc() {
+			if s.Len() > 0 {
+				s.WriteRune('\n')
+			}
+			s.Write(comment.Text())
+		}
+	}
+	return s.String()
 }
 
 func (d *FunctionDeclaration) Doc() prettier.Doc {
@@ -200,16 +243,18 @@ func (d *FunctionDeclaration) MarshalJSON() ([]byte, error) {
 		*Alias
 		Type string
 		Range
-		IsStatic bool
-		IsNative bool
-		Flags    FunctionDeclarationFlags `json:",omitempty"`
+		IsStatic  bool
+		IsNative  bool
+		Flags     FunctionDeclarationFlags `json:",omitempty"`
+		DocString string
 	}{
-		Type:     "FunctionDeclaration",
-		Range:    NewUnmeteredRangeFromPositioned(d),
-		IsStatic: d.IsStatic(),
-		IsNative: d.IsNative(),
-		Alias:    (*Alias)(d),
-		Flags:    0,
+		Type:      "FunctionDeclaration",
+		Range:     NewUnmeteredRangeFromPositioned(d),
+		IsStatic:  d.IsStatic(),
+		IsNative:  d.IsNative(),
+		Alias:     (*Alias)(d),
+		Flags:     0,
+		DocString: d.DeclarationDocString(),
 	})
 }
 

runtime/ast/function_declaration_test.go

@@ -113,8 +113,10 @@ func TestFunctionDeclaration_MarshalJSON(t *testing.T) {
 				},
 			},
 		},
-		DocString: "test",
-		StartPos:  Position{Offset: 34, Line: 35, Column: 36},
+		Comments: Comments{
+			Leading: []Comment{NewComment(nil, []byte("///test"))},
+		},
+		StartPos: Position{Offset: 34, Line: 35, Column: 36},
 	}
 
 	actual, err := json.Marshal(decl)
@@ -584,8 +586,10 @@ func TestSpecialFunctionDeclaration_MarshalJSON(t *testing.T) {
 					},
 				},
 			},
-			DocString: "test",
-			StartPos:  Position{Offset: 34, Line: 35, Column: 36},
+			Comments: Comments{
+				Leading: []Comment{NewComment(nil, []byte("///test"))},
+			},
+			StartPos: Position{Offset: 34, Line: 35, Column: 36},
 		},
 	}
 

runtime/ast/identifier.go

@@ -20,7 +20,6 @@ package ast
 
 import (
 	"encoding/json"
-
 	"github.com/onflow/cadence/runtime/common"
 )
 

runtime/ast/position.go

@@ -143,6 +143,12 @@ func (e Range) EndPosition(common.MemoryGauge) Position {
 	return e.EndPos
 }
 
+func (e Range) Source(input []byte) []byte {
+	startOffset := e.StartPos.Offset
+	endOffset := e.EndPos.Offset + 1
+	return input[startOffset:endOffset]
+}
+
 // NewRangeFromPositioned
 
 func NewRangeFromPositioned(memoryGauge common.MemoryGauge, hasPosition HasPosition) Range {

runtime/old_parser/declaration_test.go

@@ -816,8 +816,10 @@ func TestParseFunctionDeclaration(t *testing.T) {
 							},
 						},
 					},
-					DocString: " Test",
-					StartPos:  ast.Position{Line: 2, Column: 0, Offset: 9},
+					Comments: ast.Comments{
+						Leading: []ast.Comment{ast.NewComment(nil, []byte("/// Test"))},
+					},
+					StartPos: ast.Position{Line: 2, Column: 0, Offset: 9},
 				},
 			},
 			result,
@@ -854,8 +856,13 @@ func TestParseFunctionDeclaration(t *testing.T) {
 							},
 						},
 					},
-					DocString: " First line\n Second line",
-					StartPos:  ast.Position{Line: 7, Column: 0, Offset: 39},
+					Comments: ast.Comments{
+						Leading: []ast.Comment{
+							ast.NewComment(nil, []byte("/// First line")),
+							ast.NewComment(nil, []byte("/// Second line")),
+						},
+					},
+					StartPos: ast.Position{Line: 7, Column: 0, Offset: 39},
 				},
 			},
 			result,
@@ -892,8 +899,12 @@ func TestParseFunctionDeclaration(t *testing.T) {
 							},
 						},
 					},
-					DocString: " Cool dogs.\n\n Cool cats!! ",
-					StartPos:  ast.Position{Line: 7, Column: 0, Offset: 39},
+					Comments: ast.Comments{
+						Leading: []ast.Comment{
+							ast.NewComment(nil, []byte("/** Cool dogs.\n\n Cool cats!! */")),
+						},
+					},
+					StartPos: ast.Position{Line: 7, Column: 0, Offset: 39},
 				},
 			},
 			result,
@@ -6870,8 +6881,12 @@ func TestParseMemberDocStrings(t *testing.T) {
 					Members: ast.NewUnmeteredMembers(
 						[]ast.Declaration{
 							&ast.FunctionDeclaration{
-								Access:    ast.AccessNotSpecified,
-								DocString: " noReturnNoBlock",
+								Access: ast.AccessNotSpecified,
+								Comments: ast.Comments{
+									Leading: []ast.Comment{
+										ast.NewComment(nil, []byte("/// noReturnNoBlock")),
+									},
+								},
 								Identifier: ast.Identifier{
 									Identifier: "noReturnNoBlock",
 									Pos:        ast.Position{Offset: 78, Line: 5, Column: 18},
@@ -6885,8 +6900,12 @@ func TestParseMemberDocStrings(t *testing.T) {
 								StartPos: ast.Position{Offset: 74, Line: 5, Column: 14},
 							},
 							&ast.FunctionDeclaration{
-								Access:    ast.AccessNotSpecified,
-								DocString: " returnNoBlock",
+								Access: ast.AccessNotSpecified,
+								Comments: ast.Comments{
+									Leading: []ast.Comment{
+										ast.NewComment(nil, []byte("/// returnNoBlock")),
+									},
+								},
 								Identifier: ast.Identifier{
 									Identifier: "returnNoBlock",
 									Pos:        ast.Position{Offset: 147, Line: 8, Column: 18},
@@ -6910,8 +6929,12 @@ func TestParseMemberDocStrings(t *testing.T) {
 								StartPos: ast.Position{Offset: 143, Line: 8, Column: 14},
 							},
 							&ast.FunctionDeclaration{
-								Access:    ast.AccessNotSpecified,
-								DocString: " returnAndBlock",
+								Access: ast.AccessNotSpecified,
+								Comments: ast.Comments{
+									Leading: []ast.Comment{
+										ast.NewComment(nil, []byte("/// returnAndBlock")),
+									},
+								},
 								Identifier: ast.Identifier{
 									Identifier: "returnAndBlock",
 									Pos:        ast.Position{Offset: 220, Line: 11, Column: 18},
@@ -6938,6 +6961,10 @@ func TestParseMemberDocStrings(t *testing.T) {
 											StartPos: ast.Position{Offset: 245, Line: 11, Column: 43},
 											EndPos:   ast.Position{Offset: 246, Line: 11, Column: 44},
 										},
+										Comments: ast.Comments{
+											Leading:  []ast.Comment{},
+											Trailing: []ast.Comment{},
+										},
 									},
 								},
 								StartPos: ast.Position{Offset: 216, Line: 11, Column: 14},
@@ -6988,8 +7015,12 @@ func TestParseMemberDocStrings(t *testing.T) {
 							&ast.SpecialFunctionDeclaration{
 								Kind: common.DeclarationKindUnknown,
 								FunctionDeclaration: &ast.FunctionDeclaration{
-									Access:    ast.AccessNotSpecified,
-									DocString: " unknown",
+									Access: ast.AccessNotSpecified,
+									Comments: ast.Comments{
+										Leading: []ast.Comment{
+											ast.NewComment(nil, []byte("/// unknown")),
+										},
+									},
 									Identifier: ast.Identifier{
 										Identifier: "unknown",
 										Pos:        ast.Position{Offset: 66, Line: 5, Column: 14},
@@ -7006,8 +7037,12 @@ func TestParseMemberDocStrings(t *testing.T) {
 							&ast.SpecialFunctionDeclaration{
 								Kind: common.DeclarationKindInitializer,
 								FunctionDeclaration: &ast.FunctionDeclaration{
-									Access:    ast.AccessNotSpecified,
-									DocString: " initNoBlock",
+									Access: ast.AccessNotSpecified,
+									Comments: ast.Comments{
+										Leading: []ast.Comment{
+											ast.NewComment(nil, []byte("/// initNoBlock")),
+										},
+									},
 									Identifier: ast.Identifier{
 										Identifier: "init",
 										Pos:        ast.Position{Offset: 121, Line: 8, Column: 14},
@@ -7024,8 +7059,12 @@ func TestParseMemberDocStrings(t *testing.T) {
 							&ast.SpecialFunctionDeclaration{
 								Kind: common.DeclarationKindDestructorLegacy,
 								FunctionDeclaration: &ast.FunctionDeclaration{
-									Access:    ast.AccessNotSpecified,
-									DocString: " destroyWithBlock",
+									Access: ast.AccessNotSpecified,
+									Comments: ast.Comments{
+										Leading: []ast.Comment{
+											ast.NewComment(nil, []byte("/// destroyWithBlock")),
+										},
+									},
 									Identifier: ast.Identifier{
 										Identifier: "destroy",
 										Pos:        ast.Position{Offset: 178, Line: 11, Column: 14},