8359936: StableValues can release the underlying function after complete computation

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

@@ -45,6 +45,7 @@
 import jdk.internal.access.SharedSecrets;
 import jdk.internal.lang.stable.StableUtil;
 import jdk.internal.lang.stable.StableValueImpl;
+import jdk.internal.lang.stable.UnderlyingHolder;
 import jdk.internal.misc.CDS;
 import jdk.internal.util.ArraysSupport;
 import jdk.internal.vm.annotation.ForceInline;
@@ -137,11 +138,11 @@ public <E> List<E> listFromTrustedArrayNullsAllowed(Object[] array) {
                 }
                 public <E> List<E> stableList(int size, IntFunction<? extends E> mapper) {
                     // A stable list is not Serializable, so we cannot return `List.of()` if `size == 0`
-                    return new StableList<>(size, mapper);
+                    return new StableList<>(size, new UnderlyingHolder<>(mapper, size));
                 }
                 public <K, V> Map<K, V> stableMap(Set<K> keys, Function<? super K, ? extends V> mapper) {
                     // A stable map is not Serializable, so we cannot return `Map.of()` if `keys.isEmpty()`
-                    return new StableMap<>(keys, mapper);
+                    return new StableMap<>(keys, new UnderlyingHolder<>(mapper, keys.size()));
                 }
             });
         }
@@ -794,13 +795,13 @@ static final class StableList<E>
             extends AbstractImmutableList<E>
             implements HasStableDelegates<E> {
 
-        @Stable
-        private final IntFunction<? extends E> mapper;
         @Stable
         final StableValueImpl<E>[] delegates;
+        @Stable
+        private final UnderlyingHolder<IntFunction<? extends E>> underlyingHolder;
 
-        StableList(int size, IntFunction<? extends E> mapper) {
-            this.mapper = mapper;
+        StableList(int size, UnderlyingHolder<IntFunction<? extends E>> underlyingHolder) {
+            this.underlyingHolder = underlyingHolder;
             this.delegates = StableUtil.array(size);
         }
 
@@ -818,7 +819,7 @@ public E get(int i) {
                 throw new IndexOutOfBoundsException(i);
             }
             return delegate.orElseSet(new Supplier<E>() {
-                        @Override  public E get() { return mapper.apply(i); }});
+                        @Override  public E get() { return underlyingHolder.underlying().apply(i); }}, underlyingHolder);
         }
 
         @Override
@@ -1606,14 +1607,14 @@ private Object writeReplace() {
     static final class StableMap<K, V>
             extends AbstractImmutableMap<K, V> {
 
-        @Stable
-        private final Function<? super K, ? extends V> mapper;
         @Stable
         private final Map<K, StableValueImpl<V>> delegate;
+        @Stable
+        private final UnderlyingHolder<Function<? super K, ? extends V>> underlyingHolder;
 
-        StableMap(Set<K> keys, Function<? super K, ? extends V> mapper) {
-            this.mapper = mapper;
+        StableMap(Set<K> keys, UnderlyingHolder<Function<? super K, ? extends V>> underlyingHolder) {
             this.delegate = StableUtil.map(keys);
+            this.underlyingHolder = underlyingHolder;
         }
 
         @Override public boolean              containsKey(Object o) { return delegate.containsKey(o); }
@@ -1636,7 +1637,7 @@ public V getOrDefault(Object key, V defaultValue) {
             @SuppressWarnings("unchecked")
             final K k = (K) key;
             return stable.orElseSet(new Supplier<V>() {
-                @Override public V get() { return mapper.apply(k); }});
+                @Override public V get() { return underlyingHolder.underlying().apply(k); }}, underlyingHolder);
         }
 
         @jdk.internal.ValueBased
@@ -1692,7 +1693,7 @@ public Entry<K, V> next() {
                     final Map.Entry<K, StableValueImpl<V>> inner = delegateIterator.next();
                     final K k = inner.getKey();
                     return new StableEntry<>(k, inner.getValue(), new Supplier<V>() {
-                        @Override public V get() { return outer.outer.mapper.apply(k); }});
+                        @Override public V get() { return outer.outer.underlyingHolder.underlying().apply(k); }}, outer.outer.underlyingHolder);
                 }
 
                 @Override
@@ -1703,7 +1704,7 @@ public void forEachRemaining(Consumer<? super Map.Entry<K, V>> action) {
                                 public void accept(Entry<K, StableValueImpl<V>> inner) {
                                     final K k = inner.getKey();
                                     action.accept(new StableEntry<>(k, inner.getValue(), new Supplier<V>() {
-                                        @Override public V get() { return outer.outer.mapper.apply(k); }}));
+                                        @Override public V get() { return outer.outer.underlyingHolder.underlying().apply(k); }}, outer.outer.underlyingHolder));
                                 }
                             };
                     delegateIterator.forEachRemaining(innerAction);
@@ -1719,10 +1720,11 @@ private static  <K, V> LazyMapIterator<K, V> of(StableMapEntrySet<K, V> outer) {
 
         private record StableEntry<K, V>(K getKey, // trick
                                          StableValueImpl<V> stableValue,
-                                         Supplier<? extends V> supplier) implements Map.Entry<K, V> {
+                                         Supplier<? extends V> supplier,
+                                         UnderlyingHolder<?> underlyingHolder) implements Map.Entry<K, V> {
 
             @Override public V setValue(V value) { throw uoe(); }
-            @Override public V getValue() { return stableValue.orElseSet(supplier); }
+            @Override public V getValue() { return stableValue.orElseSet(supplier, underlyingHolder); }
             @Override public int hashCode() { return hash(getKey()) ^ hash(getValue()); }
             @Override public String toString() { return getKey() + "=" + stableValue.toString(); }
             @Override public boolean equals(Object o) {

src/java.base/share/classes/jdk/internal/lang/stable/StableEnumFunction.java

@@ -33,7 +33,6 @@
 import java.util.ArrayList;
 import java.util.BitSet;
 import java.util.Collection;
-import java.util.EnumSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Function;
@@ -46,20 +45,21 @@
  * @implNote This implementation can be used early in the boot sequence as it does not
  *           rely on reflection, MethodHandles, Streams etc.
  *
- * @param enumType     the class type of the Enum
- * @param firstOrdinal the lowest ordinal used
- * @param member       an int predicate that can be used to test if an enum is a member
- *                     of the valid inputs (as there might be "holes")
- * @param delegates    a delegate array of inputs to StableValue mappings
- * @param original     the original Function
- * @param <E>          the type of the input to the function
- * @param <R>          the type of the result of the function
+ * @param enumType         the class type of the Enum
+ * @param firstOrdinal     the lowest ordinal used
+ * @param member           an int predicate that can be used to test if an enum is a member
+ *                         of the valid inputs (as there might be "holes")
+ * @param delegates        a delegate array of inputs to StableValue mappings
+ * @param underlyingHolder of the original underlying Function
+ * @param <E>              the type of the input to the function
+ * @param <R>              the type of the result of the function
  */
 public record StableEnumFunction<E extends Enum<E>, R>(Class<E> enumType,
                                                        int firstOrdinal,
                                                        IntPredicate member,
                                                        @Stable StableValueImpl<R>[] delegates,
-                                                       Function<? super E, ? extends R> original) implements Function<E, R> {
+                                                       UnderlyingHolder<Function<? super E, ? extends R>> underlyingHolder) implements Function<E, R> {
+
     @ForceInline
     @Override
     public R apply(E value) {
@@ -69,7 +69,7 @@ public R apply(E value) {
         final int index = value.ordinal() - firstOrdinal;
         // Since we did the member.test above, we know the index is in bounds
         return delegates[index].orElseSet(new Supplier<R>() {
-                    @Override public R get() { return original.apply(value); }});
+            @Override public R get() { return underlyingHolder.underlying().apply(value); }}, underlyingHolder);
 
     }
 
@@ -96,9 +96,10 @@ public String toString() {
         return StableUtil.renderMappings(this, "StableFunction", entries, true);
     }
 
+
     @SuppressWarnings("unchecked")
     public static <T, E extends Enum<E>, R> Function<T, R> of(Set<? extends T> inputs,
-                                                              Function<? super T, ? extends R> original) {
+                                                              Function<? super T, ? extends R> underlying) {
         // The input set is not empty
         final Class<E> enumType = ((E) inputs.iterator().next()).getDeclaringClass();
         final BitSet bitSet = new BitSet(enumType.getEnumConstants().length);
@@ -112,7 +113,7 @@ public static <T, E extends Enum<E>, R> Function<T, R> of(Set<? extends T> input
         }
         final int size = max - min + 1;
         final IntPredicate member = ImmutableBitSetPredicate.of(bitSet);
-        return (Function<T, R>) new StableEnumFunction<E, R>(enumType, min, member, StableUtil.array(size), (Function<E, R>) original);
+        return (Function<T, R>) new StableEnumFunction<E, R>(enumType, min, member, StableUtil.array(size), new UnderlyingHolder<>((Function<E, R>) underlying, bitSet.cardinality()));
     }
 
 }

src/java.base/share/classes/jdk/internal/lang/stable/StableFunction.java

@@ -35,19 +35,20 @@
 // Note: It would be possible to just use `StableMap::get` with some additional logic
 // instead of this class but explicitly providing a class like this provides better
 // debug capability, exception handling, and may provide better performance.
+
 /**
  * Implementation of a stable Function.
  *
  * @implNote This implementation can be used early in the boot sequence as it does not
  *           rely on reflection, MethodHandles, Streams etc.
  *
- * @param values   a delegate map of inputs to StableValue mappings
- * @param original the original Function
- * @param <T>      the type of the input to the function
- * @param <R>      the type of the result of the function
+ * @param values           a delegate map of inputs to StableValue mappings
+ * @param underlyingHolder of the original underlying Function
+ * @param <T>              the type of the input to the function
+ * @param <R>              the type of the result of the function
  */
 public record StableFunction<T, R>(Map<? extends T, StableValueImpl<R>> values,
-                                   Function<? super T, ? extends R> original) implements Function<T, R> {
+                                   UnderlyingHolder<Function<? super T, ? extends R>> underlyingHolder) implements Function<T, R> {
 
     @ForceInline
     @Override
@@ -57,7 +58,7 @@ public R apply(T value) {
             throw new IllegalArgumentException("Input not allowed: " + value);
         }
         return stable.orElseSet(new Supplier<R>() {
-            @Override  public R get() { return original.apply(value); }});
+            @Override public R get() { return underlyingHolder.underlying().apply(value); }}, underlyingHolder);
     }
 
     @Override
@@ -76,8 +77,8 @@ public String toString() {
     }
 
     public static <T, R> StableFunction<T, R> of(Set<? extends T> inputs,
-                                                 Function<? super T, ? extends R> original) {
-        return new StableFunction<>(StableUtil.map(inputs), original);
+                                                 Function<? super T, ? extends R> underlying) {
+        return new StableFunction<>(StableUtil.map(inputs), new UnderlyingHolder<>(underlying, inputs.size()));
     }
 
 }

src/java.base/share/classes/jdk/internal/lang/stable/StableIntFunction.java

@@ -43,19 +43,19 @@
  * @param <R> the return type
  */
 public record StableIntFunction<R>(@Stable StableValueImpl<R>[] delegates,
-                                   IntFunction<? extends R> original) implements IntFunction<R> {
+                                   UnderlyingHolder<IntFunction<? extends R>> underlyingHolder) implements IntFunction<R> {
 
     @ForceInline
     @Override
     public R apply(int index) {
         final StableValueImpl<R> delegate;
         try {
-            delegate =  delegates[index];
+            delegate = delegates[index];
         } catch (ArrayIndexOutOfBoundsException ioob) {
             throw new IllegalArgumentException("Input not allowed: " + index, ioob);
         }
         return delegate.orElseSet(new Supplier<R>() {
-                    @Override public R get() { return original.apply(index); }});
+            @Override public R get() { return underlyingHolder.underlying().apply(index); }}, underlyingHolder);
     }
 
     @Override
@@ -73,8 +73,8 @@ public String toString() {
         return StableUtil.renderElements(this, "StableIntFunction", delegates);
     }
 
-    public static <R> StableIntFunction<R> of(int size, IntFunction<? extends R> original) {
-        return new StableIntFunction<>(StableUtil.array(size), original);
+    public static <R> StableIntFunction<R> of(int size, IntFunction<? extends R> underlying) {
+        return new StableIntFunction<>(StableUtil.array(size), new UnderlyingHolder<>(underlying, size));
     }
 
 }

src/java.base/share/classes/jdk/internal/lang/stable/StableSupplier.java

@@ -38,12 +38,12 @@
  * @param <T> the return type
  */
 public record StableSupplier<T>(StableValueImpl<T> delegate,
-                                Supplier<? extends T> original) implements Supplier<T> {
+                                UnderlyingHolder<Supplier<? extends T>> underlyingHolder) implements Supplier<T> {
 
     @ForceInline
     @Override
     public T get() {
-        return delegate.orElseSet(original);
+        return delegate.orElseSet(underlyingHolder.underlying(), underlyingHolder);
     }
 
     @Override
@@ -62,8 +62,8 @@ public String toString() {
         return t == this ? "(this StableSupplier)" : StableValueImpl.renderWrapped(t);
     }
 
-    public static <T> StableSupplier<T> of(Supplier<? extends T> original) {
-        return new StableSupplier<>(StableValueImpl.of(), original);
+    public static <T> StableSupplier<T> of(Supplier<? extends T> underlying) {
+        return new StableSupplier<>(StableValueImpl.of(), new UnderlyingHolder<>(underlying, 1));
     }
 
 }

src/java.base/share/classes/jdk/internal/lang/stable/StableValueImpl.java

@@ -126,19 +126,32 @@ public boolean isSet() {
     @Override
     public T orElseSet(Supplier<? extends T> supplier) {
         Objects.requireNonNull(supplier);
+        return orElseSet(supplier, null);
+    }
+
+    // `supplier` can be null if the `underlyingHolder` released it.
+    @ForceInline
+    public T orElseSet(Supplier<? extends T> supplier,
+                       UnderlyingHolder<?> underlyingHolder) {
         final Object t = wrappedContentsAcquire();
-        return (t == null) ? orElseSetSlowPath(supplier) : unwrap(t);
+        return (t == null) ? orElseSetSlowPath(supplier, underlyingHolder) : unwrap(t);
     }
 
     @DontInline
-    private T orElseSetSlowPath(Supplier<? extends T> supplier) {
+    private T orElseSetSlowPath(Supplier<? extends T> supplier,
+                                UnderlyingHolder<?> underlyingHolder) {
         preventReentry();
         synchronized (this) {
             final Object t = contents;  // Plain semantics suffice here
             if (t == null) {
                 final T newValue = supplier.get();
                 // The mutex is not reentrant so we know newValue should be returned
                 wrapAndSet(newValue);
+                if (underlyingHolder != null) {
+                    // Reduce the counter and if it reaches zero, clear the reference
+                    // to the underlying holder.
+                    underlyingHolder.countDown();
+                }
                 return newValue;
             }
             return unwrap(t);

src/java.base/share/classes/jdk/internal/lang/stable/UnderlyingHolder.java

@@ -0,0 +1,55 @@
+package jdk.internal.lang.stable;
+
+import jdk.internal.misc.Unsafe;
+import jdk.internal.vm.annotation.ForceInline;
+
+import java.util.Objects;
+import java.util.stream.Stream;
+
+import static java.util.stream.Collectors.joining;
+
+/**
+ * This class is thread safe.
+ *
+ * @param <U> the underlying type
+ */
+public final class UnderlyingHolder<U> {
+
+    private static final Unsafe UNSAFE = Unsafe.getUnsafe();
+
+    private static final long COUNTER_OFFSET =
+            UNSAFE.objectFieldOffset(UnderlyingHolder.class, "counter");
+
+    // Used reflectively. This field can only transition at most once from being set to a
+    // non-null reference to being `null`. Once `null`, it is never read. This allows
+    // the field to be non-volatile, which is crucial for getting optimum performance.
+    private U underlying;
+
+    // Used reflectively
+    private int counter;
+
+    public UnderlyingHolder(U underlying, int counter) {
+        this.underlying = underlying;
+        this.counter = counter;
+        // Safe publication
+        UNSAFE.storeStoreFence();
+    }
+
+    @ForceInline
+    public U underlying() {
+        return underlying;
+    }
+
+    // For testing only
+    public int counter() {
+        return UNSAFE.getIntVolatile(this, COUNTER_OFFSET);
+    }
+
+    public void countDown() {
+        if (UNSAFE.getAndAddInt(this, COUNTER_OFFSET, -1) == 1) {
+            // Do not reference the underlying function anymore so it can be collected.
+            underlying = null;
+        }
+    }
+
+}