From a26bbb3f1f897b19e6e5347a7fecaed332c0e59e Mon Sep 17 00:00:00 2001
From: Brandon Roberts <robertsbt@gmail.com>
Date: Mon, 23 Sep 2024 14:21:21 -0500
Subject: [PATCH] chore: review feedback fixes

---
 modules/signals/src/state-source.ts | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/modules/signals/src/state-source.ts b/modules/signals/src/state-source.ts
index 194d698488..4db6ec3a68 100644
--- a/modules/signals/src/state-source.ts
+++ b/modules/signals/src/state-source.ts
@@ -9,7 +9,7 @@ import {
 } from '@angular/core';
 import { Prettify } from './ts-helpers';
 
-const STATE_WATCHERS = new WeakMap<object, Array<StateWatcher<any>>>();
+const STATE_WATCHERS = new WeakMap<Signal<object>, Array<StateWatcher<any>>>();
 
 export const STATE_SOURCE = Symbol('STATE_SOURCE');
 
@@ -78,7 +78,7 @@ export function watchState<State extends object>(
 function getWatchers<State extends object>(
   stateSource: StateSource<State>
 ): Array<StateWatcher<State>> {
-  return STATE_WATCHERS.get(stateSource[STATE_SOURCE] as object) || [];
+  return STATE_WATCHERS.get(stateSource[STATE_SOURCE]) || [];
 }
 
 function notifyWatchers<State extends object>(
@@ -97,10 +97,7 @@ function addWatcher<State extends object>(
   watcher: StateWatcher<State>
 ): void {
   const watchers = getWatchers(stateSource);
-  STATE_WATCHERS.set(stateSource[STATE_SOURCE] as object, [
-    ...watchers,
-    watcher,
-  ]);
+  STATE_WATCHERS.set(stateSource[STATE_SOURCE], [...watchers, watcher]);
 }
 
 function removeWatcher<State extends object>(
@@ -109,7 +106,7 @@ function removeWatcher<State extends object>(
 ): void {
   const watchers = getWatchers(stateSource);
   STATE_WATCHERS.set(
-    stateSource[STATE_SOURCE] as object,
+    stateSource[STATE_SOURCE],
     watchers.filter((w) => w !== watcher)
   );
 }