x/tools/container/intsets: add LowerBound

Fixes golang/go#21310.

Change-Id: Id3f23a66b9889a5087c1f83e7d672d14c41a59e3
Reviewed-on: https://go-review.googlesource.com/53432
Reviewed-by: Alan Donovan <adonovan@google.com>

container/intsets/sparse.go

@@ -170,6 +170,26 @@ func (b *block) min(take bool) int {
 	panic("BUG: empty block")
 }
 
+// lowerBound returns the smallest element of the block that is greater than or
+// equal to the element corresponding to the ith bit. If there is no such
+// element, the second return value is false.
+func (b *block) lowerBound(i uint) (int, bool) {
+	w := i / bitsPerWord
+	bit := i % bitsPerWord
+
+	if val := b.bits[w] >> bit; val != 0 {
+		return b.offset + int(i) + ntz(val), true
+	}
+
+	for w++; w < wordsPerBlock; w++ {
+		if val := b.bits[w]; val != 0 {
+			return b.offset + int(w*bitsPerWord) + ntz(val), true
+		}
+	}
+
+	return 0, false
+}
+
 // forEach calls f for each element of block b.
 // f must not mutate b's enclosing Sparse.
 func (b *block) forEach(f func(int)) {
@@ -280,8 +300,26 @@ func (s *Sparse) Min() int {
 	return s.root.min(false)
 }
 
+// LowerBound returns the smallest element >= x, or MaxInt if there is no such
+// element.
+func (s *Sparse) LowerBound(x int) int {
+	offset, i := offsetAndBitIndex(x)
+	for b := s.first(); b != &none; b = s.next(b) {
+		if b.offset > offset {
+			return b.min(false)
+		}
+		if b.offset == offset {
+			if y, ok := b.lowerBound(i); ok {
+				return y
+			}
+		}
+	}
+	return MaxInt
+}
+
 // block returns the block that would contain offset,
 // or nil if s contains no such block.
+// Precondition: offset is a multiple of bitsPerBlock.
 func (s *Sparse) block(offset int) *block {
 	for b := s.first(); b != &none && b.offset <= offset; b = s.next(b) {
 		if b.offset == offset {

container/intsets/sparse_test.go

@@ -275,11 +275,30 @@ func TestRandomMutations(t *testing.T) {
 	}
 }
 
+func TestLowerBound(t *testing.T) {
+	// Use random sets of sizes from 0 to about 4000.
+	prng := rand.New(rand.NewSource(0))
+	for i := uint(0); i < 12; i++ {
+		x := randomPset(prng, 1<<i)
+		for j := 0; j < 10000; j++ {
+			found := intsets.MaxInt
+			for e := range x.hash {
+				if e >= j && e < found {
+					found = e
+				}
+			}
+			if res := x.bits.LowerBound(j); res != found {
+				t.Errorf("%s: LowerBound(%d)=%d, expected %d", &x.bits, j, res, found)
+			}
+		}
+	}
+}
+
 // TestSetOperations exercises classic set operations: ∩ , ∪, \.
 func TestSetOperations(t *testing.T) {
 	prng := rand.New(rand.NewSource(0))
 
-	// Use random sets of sizes from 0 to about 1000.
+	// Use random sets of sizes from 0 to about 4000.
 	// For each operator, we test variations such as
 	// Z.op(X, Y), Z.op(X, Z) and Z.op(Z, Y) to exercise
 	// the degenerate cases of each method implementation.