Inline NonEmpty maps (#3120)

arrow-kt · Aug 30, 2023 · a7dc178 · a7dc178
1 parent 02215d4
commit a7dc178
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 2 deletions.
diff --git a/arrow-libs/core/arrow-core/src/commonMain/kotlin/arrow/core/NonEmptyList.kt b/arrow-libs/core/arrow-core/src/commonMain/kotlin/arrow/core/NonEmptyList.kt
@@ -183,9 +183,18 @@ public class NonEmptyList<out A>(
   public override inline fun <B> map(transform: (A) -> B): NonEmptyList<B> =
     NonEmptyList(transform(head), tail.map(transform))
 
-  override fun <B> flatMap(transform: (A) -> NonEmptyCollection<B>): NonEmptyList<B> =
+  @Suppress("OVERRIDE_BY_INLINE")
+  public override inline fun <B> mapIndexed(transform: (index: Int, A) -> B): NonEmptyList<B> =
+    NonEmptyList(transform(0, head), tail.mapIndexed { ix, e -> transform(ix + 1, e) })
+
+  @Suppress("OVERRIDE_BY_INLINE")
+  public override inline fun <B> flatMap(transform: (A) -> NonEmptyCollection<B>): NonEmptyList<B> =
     transform(head).toNonEmptyList() + tail.flatMap(transform)
 
+  @Suppress("OVERRIDE_BY_INLINE")
+  public override inline fun <K> distinctBy(selector: (A) -> K): NonEmptyList<A> =
+    all.distinctBy(selector).toNonEmptyListOrNull()!!
+
   public operator fun plus(l: NonEmptyList<@UnsafeVariance A>): NonEmptyList<A> =
     this + l.all
 

diff --git a/arrow-libs/core/arrow-core/src/commonMain/kotlin/arrow/core/NonEmptySet.kt b/arrow-libs/core/arrow-core/src/commonMain/kotlin/arrow/core/NonEmptySet.kt
@@ -4,7 +4,7 @@ import kotlin.jvm.JvmInline
 
 @JvmInline
 public value class NonEmptySet<out A> private constructor(
-  private val elements: Set<A>
+  @PublishedApi internal val elements: Set<A>
 ) : Set<A> by elements, NonEmptyCollection<A> {
 
   public constructor(first: A, rest: Set<A>) : this(setOf(first) + rest)
@@ -21,6 +21,25 @@ public value class NonEmptySet<out A> private constructor(
 
   override fun lastOrNull(): A = elements.last()
 
+  @Suppress("OVERRIDE_BY_INLINE")
+  public override inline fun <B> map(transform: (A) -> B): NonEmptyList<B> =
+    elements.map(transform).toNonEmptyListOrNull()!!
+
+  @Suppress("OVERRIDE_BY_INLINE")
+  public override inline fun <B> mapIndexed(transform: (index: Int, A) -> B): NonEmptyList<B> =
+    elements.mapIndexed(transform).toNonEmptyListOrNull()!!
+
+  @Suppress("OVERRIDE_BY_INLINE")
+  public override inline fun <B> flatMap(transform: (A) -> NonEmptyCollection<B>): NonEmptyList<B> =
+    elements.flatMap(transform).toNonEmptyListOrNull()!!
+
+  override fun distinct(): NonEmptyList<A> =
+    toNonEmptyList()
+
+  @Suppress("OVERRIDE_BY_INLINE")
+  public override inline fun <K> distinctBy(selector: (A) -> K): NonEmptyList<A> =
+    elements.distinctBy(selector).toNonEmptyListOrNull()!!
+
   override fun toString(): String = "NonEmptySet(${this.joinToString()})"
 
   @Suppress("RESERVED_MEMBER_INSIDE_VALUE_CLASS")