Merge pull request #491 from onflow/bastian/port-internal-fixes

Port updates from atree-internal

array.go

@@ -2772,11 +2772,20 @@ func (a *Array) setParentUpdater(f parentUpdater) {
 // setCallbackWithChild sets up callback function with child value (child)
 // so parent array (a) can be notified when child value is modified.
 func (a *Array) setCallbackWithChild(i uint64, child Value, maxInlineSize uint64) {
-	c, ok := child.(mutableValueNotifier)
+	// Unwrap child value if needed (e.g. interpreter.SomeValue)
+	unwrappedChild, wrapperSize := unwrapValue(child)
+
+	c, ok := unwrappedChild.(mutableValueNotifier)
 	if !ok {
 		return
 	}
 
+	if maxInlineSize < wrapperSize {
+		maxInlineSize = 0
+	} else {
+		maxInlineSize -= wrapperSize
+	}
+
 	vid := c.ValueID()
 
 	// mutableElementIndex is lazily initialized.
@@ -2809,6 +2818,8 @@ func (a *Array) setCallbackWithChild(i uint64, child Value, maxInlineSize uint64
 			return false, err
 		}
 
+		storable = unwrapStorable(storable)
+
 		// Verify retrieved element is either SlabIDStorable or Slab, with identical value ID.
 		switch storable := storable.(type) {
 		case SlabIDStorable:
@@ -2827,15 +2838,19 @@ func (a *Array) setCallbackWithChild(i uint64, child Value, maxInlineSize uint64
 			return false, nil
 		}
 
+		// NOTE: Must reset child using original child (not unwrapped child)
+
 		// Set child value with parent array using updated index.
-		// Set() calls c.Storable() which returns inlined or not-inlined child storable.
-		existingValueStorable, err := a.set(adjustedIndex, c)
+		// Set() calls child.Storable() which returns inlined or not-inlined child storable.
+		existingValueStorable, err := a.set(adjustedIndex, child)
 		if err != nil {
 			return false, err
 		}
 
 		// Verify overwritten storable has identical value ID.
 
+		existingValueStorable = unwrapStorable(existingValueStorable)
+
 		switch existingValueStorable := existingValueStorable.(type) {
 		case SlabIDStorable:
 			sid := SlabID(existingValueStorable)
@@ -2923,38 +2938,87 @@ func (a *Array) Set(index uint64, value Value) (Storable, error) {
 	// If overwritten storable is an inlined slab, uninline the slab and store it in storage.
 	// This is to prevent potential data loss because the overwritten inlined slab was not in
 	// storage and any future changes to it would have been lost.
-	switch s := existingStorable.(type) {
+	existingStorable, existingValueID, _, err = uninlineStorableIfNeeded(a.Storage, existingStorable)
+	if err != nil {
+		return nil, err
+	}
+
+	// Remove overwritten array/map's ValueID from mutableElementIndex if:
+	// - new value isn't array/map, or
+	// - new value is array/map with different value ID
+	if existingValueID != emptyValueID {
+		unwrappedValue, _ := unwrapValue(value)
+		newValue, ok := unwrappedValue.(mutableValueNotifier)
+		if !ok || existingValueID != newValue.ValueID() {
+			delete(a.mutableElementIndex, existingValueID)
+		}
+	}
+
+	return existingStorable, nil
+}
+
+// uninlineStorableIfNeeded uninlines given storable if needed, and
+// returns uninlined Storable and its ValueID.
+// If given storable is a WrapperStorable, this function uninlines
+// wrapped storable if needed and returns a new WrapperStorable
+// with wrapped uninlined storable and its ValidID.
+func uninlineStorableIfNeeded(storage SlabStorage, storable Storable) (Storable, ValueID, bool, error) {
+	if storable == nil {
+		return storable, emptyValueID, false, nil
+	}
+
+	switch s := storable.(type) {
 	case ArraySlab: // inlined array slab
-		err = s.Uninline(a.Storage)
+		err := s.Uninline(storage)
 		if err != nil {
-			return nil, err
+			return nil, emptyValueID, false, err
 		}
-		existingStorable = SlabIDStorable(s.SlabID())
-		existingValueID = slabIDToValueID(s.SlabID())
+
+		slabID := s.SlabID()
+
+		newStorable := SlabIDStorable(slabID)
+		valueID := slabIDToValueID(slabID)
+
+		return newStorable, valueID, true, nil
 
 	case MapSlab: // inlined map slab
-		err = s.Uninline(a.Storage)
+		err := s.Uninline(storage)
 		if err != nil {
-			return nil, err
+			return nil, emptyValueID, false, err
 		}
-		existingStorable = SlabIDStorable(s.SlabID())
-		existingValueID = slabIDToValueID(s.SlabID())
+
+		slabID := s.SlabID()
+
+		newStorable := SlabIDStorable(slabID)
+		valueID := slabIDToValueID(slabID)
+
+		return newStorable, valueID, true, nil
 
 	case SlabIDStorable: // uninlined slab
-		existingValueID = slabIDToValueID(SlabID(s))
-	}
+		valueID := slabIDToValueID(SlabID(s))
 
-	// Remove overwritten array/map's ValueID from mutableElementIndex if:
-	// - new value isn't array/map, or
-	// - new value is array/map with different value ID
-	if existingValueID != emptyValueID {
-		newValue, ok := value.(mutableValueNotifier)
-		if !ok || existingValueID != newValue.ValueID() {
-			delete(a.mutableElementIndex, existingValueID)
+		return storable, valueID, false, nil
+
+	case WrapperStorable:
+		unwrappedStorable := unwrapStorable(s)
+
+		// Uninline wrapped storable if needed.
+		uninlinedWrappedStorable, valueID, uninlined, err := uninlineStorableIfNeeded(storage, unwrappedStorable)
+		if err != nil {
+			return nil, emptyValueID, false, err
+		}
+
+		if !uninlined {
+			return storable, valueID, uninlined, nil
 		}
+
+		// Create a new WrapperStorable with uninlinedWrappedStorable
+		newStorable := s.WrapAtreeStorable(uninlinedWrappedStorable)
+
+		return newStorable, valueID, uninlined, nil
 	}
 
-	return existingStorable, nil
+	return storable, emptyValueID, false, nil
 }
 
 func (a *Array) set(index uint64, value Value) (Storable, error) {
@@ -3068,39 +3132,20 @@ func (a *Array) Remove(index uint64) (Storable, error) {
 		return nil, err
 	}
 
-	// If overwritten storable is an inlined slab, uninline the slab and store it in storage.
+	// If removed storable is an inlined slab, uninline the slab and store it in storage.
 	// This is to prevent potential data loss because the overwritten inlined slab was not in
 	// storage and any future changes to it would have been lost.
-	switch s := storable.(type) {
-	case ArraySlab:
-		err = s.Uninline(a.Storage)
-		if err != nil {
-			return nil, err
-		}
-		storable = SlabIDStorable(s.SlabID())
-
-		// Delete removed element ValueID from mutableElementIndex
-		removedValueID := slabIDToValueID(s.SlabID())
-		delete(a.mutableElementIndex, removedValueID)
-
-	case MapSlab:
-		err = s.Uninline(a.Storage)
-		if err != nil {
-			return nil, err
-		}
-		storable = SlabIDStorable(s.SlabID())
-
-		// Delete removed element ValueID from mutableElementIndex
-		removedValueID := slabIDToValueID(s.SlabID())
-		delete(a.mutableElementIndex, removedValueID)
+	removedStorable, removedValueID, _, err := uninlineStorableIfNeeded(a.Storage, storable)
+	if err != nil {
+		return nil, err
+	}
 
-	case SlabIDStorable:
-		// Delete removed element ValueID from mutableElementIndex
-		removedValueID := slabIDToValueID(SlabID(s))
+	// Delete removed element ValueID from mutableElementIndex
+	if removedValueID != emptyValueID {
 		delete(a.mutableElementIndex, removedValueID)
 	}
 
-	return storable, nil
+	return removedStorable, nil
 }
 
 func (a *Array) remove(index uint64) (Storable, error) {
@@ -3361,7 +3406,10 @@ var defaultReadOnlyArrayIteratorMutatinCallback ReadOnlyArrayIteratorMutationCal
 var _ ArrayIterator = &readOnlyArrayIterator{}
 
 func (i *readOnlyArrayIterator) setMutationCallback(value Value) {
-	if v, ok := value.(mutableValueNotifier); ok {
+
+	unwrappedChild, _ := unwrapValue(value)
+
+	if v, ok := unwrappedChild.(mutableValueNotifier); ok {
 		v.setParentUpdater(func() (found bool, err error) {
 			i.valueMutationCallback(value)
 			return true, NewReadOnlyIteratorElementMutationError(i.array.ValueID(), v.ValueID())

array_debug.go

@@ -701,41 +701,62 @@ func (v *serializationVerifier) arrayDataSlabEqual(expected, actual *ArrayDataSl
 		ee := expected.elements[i]
 		ae := actual.elements[i]
 
-		switch ee := ee.(type) {
+		err := v.compareStorable(ee, ae)
+		if err != nil {
+			return NewFatalError(fmt.Errorf("failed to compare element %d: %s", i, err))
+		}
+	}
 
-		case SlabIDStorable: // Compare not-inlined element
-			if !v.compare(ee, ae) {
-				return NewFatalError(fmt.Errorf("element %d %+v is wrong, want %+v", i, ae, ee))
-			}
+	return nil
+}
 
-			ev, err := ee.StoredValue(v.storage)
-			if err != nil {
-				// Don't need to wrap error as external error because err is already categorized by SlabIDStorable.StoredValue().
-				return err
-			}
+func (v *serializationVerifier) compareStorable(expected, actual Storable) error {
 
-			return v.verifyValue(ev)
+	switch expected := expected.(type) {
 
-		case *ArrayDataSlab: // Compare inlined array
-			ae, ok := ae.(*ArrayDataSlab)
-			if !ok {
-				return NewFatalError(fmt.Errorf("expect element as inlined *ArrayDataSlab, actual %T", ae))
-			}
+	case SlabIDStorable: // Compare not-inlined element
+		if !v.compare(expected, actual) {
+			return NewFatalError(fmt.Errorf("failed to compare SlabIDStorable: %+v is wrong, want %+v", actual, expected))
+		}
 
-			return v.arrayDataSlabEqual(ee, ae)
+		actualValue, err := actual.StoredValue(v.storage)
+		if err != nil {
+			// Don't need to wrap error as external error because err is already categorized by SlabIDStorable.StoredValue().
+			return err
+		}
 
-		case *MapDataSlab: // Compare inlined map
-			ae, ok := ae.(*MapDataSlab)
-			if !ok {
-				return NewFatalError(fmt.Errorf("expect element as inlined *MapDataSlab, actual %T", ae))
-			}
+		return v.verifyValue(actualValue)
 
-			return v.mapDataSlabEqual(ee, ae)
+	case *ArrayDataSlab: // Compare inlined array
+		actual, ok := actual.(*ArrayDataSlab)
+		if !ok {
+			return NewFatalError(fmt.Errorf("expect storable as inlined *ArrayDataSlab, actual %T", actual))
+		}
 
-		default:
-			if !v.compare(ee, ae) {
-				return NewFatalError(fmt.Errorf("element %d %+v is wrong, want %+v", i, ae, ee))
-			}
+		return v.arrayDataSlabEqual(expected, actual)
+
+	case *MapDataSlab: // Compare inlined map
+		actual, ok := actual.(*MapDataSlab)
+		if !ok {
+			return NewFatalError(fmt.Errorf("expect storable as inlined *MapDataSlab, actual %T", actual))
+		}
+
+		return v.mapDataSlabEqual(expected, actual)
+
+	case WrapperStorable: // Compare wrapper storable
+		actual, ok := actual.(WrapperStorable)
+		if !ok {
+			return NewFatalError(fmt.Errorf("expect storable as WrapperStorable, actual %T", actual))
+		}
+
+		unwrappedExpected := expected.UnwrapAtreeStorable()
+		unwrappedActual := actual.UnwrapAtreeStorable()
+
+		return v.compareStorable(unwrappedExpected, unwrappedActual)
+
+	default:
+		if !v.compare(expected, actual) {
+			return NewFatalError(fmt.Errorf("%+v is wrong, want %+v", actual, expected))
 		}
 	}