From 25b163bb48045f1ec8edc64266f3ba87274d5d83 Mon Sep 17 00:00:00 2001
From: Tony Won <tony@daangn.com>
Date: Wed, 15 Jan 2025 19:17:42 +0900
Subject: [PATCH] fix

---
 .../react/src/future/loader/loaderPlugin.tsx  | 154 +++++++-----------
 1 file changed, 58 insertions(+), 96 deletions(-)

diff --git a/integrations/react/src/future/loader/loaderPlugin.tsx b/integrations/react/src/future/loader/loaderPlugin.tsx
index facb33cc..a6aa3c88 100644
--- a/integrations/react/src/future/loader/loaderPlugin.tsx
+++ b/integrations/react/src/future/loader/loaderPlugin.tsx
@@ -57,101 +57,63 @@ export function loaderPlugin<
         };
       });
     },
-    onBeforePush({
-      actionParams,
-      actions: { overrideActionParams, pause, resume },
-    }) {
-      const { activityName, activityParams, activityContext } = actionParams;
-
-      const matchActivity = input.config.activities.find(
-        (activity) => activity.name === activityName,
-      );
-      const matchActivityComponent =
-        input.components[activityName as T["name"]];
-
-      const loader = matchActivity?.loader;
-
-      if (!loader || !matchActivityComponent) {
-        return;
-      }
-
-      const loaderData = loader({
-        params: activityParams,
-        config: input.config,
-      });
-
-      if (loaderData instanceof Promise || "_load" in matchActivityComponent) {
-        pause();
-      }
-
-      const promises: Array<Promise<any>> = [];
-
-      if (loaderData instanceof Promise) {
-        promises.push(loaderData);
-      }
-      if ("_load" in matchActivityComponent && matchActivityComponent._load) {
-        promises.push(matchActivityComponent._load());
-      }
-
-      Promise.all(promises).finally(() => {
-        resume();
-      });
-
-      overrideActionParams({
-        ...actionParams,
-        activityContext: {
-          ...activityContext,
-          loaderData,
-        },
-      });
-    },
-    onBeforeReplace({
-      actionParams,
-      actions: { overrideActionParams, pause, resume },
-    }) {
-      const { activityName, activityParams, activityContext } = actionParams;
-
-      const matchActivity = input.config.activities.find(
-        (activity) => activity.name === activityName,
-      );
-      const matchActivityComponent =
-        input.components[activityName as T["name"]];
-
-      const loader = matchActivity?.loader;
-
-      if (!loader || !matchActivityComponent) {
-        return;
-      }
-
-      const loaderData = loader({
-        params: activityParams,
-        config: input.config,
-      });
-
-      if (loaderData instanceof Promise || "_load" in matchActivityComponent) {
-        pause();
-      }
-
-      const promises: Array<Promise<any>> = [];
-
-      if (loaderData instanceof Promise) {
-        promises.push(loaderData);
-      }
-      if ("_load" in matchActivityComponent && matchActivityComponent._load) {
-        promises.push(matchActivityComponent._load());
-      }
-
-      Promise.all(promises).finally(() => {
-        resume();
-      });
-
-      overrideActionParams({
-        ...actionParams,
-        activityContext: {
-          ...activityContext,
-          loaderData,
-        },
-      });
-    },
+    onBeforePush: createBeforeRouteHandler(input),
+    onBeforeReplace: createBeforeRouteHandler(input),
   });
 }
+
+type OnBeforeRoute = NonNullable<
+  | ReturnType<StackflowReactPlugin>["onBeforePush"]
+  | ReturnType<StackflowReactPlugin>["onBeforeReplace"]
+>;
+function createBeforeRouteHandler<
+  T extends ActivityDefinition<string>,
+  R extends {
+    [activityName in T["name"]]: ActivityComponentType<any>;
+  },
+>(input: StackflowInput<T, R>): OnBeforeRoute {
+  return ({
+    actionParams,
+    actions: { overrideActionParams, pause, resume },
+  }) => {
+    const { activityName, activityParams, activityContext } = actionParams;
+
+    const matchActivity = input.config.activities.find(
+      (activity) => activity.name === activityName,
+    );
+    const matchActivityComponent = input.components[activityName as T["name"]];
+
+    const loader = matchActivity?.loader;
+
+    if (!loader || !matchActivityComponent) {
+      return;
+    }
+
+    const loaderData = loader({
+      params: activityParams,
+      config: input.config,
+    });
+
+    const loaderDataPromise =
+      loaderData instanceof Promise ? loaderData : undefined;
+    const lazyComponentPromise =
+      "_load" in matchActivityComponent
+        ? matchActivityComponent._load?.()
+        : undefined;
+
+    if (loaderDataPromise || lazyComponentPromise) {
+      pause();
+    }
+    Promise.all([loaderDataPromise, lazyComponentPromise]).finally(() => {
+      resume();
+    });
+
+    overrideActionParams({
+      ...actionParams,
+      activityContext: {
+        ...activityContext,
+        loaderData,
+      },
+    });
+  };
+}