flatten visualizenodes, comment out rebuildnodes

11 files changed, 265 insertions, 353 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuildnodes.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/bak_rebuildnodes.go
index 34b470e..1a9f487 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuildnodes.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/bak_rebuildnodes.go
@@ -4,8 +4,10 @@
 
 package rebuildnodes
 
+/*
 import (
 	"context"
+	"errors"
 
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
@@ -13,38 +15,37 @@ import (
 	"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/btrfsutil"
-	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
 	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
 )
 
 func RebuildNodes(ctx context.Context, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) (map[btrfsvol.LogicalAddr]*RebuiltNode, error) {
-	uuidMap, err := buildUUIDMap(ctx, fs, nodeScanResults)
-	if err != nil {
-		return nil, err
-	}
+		uuidMap, err := buildUUIDMap(ctx, fs, nodeScanResults)
+		if err != nil {
+			return nil, err
+		}
 
-	nfs := &RebuiltTrees{
-		inner:   fs,
-		uuidMap: uuidMap,
-	}
+		nfs := &RebuiltTrees{
+			inner:   fs,
+			uuidMap: uuidMap,
+		}
 
-	orphanedNodes, badNodes, treeAncestors, err := classifyNodes(ctx, nfs, nodeScanResults)
-	if err != nil {
-		return nil, err
-	}
+		orphanedNodes, badNodes, treeAncestors, err := classifyNodes(ctx, nfs, nodeScanResults)
+		if err != nil {
+			return nil, err
+		}
 
-	uuidMap.considerAncestors(ctx, treeAncestors)
+		uuidMap.considerAncestors(ctx, treeAncestors)
 
-	rebuiltNodes, err := reInitBrokenNodes(ctx, nfs, badNodes)
-	if err != nil {
-		return nil, err
-	}
+		rebuiltNodes, err := reInitBrokenNodes(ctx, nfs, badNodes)
+		if err != nil {
+			return nil, err
+		}
 
-	if err := reAttachNodes(ctx, nfs, orphanedNodes, rebuiltNodes); err != nil {
-		return nil, err
-	}
+		if err := reAttachNodes(ctx, nfs, orphanedNodes, rebuiltNodes); err != nil {
+			return nil, err
+		}
 
-	return rebuiltNodes, nil
+		return rebuiltNodes, nil
 }
 
 func spanOfTreePath(fs _FS, path btrfstree.TreePath) (btrfsprim.Key, btrfsprim.Key) {
@@ -79,23 +80,6 @@ func spanOfTreePath(fs _FS, path btrfstree.TreePath) (btrfsprim.Key, btrfsprim.K
 	return low, high
 }
 
-func walkFromNode(ctx context.Context, fs _FS, nodeAddr btrfsvol.LogicalAddr, errHandle func(*btrfstree.TreeError), cbs btrfstree.TreeWalkHandler) {
-	sb, _ := fs.Superblock()
-	nodeRef, _ := btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, *sb, nodeAddr, btrfstree.NodeExpectations{
-		LAddr: containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: nodeAddr},
-	})
-	if nodeRef == nil {
-		return
-	}
-	treeInfo := btrfstree.TreeRoot{
-		TreeID:     nodeRef.Data.Head.Owner,
-		RootNode:   nodeAddr,
-		Level:      nodeRef.Data.Head.Level,
-		Generation: nodeRef.Data.Head.Generation,
-	}
-	btrfstree.TreeOperatorImpl{NodeSource: fs}.RawTreeWalk(ctx, treeInfo, errHandle, cbs)
-}
-
 func getChunkTreeUUID(ctx context.Context, fs _FS) (btrfsprim.UUID, bool) {
 	ctx, cancel := context.WithCancel(ctx)
 	var ret btrfsprim.UUID
@@ -115,3 +99,4 @@ func getChunkTreeUUID(ctx context.Context, fs _FS) (btrfsprim.UUID, bool) {
 	})
 	return ret, retOK
 }
+*/
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuilttrees.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/bak_rebuilttrees.go
index da956de..7e1a7e9 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuilttrees.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/bak_rebuilttrees.go
@@ -4,6 +4,7 @@
 
 package rebuildnodes
 
+/*
 import (
 	"context"
 	"fmt"
@@ -85,3 +86,4 @@ func (fs *RebuiltTrees) TreeSearch(treeID btrfsprim.ObjID, fn func(key btrfsprim
 func (fs *RebuiltTrees) TreeSearchAll(treeID btrfsprim.ObjID, fn func(key btrfsprim.Key, size uint32) int) ([]btrfstree.Item, error) {
 	return btrfstree.TreeOperatorImpl{NodeSource: fs}.TreeSearchAll(treeID, fn)
 }
+*/
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s3_reinit.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/bak_s3_reinit.go
index 7b43d28..5983e2f 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s3_reinit.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/bak_s3_reinit.go
@@ -4,17 +4,13 @@
 
 package rebuildnodes
 
+/*
 import (
-	"context"
 	"fmt"
 	"reflect"
-	"sort"
-
-	"github.com/datawire/dlib/dlog"
 
 	"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/slices"
 )
@@ -140,3 +136,4 @@ func reInitBrokenNodes(ctx context.Context, fs _FS, badNodes []badNode) (map[btr
 	dlog.Infof(ctx, "... initialized %d nodes", len(rebuiltNodes))
 	return rebuiltNodes, nil
 }
+*/
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s3_reinit_test.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/bak_s3_reinit_test.go
index badfdfc..865cd7e 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s3_reinit_test.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/bak_s3_reinit_test.go
@@ -4,6 +4,7 @@
 
 package rebuildnodes_test
 
+/*
 import (
 	"strings"
 	"testing"
@@ -78,3 +79,4 @@ func TestEncodeRebuiltNodes(t *testing.T) {
 }
 `, buf.String())
 }
+*/
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s4_reattach.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/bak_s4_reattach.go
index ef7d284..a78d964 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s4_reattach.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/bak_s4_reattach.go
@@ -4,6 +4,7 @@
 
 package rebuildnodes
 
+/*
 import (
 	"context"
 	"sort"
@@ -157,3 +158,4 @@ func reAttachNodes(ctx context.Context, fs _FS, orphanedNodes containers.Set[btr
 		numAttached, int(100*float64(numAttached)/float64(len(orphanedNodes))))
 	return nil
 }
+*/
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s1_uuidmap.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s1_uuidmap.go
deleted file mode 100644
index 8a5e745..0000000
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s1_uuidmap.go
+++ /dev/null
@@ -1,20 +0,0 @@
-// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
-//
-// SPDX-License-Identifier: GPL-2.0-or-later
-
-package rebuildnodes
-
-import (
-	"context"
-
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect"
-)
-
-func buildUUIDMap(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.ScanDevicesResult) (uuidMap, error) {
-	ret, err := ScanDevices(ctx, fs, scanResults)
-	if err != nil {
-		return uuidMap{}, err
-	}
-	return ret.uuidMap, nil
-}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_classify.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_classify.go
deleted file mode 100644
index 6adddc3..0000000
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/s2_classify.go
+++ /dev/null
@@ -1,141 +0,0 @@
-// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
-//
-// SPDX-License-Identifier: GPL-2.0-or-later
-
-package rebuildnodes
-
-import (
-	"context"
-	"errors"
-	iofs "io/fs"
-
-	"github.com/datawire/dlib/dlog"
-
-	"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/btrfsutil"
-	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
-	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
-	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
-)
-
-type badNode struct {
-	Err  string
-	Path btrfstree.TreePath
-}
-
-// classifyNodes returns
-//
-//  1. the set of nodes don't have another node claiming it as a child, and
-//  2. the list of bad nodes (in no particular order)
-//  3. tree ancestors thing
-func classifyNodes(ctx context.Context, fs _FS, scanResults btrfsinspect.ScanDevicesResult) (
-	orphanedNodes containers.Set[btrfsvol.LogicalAddr],
-	badNodes []badNode,
-	treeAncestors map[btrfsprim.ObjID]containers.Set[btrfsprim.ObjID],
-	err error,
-) {
-	dlog.Info(ctx, "Walking trees to identify orphan and broken nodes...")
-
-	lastPct := -1
-	total := countNodes(scanResults)
-	visitedNodes := make(containers.Set[btrfsvol.LogicalAddr])
-	progress := func() {
-		done := len(visitedNodes)
-		pct := int(100 * float64(done) / float64(total))
-		if pct != lastPct || done == total {
-			dlog.Infof(ctx, "... %v%% (%v/%v)",
-				pct, done, total)
-			lastPct = pct
-		}
-	}
-
-	treeAncestors = make(map[btrfsprim.ObjID]containers.Set[btrfsprim.ObjID])
-
-	var potentialRoot btrfsvol.LogicalAddr // zero for non-lost nodes, non-zero for lost nodes
-	nodeHandler := func(path btrfstree.TreePath, nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node], err error) error {
-		if err != nil && (errors.Is(err, btrfstree.ErrNotANode) || errors.As(err, new(*btrfstree.IOError))) {
-			badNodes = append(badNodes, badNode{
-				Err:  err.Error(),
-				Path: path.DeepCopy(),
-			})
-			return err
-		}
-		addr := path.Node(-1).ToNodeAddr
-		visitedNodes.Insert(addr)
-		if potentialRoot != 0 {
-			// lost node
-			if addr != potentialRoot {
-				delete(orphanedNodes, addr)
-			}
-		}
-		if err == nil && nodeRef.Data.Head.Owner != path.Node(-1).FromTree {
-			if treeAncestors[path.Node(-1).FromTree] == nil {
-				treeAncestors[path.Node(-1).FromTree] = make(containers.Set[btrfsprim.ObjID])
-			}
-			treeAncestors[path.Node(-1).FromTree].Insert(nodeRef.Data.Head.Owner)
-		}
-		progress()
-		return nil
-	}
-
-	walkHandler := btrfstree.TreeWalkHandler{
-		PreNode: func(path btrfstree.TreePath) error {
-			addr := path.Node(-1).ToNodeAddr
-			if visitedNodes.Has(addr) {
-				// Can happen because of COW subvolumes;
-				// this is really a DAG not a tree.
-				return iofs.SkipDir
-			}
-			return nil
-		},
-		Node: func(path btrfstree.TreePath, nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) error {
-			return nodeHandler(path, nodeRef, nil)
-		},
-		BadNode: func(path btrfstree.TreePath, nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node], err error) error {
-			return nodeHandler(path, nodeRef, err)
-		},
-	}
-
-	progress()
-	btrfsutil.WalkAllTrees(ctx, fs, btrfsutil.WalkAllTreesHandler{
-		TreeWalkHandler: walkHandler,
-		Err: func(err *btrfsutil.WalkError) {
-			// do nothing
-		},
-	})
-
-	// Start with 'orphanedRoots' as a complete set of all orphaned nodes, and then delete
-	// non-root entries from it.
-	orphanedNodes = make(containers.Set[btrfsvol.LogicalAddr])
-	for _, devResults := range scanResults {
-		for laddr := range devResults.FoundNodes {
-			if !visitedNodes.Has(laddr) {
-				orphanedNodes.Insert(laddr)
-			}
-		}
-	}
-	if len(visitedNodes)+len(orphanedNodes) != total {
-		panic("should not happen")
-	}
-	dlog.Infof(ctx,
-		"... (finished processing %v attached nodes, proceeding to process %v lost nodes, for a total of %v)",
-		len(visitedNodes), len(orphanedNodes), len(visitedNodes)+len(orphanedNodes))
-	for _, potentialRoot = range maps.SortedKeys(orphanedNodes) {
-		walkFromNode(ctx, fs, potentialRoot,
-			func(err *btrfstree.TreeError) {
-				// do nothing
-			},
-			walkHandler,
-		)
-	}
-
-	if len(visitedNodes) != total {
-		panic("should not happen")
-	}
-
-	dlog.Infof(ctx, "... identified %d orphaned nodes and %d bad nodes", len(orphanedNodes), len(badNodes))
-	return orphanedNodes, badNodes, treeAncestors, nil
-}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
index cd3f59e..5070e7e 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/scan.go
@@ -6,6 +6,7 @@ package rebuildnodes
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/datawire/dlib/dlog"
 
@@ -20,32 +21,55 @@ import (
 )
 
 type nodeData struct {
-	Level      uint8                // 0+1=1
-	Generation btrfsprim.Generation // 1+8=9
-	Owner      btrfsprim.ObjID      // 9+8=17
-	MinItem    btrfsprim.Key        // 17+17=34
-	MaxItem    btrfsprim.Key        // 34+17=51
+	Level      uint8
+	Generation btrfsprim.Generation
+	Owner      btrfsprim.ObjID
+	NumItems   uint32
+	MinItem    btrfsprim.Key
+	MaxItem    btrfsprim.Key
 }
 
 type kpData struct {
-	From, To   btrfsvol.LogicalAddr // 0+(2*8)=16
-	Key        btrfsprim.Key        // 16+17=33
-	Generation btrfsprim.Generation // 33+8=41
+	FromTree btrfsprim.ObjID
+	FromNode btrfsvol.LogicalAddr
+	FromItem int
+
+	ToNode       btrfsvol.LogicalAddr
+	ToLevel      uint8
+	ToKey        btrfsprim.Key
+	ToGeneration btrfsprim.Generation
 }
 
 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
+	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, _ := btrfstree.LookupTreeRoot(nil, sb, treeID)
+	g.insertEdge(kpData{
+		FromTree:     treeID,
+		ToNode:       treeInfo.RootNode,
+		ToLevel:      treeInfo.Level,
+		ToGeneration: treeInfo.Generation,
+	})
+}
+
 type scanResult struct {
 	uuidMap
 	nodeGraph
 }
 
 func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.ScanDevicesResult) (*scanResult, error) {
-	dlog.Infof(ctx, "Building table of ObjID←→UUID...")
+	dlog.Infof(ctx, "Reading node data from FS...")
 
 	sb, err := fs.Superblock()
 	if err != nil {
@@ -74,6 +98,7 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 		},
 		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),
 		},
@@ -81,10 +106,19 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 
 	// 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()
 	for _, devResults := range scanResults {
@@ -112,6 +146,13 @@ 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 {
@@ -132,15 +173,16 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 				MinItem:    func() btrfsprim.Key { k, _ := nodeRef.Data.MinItem(); return k }(),
 				MaxItem:    func() btrfsprim.Key { k, _ := nodeRef.Data.MaxItem(); return k }(),
 			}
-			for _, kp := range nodeRef.Data.BodyInternal {
-				dat := &kpData{
-					From:       laddr,
-					To:         kp.BlockPtr,
-					Key:        kp.Key,
-					Generation: kp.Generation,
-				}
-				ret.EdgesFrom[laddr] = append(ret.EdgesFrom[laddr], dat)
-				ret.EdgesTo[laddr] = append(ret.EdgesTo[laddr], dat)
+			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,
+				})
 			}
 
 			done++
@@ -157,6 +199,26 @@ func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults btrfsinspect.Sca
 		dlog.Errorf(ctx, "... could not find root items for %d trees: %v", len(missing), maps.SortedKeys(missing))
 	}
 
-	dlog.Info(ctx, "... done building table")
+	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()
+	}
+	dlog.Info(ctx, "... done checking keypointers")
+
 	return ret, nil
 }
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/treeancestors.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/treeancestors.go
new file mode 100644
index 0000000..2e1d856
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/treeancestors.go
@@ -0,0 +1,25 @@
+// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package rebuildnodes
+
+import (
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
+	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
+)
+
+func getTreeAncestors(scanData scanResult) map[btrfsprim.ObjID]containers.Set[btrfsprim.ObjID] {
+	treeAncestors := make(map[btrfsprim.ObjID]containers.Set[btrfsprim.ObjID])
+
+	for laddr, node := range scanData.Nodes {
+		if _, ok := treeAncestors[node.Owner]; !ok {
+			treeAncestors[node.Owner] = make(containers.Set[btrfsprim.ObjID])
+		}
+		for _, edge := range scanData.EdgesTo[laddr] {
+			treeAncestors[node.Owner].Insert(edge.FromTree)
+		}
+	}
+
+	return treeAncestors
+}
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/util.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/util.go
index f9e4e95..7e0eab1 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/util.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/util.go
@@ -6,9 +6,7 @@ package rebuildnodes
 
 import (
 	"fmt"
-	"math"
 
-	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfsprogs/btrfsinspect"
 )
 
@@ -24,6 +22,7 @@ func maybeSet[K, V comparable](name string, m map[K]V, k K, v V) error {
 	return nil
 }
 
+/*
 var maxKey = btrfsprim.Key{
 	ObjectID: math.MaxUint64,
 	ItemType: math.MaxUint8,
@@ -41,6 +40,7 @@ func keyMm(key btrfsprim.Key) btrfsprim.Key {
 	}
 	return key
 }
+*/
 
 func countNodes(nodeScanResults btrfsinspect.ScanDevicesResult) int {
 	var cnt int
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/visualizenodes.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/visualizenodes.go
index 5b1844d..368fafd 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/visualizenodes.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/visualizenodes.go
@@ -7,23 +7,18 @@ package rebuildnodes
 import (
 	"archive/zip"
 	"context"
-	"errors"
 	"fmt"
 	"html"
 	"io"
-	iofs "io/fs"
 	"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/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/btrfsutil"
 	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
-	"git.lukeshu.com/btrfs-progs-ng/lib/diskio"
 	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 )
 
@@ -53,166 +48,164 @@ func getCliqueID(cliques map[btrfsprim.ObjID]*containers.Set[btrfsprim.ObjID], t
 }
 
 func VisualizeNodes(ctx context.Context, out io.Writer, fs *btrfs.FS, nodeScanResults btrfsinspect.ScanDevicesResult) error {
-	uuidMap, err := buildUUIDMap(ctx, fs, nodeScanResults)
+	scanData, err := ScanDevices(ctx, fs, nodeScanResults)
 	if err != nil {
 		return err
 	}
 
-	nfs := &RebuiltTrees{
-		inner:   fs,
-		uuidMap: uuidMap,
-	}
-
-	orphanedNodes, _, treeAncestors, err := classifyNodes(ctx, nfs, nodeScanResults)
-	if err != nil {
-		return err
-	}
-
-	uuidMap.considerAncestors(ctx, treeAncestors)
+	dlog.Info(ctx, "Walking trees to rebuild root items...")
+	treeAncestors := getTreeAncestors(*scanData)
+	scanData.considerAncestors(ctx, treeAncestors)
 
 	////////////////////////////////////////////////////////////////////////////////////////////
 
-	cliques := getCliques(uuidMap, treeAncestors)
+	cliques := getCliques(scanData.uuidMap, treeAncestors)
 
 	dlog.Info(ctx, "Building graphviz graph...")
 
-	nodes := make(map[btrfsprim.ObjID]containers.Set[string]) // grouped by treeID
-	edges := make(map[btrfsprim.ObjID]containers.Set[string]) // grouped by cliqueID
-	visitedNodes := make(containers.Set[btrfsvol.LogicalAddr])
-	var isOrphan bool
-
-	nodeHandler := func(path btrfstree.TreePath, nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node], err error) error {
-		addr := path.Node(-1).ToNodeAddr
+	type graph struct {
+		nodes    map[btrfsprim.ObjID]containers.Set[string] // keyed by treeID
+		badNodes containers.Set[string]
+		edges    containers.Set[string]
+	}
+	graphs := make(map[btrfsprim.ObjID]graph) // keyed by cliqueID
+
+	for _, laddr := range maps.SortedKeys(scanData.Nodes) {
+		nodeData := scanData.Nodes[laddr]
+		cliqueID := getCliqueID(cliques, nodeData.Owner)
+		graph, ok := graphs[cliqueID]
+		if !ok {
+			graph.nodes = make(map[btrfsprim.ObjID]containers.Set[string])
+			graph.badNodes = make(containers.Set[string])
+			graph.edges = make(containers.Set[string])
+		}
+		if graph.nodes[nodeData.Owner] == nil {
+			graph.nodes[nodeData.Owner] = make(containers.Set[string])
+		}
 
-		// Node
-		var treeID btrfsprim.ObjID
-		var nodeStr string
-		if err != nil && (errors.Is(err, btrfstree.ErrNotANode) || errors.As(err, new(*btrfstree.IOError))) {
-			treeID = 0
-			nodeStr = fmt.Sprintf(`n%d [shape=star color=red label="%v"]`, addr, addr)
+		var buf strings.Builder
+		fmt.Fprintf(&buf, `n%d [shape=record label="%v\ngen=%v\nlvl=%v|{`,
+			laddr,
+			laddr,
+			nodeData.Generation,
+			nodeData.Level)
+		if nodeData.NumItems == 0 {
+			buf.WriteString("(no items)")
 		} else {
-			treeID = nodeRef.Data.Head.Owner
-			var buf strings.Builder
-			fmt.Fprintf(&buf, `n%d [shape=record label="%v\ngen=%v\nlvl=%v|{`,
-				addr,
-				nodeRef.Data.Head.Addr,
-				nodeRef.Data.Head.Generation,
-				nodeRef.Data.Head.Level)
-			if nodeRef.Data.Head.NumItems == 0 {
-				buf.WriteString("(no items)")
-			} else {
-				for i := uint32(0); i < nodeRef.Data.Head.NumItems; i++ {
-					if i == 0 {
-						fmt.Fprintf(&buf, "<p%[1]d>%[1]d", i)
-					} else {
-						fmt.Fprintf(&buf, "|<p%[1]d>%[1]d", i)
-					}
+			for i := uint32(0); i < nodeData.NumItems; i++ {
+				if i == 0 {
+					fmt.Fprintf(&buf, "<p%[1]d>%[1]d", i)
+				} else {
+					fmt.Fprintf(&buf, "|<p%[1]d>%[1]d", i)
 				}
 			}
-			buf.WriteString(`}"]`)
-			nodeStr = buf.String()
 		}
-		if _, ok := nodes[treeID]; !ok {
-			nodes[treeID] = make(containers.Set[string])
-			nodes[treeID].Insert(fmt.Sprintf(`t%d [label="%s"]`, treeID, html.EscapeString(treeID.String())))
-		}
-		nodes[treeID].Insert(nodeStr)
+		buf.WriteString(`}"]`)
+		graph.nodes[nodeData.Owner].Insert(buf.String())
 
-		// Edge
-		var edge strings.Builder
-		if len(path) == 1 {
-			if isOrphan {
-				edge.WriteString("orphanRoot")
-			} else {
-				fmt.Fprintf(&edge, "t%d", path[0].FromTree)
-			}
-		} else {
-			fmt.Fprintf(&edge, "n%d:p%d", path.Node(-2).ToNodeAddr, path.Node(-1).FromItemIdx)
-		}
-		fmt.Fprintf(&edge, ` -> n%d [label="`, addr)
-		if path.Node(-1).FromItemIdx >= 0 {
-			fmt.Fprintf(&edge, "%d: key=(%d,%v,%d) gen=%v",
-				path.Node(-1).FromItemIdx,
-				path.Node(-1).ToKey.ObjectID,
-				path.Node(-1).ToKey.ItemType,
-				path.Node(-1).ToKey.Offset,
-				path.Node(-1).ToNodeGeneration)
-		}
-		if err != nil {
-			fmt.Fprintf(&edge, `\n\n%s" color=red]`, html.EscapeString(err.Error()))
-		} else {
-			edge.WriteString(`"]`)
+		if len(scanData.EdgesTo[laddr]) == 0 {
+			graph.edges.Insert(fmt.Sprintf("orphanRoot -> n%d", laddr))
 		}
-		cliqueID := getCliqueID(cliques, path[0].FromTree)
-		if treeID != 0 && getCliqueID(cliques, treeID) != cliqueID {
-			panic(fmt.Errorf("tree %v is not in clique %v", treeID, maps.SortedKeys(*cliques[cliqueID])))
-		}
-		if !cliques[cliqueID].Has(cliqueID) {
-			panic(fmt.Errorf("clique %v does not contain supposed-member %v", maps.SortedKeys(*cliques[cliqueID]), cliqueID))
+
+		graphs[cliqueID] = graph
+	}
+
+	for _, laddr := range maps.SortedKeys(scanData.BadNodes) {
+		cliqueIDs := make(containers.Set[btrfsprim.ObjID])
+		for _, edge := range scanData.EdgesTo[laddr] {
+			cliqueIDs.Insert(getCliqueID(cliques, edge.FromTree))
 		}
-		if _, ok := edges[cliqueID]; !ok {
-			edges[cliqueID] = make(containers.Set[string])
+		if len(cliqueIDs) != 1 {
+			dlog.Errorf(ctx, "couldn't assign bad node@%v to a clique: %v", laddr, maps.SortedKeys(cliqueIDs))
+			continue
 		}
-		edges[cliqueID].Insert(edge.String())
 
-		// Return
-		if visitedNodes.Has(addr) {
-			return iofs.SkipDir
-		}
-		visitedNodes.Insert(addr)
-		return err
+		cliqueID := cliqueIDs.TakeOne()
+		graph := graphs[cliqueID]
+		graph.badNodes.Insert(fmt.Sprintf(`n%d [shape=star color=red label="%v"]`, laddr, laddr))
+		graphs[cliqueID] = graph
 	}
 
-	walkHandler := btrfstree.TreeWalkHandler{
-		Node: func(path btrfstree.TreePath, nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node]) error {
-			return nodeHandler(path, nodeRef, nil)
-		},
-		BadNode: func(path btrfstree.TreePath, nodeRef *diskio.Ref[btrfsvol.LogicalAddr, btrfstree.Node], err error) error {
-			return nodeHandler(path, nodeRef, err)
-		},
-	}
+	for _, laddr := range maps.SortedKeys(scanData.EdgesFrom) {
+		for _, kp := range scanData.EdgesFrom[laddr] {
+			cliqueID := getCliqueID(cliques, kp.FromTree)
+			graph := graphs[cliqueID]
 
-	btrfsutil.WalkAllTrees(ctx, nfs, btrfsutil.WalkAllTreesHandler{
-		TreeWalkHandler: walkHandler,
-		Err: func(err *btrfsutil.WalkError) {
-			// do nothing
-		},
-	})
-	isOrphan = true
-	for _, potentialRoot := range maps.SortedKeys(orphanedNodes) {
-		walkFromNode(ctx, nfs, potentialRoot,
-			func(err *btrfstree.TreeError) {
-				// do nothing
-			},
-			walkHandler,
-		)
-	}
+			var buf strings.Builder
+			if kp.FromNode == 0 {
+				graph.nodes[kp.FromTree].Insert(fmt.Sprintf(`t%d [label="%s"]`, kp.FromTree, html.EscapeString(kp.FromTree.String())))
+				fmt.Fprintf(&buf, "t%d", kp.FromTree)
+			} else {
+				fmt.Fprintf(&buf, "n%d", kp.FromNode)
+			}
+			fmt.Fprintf(&buf, ` -> n%d [label="%d: key=(%d,%v,%d) gen=%v`,
+				// dst node
+				kp.ToNode,
+				// label
+				kp.FromItem,
+				kp.ToKey.ObjectID,
+				kp.ToKey.ItemType,
+				kp.ToKey.Offset,
+				kp.ToGeneration)
+			toNode, ok := scanData.Nodes[kp.ToNode]
+			var err error
+			if !ok {
+				err = scanData.BadNodes[kp.ToNode]
+			} else {
+				var errs derror.MultiError
+				if toNode.Level != kp.ToLevel {
+					errs = append(errs, fmt.Errorf("node.level=%v != kp.level=%v",
+						toNode.Level, kp.ToLevel))
+				}
+				if toNode.Generation != kp.ToGeneration {
+					errs = append(errs, fmt.Errorf("node.generation=%v != kp.generation=%v",
+						toNode.Generation, kp.ToGeneration))
+				}
+				if toNode.NumItems == 0 {
+					errs = append(errs, fmt.Errorf("node.num_items=0"))
+				} else if toNode.MinItem != kp.ToKey {
+					errs = append(errs, fmt.Errorf("node.items[0].key=%v != kp.key=%v",
+						toNode.MinItem, kp.ToKey))
+				}
+				if len(errs) > 0 {
+					err = errs
+				}
+			}
+			if err != nil {
+				fmt.Fprintf(&buf, `\n\n%s" color=red]`, html.EscapeString(err.Error()))
+			} else {
+				buf.WriteString(`"]`)
+			}
 
-	dlog.Info(ctx, "... done building")
+			graph.edges.Insert(buf.String())
+			graphs[cliqueID] = graph
+		}
+	}
 
 	////////////////////////////////////////////////////////////////////////////////////////////
 
 	dlog.Info(ctx, "Writing graphviz output...")
 
-	cliqueIDs := maps.SortedKeys(edges)
-
 	zw := zip.NewWriter(out)
-	for _, cliqueID := range cliqueIDs {
+	for _, cliqueID := range maps.SortedKeys(graphs) {
 		if err := func() (err error) {
+			graph := graphs[cliqueID]
+
 			buf, err := zw.Create(fmt.Sprintf("%d.dot", cliqueID))
 			if err != nil {
 				return err
 			}
-
 			if _, err := fmt.Fprintf(buf, "strict digraph clique%d {\n", cliqueID); err != nil {
 				return err
 			}
-			clique := cliques[cliqueID]
-			for _, treeID := range maps.SortedKeys(*clique) {
+
+			for _, treeID := range maps.SortedKeys(graph.nodes) {
+				nodes := graph.nodes[treeID]
+
 				if _, err := fmt.Fprintf(buf, "  subgraph cluster_tree%d {\n", treeID); err != nil {
 					return err
 				}
-				for _, node := range maps.SortedKeys(nodes[treeID]) {
+				for _, node := range maps.SortedKeys(nodes) {
 					if _, err := fmt.Fprintf(buf, "    %s;\n", node); err != nil {
 						return err
 					}
@@ -221,15 +214,20 @@ func VisualizeNodes(ctx context.Context, out io.Writer, fs *btrfs.FS, nodeScanRe
 					return err
 				}
 			}
-			for _, edge := range maps.SortedKeys(edges[cliqueID]) {
+			for _, node := range maps.SortedKeys(graph.badNodes) {
+				if _, err := fmt.Fprintf(buf, "   %s;\n", node); err != nil {
+					return err
+				}
+			}
+			for _, edge := range maps.SortedKeys(graph.edges) {
 				if _, err := fmt.Fprintf(buf, "  %s;\n", edge); err != nil {
 					return err
 				}
 			}
+
 			if _, err := fmt.Fprintln(buf, "}"); err != nil {
 				return err
 			}
-
 			return nil
 		}(); err != nil {
 			return err