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

Author: UberMouse

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", () => {

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;
 }
 
 /**

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,