ls-trees: Add an option to include lost+found nodes

2 files changed, 71 insertions, 15 deletions

diff --git a/cmd/btrfs-rec/inspect_lstrees.go b/cmd/btrfs-rec/inspect_lstrees.go
index 003e7ae..c5e78fc 100644
--- a/cmd/btrfs-rec/inspect_lstrees.go
+++ b/cmd/btrfs-rec/inspect_lstrees.go
@@ -5,6 +5,7 @@
 package main
 
 import (
+	"encoding/json"
 	"fmt"
 	"os"
 	"strconv"
@@ -15,21 +16,52 @@ import (
 
 	"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/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/diskio"
 	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 	"git.lukeshu.com/btrfs-progs-ng/lib/slices"
 )
 
 func init() {
-	inspectors = append(inspectors, subcommand{
+	var nodescanFilename string
+	cmd := subcommand{
 		Command: cobra.Command{
 			Use:   "ls-trees",
 			Short: "A brief view what types of items are in each tree",
 			Args:  cliutil.WrapPositionalArgs(cobra.NoArgs),
 		},
 		RunE: func(fs *btrfs.FS, cmd *cobra.Command, _ []string) error {
+			var scanResults map[btrfsvol.DeviceID]btrfsinspect.ScanOneDevResult
+			if nodescanFilename != "" {
+				scanResultsBytes, err := os.ReadFile(nodescanFilename)
+				if err != nil {
+					return err
+				}
+				if err := json.Unmarshal(scanResultsBytes, &scanResults); err != nil {
+					return err
+				}
+			}
+
 			var treeErrCnt int
 			var treeItemCnt map[btrfsitem.Type]int
+			flush := func() {
+				totalItems := 0
+				for _, cnt := range treeItemCnt {
+					totalItems += cnt
+				}
+				numWidth := len(strconv.Itoa(slices.Max(treeErrCnt, totalItems)))
+
+				table := tabwriter.NewWriter(os.Stdout, 0, 8, 2, ' ', 0)
+				fmt.Fprintf(table, "        errors\t% *s\n", numWidth, strconv.Itoa(treeErrCnt))
+				for _, typ := range maps.SortedKeys(treeItemCnt) {
+					fmt.Fprintf(table, "        %v items\t% *s\n", typ, numWidth, strconv.Itoa(treeItemCnt[typ]))
+				}
+				fmt.Fprintf(table, "        total items\t% *s\n", numWidth, strconv.Itoa(totalItems))
+				table.Flush()
+			}
+			visitedNodes := make(map[btrfsvol.LogicalAddr]struct{})
 			btrfsutil.WalkAllTrees(cmd.Context(), fs, btrfsutil.WalkAllTreesHandler{
 				PreTree: func(name string, treeID btrfs.ObjID) {
 					treeErrCnt = 0
@@ -40,6 +72,10 @@ func init() {
 					treeErrCnt++
 				},
 				TreeWalkHandler: btrfs.TreeWalkHandler{
+					Node: func(_ btrfs.TreePath, ref *diskio.Ref[btrfsvol.LogicalAddr, btrfs.Node]) error {
+						visitedNodes[ref.Addr] = struct{}{}
+						return nil
+					},
 					Item: func(_ btrfs.TreePath, item btrfs.Item) error {
 						typ := item.Key.ItemType
 						treeItemCnt[typ] = treeItemCnt[typ] + 1
@@ -52,23 +88,40 @@ func init() {
 					},
 				},
 				PostTree: func(_ string, _ btrfs.ObjID) {
-					totalItems := 0
-					for _, cnt := range treeItemCnt {
-						totalItems += cnt
-					}
-					numWidth := len(strconv.Itoa(slices.Max(treeErrCnt, totalItems)))
-
-					table := tabwriter.NewWriter(os.Stdout, 0, 8, 2, ' ', 0)
-					fmt.Fprintf(table, "        errors\t% *s\n", numWidth, strconv.Itoa(treeErrCnt))
-					for _, typ := range maps.SortedKeys(treeItemCnt) {
-						fmt.Fprintf(table, "        %v items\t% *s\n", typ, numWidth, strconv.Itoa(treeItemCnt[typ]))
-					}
-					fmt.Fprintf(table, "        total items\t% *s\n", numWidth, strconv.Itoa(totalItems))
-					table.Flush()
+					flush()
 				},
 			})
 
+			if nodescanFilename != "" {
+				treeErrCnt = 0
+				treeItemCnt = make(map[btrfsitem.Type]int)
+				fmt.Printf("lost+found\n")
+				for _, devResults := range scanResults {
+					for laddr := range devResults.FoundNodes {
+						if _, visited := visitedNodes[laddr]; visited {
+							continue
+						}
+						visitedNodes[laddr] = struct{}{}
+						node, err := fs.ReadNode(laddr)
+						if err != nil {
+							treeErrCnt++
+							continue
+						}
+						for _, item := range node.Data.BodyLeaf {
+							typ := item.Key.ItemType
+							treeItemCnt[typ] = treeItemCnt[typ] + 1
+						}
+					}
+				}
+				flush()
+			}
+
 			return nil
 		},
-	})
+	}
+	cmd.Command.Flags().StringVar(&nodescanFilename, "nodescan", "", "Output of scan-for-nodes to use for a lost+found tree")
+	if err := cmd.Command.MarkFlagFilename("nodescan"); err != nil {
+		panic(err)
+	}
+	inspectors = append(inspectors, cmd)
 }
diff --git a/scripts/main.sh b/scripts/main.sh
index c490841..44bb879 100755
--- a/scripts/main.sh
+++ b/scripts/main.sh
@@ -35,3 +35,6 @@ gen $b.gen/3.mappings.json \
 gen $b.gen/4.ls-files.txt \
     ./btrfs-rec --pv=$b.img --mappings=$b.gen/3.mappings.json \
     inspect ls-files
+gen $b.gen/4.ls-trees.txt \
+    ./btrfs-rec --pv=$b.img --mappings=$b.gen/3.mappings.json \
+    inspect ls-trees --nodescan=$b.gen/0.scan-for-nodes.json