From a69f4c38d8702c3c46ef294179a7d1d1cb8def96 Mon Sep 17 00:00:00 2001
From: Alex <alex.simkin@outlook.com>
Date: Thu, 16 Jan 2025 02:14:13 +1100
Subject: [PATCH] IdentityFlatten#when and IdentityFlatten#unless. (#1453)

---
 .../zio/prelude/IdentityFlattenSpec.scala     | 25 +++++++++++++++++++
 .../scala/zio/prelude/IdentityFlatten.scala   | 14 +++++++++++
 .../src/main/scala/zio/prelude/package.scala  |  1 +
 3 files changed, 40 insertions(+)

diff --git a/core-tests/shared/src/test/scala/zio/prelude/IdentityFlattenSpec.scala b/core-tests/shared/src/test/scala/zio/prelude/IdentityFlattenSpec.scala
index 4736b1e59..9a35ca76e 100644
--- a/core-tests/shared/src/test/scala/zio/prelude/IdentityFlattenSpec.scala
+++ b/core-tests/shared/src/test/scala/zio/prelude/IdentityFlattenSpec.scala
@@ -7,6 +7,15 @@ import zio.test.laws._
 object IdentityFlattenSpec extends ZIOBaseSpec {
   import zio.prelude.Fixtures._
 
+  val genBoolean: Gen[Any, Boolean] =
+    Gen.boolean
+
+  val genInt: Gen[Any, Int] =
+    Gen.int
+
+  val genList: Gen[Sized, List[Int]] =
+    Gen.listOf(genInt)
+
   def spec: Spec[Environment, Any] =
     suite("IdentityFlattenSpec")(
       suite("laws")(
@@ -16,6 +25,22 @@ object IdentityFlattenSpec extends ZIOBaseSpec {
         test("option")(checkAllLaws(IdentityFlattenLaws)(GenF.option, Gen.int)),
         test("optional")(checkAllLaws(IdentityFlattenLaws)(optionalGenF, Gen.int)),
         test("vector")(checkAllLaws(IdentityFlattenLaws)(GenF.vector, Gen.int))
+      ),
+      suite("combinators")(
+        test("when") {
+          check(genList, genBoolean) { (as, b) =>
+            val actual   = as.when(b)
+            val expected = if (b) as.map(Some(_)) else List(None)
+            assert(actual)(equalTo(expected))
+          }
+        },
+        test("unless") {
+          check(genList, genBoolean) { (as, b) =>
+            val actual   = as.unless(b)
+            val expected = if (b) List(None) else as.map(Some(_))
+            assert(actual)(equalTo(expected))
+          }
+        }
       )
     )
 
diff --git a/core/shared/src/main/scala/zio/prelude/IdentityFlatten.scala b/core/shared/src/main/scala/zio/prelude/IdentityFlatten.scala
index 2ac0c69c1..db10175e2 100644
--- a/core/shared/src/main/scala/zio/prelude/IdentityFlatten.scala
+++ b/core/shared/src/main/scala/zio/prelude/IdentityFlatten.scala
@@ -43,3 +43,17 @@ object IdentityFlatten {
     identityFlatten
 
 }
+
+trait IdentityFlattenSyntax {
+
+  /**
+   * Provides infix syntax for identity operations for covariant types.
+   */
+  implicit class IdentityFlattenCovariantOps[F[+_], A](fa: F[A]) {
+    def unless(b: => Boolean)(implicit identity: IdentityFlatten[F], covariant: Covariant[F]): F[Option[A]] =
+      if (b) identity.any.as(None) else fa.map(Some(_))
+
+    def when(b: => Boolean)(implicit identity: IdentityFlatten[F], covariant: Covariant[F]): F[Option[A]] =
+      if (b) fa.map(Some(_)) else identity.any.as(None)
+  }
+}
diff --git a/core/shared/src/main/scala/zio/prelude/package.scala b/core/shared/src/main/scala/zio/prelude/package.scala
index 47be34c5e..ceb854989 100644
--- a/core/shared/src/main/scala/zio/prelude/package.scala
+++ b/core/shared/src/main/scala/zio/prelude/package.scala
@@ -38,6 +38,7 @@ package object prelude
     with IdExports
     with IdentityBothSyntax
     with IdentityEitherSyntax
+    with IdentityFlattenSyntax
     with IdentitySyntax
     with InvariantSyntax
     with InverseSyntax