rebuildnodes: Don't have wantKeys be stringly-typed

4 files changed, 250 insertions, 153 deletions

diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
index 902d810..6d610ee 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild.go
@@ -9,7 +9,6 @@ import (
 	"fmt"
 	"runtime"
 	"sort"
-	"strings"
 	"time"
 
 	"github.com/datawire/dlib/dgroup"
@@ -64,10 +63,9 @@ type rebuilder struct {
 }
 
 type treeAugmentQueue struct {
-	keyBuf strings.Builder
-	zero   map[string]struct{}
-	single map[string]btrfsvol.LogicalAddr
-	multi  map[string]containers.Set[btrfsvol.LogicalAddr]
+	zero   map[Want]struct{}
+	single map[Want]btrfsvol.LogicalAddr
+	multi  map[Want]containers.Set[btrfsvol.LogicalAddr]
 }
 
 type Rebuilder interface {
@@ -401,7 +399,9 @@ func (o *rebuilder) resolveTreeAugments(ctx context.Context, treeID btrfsprim.Ob
 	wantKeys := append(
 		maps.Keys(o.augmentQueue[treeID].single),
 		maps.Keys(o.augmentQueue[treeID].multi)...)
-	sort.Strings(wantKeys)
+	sort.Slice(wantKeys, func(i, j int) bool {
+		return wantKeys[i].Cmp(wantKeys[j]) < 0
+	})
 	for _, wantKey := range wantKeys {
 		list, ok := o.augmentQueue[treeID].multi[wantKey]
 		if !ok {
@@ -418,27 +418,14 @@ func (o *rebuilder) resolveTreeAugments(ctx context.Context, treeID btrfsprim.Ob
 	// Free some memory
 	o.augmentQueue[treeID].single = nil
 	o.augmentQueue[treeID].multi = nil
-	o.augmentQueue[treeID].keyBuf.Reset()
 
 	return ret
 }
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
-func shortenWantKey(wantKey string) string {
-	if !strings.HasPrefix(wantKey, "tree=") {
-		panic("should not happen")
-	}
-	sp := strings.IndexByte(wantKey, ' ')
-	if sp < 0 {
-		panic("should not happen")
-	}
-	return wantKey[sp+1:]
-}
-
-func (queue *treeAugmentQueue) has(wantKey string) bool {
+func (queue *treeAugmentQueue) has(wantKey Want) bool {
 	if queue != nil {
-		wantKey = shortenWantKey(wantKey)
 		if queue.zero != nil {
 			if _, ok := queue.zero[wantKey]; ok {
 				return true
@@ -458,49 +445,40 @@ func (queue *treeAugmentQueue) has(wantKey string) bool {
 	return false
 }
 
-func (queue *treeAugmentQueue) store(wantKey string, choices containers.Set[btrfsvol.LogicalAddr]) {
-	if len(choices) == 0 && (strings.Contains(wantKey, " name=") || strings.Contains(wantKey, "-")) {
-		// This wantKey is unlikely to come up again, so it's not worth storing a negative result.
+func (queue *treeAugmentQueue) store(wantKey Want, choices containers.Set[btrfsvol.LogicalAddr]) {
+	if len(choices) == 0 && wantKey.OffsetType > offsetExact {
+		// This wantKey is unlikely to come up again, so it's
+		// not worth the RAM of storing a negative result.
 		return
 	}
-	wantKey = shortenWantKey(wantKey)
-	beg := queue.keyBuf.Len()
-	if queue.keyBuf.Cap()-beg < len(wantKey) {
-		queue.keyBuf.Reset()
-		queue.keyBuf.Grow(textui.Tunable(4096))
-		beg = 0
-	}
-	queue.keyBuf.WriteString(wantKey)
-	wantKey = queue.keyBuf.String()[beg:]
-
 	switch len(choices) {
 	case 0:
 		if queue.zero == nil {
-			queue.zero = make(map[string]struct{})
+			queue.zero = make(map[Want]struct{})
 		}
 		queue.zero[wantKey] = struct{}{}
 	case 1:
 		if queue.single == nil {
-			queue.single = make(map[string]btrfsvol.LogicalAddr)
+			queue.single = make(map[Want]btrfsvol.LogicalAddr)
 		}
 		queue.single[wantKey] = choices.TakeOne()
 	default:
 		if queue.multi == nil {
-			queue.multi = make(map[string]containers.Set[btrfsvol.LogicalAddr])
+			queue.multi = make(map[Want]containers.Set[btrfsvol.LogicalAddr])
 		}
 		queue.multi[wantKey] = choices
 	}
 }
 
-func (o *rebuilder) hasAugment(treeID btrfsprim.ObjID, wantKey string) bool {
-	return o.augmentQueue[treeID].has(wantKey)
+func (o *rebuilder) hasAugment(wantKey WantWithTree) bool {
+	return o.augmentQueue[wantKey.TreeID].has(wantKey.Key)
 }
 
-func (o *rebuilder) wantAugment(ctx context.Context, treeID btrfsprim.ObjID, wantKey string, choices containers.Set[btrfsvol.LogicalAddr]) {
-	if o.augmentQueue[treeID] == nil {
-		o.augmentQueue[treeID] = new(treeAugmentQueue)
+func (o *rebuilder) wantAugment(ctx context.Context, wantKey WantWithTree, choices containers.Set[btrfsvol.LogicalAddr]) {
+	if o.augmentQueue[wantKey.TreeID] == nil {
+		o.augmentQueue[wantKey.TreeID] = new(treeAugmentQueue)
 	}
-	o.augmentQueue[treeID].store(wantKey, choices)
+	o.augmentQueue[wantKey.TreeID].store(wantKey.Key, choices)
 	if len(choices) == 0 {
 		dlog.Error(ctx, "could not find wanted item")
 		o.numAugmentFailures++
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_treecb.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_treecb.go
index 7266777..827aa5c 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_treecb.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_treecb.go
@@ -8,8 +8,6 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/datawire/dlib/dlog"
-
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsitem"
 	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
 )
@@ -25,26 +23,25 @@ func (o *rebuilder) cbAddedItem(ctx context.Context, tree btrfsprim.ObjID, key b
 }
 
 func (o *rebuilder) cbLookupRoot(ctx context.Context, tree btrfsprim.ObjID) (offset btrfsprim.Generation, item btrfsitem.Root, ok bool) {
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-tree.want.reason", "tree Root")
-	key := keyAndTree{
+	wantKey := WantWithTree{
 		TreeID: btrfsprim.ROOT_TREE_OBJECTID,
-		Key: btrfsprim.Key{
-			ObjectID: tree,
-			ItemType: btrfsitem.ROOT_ITEM_KEY,
+		Key: Want{
+			ObjectID:   tree,
+			ItemType:   btrfsitem.ROOT_ITEM_KEY,
+			OffsetType: offsetAny,
 		},
 	}
-	wantKey := fmt.Sprintf("tree=%v key={%v %v ?}", key.TreeID, key.ObjectID, key.ItemType)
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-tree.want.key", wantKey)
-	key.Key, ok = o._want(ctx, key.TreeID, wantKey, key.ObjectID, key.ItemType)
+	ctx = withWant(ctx, logFieldTreeWant, "tree Root", wantKey)
+	foundKey, ok := o._want(ctx, wantKey)
 	if !ok {
 		o.enqueueRetry()
 		return 0, btrfsitem.Root{}, false
 	}
-	switch itemBody := o.rebuilt.Tree(ctx, key.TreeID).ReadItem(ctx, key.Key).(type) {
+	switch itemBody := o.rebuilt.Tree(ctx, wantKey.TreeID).ReadItem(ctx, foundKey).(type) {
 	case *btrfsitem.Root:
-		return btrfsprim.Generation(key.Offset), *itemBody, true
+		return btrfsprim.Generation(foundKey.Offset), *itemBody, true
 	case *btrfsitem.Error:
-		o.fsErr(ctx, fmt.Errorf("error decoding item: %v: %w", key, itemBody.Err))
+		o.fsErr(ctx, fmt.Errorf("error decoding item: %v: %w", foundKey, itemBody.Err))
 		return 0, btrfsitem.Root{}, false
 	default:
 		// This is a panic because the item decoder should not emit ROOT_ITEM items as anything but
@@ -54,18 +51,20 @@ func (o *rebuilder) cbLookupRoot(ctx context.Context, tree btrfsprim.ObjID) (off
 }
 
 func (o *rebuilder) cbLookupUUID(ctx context.Context, uuid btrfsprim.UUID) (id btrfsprim.ObjID, ok bool) {
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-tree.want.reason", "resolve parent UUID")
-	key := keyAndTree{TreeID: btrfsprim.UUID_TREE_OBJECTID, Key: btrfsitem.UUIDToKey(uuid)}
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.add-tree.want.key", key.String())
-	if !o._wantOff(ctx, key.TreeID, key.String(), key.Key) {
+	wantKey := WantWithTree{
+		TreeID: btrfsprim.UUID_TREE_OBJECTID,
+		Key:    wantFromKey(btrfsitem.UUIDToKey(uuid)),
+	}
+	ctx = withWant(ctx, logFieldTreeWant, "resolve parent UUID", wantKey)
+	if !o._wantOff(ctx, wantKey) {
 		o.enqueueRetry()
 		return 0, false
 	}
-	switch itemBody := o.rebuilt.Tree(ctx, key.TreeID).ReadItem(ctx, key.Key).(type) {
+	switch itemBody := o.rebuilt.Tree(ctx, wantKey.TreeID).ReadItem(ctx, wantKey.Key.Key()).(type) {
 	case *btrfsitem.UUIDMap:
 		return itemBody.ObjID, true
 	case *btrfsitem.Error:
-		o.fsErr(ctx, fmt.Errorf("error decoding item: %v: %w", key, itemBody.Err))
+		o.fsErr(ctx, fmt.Errorf("error decoding item: %v: %w", wantKey, itemBody.Err))
 		return 0, false
 	default:
 		// This is a panic because the item decoder should not emit UUID_SUBVOL items as anything but
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_wantcb.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_wantcb.go
index 7ea9c85..c57b2bb 100644
--- a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_wantcb.go
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_wantcb.go
@@ -26,25 +26,28 @@ func (o *rebuilder) fsErr(ctx context.Context, e error) {
 
 // want implements rebuildCallbacks.
 func (o *rebuilder) want(ctx context.Context, reason string, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType) {
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.reason", reason)
-	wantKey := fmt.Sprintf("tree=%v key={%v %v ?}", treeID, objID, typ)
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.key", wantKey)
-	o._want(ctx, treeID, wantKey, objID, typ)
+	wantKey := WantWithTree{
+		TreeID: treeID,
+		Key: Want{
+			ObjectID:   objID,
+			ItemType:   typ,
+			OffsetType: offsetAny,
+		},
+	}
+	ctx = withWant(ctx, logFieldItemWant, reason, wantKey)
+	o._want(ctx, wantKey)
 }
 
-func (o *rebuilder) _want(ctx context.Context, treeID btrfsprim.ObjID, wantKey string, objID btrfsprim.ObjID, typ btrfsprim.ItemType) (key btrfsprim.Key, ok bool) {
-	if o.rebuilt.Tree(ctx, treeID) == nil {
+func (o *rebuilder) _want(ctx context.Context, wantKey WantWithTree) (key btrfsprim.Key, ok bool) {
+	if o.rebuilt.Tree(ctx, wantKey.TreeID) == nil {
 		o.enqueueRetry()
 		return btrfsprim.Key{}, false
 	}
 
 	// check if we already have it
 
-	tgt := btrfsprim.Key{
-		ObjectID: objID,
-		ItemType: typ,
-	}
-	if key, _, ok := o.rebuilt.Tree(ctx, treeID).Items(ctx).Search(func(key btrfsprim.Key, _ keyio.ItemPtr) int {
+	tgt := wantKey.Key.Key()
+	if key, _, ok := o.rebuilt.Tree(ctx, wantKey.TreeID).Items(ctx).Search(func(key btrfsprim.Key, _ keyio.ItemPtr) int {
 		key.Offset = 0
 		return tgt.Compare(key)
 	}); ok {
@@ -53,62 +56,80 @@ func (o *rebuilder) _want(ctx context.Context, treeID btrfsprim.ObjID, wantKey s
 
 	// OK, we need to insert it
 
-	if o.hasAugment(treeID, wantKey) {
+	if o.hasAugment(wantKey) {
 		return btrfsprim.Key{}, false
 	}
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.rebuilt.Tree(ctx, treeID).PotentialItems(ctx).Subrange(
-		func(k btrfsprim.Key, _ keyio.ItemPtr) int { k.Offset = 0; return tgt.Compare(k) },
+	o.rebuilt.Tree(ctx, wantKey.TreeID).PotentialItems(ctx).Subrange(
+		func(k btrfsprim.Key, _ keyio.ItemPtr) int {
+			k.Offset = 0
+			return tgt.Compare(k)
+		},
 		func(_ btrfsprim.Key, v keyio.ItemPtr) bool {
-			wants.InsertFrom(o.rebuilt.Tree(ctx, treeID).LeafToRoots(ctx, v.Node))
+			wants.InsertFrom(o.rebuilt.Tree(ctx, wantKey.TreeID).LeafToRoots(ctx, v.Node))
 			return true
 		})
-	o.wantAugment(ctx, treeID, wantKey, wants)
+	o.wantAugment(ctx, wantKey, wants)
 	return btrfsprim.Key{}, false
 }
 
 // wantOff implements rebuildCallbacks.
 func (o *rebuilder) wantOff(ctx context.Context, reason string, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType, off uint64) {
-	key := btrfsprim.Key{
-		ObjectID: objID,
-		ItemType: typ,
-		Offset:   off,
+	wantKey := WantWithTree{
+		TreeID: treeID,
+		Key: Want{
+			ObjectID:   objID,
+			ItemType:   typ,
+			OffsetType: offsetExact,
+			OffsetLow:  off,
+		},
 	}
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.reason", reason)
-	wantKey := keyAndTree{TreeID: treeID, Key: key}.String()
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.key", wantKey)
-	o._wantOff(ctx, treeID, wantKey, key)
+	ctx = withWant(ctx, logFieldItemWant, reason, wantKey)
+	o._wantOff(ctx, wantKey)
 }
 
-func (o *rebuilder) _wantOff(ctx context.Context, treeID btrfsprim.ObjID, wantKey string, tgt btrfsprim.Key) (ok bool) {
-	if o.rebuilt.Tree(ctx, treeID) == nil {
+func (o *rebuilder) _wantOff(ctx context.Context, wantKey WantWithTree) (ok bool) {
+	if o.rebuilt.Tree(ctx, wantKey.TreeID) == nil {
 		o.enqueueRetry()
 		return false
 	}
 
 	// check if we already have it
 
-	if _, ok := o.rebuilt.Tree(ctx, treeID).Items(ctx).Load(tgt); ok {
+	tgt := wantKey.Key.Key()
+	if _, ok := o.rebuilt.Tree(ctx, wantKey.TreeID).Items(ctx).Load(tgt); ok {
 		return true
 	}
 
 	// OK, we need to insert it
 
-	if o.hasAugment(treeID, wantKey) {
+	if o.hasAugment(wantKey) {
 		return false
 	}
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
-	o.rebuilt.Tree(ctx, treeID).PotentialItems(ctx).Subrange(
+	o.rebuilt.Tree(ctx, wantKey.TreeID).PotentialItems(ctx).Subrange(
 		func(k btrfsprim.Key, _ keyio.ItemPtr) int { return tgt.Compare(k) },
 		func(_ btrfsprim.Key, v keyio.ItemPtr) bool {
-			wants.InsertFrom(o.rebuilt.Tree(ctx, treeID).LeafToRoots(ctx, v.Node))
+			wants.InsertFrom(o.rebuilt.Tree(ctx, wantKey.TreeID).LeafToRoots(ctx, v.Node))
 			return true
 		})
-	o.wantAugment(ctx, treeID, wantKey, wants)
+	o.wantAugment(ctx, wantKey, wants)
 	return false
 }
 
-func (o *rebuilder) _wantFunc(ctx context.Context, treeID btrfsprim.ObjID, wantKey string, objID btrfsprim.ObjID, typ btrfsprim.ItemType, fn func(keyio.ItemPtr) bool) {
+// wantDirIndex implements rebuildCallbacks.
+func (o *rebuilder) wantDirIndex(ctx context.Context, reason string, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, name []byte) {
+	wantKey := WantWithTree{
+		TreeID: treeID,
+		Key: Want{
+			ObjectID:   objID,
+			ItemType:   btrfsitem.DIR_INDEX_KEY,
+			OffsetType: offsetName,
+			OffsetName: string(name),
+		},
+	}
+	ctx = withWant(ctx, logFieldItemWant, reason, wantKey)
+
 	if o.rebuilt.Tree(ctx, treeID) == nil {
 		o.enqueueRetry()
 		return
@@ -116,18 +137,15 @@ func (o *rebuilder) _wantFunc(ctx context.Context, treeID btrfsprim.ObjID, wantK
 
 	// check if we already have it
 
-	tgt := btrfsprim.Key{
-		ObjectID: objID,
-		ItemType: typ,
-	}
+	tgt := wantKey.Key.Key()
 	found := false
 	o.rebuilt.Tree(ctx, treeID).Items(ctx).Subrange(
 		func(key btrfsprim.Key, _ keyio.ItemPtr) int {
 			key.Offset = 0
 			return tgt.Compare(key)
 		},
-		func(_ btrfsprim.Key, itemPtr keyio.ItemPtr) bool {
-			if fn(itemPtr) {
+		func(_ btrfsprim.Key, ptr keyio.ItemPtr) bool {
+			if itemName, ok := o.keyIO.Names[ptr]; ok && bytes.Equal(itemName, name) {
 				found = true
 			}
 			return !found
@@ -138,31 +156,22 @@ func (o *rebuilder) _wantFunc(ctx context.Context, treeID btrfsprim.ObjID, wantK
 
 	// OK, we need to insert it
 
-	if o.hasAugment(treeID, wantKey) {
+	if o.hasAugment(wantKey) {
 		return
 	}
 	wants := make(containers.Set[btrfsvol.LogicalAddr])
 	o.rebuilt.Tree(ctx, treeID).PotentialItems(ctx).Subrange(
-		func(k btrfsprim.Key, _ keyio.ItemPtr) int { k.Offset = 0; return tgt.Compare(k) },
-		func(k btrfsprim.Key, v keyio.ItemPtr) bool {
-			if fn(v) {
-				wants.InsertFrom(o.rebuilt.Tree(ctx, treeID).LeafToRoots(ctx, v.Node))
+		func(key btrfsprim.Key, _ keyio.ItemPtr) int {
+			key.Offset = 0
+			return tgt.Compare(key)
+		},
+		func(_ btrfsprim.Key, ptr keyio.ItemPtr) bool {
+			if itemName, ok := o.keyIO.Names[ptr]; ok && bytes.Equal(itemName, name) {
+				wants.InsertFrom(o.rebuilt.Tree(ctx, treeID).LeafToRoots(ctx, ptr.Node))
 			}
 			return true
 		})
-	o.wantAugment(ctx, treeID, wantKey, wants)
-}
-
-// wantDirIndex implements rebuildCallbacks.
-func (o *rebuilder) wantDirIndex(ctx context.Context, reason string, treeID btrfsprim.ObjID, objID btrfsprim.ObjID, name []byte) {
-	typ := btrfsitem.DIR_INDEX_KEY
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.reason", reason)
-	wantKey := fmt.Sprintf("tree=%v key={%v %v ?} name=%q", treeID, objID, typ, name)
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.key", wantKey)
-	o._wantFunc(ctx, treeID, wantKey, objID, typ, func(ptr keyio.ItemPtr) bool {
-		itemName, ok := o.keyIO.Names[ptr]
-		return ok && bytes.Equal(itemName, name)
-	})
+	o.wantAugment(ctx, wantKey, wants)
 }
 
 func (o *rebuilder) _walkRange(
@@ -231,12 +240,20 @@ func (a gap) Compare(b gap) int {
 }
 
 func (o *rebuilder) _wantRange(
-	ctx context.Context,
+	ctx context.Context, reason string,
 	treeID btrfsprim.ObjID, objID btrfsprim.ObjID, typ btrfsprim.ItemType,
 	beg, end uint64,
 ) {
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.key",
-		fmt.Sprintf("tree=%v key={%v %v ?}", treeID, objID, typ))
+	wantKey := WantWithTree{
+		TreeID: treeID,
+		Key: Want{
+			ObjectID:   objID,
+			ItemType:   typ,
+			OffsetType: offsetAny,
+		},
+	}
+	ctx = withWant(ctx, logFieldItemWant, reason, wantKey)
+	wantKey.Key.OffsetType = offsetRange
 
 	if o.rebuilt.Tree(ctx, treeID) == nil {
 		o.enqueueRetry()
@@ -311,26 +328,23 @@ func (o *rebuilder) _wantRange(
 				// TODO: This is dumb and greedy.
 				if last < runBeg {
 					// log an error
-					wantKey := fmt.Sprintf("tree=%v key={%v %v %v-%v}", treeID, objID, typ, last, runBeg)
-					if !o.hasAugment(treeID, wantKey) {
-						wantCtx := dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.key", wantKey)
-						o.wantAugment(wantCtx, treeID, wantKey, nil)
-					}
-				}
-				wantKey := fmt.Sprintf("tree=%v key={%v %v %v-%v}", treeID, objID, typ, gap.Beg, gap.End)
-				if !o.hasAugment(treeID, wantKey) {
-					wantCtx := dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.key", wantKey)
-					o.wantAugment(wantCtx, treeID, wantKey, o.rebuilt.Tree(wantCtx, treeID).LeafToRoots(wantCtx, v.Node))
+					wantKey.Key.OffsetLow = last
+					wantKey.Key.OffsetHigh = runBeg
+					wantCtx := withWant(ctx, logFieldItemWant, reason, wantKey)
+					o.wantAugment(wantCtx, wantKey, nil)
 				}
+				wantKey.Key.OffsetLow = gap.Beg
+				wantKey.Key.OffsetHigh = gap.End
+				wantCtx := withWant(ctx, logFieldItemWant, reason, wantKey)
+				o.wantAugment(wantCtx, wantKey, o.rebuilt.Tree(wantCtx, treeID).LeafToRoots(wantCtx, v.Node))
 				last = runEnd
 			})
 		if last < gap.End {
 			// log an error
-			wantKey := fmt.Sprintf("tree=%v key={%v, %v, %v-%v}", treeID, objID, typ, last, gap.End)
-			if !o.hasAugment(treeID, wantKey) {
-				wantCtx := dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.key", wantKey)
-				o.wantAugment(wantCtx, treeID, wantKey, nil)
-			}
+			wantKey.Key.OffsetLow = last
+			wantKey.Key.OffsetHigh = gap.End
+			wantCtx := withWant(ctx, logFieldItemWant, reason, wantKey)
+			o.wantAugment(wantCtx, wantKey, nil)
 		}
 		return true
 	})
@@ -340,24 +354,23 @@ func (o *rebuilder) _wantRange(
 //
 // interval is [beg, end)
 func (o *rebuilder) wantCSum(ctx context.Context, reason string, inodeTree, inode btrfsprim.ObjID, beg, end btrfsvol.LogicalAddr) {
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.reason", reason)
-
-	inodeKey := keyAndTree{
+	inodeWant := WantWithTree{
 		TreeID: inodeTree,
-		Key: btrfsprim.Key{
-			ObjectID: inode,
-			ItemType: btrfsitem.INODE_ITEM_KEY,
-			Offset:   0,
+		Key: Want{
+			ObjectID:   inode,
+			ItemType:   btrfsitem.INODE_ITEM_KEY,
+			OffsetType: offsetExact,
+			OffsetLow:  0,
 		},
 	}
-	inodeCtx := dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.key", inodeKey.String())
-	if !o._wantOff(inodeCtx, inodeKey.TreeID, inodeKey.String(), inodeKey.Key) {
+	inodeCtx := withWant(ctx, logFieldItemWant, reason, inodeWant)
+	if !o._wantOff(inodeCtx, inodeWant) {
 		o.enqueueRetry()
 		return
 	}
-	inodePtr, ok := o.rebuilt.Tree(inodeCtx, inodeKey.TreeID).Items(inodeCtx).Load(inodeKey.Key)
+	inodePtr, ok := o.rebuilt.Tree(inodeCtx, inodeTree).Items(inodeCtx).Load(inodeWant.Key.Key())
 	if !ok {
-		panic(fmt.Errorf("should not happen: could not load key: %v", inodeKey))
+		panic(fmt.Errorf("should not happen: could not load key: %v", inodeWant))
 	}
 	inodeFlags, ok := o.keyIO.Flags[inodePtr]
 	if !ok {
@@ -372,16 +385,16 @@ func (o *rebuilder) wantCSum(ctx context.Context, reason string, inodeTree, inod
 		return
 	}
 
-	beg = roundDown(beg, btrfssum.BlockSize)
-	end = roundUp(end, btrfssum.BlockSize)
-	const treeID = btrfsprim.CSUM_TREE_OBJECTID
-	o._wantRange(ctx, treeID, btrfsprim.EXTENT_CSUM_OBJECTID, btrfsprim.EXTENT_CSUM_KEY,
-		uint64(beg), uint64(end))
+	o._wantRange(
+		ctx, reason,
+		btrfsprim.CSUM_TREE_OBJECTID, btrfsprim.EXTENT_CSUM_OBJECTID, btrfsprim.EXTENT_CSUM_KEY,
+		uint64(roundDown(beg, btrfssum.BlockSize)), uint64(roundUp(end, btrfssum.BlockSize)))
 }
 
 // wantFileExt implements rebuildCallbacks.
 func (o *rebuilder) wantFileExt(ctx context.Context, reason string, treeID btrfsprim.ObjID, ino btrfsprim.ObjID, size int64) {
-	ctx = dlog.WithField(ctx, "btrfsinspect.rebuild-nodes.rebuild.want.reason", reason)
-	o._wantRange(ctx, treeID, ino, btrfsprim.EXTENT_DATA_KEY,
+	o._wantRange(
+		ctx, reason,
+		treeID, ino, btrfsprim.EXTENT_DATA_KEY,
 		0, uint64(size))
 }
diff --git a/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_wanttyp.go b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_wanttyp.go
new file mode 100644
index 0000000..eb1bb49
--- /dev/null
+++ b/lib/btrfsprogs/btrfsinspect/rebuildnodes/rebuild_wanttyp.go
@@ -0,0 +1,107 @@
+// Copyright (C) 2022-2023  Luke Shumaker <lukeshu@lukeshu.com>
+//
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+package rebuildnodes
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/datawire/dlib/dlog"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/btrfs/btrfsprim"
+	"git.lukeshu.com/btrfs-progs-ng/lib/containers"
+)
+
+type wantOffsetType int8
+
+const (
+	offsetAny = wantOffsetType(iota)
+	offsetExact
+	offsetRange
+	offsetName
+)
+
+type Want struct {
+	ObjectID   btrfsprim.ObjID
+	ItemType   btrfsprim.ItemType
+	OffsetType wantOffsetType
+	OffsetLow  uint64
+	OffsetHigh uint64
+	OffsetName string
+}
+
+func (a Want) Cmp(b Want) int {
+	if d := containers.NativeCompare(a.ObjectID, b.ObjectID); d != 0 {
+		return d
+	}
+	if d := containers.NativeCompare(a.ItemType, b.ItemType); d != 0 {
+		return d
+	}
+	if d := containers.NativeCompare(a.OffsetType, b.OffsetType); d != 0 {
+		return d
+	}
+	if d := containers.NativeCompare(a.OffsetLow, b.OffsetLow); d != 0 {
+		return d
+	}
+	if d := containers.NativeCompare(a.OffsetHigh, b.OffsetHigh); d != 0 {
+		return d
+	}
+	if d := containers.NativeCompare(a.OffsetName, b.OffsetName); d != 0 {
+		return d
+	}
+	return 0
+}
+
+func (o Want) Key() btrfsprim.Key {
+	return btrfsprim.Key{
+		ObjectID: o.ObjectID,
+		ItemType: o.ItemType,
+		Offset:   o.OffsetLow,
+	}
+}
+
+func wantFromKey(k btrfsprim.Key) Want {
+	return Want{
+		ObjectID:   k.ObjectID,
+		ItemType:   k.ItemType,
+		OffsetType: offsetExact,
+		OffsetLow:  k.Offset,
+	}
+}
+
+func (o Want) String() string {
+	switch o.OffsetType {
+	case offsetAny:
+		return fmt.Sprintf("{%v %v ?}", o.ObjectID, o.ItemType)
+	case offsetExact:
+		return fmt.Sprintf("{%v %v %v}", o.ObjectID, o.ItemType, o.OffsetLow)
+	case offsetRange:
+		return fmt.Sprintf("{%v %v %v-%v}", o.ObjectID, o.ItemType, o.OffsetLow, o.OffsetHigh)
+	case offsetName:
+		return fmt.Sprintf("{%v %v name=%q}", o.ObjectID, o.ItemType, o.OffsetName)
+	default:
+		panic(fmt.Errorf("should not happen: OffsetType=%#v", o.OffsetType))
+	}
+}
+
+type WantWithTree struct {
+	TreeID btrfsprim.ObjID
+	Key    Want
+}
+
+func (o WantWithTree) String() string {
+	return fmt.Sprintf("tree=%v key=%v", o.TreeID, o.Key)
+}
+
+const (
+	logFieldItemWant = "btrfsinspect.rebuild-nodes.rebuild.want"
+	logFieldTreeWant = "btrfsinspect.rebuild-nodes.rebuild.add-tree.want"
+)
+
+func withWant(ctx context.Context, logField, reason string, wantKey WantWithTree) context.Context {
+	ctx = dlog.WithField(ctx, logField+".reason", reason)
+	ctx = dlog.WithField(ctx, logField+".key", wantKey)
+	return ctx
+}