Add support iterating references to iterables

runtime/interpreter/interpreter_statement.go

@@ -313,7 +313,7 @@ func (interpreter *Interpreter) VisitWhileStatement(statement *ast.WhileStatemen
 
 var intOne = NewUnmeteredIntValueFromInt64(1)
 
-func (interpreter *Interpreter) VisitForStatement(statement *ast.ForStatement) StatementResult {
+func (interpreter *Interpreter) VisitForStatement(statement *ast.ForStatement) (result StatementResult) {
 
 	interpreter.activations.PushNewWithCurrent()
 	defer interpreter.activations.Pop()
@@ -338,28 +338,36 @@ func (interpreter *Interpreter) VisitForStatement(statement *ast.ForStatement) S
 		panic(errors.NewUnreachableError())
 	}
 
-	iterator := iterable.Iterator(interpreter)
+	forStmtTypes := interpreter.Program.Elaboration.ForStatementType(statement)
 
 	var index IntValue
 	if statement.Index != nil {
 		index = NewIntValueFromInt64(interpreter, 0)
 	}
 
-	for {
-		value := iterator.Next(interpreter)
-		if value == nil {
-			return nil
-		}
-
+	executeBody := func(value Value) (resume bool) {
 		statementResult, done := interpreter.visitForStatementBody(statement, index, value)
 		if done {
-			return statementResult
+			result = statementResult
 		}
 
+		resume = !done
+
 		if statement.Index != nil {
 			index = index.Plus(interpreter, intOne, locationRange).(IntValue)
 		}
+
+		return
 	}
+
+	iterable.ForEach(
+		interpreter,
+		forStmtTypes.ValueVariableType,
+		executeBody,
+		locationRange,
+	)
+
+	return
 }
 
 func (interpreter *Interpreter) visitForStatementBody(

runtime/interpreter/value.go

@@ -235,6 +235,12 @@ type ContractValue interface {
 type IterableValue interface {
 	Value
 	Iterator(interpreter *Interpreter) ValueIterator
+	ForEach(
+		interpreter *Interpreter,
+		elementType sema.Type,
+		function func(value Value) (resume bool),
+		locationRange LocationRange,
+	)
 }
 
 // ValueIterator is an iterator which returns values.
@@ -1586,6 +1592,25 @@ func (v *StringValue) Iterator(_ *Interpreter) ValueIterator {
 	}
 }
 
+func (v *StringValue) ForEach(
+	interpreter *Interpreter,
+	_ sema.Type,
+	function func(value Value) (resume bool),
+	_ LocationRange,
+) {
+	iterator := v.Iterator(interpreter)
+	for {
+		value := iterator.Next(interpreter)
+		if value == nil {
+			return
+		}
+
+		if !function(value) {
+			return
+		}
+	}
+}
+
 type StringValueIterator struct {
 	graphemes *uniseg.Graphemes
 }
@@ -3244,6 +3269,15 @@ func (v *ArrayValue) Map(
 	)
 }
 
+func (v *ArrayValue) ForEach(
+	interpreter *Interpreter,
+	_ sema.Type,
+	function func(value Value) (resume bool),
+	_ LocationRange,
+) {
+	v.Iterate(interpreter, function)
+}
+
 // NumberValue
 type NumberValue interface {
 	ComparableValue
@@ -19844,6 +19878,7 @@ var _ TypeIndexableValue = &StorageReferenceValue{}
 var _ MemberAccessibleValue = &StorageReferenceValue{}
 var _ AuthorizedValue = &StorageReferenceValue{}
 var _ ReferenceValue = &StorageReferenceValue{}
+var _ IterableValue = &StorageReferenceValue{}
 
 func NewUnmeteredStorageReferenceValue(
 	authorization Authorization,
@@ -20196,6 +20231,62 @@ func (*StorageReferenceValue) DeepRemove(_ *Interpreter) {
 
 func (*StorageReferenceValue) isReference() {}
 
+func (v *StorageReferenceValue) Iterator(_ *Interpreter) ValueIterator {
+	// Not used for now
+	panic(errors.NewUnreachableError())
+}
+
+func (v *StorageReferenceValue) ForEach(
+	interpreter *Interpreter,
+	elementType sema.Type,
+	function func(value Value) (resume bool),
+	locationRange LocationRange,
+) {
+	referencedValue := v.mustReferencedValue(interpreter, locationRange)
+	forEachReference(
+		interpreter,
+		referencedValue,
+		elementType,
+		function,
+		locationRange,
+	)
+}
+
+func forEachReference(
+	interpreter *Interpreter,
+	referencedValue Value,
+	elementType sema.Type,
+	function func(value Value) (resume bool),
+	locationRange LocationRange,
+) {
+	referencedIterable, ok := referencedValue.(IterableValue)
+	if !ok {
+		panic(errors.NewUnreachableError())
+	}
+
+	referenceType, isResultReference := sema.MaybeReferenceType(elementType)
+
+	updatedFunction := func(value Value) (resume bool) {
+		if isResultReference {
+			value = interpreter.getReferenceValue(value, elementType)
+		}
+
+		return function(value)
+	}
+
+	referencedElementType := elementType
+	if isResultReference {
+		referencedElementType = referenceType.Type
+	}
+
+	referencedIterable.ForEach(
+		interpreter,
+		referencedElementType,
+		updatedFunction,
+		locationRange,
+	)
+}
+
 // EphemeralReferenceValue
 
 type EphemeralReferenceValue struct {
@@ -20212,6 +20303,7 @@ var _ TypeIndexableValue = &EphemeralReferenceValue{}
 var _ MemberAccessibleValue = &EphemeralReferenceValue{}
 var _ AuthorizedValue = &EphemeralReferenceValue{}
 var _ ReferenceValue = &EphemeralReferenceValue{}
+var _ IterableValue = &EphemeralReferenceValue{}
 
 func NewUnmeteredEphemeralReferenceValue(
 	authorization Authorization,
@@ -20541,6 +20633,27 @@ func (*EphemeralReferenceValue) DeepRemove(_ *Interpreter) {
 
 func (*EphemeralReferenceValue) isReference() {}
 
+func (v *EphemeralReferenceValue) Iterator(_ *Interpreter) ValueIterator {
+	// Not used for now
+	panic(errors.NewUnreachableError())
+}
+
+func (v *EphemeralReferenceValue) ForEach(
+	interpreter *Interpreter,
+	elementType sema.Type,
+	function func(value Value) (resume bool),
+	locationRange LocationRange,
+) {
+	referencedValue := v.MustReferencedValue(interpreter, locationRange)
+	forEachReference(
+		interpreter,
+		referencedValue,
+		elementType,
+		function,
+		locationRange,
+	)
+}
+
 // AddressValue
 type AddressValue common.Address
 

runtime/sema/check_assignment.go

@@ -329,7 +329,7 @@ func (checker *Checker) visitIndexExpressionAssignment(
 	elementType = checker.visitIndexExpression(indexExpression, true)
 
 	indexExprTypes := checker.Elaboration.IndexExpressionTypes(indexExpression)
-	indexedRefType, isReference := referenceType(indexExprTypes.IndexedType)
+	indexedRefType, isReference := MaybeReferenceType(indexExprTypes.IndexedType)
 
 	if isReference &&
 		!mutableEntitledAccess.PermitsAccess(indexedRefType.Authorization) &&

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,
@@ -90,11 +66,14 @@ func (checker *Checker) VisitForStatement(statement *ast.ForStatement) (_ struct
 		checker.recordVariableDeclarationOccurrence(identifier, variable)
 	}
 
+	var indexType Type
+
 	if statement.Index != nil {
 		index := statement.Index.Identifier
+		indexType = IntType
 		indexVariable, err := checker.valueActivations.declare(variableDeclaration{
 			identifier:               index,
-			ty:                       IntType,
+			ty:                       indexType,
 			kind:                     common.DeclarationKindConstant,
 			pos:                      statement.Index.Pos,
 			isConstant:               true,
@@ -108,6 +87,11 @@ func (checker *Checker) VisitForStatement(statement *ast.ForStatement) (_ struct
 		}
 	}
 
+	checker.Elaboration.SetForStatementType(statement, ForStatementTypes{
+		IndexVariableType: indexType,
+		ValueVariableType: loopVariableType,
+	})
+
 	// The body of the loop will maybe be evaluated.
 	// That means that resource invalidations and
 	// returns are not definite, but only potential.
@@ -123,3 +107,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 checker.getReferenceType(referencedIterableElementType, false, UnauthorizedAccess)
+		}
+
+		// 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/sema/check_member_expression.go

@@ -116,14 +116,14 @@ func shouldReturnReference(parentType, memberType Type, isAssignment bool) bool
 		return false
 	}
 
-	if _, isReference := referenceType(parentType); !isReference {
+	if _, isReference := MaybeReferenceType(parentType); !isReference {
 		return false
 	}
 
 	return memberType.ContainFieldsOrElements()
 }
 
-func referenceType(typ Type) (*ReferenceType, bool) {
+func MaybeReferenceType(typ Type) (*ReferenceType, bool) {
 	unwrappedType := UnwrapOptionalType(typ)
 	refType, isReference := unwrappedType.(*ReferenceType)
 	return refType, isReference

runtime/sema/check_variable_declaration.go

@@ -264,7 +264,7 @@ func (checker *Checker) recordReference(targetVariable *Variable, expr ast.Expre
 		return
 	}
 
-	if _, isReference := referenceType(targetVariable.Type); !isReference {
+	if _, isReference := MaybeReferenceType(targetVariable.Type); !isReference {
 		return
 	}