From a6010402746a2b9ab3f21a96fe2c2d7e64b1856b Mon Sep 17 00:00:00 2001
From: Julian Hyde <jhyde@apache.org>
Date: Mon, 13 Jan 2025 23:49:07 -0800
Subject: [PATCH] Add method ImmutableBitSet.stream()

---
 .../apache/calcite/util/ImmutableBitSet.java  | 18 +++++++++++++++-
 .../calcite/util/ImmutableBitSetTest.java     | 21 +++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/core/src/main/java/org/apache/calcite/util/ImmutableBitSet.java b/core/src/main/java/org/apache/calcite/util/ImmutableBitSet.java
index cd688476f8f4..29aac2d42120 100644
--- a/core/src/main/java/org/apache/calcite/util/ImmutableBitSet.java
+++ b/core/src/main/java/org/apache/calcite/util/ImmutableBitSet.java
@@ -48,6 +48,7 @@
 import java.util.function.IntConsumer;
 import java.util.function.IntPredicate;
 import java.util.stream.Collector;
+import java.util.stream.IntStream;
 
 import static org.apache.calcite.linq4j.Nullness.castNonNull;
 
@@ -76,7 +77,7 @@ public class ImmutableBitSet
   public static final Ordering<ImmutableBitSet> ORDERING =
       Ordering.from(COMPARATOR);
 
-  // BitSets are packed into arrays of "words."  Currently a word is
+  // BitSets are packed into arrays of "words."  Currently, a word is
   // a long, which consists of 64 bits, requiring 6 address bits.
   // The choice of word size is determined purely by performance concerns.
   private static final int ADDRESS_BITS_PER_WORD = 6;
@@ -142,6 +143,8 @@ public static ImmutableBitSet of(int... bits) {
     return new ImmutableBitSet(words);
   }
 
+  /** Creates an ImmutableBitSet whose contents are the bits denoted by a
+   * given collection of integers. */
   public static ImmutableBitSet of(Iterable<Integer>  bits) {
     if (bits instanceof ImmutableBitSet) {
       return (ImmutableBitSet) bits;
@@ -514,6 +517,19 @@ public int size() {
     }
   }
 
+  /**
+   * Returns a stream of indices for which this {@code ImmutableBitSet}
+   * contains a bit in the set state. The indices are returned
+   * in order, from lowest to highest. The size of the stream
+   * is the number of bits in the set state, equal to the value
+   * returned by the {@link #cardinality()} method.
+   *
+   * @return a stream of integers representing set indices
+   */
+  public IntStream stream() {
+    return toList().stream().mapToInt(i -> i);
+  }
+
   /**
    * Returns the index of the first bit that is set to {@code true}
    * that occurs on or after the specified starting index. If no such
diff --git a/core/src/test/java/org/apache/calcite/util/ImmutableBitSetTest.java b/core/src/test/java/org/apache/calcite/util/ImmutableBitSetTest.java
index 260af704c55b..5e64c555d81e 100644
--- a/core/src/test/java/org/apache/calcite/util/ImmutableBitSetTest.java
+++ b/core/src/test/java/org/apache/calcite/util/ImmutableBitSetTest.java
@@ -38,6 +38,7 @@
 import java.util.function.Consumer;
 import java.util.function.IntConsumer;
 import java.util.function.IntPredicate;
+import java.util.stream.Collectors;
 
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.CoreMatchers.is;
@@ -123,6 +124,16 @@ private void assertToIterBitSet(String expected, ImmutableBitSet bitSet) {
       buf.append(i);
     }
     assertThat(buf, hasToString(expected));
+
+    // Now check that bitSet.stream() does the same as bitSet.iterator().
+    buf.setLength(0);
+    bitSet.stream().forEach(i -> {
+      if (buf.length() > 0) {
+        buf.append(", ");
+      }
+      buf.append(i);
+    });
+    assertThat(buf, hasToString(expected));
   }
 
   /**
@@ -133,6 +144,16 @@ private void assertToIterBitSet(String expected, ImmutableBitSet bitSet) {
     check((bitSet, list) -> assertThat(bitSet.toList(), equalTo(list)));
   }
 
+  /**
+   * Tests the method
+   * {@link org.apache.calcite.util.ImmutableBitSet#stream()}.
+   */
+  @Test void testStream() {
+    check((bitSet, list) ->
+        assertThat(bitSet.stream().boxed().collect(Collectors.toList()),
+            equalTo(list)));
+  }
+
   /**
    * Tests the method
    * {@link org.apache.calcite.util.ImmutableBitSet#forEachInt}.