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/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/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)
 	}