From 4efe1ae166fffbaf7531549dac44b5d355aa1098 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=A9ry=20Debongnie?= <ged@odoo.com>
Date: Sat, 13 Nov 2021 08:34:43 +0100
Subject: [PATCH] [IMP] reactivity: slightly improve code and typing

---
 src/reactivity.ts        | 13 +++++--------
 tests/reactivity.test.ts |  2 +-
 2 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/src/reactivity.ts b/src/reactivity.ts
index 31edbe81..7b5b1336 100644
--- a/src/reactivity.ts
+++ b/src/reactivity.ts
@@ -13,9 +13,9 @@ const OBSERVER = Symbol("observer");
 const KEYS = Symbol("keys");
 const ROOT = Symbol("root");
 
-export function atom(source: any, observer: Observer) {
+export function atom<T>(source: T, observer: Observer): T {
   if (isTrackable(source) && observerSourceAtom.has(observer)) {
-    source = source[SOURCE] || source;
+    source = (source as any)[SOURCE] || source;
     const oldAtom = observerSourceAtom.get(observer)!.get(source);
     if (oldAtom) {
       return oldAtom;
@@ -33,8 +33,7 @@ export function atom(source: any, observer: Observer) {
 
 function createAtom(source: Source, observer: Observer): Atom {
   const keys: Set<any> = new Set();
-  let self: Atom;
-  const newAtom: Atom = new Proxy(source as any, {
+  return new Proxy(source as any, {
     set(target: any, key: string, value: any): boolean {
       if (!(key in target)) {
         target[key] = value;
@@ -65,7 +64,7 @@ function createAtom(source: Source, observer: Observer): Atom {
       }
       return true;
     },
-    get(target: any, key: any): any {
+    get(target: any, key: any, proxy: any): any {
       switch (key) {
         case OBSERVER:
           return observer;
@@ -81,7 +80,7 @@ function createAtom(source: Source, observer: Observer): Atom {
             if (!atoms.has(key)) {
               atoms.set(key, new ObserverSet());
             }
-            atoms.get(key)!.add(self);
+            atoms.get(key)!.add(proxy);
             keys.add(key);
           }
           //
@@ -89,8 +88,6 @@ function createAtom(source: Source, observer: Observer): Atom {
       }
     },
   });
-  self = newAtom;
-  return newAtom;
 }
 
 function isTrackable(value: any): boolean {
diff --git a/tests/reactivity.test.ts b/tests/reactivity.test.ts
index 4871fe8e..24b69e73 100644
--- a/tests/reactivity.test.ts
+++ b/tests/reactivity.test.ts
@@ -871,7 +871,7 @@ describe("Reactivity: atom", () => {
     let n = 0;
     const observer = () => n++;
     const unregisterObserver = registerObserver(observer);
-    const state = atom({ a: 1 }, observer);
+    const state = atom({ a: 1 } as any, observer);
 
     state.a = state.a;
     await nextMicroTick();