From e12fe314ece8f30b36a583cdb6ff88fedadcdfde Mon Sep 17 00:00:00 2001
From: David Sancho Moreno <dsnxmoreno@gmail.com>
Date: Mon, 25 Nov 2024 18:33:57 +0100
Subject: [PATCH] Implement ref with cleanup

---
 src/ReactDOM.re     |  4 ++++
 src/ReactDOM.rei    |  3 +++
 test/Ref__test.re   | 51 ++++++++++++++++++++++++++++++++++++++++++++-
 test/jest/Expect.re |  6 ++++++
 test/jest/Mock.re   |  2 +-
 5 files changed, 64 insertions(+), 2 deletions(-)

diff --git a/src/ReactDOM.re b/src/ReactDOM.re
index bedb622f2..01d852bf7 100644
--- a/src/ReactDOM.re
+++ b/src/ReactDOM.re
@@ -487,9 +487,13 @@ module Ref = {
   type t = domRef;
   type currentDomRef = React.ref(Js.nullable(Dom.element));
   type callbackDomRef = Js.nullable(Dom.element) => unit;
+  type callbackRefWithCleanup = (Js.nullable(Dom.element), unit) => unit;
 
   external domRef: currentDomRef => domRef = "%identity";
   external callbackDomRef: callbackDomRef => domRef = "%identity";
+  external callbackRefWithCleanup:
+    [@mel.unwrap] (callbackRefWithCleanup => domRef) =
+    "%identity";
 };
 
 /* This list isn't exhaustive. We'll add more as we go. */
diff --git a/src/ReactDOM.rei b/src/ReactDOM.rei
index 9a7eee88c..d8212604c 100644
--- a/src/ReactDOM.rei
+++ b/src/ReactDOM.rei
@@ -488,9 +488,12 @@ module Ref: {
   type t = domRef;
   type currentDomRef = React.ref(Js.nullable(Dom.element));
   type callbackDomRef = Js.nullable(Dom.element) => unit;
+  type callbackRefWithCleanup = (Js.nullable(Dom.element), unit) => unit;
 
   external domRef: currentDomRef => domRef = "%identity";
   external callbackDomRef: callbackDomRef => domRef = "%identity";
+  external callbackRefWithCleanup: callbackRefWithCleanup => domRef =
+    "%identity";
 };
 
 /* This list isn't exhaustive. We'll add more as we go. */
diff --git a/test/Ref__test.re b/test/Ref__test.re
index 223a5f348..484aa049a 100644
--- a/test/Ref__test.re
+++ b/test/Ref__test.re
@@ -31,5 +31,54 @@ describe("ref", () => {
       | None => failwith("No element found")
       };
     expect(content)->toBe("Click me");
-  })
+  });
+
+  test("with callback ref", () => {
+    let callbackRef = Mock.fn();
+
+    let refWithCleanup =
+      ReactDOM.Ref.callbackDomRef(_ref => {ignore(callbackRef())});
+
+    let _container =
+      ReactTestingLibrary.render(
+        <Button_with_ref ref=refWithCleanup>
+          {React.string("Click me")}
+        </Button_with_ref>,
+      );
+
+    expect(callbackRef->Mock.getMock)->toHaveBeenCalled();
+  });
+
+  let (let.await) = (p, f) => Js.Promise.then_(f, p);
+
+  testPromise("with callback ref and cleanup", finish => {
+    let callbackRef = Mock.fnWithImplementation(_ => ());
+    let callbackCleanupRef = Mock.fnWithImplementation(_ => ());
+
+    let refWithCleanup =
+      ReactDOM.Ref.callbackRefWithCleanup(_ref => {
+        callbackRef();
+        () => {
+          callbackCleanupRef();
+        };
+      });
+
+    let.await _ =
+      ReactTestingLibrary.actAsync(() => {
+        let _container =
+          ReactTestingLibrary.render(
+            <Button_with_ref ref=refWithCleanup>
+              {React.string("Click me")}
+            </Button_with_ref>,
+          );
+        Js.Promise.resolve();
+      });
+
+    expect(callbackRef->Mock.getMock)->toHaveBeenCalledTimes(1);
+    expect(callbackCleanupRef->Mock.getMock)->toHaveBeenCalledTimes(0);
+
+    ReactTestingLibrary.cleanup();
+    expect(callbackCleanupRef->Mock.getMock)->toHaveBeenCalledTimes(1);
+    finish();
+  });
 });
diff --git a/test/jest/Expect.re b/test/jest/Expect.re
index 0b1e0a6a0..bd0f1e37c 100644
--- a/test/jest/Expect.re
+++ b/test/jest/Expect.re
@@ -20,6 +20,12 @@ external toBeLessThanOrEqual: (t('a), 'a) => unit = "toBeLessThanOrEqual";
 [@mel.send]
 external toHaveLength: (t(array('a)), 'a) => unit = "toHaveLength";
 
+[@mel.send]
+external toHaveBeenCalled: (t('a), unit) => unit = "toHaveBeenCalled";
+[@mel.send]
+external toHaveBeenCalledTimes: (t('a), int) => unit =
+  "toHaveBeenCalledTimes";
+
 [@mel.get]
 external rejects: t(Js.Promise.t('a)) => t(unit => 'a) = "rejects";
 
diff --git a/test/jest/Mock.re b/test/jest/Mock.re
index 2301b7564..04401920f 100644
--- a/test/jest/Mock.re
+++ b/test/jest/Mock.re
@@ -5,7 +5,7 @@ let undefined: undefined = Js.Undefined.empty;
 type mock('a);
 
 type t('a) = mock('a);
-[@mel.scope "jest"] external fn: unit => 'a = "fn";
+[@mel.scope "jest"] external fn: [@mel.unwrap] (unit => 'a) = "fn";
 [@mel.scope "jest"] external fnWithImplementation: 'a => 'a = "fn";
 [@mel.scope "jest"] external mockModule: string => unit = "mock";
 external getMock: 'a => t('a) = "%identity";