From 7ade36e83aa7344586bd3780a8333a189dfe95ec Mon Sep 17 00:00:00 2001
From: Gabriel Nordeborn <gabbe.nord@gmail.com>
Date: Fri, 2 Aug 2024 17:32:13 +0200
Subject: [PATCH] reuse main router transition in links

---
 .changeset/twenty-deers-hope.md                |  5 +++++
 .../rescript-relay-router/src/RelayRouter.res  | 18 ++++++++----------
 .../src/RelayRouter__Internal.res              | 12 ++++++++++++
 .../src/RelayRouter__Internal.resi             | 10 ++++++++++
 .../src/RelayRouter__Link.res                  |  2 +-
 5 files changed, 36 insertions(+), 11 deletions(-)
 create mode 100644 .changeset/twenty-deers-hope.md

diff --git a/.changeset/twenty-deers-hope.md b/.changeset/twenty-deers-hope.md
new file mode 100644
index 0000000..1109945
--- /dev/null
+++ b/.changeset/twenty-deers-hope.md
@@ -0,0 +1,5 @@
+---
+"rescript-relay-router": patch
+---
+
+Make sure links also reuse the same main router transition
diff --git a/packages/rescript-relay-router/src/RelayRouter.res b/packages/rescript-relay-router/src/RelayRouter.res
index 876d75d..41f7367 100644
--- a/packages/rescript-relay-router/src/RelayRouter.res
+++ b/packages/rescript-relay-router/src/RelayRouter.res
@@ -160,14 +160,12 @@ module Router = {
             ~pathParams=match.params,
             ~queryParams,
             ~location,
-          )->(
-            Js.Promise.then_(
-              assetsToPreload => {
-                assetsToPreload->Belt.Array.forEach(a => a->preloadAsset(~priority))
-                Js.Promise.resolve()
-              },
-              _,
-            )
+          )->Js.Promise.then_(
+            assetsToPreload => {
+              assetsToPreload->Belt.Array.forEach(a => a->preloadAsset(~priority))
+              Js.Promise.resolve()
+            },
+            _,
           )
         }, ~priority)
       })
@@ -280,13 +278,13 @@ module RouteRenderer = {
       </RouteComponent>
     })
 
-    <>
+    <RelayRouter__Internal.RouterTransitionContext.Provider value=startTransition>
       {switch renderPending {
       | Some(renderPending) => renderPending(isPending)
       | None => React.null
       }}
       {renderedContent.contents}
-    </>
+    </RelayRouter__Internal.RouterTransitionContext.Provider>
   }
 }
 
diff --git a/packages/rescript-relay-router/src/RelayRouter__Internal.res b/packages/rescript-relay-router/src/RelayRouter__Internal.res
index f119440..737b049 100644
--- a/packages/rescript-relay-router/src/RelayRouter__Internal.res
+++ b/packages/rescript-relay-router/src/RelayRouter__Internal.res
@@ -137,3 +137,15 @@ let runAtPriority = (cb, ~priority) => {
     None
   }
 }
+
+module RouterTransitionContext = {
+  type transitionFn = (unit => unit) => unit
+
+  let context = React.createContext(_cb => ())
+
+  module Provider = {
+    let make = React.Context.provider(context)
+  }
+
+  let use = (): transitionFn => React.useContext(context)
+}
diff --git a/packages/rescript-relay-router/src/RelayRouter__Internal.resi b/packages/rescript-relay-router/src/RelayRouter__Internal.resi
index 1210320..f2a2dce 100644
--- a/packages/rescript-relay-router/src/RelayRouter__Internal.resi
+++ b/packages/rescript-relay-router/src/RelayRouter__Internal.resi
@@ -31,3 +31,13 @@ let runAtPriority: (
   RelayRouter__Types.callback,
   ~priority: RelayRouter__Types.preloadPriority,
 ) => option<unit => unit>
+
+module RouterTransitionContext: {
+  type transitionFn = (unit => unit) => unit
+  let context: React.Context.t<transitionFn>
+  let use: unit => transitionFn
+
+  module Provider: {
+    let make: React.component<React.Context.props<transitionFn>>
+  }
+}
diff --git a/packages/rescript-relay-router/src/RelayRouter__Link.res b/packages/rescript-relay-router/src/RelayRouter__Link.res
index 57672a7..a36100d 100644
--- a/packages/rescript-relay-router/src/RelayRouter__Link.res
+++ b/packages/rescript-relay-router/src/RelayRouter__Link.res
@@ -52,10 +52,10 @@ let make = (
 ) => {
   let linkElement = React.useRef(Js.Nullable.null)
   let hasPreloaded = React.useRef(false)
-  let (_, startTransition) = React.useTransition()
   let router = RelayRouter__Context.useRouterContext()
   let {history} = router
   let targetElementRef = TargetScrollElement.useTargetElement()
+  let startTransition = RelayRouter__Internal.RouterTransitionContext.use()
 
   let changeRoute = React.useCallback(e =>
     startTransition(() => {