zio · jdegoes · Aug 17, 2023 · Aug 16, 2023 · Aug 16, 2023 · Aug 16, 2023
diff --git a/core/shared/src/main/scala/zio/FiberRef.scala b/core/shared/src/main/scala/zio/FiberRef.scala
@@ -362,6 +362,17 @@ object FiberRef {
   ): ZIO[Scope, Nothing, FiberRef.WithPatch[ZEnvironment[A], ZEnvironment.Patch[A, A]]] =
     makeWith(unsafe.makeEnvironment(initial)(Unsafe.unsafe))
 
+  /**
+   * Creates a new `FiberRef` with the specified initial value, using the
+   * specified patch type to combine updates to the value in a compositional
+   * way.
+   */
+  def makePatch[Value, Patch](
+    initial: Value,
+    differ: Differ[Value, Patch]
+  )(implicit trace: Trace): ZIO[Scope, Nothing, FiberRef.WithPatch[Value, Patch]] =
+    makePatch(initial, differ, differ.empty)
+
   /**
    * Creates a new `FiberRef` with the specified initial value, using the
    * specified patch type to combine updates to the value in a compositional

diff --git a/core/shared/src/main/scala/zio/ZIO.scala b/core/shared/src/main/scala/zio/ZIO.scala
@@ -4601,6 +4601,14 @@ object ZIO extends ZIOCompanionPlatformSpecific with ZIOCompanionVersionSpecific
   def stateful[R]: StatefulPartiallyApplied[R] =
     new StatefulPartiallyApplied[R]
 
+  /**
+   * Provides a stateful ZIO workflow with its initial state, using the
+   * specified patch type to combine updates to the state by different fibers in
+   * a compositional way.
+   */
+  def statefulPatch[R]: StatefulPatchPartiallyApplied[R] =
+    new StatefulPatchPartiallyApplied[R]
+
   def succeedBlockingUnsafe[A](a: Unsafe => A)(implicit trace: Trace): UIO[A] =
     ZIO.blocking(ZIO.succeedUnsafe(a))
 
@@ -5383,6 +5391,17 @@ object ZIO extends ZIOCompanionPlatformSpecific with ZIOCompanionVersionSpecific
       zio.provideSomeLayer[R](ZState.initial(s))
   }
 
+  final class StatefulPatchPartiallyApplied[R](private val dummy: Boolean = true) extends AnyVal {
+    def apply[State, Patch, E, A](
+      state: State,
+      differ: Differ[State, Patch]
+    )(zio: => ZIO[ZState[State] with R, E, A])(implicit
+      tag: EnvironmentTag[State],
+      trace: Trace
+    ): ZIO[R, E, A] =
+      zio.provideSomeLayer[R](ZState.initialPatch(state, differ))
+  }
+
   final class GetStateWithPartiallyApplied[S](private val dummy: Boolean = true) extends AnyVal {
     def apply[A](f: S => A)(implicit tag: EnvironmentTag[S], trace: Trace): ZIO[ZState[S], Nothing, A] =
       ZIO.serviceWithZIO(_.get.map(f))

diff --git a/core/shared/src/main/scala/zio/ZState.scala b/core/shared/src/main/scala/zio/ZState.scala
@@ -63,15 +63,25 @@ object ZState {
    * A layer that allocates the initial state of a stateful workflow.
    */
   def initial[S: EnvironmentTag](s: => S)(implicit trace: Trace): ZLayer[Any, Nothing, ZState[S]] =
+    initialPatch(s, Differ.update[S])
+
+  /**
+   * A layer that allocates the initial state of a stateful workflow, using the
+   * specified patch type to combine updates to the state by different fibers in
+   * a compositional way.
+   */
+  def initialPatch[State: EnvironmentTag, Patch](state: => State, differ: => Differ[State, Patch])(implicit
+    trace: Trace
+  ): ZLayer[Any, Nothing, ZState[State]] =
     ZLayer.scoped {
       for {
-        fiberRef <- FiberRef.make(s)
-      } yield new ZState[S] {
-        def get(implicit trace: Trace): UIO[S] =
+        fiberRef <- FiberRef.makePatch(state, differ)
+      } yield new ZState[State] {
+        def get(implicit trace: Trace): UIO[State] =
           fiberRef.get
-        def set(s: S)(implicit trace: Trace): UIO[Unit] =
-          fiberRef.set(s)
-        def update(f: S => S)(implicit trace: Trace): UIO[Unit] =
+        def set(state: State)(implicit trace: Trace): UIO[Unit] =
+          fiberRef.set(state)
+        def update(f: State => State)(implicit trace: Trace): UIO[Unit] =
           fiberRef.update(f)
       }
     }