Merge pull request #336 from onflow/fxamacker/fix-slab-size-when-rese…

…tting-same-storable-in-array

Fix slab size when resetting mutable storable in Array

array.go

@@ -609,6 +609,13 @@ func (a *ArrayDataSlab) ChildStorables() []Storable {
 	return s
 }
 
+func (a *ArrayDataSlab) getPrefixSize() uint32 {
+	if a.extraData != nil {
+		return arrayRootDataSlabPrefixSize
+	}
+	return arrayDataSlabPrefixSize
+}
+
 func (a *ArrayDataSlab) Get(_ SlabStorage, index uint64) (Storable, error) {
 	if index >= uint64(len(a.elements)) {
 		return nil, NewIndexOutOfBoundsError(index, 0, uint64(len(a.elements)))
@@ -622,7 +629,6 @@ func (a *ArrayDataSlab) Set(storage SlabStorage, address Address, index uint64,
 	}
 
 	oldElem := a.elements[index]
-	oldSize := oldElem.ByteSize()
 
 	storable, err := value.Storable(storage, address, maxInlineArrayElementSize)
 	if err != nil {
@@ -631,7 +637,16 @@ func (a *ArrayDataSlab) Set(storage SlabStorage, address Address, index uint64,
 	}
 
 	a.elements[index] = storable
-	a.header.size = a.header.size - oldSize + storable.ByteSize()
+
+	// Recompute slab size by adding all element sizes instead of using the size diff of old and new element because
+	// oldElem can be the same storable when the same value is reset and oldElem.ByteSize() can equal storable.ByteSize().
+	// Given this, size diff of the old and new element can be 0 even when its actual size changed.
+	size := a.getPrefixSize()
+	for _, e := range a.elements {
+		size += e.ByteSize()
+	}
+
+	a.header.size = size
 
 	err = storage.Store(a.header.slabID, a)
 	if err != nil {
@@ -973,7 +988,7 @@ func (a *ArrayDataSlab) PopIterate(_ SlabStorage, fn ArrayPopIterationFunc) erro
 	// Reset data slab
 	a.elements = nil
 	a.header.count = 0
-	a.header.size = arrayDataSlabPrefixSize
+	a.header.size = a.getPrefixSize()
 
 	return nil
 }

array_test.go

@@ -3666,3 +3666,52 @@ func TestArrayID(t *testing.T) {
 	require.Equal(t, sid.address[:], id[:8])
 	require.Equal(t, sid.index[:], id[8:])
 }
+
+func TestSlabSizeWhenResettingMutableStorable(t *testing.T) {
+	const (
+		arraySize           = 3
+		initialStorableSize = 1
+		mutatedStorableSize = 5
+	)
+
+	typeInfo := testTypeInfo{42}
+	storage := newTestPersistentStorage(t)
+	address := Address{1, 2, 3, 4, 5, 6, 7, 8}
+
+	array, err := NewArray(storage, address, typeInfo)
+	require.NoError(t, err)
+
+	values := make([]*mutableValue, arraySize)
+	for i := uint64(0); i < arraySize; i++ {
+		v := newMutableValue(initialStorableSize)
+		values[i] = v
+
+		err := array.Append(v)
+		require.NoError(t, err)
+	}
+
+	require.True(t, array.root.IsData())
+
+	expectedArrayRootDataSlabSize := arrayRootDataSlabPrefixSize + initialStorableSize*arraySize
+	require.Equal(t, uint32(expectedArrayRootDataSlabSize), array.root.ByteSize())
+
+	err = ValidArray(array, typeInfo, typeInfoComparator, hashInputProvider)
+	require.NoError(t, err)
+
+	for i := uint64(0); i < arraySize; i++ {
+		mv := values[i]
+		mv.updateStorableSize(mutatedStorableSize)
+
+		existingStorable, err := array.Set(i, mv)
+		require.NoError(t, err)
+		require.NotNil(t, existingStorable)
+	}
+
+	require.True(t, array.root.IsData())
+
+	expectedArrayRootDataSlabSize = arrayRootDataSlabPrefixSize + mutatedStorableSize*arraySize
+	require.Equal(t, uint32(expectedArrayRootDataSlabSize), array.root.ByteSize())
+
+	err = ValidArray(array, typeInfo, typeInfoComparator, hashInputProvider)
+	require.NoError(t, err)
+}

map.go

@@ -1374,8 +1374,6 @@ func (e *hkeyElements) Set(storage SlabStorage, address Address, b DigesterBuild
 			}
 		}
 
-		oldElemSize := elem.Size()
-
 		elem, existingValue, err := elem.Set(storage, address, b, digester, level, hkey, comparator, hip, key, value)
 		if err != nil {
 			// Don't need to wrap error as external error because err is already categorized by element.Set().
@@ -1384,7 +1382,14 @@ func (e *hkeyElements) Set(storage SlabStorage, address Address, b DigesterBuild
 
 		e.elems[equalIndex] = elem
 
-		e.size += elem.Size() - oldElemSize
+		// Recompute slab size by adding all element sizes instead of using the size diff of old and new element because
+		// oldElem can be the same storable when the same value is reset and oldElem.ByteSize() can equal storable.ByteSize().
+		// Given this, size diff of the old and new element can be 0 even when its actual size changed.
+		size := uint32(hkeyElementsPrefixSize)
+		for _, element := range e.elems {
+			size += element.Size() + digestSize
+		}
+		e.size = size
 
 		return existingValue, nil
 	}
@@ -1879,8 +1884,6 @@ func (e *singleElements) Set(storage SlabStorage, address Address, _ DigesterBui
 		if equal {
 			existingValue := elem.value
 
-			oldSize := elem.Size()
-
 			vs, err := value.Storable(storage, address, maxInlineMapValueSize(uint64(elem.key.ByteSize())))
 			if err != nil {
 				// Wrap err as external error (if needed) because err is returned by Value interface.
@@ -1890,7 +1893,14 @@ func (e *singleElements) Set(storage SlabStorage, address Address, _ DigesterBui
 			elem.value = vs
 			elem.size = singleElementPrefixSize + elem.key.ByteSize() + elem.value.ByteSize()
 
-			e.size += elem.Size() - oldSize
+			// Recompute slab size by adding all element sizes instead of using the size diff of old and new element because
+			// oldElem can be the same storable when the same value is reset and oldElem.ByteSize() can equal storable.ByteSize().
+			// Given this, size diff of the old and new element can be 0 even when its actual size changed.
+			size := uint32(singleElementsPrefixSize)
+			for _, element := range e.elems {
+				size += element.Size()
+			}
+			e.size = size
 
 			return existingValue, nil
 		}
@@ -2372,6 +2382,13 @@ func (m *MapDataSlab) ChildStorables() []Storable {
 	return elementsStorables(m.elements, nil)
 }
 
+func (m *MapDataSlab) getPrefixSize() uint32 {
+	if m.extraData != nil {
+		return mapRootDataSlabPrefixSize
+	}
+	return mapDataSlabPrefixSize
+}
+
 func elementsStorables(elems elements, childStorables []Storable) []Storable {
 
 	switch v := elems.(type) {
@@ -2436,11 +2453,7 @@ func (m *MapDataSlab) Set(storage SlabStorage, b DigesterBuilder, digester Diges
 	m.header.firstKey = m.elements.firstKey()
 
 	// Adjust header's slab size
-	if m.extraData == nil {
-		m.header.size = mapDataSlabPrefixSize + m.elements.Size()
-	} else {
-		m.header.size = mapRootDataSlabPrefixSize + m.elements.Size()
-	}
+	m.header.size = m.getPrefixSize() + m.elements.Size()
 
 	// Store modified slab
 	err = storage.Store(m.header.slabID, m)
@@ -2464,11 +2477,7 @@ func (m *MapDataSlab) Remove(storage SlabStorage, digester Digester, level uint,
 	m.header.firstKey = m.elements.firstKey()
 
 	// Adjust header's slab size
-	if m.extraData == nil {
-		m.header.size = mapDataSlabPrefixSize + m.elements.Size()
-	} else {
-		m.header.size = mapRootDataSlabPrefixSize + m.elements.Size()
-	}
+	m.header.size = m.getPrefixSize() + m.elements.Size()
 
 	// Store modified slab
 	err = storage.Store(m.header.slabID, m)
@@ -2668,7 +2677,7 @@ func (m *MapDataSlab) PopIterate(storage SlabStorage, fn MapPopIterationFunc) er
 	}
 
 	// Reset data slab
-	m.header.size = mapDataSlabPrefixSize + hkeyElementsPrefixSize
+	m.header.size = m.getPrefixSize() + hkeyElementsPrefixSize
 	m.header.firstKey = 0
 	return nil
 }

map_test.go

@@ -6817,3 +6817,59 @@ func TestMapID(t *testing.T) {
 	require.Equal(t, sid.address[:], id[:8])
 	require.Equal(t, sid.index[:], id[8:])
 }
+
+func TestSlabSizeWhenResettingMutableStorableInMap(t *testing.T) {
+	const (
+		mapSize             = 3
+		keyStringSize       = 16
+		initialStorableSize = 1
+		mutatedStorableSize = 5
+	)
+
+	keyValues := make(map[Value]*mutableValue, mapSize)
+	for i := 0; i < mapSize; i++ {
+		k := Uint64Value(i)
+		v := newMutableValue(initialStorableSize)
+		keyValues[k] = v
+	}
+
+	typeInfo := testTypeInfo{42}
+	address := Address{1, 2, 3, 4, 5, 6, 7, 8}
+	storage := newTestPersistentStorage(t)
+
+	m, err := NewMap(storage, address, newBasicDigesterBuilder(), typeInfo)
+	require.NoError(t, err)
+
+	for k, v := range keyValues {
+		existingStorable, err := m.Set(compare, hashInputProvider, k, v)
+		require.NoError(t, err)
+		require.Nil(t, existingStorable)
+	}
+
+	require.True(t, m.root.IsData())
+
+	expectedElementSize := singleElementPrefixSize + digestSize + Uint64Value(0).ByteSize() + initialStorableSize
+	expectedMapRootDataSlabSize := mapRootDataSlabPrefixSize + hkeyElementsPrefixSize + expectedElementSize*mapSize
+	require.Equal(t, expectedMapRootDataSlabSize, m.root.ByteSize())
+
+	err = ValidMap(m, typeInfo, typeInfoComparator, hashInputProvider)
+	require.NoError(t, err)
+
+	// Reset mutable values after changing its storable size
+	for k, v := range keyValues {
+		v.updateStorableSize(mutatedStorableSize)
+
+		existingStorable, err := m.Set(compare, hashInputProvider, k, v)
+		require.NoError(t, err)
+		require.NotNil(t, existingStorable)
+	}
+
+	require.True(t, m.root.IsData())
+
+	expectedElementSize = singleElementPrefixSize + digestSize + Uint64Value(0).ByteSize() + mutatedStorableSize
+	expectedMapRootDataSlabSize = mapRootDataSlabPrefixSize + hkeyElementsPrefixSize + expectedElementSize*mapSize
+	require.Equal(t, expectedMapRootDataSlabSize, m.root.ByteSize())
+
+	err = ValidMap(m, typeInfo, typeInfoComparator, hashInputProvider)
+	require.NoError(t, err)
+}

storable_test.go

@@ -676,3 +676,48 @@ func (v SomeStorable) StoredValue(storage SlabStorage) (Value, error) {
 func (v SomeStorable) String() string {
 	return fmt.Sprintf("%s", v.Storable)
 }
+
+type mutableValue struct {
+	storable *mutableStorable
+}
+
+var _ Value = &mutableValue{}
+
+func newMutableValue(storableSize uint32) *mutableValue {
+	return &mutableValue{
+		storable: &mutableStorable{
+			size: storableSize,
+		},
+	}
+}
+
+func (v *mutableValue) Storable(SlabStorage, Address, uint64) (Storable, error) {
+	return v.storable, nil
+}
+
+func (v *mutableValue) updateStorableSize(n uint32) {
+	v.storable.size = n
+}
+
+type mutableStorable struct {
+	size uint32
+}
+
+var _ Storable = &mutableStorable{}
+
+func (s *mutableStorable) ByteSize() uint32 {
+	return s.size
+}
+
+func (s *mutableStorable) StoredValue(SlabStorage) (Value, error) {
+	return &mutableValue{s}, nil
+}
+
+func (*mutableStorable) ChildStorables() []Storable {
+	return nil
+}
+
+func (*mutableStorable) Encode(*Encoder) error {
+	// no-op for testing
+	return nil
+}