From e1b128f329f1e88d027ca99f949feb135858a276 Mon Sep 17 00:00:00 2001
From: Taylor Lodge <taylor.lodge@koordinates.com>
Date: Fri, 15 Mar 2024 14:37:26 +1300
Subject: [PATCH] fix(builders): patch provide to function correctly

Due to how provide works it loses the _xstateTree property that is attached to xstate-tree machines, making them unusable.

It also fixes the return type to return an XstateTreeMachine instead of a StateMachine
---
 src/builders.spec.tsx | 26 ++++++++++++++++++++++++++
 src/builders.tsx      | 18 +++++++++++++++++-
 src/lazy.tsx          |  4 ++--
 src/types.ts          | 24 +++++++++++++++++++++++-
 xstate-tree.api.md    |  7 ++++---
 5 files changed, 72 insertions(+), 7 deletions(-)

diff --git a/src/builders.spec.tsx b/src/builders.spec.tsx
index de7249d..a4922b7 100644
--- a/src/builders.spec.tsx
+++ b/src/builders.spec.tsx
@@ -53,6 +53,32 @@ describe("xstate-tree builders", () => {
         }}
       />;
     });
+
+    it("repairs the provide function to not lose the _xstateTree property and return an XstateTreeMachine", () => {
+      const machine = setup({
+        actions: {
+          someAction: () => {},
+        },
+      }).createMachine({
+        initial: "idle",
+        states: {
+          idle: {},
+        },
+      });
+
+      const xstateTreeMachine = createXStateTreeMachine(machine, {
+        View() {
+          return <div>hello world</div>;
+        },
+      }).provide({
+        actions: {
+          someAction: () => {},
+        },
+      });
+
+      const Root = buildRootComponent({ machine: xstateTreeMachine });
+      render(<Root />);
+    });
   });
 
   describe("viewToMachine", () => {
diff --git a/src/builders.tsx b/src/builders.tsx
index 2219e02..4616a25 100644
--- a/src/builders.tsx
+++ b/src/builders.tsx
@@ -54,7 +54,23 @@ export function createXStateTreeMachine<
     slots: (options.slots ?? []) as any,
   };
 
-  return machineWithMeta;
+  return fixProvideLosingXstateTreeMeta(machineWithMeta);
+}
+
+function fixProvideLosingXstateTreeMeta<
+  T extends XstateTreeMachine<any, any, any, any>
+>(machine: T): T {
+  const originalProvide = machine.provide.bind(machine);
+  (machine as any).provide = (impl: any) => {
+    const result = originalProvide(impl) as T;
+
+    result._xstateTree = machine._xstateTree;
+    fixProvideLosingXstateTreeMeta(result);
+
+    return result;
+  };
+
+  return machine;
 }
 
 /**
diff --git a/src/lazy.tsx b/src/lazy.tsx
index bf4b5d5..92a0756 100644
--- a/src/lazy.tsx
+++ b/src/lazy.tsx
@@ -1,11 +1,11 @@
 import React from "react";
-import { AnyStateMachine, assign, setup, fromPromise, InputFrom } from "xstate";
+import { assign, setup, fromPromise, InputFrom } from "xstate";
 
 import { createXStateTreeMachine } from "./builders";
 import { singleSlot } from "./slots";
 import { AnyXstateTreeMachine } from "./types";
 
-type Options<TStateMachine extends AnyStateMachine> = {
+type Options<TStateMachine extends AnyXstateTreeMachine> = {
   /**
    * Displayed while the promise is resolving, defaults to returning null
    */
diff --git a/src/types.ts b/src/types.ts
index 320ddf4..0e7d931 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -102,6 +102,23 @@ export type XstateTreeMachineInjection<
   >;
 };
 
+/**
+ * Repairs the return type of the `provide` function on XstateTreeMachines to correctly return
+ * an XstateTreeMachine type instead of an xstate StateMachine
+ */
+type RepairProvideReturnType<
+  T extends AnyStateMachine,
+  TSelectorsOutput,
+  TActionsOutput,
+  TSlots extends readonly Slot[]
+> = {
+  [K in keyof T]: K extends "provide"
+    ? (
+        ...args: Parameters<T[K]>
+      ) => XstateTreeMachine<T, TSelectorsOutput, TActionsOutput, TSlots>
+    : T[K];
+};
+
 /**
  * @public
  */
@@ -110,7 +127,12 @@ export type XstateTreeMachine<
   TSelectorsOutput = ContextFrom<TMachine>,
   TActionsOutput = Record<never, string>,
   TSlots extends readonly Slot[] = Slot[]
-> = TMachine &
+> = RepairProvideReturnType<
+  TMachine,
+  TSelectorsOutput,
+  TActionsOutput,
+  TSlots
+> &
   XstateTreeMachineInjection<
     TMachine,
     TSelectorsOutput,
diff --git a/xstate-tree.api.md b/xstate-tree.api.md
index 6de2ab7..103381b 100644
--- a/xstate-tree.api.md
+++ b/xstate-tree.api.md
@@ -380,10 +380,11 @@ export type XstateTreeHistory<T = unknown> = History_2<{
     previousUrl?: string;
 }>;
 
+// Warning: (ae-forgotten-export) The symbol "RepairProvideReturnType" needs to be exported by the entry point index.d.ts
 // Warning: (ae-incompatible-release-tags) The symbol "XstateTreeMachine" is marked as @public, but its signature references "XstateTreeMachineInjection" which is marked as @internal
 //
 // @public (undocumented)
-export type XstateTreeMachine<TMachine extends AnyStateMachine, TSelectorsOutput = ContextFrom<TMachine>, TActionsOutput = Record<never, string>, TSlots extends readonly Slot[] = Slot[]> = TMachine & XstateTreeMachineInjection<TMachine, TSelectorsOutput, TActionsOutput, TSlots>;
+export type XstateTreeMachine<TMachine extends AnyStateMachine, TSelectorsOutput = ContextFrom<TMachine>, TActionsOutput = Record<never, string>, TSlots extends readonly Slot[] = Slot[]> = RepairProvideReturnType<TMachine, TSelectorsOutput, TActionsOutput, TSlots> & XstateTreeMachineInjection<TMachine, TSelectorsOutput, TActionsOutput, TSlots>;
 
 // Warning: (ae-internal-missing-underscore) The name "XstateTreeMachineInjection" should be prefixed with an underscore because the declaration is marked as @internal
 //
@@ -400,8 +401,8 @@ export type XstateTreeMachineStateSchemaV2<TMachine extends AnyStateMachine, TSe
 // src/routing/createRoute/createRoute.ts:285:19 - (ae-forgotten-export) The symbol "MergeRouteTypes" needs to be exported by the entry point index.d.ts
 // src/routing/createRoute/createRoute.ts:285:19 - (ae-forgotten-export) The symbol "ResolveZodType" needs to be exported by the entry point index.d.ts
 // src/routing/createRoute/createRoute.ts:322:9 - (ae-forgotten-export) The symbol "RouteRedirect" needs to be exported by the entry point index.d.ts
-// src/types.ts:142:3 - (ae-incompatible-release-tags) The symbol "canHandleEvent" is marked as @public, but its signature references "CanHandleEvent" which is marked as @internal
-// src/types.ts:143:3 - (ae-incompatible-release-tags) The symbol "inState" is marked as @public, but its signature references "MatchesFrom" which is marked as @internal
+// src/types.ts:164:3 - (ae-incompatible-release-tags) The symbol "canHandleEvent" is marked as @public, but its signature references "CanHandleEvent" which is marked as @internal
+// src/types.ts:165:3 - (ae-incompatible-release-tags) The symbol "inState" is marked as @public, but its signature references "MatchesFrom" which is marked as @internal
 
 // (No @packageDocumentation comment for this package)