From 6ad666c0615649760eb28c671c68cfdff30bda72 Mon Sep 17 00:00:00 2001
From: Maksim Ochenashko <maxochenashko@gmail.com>
Date: Fri, 13 Jan 2023 12:05:21 +0200
Subject: [PATCH] IOLocal - generalize `scope` function

---
 core/shared/src/main/scala/cats/effect/IOLocal.scala | 12 ++++++------
 .../src/test/scala/cats/effect/IOLocalSpec.scala     |  2 +-
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/core/shared/src/main/scala/cats/effect/IOLocal.scala b/core/shared/src/main/scala/cats/effect/IOLocal.scala
index 4d781a1a50d..eb2f6d3e589 100644
--- a/core/shared/src/main/scala/cats/effect/IOLocal.scala
+++ b/core/shared/src/main/scala/cats/effect/IOLocal.scala
@@ -201,15 +201,15 @@ sealed trait IOLocal[A] {
    *   for {
    *     local <- IOLocal(42)
    *     _     <- local.get // returns 42
-   *     _     <- local.scope(0).surround(local.getAndSet(1)) // returns 0
+   *     _     <- local.scope(current => current + 1).surround(local.getAndSet(1)) // returns 43
    *     _     <- local.get // returns 42, even though 1 was set inside of the resource
    *   } yield ()
    * }}}
    *
-   * @param value
-   *   the value to make a scope with
+   * @param f
+   *   the function to make a new scope
    */
-  def scope(value: A): Resource[IO, Unit]
+  def scope(f: A => A): Resource[IO, Unit]
 
 }
 
@@ -256,8 +256,8 @@ object IOLocal {
         override def getAndReset: IO[A] =
           IO.Local(state => (state - self, getOrDefault(state)))
 
-        override def scope(value: A): Resource[IO, Unit] =
-          Resource.make(getAndSet(value))(p => set(p)).void
+        override def scope(f: A => A): Resource[IO, Unit] =
+          Resource.make(modify(a => (f(a), a)))(p => set(p)).void
       }
     }
 
diff --git a/tests/shared/src/test/scala/cats/effect/IOLocalSpec.scala b/tests/shared/src/test/scala/cats/effect/IOLocalSpec.scala
index c525f9a252d..d8b642c9829 100644
--- a/tests/shared/src/test/scala/cats/effect/IOLocalSpec.scala
+++ b/tests/shared/src/test/scala/cats/effect/IOLocalSpec.scala
@@ -85,7 +85,7 @@ class IOLocalSpec extends BaseSpec {
     "do not leak internal updates outside of a scope" in ticked { implicit ticker =>
       val io = for {
         local <- IOLocal(0)
-        inside <- local.scope(1).surround(local.getAndSet(2))
+        inside <- local.scope(_ => 1).surround(local.getAndSet(2))
         outside <- local.get
       } yield (inside, outside)