1 files changed, 77 insertions, 0 deletions
diff --git a/cmd/btrfs-rec/inspect/rebuildtrees/scan.go b/cmd/btrfs-rec/inspect/rebuildtrees/scan.go
new file mode 100644
index 0000000..2995a2e
--- /dev/null
+++ b/cmd/btrfs-rec/inspect/rebuildtrees/scan.go
@@ -0,0 +1,77 @@
+// Copyright (C) 2022-2023  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package rebuildtrees
+
+import (
+	"context"
+	"time"
+
+	"github.com/datawire/dlib/dlog"
+
+	"git.lukeshu.com/btrfs-progs-ng/cmd/btrfs-rec/inspect/rebuildmappings"
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs"
+	"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/btrfsutil"
+	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
+	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
+	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
+)
+
+func ScanDevices(ctx context.Context, fs *btrfs.FS, scanResults rebuildmappings.ScanDevicesResult) (btrfstree.Superblock, btrfsutil.Graph, *btrfsutil.KeyIO, error) {
+	dlog.Info(ctx, "Reading superblock...")
+	sb, err := fs.Superblock()
+	if err != nil {
+		return btrfstree.Superblock{}, btrfsutil.Graph{}, nil, err
+	}
+
+	dlog.Infof(ctx, "Reading node data from FS...")
+
+	var stats textui.Portion[int]
+	stats.D = countNodes(scanResults)
+	progressWriter := textui.NewProgress[textui.Portion[int]](
+		dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.read.substep", "read-nodes"),
+		dlog.LogLevelInfo, textui.Tunable(1*time.Second))
+
+	nodeGraph := btrfsutil.NewGraph(*sb)
+	keyIO := btrfsutil.NewKeyIO(fs, *sb)
+
+	progressWriter.Set(stats)
+	for _, devResults := range scanResults {
+		for _, laddr := range maps.SortedKeys(devResults.FoundNodes) {
+			if err := ctx.Err(); err != nil {
+				return btrfstree.Superblock{}, btrfsutil.Graph{}, nil, err
+			}
+			nodeRef, err := btrfstree.ReadNode[btrfsvol.LogicalAddr](fs, *sb, laddr, btrfstree.NodeExpectations{
+				LAddr: containers.Optional[btrfsvol.LogicalAddr]{OK: true, Val: laddr},
+			})
+			if err != nil {
+				btrfstree.FreeNodeRef(nodeRef)
+				return btrfstree.Superblock{}, btrfsutil.Graph{}, nil, err
+			}
+
+			nodeGraph.InsertNode(nodeRef)
+			keyIO.InsertNode(nodeRef)
+
+			btrfstree.FreeNodeRef(nodeRef)
+
+			stats.N++
+			progressWriter.Set(stats)
+		}
+	}
+	if stats.N != stats.D {
+		panic("should not happen")
+	}
+	progressWriter.Done()
+	dlog.Info(ctx, "... done reading node data")
+
+	ctx = dlog.WithField(ctx, "btrfs.inspect.rebuild-trees.read.substep", "check")
+	if err := nodeGraph.FinalCheck(ctx, fs, *sb); err != nil {
+		return btrfstree.Superblock{}, btrfsutil.Graph{}, nil, err
+	}
+	keyIO.SetGraph(*nodeGraph)
+
+	return *sb, *nodeGraph, keyIO, nil
+}