containers: Rethink the RBTree interface to be simpler

3 files changed, 39 insertions, 38 deletions

diff --git a/lib/btrfs/btrfsvol/chunk.go b/lib/btrfs/btrfsvol/chunk.go
index 08a0b2b..a112fd3 100644
--- a/lib/btrfs/btrfsvol/chunk.go
+++ b/lib/btrfs/btrfsvol/chunk.go
@@ -22,7 +22,10 @@ type chunkMapping struct {
 	Flags      containers.Optional[BlockGroupFlags]
 }
 
-type ChunkMapping = chunkMapping
+// Compare implements containers.Ordered.
+func (a chunkMapping) Compare(b chunkMapping) int {
+	return containers.NativeCompare(a.LAddr, b.LAddr)
+}
 
 // return -1 if 'a' is wholly to the left of 'b'
 // return 0 if there is some overlap between 'a' and 'b'
diff --git a/lib/btrfs/btrfsvol/devext.go b/lib/btrfs/btrfsvol/devext.go
index 037021c..3324476 100644
--- a/lib/btrfs/btrfsvol/devext.go
+++ b/lib/btrfs/btrfsvol/devext.go
@@ -20,6 +20,11 @@ type devextMapping struct {
 	Flags      containers.Optional[BlockGroupFlags]
 }
 
+// Compare implements containers.Ordered.
+func (a devextMapping) Compare(b devextMapping) int {
+	return containers.NativeCompare(a.PAddr, b.PAddr)
+}
+
 // return -1 if 'a' is wholly to the left of 'b'
 // return 0 if there is some overlap between 'a' and 'b'
 // return 1 if 'a is wholly to the right of 'b'
diff --git a/lib/btrfs/btrfsvol/lvm.go b/lib/btrfs/btrfsvol/lvm.go
index 59ca609..93ec438 100644
--- a/lib/btrfs/btrfsvol/lvm.go
+++ b/lib/btrfs/btrfsvol/lvm.go
@@ -22,8 +22,8 @@ type LogicalVolume[PhysicalVolume diskio.File[PhysicalAddr]] struct {
 
 	id2pv map[DeviceID]PhysicalVolume
 
-	logical2physical *containers.RBTree[containers.NativeOrdered[LogicalAddr], chunkMapping]
-	physical2logical map[DeviceID]*containers.RBTree[containers.NativeOrdered[PhysicalAddr], devextMapping]
+	logical2physical *containers.RBTree[chunkMapping]
+	physical2logical map[DeviceID]*containers.RBTree[devextMapping]
 }
 
 var _ diskio.File[LogicalAddr] = (*LogicalVolume[diskio.File[PhysicalAddr]])(nil)
@@ -33,22 +33,14 @@ func (lv *LogicalVolume[PhysicalVolume]) init() {
 		lv.id2pv = make(map[DeviceID]PhysicalVolume)
 	}
 	if lv.logical2physical == nil {
-		lv.logical2physical = &containers.RBTree[containers.NativeOrdered[LogicalAddr], chunkMapping]{
-			KeyFn: func(chunk chunkMapping) containers.NativeOrdered[LogicalAddr] {
-				return containers.NativeOrdered[LogicalAddr]{Val: chunk.LAddr}
-			},
-		}
+		lv.logical2physical = new(containers.RBTree[chunkMapping])
 	}
 	if lv.physical2logical == nil {
-		lv.physical2logical = make(map[DeviceID]*containers.RBTree[containers.NativeOrdered[PhysicalAddr], devextMapping], len(lv.id2pv))
+		lv.physical2logical = make(map[DeviceID]*containers.RBTree[devextMapping], len(lv.id2pv))
 	}
 	for devid := range lv.id2pv {
 		if _, ok := lv.physical2logical[devid]; !ok {
-			lv.physical2logical[devid] = &containers.RBTree[containers.NativeOrdered[PhysicalAddr], devextMapping]{
-				KeyFn: func(ext devextMapping) containers.NativeOrdered[PhysicalAddr] {
-					return containers.NativeOrdered[PhysicalAddr]{Val: ext.PAddr}
-				},
-			}
+			lv.physical2logical[devid] = new(containers.RBTree[devextMapping])
 		}
 	}
 }
@@ -90,11 +82,7 @@ func (lv *LogicalVolume[PhysicalVolume]) AddPhysicalVolume(id DeviceID, dev Phys
 			lv, dev.Name(), other.Name(), id)
 	}
 	lv.id2pv[id] = dev
-	lv.physical2logical[id] = &containers.RBTree[containers.NativeOrdered[PhysicalAddr], devextMapping]{
-		KeyFn: func(ext devextMapping) containers.NativeOrdered[PhysicalAddr] {
-			return containers.NativeOrdered[PhysicalAddr]{Val: ext.PAddr}
-		},
-	}
+	lv.physical2logical[id] = new(containers.RBTree[devextMapping])
 	return nil
 }
 
@@ -143,11 +131,13 @@ func (lv *LogicalVolume[PhysicalVolume]) addMapping(m Mapping, dryRun bool) erro
 		SizeLocked: m.SizeLocked,
 		Flags:      m.Flags,
 	}
-	logicalOverlaps := lv.logical2physical.SearchRange(newChunk.compareRange)
+	var logicalOverlaps []chunkMapping
 	numOverlappingStripes := 0
-	for _, chunk := range logicalOverlaps {
-		numOverlappingStripes += len(chunk.PAddrs)
-	}
+	lv.logical2physical.Subrange(newChunk.compareRange, func(node *containers.RBNode[chunkMapping]) bool {
+		logicalOverlaps = append(logicalOverlaps, node.Value)
+		numOverlappingStripes += len(node.Value.PAddrs)
+		return true
+	})
 	var err error
 	newChunk, err = newChunk.union(logicalOverlaps...)
 	if err != nil {
@@ -162,7 +152,11 @@ func (lv *LogicalVolume[PhysicalVolume]) addMapping(m Mapping, dryRun bool) erro
 		SizeLocked: m.SizeLocked,
 		Flags:      m.Flags,
 	}
-	physicalOverlaps := lv.physical2logical[m.PAddr.Dev].SearchRange(newExt.compareRange)
+	var physicalOverlaps []devextMapping
+	lv.physical2logical[m.PAddr.Dev].Subrange(newExt.compareRange, func(node *containers.RBNode[devextMapping]) bool {
+		physicalOverlaps = append(physicalOverlaps, node.Value)
+		return true
+	})
 	newExt, err = newExt.union(physicalOverlaps...)
 	if err != nil {
 		return fmt.Errorf("(%p).AddMapping: %w", lv, err)
@@ -202,13 +196,13 @@ func (lv *LogicalVolume[PhysicalVolume]) addMapping(m Mapping, dryRun bool) erro
 
 	// logical2physical
 	for _, chunk := range logicalOverlaps {
-		lv.logical2physical.Delete(containers.NativeOrdered[LogicalAddr]{Val: chunk.LAddr})
+		lv.logical2physical.Delete(lv.logical2physical.Search(chunk.Compare))
 	}
 	lv.logical2physical.Insert(newChunk)
 
 	// physical2logical
 	for _, ext := range physicalOverlaps {
-		lv.physical2logical[m.PAddr.Dev].Delete(containers.NativeOrdered[PhysicalAddr]{Val: ext.PAddr})
+		lv.physical2logical[m.PAddr.Dev].Delete(lv.physical2logical[m.PAddr.Dev].Search(ext.Compare))
 	}
 	lv.physical2logical[m.PAddr.Dev].Insert(newExt)
 
@@ -227,20 +221,18 @@ func (lv *LogicalVolume[PhysicalVolume]) addMapping(m Mapping, dryRun bool) erro
 }
 
 func (lv *LogicalVolume[PhysicalVolume]) fsck() error {
-	physical2logical := make(map[DeviceID]*containers.RBTree[containers.NativeOrdered[PhysicalAddr], devextMapping])
-	if err := lv.logical2physical.Walk(func(node *containers.RBNode[chunkMapping]) error {
+	physical2logical := make(map[DeviceID]*containers.RBTree[devextMapping])
+	var err 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 {
-				return fmt.Errorf("(%p).fsck: chunk references physical volume %v which does not exist",
+				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 {
-				physical2logical[stripe.Dev] = &containers.RBTree[containers.NativeOrdered[PhysicalAddr], devextMapping]{
-					KeyFn: func(ext devextMapping) containers.NativeOrdered[PhysicalAddr] {
-						return containers.NativeOrdered[PhysicalAddr]{Val: ext.PAddr}
-					},
-				}
+				physical2logical[stripe.Dev] = new(containers.RBTree[devextMapping])
 			}
 			physical2logical[stripe.Dev].Insert(devextMapping{
 				PAddr: stripe.Addr,
@@ -249,8 +241,9 @@ func (lv *LogicalVolume[PhysicalVolume]) fsck() error {
 				Flags: chunk.Flags,
 			})
 		}
-		return nil
-	}); err != nil {
+		return true
+	})
+	if err != nil {
 		return err
 	}
 
@@ -270,7 +263,7 @@ func (lv *LogicalVolume[PhysicalVolume]) fsck() error {
 
 func (lv *LogicalVolume[PhysicalVolume]) Mappings() []Mapping {
 	var ret []Mapping
-	_ = lv.logical2physical.Walk(func(node *containers.RBNode[chunkMapping]) error {
+	lv.logical2physical.Range(func(node *containers.RBNode[chunkMapping]) bool {
 		chunk := node.Value
 		for _, slice := range chunk.PAddrs {
 			ret = append(ret, Mapping{
@@ -280,7 +273,7 @@ func (lv *LogicalVolume[PhysicalVolume]) Mappings() []Mapping {
 				Flags: chunk.Flags,
 			})
 		}
-		return nil
+		return true
 	})
 	return ret
 }