Improve pretty printing

go.mod

@@ -15,7 +15,7 @@ require (
 	github.com/rivo/uniseg v0.2.1-0.20211004051800-57c86be7915a
 	github.com/schollz/progressbar/v3 v3.8.3
 	github.com/stretchr/testify v1.7.1
-	github.com/turbolent/prettier v0.0.0-20210613180524-3a3f5a5b49ba
+	github.com/turbolent/prettier v0.0.0-20220320183459-661cc755135d
 	go.uber.org/goleak v1.1.10
 	golang.org/x/crypto v0.0.0-20210921155107-089bfa567519
 	golang.org/x/text v0.3.7

go.sum

@@ -60,8 +60,8 @@ github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UV
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 github.com/stretchr/testify v1.7.1 h1:5TQK59W5E3v0r2duFAb7P95B6hEeOyEnHRa8MjYSMTY=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/turbolent/prettier v0.0.0-20210613180524-3a3f5a5b49ba h1:GPg+SVJURgCt6b4IwuRQupixdBM+KzjXPGvawnaQ15E=
-github.com/turbolent/prettier v0.0.0-20210613180524-3a3f5a5b49ba/go.mod h1:Nlx5Y115XQvNcIdIy7dZXaNSUpzwBSge4/Ivk93/Yog=
+github.com/turbolent/prettier v0.0.0-20220320183459-661cc755135d h1:5JInRQbk5UBX8JfUvKh2oYTLMVwj3p6n+wapDDm7hko=
+github.com/turbolent/prettier v0.0.0-20220320183459-661cc755135d/go.mod h1:Nlx5Y115XQvNcIdIy7dZXaNSUpzwBSge4/Ivk93/Yog=
 github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM=
 github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg=
 github.com/zeebo/assert v1.1.0 h1:hU1L1vLTHsnO8x8c9KAR5GmM5QscxHg5RNU5z5qbUWY=

runtime/ast/argument.go

@@ -20,7 +20,8 @@ package ast
 
 import (
 	"encoding/json"
-	"strings"
+
+	"github.com/turbolent/prettier"
 
 	"github.com/onflow/cadence/runtime/common"
 )
@@ -68,13 +69,7 @@ func (a *Argument) EndPosition(memoryGauge common.MemoryGauge) Position {
 }
 
 func (a *Argument) String() string {
-	var builder strings.Builder
-	if a.Label != "" {
-		builder.WriteString(a.Label)
-		builder.WriteString(": ")
-	}
-	builder.WriteString(a.Expression.String())
-	return builder.String()
+	return Prettier(a)
 }
 
 func (a *Argument) MarshalJSON() ([]byte, error) {
@@ -87,3 +82,14 @@ func (a *Argument) MarshalJSON() ([]byte, error) {
 		Alias: (*Alias)(a),
 	})
 }
+
+func (a *Argument) Doc() prettier.Doc {
+	argumentDoc := a.Expression.Doc()
+	if a.Label == "" {
+		return argumentDoc
+	}
+	return prettier.Concat{
+		prettier.Text(a.Label + ": "),
+		argumentDoc,
+	}
+}

runtime/ast/argument_test.go

@@ -24,6 +24,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/turbolent/prettier"
 )
 
 func TestArgument_MarshalJSON(t *testing.T) {
@@ -108,3 +109,86 @@ func TestArgument_MarshalJSON(t *testing.T) {
 		)
 	})
 }
+
+func TestArgument_Doc(t *testing.T) {
+
+	t.Parallel()
+
+	t.Run("without label", func(t *testing.T) {
+
+		t.Parallel()
+
+		argument := &Argument{
+			Expression: &BoolExpression{
+				Value: false,
+			},
+		}
+
+		require.Equal(
+			t,
+			prettier.Text("false"),
+			argument.Doc(),
+		)
+	})
+
+	t.Run("with label", func(t *testing.T) {
+
+		t.Parallel()
+
+		argument := &Argument{
+			Label: "ok",
+			Expression: &BoolExpression{
+				Value: false,
+			},
+		}
+
+		require.Equal(
+			t,
+			prettier.Concat{
+				prettier.Text("ok: "),
+				prettier.Text("false"),
+			},
+			argument.Doc(),
+		)
+	})
+}
+
+func TestArgument_String(t *testing.T) {
+
+	t.Parallel()
+
+	t.Run("without label", func(t *testing.T) {
+
+		t.Parallel()
+
+		argument := &Argument{
+			Expression: &BoolExpression{
+				Value: false,
+			},
+		}
+
+		require.Equal(
+			t,
+			"false",
+			argument.String(),
+		)
+	})
+
+	t.Run("with label", func(t *testing.T) {
+
+		t.Parallel()
+
+		argument := &Argument{
+			Label: "ok",
+			Expression: &BoolExpression{
+				Value: false,
+			},
+		}
+
+		require.Equal(
+			t,
+			"ok: false",
+			argument.String(),
+		)
+	})
+}

runtime/ast/block.go

@@ -78,23 +78,21 @@ func (b *Block) Doc() prettier.Doc {
 }
 
 func StatementsDoc(statements []Statement) prettier.Doc {
-	var statementsDoc prettier.Concat
+	var doc prettier.Concat
 
 	for _, statement := range statements {
-		// TODO: replace once Statement implements Doc
-		hasDoc, ok := statement.(interface{ Doc() prettier.Doc })
-		if !ok {
-			continue
-		}
-
-		statementsDoc = append(
-			statementsDoc,
+		doc = append(
+			doc,
 			prettier.HardLine{},
-			hasDoc.Doc(),
+			statement.Doc(),
 		)
 	}
 
-	return statementsDoc
+	return doc
+}
+
+func (b *Block) String() string {
+	return Prettier(b)
 }
 
 func (b *Block) MarshalJSON() ([]byte, error) {
@@ -174,6 +172,61 @@ func (b *FunctionBlock) EndPosition(common.MemoryGauge) Position {
 	return b.Block.EndPos
 }
 
+var preConditionsKeywordDoc = prettier.Text("pre")
+var postConditionsKeywordDoc = prettier.Text("post")
+
+func (b *FunctionBlock) Doc() prettier.Doc {
+	if b.IsEmpty() {
+		return blockEmptyDoc
+	}
+
+	var conditionDocs []prettier.Doc
+
+	if conditionsDoc := b.PreConditions.Doc(preConditionsKeywordDoc); conditionsDoc != nil {
+		conditionDocs = append(
+			conditionDocs,
+			prettier.HardLine{},
+			conditionsDoc,
+		)
+	}
+
+	if conditionsDoc := b.PostConditions.Doc(postConditionsKeywordDoc); conditionsDoc != nil {
+		conditionDocs = append(
+			conditionDocs,
+			prettier.HardLine{},
+			conditionsDoc,
+		)
+	}
+
+	var bodyDoc prettier.Doc
+
+	statementsDoc := StatementsDoc(b.Block.Statements)
+
+	if len(conditionDocs) > 0 {
+		bodyConcatDoc := prettier.Concat(conditionDocs)
+		bodyConcatDoc = append(
+			bodyConcatDoc,
+			statementsDoc,
+		)
+		bodyDoc = bodyConcatDoc
+	} else {
+		bodyDoc = statementsDoc
+	}
+
+	return prettier.Concat{
+		blockStartDoc,
+		prettier.Indent{
+			Doc: bodyDoc,
+		},
+		prettier.HardLine{},
+		blockEndDoc,
+	}
+}
+
+func (b *FunctionBlock) String() string {
+	return Prettier(b)
+}
+
 // Condition
 
 type Condition struct {
@@ -182,10 +235,59 @@ type Condition struct {
 	Message Expression
 }
 
+func (c Condition) Doc() prettier.Doc {
+	doc := c.Test.Doc()
+	if c.Message != nil {
+		doc = prettier.Concat{
+			doc,
+			prettier.Text(":"),
+			prettier.Indent{
+				Doc: prettier.Concat{
+					prettier.HardLine{},
+					c.Message.Doc(),
+				},
+			},
+		}
+	}
+
+	return prettier.Group{
+		Doc: doc,
+	}
+}
+
 // Conditions
 
 type Conditions []*Condition
 
 func (c *Conditions) IsEmpty() bool {
 	return c == nil || len(*c) == 0
 }
+
+func (c *Conditions) Doc(keywordDoc prettier.Doc) prettier.Doc {
+	if c.IsEmpty() {
+		return nil
+	}
+
+	var doc prettier.Concat
+
+	for _, condition := range *c {
+		doc = append(
+			doc,
+			prettier.HardLine{},
+			condition.Doc(),
+		)
+	}
+
+	return prettier.Group{
+		Doc: prettier.Concat{
+			keywordDoc,
+			prettier.Space,
+			blockStartDoc,
+			prettier.Indent{
+				Doc: doc,
+			},
+			prettier.HardLine{},
+			blockEndDoc,
+		},
+	}
+}