add binary search tools to lib/slices

2 files changed, 121 insertions, 32 deletions

diff --git a/lib/btrfs/io3_btree.go b/lib/btrfs/io3_btree.go
index ca528d7..73f9804 100644
--- a/lib/btrfs/io3_btree.go
+++ b/lib/btrfs/io3_btree.go
@@ -433,19 +433,11 @@ func (fs *FS) treeSearch(treeRoot TreeRoot, fn func(Key, uint32) int) (TreePath,
 			//
 			// There may or may not be a value that returns '0'.
 			//
-			// Implement this search as a binary search.
-			lastGood := -1
-			firstBad := len(node.Data.BodyInternal)
-			for firstBad > lastGood+1 {
-				midpoint := (lastGood + firstBad) / 2
-				direction := fn(node.Data.BodyInternal[midpoint].Key, math.MaxUint32)
-				if direction < 0 {
-					firstBad = midpoint
-				} else {
-					lastGood = midpoint
-				}
-			}
-			if lastGood < 0 {
+			// i.e. find the highest value that isn't too high.
+			lastGood, ok := slices.SearchHighest(node.Data.BodyInternal, func(kp KeyPointer) int {
+				return slices.Min(fn(kp.Key, math.MaxUint32), 0) // don't return >0; a key can't be "too low"
+			})
+			if !ok {
 				return TreePath{}, nil, iofs.ErrNotExist
 			}
 			path = path.Append(TreePathElem{
@@ -466,26 +458,16 @@ func (fs *FS) treeSearch(treeRoot TreeRoot, fn func(Key, uint32) int) (TreePath,
 			// is returned.
 			//
 			// Implement this search as a binary search.
-			beg := 0
-			end := len(node.Data.BodyLeaf)
-			for beg < end {
-				midpoint := (beg + end) / 2
-				direction := fn(
-					node.Data.BodyLeaf[midpoint].Key,
-					node.Data.BodyLeaf[midpoint].BodySize)
-				switch {
-				case direction < 0:
-					end = midpoint
-				case direction > 0:
-					beg = midpoint + 1
-				case direction == 0:
-					path = path.Append(TreePathElem{
-						ItemIdx: midpoint,
-					})
-					return path, node, nil
-				}
+			idx, ok := slices.Search(node.Data.BodyLeaf, func(item Item) int {
+				return fn(item.Key, item.BodySize)
+			})
+			if !ok {
+				return TreePath{}, nil, iofs.ErrNotExist
 			}
-			return TreePath{}, nil, iofs.ErrNotExist
+			path = path.Append(TreePathElem{
+				ItemIdx: idx,
+			})
+			return path, node, nil
 		}
 	}
 }
diff --git a/lib/slices/sliceutil.go b/lib/slices/sliceutil.go
index faaffcb..f95c7e7 100644
--- a/lib/slices/sliceutil.go
+++ b/lib/slices/sliceutil.go
@@ -73,3 +73,110 @@ func Sort[T constraints.Ordered](slice []T) {
 		return slice[i] < slice[j]
 	})
 }
+
+// returns (a+b)/2, but avoids overflow
+func avg(a, b int) int {
+	return int(uint(a+b) >> 1)
+}
+
+func min(a, b int) int {
+	if a < b {
+		return a
+	}
+	return b
+}
+
+func max(a, b int) int {
+	if a > b {
+		return a
+	}
+	return b
+}
+
+// Search the slice for a value for which `fn(slice[i]) = 0`.
+//
+//    + + + 0 0 0 - - -
+//          ^ ^ ^
+//          any of
+//
+// You can conceptualize `fn` as subtraction:
+//
+//     func(straw T) int {
+//         return needle - straw
+//     }
+func Search[T any](slice []T, fn func(T) int) (int, bool) {
+	beg, end := 0, len(slice)
+	for beg < end {
+		midpoint := avg(beg, end)
+		direction := fn(slice[midpoint])
+		switch {
+		case direction < 0:
+			end = midpoint
+		case direction > 0:
+			beg = midpoint + 1
+		case direction == 0:
+			return midpoint, true
+		}
+	}
+	return 0, false
+}
+
+// Search the slice for the left-most value for which `fn(slice[i]) = 0`.
+//
+//    + + + 0 0 0 - - -
+//          ^
+//
+// You can conceptualize `fn` as subtraction:
+//
+//     func(straw T) int {
+//         return needle - straw
+//     }
+func SearchLowest[T any](slice []T, fn func(T) int) (int, bool) {
+	lastBad, firstGood, firstBad := -1, len(slice), len(slice)
+	for lastBad+1 < min(firstGood, firstBad) {
+		midpoint := avg(lastBad, min(firstGood, firstBad))
+		direction := fn(slice[midpoint])
+		switch {
+		case direction < 0:
+			firstBad = midpoint
+		case direction > 0:
+			lastBad = midpoint
+		default:
+			firstGood = midpoint
+		}
+	}
+	if firstGood == len(slice) {
+		return 0, false
+	}
+	return firstGood, true
+}
+
+// Search the slice for the right-most value for which `fn(slice[i]) = 0`.
+//
+//    + + + 0 0 0 - - -
+//              ^
+//
+// You can conceptualize `fn` as subtraction:
+//
+//     func(straw T) int {
+//         return needle - straw
+//     }
+func SearchHighest[T any](slice []T, fn func(T) int) (int, bool) {
+	lastBad, lastGood, firstBad := -1, -1, len(slice)
+	for max(lastBad, lastGood)+1 < firstBad {
+		midpoint := avg(max(lastBad, lastGood), firstBad)
+		direction := fn(slice[midpoint])
+		switch {
+		case direction < 0:
+			firstBad = midpoint
+		case direction > 0:
+			lastBad = midpoint
+		default:
+			lastGood = midpoint
+		}
+	}
+	if lastGood < 0 {
+		return 0, false
+	}
+	return lastGood, true
+}