rebuildnodes: Refactor uuidMap in to a separate sub-package

4 files changed, 83 insertions, 59 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
index 40ace46..9174035 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
@@ -10,18 +10,17 @@ import (
 	"github.com/datawire/dlib/dlog"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfstree"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 )
 
 type scanResult struct {
-	uuidMap
+	uuidMap   *uuidmap.UUIDMap
 	nodeGraph *graph.Graph
 }
 
@@ -46,23 +45,10 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 	}
 
 	ret := &scanResult{
-		uuidMap: uuidMap{
-			ObjID2UUID: make(map[btrfsprim.ObjID]btrfsprim.UUID),
-			UUID2ObjID: make(map[btrfsprim.UUID]btrfsprim.ObjID),
-			TreeParent: make(map[btrfsprim.ObjID]btrfsprim.UUID),
-
-			SeenTrees: make(containers.Set[btrfsprim.ObjID]),
-		},
+		uuidMap:   uuidmap.New(),
 		nodeGraph: graph.New(*sb),
 	}
 
-	// These 4 trees are mentioned directly in the superblock, so
-	// they are always seen.
-	ret.SeenTrees.Insert(btrfsprim.ROOT_TREE_OBJECTID)
-	ret.SeenTrees.Insert(btrfsprim.CHUNK_TREE_OBJECTID)
-	ret.SeenTrees.Insert(btrfsprim.TREE_LOG_OBJECTID)
-	ret.SeenTrees.Insert(btrfsprim.BLOCK_GROUP_TREE_OBJECTID)
-
 	progress(done, total)
 	for _, devResults := range scanResults {
 		for laddr := range devResults.FoundNodes {
@@ -73,35 +59,10 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 				return nil, err
 			}
 
-			// UUID map rebuilding
-			for _, item := range nodeRef.Data.BodyLeaf {
-				switch itemBody := item.Body.(type) {
-				case btrfsitem.Root:
-					if err := maybeSet("ObjID2UUID", ret.ObjID2UUID, item.Key.ObjectID, itemBody.UUID); err != nil {
-						return nil, err
-					}
-					if itemBody.UUID != (btrfsprim.UUID{}) {
-						if err := maybeSet("UUID2ObjID", ret.UUID2ObjID, itemBody.UUID, item.Key.ObjectID); err != nil {
-							return nil, err
-						}
-					}
-					if err := maybeSet("ParentUUID", ret.TreeParent, item.Key.ObjectID, itemBody.ParentUUID); err != nil {
-						return nil, err
-					}
-					ret.SeenTrees.Insert(item.Key.ObjectID)
-				case btrfsitem.UUIDMap:
-					uuid := btrfsitem.KeyToUUID(item.Key)
-					if err := maybeSet("ObjID2UUID", ret.ObjID2UUID, itemBody.ObjID, uuid); err != nil {
-						return nil, err
-					}
-					if err := maybeSet("UUID2ObjID", ret.UUID2ObjID, uuid, itemBody.ObjID); err != nil {
-						return nil, err
-					}
-				}
+			if err := ret.uuidMap.InsertNode(nodeRef); err != nil {
+				return nil, err
 			}
-			ret.SeenTrees.Insert(nodeRef.Data.Head.Owner)
 
-			// graph building
 			ret.nodeGraph.InsertNode(nodeRef)
 
 			done++
@@ -113,7 +74,7 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 		panic("should not happen")
 	}
 
-	missing := ret.missingRootItems()
+	missing := ret.uuidMap.MissingRootItems()
 	if len(missing) > 0 {
 		dlog.Errorf(ctx, "... could not find root items for %d trees: %v", len(missing), maps.SortedKeys(missing))
 	}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/util.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/util.go
index bdaa0c4..8c43dad 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/util.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/util.go
@@ -5,21 +5,11 @@
 package rebuildnodes
 
 import (
-	"fmt"
-
 	"golang.org/x/exp/constraints"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect"
 )
 
-func maybeSet[K, V comparable](name string, m map[K]V, k K, v V) error {
-	if other, conflict := m[k]; conflict && other != v {
-		return fmt.Errorf("conflict: %s %v can't have both %v and %v", name, k, other, v)
-	}
-	m[k] = v
-	return nil
-}
-
 func countNodes(nodeScanResults btrfsinspect.ScanDevicesResult) int {
 	var cnt int
 	for _, devResults := range nodeScanResults {
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap/scan.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap/scan.go
new file mode 100644
index 0000000..98ed097
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap/scan.go
@@ -0,0 +1,73 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package uuidmap
+
+import (
+	"fmt"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfstree"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
+	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
+	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
+)
+
+func maybeSet[K, V comparable](name string, m map[K]V, k K, v V) error {
+	if other, conflict := m[k]; conflict && other != v {
+		return fmt.Errorf("conflict: %s %v can't have both %v and %v", name, k, other, v)
+	}
+	m[k] = v
+	return nil
+}
+
+func New() *UUIDMap {
+	ret := &UUIDMap{
+		ObjID2UUID: make(map[btrfsprim.ObjID]btrfsprim.UUID),
+		UUID2ObjID: make(map[btrfsprim.UUID]btrfsprim.ObjID),
+		TreeParent: make(map[btrfsprim.ObjID]btrfsprim.UUID),
+
+		SeenTrees: make(containers.Set[btrfsprim.ObjID]),
+	}
+
+	// These 4 trees are mentioned directly in the superblock, so
+	// they are always seen.
+	ret.SeenTrees.Insert(btrfsprim.ROOT_TREE_OBJECTID)
+	ret.SeenTrees.Insert(btrfsprim.CHUNK_TREE_OBJECTID)
+	ret.SeenTrees.Insert(btrfsprim.TREE_LOG_OBJECTID)
+	ret.SeenTrees.Insert(btrfsprim.BLOCK_GROUP_TREE_OBJECTID)
+
+	return ret
+}
+
+func (m *UUIDMap) InsertNode(nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) error {
+	for _, item := range nodeRef.Data.BodyLeaf {
+		switch itemBody := item.Body.(type) {
+		case btrfsitem.Root:
+			if err := maybeSet("ObjID2UUID", m.ObjID2UUID, item.Key.ObjectID, itemBody.UUID); err != nil {
+				return err
+			}
+			if itemBody.UUID != (btrfsprim.UUID{}) {
+				if err := maybeSet("UUID2ObjID", m.UUID2ObjID, itemBody.UUID, item.Key.ObjectID); err != nil {
+					return err
+				}
+			}
+			if err := maybeSet("ParentUUID", m.TreeParent, item.Key.ObjectID, itemBody.ParentUUID); err != nil {
+				return err
+			}
+			m.SeenTrees.Insert(item.Key.ObjectID)
+		case btrfsitem.UUIDMap:
+			uuid := btrfsitem.KeyToUUID(item.Key)
+			if err := maybeSet("ObjID2UUID", m.ObjID2UUID, itemBody.ObjID, uuid); err != nil {
+				return err
+			}
+			if err := maybeSet("UUID2ObjID", m.UUID2ObjID, uuid, itemBody.ObjID); err != nil {
+				return err
+			}
+		}
+	}
+	m.SeenTrees.Insert(nodeRef.Data.Head.Owner)
+	return nil
+}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap/uuidmap.go
index 4d93916..32a34d1 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/uuidmap/uuidmap.go
@@ -2,14 +2,14 @@
 //
 // SPDX-License-Identifier: GPL-2.0-or-later
 
-package rebuildnodes
+package uuidmap
 
 import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 )
 
-type uuidMap struct {
+type UUIDMap struct {
 	ObjID2UUID map[btrfsprim.ObjID]btrfsprim.UUID
 	UUID2ObjID map[btrfsprim.UUID]btrfsprim.ObjID
 	TreeParent map[btrfsprim.ObjID]btrfsprim.UUID
@@ -17,7 +17,7 @@ type uuidMap struct {
 	SeenTrees containers.Set[btrfsprim.ObjID]
 }
 
-func (m uuidMap) missingRootItems() containers.Set[btrfsprim.ObjID] {
+func (m UUIDMap) MissingRootItems() containers.Set[btrfsprim.ObjID] {
 	missing := make(containers.Set[btrfsprim.ObjID])
 	for treeID := range m.SeenTrees {
 		if _, ok := m.ObjID2UUID[treeID]; !ok && treeID != btrfsprim.ROOT_TREE_OBJECTID {
@@ -32,7 +32,7 @@ func (m uuidMap) missingRootItems() containers.Set[btrfsprim.ObjID] {
 }
 
 // ParentTree implements btrfstree.NodeFile.
-func (m uuidMap) ParentTree(tree btrfsprim.ObjID) (btrfsprim.ObjID, bool) {
+func (m UUIDMap) ParentTree(tree btrfsprim.ObjID) (btrfsprim.ObjID, bool) {
 	if tree < btrfsprim.FIRST_FREE_OBJECTID || tree > btrfsprim.LAST_FREE_OBJECTID {
 		// no parent
 		return 0, true