From 8764efebf01ca767e1f4bb627663d96808d2665e Mon Sep 17 00:00:00 2001
From: Leon Senft <leonsenft@users.noreply.github.com>
Date: Wed, 15 May 2024 14:11:35 -0700
Subject: [PATCH] fix(core): treat reads/exceptions in `equal` as part of
 computation

Prevent leaking signal reads and exceptions from a custom `equal`
function of a producer `computed()` to a consumer.

Upstream https://github.com/tc39/proposal-signals/pull/90.
---
 .../core/primitives/signals/src/computed.ts   | 13 +--
 packages/core/test/signals/computed_spec.ts   | 90 +++++++++++++++++++
 2 files changed, 97 insertions(+), 6 deletions(-)

diff --git a/packages/core/primitives/signals/src/computed.ts b/packages/core/primitives/signals/src/computed.ts
index 94fa26e5ed437..5544ffa5d2750 100644
--- a/packages/core/primitives/signals/src/computed.ts
+++ b/packages/core/primitives/signals/src/computed.ts
@@ -119,8 +119,14 @@ const COMPUTED_NODE = /* @__PURE__ */ (() => {
 
       const prevConsumer = consumerBeforeComputation(node);
       let newValue: unknown;
+      let wasEqual = false;
       try {
         newValue = node.computation();
+        wasEqual =
+          oldValue !== UNSET &&
+          oldValue !== ERRORED &&
+          newValue !== ERRORED &&
+          node.equal(oldValue, newValue);
       } catch (err) {
         newValue = ERRORED;
         node.error = err;
@@ -128,12 +134,7 @@ const COMPUTED_NODE = /* @__PURE__ */ (() => {
         consumerAfterComputation(node, prevConsumer);
       }
 
-      if (
-        oldValue !== UNSET &&
-        oldValue !== ERRORED &&
-        newValue !== ERRORED &&
-        node.equal(oldValue, newValue)
-      ) {
+      if (wasEqual) {
         // No change to `valueVersion` - old and new values are
         // semantically equivalent.
         node.value = oldValue;
diff --git a/packages/core/test/signals/computed_spec.ts b/packages/core/test/signals/computed_spec.ts
index 700dee7bf39e2..a687c5de02c05 100644
--- a/packages/core/test/signals/computed_spec.ts
+++ b/packages/core/test/signals/computed_spec.ts
@@ -193,4 +193,94 @@ describe('computed', () => {
     const double = computed(() => counter() * 2);
     expect(double + '').toBe('[Computed: 2]');
   });
+
+  describe('with custom equal', () => {
+    it('should cache exceptions thrown by equal', () => {
+      const s = signal(0);
+
+      let computedRunCount = 0;
+      let equalRunCount = 0;
+      const c = computed(
+        () => {
+          computedRunCount++;
+          return s();
+        },
+        {
+          equal: () => {
+            equalRunCount++;
+            throw new Error('equal');
+          },
+        },
+      );
+
+      // equal() isn't run for the initial computation.
+      expect(c()).toBe(0);
+      expect(computedRunCount).toBe(1);
+      expect(equalRunCount).toBe(0);
+
+      s.set(1);
+
+      // Error is thrown by equal().
+      expect(() => c()).toThrowError('equal');
+      expect(computedRunCount).toBe(2);
+      expect(equalRunCount).toBe(1);
+
+      // Error is cached; c throws again without needing to rerun computation or equal().
+      expect(() => c()).toThrowError('equal');
+      expect(computedRunCount).toBe(2);
+      expect(equalRunCount).toBe(1);
+    });
+
+    it('does not leak tracking information', () => {
+      const exact = signal(1);
+      const epsilon = signal(0.1);
+      const counter = signal(0);
+
+      let innerRunCount = 0;
+      let equalRunCount = 0;
+      const inner = computed(
+        () => {
+          innerRunCount++;
+          return exact();
+        },
+        {
+          equal: (a, b) => {
+            equalRunCount++;
+            return Math.abs(a - b) < epsilon();
+          },
+        },
+      );
+
+      let outerRunCount = 0;
+      const outer = computed(() => {
+        outerRunCount++;
+        counter();
+        return inner();
+      });
+
+      // Everything runs the first time.
+      expect(outer()).toBe(1);
+      expect(innerRunCount).toBe(1);
+      expect(outerRunCount).toBe(1);
+
+      exact.set(2);
+      counter.set(2);
+
+      // `outer` reruns because `counter` changed, `inner` reruns when called by `outer`, and
+      // `cutoff` is called for the first time.
+      expect(outer()).toBe(2);
+      expect(innerRunCount).toBe(2);
+      expect(outerRunCount).toBe(2);
+      expect(equalRunCount).toBe(1);
+
+      epsilon.set(0.2);
+
+      // Changing something the custom `equal` depends on makes the `inner` need to rerun, but not
+      // `outer` (since the new and old values are equal).
+      expect(outer()).toBe(2);
+      expect(innerRunCount).toBe(3);
+      expect(outerRunCount).toBe(2);
+      expect(equalRunCount).toBe(2);
+    });
+  });
 });