maps: Add HasKey and HaveAnyKeysInCommon functions, use them

8 files changed, 37 insertions, 38 deletions

diff --git a/cmd/btrfs-rec/inspect/rebuildmappings/process.go b/cmd/btrfs-rec/inspect/rebuildmappings/process.go
index 7a49cc6..2e694b5 100644
--- a/cmd/btrfs-rec/inspect/rebuildmappings/process.go
+++ b/cmd/btrfs-rec/inspect/rebuildmappings/process.go
@@ -38,7 +38,7 @@ func RebuildMappings(ctx context.Context, fs *btrfs.FS, scanResults ScanDevicesR
 	devIDs := maps.SortedKeys(scanResults)
 	devices := fs.LV.PhysicalVolumes()
 	for _, devID := range devIDs {
-		if _, ok := devices[devID]; !ok {
+		if !maps.HasKey(devices, devID) {
 			return fmt.Errorf("device ID %v mentioned in scan results is not part of the filesystem", devID)
 		}
 		devResults := scanResults[devID]
diff --git a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
index 0d25ac3..2160969 100644
--- a/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
+++ b/cmd/btrfs-rec/inspect/rebuildtrees/rebuild.go
@@ -517,24 +517,12 @@ func (o *rebuilder) resolveTreeAugments(ctx context.Context, treeID btrfsprim.Ob
 ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
 func (queue *treeAugmentQueue) has(wantKey want) bool {
-	if queue != nil {
-		if queue.zero != nil {
-			if _, ok := queue.zero[wantKey]; ok {
-				return true
-			}
-		}
-		if queue.single != nil {
-			if _, ok := queue.single[wantKey]; ok {
-				return true
-			}
-		}
-		if queue.multi != nil {
-			if _, ok := queue.multi[wantKey]; ok {
-				return true
-			}
-		}
+	if queue == nil {
+		return false
 	}
-	return false
+	return (queue.zero != nil && maps.HasKey(queue.zero, wantKey)) ||
+		(queue.single != nil && maps.HasKey(queue.single, wantKey)) ||
+		(queue.multi != nil && maps.HasKey(queue.multi, wantKey))
 }
 
 func (queue *treeAugmentQueue) store(wantKey want, choices containers.Set[btrfsvol.LogicalAddr]) {
diff --git a/lib/btrfs/btrfsvol/lvm.go b/lib/btrfs/btrfsvol/lvm.go
index 3834345..7ed58a0 100644
--- a/lib/btrfs/btrfsvol/lvm.go
+++ b/lib/btrfs/btrfsvol/lvm.go
@@ -17,6 +17,7 @@ import (
 
 	"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 LogicalVolume[PhysicalVolume diskio.File[PhysicalAddr]] struct {
@@ -41,7 +42,7 @@ func (lv *LogicalVolume[PhysicalVolume]) init() {
 		lv.physical2logical = make(map[DeviceID]*containers.RBTree[devextMapping], len(lv.id2pv))
 	}
 	for devid := range lv.id2pv {
-		if _, ok := lv.physical2logical[devid]; !ok {
+		if !maps.HasKey(lv.physical2logical, devid) {
 			lv.physical2logical[devid] = new(containers.RBTree[devextMapping])
 		}
 	}
@@ -120,7 +121,7 @@ func (lv *LogicalVolume[PhysicalVolume]) AddMapping(m Mapping) error {
 func (lv *LogicalVolume[PhysicalVolume]) addMapping(m Mapping, dryRun bool) error {
 	lv.init()
 	// sanity check
-	if _, haveDev := lv.id2pv[m.PAddr.Dev]; !haveDev {
+	if !maps.HasKey(lv.id2pv, m.PAddr.Dev) {
 		return fmt.Errorf("(%p).AddMapping: do not have a physical volume with id=%v",
 			lv, m.PAddr.Dev)
 	}
@@ -228,12 +229,12 @@ func (lv *LogicalVolume[PhysicalVolume]) fsck() error {
 	lv.logical2physical.Range(func(node *containers.RBNode[chunkMapping]) bool {
 		chunk := node.Value
 		for _, stripe := range chunk.PAddrs {
-			if _, devOK := lv.id2pv[stripe.Dev]; !devOK {
+			if !maps.HasKey(lv.id2pv, stripe.Dev) {
 				err = fmt.Errorf("(%p).fsck: chunk references physical volume %v which does not exist",
 					lv, stripe.Dev)
 				return false
 			}
-			if _, exists := physical2logical[stripe.Dev]; !exists {
+			if !maps.HasKey(physical2logical, stripe.Dev) {
 				physical2logical[stripe.Dev] = new(containers.RBTree[devextMapping])
 			}
 			physical2logical[stripe.Dev].Insert(devextMapping{
diff --git a/lib/btrfsutil/graph.go b/lib/btrfsutil/graph.go
index 7863e0d..db036d8 100644
--- a/lib/btrfsutil/graph.go
+++ b/lib/btrfsutil/graph.go
@@ -226,7 +226,7 @@ func (g Graph) FinalCheck(ctx context.Context, fs btrfstree.NodeSource) error {
 	stats.D = len(g.EdgesTo)
 	progressWriter.Set(stats)
 	for laddr := range g.EdgesTo {
-		if _, ok := g.Nodes[laddr]; !ok {
+		if !maps.HasKey(g.Nodes, laddr) {
 			node, err := fs.AcquireNode(ctx, laddr, btrfstree.NodeExpectations{
 				LAddr: containers.OptionalValue(laddr),
 			})
diff --git a/lib/btrfsutil/rebuilt_tree.go b/lib/btrfsutil/rebuilt_tree.go
index b996bdb..016bb1d 100644
--- a/lib/btrfsutil/rebuilt_tree.go
+++ b/lib/btrfsutil/rebuilt_tree.go
@@ -85,7 +85,7 @@ func (tree *RebuiltTree) indexNode(ctx context.Context, node btrfsvol.LogicalAdd
 	if err := ctx.Err(); err != nil {
 		return
 	}
-	if _, done := index[node]; done {
+	if maps.HasKey(index, node) {
 		return
 	}
 	if slices.Contains(node, stack) {
diff --git a/lib/containers/set.go b/lib/containers/set.go
index 074d126..af13d50 100644
--- a/lib/containers/set.go
+++ b/lib/containers/set.go
@@ -138,20 +138,11 @@ func (o Set[T]) TakeOne() T {
 }
 
 func (o Set[T]) Has(v T) bool {
-	_, has := o[v]
-	return has
+	return maps.HasKey(o, v)
 }
 
-func (small Set[T]) HasAny(big Set[T]) bool {
-	if len(big) < len(small) {
-		small, big = big, small
-	}
-	for v := range small {
-		if _, ok := big[v]; ok {
-			return true
-		}
-	}
-	return false
+func (a Set[T]) HasAny(b Set[T]) bool {
+	return maps.HaveAnyKeysInCommon(a, b)
 }
 
 func (small Set[T]) Intersection(big Set[T]) Set[T] {
@@ -160,7 +151,7 @@ func (small Set[T]) Intersection(big Set[T]) Set[T] {
 	}
 	ret := make(Set[T])
 	for v := range small {
-		if _, ok := big[v]; ok {
+		if maps.HasKey(big, v) {
 			ret.Insert(v)
 		}
 	}
diff --git a/lib/maps/maputil.go b/lib/maps/maputil.go
index d409e70..63e52a0 100644
--- a/lib/maps/maputil.go
+++ b/lib/maps/maputil.go
@@ -25,3 +25,20 @@ func SortedKeys[K constraints.Ordered, V any](m map[K]V) []K {
 	slices.Sort(ret)
 	return ret
 }
+
+func HasKey[K comparable, V any](m map[K]V, k K) bool {
+	_, has := m[k]
+	return has
+}
+
+func HaveAnyKeysInCommon[K comparable, V1, V2 any](small map[K]V1, big map[K]V2) bool {
+	if len(big) < len(small) {
+		return HaveAnyKeysInCommon(big, small)
+	}
+	for v := range small {
+		if _, ok := big[v]; ok {
+			return true
+		}
+	}
+	return false
+}
diff --git a/lib/textui/log.go b/lib/textui/log.go
index e8325c2..e5d3f60 100644
--- a/lib/textui/log.go
+++ b/lib/textui/log.go
@@ -26,6 +26,8 @@ import (
 	"git.lukeshu.com/go/typedsync"
 	"github.com/datawire/dlib/dlog"
 	"github.com/spf13/pflag"
+
+	"git.lukeshu.com/btrfs-progs-ng/lib/maps"
 )
 
 type LogLevelFlag struct {
@@ -201,7 +203,7 @@ func (l *logger) log(lvl dlog.LogLevel, writeMsg func(io.Writer)) {
 	fields := make(map[string]any)
 	var fieldKeys []string
 	for f := l; f.parent != nil; f = f.parent {
-		if _, exists := fields[f.fieldKey]; exists {
+		if maps.HasKey(fields, f.fieldKey) {
 			continue
 		}
 		fields[f.fieldKey] = f.fieldVal