Fix references to resource-kinded values

runtime/interpreter/interpreter.go

@@ -279,6 +279,8 @@ type Storage interface {
 	CheckHealth() error
 }
 
+type ReferencedResourceKindedValues map[atree.StorageID]map[ReferenceTrackedResourceKindedValue]struct{}
+
 type Interpreter struct {
 	Program                        *Program
 	Location                       common.Location
@@ -315,6 +317,8 @@ type Interpreter struct {
 	atreeValueValidationEnabled    bool
 	atreeStorageValidationEnabled  bool
 	tracingEnabled                 bool
+	// TODO: ideally this would be a weak map, but Go has no weak references
+	referencedResourceKindedValues ReferencedResourceKindedValues
 }
 
 type Option func(*Interpreter) error
@@ -577,6 +581,15 @@ func withTypeCodes(typeCodes TypeCodes) Option {
 	}
 }
 
+// withReferencedResourceKindedValues returns an interpreter option which sets the referenced values.
+//
+func withReferencedResourceKindedValues(referencedResourceKindedValues ReferencedResourceKindedValues) Option {
+	return func(interpreter *Interpreter) error {
+		interpreter.referencedResourceKindedValues = referencedResourceKindedValues
+		return nil
+	}
+}
+
 // WithDebugger returns an interpreter option which sets the given debugger
 //
 func WithDebugger(debugger *Debugger) Option {
@@ -615,6 +628,7 @@ func NewInterpreter(program *Program, location common.Location, options ...Optio
 			InterfaceCodes:       map[sema.TypeID]WrapperCode{},
 			TypeRequirementCodes: map[sema.TypeID]WrapperCode{},
 		}),
+		withReferencedResourceKindedValues(map[atree.StorageID]map[ReferenceTrackedResourceKindedValue]struct{}{}),
 	}
 
 	for _, option := range defaultOptions {
@@ -1230,7 +1244,7 @@ func (interpreter *Interpreter) visitFunctionBody(
 	}
 
 	// If there is a return type, declare the constant `result`.
-	// If it is a resource type, the constant has the same type as a referecne to the return type.
+	// If it is a resource type, the constant has the same type as a reference to the return type.
 	// If it is not a resource type, the constant has the same type as the return type.
 
 	if returnType != sema.VoidType {
@@ -2506,6 +2520,7 @@ func (interpreter *Interpreter) NewSubInterpreter(
 		WithAtreeValueValidationEnabled(interpreter.atreeValueValidationEnabled),
 		WithAtreeStorageValidationEnabled(interpreter.atreeStorageValidationEnabled),
 		withTypeCodes(interpreter.typeCodes),
+		withReferencedResourceKindedValues(interpreter.referencedResourceKindedValues),
 		WithPublicAccountHandler(interpreter.publicAccountHandler),
 		WithPublicKeyValidationHandler(interpreter.PublicKeyValidationHandler),
 		WithSignatureVerificationHandler(interpreter.SignatureVerificationHandler),
@@ -4354,5 +4369,34 @@ func (interpreter *Interpreter) ValidateAtreeValue(v atree.Value) {
 			}
 		}
 	}
+}
 
+func (interpreter *Interpreter) trackReferencedResourceKindedValue(
+	id atree.StorageID,
+	value ReferenceTrackedResourceKindedValue,
+) {
+	values := interpreter.referencedResourceKindedValues[id]
+	if values == nil {
+		values = map[ReferenceTrackedResourceKindedValue]struct{}{}
+		interpreter.referencedResourceKindedValues[id] = values
+	}
+	values[value] = struct{}{}
+}
+
+func (interpreter *Interpreter) updateReferencedResource(
+	currentStorageID atree.StorageID,
+	newStorageID atree.StorageID,
+	updateFunc func(value ReferenceTrackedResourceKindedValue),
+) {
+	values := interpreter.referencedResourceKindedValues[currentStorageID]
+	if values == nil {
+		return
+	}
+	for value := range values { //nolint:maprangecheck
+		updateFunc(value)
+	}
+	if newStorageID != currentStorageID {
+		interpreter.referencedResourceKindedValues[newStorageID] = values
+		interpreter.referencedResourceKindedValues[currentStorageID] = nil
+	}
 }

runtime/interpreter/interpreter_expression.go

@@ -819,6 +819,10 @@ func (interpreter *Interpreter) VisitReferenceExpression(referenceExpression *as
 
 	result := interpreter.evalExpression(referenceExpression.Expression)
 
+	if result, ok := result.(ReferenceTrackedResourceKindedValue); ok {
+		interpreter.trackReferencedResourceKindedValue(result.StorageID(), result)
+	}
+
 	return &EphemeralReferenceValue{
 		Authorized:   borrowType.Authorized,
 		Value:        result,

runtime/interpreter/value.go

@@ -167,6 +167,15 @@ func maybeDestroy(interpreter *Interpreter, getLocationRange func() LocationRang
 	resourceKindedValue.Destroy(interpreter, getLocationRange)
 }
 
+// ReferenceTrackedResourceKindedValue is a resource-kinded value
+// that must be tracked when a reference of it is taken.
+//
+type ReferenceTrackedResourceKindedValue interface {
+	ResourceKindedValue
+	IsReferenceTrackedResourceKindedValue()
+	StorageID() atree.StorageID
+}
+
 // TypeValue
 
 type TypeValue struct {
@@ -995,6 +1004,7 @@ var _ atree.Value = &ArrayValue{}
 var _ EquatableValue = &ArrayValue{}
 var _ ValueIndexableValue = &ArrayValue{}
 var _ MemberAccessibleValue = &ArrayValue{}
+var _ ReferenceTrackedResourceKindedValue = &ArrayValue{}
 
 func (*ArrayValue) IsValue() {}
 
@@ -1528,6 +1538,8 @@ func (v *ArrayValue) Storable(_ atree.SlabStorage, _ atree.Address, _ uint64) (a
 	return atree.StorageIDStorable(v.StorageID()), nil
 }
 
+func (v *ArrayValue) IsReferenceTrackedResourceKindedValue() {}
+
 func (v *ArrayValue) Transfer(
 	interpreter *Interpreter,
 	getLocationRange func() LocationRange,
@@ -1543,9 +1555,12 @@ func (v *ArrayValue) Transfer(
 		}()
 	}
 
+	currentStorageID := v.StorageID()
+	currentAddress := currentStorageID.Address
+
 	array := v.array
 
-	needsStoreTo := v.NeedsStoreTo(address)
+	needsStoreTo := address != currentAddress
 	isResourceKinded := v.IsResourceKinded(interpreter)
 
 	if needsStoreTo || !isResourceKinded {
@@ -1592,7 +1607,25 @@ func (v *ArrayValue) Transfer(
 	}
 
 	if isResourceKinded {
+		// Update the resource in-place,
+		// and also update all values that are referencing the same value
+		// (but currently point to an outdated Go instance of the value)
+
 		v.array = array
+
+		newStorageID := array.StorageID()
+		interpreter.updateReferencedResource(
+			currentStorageID,
+			newStorageID,
+			func(value ReferenceTrackedResourceKindedValue) {
+				arrayValue, ok := value.(*ArrayValue)
+				if !ok {
+					panic(errors.NewUnreachableError())
+				}
+				arrayValue.array = array
+			},
+		)
+
 		return v
 	} else {
 		return &ArrayValue{
@@ -11152,6 +11185,7 @@ var _ Value = &CompositeValue{}
 var _ EquatableValue = &CompositeValue{}
 var _ HashableValue = &CompositeValue{}
 var _ MemberAccessibleValue = &CompositeValue{}
+var _ ReferenceTrackedResourceKindedValue = &CompositeValue{}
 
 func (*CompositeValue) IsValue() {}
 
@@ -11672,6 +11706,8 @@ func (v *CompositeValue) IsResourceKinded(_ *Interpreter) bool {
 	return v.Kind == common.CompositeKindResource
 }
 
+func (v *CompositeValue) IsReferenceTrackedResourceKindedValue() {}
+
 func (v *CompositeValue) Transfer(
 	interpreter *Interpreter,
 	getLocationRange func() LocationRange,
@@ -11680,9 +11716,10 @@ func (v *CompositeValue) Transfer(
 	storable atree.Storable,
 ) Value {
 
-	dictionary := v.dictionary
+	currentStorageID := v.StorageID()
+	currentAddress := currentStorageID.Address
 
-	currentAddress := v.StorageID().Address
+	dictionary := v.dictionary
 
 	needsStoreTo := address != currentAddress
 	isResourceKinded := v.IsResourceKinded(interpreter)
@@ -11740,7 +11777,26 @@ func (v *CompositeValue) Transfer(
 
 	var res *CompositeValue
 	if isResourceKinded {
+		// Update the resource in-place,
+		// and also update all values that are referencing the same value
+		// (but currently point to an outdated Go instance of the value)
+
 		v.dictionary = dictionary
+
+		newStorageID := dictionary.StorageID()
+
+		interpreter.updateReferencedResource(
+			currentStorageID,
+			newStorageID,
+			func(value ReferenceTrackedResourceKindedValue) {
+				compositeValue, ok := value.(*CompositeValue)
+				if !ok {
+					panic(errors.NewUnreachableError())
+				}
+				compositeValue.dictionary = dictionary
+			},
+		)
+
 		res = v
 	} else {
 		res = &CompositeValue{
@@ -12011,6 +12067,7 @@ var _ atree.Value = &DictionaryValue{}
 var _ EquatableValue = &DictionaryValue{}
 var _ ValueIndexableValue = &DictionaryValue{}
 var _ MemberAccessibleValue = &DictionaryValue{}
+var _ ReferenceTrackedResourceKindedValue = &DictionaryValue{}
 
 func (*DictionaryValue) IsValue() {}
 
@@ -12575,6 +12632,8 @@ func (v *DictionaryValue) Storable(_ atree.SlabStorage, _ atree.Address, _ uint6
 	return atree.StorageIDStorable(v.StorageID()), nil
 }
 
+func (v *DictionaryValue) IsReferenceTrackedResourceKindedValue() {}
+
 func (v *DictionaryValue) Transfer(
 	interpreter *Interpreter,
 	getLocationRange func() LocationRange,
@@ -12590,9 +12649,12 @@ func (v *DictionaryValue) Transfer(
 		}()
 	}
 
+	currentStorageID := v.StorageID()
+	currentAddress := currentStorageID.Address
+
 	dictionary := v.dictionary
 
-	needsStoreTo := v.NeedsStoreTo(address)
+	needsStoreTo := address != currentAddress
 	isResourceKinded := v.IsResourceKinded(interpreter)
 
 	if needsStoreTo || !isResourceKinded {
@@ -12651,7 +12713,26 @@ func (v *DictionaryValue) Transfer(
 	}
 
 	if isResourceKinded {
+		// Update the resource in-place,
+		// and also update all values that are referencing the same value
+		// (but currently point to an outdated Go instance of the value)
+
 		v.dictionary = dictionary
+
+		newStorageID := dictionary.StorageID()
+
+		interpreter.updateReferencedResource(
+			currentStorageID,
+			newStorageID,
+			func(value ReferenceTrackedResourceKindedValue) {
+				dictionaryValue, ok := value.(*DictionaryValue)
+				if !ok {
+					panic(errors.NewUnreachableError())
+				}
+				dictionaryValue.dictionary = dictionary
+			},
+		)
+
 		return v
 	} else {
 		return &DictionaryValue{