Soviet2.scala

object Soviet {
  def repeated[F[_]: Async, A, B, R](runner: (A => B) => F[R])(body: A => F[B]): F[R] =
    for {
      ref <- Sync[F].delay(new AtomicReference[A => B])
      result <- Sync[F].delay(new AtomicReference[B])

      latch <- Sync[F].delay(new Semaphore(1))
      _ <- Sync[F].delay(latch.acquire())

      r <- {
        def adapt(cb: Either[Throwable, A] => Unit): A => B = { a =>
          cb(Right(a))
          latch.acquire()
          result.get()
        }

        def redirect(a: A): B = {
          var f: A => B = ref.getAndSet(null)
          while (f == null) {
            ref.await()
            f = ref.getAndSet(null)
          }

          f(a)
        }

        val registerF = {
          val cont: F[A] = Async[F] async_ { cb =>
            ref.set(adapt(cb))
            ref.notifyAll()
          }

          cont.flatMap(body) flatMap { b =>
            Sync[F] delay {
              result.set(b)
              latch.release()
            }
          }
        }

        registerF.foreverM.start.bracket({ _ =>
          runner(redirect _)
        })(_.cancel)
      }
    } yield r
}