btrfsvol: use rbtree

4 files changed, 110 insertions, 65 deletions

diff --git a/pkg/btrfs/btrfsvol/lvm.go b/pkg/btrfs/btrfsvol/lvm.go
index 01ecaf3..078b18e 100644
--- a/pkg/btrfs/btrfsvol/lvm.go
+++ b/pkg/btrfs/btrfsvol/lvm.go
@@ -3,9 +3,7 @@ package btrfsvol
 import (
 	"bytes"
 	"fmt"
-	"math"
 	"reflect"
-	"sort"
 
 	"lukeshu.com/btrfs-tools/pkg/rbtree"
 	"lukeshu.com/btrfs-tools/pkg/util"
@@ -33,10 +31,13 @@ func (lv *LogicalVolume[PhysicalVolume]) init() {
 			},
 		}
 	}
+	if lv.physical2logical == nil {
+		lv.physical2logical = make(map[DeviceID]*rbtree.Tree[PhysicalAddr, devextMapping], len(lv.id2pv))
+	}
 	for devid := range lv.id2pv {
 		if _, ok := lv.physical2logical[devid]; !ok {
-			lv.physical2logical[devid] = &rbtree.Tree[LogicalAddr, chunkMapping]{
-				KeyFn: func(ext devextMapping) LogicalAddr {
+			lv.physical2logical[devid] = &rbtree.Tree[PhysicalAddr, devextMapping]{
+				KeyFn: func(ext devextMapping) PhysicalAddr {
 					return ext.PAddr
 				},
 			}
@@ -68,8 +69,8 @@ func (lv *LogicalVolume[PhysicalVolume]) AddPhysicalVolume(id DeviceID, dev Phys
 			lv, dev.Name(), other.Name(), id)
 	}
 	lv.id2pv[id] = dev
-	lv.physical2logical[id] = &rbtree.Tree[LogicalAddr, chunkMapping]{
-		KeyFn: func(ext devextMapping) LogicalAddr {
+	lv.physical2logical[id] = &rbtree.Tree[PhysicalAddr, devextMapping]{
+		KeyFn: func(ext devextMapping) PhysicalAddr {
 			return ext.PAddr
 		},
 	}
@@ -139,23 +140,15 @@ func (lv *LogicalVolume[PhysicalVolume]) AddMapping(m Mapping) error {
 
 	// logical2physical
 	for _, chunk := range logicalOverlaps {
-		lv.logical2physical = util.RemoveAllFromSliceFunc(lv.logical2physical, func(otherChunk chunkMapping) bool {
-			return otherChunk.LAddr == chunk.LAddr
-		})
+		lv.logical2physical.Delete(chunk.LAddr)
 	}
-	lv.logical2physical = append(lv.logical2physical, newChunk)
-	sort.Slice(lv.logical2physical, func(i, j int) bool {
-		return lv.logical2physical[i].LAddr < lv.logical2physical[j].LAddr
-	})
+	lv.logical2physical.Insert(newChunk)
 
 	// physical2logical
 	for _, ext := range physicalOverlaps {
-		lv.physical2logical[m.PAddr.Dev] = util.RemoveAllFromSlice(lv.physical2logical[m.PAddr.Dev], ext)
+		lv.physical2logical[m.PAddr.Dev].Delete(ext.PAddr)
 	}
-	lv.physical2logical[m.PAddr.Dev] = append(lv.physical2logical[m.PAddr.Dev], newExt)
-	sort.Slice(lv.physical2logical[m.PAddr.Dev], func(i, j int) bool {
-		return lv.physical2logical[m.PAddr.Dev][i].PAddr < lv.physical2logical[m.PAddr.Dev][j].PAddr
-	})
+	lv.physical2logical[m.PAddr.Dev].Insert(newExt)
 
 	// sanity check
 	//
@@ -170,38 +163,51 @@ func (lv *LogicalVolume[PhysicalVolume]) AddMapping(m Mapping) error {
 }
 
 func (lv *LogicalVolume[PhysicalVolume]) fsck() error {
-	physical2logical := make(map[DeviceID][]devextMapping)
-	for _, chunk := range lv.logical2physical {
+	physical2logical := make(map[DeviceID]*rbtree.Tree[PhysicalAddr, devextMapping])
+	if err := lv.logical2physical.Walk(func(node *rbtree.Node[chunkMapping]) error {
+		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",
 					lv, stripe.Dev)
 			}
-			physical2logical[stripe.Dev] = append(physical2logical[stripe.Dev], devextMapping{
+			if _, exists := physical2logical[stripe.Dev]; !exists {
+				physical2logical[stripe.Dev] = &rbtree.Tree[PhysicalAddr, devextMapping]{
+					KeyFn: func(ext devextMapping) PhysicalAddr {
+						return ext.PAddr
+					},
+				}
+			}
+			physical2logical[stripe.Dev].Insert(devextMapping{
 				PAddr: stripe.Addr,
 				LAddr: chunk.LAddr,
 				Size:  chunk.Size,
 				Flags: chunk.Flags,
 			})
 		}
-	}
-	for _, exts := range physical2logical {
-		sort.Slice(exts, func(i, j int) bool {
-			return exts[i].PAddr < exts[j].PAddr
-		})
+		return nil
+	}); err != nil {
+		return err
 	}
 
-	if !reflect.DeepEqual(lv.physical2logical, physical2logical) {
+	if len(lv.physical2logical) != len(physical2logical) {
 		return fmt.Errorf("(%p).fsck: skew between chunk tree and devext tree",
 			lv)
 	}
+	for devid := range lv.physical2logical {
+		if !lv.physical2logical[devid].Equal(physical2logical[devid]) {
+			return fmt.Errorf("(%p).fsck: skew between chunk tree and devext tree",
+				lv)
+		}
+	}
 
 	return nil
 }
 
 func (lv *LogicalVolume[PhysicalVolume]) Mappings() []Mapping {
 	var ret []Mapping
-	for _, chunk := range lv.logical2physical {
+	lv.logical2physical.Walk(func(node *rbtree.Node[chunkMapping]) error {
+		chunk := node.Value
 		var flags *BlockGroupFlags
 		if chunk.Flags != nil {
 			val := *chunk.Flags
@@ -215,42 +221,46 @@ func (lv *LogicalVolume[PhysicalVolume]) Mappings() []Mapping {
 				Flags: flags,
 			})
 		}
-	}
+		return nil
+	})
 	return ret
 }
 
 func (lv *LogicalVolume[PhysicalVolume]) Resolve(laddr LogicalAddr) (paddrs map[QualifiedPhysicalAddr]struct{}, maxlen AddrDelta) {
+	node := lv.logical2physical.Search(func(chunk chunkMapping) int {
+		return chunkMapping{LAddr: laddr, Size: 1}.cmpRange(chunk)
+	})
+	if node == nil {
+		return nil, 0
+	}
+
+	chunk := node.Value
+
+	offsetWithinChunk := laddr.Sub(chunk.LAddr)
 	paddrs = make(map[QualifiedPhysicalAddr]struct{})
-	maxlen = math.MaxInt64
-
-	for _, chunk := range lv.logical2physical {
-		low := chunk.LAddr
-		high := low.Add(chunk.Size)
-		if low <= laddr && laddr < high {
-			offsetWithinChunk := laddr.Sub(low)
-			maxlen = util.Min(maxlen, chunk.Size-offsetWithinChunk)
-			for _, stripe := range chunk.PAddrs {
-				paddrs[QualifiedPhysicalAddr{
-					Dev:  stripe.Dev,
-					Addr: stripe.Addr.Add(offsetWithinChunk),
-				}] = struct{}{}
-			}
-		}
+	maxlen = chunk.Size - offsetWithinChunk
+	for _, stripe := range chunk.PAddrs {
+		paddrs[QualifiedPhysicalAddr{
+			Dev:  stripe.Dev,
+			Addr: stripe.Addr.Add(offsetWithinChunk),
+		}] = struct{}{}
 	}
 
 	return paddrs, maxlen
 }
 
 func (lv *LogicalVolume[PhysicalVolume]) UnResolve(paddr QualifiedPhysicalAddr) LogicalAddr {
-	for _, ext := range lv.physical2logical[paddr.Dev] {
-		low := ext.PAddr
-		high := low.Add(ext.Size)
-		if low <= paddr.Addr && paddr.Addr < high {
-			offsetWithinExt := paddr.Addr.Sub(low)
-			return ext.LAddr.Add(offsetWithinExt)
-		}
+	node := lv.physical2logical[paddr.Dev].Search(func(ext devextMapping) int {
+		return devextMapping{PAddr: paddr.Addr, Size: 1}.cmpRange(ext)
+	})
+	if node == nil {
+		return -1
 	}
-	return -1
+
+	ext := node.Value
+
+	offsetWithinExt := paddr.Addr.Sub(ext.PAddr)
+	return ext.LAddr.Add(offsetWithinExt)
 }
 
 func (lv *LogicalVolume[PhysicalVolume]) ReadAt(dat []byte, laddr LogicalAddr) (int, error) {
diff --git a/pkg/rbtree/rbtree.go b/pkg/rbtree/rbtree.go
index 13d9602..7927307 100644
--- a/pkg/rbtree/rbtree.go
+++ b/pkg/rbtree/rbtree.go
@@ -33,17 +33,24 @@ type Tree[K constraints.Ordered, V any] struct {
 	root  *Node[V]
 }
 
-func (t *Tree[K, V]) Walk(fn func(*Node[V])) {
-	t.root.walk(fn)
+func (t *Tree[K, V]) Walk(fn func(*Node[V]) error) error {
+	return t.root.walk(fn)
 }
 
-func (node *Node[V]) walk(fn func(*Node[V])) {
+func (node *Node[V]) walk(fn func(*Node[V]) error) error {
 	if node == nil {
-		return
+		return nil
+	}
+	if err := node.Left.walk(fn); err != nil {
+		return err
+	}
+	if err := fn(node); err != nil {
+		return err
+	}
+	if err := node.Right.walk(fn); err != nil {
+		return err
 	}
-	node.Left.walk(fn)
-	fn(node)
-	node.Right.walk(fn)
+	return nil
 }
 
 // Search the tree for a value that satisfied the given callbackk
diff --git a/pkg/rbtree/rbtree_test.go b/pkg/rbtree/rbtree_test.go
index 39307f4..9b8e02c 100644
--- a/pkg/rbtree/rbtree_test.go
+++ b/pkg/rbtree/rbtree_test.go
@@ -17,12 +17,13 @@ func checkTree[K constraints.Ordered, V any](t *testing.T, expectedSet map[K]str
 	// 3. Every nil is black.
 
 	// 4. If a node is red, then both its children are black.
-	tree.Walk(func(node *Node[V]) {
+	require.NoError(t, tree.Walk(func(node *Node[V]) error {
 		if node.getColor() == Red {
 			require.Equal(t, Black, node.Left.getColor())
 			require.Equal(t, Black, node.Right.getColor())
 		}
-	})
+		return nil
+	}))
 
 	// 5. For each node, all simple paths from the node to
 	//    descendent leaves contain the same number of black
@@ -39,7 +40,7 @@ func checkTree[K constraints.Ordered, V any](t *testing.T, expectedSet map[K]str
 		walkCnt(node.Left, cnt, leafFn)
 		walkCnt(node.Right, cnt, leafFn)
 	}
-	tree.Walk(func(node *Node[V]) {
+	require.NoError(t, tree.Walk(func(node *Node[V]) error {
 		var cnts []int
 		walkCnt(node, 0, func(cnt int) {
 			cnts = append(cnts, cnt)
@@ -50,7 +51,8 @@ func checkTree[K constraints.Ordered, V any](t *testing.T, expectedSet map[K]str
 				break
 			}
 		}
-	})
+		return nil
+	}))
 
 	// expected contents
 	expectedOrder := make([]K, 0, len(expectedSet))
@@ -64,9 +66,10 @@ func checkTree[K constraints.Ordered, V any](t *testing.T, expectedSet map[K]str
 		return expectedOrder[i] < expectedOrder[j]
 	})
 	actOrder := make([]K, 0, len(expectedSet))
-	tree.Walk(func(node *Node[V]) {
+	require.NoError(t, tree.Walk(func(node *Node[V]) error {
 		actOrder = append(actOrder, tree.KeyFn(node.Value))
-	})
+		return nil
+	}))
 	require.Equal(t, expectedOrder, actOrder)
 }
 
diff --git a/pkg/rbtree/range.go b/pkg/rbtree/rbtree_util.go
index 9f67078..db7fd32 100644
--- a/pkg/rbtree/range.go
+++ b/pkg/rbtree/rbtree_util.go
@@ -1,6 +1,8 @@
 package rbtree
 
 import (
+	"reflect"
+
 	"lukeshu.com/btrfs-tools/pkg/util"
 )
 
@@ -21,3 +23,26 @@ func (t *Tree[K, V]) SearchRange(fn func(V) int) []V {
 	}
 	return ret
 }
+
+func (t *Tree[K, V]) Equal(u *Tree[K, V]) bool {
+	if (t == nil) != (u == nil) {
+		return false
+	}
+	if t == nil {
+		return true
+	}
+
+	var tSlice []V
+	t.Walk(func(node *Node[V]) error {
+		tSlice = append(tSlice, node.Value)
+		return nil
+	})
+
+	var uSlice []V
+	u.Walk(func(node *Node[V]) error {
+		uSlice = append(uSlice, node.Value)
+		return nil
+	})
+
+	return reflect.DeepEqual(tSlice, uSlice)
+}