Add support iterating references to iterables

runtime/interpreter/interpreter_statement.go

@@ -338,7 +338,7 @@ func (interpreter *Interpreter) VisitForStatement(statement *ast.ForStatement) S
 		panic(errors.NewUnreachableError())
 	}
 
-	iterator := iterable.Iterator(interpreter)
+	iterator := iterable.Iterator(interpreter, locationRange)
 
 	var index IntValue
 	if statement.Index != nil {

runtime/interpreter/value.go

@@ -234,7 +234,7 @@ type ContractValue interface {
 // IterableValue is a value which can be iterated over, e.g. with a for-loop
 type IterableValue interface {
 	Value
-	Iterator(interpreter *Interpreter) ValueIterator
+	Iterator(interpreter *Interpreter, locationRange LocationRange) ValueIterator
 }
 
 // ValueIterator is an iterator which returns values.
@@ -1580,7 +1580,7 @@ func (v *StringValue) ConformsToStaticType(
 	return true
 }
 
-func (v *StringValue) Iterator(_ *Interpreter) ValueIterator {
+func (v *StringValue) Iterator(_ *Interpreter, _ LocationRange) ValueIterator {
 	return StringValueIterator{
 		graphemes: uniseg.NewGraphemes(v.Str),
 	}
@@ -1614,7 +1614,7 @@ type ArrayValueIterator struct {
 	atreeIterator *atree.ArrayIterator
 }
 
-func (v *ArrayValue) Iterator(_ *Interpreter) ValueIterator {
+func (v *ArrayValue) Iterator(_ *Interpreter, _ LocationRange) ValueIterator {
 	arrayIterator, err := v.array.Iterator()
 	if err != nil {
 		panic(errors.NewExternalError(err))
@@ -19844,6 +19844,7 @@ var _ TypeIndexableValue = &StorageReferenceValue{}
 var _ MemberAccessibleValue = &StorageReferenceValue{}
 var _ AuthorizedValue = &StorageReferenceValue{}
 var _ ReferenceValue = &StorageReferenceValue{}
+var _ IterableValue = &StorageReferenceValue{}
 
 func NewUnmeteredStorageReferenceValue(
 	authorization Authorization,
@@ -20196,6 +20197,37 @@ func (*StorageReferenceValue) DeepRemove(_ *Interpreter) {
 
 func (*StorageReferenceValue) isReference() {}
 
+func (v *StorageReferenceValue) Iterator(interpreter *Interpreter, locationRange LocationRange) ValueIterator {
+	referencedValue := v.mustReferencedValue(interpreter, locationRange)
+	return referenceValueIterator(interpreter, referencedValue, v.BorrowedType, locationRange)
+}
+
+func referenceValueIterator(
+	interpreter *Interpreter,
+	referencedValue Value,
+	borrowedType sema.Type,
+	locationRange LocationRange,
+) ValueIterator {
+	referencedIterable, ok := referencedValue.(IterableValue)
+	if !ok {
+		panic(errors.NewUnreachableError())
+	}
+
+	referencedValueIterator := referencedIterable.Iterator(interpreter, locationRange)
+
+	referencedType, ok := borrowedType.(sema.ValueIndexableType)
+	if !ok {
+		panic(errors.NewUnreachableError())
+	}
+
+	elementType := referencedType.ElementType(false)
+
+	return ReferenceValueIterator{
+		iterator:    referencedValueIterator,
+		elementType: elementType,
+	}
+}
+
 // EphemeralReferenceValue
 
 type EphemeralReferenceValue struct {
@@ -20212,6 +20244,7 @@ var _ TypeIndexableValue = &EphemeralReferenceValue{}
 var _ MemberAccessibleValue = &EphemeralReferenceValue{}
 var _ AuthorizedValue = &EphemeralReferenceValue{}
 var _ ReferenceValue = &EphemeralReferenceValue{}
+var _ IterableValue = &EphemeralReferenceValue{}
 
 func NewUnmeteredEphemeralReferenceValue(
 	authorization Authorization,
@@ -20541,6 +20574,35 @@ func (*EphemeralReferenceValue) DeepRemove(_ *Interpreter) {
 
 func (*EphemeralReferenceValue) isReference() {}
 
+func (v *EphemeralReferenceValue) Iterator(interpreter *Interpreter, locationRange LocationRange) ValueIterator {
+	referencedValue := v.MustReferencedValue(interpreter, locationRange)
+	return referenceValueIterator(interpreter, referencedValue, v.BorrowedType, locationRange)
+}
+
+// ReferenceValueIterator
+
+type ReferenceValueIterator struct {
+	iterator    ValueIterator
+	elementType sema.Type
+}
+
+var _ ValueIterator = ReferenceValueIterator{}
+
+func (i ReferenceValueIterator) Next(interpreter *Interpreter) Value {
+	element := i.iterator.Next(interpreter)
+
+	if element == nil {
+		return nil
+	}
+
+	// For non-primitive values, return a reference.
+	if i.elementType.ContainFieldsOrElements() {
+		return NewEphemeralReferenceValue(interpreter, UnauthorizedAccess, element, i.elementType)
+	}
+
+	return element
+}
+
 // AddressValue
 type AddressValue common.Address
 

runtime/sema/check_for.go

@@ -43,41 +43,17 @@ func (checker *Checker) VisitForStatement(statement *ast.ForStatement) (_ struct
 
 	valueType := checker.VisitExpression(valueExpression, expectedType)
 
-	var elementType Type = InvalidType
-
-	if !valueType.IsInvalidType() {
-
-		// Only get the element type if the array is not a resource array.
-		// Otherwise, in addition to the `UnsupportedResourceForLoopError`,
-		// the loop variable will be declared with the resource-typed element type,
-		// leading to an additional `ResourceLossError`.
-
-		if valueType.IsResourceType() {
-			checker.report(
-				&UnsupportedResourceForLoopError{
-					Range: ast.NewRangeFromPositioned(checker.memoryGauge, valueExpression),
-				},
-			)
-		} else if arrayType, ok := valueType.(ArrayType); ok {
-			elementType = arrayType.ElementType(false)
-		} else if valueType == StringType {
-			elementType = CharacterType
-		} else {
-			checker.report(
-				&TypeMismatchWithDescriptionError{
-					ExpectedTypeDescription: "array",
-					ActualType:              valueType,
-					Range:                   ast.NewRangeFromPositioned(checker.memoryGauge, valueExpression),
-				},
-			)
-		}
-	}
+	// Only get the element type if the array is not a resource array.
+	// Otherwise, in addition to the `UnsupportedResourceForLoopError`,
+	// the loop variable will be declared with the resource-typed element type,
+	// leading to an additional `ResourceLossError`.
+	loopVariableType := checker.loopVariableType(valueType, valueExpression)
 
 	identifier := statement.Identifier.Identifier
 
 	variable, err := checker.valueActivations.declare(variableDeclaration{
 		identifier:               identifier,
-		ty:                       elementType,
+		ty:                       loopVariableType,
 		kind:                     common.DeclarationKindConstant,
 		pos:                      statement.Identifier.Pos,
 		isConstant:               true,
@@ -123,3 +99,68 @@ func (checker *Checker) VisitForStatement(statement *ast.ForStatement) (_ struct
 
 	return
 }
+
+func (checker *Checker) loopVariableType(valueType Type, hasPosition ast.HasPosition) Type {
+	if valueType.IsInvalidType() {
+		return InvalidType
+	}
+
+	// Resources cannot be looped.
+	if valueType.IsResourceType() {
+		checker.report(
+			&UnsupportedResourceForLoopError{
+				Range: ast.NewRangeFromPositioned(checker.memoryGauge, hasPosition),
+			},
+		)
+		return InvalidType
+	}
+
+	// If it's a reference, check whether the referenced type is iterable.
+	// If yes, then determine the loop-var type depending on the
+	// element-type of the referenced type.
+	// If that element type is:
+	//  a) A container type, then the loop-var is also a reference-type.
+	//  b) A primitive type, then the loop-var is the concrete type itself.
+
+	if referenceType, ok := valueType.(*ReferenceType); ok {
+		referencedType := referenceType.Type
+		referencedIterableElementType := checker.iterableElementType(referencedType, hasPosition)
+
+		if referencedIterableElementType.IsInvalidType() {
+			return referencedIterableElementType
+		}
+
+		// Case (a): Element type is a container type.
+		// Then the loop-var must also be a reference type.
+		if referencedIterableElementType.ContainFieldsOrElements() {
+			return NewReferenceType(checker.memoryGauge, UnauthorizedAccess, referencedIterableElementType)
+		}
+
+		// Case (b): Element type is a primitive type.
+		// Then the loop-var must be the concrete type.
+		return referencedIterableElementType
+	}
+
+	// If it's not a reference, then simply get the element type.
+	return checker.iterableElementType(valueType, hasPosition)
+}
+
+func (checker *Checker) iterableElementType(valueType Type, hasPosition ast.HasPosition) Type {
+	if arrayType, ok := valueType.(ArrayType); ok {
+		return arrayType.ElementType(false)
+	}
+
+	if valueType == StringType {
+		return CharacterType
+	}
+
+	checker.report(
+		&TypeMismatchWithDescriptionError{
+			ExpectedTypeDescription: "array",
+			ActualType:              valueType,
+			Range:                   ast.NewRangeFromPositioned(checker.memoryGauge, hasPosition),
+		},
+	)
+
+	return InvalidType
+}

runtime/tests/checker/for_test.go

@@ -22,6 +22,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	"github.com/onflow/cadence/runtime/sema"
 )
@@ -273,3 +274,122 @@ func TestCheckInvalidForShadowing(t *testing.T) {
 
 	assert.IsType(t, &sema.RedeclarationError{}, errs[0])
 }
+
+func TestCheckReferencesInForLoop(t *testing.T) {
+
+	t.Parallel()
+
+	t.Run("Primitive array", func(t *testing.T) {
+		t.Parallel()
+
+		_, err := ParseAndCheck(t, `
+            fun main() {
+                var array = ["Hello", "World", "Foo", "Bar"]
+                var arrayRef = &array as &[String]
+
+                for element in arrayRef {
+                    let e: String = element
+                }
+            }
+        `)
+
+		require.NoError(t, err)
+	})
+
+	t.Run("Struct array", func(t *testing.T) {
+		t.Parallel()
+
+		_, err := ParseAndCheck(t, `
+            struct Foo{}
+
+            fun main() {
+                var array = [Foo(), Foo()]
+                var arrayRef = &array as &[Foo]
+
+                for element in arrayRef {
+                    let e: &Foo = element
+                }
+            }
+        `)
+
+		require.NoError(t, err)
+	})
+
+	t.Run("Resource array", func(t *testing.T) {
+		t.Parallel()
+
+		_, err := ParseAndCheck(t, `
+            resource Foo{}
+
+            fun main() {
+                var array <- [ <- create Foo(), <- create Foo()]
+                var arrayRef = &array as &[Foo]
+
+                for element in arrayRef {
+                    let e: &Foo = element
+                }
+
+                destroy array
+            }
+        `)
+
+		require.NoError(t, err)
+	})
+
+	t.Run("Dictionary", func(t *testing.T) {
+		t.Parallel()
+
+		_, err := ParseAndCheck(t, `
+            struct Foo{}
+
+            fun main() {
+                var foo = {"foo": Foo()}
+                var fooRef = &foo as &{String: Foo}
+
+                for element in fooRef {
+                    let e: &Foo = element
+                }
+            }
+        `)
+
+		errors := RequireCheckerErrors(t, err, 1)
+		assert.IsType(t, &sema.TypeMismatchWithDescriptionError{}, errors[0])
+	})
+
+	t.Run("Non iterable", func(t *testing.T) {
+		t.Parallel()
+
+		_, err := ParseAndCheck(t, `
+            struct Foo{}
+
+            fun main() {
+                var foo = Foo()
+                var fooRef = &foo as &Foo
+
+                for element in fooRef {
+                    let e: &Foo = element
+                }
+            }
+        `)
+
+		errors := RequireCheckerErrors(t, err, 1)
+		assert.IsType(t, &sema.TypeMismatchWithDescriptionError{}, errors[0])
+	})
+
+	t.Run("Non existing type", func(t *testing.T) {
+		t.Parallel()
+
+		_, err := ParseAndCheck(t, `
+            fun main() {
+                var foo = Foo()
+                var fooRef = &foo as &Foo
+
+                for element in fooRef {}
+            }
+        `)
+
+		errors := RequireCheckerErrors(t, err, 2)
+		assert.IsType(t, &sema.NotDeclaredError{}, errors[0])
+		assert.IsType(t, &sema.NotDeclaredError{}, errors[1])
+	})
+}