From 14af8170374d68f742e287c77a5d4a3b8567adaf Mon Sep 17 00:00:00 2001
From: R Midhun Suresh <hi@midhun.dev>
Date: Tue, 12 Aug 2025 13:33:59 +0530
Subject: [PATCH] Add hook to auto dispose viewmodels

---
 .../base/useAutoDisposedViewModel.ts          | 64 +++++++++++++++++++
 .../base/useAutoDisposedViewModel-test.ts     | 47 ++++++++++++++
 2 files changed, 111 insertions(+)
 create mode 100644 src/viewmodels/base/useAutoDisposedViewModel.ts
 create mode 100644 test/viewmodels/base/useAutoDisposedViewModel-test.ts

diff --git a/src/viewmodels/base/useAutoDisposedViewModel.ts b/src/viewmodels/base/useAutoDisposedViewModel.ts
new file mode 100644
index 0000000000..64300fcd73
--- /dev/null
+++ b/src/viewmodels/base/useAutoDisposedViewModel.ts
@@ -0,0 +1,64 @@
+/*
+Copyright 2024 New Vector Ltd.
+Copyright 2022 The Matrix.org Foundation C.I.C.
+
+SPDX-License-Identifier: AGPL-3.0-only OR GPL-3.0-only OR LicenseRef-Element-Commercial
+Please see LICENSE files in the repository root for full details.
+*/
+
+import { useEffect, useState } from "react";
+
+import type { BaseViewModel } from "./BaseViewModel";
+
+type VmCreator<B extends BaseViewModel<unknown, unknown>> = () => B;
+
+/**
+ * Instantiate a view-model that gets disposed when the calling react component unmounts.
+ * In other words, this hook ties the lifecycle of a view-model to the lifecycle of a
+ * react component.
+ *
+ * @param vmCreator A function that returns a view-model instance
+ * @returns view-model instance from vmCreator
+ * @example
+ * const vm = useAutoDisposedViewModel(() => new FooViewModel({prop1, prop2, ...});
+ */
+export function useAutoDisposedViewModel<B extends BaseViewModel<unknown, unknown>>(vmCreator: VmCreator<B>): B {
+    /**
+     * The view-model instance may be replaced by a different instance in some scenarios.
+     * We want to be sure that whatever react component called this hook gets re-rendered
+     * when this happens, hence the state.
+     */
+    const [viewModel, setViewModel] = useState<B>(vmCreator);
+
+    /**
+     * Our intention here is to ensure that the dispose method of the view-model gets called
+     * when the component that uses this hook unmounts.
+     * We can do that by combining a useEffect cleanup with an empty dependency array.
+     */
+    useEffect(() => {
+        let toDispose = viewModel;
+
+        /**
+         * Because we use react strict mode, react will run our effects twice in dev mode to make
+         * sure that they are pure.
+         * This presents a complication - the vm instance that we created in our state initializer
+         * will get disposed on the first cleanup.
+         * So we'll recreate the view-model if it's already disposed.
+         */
+        if (viewModel.isDisposed) {
+            const newViewModel = vmCreator();
+            // Change toDispose so that we don't end up disposing the already disposed vm.
+            toDispose = newViewModel;
+            setViewModel(newViewModel);
+        }
+        return () => {
+            // Dispose the view-model when this component unmounts
+            toDispose.dispose();
+        };
+
+        // eslint-disable-next-line react-compiler/react-compiler
+        // eslint-disable-next-line react-hooks/exhaustive-deps
+    }, []);
+
+    return viewModel;
+}
diff --git a/test/viewmodels/base/useAutoDisposedViewModel-test.ts b/test/viewmodels/base/useAutoDisposedViewModel-test.ts
new file mode 100644
index 0000000000..d77b3a9c37
--- /dev/null
+++ b/test/viewmodels/base/useAutoDisposedViewModel-test.ts
@@ -0,0 +1,47 @@
+/*
+Copyright 2025 New Vector Ltd.
+SPDX-License-Identifier: AGPL-3.0-only OR GPL-3.0-only OR LicenseRef-Element-Commercial
+Please see LICENSE files in the repository root for full details.
+*/
+
+import { renderHook } from "jest-matrix-react";
+
+import { BaseViewModel } from "../../../src/viewmodels/base/BaseViewModel";
+import { useAutoDisposedViewModel } from "../../../src/viewmodels/base/useAutoDisposedViewModel";
+
+class TestViewModel extends BaseViewModel<{ count: number }, { initial: number }> {
+    constructor(props: { initial: number }) {
+        super(props, { count: props.initial });
+    }
+
+    public increment() {
+        const newCount = this.getSnapshot().count + 1;
+        this.snapshot.set({ count: newCount });
+    }
+}
+
+describe("useAutoDisposedViewModel", () => {
+    it("should return view-model", () => {
+        const vmCreator = () => new TestViewModel({ initial: 0 });
+        const { result } = renderHook(() => useAutoDisposedViewModel(vmCreator));
+        const vm = result.current;
+        expect(vm).toBeInstanceOf(TestViewModel);
+        expect(vm.isDisposed).toStrictEqual(false);
+    });
+
+    it("should dispose view-model on unmount", () => {
+        const vmCreator = () => new TestViewModel({ initial: 0 });
+        const { result, unmount } = renderHook(() => useAutoDisposedViewModel(vmCreator));
+        const vm = result.current;
+        vm.increment();
+        unmount();
+        expect(vm.isDisposed).toStrictEqual(true);
+    });
+
+    it("should recreate view-model on react strict mode", async () => {
+        const vmCreator = () => new TestViewModel({ initial: 0 });
+        const output = renderHook(() => useAutoDisposedViewModel(vmCreator), { reactStrictMode: true });
+        const vm = output.result.current;
+        expect(vm.isDisposed).toStrictEqual(false);
+    });
+});