8303621: BitMap::iterate should support lambdas and other function objects

Reviewed-by: aboldtch, tschatzl, stefank

src/hotspot/share/c1/c1_LinearScan.cpp

@@ -1329,10 +1329,9 @@ void LinearScan::build_intervals() {
     assert(block_to   == instructions->at(instructions->length() - 1)->id(), "must be");
 
     // Update intervals for registers live at the end of this block;
-    ResourceBitMap live = block->live_out();
-    int size = (int)live.size();
-    for (int number = (int)live.get_next_one_offset(0, size); number < size; number = (int)live.get_next_one_offset(number + 1, size)) {
-      assert(live.at(number), "should not stop here otherwise");
+    ResourceBitMap& live = block->live_out();
+    auto updater = [&](BitMap::idx_t index) {
+      int number = static_cast<int>(index);
       assert(number >= LIR_Opr::vreg_base, "fixed intervals must not be live on block bounds");
       TRACE_LINEAR_SCAN(2, tty->print_cr("live in %d to %d", number, block_to + 2));
 
@@ -1347,7 +1346,8 @@ void LinearScan::build_intervals() {
           is_interval_in_loop(number, block->loop_index())) {
         interval_at(number)->add_use_pos(block_to + 1, loopEndMarker);
       }
-    }
+    };
+    live.iterate(updater);
 
     // iterate all instructions of the block in reverse order.
     // skip the first instruction because it is always a label
@@ -1738,11 +1738,10 @@ Interval* LinearScan::interval_at_op_id(int reg_num, int op_id) {
 void LinearScan::resolve_collect_mappings(BlockBegin* from_block, BlockBegin* to_block, MoveResolver &move_resolver) {
   DEBUG_ONLY(move_resolver.check_empty());
 
-  const int size = live_set_size();
-  const ResourceBitMap live_at_edge = to_block->live_in();
-
   // visit all registers where the live_at_edge bit is set
-  for (int r = (int)live_at_edge.get_next_one_offset(0, size); r < size; r = (int)live_at_edge.get_next_one_offset(r + 1, size)) {
+  const ResourceBitMap& live_at_edge = to_block->live_in();
+  auto visitor = [&](BitMap::idx_t index) {
+    int r = static_cast<int>(index);
     assert(r < num_virtual_regs(), "live information set for not existing interval");
     assert(from_block->live_out().at(r) && to_block->live_in().at(r), "interval not live at this edge");
 
@@ -1753,7 +1752,8 @@ void LinearScan::resolve_collect_mappings(BlockBegin* from_block, BlockBegin* to
       // need to insert move instruction
       move_resolver.add_mapping(from_interval, to_interval);
     }
-  }
+  };
+  live_at_edge.iterate(visitor, 0, live_set_size());
 }
 
 
@@ -1913,10 +1913,11 @@ void LinearScan::resolve_exception_entry(BlockBegin* block, MoveResolver &move_r
   DEBUG_ONLY(move_resolver.check_empty());
 
   // visit all registers where the live_in bit is set
-  int size = live_set_size();
-  for (int r = (int)block->live_in().get_next_one_offset(0, size); r < size; r = (int)block->live_in().get_next_one_offset(r + 1, size)) {
+  auto resolver = [&](BitMap::idx_t index) {
+    int r = static_cast<int>(index);
     resolve_exception_entry(block, r, move_resolver);
-  }
+  };
+  block->live_in().iterate(resolver, 0, live_set_size());
 
   // the live_in bits are not set for phi functions of the xhandler entry, so iterate them separately
   for_each_phi_fun(block, phi,
@@ -1986,10 +1987,11 @@ void LinearScan::resolve_exception_edge(XHandler* handler, int throwing_op_id, M
 
   // visit all registers where the live_in bit is set
   BlockBegin* block = handler->entry_block();
-  int size = live_set_size();
-  for (int r = (int)block->live_in().get_next_one_offset(0, size); r < size; r = (int)block->live_in().get_next_one_offset(r + 1, size)) {
+  auto resolver = [&](BitMap::idx_t index) {
+    int r = static_cast<int>(index);
     resolve_exception_edge(handler, throwing_op_id, r, NULL, move_resolver);
-  }
+  };
+  block->live_in().iterate(resolver, 0, live_set_size());
 
   // the live_in bits are not set for phi functions of the xhandler entry, so iterate them separately
   for_each_phi_fun(block, phi,
@@ -3466,10 +3468,11 @@ void LinearScan::verify_constants() {
 
   for (int i = 0; i < num_blocks; i++) {
     BlockBegin* block = block_at(i);
-    ResourceBitMap live_at_edge = block->live_in();
+    ResourceBitMap& live_at_edge = block->live_in();
 
     // visit all registers where the live_at_edge bit is set
-    for (int r = (int)live_at_edge.get_next_one_offset(0, size); r < size; r = (int)live_at_edge.get_next_one_offset(r + 1, size)) {
+    auto visitor = [&](BitMap::idx_t index) {
+      int r = static_cast<int>(index);
       TRACE_LINEAR_SCAN(4, tty->print("checking interval %d of block B%d", r, block->block_id()));
 
       Value value = gen()->instruction_for_vreg(r);
@@ -3478,7 +3481,8 @@ void LinearScan::verify_constants() {
       assert(value->operand()->is_register() && value->operand()->is_virtual(), "value must have virtual operand");
       assert(value->operand()->vreg_number() == r, "register number must match");
       // TKR assert(value->as_Constant() == NULL || value->is_pinned(), "only pinned constants can be alive across block boundaries");
-    }
+    };
+    live_at_edge.iterate(visitor, 0, size);
   }
 }
 

src/hotspot/share/gc/g1/g1CardSetContainers.inline.hpp

@@ -249,11 +249,7 @@ inline G1AddCardResult G1CardSetBitMap::add(uint card_idx, size_t threshold, siz
 template <class CardVisitor>
 inline void G1CardSetBitMap::iterate(CardVisitor& found, size_t size_in_bits, uint offset) {
   BitMapView bm(_bits, size_in_bits);
-  BitMap::idx_t idx = bm.get_next_one_offset(0);
-  while (idx != size_in_bits) {
-    found((offset | (uint)idx));
-    idx = bm.get_next_one_offset(idx + 1);
-  }
+  bm.iterate([&](BitMap::idx_t idx) { found(offset | (uint)idx); });
 }
 
 inline size_t G1CardSetBitMap::header_size_in_bytes() {

src/hotspot/share/utilities/bitMap.hpp

@@ -105,6 +105,8 @@ class BitMap {
   static const bm_word_t find_ones_flip = 0;
   static const bm_word_t find_zeros_flip = ~(bm_word_t)0;
 
+  template<typename ReturnType> struct IterateInvoker;
+
   // Threshold for performing small range operation, even when large range
   // operation was requested. Measured in words.
   static const size_t small_range_words = 32;
@@ -252,19 +254,38 @@ class BitMap {
   // Verify [beg,end) is a valid range, e.g. beg <= end <= size().
   void verify_range(idx_t beg, idx_t end) const NOT_DEBUG_RETURN;
 
-  // Iteration support.  Applies the closure to the index for each set bit,
-  // starting from the least index in the range to the greatest, in order.
-  // The iteration terminates if the closure returns false.  Returns true if
-  // the iteration completed, false if terminated early because the closure
-  // returned false.  If the closure modifies the bitmap, modifications to
-  // bits at indices greater than the current index will affect which further
-  // indices the closure will be applied to.
-  // precondition: beg and end form a valid range.
-  template <class BitMapClosureType>
-  bool iterate(BitMapClosureType* cl, idx_t beg, idx_t end);
+  // Applies an operation to the index of each set bit in [beg, end), in
+  // increasing order.
+  //
+  // If i is an index of the bitmap, the operation is either
+  // - function(i)
+  // - cl->do_bit(i)
+  // The result of an operation must be either void or convertible to bool.
+  //
+  // If an operation returns false then the iteration stops at that index.
+  // The result of the iteration is true unless the iteration was stopped by
+  // an operation returning false.
+  //
+  // If an operation modifies the bitmap, modifications to bits at indices
+  // greater than the current index will affect which further indices the
+  // operation will be applied to.
+  //
+  // precondition: beg and end form a valid range for the bitmap.
+  template<typename Function>
+  bool iterate(Function function, idx_t beg, idx_t end) const;
 
-  template <class BitMapClosureType>
-  bool iterate(BitMapClosureType* cl);
+  template<typename BitMapClosureType>
+  bool iterate(BitMapClosureType* cl, idx_t beg, idx_t end) const;
+
+  template<typename Function>
+  bool iterate(Function function) const {
+    return iterate(function, 0, size());
+  }
+
+  template<typename BitMapClosureType>
+  bool iterate(BitMapClosureType* cl) const {
+    return iterate(cl, 0, size());
+  }
 
   // Looking for 1's and 0's at indices equal to or greater than "beg",
   // stopping if none has been found before "end", and returning

src/hotspot/share/utilities/bitMap.inline.hpp

@@ -243,21 +243,47 @@ BitMap::get_next_one_offset_aligned_right(idx_t beg, idx_t end) const {
   return get_next_bit_impl<find_ones_flip, true>(beg, end);
 }
 
-template <typename BitMapClosureType>
-inline bool BitMap::iterate(BitMapClosureType* cl, idx_t beg, idx_t end) {
+// IterateInvoker supports conditionally stopping iteration early.  The
+// invoker is called with the function to apply to each set index, along with
+// the current index.  If the function returns void then the invoker always
+// returns true, so no early stopping.  Otherwise, the result of the function
+// is returned by the invoker.  Iteration stops early if conversion of that
+// result to bool is false.
+
+template<typename ReturnType>
+struct BitMap::IterateInvoker {
+  template<typename Function>
+  bool operator()(Function function, idx_t index) const {
+    return function(index);     // Stop early if converting to bool is false.
+  }
+};
+
+template<>
+struct BitMap::IterateInvoker<void> {
+  template<typename Function>
+  bool operator()(Function function, idx_t index) const {
+    function(index);            // Result is void.
+    return true;                // Never stop early.
+  }
+};
+
+template <typename Function>
+inline bool BitMap::iterate(Function function, idx_t beg, idx_t end) const {
+  auto invoke = IterateInvoker<decltype(function(beg))>();
   for (idx_t index = beg; true; ++index) {
     index = get_next_one_offset(index, end);
     if (index >= end) {
       return true;
-    } else if (!cl->do_bit(index)) {
+    } else if (!invoke(function, index)) {
       return false;
     }
   }
 }
 
 template <typename BitMapClosureType>
-inline bool BitMap::iterate(BitMapClosureType* cl) {
-  return iterate(cl, 0, size());
+inline bool BitMap::iterate(BitMapClosureType* cl, idx_t beg, idx_t end) const {
+  auto function = [&](idx_t index) { return cl->do_bit(index); };
+  return iterate(function, beg, end);
 }
 
 // Returns a bit mask for a range of bits [beg, end) within a single word.  Each

test/hotspot/gtest/utilities/test_bitMap_iterate.cpp

@@ -0,0 +1,213 @@
+/*
+ * Copyright (c) 2023, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ *
+ */
+
+#include "precompiled.hpp"
+#include "utilities/align.hpp"
+#include "utilities/bitMap.inline.hpp"
+#include "utilities/debug.hpp"
+#include "utilities/globalDefinitions.hpp"
+#include "unittest.hpp"
+
+using idx_t = BitMap::idx_t;
+using bm_word_t = BitMap::bm_word_t;
+
+static const idx_t BITMAP_SIZE = 1024;
+static const idx_t BITMAP_WORD_SIZE = align_up(BITMAP_SIZE, BitsPerWord) / BitsPerWord;
+
+
+static void test_iterate_step(const BitMap& map,
+                              idx_t index,
+                              const idx_t* positions,
+                              size_t positions_index,
+                              size_t positions_size) {
+  ASSERT_LT(positions_index, positions_size);
+  ASSERT_EQ(index, positions[positions_index]);
+  ASSERT_TRUE(map.at(index));
+}
+
+// Test lambda returning void.
+static void test_iterate_lambda(const BitMap& map,
+                                const idx_t* positions,
+                                size_t positions_size) {
+  SCOPED_TRACE("iterate with lambda");
+  size_t positions_index = 0;
+  auto f = [&](idx_t i) {
+    test_iterate_step(map, i, positions, positions_index++, positions_size);
+  };
+  ASSERT_TRUE(map.iterate(f));
+  ASSERT_EQ(positions_index, positions_size);
+}
+
+// Test closure returning bool.  Also tests lambda returning bool.
+static void test_iterate_closure(const BitMap& map,
+                                 const idx_t* positions,
+                                 size_t positions_size) {
+  SCOPED_TRACE("iterate with BitMapClosure");
+  struct Closure : public BitMapClosure {
+    const BitMap& _map;
+    const idx_t* _positions;
+    size_t _positions_index;
+    size_t _positions_size;
+
+    Closure(const BitMap& map, const idx_t* positions, size_t positions_size)
+      : _map(map),
+        _positions(positions),
+        _positions_index(0),
+        _positions_size(positions_size)
+    {}
+
+    bool do_bit(idx_t i) override {
+      test_iterate_step(_map, i, _positions, _positions_index++, _positions_size);
+      return true;
+    }
+  } closure{map, positions, positions_size};
+  ASSERT_TRUE(map.iterate(&closure));
+  ASSERT_EQ(closure._positions_index, positions_size);
+}
+
+// Test closure returning void.  Also tests lambda returning bool.
+static void test_iterate_non_closure(const BitMap& map,
+                                     const idx_t* positions,
+                                     size_t positions_size) {
+  SCOPED_TRACE("iterate with non-BitMapClosure");
+  struct Closure {
+    const BitMap& _map;
+    const idx_t* _positions;
+    size_t _positions_index;
+    size_t _positions_size;
+
+    Closure(const BitMap& map, const idx_t* positions, size_t positions_size)
+      : _map(map),
+        _positions(positions),
+        _positions_index(0),
+        _positions_size(positions_size)
+    {}
+
+    void do_bit(idx_t i) {
+      test_iterate_step(_map, i, _positions, _positions_index++, _positions_size);
+    }
+  } closure{map, positions, positions_size};
+  ASSERT_TRUE(map.iterate(&closure));
+  ASSERT_EQ(closure._positions_index, positions_size);
+}
+
+static void fill_iterate_map(BitMap& map,
+                             const idx_t* positions,
+                             size_t positions_size) {
+  map.clear_range(0, map.size());
+  for (size_t i = 0; i < positions_size; ++i) {
+    map.set_bit(positions[i]);
+  }
+}
+
+static void test_iterate(BitMap& map,
+                         const idx_t* positions,
+                         size_t positions_size) {
+  fill_iterate_map(map, positions, positions_size);
+  test_iterate_lambda(map, positions, positions_size);
+  test_iterate_closure(map, positions, positions_size);
+  test_iterate_non_closure(map, positions, positions_size);
+}
+
+TEST(BitMap, iterate_empty) {
+  bm_word_t test_data[BITMAP_WORD_SIZE];
+  BitMapView test_map{test_data, BITMAP_SIZE};
+  idx_t positions[1] = {};
+  test_iterate(test_map, positions, 0);
+}
+
+TEST(BitMap, iterate_with_endpoints) {
+  bm_word_t test_data[BITMAP_WORD_SIZE];
+  BitMapView test_map{test_data, BITMAP_SIZE};
+  idx_t positions[] = { 0, 2, 6, 31, 61, 131, 247, 578, BITMAP_SIZE - 1 };
+  test_iterate(test_map, positions, ARRAY_SIZE(positions));
+}
+
+TEST(BitMap, iterate_without_endpoints) {
+  bm_word_t test_data[BITMAP_WORD_SIZE];
+  BitMapView test_map{test_data, BITMAP_SIZE};
+  idx_t positions[] = { 1, 2, 6, 31, 61, 131, 247, 578, BITMAP_SIZE - 2 };
+  test_iterate(test_map, positions, ARRAY_SIZE(positions));
+}
+
+TEST(BitMap, iterate_full) {
+  bm_word_t test_data[BITMAP_WORD_SIZE];
+  BitMapView test_map{test_data, BITMAP_SIZE};
+  static idx_t positions[BITMAP_SIZE]; // static to avoid large stack allocation.
+  for (idx_t i = 0; i < BITMAP_SIZE; ++i) {
+    positions[i] = i;
+  }
+  test_iterate(test_map, positions, ARRAY_SIZE(positions));
+}
+
+TEST(BitMap, iterate_early_termination) {
+  bm_word_t test_data[BITMAP_WORD_SIZE];
+  BitMapView test_map{test_data, BITMAP_SIZE};
+  idx_t positions[] = { 1, 2, 6, 31, 61, 131, 247, 578, BITMAP_SIZE - 2 };
+  size_t positions_size = ARRAY_SIZE(positions);
+  size_t positions_index = 0;
+  fill_iterate_map(test_map, positions, positions_size);
+  idx_t stop_at = 131;
+  auto f = [&](idx_t i) {
+    test_iterate_step(test_map, i, positions, positions_index, positions_size);
+    if (positions[positions_index] == stop_at) {
+      return false;
+    } else {
+      positions_index += 1;
+      return true;
+    }
+  };
+  ASSERT_FALSE(test_map.iterate(f));
+  ASSERT_LT(positions_index, positions_size);
+  ASSERT_EQ(positions[positions_index], stop_at);
+
+  struct Closure : public BitMapClosure {
+    const BitMap& _map;
+    const idx_t* _positions;
+    size_t _positions_index;
+    size_t _positions_size;
+    idx_t _stop_at;
+
+    Closure(const BitMap& map, const idx_t* positions, size_t positions_size, idx_t stop_at)
+      : _map(map),
+        _positions(positions),
+        _positions_index(0),
+        _positions_size(positions_size),
+        _stop_at(stop_at)
+    {}
+
+    bool do_bit(idx_t i) override {
+      test_iterate_step(_map, i, _positions, _positions_index, _positions_size);
+      if (_positions[_positions_index] == _stop_at) {
+        return false;
+      } else {
+        _positions_index += 1;
+        return true;
+      }
+    }
+  } closure{test_map, positions, positions_size, stop_at};
+  ASSERT_FALSE(test_map.iterate(&closure));
+  ASSERT_LT(closure._positions_index, positions_size);
+  ASSERT_EQ(positions[closure._positions_index], stop_at);
+}