putil: Add BitArray::find_off_range()

panda/src/putil/bitArray.cxx

@@ -443,6 +443,42 @@ get_next_higher_different_bit(int low_bit) const {
   return w2 * num_bits_per_word + b2;
 }
 
+/**
+ * Returns the index of the first off bit of a range of off bits of the given
+ * size.  This is useful to find a free range of off bits in the array.
+ *
+ * If low_bit is not 0, it indicates which bit to start looking at.
+ */
+int BitArray::
+find_off_range(int size, int low_bit) {
+  if (_array.empty()) {
+    return _highest_bits ? -1 : low_bit;
+  }
+  if (size >= 0) {
+    low_bit = std::max(low_bit, _array[0].get_lowest_off_bit());
+    if (get_bit(low_bit)) {
+      // Make sure we're positioned at an off bit.
+      low_bit = get_next_higher_different_bit(low_bit);
+    }
+
+    size_t num_bits = get_num_bits();
+
+    while (size > 1 && low_bit >= 0 && has_any_of(low_bit, size)) {
+      // Not enough bits free, try the next open range.
+      int next_bit = get_next_higher_different_bit(low_bit);
+      if (next_bit <= low_bit || (_highest_bits && (size_t)next_bit >= num_bits)) {
+        return -1;
+      }
+      low_bit = get_next_higher_different_bit(next_bit);
+
+      if ((size_t)low_bit >= num_bits) {
+        return _highest_bits ? -1 : low_bit;
+      }
+    }
+  }
+  return low_bit;
+}
+
 /**
  * Inverts all the bits in the BitArray.  This is equivalent to array =
  * ~array.

panda/src/putil/bitArray.h

@@ -86,6 +86,8 @@ PUBLISHED:
   int get_highest_off_bit() const;
   int get_next_higher_different_bit(int low_bit) const;
 
+  int find_off_range(int size, int low_bit = 0);
+
   INLINE size_t get_num_words() const;
   INLINE MaskType get_word(size_t n) const;
   INLINE void set_word(size_t n, WordType value);

tests/putil/test_bitarray.py

@@ -140,3 +140,42 @@ def test_bitarray_has_any_of():
     assert ba.has_any_of(0, 1)
     assert ba.has_any_of(53, 45)
     assert ba.has_any_of(0, 100)
+
+
+def test_bitarray_find_off_range():
+    ba = BitArray()
+    assert ba.find_off_range(1, 0) == 0
+    assert ba.find_off_range(100, 0) == 0
+    assert ba.find_off_range(1, 150) == 150
+    assert ba.find_off_range(100, 150) == 150
+
+    ba = BitArray(0b1)
+    assert ba.find_off_range(1, 0) == 1
+    assert ba.find_off_range(1, 1) == 1
+    assert ba.find_off_range(2, 0) == 1
+    assert ba.find_off_range(2, 1) == 1
+    assert ba.find_off_range(2, 2) == 2
+
+    ba = BitArray(0b1000100)
+    assert ba.find_off_range(2) == 0
+    assert ba.find_off_range(3, 0) == 3
+    assert ba.find_off_range(3, 1) == 3
+    assert ba.find_off_range(3, 2) == 3
+    assert ba.find_off_range(3, 3) == 3
+    assert ba.find_off_range(3, 4) == 7
+    assert ba.find_off_range(4) == 7
+
+    ba = BitArray(3)
+    ba.invert_in_place()
+    assert ba.find_off_range(2) == 0
+    assert ba.find_off_range(3) == -1
+
+    ba = BitArray()
+    ba.set_range(0, 64)
+    assert ba.find_off_range(2) == 64
+
+    ba = BitArray()
+    ba.set_range(1, 63)
+    ba.invert_in_place()
+    assert ba.find_off_range(63) == 1
+    assert ba.find_off_range(64) == -1