Optimize operations

Glavo · May 29, 2024 · 800c52c · 800c52c
1 parent 29b29ca
commit 800c52c
Show file tree

Hide file tree

Showing 6 changed files with 90 additions and 28 deletions.
diff --git a/kala-collection/src/main/java/kala/collection/ArraySeq.java b/kala-collection/src/main/java/kala/collection/ArraySeq.java
@@ -462,6 +462,23 @@ public final E get(int index) {
         return ImmutableArray.Unsafe.wrap(newValues);
     }
 
+    @Override
+    public @NotNull ImmutableSeq<E> removedAt(int index) {
+        final Object[] elements = this.elements;
+        final int size = elements.length;
+
+        Conditions.checkElementIndex(index, size);
+
+        if (size == 1) {
+            return ImmutableArray.empty();
+        }
+
+        Object[] newValues = new Object[size - 1];
+        System.arraycopy(elements, 0, newValues, 0, index);
+        System.arraycopy(elements, index + 1, newValues, index, size - index - 1);
+        return ImmutableArray.Unsafe.wrap(newValues);
+    }
+
     @Override
     public @NotNull ImmutableSeq<E> concat(@NotNull SeqLike<? extends E> other) {
         return appendedAll(other);

diff --git a/kala-collection/src/main/java/kala/collection/Seq.java b/kala-collection/src/main/java/kala/collection/Seq.java
@@ -249,7 +249,7 @@ static boolean equals(@NotNull Seq<?> seq1, @NotNull AnySeq<?> seq2) {
 
     @Override
     @Contract(pure = true)
-    default @NotNull SeqLike<E> removedAt(int index) {
+    default @NotNull ImmutableSeq<E> removedAt(int index) {
         Conditions.checkElementIndex(index, size());
         return view().removedAt(index).toImmutableSeq();
     }

diff --git a/kala-collection/src/main/java/kala/collection/immutable/ImmutableSeq.java b/kala-collection/src/main/java/kala/collection/immutable/ImmutableSeq.java
@@ -1,3 +1,18 @@
+/*
+ * Copyright 2024 Glavo
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
 package kala.collection.immutable;
 
 import kala.annotations.DelegateBy;
@@ -281,11 +296,6 @@ static <E> ImmutableSeq<E> narrow(ImmutableSeq<? extends E> seq) {
         return AbstractImmutableSeq.takeWhile(this, predicate, iterableFactory());
     }
 
-    @Contract(pure = true)
-    default @NotNull ImmutableSeq<E> updated(int index, E newValue) {
-        return AbstractImmutableSeq.updated(this, index, newValue, iterableFactory());
-    }
-
     @Contract(pure = true)
     default @NotNull ImmutableSeq<E> concat(@NotNull SeqLike<? extends E> other) {
         return AbstractImmutableSeq.concat(this, other, iterableFactory());
@@ -296,6 +306,11 @@ static <E> ImmutableSeq<E> narrow(ImmutableSeq<? extends E> seq) {
         return AbstractImmutableSeq.concat(this, other, iterableFactory());
     }
 
+    @Contract(pure = true)
+    default @NotNull ImmutableSeq<E> updated(int index, E newValue) {
+        return AbstractImmutableSeq.updated(this, index, newValue, iterableFactory());
+    }
+
     @Contract(pure = true)
     default @NotNull ImmutableSeq<E> prepended(E value) {
         return AbstractImmutableSeq.prepended(this, value, iterableFactory());
@@ -327,8 +342,7 @@ static <E> ImmutableSeq<E> narrow(ImmutableSeq<? extends E> seq) {
     }
 
     @Override
-    @NotNull
-    default SeqLike<E> removedAt(int index) {
+    default @NotNull ImmutableSeq<E> removedAt(int index) {
         return AbstractImmutableSeq.removedAt(this, index, iterableFactory());
     }
 

diff --git a/kala-collection/src/main/java/kala/collection/immutable/ImmutableSeqs.java b/kala-collection/src/main/java/kala/collection/immutable/ImmutableSeqs.java
@@ -218,6 +218,14 @@ public E get(int index) {
             return ImmutableSeq.of(value1, value);
         }
 
+        @Override
+        public @NotNull ImmutableSeq<E> removedAt(int index) {
+            if (index != 0) {
+                throw new IndexOutOfBoundsException(index);
+            }
+            return ImmutableSeq.empty();
+        }
+
         @Override
         public @NotNull ImmutableSeq<E> take(int n) {
             if (n < 0) {
@@ -278,7 +286,7 @@ public E get(int index) {
         @Override
         public @NotNull ImmutableSeq<E> updated(int index, E newValue) {
             if (index != 0) {
-                throw new IndexOutOfBoundsException("index: " + index);
+                throw new IndexOutOfBoundsException(index);
             }
             return ImmutableSeq.of(newValue);
         }
@@ -351,6 +359,15 @@ public E get(int index) {
             };
         }
 
+        @Override
+        public @NotNull ImmutableSeq<E> updated(int index, E newValue) {
+            return switch (index) {
+                case 0 -> ImmutableSeq.of(newValue, value2);
+                case 1 -> ImmutableSeq.of(value1, newValue);
+                default -> throw new IndexOutOfBoundsException(index);
+            };
+        }
+
         @Override
         public @NotNull ImmutableSeq<E> prepended(E value) {
             return ImmutableSeq.of(value, value1, value2);
@@ -361,6 +378,15 @@ public E get(int index) {
             return ImmutableSeq.of(value1, value2, value);
         }
 
+        @Override
+        public @NotNull ImmutableSeq<E> removedAt(int index) {
+            return switch (index) {
+                case 0 -> ImmutableSeq.of(value2);
+                case 1 -> ImmutableSeq.of(value1);
+                default -> throw new IndexOutOfBoundsException(index);
+            };
+        }
+
         @Override
         public @NotNull ImmutableSeq<E> take(int n) {
             if (n < 0) {
@@ -437,15 +463,6 @@ public E get(int index) {
             return ImmutableSeq.empty();
         }
 
-        @Override
-        public @NotNull ImmutableSeq<E> updated(int index, E newValue) {
-            return switch (index) {
-                case 0 -> ImmutableSeq.of(newValue, value2);
-                case 1 -> ImmutableSeq.of(value1, newValue);
-                default -> throw new IndexOutOfBoundsException("index: " + index);
-            };
-        }
-
         @Override
         public @NotNull <U> ImmutableSeq<U> map(@NotNull Function<? super E, ? extends U> mapper) {
             return ImmutableSeq.of(

diff --git a/kala-collection/src/main/java/kala/collection/immutable/ImmutableTreeSeq.java b/kala-collection/src/main/java/kala/collection/immutable/ImmutableTreeSeq.java
@@ -266,6 +266,13 @@ public E get(int index) {
         return newRoot != root ? new ImmutableTreeSeq<>(newRoot) : this;
     }
 
+    @Override
+    public @NotNull ImmutableSeq<E> removedAt(int index) {
+        final int size = size();
+        Conditions.checkElementIndex(index, size);
+        return size > 1 ? new ImmutableTreeSeq<>(root.minus(index).changeKeysBelow(index, -1)) : empty();
+    }
+
     @Serial
     private Object writeReplace() {
         return new SerializationReplaced<>(this);

diff --git a/src/test/template/kala/collection/SeqLikeTestTemplate.java b/src/test/template/kala/collection/SeqLikeTestTemplate.java
@@ -473,15 +473,24 @@ default void removedAtTest() {
         }
 
         {
-            var seq = from(List.of("foo"));
+            var seq = from(List.of("str0"));
             assertIterableEquals(List.of(), seq.removedAt(0));
 
-            if (seq instanceof Seq<String>) {
-                assertThrows(IndexOutOfBoundsException.class, () -> seq.removedAt(1));
-                assertThrows(IndexOutOfBoundsException.class, () -> seq.removedAt(-1));
-                assertThrows(IndexOutOfBoundsException.class, () -> seq.removedAt(Integer.MAX_VALUE));
-                assertThrows(IndexOutOfBoundsException.class, () -> seq.removedAt(Integer.MIN_VALUE));
-            }
+            assertThrows(IndexOutOfBoundsException.class, () -> seq.removedAt(1));
+            assertThrows(IndexOutOfBoundsException.class, () -> seq.removedAt(-1));
+            assertThrows(IndexOutOfBoundsException.class, () -> seq.removedAt(Integer.MAX_VALUE));
+            assertThrows(IndexOutOfBoundsException.class, () -> seq.removedAt(Integer.MIN_VALUE));
+        }
+
+        {
+            var seq = from(List.of("str0", "str1"));
+            assertIterableEquals(List.of("str1"), seq.removedAt(0));
+            assertIterableEquals(List.of("str0"), seq.removedAt(1));
+
+            assertThrows(IndexOutOfBoundsException.class, () -> seq.removedAt(2));
+            assertThrows(IndexOutOfBoundsException.class, () -> seq.removedAt(-1));
+            assertThrows(IndexOutOfBoundsException.class, () -> seq.removedAt(Integer.MAX_VALUE));
+            assertThrows(IndexOutOfBoundsException.class, () -> seq.removedAt(Integer.MIN_VALUE));
         }
 
         {
@@ -490,9 +499,7 @@ default void removedAtTest() {
             assertIterableEquals(List.of("str0", "str2"), seq.removedAt(1));
             assertIterableEquals(List.of("str0", "str1"), seq.removedAt(2));
 
-            if (seq instanceof Seq<String>) {
-                assertThrows(IndexOutOfBoundsException.class, () -> seq.removedAt(3));
-            }
+            assertThrows(IndexOutOfBoundsException.class, () -> seq.removedAt(3));
         }
     }