btrfsutil: Graph: Track item sizes

2 files changed, 19 insertions, 12 deletions

diff --git a/lib/btrfsutil/graph.go b/lib/btrfsutil/graph.go
index 860a49c..1b55642 100644
--- a/lib/btrfsutil/graph.go
+++ b/lib/btrfsutil/graph.go
@@ -24,12 +24,17 @@ import (
 	"git.lukeshu.com/btrfs-progs-ng/lib/textui"
 )
 
+type KeyAndSize struct {
+	Key  btrfsprim.Key
+	Size uint32
+}
+
 type GraphNode struct {
 	Addr       btrfsvol.LogicalAddr
 	Level      uint8
 	Generation btrfsprim.Generation
 	Owner      btrfsprim.ObjID
-	Items      []btrfsprim.Key
+	Items      []KeyAndSize
 }
 
 func (n GraphNode) NumItems(g Graph) int {
@@ -47,7 +52,7 @@ func (n GraphNode) MinItem(g Graph) btrfsprim.Key {
 	}
 	switch n.Level {
 	case 0:
-		return n.Items[0]
+		return n.Items[0].Key
 	default:
 		return g.EdgesFrom[n.Addr][0].ToKey
 	}
@@ -59,7 +64,7 @@ func (n GraphNode) MaxItem(g Graph) btrfsprim.Key {
 	}
 	switch n.Level {
 	case 0:
-		return n.Items[len(n.Items)-1]
+		return n.Items[len(n.Items)-1].Key
 	default:
 		return g.EdgesFrom[n.Addr][len(g.EdgesFrom[n.Addr])-1].ToKey
 	}
@@ -225,11 +230,13 @@ func (g Graph) InsertNode(node *btrfstree.Node) {
 			}
 		}
 		kps := make([]GraphEdge, 0, cnt)
-		keys := make([]btrfsprim.Key, len(node.BodyLeaf))
-		nodeData.Items = keys
+		nodeData.Items = make([]KeyAndSize, len(node.BodyLeaf))
 		g.Nodes[node.Head.Addr] = nodeData
 		for i, item := range node.BodyLeaf {
-			keys[i] = item.Key
+			nodeData.Items[i] = KeyAndSize{
+				Key:  item.Key,
+				Size: item.BodySize,
+			}
 			if itemBody, ok := item.Body.(*btrfsitem.Root); ok {
 				kps = append(kps, GraphEdge{
 					FromRoot:     node.Head.Addr,
diff --git a/lib/btrfsutil/rebuilt_tree.go b/lib/btrfsutil/rebuilt_tree.go
index 177b859..12d5978 100644
--- a/lib/btrfsutil/rebuilt_tree.go
+++ b/lib/btrfsutil/rebuilt_tree.go
@@ -329,17 +329,17 @@ func (tree *RebuiltTree) items(ctx context.Context, inc bool) containers.SortedM
 	for i, leaf := range leafs {
 		stats.Leafs.N = i
 		progressWriter.Set(stats)
-		for j, itemKey := range tree.forrest.graph.Nodes[leaf].Items {
+		for j, itemKeyAndSize := range tree.forrest.graph.Nodes[leaf].Items {
 			newPtr := ItemPtr{
 				Node: leaf,
 				Slot: j,
 			}
-			if oldPtr, exists := index.Load(itemKey); !exists {
-				index.Store(itemKey, newPtr)
+			if oldPtr, exists := index.Load(itemKeyAndSize.Key); !exists {
+				index.Store(itemKeyAndSize.Key, newPtr)
 				stats.NumItems++
 			} else {
 				if tree.RebuiltShouldReplace(oldPtr.Node, newPtr.Node) {
-					index.Store(itemKey, newPtr)
+					index.Store(itemKeyAndSize.Key, newPtr)
 				}
 				stats.NumDups++
 			}
@@ -425,8 +425,8 @@ func (tree *RebuiltTree) RebuiltAddRoot(ctx context.Context, rootNode btrfsvol.L
 			stats.AddedLeafs++
 			progressWriter.Set(stats)
 
-			for _, itemKey := range tree.forrest.graph.Nodes[leaf].Items {
-				extCB.AddedItem(ctx, tree.ID, itemKey)
+			for _, itemKeyAndSize := range tree.forrest.graph.Nodes[leaf].Items {
+				extCB.AddedItem(ctx, tree.ID, itemKeyAndSize.Key)
 				stats.AddedItems++
 				progressWriter.Set(stats)
 			}