rebuildnodes: Refactor the node-graph stuff in to a separate sub-package

6 files changed, 337 insertions, 282 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go
new file mode 100644
index 0000000..7ae3b89
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/graph.go
@@ -0,0 +1,157 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package graph
+
+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"
+)
+
+type Node struct {
+	Level      uint8
+	Generation btrfsprim.Generation
+	Owner      btrfsprim.ObjID
+	NumItems   uint32
+	MinItem    btrfsprim.Key
+	MaxItem    btrfsprim.Key
+}
+
+type Edge struct {
+	FromTree btrfsprim.ObjID
+	FromNode btrfsvol.LogicalAddr
+	FromItem int
+
+	ToNode       btrfsvol.LogicalAddr
+	ToLevel      uint8
+	ToKey        btrfsprim.Key
+	ToGeneration btrfsprim.Generation
+}
+
+func (kp Edge) String() string {
+	return fmt.Sprintf(`{t:%v,n:%v}[%d]->{n:%v,l:%v,g:%v,k:(%d,%v,%d)}`,
+		kp.FromTree, kp.FromNode, kp.FromItem,
+		kp.ToNode, kp.ToLevel, kp.ToGeneration,
+		kp.ToKey.ObjectID,
+		kp.ToKey.ItemType,
+		kp.ToKey.Offset)
+}
+
+type Graph struct {
+	Nodes     map[btrfsvol.LogicalAddr]Node
+	BadNodes  map[btrfsvol.LogicalAddr]error
+	EdgesFrom map[btrfsvol.LogicalAddr][]*Edge
+	EdgesTo   map[btrfsvol.LogicalAddr][]*Edge
+}
+
+func (g Graph) insertEdge(kp Edge) {
+	ptr := &kp
+	if kp.ToNode == 0 {
+		panic("kp.ToNode should not be zero")
+	}
+	g.EdgesFrom[kp.FromNode] = append(g.EdgesFrom[kp.FromNode], ptr)
+	g.EdgesTo[kp.ToNode] = append(g.EdgesTo[kp.ToNode], ptr)
+}
+
+func (g Graph) insertTreeRoot(sb btrfstree.Superblock, treeID btrfsprim.ObjID) {
+	treeInfo, err := btrfstree.LookupTreeRoot(nil, sb, treeID)
+	if err != nil {
+		// This shouldn't ever happen for treeIDs that are
+		// mentioned directly in the superblock; which are the
+		// only trees for which we should call
+		// .insertTreeRoot().
+		panic(fmt.Errorf("LookupTreeRoot(%v): %w", treeID, err))
+	}
+	if treeInfo.RootNode == 0 {
+		return
+	}
+	g.insertEdge(Edge{
+		FromTree:     treeID,
+		ToNode:       treeInfo.RootNode,
+		ToLevel:      treeInfo.Level,
+		ToGeneration: treeInfo.Generation,
+	})
+}
+
+func New(sb btrfstree.Superblock) *Graph {
+	g := &Graph{
+		Nodes:     make(map[btrfsvol.LogicalAddr]Node),
+		BadNodes:  make(map[btrfsvol.LogicalAddr]error),
+		EdgesFrom: make(map[btrfsvol.LogicalAddr][]*Edge),
+		EdgesTo:   make(map[btrfsvol.LogicalAddr][]*Edge),
+	}
+
+	// These 4 trees are mentioned directly in the superblock, so
+	// they are always seen.
+	g.insertTreeRoot(sb, btrfsprim.ROOT_TREE_OBJECTID)
+	g.insertTreeRoot(sb, btrfsprim.CHUNK_TREE_OBJECTID)
+	g.insertTreeRoot(sb, btrfsprim.TREE_LOG_OBJECTID)
+	g.insertTreeRoot(sb, btrfsprim.BLOCK_GROUP_TREE_OBJECTID)
+
+	return g
+}
+
+func (g *Graph) InsertNode(nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) {
+	for _, item := range nodeRef.Data.BodyLeaf {
+		switch itemBody := item.Body.(type) {
+		case btrfsitem.Root:
+			g.insertEdge(Edge{
+				FromTree:     item.Key.ObjectID,
+				ToNode:       itemBody.ByteNr,
+				ToLevel:      itemBody.Level,
+				ToGeneration: itemBody.Generation,
+			})
+		}
+	}
+
+	g.Nodes[nodeRef.Addr] = Node{
+		Level:      nodeRef.Data.Head.Level,
+		Generation: nodeRef.Data.Head.Generation,
+		Owner:      nodeRef.Data.Head.Owner,
+		NumItems:   nodeRef.Data.Head.NumItems,
+		MinItem:    discardOK(nodeRef.Data.MinItem()),
+		MaxItem:    discardOK(nodeRef.Data.MaxItem()),
+	}
+	for i, kp := range nodeRef.Data.BodyInternal {
+		g.insertEdge(Edge{
+			FromTree:     nodeRef.Data.Head.Owner,
+			FromNode:     nodeRef.Addr,
+			FromItem:     i,
+			ToNode:       kp.BlockPtr,
+			ToLevel:      nodeRef.Data.Head.Level - 1,
+			ToKey:        kp.Key,
+			ToGeneration: kp.Generation,
+		})
+	}
+}
+
+func (g *Graph) FinalCheck(fs diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock, progress func(n, d int)) error {
+	total := len(g.EdgesTo)
+	done := 0
+	progress(done, total)
+	for laddr := range g.EdgesTo {
+		if _, ok := g.Nodes[laddr]; !ok {
+			_, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, sb, laddr, btrfstree.NodeExpectations{
+				LAddr: containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr},
+			})
+			if err == nil {
+				return fmt.Errorf("node@%v exists but was not in node scan results", laddr)
+			}
+			g.BadNodes[laddr] = err
+		}
+		done++
+		progress(done, total)
+	}
+	return nil
+}
+
+func discardOK[T any](val T, _ bool) T {
+	return val
+}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/loops.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/loops.go
new file mode 100644
index 0000000..e76985c
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/graph/loops.go
@@ -0,0 +1,165 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package graph
+
+import (
+	"fmt"
+	"io"
+	"strings"
+
+	"github.com/datawire/dlib/derror"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsvol"
+	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
+	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
+	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
+)
+
+func (g Graph) ShowLoops(out io.Writer) {
+	orphans := make(containers.Set[btrfsvol.LogicalAddr])
+	for node := range g.Nodes {
+		if len(g.EdgesTo[node]) == 0 {
+			orphans.Insert(node)
+		}
+	}
+
+	loopWalk(out, g, 0)
+	for _, orphan := range maps.SortedKeys(orphans) {
+		loopWalk(out, g, orphan)
+	}
+}
+
+func loopWalk(out io.Writer, scanData Graph, stack ...btrfsvol.LogicalAddr) {
+	for _, kp := range scanData.EdgesFrom[stack[len(stack)-1]] {
+		childStack := append(stack, kp.ToNode)
+		if slices.Contains(kp.ToNode, stack) {
+			loopRender(out, scanData, childStack...)
+		} else {
+			loopWalk(out, scanData, childStack...)
+		}
+	}
+}
+
+func nodeRender(scanData Graph, node btrfsvol.LogicalAddr) []string {
+	if node == 0 {
+		return []string{"root"}
+	} else if nodeData, ok := scanData.Nodes[node]; ok {
+		return []string{
+			fmt.Sprintf("{addr:      %v,", node),
+			fmt.Sprintf(" level:     %v,", nodeData.Level),
+			fmt.Sprintf(" gen:       %v,", nodeData.Generation),
+			fmt.Sprintf(" num_items: %v,", nodeData.NumItems),
+			fmt.Sprintf(" min_item:  {%d,%v,%d},",
+				nodeData.MinItem.ObjectID,
+				nodeData.MinItem.ItemType,
+				nodeData.MinItem.Offset),
+			fmt.Sprintf(" max_item:  {%d,%v,%d}}",
+				nodeData.MaxItem.ObjectID,
+				nodeData.MaxItem.ItemType,
+				nodeData.MaxItem.Offset),
+		}
+	} else if nodeErr, ok := scanData.BadNodes[node]; ok {
+		return []string{
+			fmt.Sprintf("{addr:%v,", node),
+			fmt.Sprintf(" err:%q}", nodeErr.Error()),
+		}
+	} else {
+		panic("should not happen")
+	}
+}
+
+func edgeRender(scanData Graph, kp Edge) []string {
+	a := fmt.Sprintf("[%d]={", kp.FromItem)
+	b := strings.Repeat(" ", len(a))
+	ret := []string{
+		a + fmt.Sprintf("ToNode:  %v,", kp.ToNode),
+		b + fmt.Sprintf("ToLevel: %v,", kp.ToLevel),
+		b + fmt.Sprintf("ToGen:   %v,", kp.ToGeneration),
+		b + fmt.Sprintf("ToKey:   {%d,%v,%d}}",
+			kp.ToKey.ObjectID,
+			kp.ToKey.ItemType,
+			kp.ToKey.Offset),
+	}
+
+	var err error
+	if toNode, ok := scanData.Nodes[kp.ToNode]; !ok {
+		err = scanData.BadNodes[kp.ToNode]
+	} else {
+		err = checkNodeExpectations(kp, toNode)
+	}
+	if err != nil {
+		c := strings.Repeat(" ", len(a)-1)
+		ret = append(ret,
+			c+"^",
+			c+"`-err="+strings.ReplaceAll(err.Error(), "\n", "\n"+c+"      "),
+		)
+	}
+	return ret
+}
+
+func loopRender(out io.Writer, scanData Graph, stack ...btrfsvol.LogicalAddr) {
+	var lines []string
+	add := func(suffixes []string) {
+		curLen := 0
+		for _, line := range lines {
+			if len(line) > curLen {
+				curLen = len(line)
+			}
+		}
+		for i, suffix := range suffixes {
+			if len(lines) <= i {
+				lines = append(lines, "")
+			}
+			if len(lines[i]) < curLen {
+				if i == 0 {
+					lines[i] += strings.Repeat("-", curLen-len(lines[i])-1) + ">"
+				} else {
+					lines[i] += strings.Repeat(" ", curLen-len(lines[i]))
+				}
+			}
+			lines[i] += suffix
+		}
+	}
+
+	for i, node := range stack {
+		if i > 0 {
+			for _, kp := range scanData.EdgesTo[node] {
+				if kp.FromNode == stack[i-1] {
+					add(edgeRender(scanData, *kp))
+					break
+				}
+			}
+		}
+		add(nodeRender(scanData, node))
+	}
+
+	fmt.Fprintln(out, "loop:")
+	for _, line := range lines {
+		fmt.Fprintln(out, "    "+line)
+	}
+}
+
+func checkNodeExpectations(kp Edge, toNode Node) error {
+	var errs derror.MultiError
+	if toNode.Level != kp.ToLevel {
+		errs = append(errs, fmt.Errorf("kp.level=%v != node.level=%v",
+			kp.ToLevel, toNode.Level))
+	}
+	if toNode.Generation != kp.ToGeneration {
+		errs = append(errs, fmt.Errorf("kp.generation=%v != node.generation=%v",
+			kp.ToGeneration, toNode.Generation))
+	}
+	if toNode.NumItems == 0 {
+		errs = append(errs, fmt.Errorf("node.num_items=0"))
+	} else if kp.ToKey != (btrfsprim.Key{}) && toNode.MinItem != kp.ToKey {
+		errs = append(errs, fmt.Errorf("kp.key=%v != node.items[0].key=%v",
+			kp.ToKey, toNode.MinItem))
+	}
+	if len(errs) > 0 {
+		return errs
+	}
+	return nil
+}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/loops.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/loops.go
index 8ee53d2..57eb139 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/loops.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/loops.go
@@ -6,20 +6,12 @@ package rebuildnodes
 
 import (
 	"context"
-	"fmt"
 	"io"
-	"strings"
 
-	"github.com/datawire/dlib/derror"
 	"github.com/datawire/dlib/dlog"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
-	"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/containers"
-	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
-	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
 )
 
 func ShowLoops(ctx context.Context, out io.Writer, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) error {
@@ -28,151 +20,8 @@ func ShowLoops(ctx context.Context, out io.Writer, fs *btrfs.FS, nodeScanResults
 		return err
 	}
 
-	dlog.Info(ctx, "Collecting orphans...")
-	orphans := make(containers.Set[btrfsvol.LogicalAddr])
-	for node := range scanData.Nodes {
-		if len(scanData.EdgesTo[node]) == 0 {
-			orphans.Insert(node)
-		}
-	}
-
 	dlog.Info(ctx, "Walking graph...")
-	loopWalk(out, *scanData, 0)
-	for _, orphan := range maps.SortedKeys(orphans) {
-		loopWalk(out, *scanData, orphan)
-	}
-
-	return nil
-}
-
-func loopWalk(out io.Writer, scanData scanResult, stack ...btrfsvol.LogicalAddr) {
-	for _, kp := range scanData.EdgesFrom[stack[len(stack)-1]] {
-		childStack := append(stack, kp.ToNode)
-		if slices.Contains(kp.ToNode, stack) {
-			loopRender(out, scanData, childStack...)
-		} else {
-			loopWalk(out, scanData, childStack...)
-		}
-	}
-}
-
-func nodeRender(scanData scanResult, node btrfsvol.LogicalAddr) []string {
-	if node == 0 {
-		return []string{"root"}
-	} else if nodeData, ok := scanData.Nodes[node]; ok {
-		return []string{
-			fmt.Sprintf("{addr:      %v,", node),
-			fmt.Sprintf(" level:     %v,", nodeData.Level),
-			fmt.Sprintf(" gen:       %v,", nodeData.Generation),
-			fmt.Sprintf(" num_items: %v,", nodeData.NumItems),
-			fmt.Sprintf(" min_item:  {%d,%v,%d},",
-				nodeData.MinItem.ObjectID,
-				nodeData.MinItem.ItemType,
-				nodeData.MinItem.Offset),
-			fmt.Sprintf(" max_item:  {%d,%v,%d}}",
-				nodeData.MaxItem.ObjectID,
-				nodeData.MaxItem.ItemType,
-				nodeData.MaxItem.Offset),
-		}
-	} else if nodeErr, ok := scanData.BadNodes[node]; ok {
-		return []string{
-			fmt.Sprintf("{addr:%v,", node),
-			fmt.Sprintf(" err:%q}", nodeErr.Error()),
-		}
-	} else {
-		panic("should not happen")
-	}
-}
-
-func edgeRender(scanData scanResult, kp kpData) []string {
-	a := fmt.Sprintf("[%d]={", kp.FromItem)
-	b := strings.Repeat(" ", len(a))
-	ret := []string{
-		a + fmt.Sprintf("ToNode:  %v,", kp.ToNode),
-		b + fmt.Sprintf("ToLevel: %v,", kp.ToLevel),
-		b + fmt.Sprintf("ToGen:   %v,", kp.ToGeneration),
-		b + fmt.Sprintf("ToKey:   {%d,%v,%d}}",
-			kp.ToKey.ObjectID,
-			kp.ToKey.ItemType,
-			kp.ToKey.Offset),
-	}
-
-	var err error
-	if toNode, ok := scanData.Nodes[kp.ToNode]; !ok {
-		err = scanData.BadNodes[kp.ToNode]
-	} else {
-		err = checkNodeExpectations(kp, toNode)
-	}
-	if err != nil {
-		c := strings.Repeat(" ", len(a)-1)
-		ret = append(ret,
-			c+"^",
-			c+"`-err="+strings.ReplaceAll(err.Error(), "\n", "\n"+c+"      "),
-		)
-	}
-	return ret
-}
-
-func loopRender(out io.Writer, scanData scanResult, stack ...btrfsvol.LogicalAddr) {
-	var lines []string
-	add := func(suffixes []string) {
-		curLen := 0
-		for _, line := range lines {
-			if len(line) > curLen {
-				curLen = len(line)
-			}
-		}
-		for i, suffix := range suffixes {
-			if len(lines) <= i {
-				lines = append(lines, "")
-			}
-			if len(lines[i]) < curLen {
-				if i == 0 {
-					lines[i] += strings.Repeat("-", curLen-len(lines[i])-1) + ">"
-				} else {
-					lines[i] += strings.Repeat(" ", curLen-len(lines[i]))
-				}
-			}
-			lines[i] += suffix
-		}
-	}
-
-	for i, node := range stack {
-		if i > 0 {
-			for _, kp := range scanData.EdgesTo[node] {
-				if kp.FromNode == stack[i-1] {
-					add(edgeRender(scanData, *kp))
-					break
-				}
-			}
-		}
-		add(nodeRender(scanData, node))
-	}
+	scanData.nodeGraph.ShowLoops(out)
 
-	fmt.Fprintln(out, "loop:")
-	for _, line := range lines {
-		fmt.Fprintln(out, "    "+line)
-	}
-}
-
-func checkNodeExpectations(kp kpData, toNode nodeData) error {
-	var errs derror.MultiError
-	if toNode.Level != kp.ToLevel {
-		errs = append(errs, fmt.Errorf("kp.level=%v != node.level=%v",
-			kp.ToLevel, toNode.Level))
-	}
-	if toNode.Generation != kp.ToGeneration {
-		errs = append(errs, fmt.Errorf("kp.generation=%v != node.generation=%v",
-			kp.ToGeneration, toNode.Generation))
-	}
-	if toNode.NumItems == 0 {
-		errs = append(errs, fmt.Errorf("node.num_items=0"))
-	} else if kp.ToKey != (btrfsprim.Key{}) && toNode.MinItem != kp.ToKey {
-		errs = append(errs, fmt.Errorf("kp.key=%v != node.items[0].key=%v",
-			kp.ToKey, toNode.MinItem))
-	}
-	if len(errs) > 0 {
-		return errs
-	}
 	return nil
 }
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
index 3d375f0..70bd128 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/orphans.go
@@ -8,11 +8,12 @@ import (
 	"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/rebuildnodes/graph"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
 )
 
-func listRoots(graph nodeGraph, leaf btrfsvol.LogicalAddr) containers.Set[btrfsvol.LogicalAddr] {
+func listRoots(graph graph.Graph, leaf btrfsvol.LogicalAddr) containers.Set[btrfsvol.LogicalAddr] {
 	kps := graph.EdgesTo[leaf]
 	if len(kps) == 0 {
 		return containers.NewSet(leaf)
@@ -24,7 +25,7 @@ func listRoots(graph nodeGraph, leaf btrfsvol.LogicalAddr) containers.Set[btrfsv
 	return ret
 }
 
-func walk(graph nodeGraph, root btrfsvol.LogicalAddr, fn func(btrfsvol.LogicalAddr) bool) {
+func walk(graph graph.Graph, root btrfsvol.LogicalAddr, fn func(btrfsvol.LogicalAddr) bool) {
 	if _, ok := graph.Nodes[root]; !ok {
 		return
 	}
@@ -48,7 +49,7 @@ func (a keyAndTree) Cmp(b keyAndTree) int {
 	return containers.NativeCmp(a.TreeID, b.TreeID)
 }
 
-func indexOrphans(fs diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock, graph nodeGraph) (
+func indexOrphans(fs diskio.File[btrfsvol.LogicalAddr], sb btrfstree.Superblock, graph graph.Graph) (
 	orphans containers.Set[btrfsvol.LogicalAddr],
 	leaf2orphans map[btrfsvol.LogicalAddr]containers.Set[btrfsvol.LogicalAddr],
 	key2leaf *containers.SortedMap[keyAndTree, btrfsvol.LogicalAddr],
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 536b61d..5a28008 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -45,7 +45,7 @@ func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspec
 	}
 
 	dlog.Info(ctx, "Indexing orphans...")
-	orphans, leaf2orphans, key2leaf, err := indexOrphans(fs, *sb, scanData.nodeGraph)
+	orphans, leaf2orphans, key2leaf, err := indexOrphans(fs, *sb, *scanData.nodeGraph)
 	if err != nil {
 		return nil, err
 	}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
index c32ae8e..40ace46 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
@@ -6,7 +6,6 @@ package rebuildnodes
 
 import (
 	"context"
-	"fmt"
 
 	"github.com/datawire/dlib/dlog"
 
@@ -16,78 +15,14 @@ import (
 	"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/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 )
 
-type nodeData struct {
-	Level      uint8
-	Generation btrfsprim.Generation
-	Owner      btrfsprim.ObjID
-	NumItems   uint32
-	MinItem    btrfsprim.Key
-	MaxItem    btrfsprim.Key
-}
-
-type kpData struct {
-	FromTree btrfsprim.ObjID
-	FromNode btrfsvol.LogicalAddr
-	FromItem int
-
-	ToNode       btrfsvol.LogicalAddr
-	ToLevel      uint8
-	ToKey        btrfsprim.Key
-	ToGeneration btrfsprim.Generation
-}
-
-func (kp kpData) String() string {
-	return fmt.Sprintf(`{t:%v,n:%v}[%d]->{n:%v,l:%v,g:%v,k:(%d,%v,%d)}`,
-		kp.FromTree, kp.FromNode, kp.FromItem,
-		kp.ToNode, kp.ToLevel, kp.ToGeneration,
-		kp.ToKey.ObjectID,
-		kp.ToKey.ItemType,
-		kp.ToKey.Offset)
-}
-
-type nodeGraph struct {
-	Nodes     map[btrfsvol.LogicalAddr]nodeData
-	BadNodes  map[btrfsvol.LogicalAddr]error
-	EdgesFrom map[btrfsvol.LogicalAddr][]*kpData
-	EdgesTo   map[btrfsvol.LogicalAddr][]*kpData
-}
-
-func (g nodeGraph) insertEdge(kp kpData) {
-	ptr := &kp
-	if kp.ToNode == 0 {
-		panic("kp.ToNode should not be zero")
-	}
-	g.EdgesFrom[kp.FromNode] = append(g.EdgesFrom[kp.FromNode], ptr)
-	g.EdgesTo[kp.ToNode] = append(g.EdgesTo[kp.ToNode], ptr)
-}
-
-func (g nodeGraph) insertTreeRoot(sb btrfstree.Superblock, treeID btrfsprim.ObjID) {
-	treeInfo, err := btrfstree.LookupTreeRoot(nil, sb, treeID)
-	if err != nil {
-		// This shouldn't ever happen for treeIDs that are
-		// mentioned directly in the superblock; which are the
-		// only trees for which we should call
-		// .insertTreeRoot().
-		panic(fmt.Errorf("LookupTreeRoot(%v): %w", treeID, err))
-	}
-	if treeInfo.RootNode == 0 {
-		return
-	}
-	g.insertEdge(kpData{
-		FromTree:     treeID,
-		ToNode:       treeInfo.RootNode,
-		ToLevel:      treeInfo.Level,
-		ToGeneration: treeInfo.Generation,
-	})
-}
-
 type scanResult struct {
 	uuidMap
-	nodeGraph
+	nodeGraph *graph.Graph
 }
 
 func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.ScanDevicesResult) (*scanResult, error) {
@@ -101,7 +36,7 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 	lastPct := -1
 	total := countNodes(scanResults)
 	done := 0
-	progress := func() {
+	progress := func(done, total int) {
 		pct := int(100 * float64(done) / float64(total))
 		if pct != lastPct || done == total {
 			dlog.Infof(ctx, "... %v%% (%v/%v)",
@@ -118,31 +53,17 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 
 			SeenTrees: make(containers.Set[btrfsprim.ObjID]),
 		},
-		nodeGraph: nodeGraph{
-			Nodes:     make(map[btrfsvol.LogicalAddr]nodeData),
-			BadNodes:  make(map[btrfsvol.LogicalAddr]error),
-			EdgesFrom: make(map[btrfsvol.LogicalAddr][]*kpData),
-			EdgesTo:   make(map[btrfsvol.LogicalAddr][]*kpData),
-		},
+		nodeGraph: graph.New(*sb),
 	}
 
 	// These 4 trees are mentioned directly in the superblock, so
 	// they are always seen.
-	//
-	// 1
 	ret.SeenTrees.Insert(btrfsprim.ROOT_TREE_OBJECTID)
-	ret.insertTreeRoot(*sb, btrfsprim.ROOT_TREE_OBJECTID)
-	// 2
 	ret.SeenTrees.Insert(btrfsprim.CHUNK_TREE_OBJECTID)
-	ret.insertTreeRoot(*sb, btrfsprim.CHUNK_TREE_OBJECTID)
-	// 3
 	ret.SeenTrees.Insert(btrfsprim.TREE_LOG_OBJECTID)
-	ret.insertTreeRoot(*sb, btrfsprim.TREE_LOG_OBJECTID)
-	// 4
 	ret.SeenTrees.Insert(btrfsprim.BLOCK_GROUP_TREE_OBJECTID)
-	ret.insertTreeRoot(*sb, btrfsprim.BLOCK_GROUP_TREE_OBJECTID)
 
-	progress()
+	progress(done, total)
 	for _, devResults := range scanResults {
 		for laddr := range devResults.FoundNodes {
 			nodeRef, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, *sb, laddr, btrfstree.NodeExpectations{
@@ -168,13 +89,6 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 						return nil, err
 					}
 					ret.SeenTrees.Insert(item.Key.ObjectID)
-					// graph building
-					ret.insertEdge(kpData{
-						FromTree:     item.Key.ObjectID,
-						ToNode:       itemBody.ByteNr,
-						ToLevel:      itemBody.Level,
-						ToGeneration: itemBody.Generation,
-					})
 				case btrfsitem.UUIDMap:
 					uuid := btrfsitem.KeyToUUID(item.Key)
 					if err := maybeSet("ObjID2UUID", ret.ObjID2UUID, itemBody.ObjID, uuid); err != nil {
@@ -188,28 +102,10 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 			ret.SeenTrees.Insert(nodeRef.Data.Head.Owner)
 
 			// graph building
-			ret.Nodes[laddr] = nodeData{
-				Level:      nodeRef.Data.Head.Level,
-				Generation: nodeRef.Data.Head.Generation,
-				Owner:      nodeRef.Data.Head.Owner,
-				NumItems:   nodeRef.Data.Head.NumItems,
-				MinItem:    func() btrfsprim.Key { k, _ := nodeRef.Data.MinItem(); return k }(),
-				MaxItem:    func() btrfsprim.Key { k, _ := nodeRef.Data.MaxItem(); return k }(),
-			}
-			for i, kp := range nodeRef.Data.BodyInternal {
-				ret.insertEdge(kpData{
-					FromTree:     nodeRef.Data.Head.Owner,
-					FromNode:     laddr,
-					FromItem:     i,
-					ToNode:       kp.BlockPtr,
-					ToLevel:      nodeRef.Data.Head.Level - 1,
-					ToKey:        kp.Key,
-					ToGeneration: kp.Generation,
-				})
-			}
+			ret.nodeGraph.InsertNode(nodeRef)
 
 			done++
-			progress()
+			progress(done, total)
 		}
 	}
 
@@ -225,21 +121,8 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 	dlog.Info(ctx, "... done reading node data")
 
 	dlog.Infof(ctx, "Checking keypointers for dead-ends...")
-	total = len(ret.EdgesTo)
-	done = 0
-	progress()
-	for laddr := range ret.EdgesTo {
-		if _, ok := ret.Nodes[laddr]; !ok {
-			_, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, *sb, laddr, btrfstree.NodeExpectations{
-				LAddr: containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr},
-			})
-			if err == nil {
-				return nil, fmt.Errorf("node@%v exists but was not in node scan results", laddr)
-			}
-			ret.BadNodes[laddr] = err
-		}
-		done++
-		progress()
+	if err := ret.nodeGraph.FinalCheck(fs, *sb, progress); err != nil {
+		return nil, err
 	}
 	dlog.Info(ctx, "... done checking keypointers")