Create SeqLike::inserted

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

Hide file tree

Showing 10 changed files with 239 additions and 102 deletions.
diff --git a/kala-base/src/main/java/kala/collection/base/Iterators.java b/kala-base/src/main/java/kala/collection/base/Iterators.java
@@ -799,8 +799,30 @@ public E next() {
         return value == null ? new PrependedNull<>(it) : new PrependedNotNull<>(it, value);
     }
 
+    public static <E> @NotNull Iterator<E> inserted(@NotNull Iterator<? extends E> it, int index, E value) {
+        return new AbstractIterator<>() {
+            private int i;
+
+            @Override
+            public boolean hasNext() {
+                return i == index || it.hasNext();
+            }
+
+            @Override
+            public E next() {
+                checkStatus();
+
+                if (i++ == index) {
+                    return value;
+                } else {
+                    return it.next();
+                }
+            }
+        };
+    }
+
     public static <E> @NotNull Iterator<E> removed(@NotNull Iterator<? extends E> it, int index) {
-        return new AbstractIterator<E>() {
+        return new AbstractIterator<>() {
             private int count = 0;
 
             @Override
@@ -969,7 +991,7 @@ public U next() {
     }
 
     public static <E, U, V> @NotNull Iterator<@NotNull Tuple3<E, U, V>> zip3(@NotNull Iterator<? extends E> it1, Iterator<? extends U> it2, Iterator<? extends V> it3) {
-        if (!it1.hasNext() || !it2.hasNext() ||!it3.hasNext()) { // implicit null check of it1 and it2
+        if (!it1.hasNext() || !it2.hasNext() || !it3.hasNext()) { // implicit null check of it1 and it2
             return Iterators.empty();
         }
         return new Zip3<>(it1, it2, it3);

diff --git a/kala-collection/src/main/java/kala/collection/Seq.java b/kala-collection/src/main/java/kala/collection/Seq.java
@@ -247,6 +247,12 @@ static boolean equals(@NotNull Seq<?> seq1, @NotNull AnySeq<?> seq2) {
         return view().appendedAll(values).toImmutableSeq();
     }
 
+    @Override
+    default @NotNull ImmutableSeq<E> inserted(int index, E value) {
+        Conditions.checkPositionIndex(index, size());
+        return view().inserted(index, value).toImmutableSeq();
+    }
+
     @Override
     @Contract(pure = true)
     default @NotNull ImmutableSeq<E> removedAt(int index) {

diff --git a/kala-collection/src/main/java/kala/collection/SeqLike.java b/kala-collection/src/main/java/kala/collection/SeqLike.java
@@ -213,6 +213,9 @@ default int binarySearch(int beginIndex, int endIndex, E value, Comparator<? sup
     @Contract(pure = true)
     @NotNull SeqLike<E> appendedAll(E @NotNull [] values);
 
+    @Contract(pure = true)
+    @NotNull SeqLike<E> inserted(int index, E value);
+
     @Contract(pure = true)
     @NotNull SeqLike<E> removedAt(int index);
 

diff --git a/kala-collection/src/main/java/kala/collection/SeqView.java b/kala-collection/src/main/java/kala/collection/SeqView.java
@@ -216,16 +216,13 @@ static <E> SeqView<E> narrow(SeqView<? extends E> view) {
         return new SeqViews.Concat<>(ArraySeq.wrap(prefix), this);
     }
 
-    default @NotNull SeqView<E> removedAt(int index) {
-        final int ks = this.knownSize();
-        if (ks < 0) {
-            if (index < 0) {
-                throw new IndexOutOfBoundsException("index(" + index + ") < 0");
-            }
-        } else {
-            Conditions.checkElementIndex(index, ks);
-        }
+    default @NotNull SeqLike<E> inserted(int index, E value) {
+        Conditions.checkPositionIndex(index, size());
+        return new SeqViews.Inserted<>(this, index, value);
+    }
 
+    default @NotNull SeqView<E> removedAt(int index) {
+        Conditions.checkElementIndex(index, size());
         return new SeqViews.RemovedAt<>(this, index);
     }
 

diff --git a/kala-collection/src/main/java/kala/collection/immutable/AbstractImmutableSeq.java b/kala-collection/src/main/java/kala/collection/immutable/AbstractImmutableSeq.java
@@ -59,22 +59,42 @@ static <E, T, Builder> T updated(
         return factory.build(builder);
     }
 
-    static <E, T, Builder> T removedAt(
+    static <E, T, Builder> T inserted(
             @NotNull ImmutableSeq<? extends E> seq,
             int index,
+            E newValue,
             @NotNull CollectionFactory<? super E, Builder, ? extends T> factory
     ) {
-        final int s = seq.size();
+        final int size = seq.size();
+        Conditions.checkPositionIndex(index, size);
+
+        Builder builder = factory.newBuilder();
+        factory.sizeHint(builder, size - 1);
+
+        Iterator<? extends E> iterator = seq.iterator();
+
+        for (int i = 0; i < index; i++) {
+            factory.addToBuilder(builder, iterator.next());
+        }
+        factory.addToBuilder(builder, newValue);
+        factory.addAllToBuilder(builder, iterator);
+        return factory.build(builder);
+    }
 
-        Conditions.checkElementIndex(index, s);
+    static <E, T, Builder> T removedAt(
+            @NotNull ImmutableSeq<? extends E> seq,
+            int index,
+            @NotNull CollectionFactory<? super E, Builder, ? extends T> factory
+    ) {
+        final int size = seq.size();
+        Conditions.checkElementIndex(index, size);
 
-        if (s == 1) {
+        if (size == 1) {
             return factory.empty();
         }
 
         Builder builder = factory.newBuilder();
-
-        factory.sizeHint(builder, s - 1);
+        factory.sizeHint(builder, size - 1);
 
         Iterator<? extends E> iterator = seq.iterator();
 

diff --git a/kala-collection/src/main/java/kala/collection/immutable/ImmutableSeq.java b/kala-collection/src/main/java/kala/collection/immutable/ImmutableSeq.java
@@ -341,7 +341,12 @@ static <E> ImmutableSeq<E> narrow(ImmutableSeq<? extends E> seq) {
         return AbstractImmutableSeq.appendedAll(this, values, iterableFactory());
     }
 
-    @Override
+    @Contract(pure = true)
+    default @NotNull ImmutableSeq<E> inserted(int index, E value) {
+        return AbstractImmutableSeq.inserted(this, index, value, iterableFactory());
+    }
+
+    @Contract(pure = true)
     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,15 @@ public E get(int index) {
             return ImmutableSeq.of(value1, value);
         }
 
+        @Override
+        public @NotNull ImmutableSeq<E> inserted(int index, E value) {
+            return switch (index) {
+                case 0 -> ImmutableSeq.of(value, value1);
+                case 1 -> ImmutableSeq.of(value1, value);
+                default -> throw new IndexOutOfBoundsException(index);
+            };
+        }
+
         @Override
         public @NotNull ImmutableSeq<E> removedAt(int index) {
             if (index != 0) {
@@ -378,6 +387,16 @@ public E get(int index) {
             return ImmutableSeq.of(value1, value2, value);
         }
 
+        @Override
+        public @NotNull ImmutableSeq<E> inserted(int index, E value) {
+            return switch (index) {
+                case 0 -> ImmutableSeq.of(value, value1, value2);
+                case 1 -> ImmutableSeq.of(value1, value, value2);
+                case 2 -> ImmutableSeq.of(value1, value2, value);
+                default -> throw new IndexOutOfBoundsException(index);
+            };
+        }
+
         @Override
         public @NotNull ImmutableSeq<E> removedAt(int index) {
             return switch (index) {

diff --git a/kala-collection/src/main/java/kala/collection/immutable/ImmutableTreeSeq.java b/kala-collection/src/main/java/kala/collection/immutable/ImmutableTreeSeq.java
@@ -249,6 +249,13 @@ public E get(int index) {
         return root.get(index);
     }
 
+    @Override
+    public @NotNull ImmutableSeq<E> updated(int index, E newValue) {
+        Conditions.checkElementIndex(index, size());
+        IndexedTree<E> newRoot = root.plus(index, newValue);
+        return newRoot != root ? new ImmutableTreeSeq<>(newRoot) : this;
+    }
+
     @Override
     public @NotNull ImmutableSeq<E> prepended(E value) {
         return new ImmutableTreeSeq<>(root.changeKeysAbove(0, 1).plus(0, value));
@@ -259,13 +266,6 @@ public E get(int index) {
         return new ImmutableTreeSeq<>(root.plus(size(), value));
     }
 
-    @Override
-    public @NotNull ImmutableSeq<E> updated(int index, E newValue) {
-        Conditions.checkElementIndex(index, size());
-        IndexedTree<E> newRoot = root.plus(index, newValue);
-        return newRoot != root ? new ImmutableTreeSeq<>(newRoot) : this;
-    }
-
     @Override
     public @NotNull ImmutableSeq<E> removedAt(int index) {
         final int size = size();