Refactoring ImmutableSet

Glavo · Jan 10, 2025 · 59d315e · 59d315e
1 parent 42dd5db
commit 59d315e
Show file tree

Hide file tree

Showing 7 changed files with 136 additions and 169 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -8,6 +8,9 @@
 
 * Add `kala.collection.base.OrderedTraversable::mapIndexed(CollectionFactory<U, ?, R>, IndexedFunction<T, U>)`
 * Add `kala.collection.base.OrderedTraversable::mapIndexedTo(G, IndexedFunction<T, U>)`
+* Add `kala.collection.Set::added(E)`
+* Add `kala.collection.Set::addedAll(E...)`
+* Add `kala.collection.Set::addedAll(Iterable<E>)`
 
 ### Breaking Changes
 

diff --git a/kala-collection/src/main/java/kala/collection/Set.java b/kala-collection/src/main/java/kala/collection/Set.java
@@ -18,13 +18,15 @@
 import kala.annotations.DelegateBy;
 import kala.collection.factory.CollectionFactory;
 import kala.collection.immutable.ImmutableSet;
+import kala.collection.internal.CollectionHelper;
 import kala.collection.internal.convert.AsJavaConvert;
 import kala.collection.internal.convert.FromJavaConvert;
 import kala.collection.internal.view.SetViews;
 import kala.function.Predicates;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
 
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.function.Predicate;
 import java.util.stream.Stream;
@@ -33,6 +35,11 @@ public interface Set<E> extends Collection<E>, SetLike<E>, AnySet<E> {
 
     //region Static Factories
 
+    @SuppressWarnings("unchecked")
+    static <E> Set<E> narrow(Set<? extends E> set) {
+        return (Set<E>) set;
+    }
+
     static <E> CollectionFactory<E, ?, Set<E>> factory() {
         return CollectionFactory.narrow(ImmutableSet.factory());
     }
@@ -132,24 +139,99 @@ static boolean equals(@NotNull Set<?> set1, @NotNull AnySet<?> set2) {
 
     @Override
     default @NotNull ImmutableSet<E> added(E value) {
-        return this.view().added(value).toImmutableSet();
-    }
+        if (this instanceof ImmutableSet<E> immutableSet && contains(value)) {
+            return immutableSet;
+        }
 
-    @Override
-    default @NotNull ImmutableSet<E> addedAll(@NotNull Iterable<? extends E> values) {
-        return this.view().addedAll(values).toImmutableSet();
+        var factory = CollectionHelper.immutableSetFactoryBy(this);
+        var builder = factory.newCollectionBuilder();
+        for (E e : this) {
+            builder.plusAssign(e);
+        }
+        builder.plusAssign(value);
+        return builder.build();
     }
 
     @Override
     @ApiStatus.NonExtendable
+    @SuppressWarnings("unchecked")
     @DelegateBy("addedAll(Iterable<E>)")
-    default @NotNull ImmutableSet<E> addedAll(E @NotNull [] values) {
+    default @NotNull ImmutableSet<E> addedAll(E... values) {
         return addedAll(ArraySeq.wrap(values));
     }
 
+    @Override
+    default @NotNull ImmutableSet<E> addedAll(@NotNull Iterable<? extends E> values) {
+        var factory = CollectionHelper.immutableSetFactoryBy(this);
+        var builder = factory.newCollectionBuilder();
+        for (E e : this) {
+            builder.plusAssign(e);
+        }
+        for (E value : values) {
+            builder.plusAssign(value);
+        }
+        return builder.build();
+    }
+
+    default @NotNull ImmutableSet<E> removed(E value) {
+        if (this instanceof ImmutableSet<E> immutableSet && !contains(value)) {
+            return immutableSet;
+        }
+
+        var factory = CollectionHelper.immutableSetFactoryBy(this);
+        var builder = factory.newCollectionBuilder();
+        builder.sizeHint(size() -1);
+        if (value == null) {
+            for (E e : this) {
+                if (null != e) {
+                    builder.plusAssign(e);
+                }
+            }
+        } else {
+            for (E e : this) {
+                if (!value.equals(e)) {
+                    builder.plusAssign(e);
+                }
+            }
+        }
+        return builder.build();
+    }
+
+    @ApiStatus.NonExtendable
+    @SuppressWarnings("unchecked")
+    @DelegateBy("removedAll(Iterable<E>)")
+    default @NotNull ImmutableSet<E> removedAll(E... values) {
+        return removedAll(ArraySeq.wrap(values));
+    }
+
+    default @NotNull ImmutableSet<E> removedAll(@NotNull Iterable<? extends E> values) {
+        var factory = CollectionHelper.immutableSetFactoryBy(this);
+
+        Predicate<? super E> contains;
+        if (values instanceof java.util.Set<? extends E> collection) {
+            contains = collection::contains;
+        } else if (values instanceof SetLike<? extends E> setLike) {
+            contains = setLike::contains;
+        } else {
+            HashSet<E> hashSet = new HashSet<>();
+            for (E value : values) {
+                hashSet.add(value);
+            }
+            contains = hashSet::contains;
+        }
+
+        var builder = factory.newCollectionBuilder();
+        for (E e : this) {
+            if (!contains.test(e)) {
+                builder.plusAssign(e);
+            }
+        }
+        return builder.build();
+    }
+
     @Override
     default @NotNull ImmutableSet<E> filter(@NotNull Predicate<? super E> predicate) {
-        return view().filter(predicate).toImmutableSet();
+        return filter(CollectionHelper.immutableSetFactoryBy(this), predicate);
     }
 
     @Override
@@ -167,10 +249,11 @@ static boolean equals(@NotNull Set<?> set1, @NotNull AnySet<?> set2) {
     }
 
     @Override
-    @SuppressWarnings("unchecked")
     @ApiStatus.NonExtendable
     @DelegateBy("filter(Predicate<E>)")
     default <U> @NotNull ImmutableSet<U> filterIsInstance(@NotNull Class<? extends U> clazz) {
-        return (ImmutableSet<U>) filter(Predicates.isInstance(clazz));
+        @SuppressWarnings("unchecked")
+        ImmutableSet<U> result = (ImmutableSet<U>) filter(Predicates.isInstance(clazz));
+        return result;
     }
 }
diff --git a/kala-collection/src/main/java/kala/collection/SetLike.java b/kala-collection/src/main/java/kala/collection/SetLike.java
@@ -35,9 +35,9 @@ default Predicate<E> asPredicate() {
 
     @NotNull SetLike<E> added(E value);
 
-    @NotNull SetLike<E> addedAll(@NotNull Iterable<? extends E> values);
+    @NotNull SetLike<E> addedAll(E... values);
 
-    @NotNull SetLike<E> addedAll(E @NotNull [] values);
+    @NotNull SetLike<E> addedAll(@NotNull Iterable<? extends E> values);
 
     @Override
     @NotNull SetLike<E> filter(@NotNull Predicate<? super E> predicate);

diff --git a/kala-collection/src/main/java/kala/collection/immutable/AbstractImmutableSet.java b/kala-collection/src/main/java/kala/collection/immutable/AbstractImmutableSet.java
@@ -17,71 +17,9 @@
 
 import kala.collection.AbstractSet;
 import kala.annotations.Covariant;
-import kala.collection.factory.CollectionFactory;
 import org.jetbrains.annotations.Debug;
-import org.jetbrains.annotations.NotNull;
 
 @Debug.Renderer(hasChildren = "isNotEmpty()", childrenArray = "toArray()")
 public abstract class AbstractImmutableSet<@Covariant E> extends AbstractSet<E> implements ImmutableSet<E> {
-    static <E, T, Builder> T added(
-            @NotNull ImmutableSet<? extends E> set,
-            E value,
-            @NotNull CollectionFactory<? super E, Builder, ? extends T> factory
-    ) {
-        Builder builder = factory.newBuilder();
-        factory.sizeHint(builder, set, 1);
-        factory.addAllToBuilder(builder, set);
-        factory.addToBuilder(builder, value);
-        return factory.build(builder);
-    }
 
-    static <E, T, Builder> T addedAll(
-            @NotNull ImmutableSet<? extends E> set,
-            @NotNull Iterable<? extends E> values,
-            @NotNull CollectionFactory<? super E, Builder, ? extends T> factory
-    ) {
-        Builder builder = factory.newBuilder();
-
-        factory.sizeHint(builder, set);
-        factory.addAllToBuilder(builder, set);
-        factory.sizeHint(builder, values);
-        factory.addAllToBuilder(builder, values);
-
-        return factory.build(builder);
-    }
-
-    static <E, T, Builder> T removed(
-            @NotNull ImmutableSet<? extends E> set,
-            E value,
-            @NotNull CollectionFactory<? super E, Builder, ? extends T> factory
-    ) {
-        Builder builder = factory.newBuilder();
-        factory.sizeHint(builder, set, -1);
-        if (value == null) {
-            for (E e : set) {
-                if (null != e)
-                    factory.addToBuilder(builder, e);
-            }
-        } else {
-            for (E e : set) {
-                if (!value.equals(e))
-                    factory.addToBuilder(builder, e);
-            }
-        }
-        return factory.build(builder);
-    }
-
-    static <E, T, Builder> T removedAll(
-            @NotNull ImmutableSet<? extends E> set,
-            @NotNull Iterable<? extends E> values,
-            @NotNull CollectionFactory<? super E, Builder, ? extends T> factory
-    ) {
-        ImmutableHashSet<? extends E> s = ImmutableHashSet.from(values);
-        Builder builder = factory.newBuilder();
-        for (E e : set) {
-            if (!s.contains(e))
-                factory.addToBuilder(builder, e);
-        }
-        return factory.build(builder);
-    }
 }
diff --git a/kala-collection/src/main/java/kala/collection/immutable/ImmutableSet.java b/kala-collection/src/main/java/kala/collection/immutable/ImmutableSet.java
@@ -15,18 +15,13 @@
  */
 package kala.collection.immutable;
 
-import kala.annotations.DelegateBy;
-import kala.collection.ArraySeq;
-import kala.collection.SortedSet;
 import kala.annotations.Covariant;
 import kala.collection.factory.CollectionFactory;
 import kala.collection.Set;
-import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.*;
-import java.util.function.Predicate;
 import java.util.stream.Stream;
 
 @SuppressWarnings("unchecked")
@@ -103,79 +98,6 @@ static <E> ImmutableSet<E> narrow(ImmutableSet<? extends E> set) {
         return factory();
     }
 
-    @Override
-    default @NotNull ImmutableSet<E> added(E value) {
-        if (contains(value))
-            return this;
-
-        CollectionFactory<E, ?, ? extends ImmutableSet<E>> factory;
-        if (this instanceof SortedSet<?> self) {
-            factory = (CollectionFactory<E, ?, ? extends ImmutableSet<E>>)
-                    self.sortedIterableFactory();
-        } else {
-            factory = iterableFactory();
-        }
-        return AbstractImmutableSet.added(this, value, factory);
-    }
-
-    @Override
-    @SuppressWarnings("unchecked")
-    default @NotNull ImmutableSet<E> addedAll(@NotNull Iterable<? extends E> values) {
-        CollectionFactory<E, ?, ? extends ImmutableSet<E>> factory;
-        if (this instanceof SortedSet<?> sortedSet) {
-            factory = (CollectionFactory<E, ?, ? extends ImmutableSet<E>>)
-                    sortedSet.sortedIterableFactory();
-        } else {
-            factory = iterableFactory();
-        }
-        return AbstractImmutableSet.addedAll(this, values, factory);
-    }
-
-    default @NotNull ImmutableSet<E> removed(E value) {
-        if (!contains(value))
-            return this;
-
-        CollectionFactory<E, ?, ? extends ImmutableSet<E>> factory;
-
-        if (this instanceof SortedSet<?> self) {
-            factory = (CollectionFactory<E, ?, ? extends ImmutableSet<E>>)
-                    self.sortedIterableFactory();
-        } else {
-            factory = iterableFactory();
-        }
-
-        return AbstractImmutableSet.removed(this, value, factory);
-    }
-
-    default @NotNull ImmutableSet<E> removedAll(@NotNull Iterable<? extends E> values) {
-        CollectionFactory<E, ?, ? extends ImmutableSet<E>> factory;
-        if (this instanceof SortedSet<?> self) {
-            factory = (CollectionFactory<E, ?, ? extends ImmutableSet<E>>)
-                    self.sortedIterableFactory();
-        } else {
-            factory = iterableFactory();
-        }
-        return AbstractImmutableSet.removedAll(this, values, factory);
-    }
-
-    @ApiStatus.NonExtendable
-    @DelegateBy("removedAll(Iterable<E>)")
-    default @NotNull ImmutableSet<E> removedAll(E @NotNull [] values) {
-        return removedAll(ArraySeq.wrap(values));
-    }
-
-    @Override
-    default @NotNull ImmutableSet<E> filter(@NotNull Predicate<? super E> predicate) {
-        CollectionFactory<E, ?, ? extends ImmutableSet<E>> factory;
-        if (this instanceof SortedSet<?> self) {
-            factory = (CollectionFactory<E, ?, ? extends ImmutableSet<E>>)
-                    self.sortedIterableFactory();
-        } else {
-            factory = iterableFactory();
-        }
-        return filter(factory, predicate);
-    }
-
     @Override
     default @NotNull ImmutableSet<E> toImmutableSet() {
         return this;