Use newly added util funcs to simplify slab ops

This commit uses slice util functions to move elements in slab operations, such as Split(), Merge(), LendToRight(), and BorrowFromRight(). This simplifies code and reduces risk of problems such as memory leak caused by not freeing elements.
onflow · Mar 4, 2025 · bdd2ecb · bdd2ecb
1 parent 1fa0272
commit bdd2ecb
Show file tree

Hide file tree

Showing 4 changed files with 103 additions and 116 deletions.
diff --git a/array_data_slab.go b/array_data_slab.go
@@ -187,7 +187,12 @@ func (a *ArrayDataSlab) Split(storage SlabStorage) (Slab, Slab, error) {
 		leftSize += elemSize
 	}
 
-	// Construct right slab
+	// Split elements
+	var rightElements []Storable
+	a.elements, rightElements = split(a.elements, leftCount)
+
+	// Create right slab
+
 	sID, err := storage.GenerateSlabID(a.header.slabID.address)
 	if err != nil {
 		// Wrap err as external error (if needed) because err is returned by SlabStorage interface.
@@ -199,20 +204,18 @@ func (a *ArrayDataSlab) Split(storage SlabStorage) (Slab, Slab, error) {
 			),
 		)
 	}
-	rightSlabCount := len(a.elements) - leftCount
+
 	rightSlab := &ArrayDataSlab{
 		header: ArraySlabHeader{
 			slabID: sID,
 			size:   arrayDataSlabPrefixSize + dataSize - leftSize,
-			count:  uint32(rightSlabCount),
+			count:  uint32(len(rightElements)),
 		},
-		next: a.next,
+		next:     a.next,
+		elements: rightElements,
 	}
 
-	rightSlab.elements = slices.Clone(a.elements[leftCount:])
-
 	// Modify left (original) slab
-	a.elements = slices.Delete(a.elements, leftCount, len(a.elements))
 	a.header.size = arrayDataSlabPrefixSize + leftSize
 	a.header.count = uint32(leftCount)
 	a.next = rightSlab.header.slabID
@@ -222,7 +225,7 @@ func (a *ArrayDataSlab) Split(storage SlabStorage) (Slab, Slab, error) {
 
 func (a *ArrayDataSlab) Merge(slab Slab) error {
 	rightSlab := slab.(*ArrayDataSlab)
-	a.elements = append(a.elements, rightSlab.elements...)
+	a.elements = merge(a.elements, rightSlab.elements)
 	a.header.size = a.header.size + rightSlab.header.size - arrayDataSlabPrefixSize
 	a.header.count += rightSlab.header.count
 	a.next = rightSlab.next
@@ -237,6 +240,7 @@ func (a *ArrayDataSlab) LendToRight(slab Slab) error {
 	count := a.header.count + rightSlab.header.count
 	size := a.header.size + rightSlab.header.size
 
+	oldLeftCount := a.header.count
 	leftCount := a.header.count
 	leftSize := a.header.size
 
@@ -252,22 +256,15 @@ func (a *ArrayDataSlab) LendToRight(slab Slab) error {
 		leftCount--
 	}
 
-	// Update the right slab
-
-	// Prepend elements from the left slab to the right slab
-	rightSlab.elements = slices.Insert(
-		rightSlab.elements,
-		0,
-		a.elements[leftCount:]...,
-	)
+	// Move elements
+	moveCount := oldLeftCount - leftCount
+	a.elements, rightSlab.elements = lendToRight(a.elements, rightSlab.elements, int(moveCount))
 
+	// Update right slab
 	rightSlab.header.size = size - leftSize
 	rightSlab.header.count = count - leftCount
 
 	// Update left slab
-	// NOTE: clear to prevent memory leak
-	clear(a.elements[leftCount:])
-	a.elements = a.elements[:leftCount]
 	a.header.size = leftSize
 	a.header.count = leftCount
 
@@ -281,6 +278,7 @@ func (a *ArrayDataSlab) BorrowFromRight(slab Slab) error {
 	count := a.header.count + rightSlab.header.count
 	size := a.header.size + rightSlab.header.size
 
+	oldLeftCount := a.header.count
 	leftCount := a.header.count
 	leftSize := a.header.size
 
@@ -300,20 +298,15 @@ func (a *ArrayDataSlab) BorrowFromRight(slab Slab) error {
 		leftCount++
 	}
 
-	rightStartIndex := leftCount - a.header.count
+	// Move elements
+	moveCount := leftCount - oldLeftCount
+	a.elements, rightSlab.elements = borrowFromRight(a.elements, rightSlab.elements, int(moveCount))
 
 	// Update left slab
-	a.elements = append(a.elements, rightSlab.elements[:rightStartIndex]...)
 	a.header.size = leftSize
 	a.header.count = leftCount
 
 	// Update right slab
-	// TODO: copy elements to front instead?
-	// NOTE: prevent memory leak
-	for i := range rightStartIndex {
-		rightSlab.elements[i] = nil
-	}
-	rightSlab.elements = rightSlab.elements[rightStartIndex:]
 	rightSlab.header.size = size - leftSize
 	rightSlab.header.count = count - leftCount
 

diff --git a/array_metadata_slab.go b/array_metadata_slab.go
@@ -631,7 +631,7 @@ func (a *ArrayMetaDataSlab) Merge(slab Slab) error {
 	leftSlabChildrenCount := len(a.childrenHeaders)
 
 	rightSlab := slab.(*ArrayMetaDataSlab)
-	a.childrenHeaders = append(a.childrenHeaders, rightSlab.childrenHeaders...)
+	a.childrenHeaders = merge(a.childrenHeaders, rightSlab.childrenHeaders)
 	a.header.size += rightSlab.header.size - arrayMetaDataSlabPrefixSize
 	a.header.count += rightSlab.header.count
 
@@ -661,7 +661,11 @@ func (a *ArrayMetaDataSlab) Split(storage SlabStorage) (Slab, Slab, error) {
 		leftCount += a.childrenHeaders[i].count
 	}
 
-	// Construct right slab
+	// Split childrenHeaders
+	var rightChildrenHeaders []ArraySlabHeader
+	a.childrenHeaders, rightChildrenHeaders = split(a.childrenHeaders, leftChildrenCount)
+
+	// Create right slab
 	sID, err := storage.GenerateSlabID(a.header.slabID.address)
 	if err != nil {
 		// Wrap err as external error (if needed) because err is returned by SlabStorage interface.
@@ -676,10 +680,9 @@ func (a *ArrayMetaDataSlab) Split(storage SlabStorage) (Slab, Slab, error) {
 			size:   a.header.size - uint32(leftSize),
 			count:  a.header.count - leftCount,
 		},
+		childrenHeaders: rightChildrenHeaders,
 	}
 
-	rightSlab.childrenHeaders = slices.Clone(a.childrenHeaders[leftChildrenCount:])
-
 	rightSlab.childrenCountSum = make([]uint32, len(rightSlab.childrenHeaders))
 	countSum := uint32(0)
 	for i := range rightSlab.childrenCountSum {
@@ -688,7 +691,6 @@ func (a *ArrayMetaDataSlab) Split(storage SlabStorage) (Slab, Slab, error) {
 	}
 
 	// Modify left (original)slab
-	a.childrenHeaders = a.childrenHeaders[:leftChildrenCount]
 	a.childrenCountSum = a.childrenCountSum[:leftChildrenCount]
 	a.header.count = leftCount
 	a.header.size = arrayMetaDataSlabPrefixSize + uint32(leftSize)
@@ -699,15 +701,13 @@ func (a *ArrayMetaDataSlab) Split(storage SlabStorage) (Slab, Slab, error) {
 func (a *ArrayMetaDataSlab) LendToRight(slab Slab) error {
 	rightSlab := slab.(*ArrayMetaDataSlab)
 
-	childrenHeadersLen := len(a.childrenHeaders) + len(rightSlab.childrenHeaders)
-	leftChildrenHeadersLen := childrenHeadersLen / 2
+	oldLeftChildrenHeaderCount := len(a.childrenHeaders)
+	childrenHeaderCount := len(a.childrenHeaders) + len(rightSlab.childrenHeaders)
+	leftChildrenHeaderCount := childrenHeaderCount / 2
 
-	// Prepend child headers from the left slab to the right slab headers
-	rightSlab.childrenHeaders = slices.Insert(
-		rightSlab.childrenHeaders,
-		0,
-		a.childrenHeaders[leftChildrenHeadersLen:]...,
-	)
+	// Move elements
+	moveCount := oldLeftChildrenHeaderCount - leftChildrenHeaderCount
+	a.childrenHeaders, rightSlab.childrenHeaders = lendToRight(a.childrenHeaders, rightSlab.childrenHeaders, moveCount)
 
 	// Rebuild right slab childrenCountSum
 	rightSlab.childrenCountSum = make([]uint32, len(rightSlab.childrenHeaders))
@@ -725,41 +725,40 @@ func (a *ArrayMetaDataSlab) LendToRight(slab Slab) error {
 	rightSlab.header.size = arrayMetaDataSlabPrefixSize + uint32(len(rightSlab.childrenHeaders))*arraySlabHeaderSize
 
 	// Update left slab (original)
-	a.childrenHeaders = a.childrenHeaders[:leftChildrenHeadersLen]
-	a.childrenCountSum = a.childrenCountSum[:leftChildrenHeadersLen]
+	a.childrenCountSum = a.childrenCountSum[:leftChildrenHeaderCount]
 
 	a.header.count = 0
 	for i := range a.childrenHeaders {
 		a.header.count += a.childrenHeaders[i].count
 	}
-	a.header.size = arrayMetaDataSlabPrefixSize + uint32(leftChildrenHeadersLen)*arraySlabHeaderSize
+	a.header.size = arrayMetaDataSlabPrefixSize + uint32(leftChildrenHeaderCount)*arraySlabHeaderSize
 
 	return nil
 }
 
 func (a *ArrayMetaDataSlab) BorrowFromRight(slab Slab) error {
-	originalLeftSlabCountSum := a.header.count
-	originalLeftSlabHeaderLen := len(a.childrenHeaders)
+	oldLeftSlabCountSum := a.header.count
+	oldLeftChildrenHeaderCount := len(a.childrenHeaders)
 
 	rightSlab := slab.(*ArrayMetaDataSlab)
 
-	childrenHeadersLen := len(a.childrenHeaders) + len(rightSlab.childrenHeaders)
-	leftSlabHeaderLen := childrenHeadersLen / 2
-	rightSlabHeaderLen := childrenHeadersLen - leftSlabHeaderLen
+	childrenHeaderCount := len(a.childrenHeaders) + len(rightSlab.childrenHeaders)
+	leftChildrenHeaderCount := childrenHeaderCount / 2
 
-	// Update left slab (original)
-	a.childrenHeaders = append(a.childrenHeaders, rightSlab.childrenHeaders[:leftSlabHeaderLen-len(a.childrenHeaders)]...)
+	// Move elements
+	moveCount := leftChildrenHeaderCount - oldLeftChildrenHeaderCount
+	a.childrenHeaders, rightSlab.childrenHeaders = borrowFromRight(a.childrenHeaders, rightSlab.childrenHeaders, moveCount)
 
-	countSum := originalLeftSlabCountSum
-	for i := originalLeftSlabHeaderLen; i < len(a.childrenHeaders); i++ {
+	// Update left slab (original)
+	countSum := oldLeftSlabCountSum
+	for i := oldLeftChildrenHeaderCount; i < len(a.childrenHeaders); i++ {
 		countSum += a.childrenHeaders[i].count
 		a.childrenCountSum = append(a.childrenCountSum, countSum)
 	}
 	a.header.count = countSum
-	a.header.size = arrayMetaDataSlabPrefixSize + uint32(leftSlabHeaderLen)*arraySlabHeaderSize
+	a.header.size = arrayMetaDataSlabPrefixSize + uint32(leftChildrenHeaderCount)*arraySlabHeaderSize
 
 	// Update right slab
-	rightSlab.childrenHeaders = rightSlab.childrenHeaders[len(rightSlab.childrenHeaders)-rightSlabHeaderLen:]
 	rightSlab.childrenCountSum = rightSlab.childrenCountSum[:len(rightSlab.childrenHeaders)]
 
 	countSum = uint32(0)
@@ -768,7 +767,7 @@ func (a *ArrayMetaDataSlab) BorrowFromRight(slab Slab) error {
 		rightSlab.childrenCountSum[i] = countSum
 	}
 	rightSlab.header.count = countSum
-	rightSlab.header.size = arrayMetaDataSlabPrefixSize + uint32(rightSlabHeaderLen)*arraySlabHeaderSize
+	rightSlab.header.size = arrayMetaDataSlabPrefixSize + uint32(len(rightSlab.childrenHeaders))*arraySlabHeaderSize
 
 	return nil
 }

diff --git a/map_elements_hashkey.go b/map_elements_hashkey.go
@@ -427,13 +427,10 @@ func (e *hkeyElements) Merge(elems elements) error {
 		return NewSlabMergeError(fmt.Errorf("cannot merge elements of different types (%T, %T)", e, elems))
 	}
 
-	e.hkeys = append(e.hkeys, rElems.hkeys...)
-	e.elems = append(e.elems, rElems.elems...)
+	e.hkeys = merge(e.hkeys, rElems.hkeys)
+	e.elems = merge(e.elems, rElems.elems)
 	e.size += rElems.Size() - hkeyElementsPrefixSize
 
-	// Set merged elements to nil to prevent memory leak.
-	clear(rElems.elems)
-
 	return nil
 }
 
@@ -461,19 +458,22 @@ func (e *hkeyElements) Split() (elements, elements, error) {
 		leftSize += elemSize
 	}
 
-	// Create right slab elements
-	rightElements := &hkeyElements{level: e.level}
-
-	rightElements.hkeys = slices.Clone(e.hkeys[leftCount:])
+	// Split elements
+	var rightKeys []Digest
+	e.hkeys, rightKeys = split(e.hkeys, leftCount)
 
-	rightElements.elems = slices.Clone(e.elems[leftCount:])
+	var rightElems []element
+	e.elems, rightElems = split(e.elems, leftCount)
 
-	rightElements.size = dataSize - leftSize + hkeyElementsPrefixSize
+	// Create right slab
+	rightElements := &hkeyElements{
+		level: e.level,
+		hkeys: rightKeys,
+		elems: rightElems,
+		size:  dataSize - leftSize + hkeyElementsPrefixSize,
+	}
 
-	e.hkeys = e.hkeys[:leftCount]
-	// NOTE: prevent memory leak
-	clear(e.elems[leftCount:])
-	e.elems = e.elems[:leftCount]
+	// Update left slab
 	e.size = hkeyElementsPrefixSize + leftSize
 
 	return e, rightElements, nil
@@ -494,6 +494,7 @@ func (e *hkeyElements) LendToRight(re elements) error {
 
 	size := e.Size() + rightElements.Size() - hkeyElementsPrefixSize*2
 
+	oldLeftCount := len(e.elems)
 	leftCount := len(e.elems)
 	leftSize := e.Size() - hkeyElementsPrefixSize
 
@@ -509,16 +510,15 @@ func (e *hkeyElements) LendToRight(re elements) error {
 		leftCount--
 	}
 
-	// Update the right elements
-	rightElements.hkeys = slices.Insert(rightElements.hkeys, 0, e.hkeys[leftCount:]...)
-	rightElements.elems = slices.Insert(rightElements.elems, 0, e.elems[leftCount:]...)
+	// Move elements
+	moveCount := oldLeftCount - leftCount
+	e.hkeys, rightElements.hkeys = lendToRight(e.hkeys, rightElements.hkeys, moveCount)
+	e.elems, rightElements.elems = lendToRight(e.elems, rightElements.elems, moveCount)
+
+	// Update right slab
 	rightElements.size = size - leftSize + hkeyElementsPrefixSize
 
 	// Update left slab
-	// NOTE: prevent memory leak
-	clear(e.elems[leftCount:])
-	e.hkeys = e.hkeys[:leftCount]
-	e.elems = e.elems[:leftCount]
 	e.size = hkeyElementsPrefixSize + leftSize
 
 	return nil
@@ -539,6 +539,7 @@ func (e *hkeyElements) BorrowFromRight(re elements) error {
 
 	size := e.Size() + rightElements.Size() - hkeyElementsPrefixSize*2
 
+	oldLeftCount := len(e.elems)
 	leftCount := len(e.elems)
 	leftSize := e.Size() - hkeyElementsPrefixSize
 
@@ -558,21 +559,15 @@ func (e *hkeyElements) BorrowFromRight(re elements) error {
 		leftCount++
 	}
 
-	rightStartIndex := leftCount - len(e.elems)
+	// Move elements
+	moveCount := leftCount - oldLeftCount
+	e.hkeys, rightElements.hkeys = borrowFromRight(e.hkeys, rightElements.hkeys, moveCount)
+	e.elems, rightElements.elems = borrowFromRight(e.elems, rightElements.elems, moveCount)
 
-	// Update left elements
-	e.hkeys = append(e.hkeys, rightElements.hkeys[:rightStartIndex]...)
-	e.elems = append(e.elems, rightElements.elems[:rightStartIndex]...)
+	// Update left slab
 	e.size = leftSize + hkeyElementsPrefixSize
 
 	// Update right slab
-	// TODO: copy elements to front instead?
-	// NOTE: prevent memory leak
-	for i := range rightStartIndex {
-		rightElements.elems[i] = nil
-	}
-	rightElements.hkeys = rightElements.hkeys[rightStartIndex:]
-	rightElements.elems = rightElements.elems[rightStartIndex:]
 	rightElements.size = size - leftSize + hkeyElementsPrefixSize
 
 	return nil