From bcbc64bf2a271d17e770fc12fa51b195015b2ba4 Mon Sep 17 00:00:00 2001
From: Daniel Urban <urban.dani@gmail.com>
Date: Fri, 24 May 2024 20:31:46 +0200
Subject: [PATCH] Make Defer#fix safer

---
 .../choam/core/RxnCompanionPlatform.scala     |  8 +++++++
 .../choam/core/RxnCompanionPlatform.scala     | 12 ++++++++++
 .../main/scala/dev/tauri/choam/core/Rxn.scala | 23 +++++++++++++++----
 3 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/core/js/src/main/scala/dev/tauri/choam/core/RxnCompanionPlatform.scala b/core/js/src/main/scala/dev/tauri/choam/core/RxnCompanionPlatform.scala
index 7a65385fd..165ea9fde 100644
--- a/core/js/src/main/scala/dev/tauri/choam/core/RxnCompanionPlatform.scala
+++ b/core/js/src/main/scala/dev/tauri/choam/core/RxnCompanionPlatform.scala
@@ -23,4 +23,12 @@ private abstract class RxnCompanionPlatform { this: Rxn.type =>
   private[core] final type ExchangerImpl[A, B] = ExchangerImplJs[A, B]
 
   private[core] final type ExStatMap = Map[Exchanger.Key, AnyRef]
+
+  @inline
+  protected[this] final def releaseFence(): Unit =
+    ()
+
+  @inline
+  protected[this] final def acquireFence(): Unit =
+    ()
 }
diff --git a/core/jvm/src/main/scala/dev/tauri/choam/core/RxnCompanionPlatform.scala b/core/jvm/src/main/scala/dev/tauri/choam/core/RxnCompanionPlatform.scala
index f0e585c45..145d38121 100644
--- a/core/jvm/src/main/scala/dev/tauri/choam/core/RxnCompanionPlatform.scala
+++ b/core/jvm/src/main/scala/dev/tauri/choam/core/RxnCompanionPlatform.scala
@@ -18,9 +18,21 @@
 package dev.tauri.choam
 package core
 
+import java.lang.invoke.VarHandle
+
 private abstract class RxnCompanionPlatform { this: Rxn.type =>
 
   private[core] final type ExchangerImpl[A, B] = ExchangerImplJvm[A, B]
 
   private[core] final type ExStatMap = ExchangerImplJvm.StatMap
+
+  @inline
+  protected[this] final def releaseFence(): Unit = {
+    VarHandle.releaseFence()
+  }
+
+  @inline
+  protected[this] final def acquireFence(): Unit = {
+    VarHandle.acquireFence()
+  }
 }
diff --git a/core/shared/src/main/scala/dev/tauri/choam/core/Rxn.scala b/core/shared/src/main/scala/dev/tauri/choam/core/Rxn.scala
index 6c3f79ac4..4c49786f9 100644
--- a/core/shared/src/main/scala/dev/tauri/choam/core/Rxn.scala
+++ b/core/shared/src/main/scala/dev/tauri/choam/core/Rxn.scala
@@ -1604,12 +1604,25 @@ private sealed abstract class RxnInstances6 extends RxnInstances7 { self: Rxn.ty
     final override def defer[A](fa: => Rxn[X, A]): Rxn[X, A] =
       self.computed[X, A] { x => fa.provide(x) }
     final override def fix[A](fn: Rxn[X, A] => Rxn[X, A]): Rxn[X, A] = {
-      // This is in effect a "thread-unsafe lazy val";
-      // we know exactly how `defer` works, so it's safe
-      // to do this here (and this way we avoid the
-      // synchronization of an actual lazy val):
+      // Instead of a `lazy val` (like in the superclass), we just
+      // do a rel/acq here, because we know exactly how `defer`
+      // works, and know that `.elem` will be initialized before
+      // we return from this method. However, we still need the
+      // fences, to make sure that if the resulting `Rxn[X, A]`
+      // is published by a race, there is an ordering between
+      // writing `ref.elem` and reading it.
+      // TODO: The point of this whole thing is to avoid a
+      // TODO: `lazy val`, which might block. This way of
+      // TODO: doing it is correct, but it's unclear if it's
+      // TODO: faster than a `lazy val`, and also, it could
+      // TODO: be that in this specific case, a `lazy val`
+      // TODO: also woudn't block.
       val ref = new scala.runtime.ObjectRef[Rxn[X, A]](null)
-      ref.elem = fn(defer { ref.elem })
+      ref.elem = fn(defer {
+        self.acquireFence()
+        ref.elem
+      })
+      self.releaseFence()
       ref.elem
     }
   }