Fixed #1312 - NonEmptySet#removeNonEmpty (#1313)

* Fixed #1312 - NonEmptySet#removeNonEmpty ... and NonEmptySortedSet#removeNonEmpty * Fixed #1312 - Formatting
zio · May 8, 2024 · 053e4b4 · 053e4b4
1 parent d70512f
commit 053e4b4
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 deletions.
diff --git a/core-tests/shared/src/test/scala/zio/prelude/NonEmptySetSpec.scala b/core-tests/shared/src/test/scala/zio/prelude/NonEmptySetSpec.scala
@@ -84,6 +84,15 @@ object NonEmptySetSpec extends ZIOBaseSpec {
             val expected = as.map(f).flatten
             actual <-> expected
           }
+        },
+        test("removeNonEmpty") {
+          check(genSet) { as =>
+            val toRemove = as.head
+            val actual   = NonEmptySet.fromSetOption(as).get.removeNonEmpty(toRemove).map(_.toSet)
+            val modified = as - toRemove
+            val expected = Option(modified).filter(_.nonEmpty)
+            actual <-> expected
+          }
         }
       ),
       suite("constructors")(

diff --git a/core/shared/src/main/scala/zio/prelude/NonEmptySet.scala b/core/shared/src/main/scala/zio/prelude/NonEmptySet.scala
@@ -101,7 +101,7 @@ final class NonEmptySet[A] private (private val set: Set[A]) { self =>
    */
   def removeNonEmpty(elem: A): Option[NonEmptySet[A]] = {
     val newSet = set - elem
-    if (newSet.nonEmpty) Some(new NonEmptySet(set)) else None
+    if (newSet.nonEmpty) Some(new NonEmptySet(newSet)) else None
   }
 
   /**

diff --git a/core/shared/src/main/scala/zio/prelude/NonEmptySortedSet.scala b/core/shared/src/main/scala/zio/prelude/NonEmptySortedSet.scala
@@ -103,7 +103,7 @@ final class NonEmptySortedSet[A] private (private val set: SortedSet[A]) { self
    */
   def removeNonEmpty(elem: A): Option[NonEmptySortedSet[A]] = {
     val newSet = set - elem
-    if (newSet.nonEmpty) Some(new NonEmptySortedSet(set)) else None
+    if (newSet.nonEmpty) Some(new NonEmptySortedSet(newSet)) else None
   }
 
   /**