From e9d2cca81baacdd1d8532b96cbae9825802b0c38 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/types.ts          | 24 +++++++++++++++++++++++-
 3 files changed, 66 insertions(+), 2 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/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,