8327858: Improve spliterator and forEach for single-element immutable collections

Reviewed-by: smarks, vklang

Author: liach

src/java.base/share/classes/java/util/ImmutableCollections.java

@@ -32,7 +32,9 @@
 import java.io.ObjectStreamException;
 import java.io.Serializable;
 import java.lang.reflect.Array;
+import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.function.UnaryOperator;
@@ -656,6 +658,24 @@ public <T> T[] toArray(T[] a) {
             }
             return array;
         }
+
+        @Override
+        @SuppressWarnings("unchecked")
+        public void forEach(Consumer<? super E> action) {
+            action.accept(e0); // implicit null check
+            if (e1 != EMPTY) {
+                action.accept((E) e1);
+            }
+        }
+
+        @Override
+        public Spliterator<E> spliterator() {
+            if (e1 == EMPTY) {
+                return Collections.singletonSpliterator(e0);
+            } else {
+                return super.spliterator();
+            }
+        }
     }
 
     @jdk.internal.ValueBased
@@ -895,6 +915,26 @@ public <T> T[] toArray(T[] a) {
             }
             return array;
         }
+
+        @Override
+        @SuppressWarnings("unchecked")
+        public void forEach(Consumer<? super E> action) {
+            if (e1 == EMPTY) {
+                action.accept(e0); // implicit null check
+            } else {
+                action.accept(REVERSE ? (E)e1 : e0); // implicit null check
+                action.accept(REVERSE ? e0 : (E)e1);
+            }
+        }
+
+        @Override
+        public Spliterator<E> spliterator() {
+            if (e1 == EMPTY) {
+                return Collections.singletonSpliterator(e0);
+            } else {
+                return super.spliterator();
+            }
+        }
     }
 
 
@@ -1158,6 +1198,11 @@ private Object writeReplace() {
         public int hashCode() {
             return k0.hashCode() ^ v0.hashCode();
         }
+
+        @Override
+        public void forEach(BiConsumer<? super K, ? super V> action) {
+            action.accept(k0, v0); // implicit null check
+        }
     }
 
     /**

test/jdk/java/util/Collection/MOAT.java

@@ -26,7 +26,7 @@
  * @bug     6207984 6272521 6192552 6269713 6197726 6260652 5073546 4137464
  *          4155650 4216399 4294891 6282555 6318622 6355327 6383475 6420753
  *          6431845 4802633 6570566 6570575 6570631 6570924 6691185 6691215
- *          4802647 7123424 8024709 8193128
+ *          4802647 7123424 8024709 8193128 8327858
  * @summary Run many tests on many Collection and Map implementations
  * @author  Martin Buchholz
  * @modules java.base/java.util:open
@@ -479,6 +479,8 @@ private static <T> void testImmutableCollection(final Collection<T> c, T t) {
                    () -> c.removeAll(singleton(first)),
                    () -> c.retainAll(emptyList()));
         }
+        testForEachMatch(c);
+        testSpliteratorMatch(c);
     }
 
     private static <T> void testImmutableSeqColl(final SequencedCollection<T> c, T t) {
@@ -540,6 +542,39 @@ private static void testCollMutatorsAlwaysThrow(Collection<Integer> c) {
                 () -> c.retainAll(c));
     }
 
+    // Ensures forEach supplies in the same order as the iterator
+    private static <T> void testForEachMatch(Collection<T> c) {
+        var itr = c.iterator();
+        int[] index = {0};
+        c.forEach(item -> {
+            T itrNext = null;
+            if (!itr.hasNext() || !Objects.equals(itrNext = itr.next(), item)) {
+                fail("forEach and iterator mismatch at " + index[0] + " forEach: " + item + ", itr: " + itrNext);
+            }
+            index[0]++;
+        });
+        if (itr.hasNext()) {
+            fail("forEach and iterator mismatch at tail, extras in itr");
+        }
+    }
+
+    // Ensures spliterator returns in the same order as the iterator
+    private static <T> void testSpliteratorMatch(Collection<T> c) {
+        var itr = c.iterator();
+        var split = c.spliterator();
+        int[] index = {0};
+        split.forEachRemaining(item -> {
+            T itrNext = null;
+            if (!itr.hasNext() || !Objects.equals(itrNext = itr.next(), item)) {
+                fail("iterator and spliterator mismatch at " + index[0] + " spliterator: " + item + ", itr: " + itrNext);
+            }
+            index[0]++;
+        });
+        if (itr.hasNext()) {
+            fail("iterator and spliterator mismatch at tail, extra item in itr");
+        }
+    }
+
     /**
      * Test that calling a mutator always throws UOE, even if the mutator
      * wouldn't actually do anything on an empty collection.

test/micro/org/openjdk/bench/java/util/ImmutableColls.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019, 2020, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2019, 2024, 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
@@ -297,6 +297,50 @@ public void iterateSet(Blackhole bh, Set<String> coll) {
         }
     }
 
+    @Benchmark
+    @CompilerControl(CompilerControl.Mode.DONT_INLINE)
+    public void forEachOverSet(Blackhole bh) {
+        forEachSet(bh, fs4);
+        forEachSet(bh, s1);
+        forEachSet(bh, s3);
+        forEachSet(bh, fs2);
+        forEachSet(bh, s0);
+    }
+
+    public void forEachSet(Blackhole bh, Set<String> coll) {
+        coll.forEach(bh::consume);
+    }
+
+    @Benchmark
+    @CompilerControl(CompilerControl.Mode.DONT_INLINE)
+    public void iterateOverList(Blackhole bh) {
+        iterateList(bh, fl4);
+        iterateList(bh, fl1);
+        iterateList(bh, l3);
+        iterateList(bh, l0);
+        iterateList(bh, fl2);
+    }
+
+    public void iterateList(Blackhole bh, List<String> coll) {
+        for (String s : coll) {
+            bh.consume(s);
+        }
+    }
+
+    @Benchmark
+    @CompilerControl(CompilerControl.Mode.DONT_INLINE)
+    public void forEachOverList(Blackhole bh) {
+        forEachList(bh, fl4);
+        forEachList(bh, fl1);
+        forEachList(bh, l3);
+        forEachList(bh, l0);
+        forEachList(bh, fl2);
+    }
+
+    public void forEachList(Blackhole bh, List<String> coll) {
+        coll.forEach(bh::consume);
+    }
+
     @Benchmark
     @CompilerControl(CompilerControl.Mode.DONT_INLINE)
     public void toArrayFromMap(Blackhole bh) {