Fix AST walking

ast/access.go

@@ -33,6 +33,7 @@ import (
 
 type Access interface {
 	isAccess()
+	Walk(walkChild func(Element))
 	Keyword() string
 	Description() string
 	String() string
@@ -172,6 +173,10 @@ func NewEntitlementAccess(entitlements EntitlementSet) EntitlementAccess {
 
 func (EntitlementAccess) isAccess() {}
 
+func (e EntitlementAccess) Walk(walkChild func(Element)) {
+	e.EntitlementSet.Walk(walkChild)
+}
+
 func (EntitlementAccess) Description() string {
 	return "entitled access"
 }
@@ -300,6 +305,8 @@ func PrimitiveAccessCount() int {
 
 func (PrimitiveAccess) isAccess() {}
 
+func (PrimitiveAccess) Walk(_ func(Element)) {}
+
 // TODO: remove.
 //   only used by tests which are not updated yet
 //   to include contract and account access

ast/access_test.go

@@ -87,6 +87,29 @@ func TestMappedAccess_Walk(t *testing.T) {
 	assert.Equal(t, []Element{mapType}, visited)
 }
 
+func TestEntitlementAccess_Walk(t *testing.T) {
+
+	t.Parallel()
+
+	e := &NominalType{
+		Identifier: Identifier{Identifier: "E"},
+	}
+	f := &NominalType{
+		Identifier: Identifier{Identifier: "F"},
+	}
+
+	access := NewEntitlementAccess(
+		NewConjunctiveEntitlementSet([]*NominalType{e, f}),
+	)
+
+	var visited []Element
+	access.Walk(func(element Element) {
+		visited = append(visited, element)
+	})
+
+	assert.Equal(t, []Element{e, f}, visited)
+}
+
 func TestConjunctiveEntitlementSet_MarshalJSON(t *testing.T) {
 
 	t.Parallel()

ast/attachment.go

@@ -71,7 +71,14 @@ func (*AttachmentDeclaration) ElementType() ElementType {
 }
 
 func (d *AttachmentDeclaration) Walk(walkChild func(Element)) {
-	walkDeclarations(walkChild, d.Members.declarations)
+	if d.Access != nil {
+		d.Access.Walk(walkChild)
+	}
+	if d.BaseType != nil {
+		walkChild(d.BaseType)
+	}
+	walkElements(walkChild, d.Conformances)
+	walkElements(walkChild, d.Members.declarations)
 }
 
 func (*AttachmentDeclaration) isDeclaration() {}

ast/attachment_test.go

@@ -252,38 +252,141 @@ func TestAttachmentDeclaration_Walk(t *testing.T) {
 
 	t.Parallel()
 
-	field := &FieldDeclaration{
-		Identifier: Identifier{Identifier: "field"},
-		TypeAnnotation: &TypeAnnotation{
-			Type: &NominalType{
-				Identifier: Identifier{Identifier: "Int"},
+	t.Run("members only", func(t *testing.T) {
+		t.Parallel()
+
+		field := &FieldDeclaration{
+			Identifier: Identifier{Identifier: "field"},
+			TypeAnnotation: &TypeAnnotation{
+				Type: &NominalType{
+					Identifier: Identifier{Identifier: "Int"},
+				},
 			},
-		},
-	}
+		}
 
-	function := &FunctionDeclaration{
-		Identifier: Identifier{Identifier: "function"},
-	}
+		function := &FunctionDeclaration{
+			Identifier: Identifier{Identifier: "function"},
+		}
 
-	decl := &AttachmentDeclaration{
-		Members: NewUnmeteredMembers([]Declaration{
-			field,
-			function,
-		}),
-	}
+		decl := &AttachmentDeclaration{
+			Members: NewUnmeteredMembers([]Declaration{
+				field,
+				function,
+			}),
+		}
 
-	var visited []Element
-	decl.Walk(func(element Element) {
-		visited = append(visited, element)
+		var visited []Element
+		decl.Walk(func(element Element) {
+			visited = append(visited, element)
+		})
+
+		assert.Equal(t,
+			[]Element{
+				field,
+				function,
+			},
+			visited,
+		)
 	})
 
-	assert.Equal(t,
-		[]Element{
-			field,
-			function,
-		},
-		visited,
-	)
+	t.Run("with base type", func(t *testing.T) {
+		t.Parallel()
+
+		baseType := &NominalType{
+			Identifier: Identifier{Identifier: "Base"},
+		}
+		field := &FieldDeclaration{
+			Identifier: Identifier{Identifier: "field"},
+		}
+
+		decl := &AttachmentDeclaration{
+			BaseType: baseType,
+			Members:  NewUnmeteredMembers([]Declaration{field}),
+		}
+
+		var visited []Element
+		decl.Walk(func(element Element) {
+			visited = append(visited, element)
+		})
+
+		assert.Equal(t,
+			[]Element{
+				baseType,
+				field,
+			},
+			visited,
+		)
+	})
+
+	t.Run("with conformances", func(t *testing.T) {
+		t.Parallel()
+
+		conformance := &NominalType{
+			Identifier: Identifier{Identifier: "Foo"},
+		}
+		field := &FieldDeclaration{
+			Identifier: Identifier{Identifier: "field"},
+		}
+
+		decl := &AttachmentDeclaration{
+			Conformances: []*NominalType{conformance},
+			Members:      NewUnmeteredMembers([]Declaration{field}),
+		}
+
+		var visited []Element
+		decl.Walk(func(element Element) {
+			visited = append(visited, element)
+		})
+
+		assert.Equal(t,
+			[]Element{
+				conformance,
+				field,
+			},
+			visited,
+		)
+	})
+
+	t.Run("with entitlement access, base type, and conformances", func(t *testing.T) {
+		t.Parallel()
+
+		entitlement := &NominalType{
+			Identifier: Identifier{Identifier: "E"},
+		}
+		baseType := &NominalType{
+			Identifier: Identifier{Identifier: "Base"},
+		}
+		conformance := &NominalType{
+			Identifier: Identifier{Identifier: "Iface"},
+		}
+		field := &FieldDeclaration{
+			Identifier: Identifier{Identifier: "field"},
+		}
+
+		decl := &AttachmentDeclaration{
+			Access: NewEntitlementAccess(
+				NewConjunctiveEntitlementSet([]*NominalType{entitlement}),
+			),
+			BaseType:     baseType,
+			Conformances: []*NominalType{conformance},
+			Members:      NewUnmeteredMembers([]Declaration{field}),
+		}
+
+		var visited []Element
+		decl.Walk(func(element Element) {
+			visited = append(visited, element)
+		})
+
+		assert.Equal(t,
+			[]Element{
+				entitlement,
+				baseType,
+				conformance,
+				field,
+			},
+			visited,
+		)
+	})
 }
 
 func TestAttachExpressionMarshallJSON(t *testing.T) {

ast/block.go

@@ -51,7 +51,7 @@ func (b *Block) IsEmpty() bool {
 }
 
 func (b *Block) Walk(walkChild func(Element)) {
-	walkStatements(walkChild, b.Statements)
+	walkElements(walkChild, b.Statements)
 }
 
 var blockStartDoc prettier.Doc = prettier.Text("{")

ast/block_test.go

@@ -870,7 +870,7 @@ func TestFunctionBlock_Walk(t *testing.T) {
 	)
 }
 
-func TestTestCondition_Doc(t *testing.T) {
+func TestCondition_Doc(t *testing.T) {
 	t.Parallel()
 
 	t.Run("with test and message", func(t *testing.T) {
@@ -972,6 +972,87 @@ func TestTestCondition_Doc(t *testing.T) {
 	})
 }
 
+func TestCondition_Walk(t *testing.T) {
+
+	t.Parallel()
+
+	t.Run("with test only", func(t *testing.T) {
+		t.Parallel()
+
+		test := &BoolExpression{Value: true}
+
+		condition := TestCondition{
+			Test: test,
+		}
+
+		var visited []Element
+		condition.Walk(func(element Element) {
+			visited = append(visited, element)
+		})
+
+		assert.Equal(t, []Element{test}, visited)
+	})
+
+	t.Run("with test and message", func(t *testing.T) {
+		t.Parallel()
+
+		test := &BoolExpression{Value: true}
+		message := &StringExpression{Value: "fail"}
+
+		condition := TestCondition{
+			Test:    test,
+			Message: message,
+		}
+
+		var visited []Element
+		condition.Walk(func(element Element) {
+			visited = append(visited, element)
+		})
+
+		assert.Equal(t, []Element{test, message}, visited)
+	})
+}
+
+func TestConditions_Walk(t *testing.T) {
+
+	t.Parallel()
+
+	t.Run("empty", func(t *testing.T) {
+		t.Parallel()
+
+		conditions := &Conditions{}
+
+		var visited []Element
+		conditions.Walk(func(element Element) {
+			visited = append(visited, element)
+		})
+
+		assert.Empty(t, visited)
+	})
+
+	t.Run("with conditions", func(t *testing.T) {
+		t.Parallel()
+
+		cond1 := TestCondition{
+			Test: &BoolExpression{Value: true},
+		}
+		cond2 := TestCondition{
+			Test: &BoolExpression{Value: false},
+		}
+
+		conditions := &Conditions{
+			Conditions: []Condition{cond1, cond2},
+		}
+
+		var visited []Element
+		conditions.Walk(func(element Element) {
+			visited = append(visited, element)
+		})
+
+		assert.Equal(t, []Element{cond1, cond2}, visited)
+	})
+}
+
 func TestEmitCondition_Walk(t *testing.T) {
 
 	t.Parallel()

ast/composite.go

@@ -92,7 +92,11 @@ func (*CompositeDeclaration) ElementType() ElementType {
 }
 
 func (d *CompositeDeclaration) Walk(walkChild func(Element)) {
-	walkDeclarations(walkChild, d.Members.declarations)
+	if d.Access != nil {
+		d.Access.Walk(walkChild)
+	}
+	walkElements(walkChild, d.Conformances)
+	walkElements(walkChild, d.Members.declarations)
 }
 
 func (*CompositeDeclaration) isDeclaration() {}
@@ -354,6 +358,9 @@ func (*FieldDeclaration) ElementType() ElementType {
 }
 
 func (d *FieldDeclaration) Walk(walkChild func(Element)) {
+	if d.Access != nil {
+		d.Access.Walk(walkChild)
+	}
 	if d.TypeAnnotation != nil {
 		walkChild(d.TypeAnnotation)
 	}
@@ -525,8 +532,10 @@ func (*EnumCaseDeclaration) ElementType() ElementType {
 	return ElementTypeEnumCaseDeclaration
 }
 
-func (*EnumCaseDeclaration) Walk(_ func(Element)) {
-	// NO-OP
+func (d *EnumCaseDeclaration) Walk(walkChild func(Element)) {
+	if d.Access != nil {
+		d.Access.Walk(walkChild)
+	}
 }
 
 func (*EnumCaseDeclaration) isDeclaration() {}