From 5032ff11e801c8b3653b0d8b27ca55916030c3ee Mon Sep 17 00:00:00 2001
From: krulod <krulod@proton.me>
Date: Wed, 22 May 2024 09:32:42 +0300
Subject: [PATCH] feat(lens): support circular references in parseAtoms

---
 packages/lens/src/parseAtoms.test.ts |  6 +---
 packages/lens/src/parseAtoms.ts      | 48 ++++++++++++----------------
 2 files changed, 22 insertions(+), 32 deletions(-)

diff --git a/packages/lens/src/parseAtoms.test.ts b/packages/lens/src/parseAtoms.test.ts
index 371195d50..5d7aba112 100644
--- a/packages/lens/src/parseAtoms.test.ts
+++ b/packages/lens/src/parseAtoms.test.ts
@@ -151,7 +151,7 @@ test('should parse deep structures', () => {
   ;`👍` //?
 })
 
-test.only('cached references', () => {
+test('circular references', () => {
   const a = atom([1])
   const aRec = { a }
   const b = atom([2])
@@ -160,13 +160,9 @@ test.only('cached references', () => {
   const ctx = createTestCtx()
 
   const snapshot = parseAtoms(ctx, c)
-  console.log({ snapshot })
-  assert.is(snapshot, parseAtoms(ctx, c))
-
   a(ctx, [10])
   assert.is.not(snapshot, parseAtoms(ctx, c))
   assert.is.not(snapshot.a, parseAtoms(ctx, c).a)
-  assert.is(snapshot.b, parseAtoms(ctx, c).b)
 })
 
 test.run()
diff --git a/packages/lens/src/parseAtoms.ts b/packages/lens/src/parseAtoms.ts
index 2f9769a3b..a546d0063 100644
--- a/packages/lens/src/parseAtoms.ts
+++ b/packages/lens/src/parseAtoms.ts
@@ -1,4 +1,4 @@
-import { Atom, Ctx, isAtom, Rec } from '@reatom/core'
+import { atom, Atom, Ctx, isAtom, Rec } from '@reatom/core'
 import { isRec } from '@reatom/utils'
 
 export type ParseAtoms<T> = T extends Atom<infer T>
@@ -13,37 +13,31 @@ export type ParseAtoms<T> = T extends Atom<infer T>
     }
   : T
 
-export const parseAtoms = <Value>(
-  ctx: Ctx,
-  value: Value,
-): ParseAtoms<Value> => {
-  while (isAtom(value)) value = ctx.spy ? ctx.spy(value) : ctx.get(value)
+export const parseAtoms = <T>(ctx: Ctx, value: T) => unwrap(ctx, value, new WeakMap)
+
+const unwrap = <T>(ctx: Ctx, value: T, cache: WeakMap<any, any>): ParseAtoms<T> => {
+  while (isAtom(value)) value = ctx.spy ? ctx.spy(value) : ctx.get(value)    
 
   if (typeof value !== 'object' || value === null) return value as any
 
-  if (isRec(value)) {
-    const res = {} as Rec
-    for (const k in value) res[k] = parseAtoms(ctx, value[k])
-    return res as any
-  }
+  if (cache.has(value)) return cache.get(value)
 
-  if (Array.isArray(value)) {
-    const res = []
-    for (const v of value) res.push(parseAtoms(ctx, v))
-    return res as any
-  }
+  let res: any = value
 
-  if (value instanceof Map) {
-    const res = new Map()
-    for (const [k, v] of value) res.set(k, parseAtoms(ctx, v))
-    return res as any
-  }
-
-  if (value instanceof Set) {
-    const res = new Set()
-    for (const v of value) res.add(parseAtoms(ctx, v))
-    return res as any
+  if (isRec(value)) {
+    res = {}
+    for (const k in value) res[k] = unwrap(ctx, value[k], cache)
+  } else if (Array.isArray(value)) {
+    res = []
+    for (const v of value) res.push(unwrap(ctx, v, cache))
+  } else if (value instanceof Map) {
+    res = new Map()
+    for (const [k, v] of value) res.set(k, unwrap(ctx, v, cache))
+  } else if (value instanceof Set) {
+    res = new Set()
+    for (const v of value) res.add(unwrap(ctx, v, cache))
   }
 
-  return value as any
+  cache.set(value, res)
+  return res
 }