Refactor useObservable to handle deeply nested values

app/routes/hook.tsx

@@ -12,10 +12,8 @@ export const meta: MetaFunction = () => {
 export default function HookPage() {
   const {
     count,
-    submodel: { title, allCaps},
+    submodel: { title, allCaps, lowercase, shift },
   } = useObservable(store);
-
-  console.log("Rendering")
 
   return (
     <div className="App bg-white p-8 min-h-screen">
@@ -25,7 +23,8 @@ export default function HookPage() {
         {count}: {title}
       </p>
       <p className="text-gray-800 my-2">This is a computed view that capitalizes the title: {allCaps}</p>
-      {/* <p>This is a lazily evaluated view that converts the title to lowercase: {lowercase()}</p> */}
+      <p className="text-gray-800 my-2">This is a lazily evaluated view that converts the title to lowercase: {lowercase()}</p>
+      <p className="text-gray-800 my-2">This is a lazily evaluated view that shifts the title by a variable amount (set to 2): {shift(2)}</p>
       <div className="flex flex-wrap gap-2">
         <button 
           onClick={() => store.increment()} 

app/routes/observer.tsx

@@ -12,7 +12,7 @@ export const meta: MetaFunction = () => {
 export default observer(function ObserverPage() {
   const {
     count,
-    submodel: { title, allCaps },
+    submodel: { title, allCaps, lowercase, shift },
   } = store;
 
   console.log("Rendering")
@@ -25,6 +25,8 @@ export default observer(function ObserverPage() {
         {count}: {title}
       </p>
       <p className="text-gray-800 my-2">This is a computed view that capitalizes the title: {allCaps}</p>
+      <p className="text-gray-800 my-2">This is a lazily evaluated view that converts the title to lowercase: {lowercase()}</p>
+      <p className="text-gray-800 my-2">This is a lazily evaluated view that shifts the title by a variable amount (set to 2): {shift(2)}</p>
       <div className="flex flex-wrap gap-2">
         <button 
           onClick={() => store.increment()} 

app/routes/state.tsx

@@ -1,5 +1,5 @@
 import type { MetaFunction } from "@remix-run/node";
-import { useState, useMemo } from "react";
+import { useState, useMemo, useCallback } from "react";
 
 export const meta: MetaFunction = () => {
   return [
@@ -17,6 +17,12 @@ export default function StatePage() {
   // Equivalent to MST computed view
   const allCaps = useMemo(() => title.toUpperCase(), [title]);
 
+  // Equivalent to MST lazily evaluated view
+  const lowercase = useCallback(() => title.toLowerCase(), [title]);
+
+  // Equivalent to MST lazily evaluated view
+  const shift = useCallback((spaces: number) => title.slice(spaces), [title]);
+
   console.log("Rendering");
 
   const increment = () => {
@@ -50,6 +56,8 @@ export default function StatePage() {
         {count}: {title}
       </p>
       <p className="text-gray-800 my-2">This is a memoized value that capitalizes the title: {allCaps}</p>
+      <p className="text-gray-800 my-2">This is a lazily evaluated view that converts the title to lowercase: {lowercase()}</p>
+      <p className="text-gray-800 my-2">This is a lazily evaluated view that shifts the title by a variable amount (set to 2): {shift(2)}</p>
       <div className="flex flex-wrap gap-2">
         <button 
           onClick={increment} 

hooks/useObservable.ts

@@ -1,117 +1,111 @@
-import { useEffect, useState } from "react";
-import { IStateTreeNode, getSnapshot } from "mobx-state-tree";
 import { reaction } from "mobx";
-
-function isMSTNode(value: unknown): value is IStateTreeNode {
-  return Boolean(value && typeof value === "object" && "toJSON" in value);
+import {
+  IAnyStateTreeNode,
+  IStateTreeNode,
+  getMembers,
+  isStateTreeNode,
+} from "mobx-state-tree";
+import { useCallback, useLayoutEffect, useRef, useState } from "react";
+
+const emptyObject = {};
+
+type ViewArgs<T> = T extends (...args: infer A) => unknown ? A : never;
+
+class Path<T extends IAnyStateTreeNode> {
+  parts: (string | symbol)[];
+  args?: ViewArgs<T>;
+
+  constructor(parts: (string | symbol)[], args?: ViewArgs<T>) {
+    this.parts = parts;
+    this.args = args;
+  }
+
+  toString() {
+    return this.parts.join(".");
+  }
 }
 
-export function useObservable<T extends IStateTreeNode>(model: T) {
-  type SnapshotType = ReturnType<typeof getSnapshot<T>>;
-
-  const [state, setState] = useState<Record<string | symbol, unknown>>({});
-
-  useEffect(() => {
-    if (!model) return;
-
-    const disposers: ReturnType<typeof reaction>[] = [];
-
-    // Set up reaction for each property
-    Object.keys(state).forEach(prop => {
-      const disposer = reaction(
-        () => {
-          const value = model[prop as keyof T];
-          if (isMSTNode(value)) {
-            // Handle nested properties
-            const nested = {} as Record<string, unknown>;
-            Object.keys(state)
-              .filter(key => key.startsWith(`${prop}.`))
-              .forEach(key => {
-                const nestedProp = key.split('.')[1];
-                nested[nestedProp] = value[nestedProp as keyof typeof value];
-              });
-            return { value, ...nested };
-          }
-          return value;
-        },
-        (newValue) => {
-          setState(prev => ({
-            ...prev,
-            // @ts-expect-error - need to figure out some of this weird typing
-            [prop]: isMSTNode(newValue) ? newValue.value : newValue
-          }));
-        },
-        { fireImmediately: true }
-      );
-      disposers.push(disposer);
-    });
-
-    return () => disposers.forEach(d => d());
-  }, [model, Object.keys(state).join(',')]);
-
-  return new Proxy({} as SnapshotType, {
-    get: (_, property: string | symbol) => {
-      // Initialize tracking for new properties
-      if (!(property in state)) {
-        const value = model[property as keyof T];
-
-        setState(prev => ({
-          ...prev,
-          [property]: value
-        }));
-
-        if (typeof value === 'function') {
-          return value.bind(model);
+export function useObservable<T extends IStateTreeNode>(
+  model: T,
+  debug?: boolean
+): T {
+  const [, setState] = useState(1);
+  const forceRender = useCallback(() => setState((n) => -n), []);
+
+  const pathsAccessedInRender = useRef<Set<Path<T>>>(new Set());
+  pathsAccessedInRender.current = new Set();
+
+  useLayoutEffect(() => {
+    const disposer = reaction(
+      function expression() {
+        if (debug) {
+          console.log("Paths accessed in render:");
+          pathsAccessedInRender.current.forEach((path) => {
+            console.log("  ", path.toString());
+            if (path.args) {
+              console.log("    This is a view fn, with args:", path.args);
+            }
+          });
         }
 
-        if (isMSTNode(value)) {
-          return new Proxy({} as SnapshotType, {
-            get: (_, nestedProp: string | symbol) => {
-              const nestedValue = value[nestedProp as keyof typeof value];
-
-              if (typeof nestedValue === 'function') {
-                return nestedValue.bind(value);
-              }
-
-              // Track nested property
-              const fullProp = `${String(property)}.${String(nestedProp)}`;
-              if (!(fullProp in state)) {
-                setState(prev => ({
-                  ...prev,
-                  [fullProp]: nestedValue
-                }));
-              }
-
-              return nestedValue;
+        for (const path of pathsAccessedInRender.current) {
+          let current: IAnyStateTreeNode = model;
+          for (const part of path.parts) {
+            current = current[part as keyof typeof current];
+
+            if (typeof current === "function" && path.args) {
+              current = current(...path.args);
             }
-          });
+          }
         }
 
-        return value;
+        return []; // If anything changes we want to re-render
+      },
+      function effect() {
+        forceRender();
+      },
+      {
+        fireImmediately: false,
       }
+    );
 
-      const value = model[property as keyof T];
+    return () => {
+      disposer();
+    };
+  }, [forceRender, model, debug]);
 
-      if (typeof value === 'function') {
-        return value.bind(model);
-      }
+  function makeProxy(instance: IAnyStateTreeNode, path: (string | symbol)[]) {
+    return new Proxy(emptyObject, {
+      get(_, key) {
+        const value = instance[key as keyof typeof instance];
 
-      if (isMSTNode(value)) {
-        return new Proxy({} as SnapshotType, {
-          get: (_, nestedProp: string | symbol) => {
-            const nestedValue = value[nestedProp as keyof typeof value];
-
-            if (typeof nestedValue === 'function') {
-              return nestedValue.bind(value);
-            }
+        const newPath = new Path([...path, key as string]);
+        pathsAccessedInRender.current.add(newPath);
 
-            const fullProp = `${String(property)}.${String(nestedProp)}`;
-            return state[fullProp] ?? nestedValue;
+        if (isStateTreeNode(value)) {
+          return makeProxy(value, [...path, key]);
+        }
+
+        if (typeof value === "function") {
+          const members = getMembers(instance);
+          if (members.views.includes(key as string)) {
+            const view = value.bind(instance);
+            const argumentWatcher = (...args: ViewArgs<typeof view>) => {
+              const newPath = new Path([...path, key as string], args);
+              pathsAccessedInRender.current.add(newPath);
+              return view(...args);
+            };
+
+            return argumentWatcher;
+          } else {
+            return value.bind(instance);
           }
-        });
-      }
+        }
+
+        return value;
+      },
+    });
+  }
 
-      return state[property] ?? value;
-    }
-  });
+  return makeProxy(model, []) as T;
 }

store/store.ts

@@ -19,6 +19,9 @@ const SubModel = t
     },
     lowercase() {
       return self.title.toLowerCase();
+    },
+    shift(spaces: number) {
+      return self.title.slice(spaces);
     }
   }))