From a1b88c95172fb4e06e5057ca85d13f0ff73348be Mon Sep 17 00:00:00 2001
From: Damian Tarnawski <gthetarnav@gmail.com>
Date: Fri, 31 Mar 2023 21:18:42 +0200
Subject: [PATCH] Improve lazy memo reads in another memo.

---
 .changeset/light-seals-impress.md |  5 +++++
 packages/memo/src/index.ts        | 20 +++++++++++++++-----
 packages/memo/test/lazy.test.ts   | 21 ++++++++++++++++++++-
 3 files changed, 40 insertions(+), 6 deletions(-)
 create mode 100644 .changeset/light-seals-impress.md

diff --git a/.changeset/light-seals-impress.md b/.changeset/light-seals-impress.md
new file mode 100644
index 000000000..ae5924e10
--- /dev/null
+++ b/.changeset/light-seals-impress.md
@@ -0,0 +1,5 @@
+---
+"@solid-primitives/memo": patch
+---
+
+Improve lazy memo reads in another memo.
diff --git a/packages/memo/src/index.ts b/packages/memo/src/index.ts
index e642b9b32..2643a5297 100644
--- a/packages/memo/src/index.ts
+++ b/packages/memo/src/index.ts
@@ -16,6 +16,7 @@ import {
   NoInfer,
   Owner,
   SignalOptions,
+  DEV,
 } from "solid-js";
 import { isServer } from "solid-js/web";
 import { debounce, throttle } from "@solid-primitives/scheduled";
@@ -338,6 +339,8 @@ export function createAsyncMemo<T>(
   return state;
 }
 
+const EQUALS_FALSE = { equals: false } as const;
+
 /**
  * Lazily evaluated `createMemo`. Will run the calculation only if is being listened to.
  *
@@ -379,16 +382,23 @@ export function createLazyMemo<T>(
     };
   }
 
-  let isStale = true;
+  let isReading = false,
+    isStale: boolean | undefined = true;
 
-  const [isDirty, setDirty] = createSignal({ v: false }),
+  const [track, trigger] = createSignal(void 0, EQUALS_FALSE),
     memo = createMemo<T>(
-      p => (isDirty().v ? ((isStale = isDirty().v = false), calc(p)) : ((isStale = true), p)),
+      p => (isReading ? calc(p) : ((isStale = !track()), p)),
       value as T,
-      { equals: false, name: options?.name },
+      DEV ? { name: options?.name, equals: false } : EQUALS_FALSE,
     );
 
-  return (): T => (isStale && setDirty({ v: true }), memo());
+  return (): T => {
+    isReading = true;
+    if (isStale) isStale = trigger();
+    const v = memo();
+    isReading = false;
+    return v;
+  };
 }
 
 /*
diff --git a/packages/memo/test/lazy.test.ts b/packages/memo/test/lazy.test.ts
index a01bae659..179480bee 100644
--- a/packages/memo/test/lazy.test.ts
+++ b/packages/memo/test/lazy.test.ts
@@ -1,6 +1,6 @@
 import { describe, it, expect } from "vitest";
 import { createLazyMemo } from "../src";
-import { createComputed, createEffect, createRoot, createSignal } from "solid-js";
+import { createComputed, createEffect, createMemo, createRoot, createSignal } from "solid-js";
 
 describe("createLazyMemo", () => {
   it("won't run if not accessed", () =>
@@ -256,4 +256,23 @@ describe("createLazyMemo", () => {
 
       dispose();
     }));
+
+  it("stays in sync when read in a memo", () => {
+    const { dispose, setA, memo } = createRoot(dispose => {
+      const [a, setA] = createSignal(1);
+      const b = createLazyMemo(() => a());
+
+      const memo = createMemo(() => {
+        expect(a()).toBe(b());
+      });
+
+      return { dispose, setA, memo };
+    });
+
+    memo();
+    setA(2);
+    memo();
+
+    dispose();
+  });
 });