Add feature to fix references to non-existent register

The intended use case is to enable migration programs in onflow/flow-go to fix
broken references.  As of April 2024, only 10 registers in testnet (not mainnet)
were found to have broken references and they seem to have resulted from a bug
that was fixed 2 years ago by onflow/cadence#1565.

A broken reference is a StorageID referencing a non-existent slab.

This commit adds FixLoadedBrokenReferences(), which traverses loaded slabs and
fixes broken references in maps.  To fix a map containing broken references,
new function replaces broken map with empty map having the same StorageID and
also removes all slabs in the old map.

Limitations:
- only fix broken references in map (this is intentional)
- only traverse loaded slabs in deltas and cache

IMPORTANT: This should not be used to silently fix unknown problems. It
should only be used by migration programs to fix known problems which
were determined to be appropriate to fix.

storage.go

@@ -21,6 +21,7 @@ package atree
 import (
 	"bytes"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"sort"
 	"strings"
@@ -1027,3 +1028,235 @@ func storeSlab(storage SlabStorage, slab Slab) error {
 	}
 	return nil
 }
+
+// FixLoadedBrokenReferences traverses loaded slabs and fixes broken references in maps.
+// A broken reference is a StorageID referencing a non-existent slab.
+// To fix a map containing broken references, this function replaces broken map with
+// empty map having the same StorageID and also removes all slabs in the old map.
+// Limitations:
+// - only fix broken references in map
+// - only traverse loaded slabs in deltas and cache
+// NOTE: The intended use case is to enable migration programs in onflow/flow-go to
+// fix broken references.  As of April 2024, only 10 registers in testnet (not mainnet)
+// were found to have broken references and they seem to have resulted from a bug
+// that was fixed 2 years ago by https://github.com/onflow/cadence/pull/1565.
+func (s *PersistentSlabStorage) FixLoadedBrokenReferences() ([]StorageID, error) {
+
+	// parentOf is used to find root slab from non-root slab.
+	// Broken reference can be in non-root slab, and we need StorageID of root slab
+	// to replace broken map by creating an empty new map with same StorageID.
+	parentOf := make(map[StorageID]StorageID)
+
+	getRootSlabID := func(id StorageID) StorageID {
+		for {
+			parentID, ok := parentOf[id]
+			if ok {
+				id = parentID
+			} else {
+				return id
+			}
+		}
+	}
+
+	hasBrokenReferenceInSlab := func(id StorageID, slab Slab) bool {
+		if slab == nil {
+			return false
+		}
+
+		var isMetaDataSlab bool
+
+		switch slab.(type) {
+		case *ArrayMetaDataSlab, *MapMetaDataSlab:
+			isMetaDataSlab = true
+		}
+
+		var foundBrokenRef bool
+		for _, childStorable := range slab.ChildStorables() {
+
+			storageIDStorable, ok := childStorable.(StorageIDStorable)
+			if !ok {
+				continue
+			}
+
+			childID := StorageID(storageIDStorable)
+
+			// Track parent-child relationship of root slabs and non-root slabs.
+			if isMetaDataSlab {
+				parentOf[childID] = id
+			}
+
+			if s.existIfLoaded(childID) {
+				continue
+			}
+
+			foundBrokenRef = true
+
+			if !isMetaDataSlab {
+				return true
+			}
+		}
+
+		return foundBrokenRef
+	}
+
+	var brokenStorageIDs []StorageID
+
+	// Iterate delta slabs.
+	for id, slab := range s.deltas {
+		if hasBrokenReferenceInSlab(id, slab) {
+			brokenStorageIDs = append(brokenStorageIDs, id)
+		}
+	}
+
+	// Iterate cache slabs.
+	for id, slab := range s.cache {
+		if _, ok := s.deltas[id]; ok {
+			continue
+		}
+		if hasBrokenReferenceInSlab(id, slab) {
+			brokenStorageIDs = append(brokenStorageIDs, id)
+		}
+	}
+
+	if len(brokenStorageIDs) == 0 {
+		return nil, nil
+	}
+
+	rootSlabStorageIDsWithBrokenData := make(map[StorageID]struct{})
+	var errs []error
+
+	// Find StorageIDs of root slab for slabs containing broken references.
+	for _, id := range brokenStorageIDs {
+		rootID := getRootSlabID(id)
+		if rootID == StorageIDUndefined {
+			errs = append(errs, fmt.Errorf("failed to get root slab id for slab %s", id))
+			continue
+		}
+		rootSlabStorageIDsWithBrokenData[rootID] = struct{}{}
+	}
+
+	for rootSlabID := range rootSlabStorageIDsWithBrokenData {
+		rootSlab := s.RetrieveIfLoaded(rootSlabID)
+		if rootSlab == nil {
+			errs = append(errs, fmt.Errorf("failed to retrieve loaded root slab %s", rootSlabID))
+			continue
+		}
+
+		switch rootSlab := rootSlab.(type) {
+		case MapSlab:
+			if rootSlab.ExtraData() == nil {
+				errs = append(errs, fmt.Errorf("failed to fix broken references because slab %s isn't root slab", rootSlab.ID()))
+				continue
+			}
+
+			err := s.fixBrokenReferencesInMap(rootSlab)
+			if err != nil {
+				errs = append(errs, err)
+			}
+
+		default:
+			// IMPORTANT: Only handle map slabs for now.  DO NOT silently fix currently unknown problems.
+			errs = append(errs, fmt.Errorf("failed to fix broken references in non-map slab %s (%T)", rootSlab.ID(), rootSlab))
+		}
+	}
+
+	return brokenStorageIDs, errors.Join(errs...)
+}
+
+// fixBrokenReferencesInMap replaces replaces broken map with empty map
+// having the same StorageID and also removes all slabs in the old map.
+func (s *PersistentSlabStorage) fixBrokenReferencesInMap(old MapSlab) error {
+	id := old.ID()
+
+	// Create an empty map with the same StorgeID, type, and seed as the old map.
+	new := &MapDataSlab{
+		header: MapSlabHeader{
+			id:   id,
+			size: mapRootDataSlabPrefixSize + hkeyElementsPrefixSize,
+		},
+		extraData: &MapExtraData{
+			TypeInfo: old.ExtraData().TypeInfo,
+			Seed:     old.ExtraData().Seed,
+		},
+		elements: newHkeyElements(0),
+	}
+
+	// Store new empty map with the same StorageID.
+	err := s.Store(id, new)
+	if err != nil {
+		return err
+	}
+
+	// Remove all slabs and references in old map.
+	references, _, err := s.getAllChildReferences(old)
+	if err != nil {
+		return err
+	}
+
+	for _, childID := range references {
+		err = s.Remove(childID)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (s *PersistentSlabStorage) existIfLoaded(id StorageID) bool {
+	// Check deltas.
+	if slab, ok := s.deltas[id]; ok {
+		return slab != nil
+	}
+
+	// Check read cache.
+	if slab, ok := s.cache[id]; ok {
+		return slab != nil
+	}
+
+	return false
+}
+
+// getAllChildReferences returns child references of given slab (all levels).
+func (s *PersistentSlabStorage) getAllChildReferences(slab Slab) (
+	references []StorageID,
+	brokenReferences []StorageID,
+	err error,
+) {
+	childStorables := slab.ChildStorables()
+
+	for len(childStorables) > 0 {
+
+		var nextChildStorables []Storable
+
+		for _, childStorable := range childStorables {
+
+			storageIDStorable, ok := childStorable.(StorageIDStorable)
+			if !ok {
+				continue
+			}
+
+			childID := StorageID(storageIDStorable)
+
+			childSlab, ok, err := s.Retrieve(childID)
+			if err != nil {
+				return nil, nil, err
+			}
+			if !ok {
+				brokenReferences = append(brokenReferences, childID)
+				continue
+			}
+
+			references = append(references, childID)
+
+			nextChildStorables = append(
+				nextChildStorables,
+				childSlab.ChildStorables()...,
+			)
+		}
+
+		childStorables = nextChildStorables
+	}
+
+	return references, brokenReferences, nil
+}