From cacc2d311e04bf6cd0e32aefed24fca5caa00269 Mon Sep 17 00:00:00 2001
From: EYHN <cneyhn@gmail.com>
Date: Wed, 13 Mar 2024 06:06:18 +0000
Subject: [PATCH] feat(infra): livedata flatten (#6083)

---
 .../src/livedata/__tests__/livedata.spec.ts   | 38 +++++++++++++++
 packages/common/infra/src/livedata/index.ts   | 48 ++++++++++++++++++-
 2 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/packages/common/infra/src/livedata/__tests__/livedata.spec.ts b/packages/common/infra/src/livedata/__tests__/livedata.spec.ts
index a48545a9a2..4eccf553de 100644
--- a/packages/common/infra/src/livedata/__tests__/livedata.spec.ts
+++ b/packages/common/infra/src/livedata/__tests__/livedata.spec.ts
@@ -185,4 +185,42 @@ describe('livedata', () => {
     });
     expect(value).toBe(1);
   });
+
+  test('flat', () => {
+    {
+      const wrapped = new LiveData(new LiveData(0));
+      const flatten = wrapped.flat();
+      expect(flatten.value).toBe(0);
+
+      wrapped.next(new LiveData(1));
+      expect(flatten.value).toBe(1);
+
+      wrapped.next(LiveData.from(of(2, 3), 0));
+      expect(flatten.value).toBe(3);
+    }
+
+    {
+      const wrapped = new LiveData(
+        new LiveData([
+          new LiveData(new LiveData(1)),
+          new LiveData(new LiveData(2)),
+        ])
+      );
+      const flatten = wrapped.flat();
+      expect(flatten.value).toStrictEqual([1, 2]);
+    }
+
+    {
+      const wrapped = new LiveData([new LiveData(0), new LiveData(1)]);
+      const flatten = wrapped.flat();
+
+      expect(flatten.value).toEqual([0, 1]);
+
+      const inner = new LiveData(2);
+      wrapped.next([inner, new LiveData(3)]);
+      expect(flatten.value).toEqual([2, 3]);
+      inner.next(4);
+      expect(flatten.value).toEqual([4, 3]);
+    }
+  });
 });
diff --git a/packages/common/infra/src/livedata/index.ts b/packages/common/infra/src/livedata/index.ts
index cc9259a9b1..53d737d235 100644
--- a/packages/common/infra/src/livedata/index.ts
+++ b/packages/common/infra/src/livedata/index.ts
@@ -1,5 +1,6 @@
 import { DebugLogger } from '@affine/debug';
 import {
+  combineLatest,
   distinctUntilChanged,
   EMPTY,
   filter,
@@ -192,7 +193,7 @@ export class LiveData<T = unknown> implements InteropObservable<T> {
     return subscription;
   }
 
-  map<R>(mapper: (v: T) => R): LiveData<R> {
+  map<R>(mapper: (v: T) => R) {
     const sub = LiveData.from(
       new Observable<R>(subscriber =>
         this.subscribe({
@@ -268,6 +269,42 @@ export class LiveData<T = unknown> implements InteropObservable<T> {
     this.upstreamSubscription?.unsubscribe();
   }
 
+  /**
+   * flatten the livedata
+   *
+   * ```
+   * new LiveData(new LiveData(0)).flat() // LiveData<number>
+   * ```
+   *
+   * ```
+   * new LiveData([new LiveData(0)]).flat() // LiveData<number[]>
+   * ```
+   */
+  flat(): Flat<this> {
+    return LiveData.from(
+      this.pipe(
+        switchMap(v => {
+          if (v instanceof LiveData) {
+            return (v as LiveData<any>).flat();
+          } else if (Array.isArray(v)) {
+            return combineLatest(
+              v.map(v => {
+                if (v instanceof LiveData) {
+                  return v.flat();
+                } else {
+                  return of(v);
+                }
+              })
+            );
+          } else {
+            return of(v);
+          }
+        })
+      ),
+      null as any
+    ) as any;
+  }
+
   reactSubscribe = (cb: () => void) => {
     this.ops.next('watch');
     const subscription = this.raw
@@ -297,3 +334,12 @@ export class LiveData<T = unknown> implements InteropObservable<T> {
 }
 
 export type LiveDataOperation = 'set' | 'get' | 'watch' | 'unwatch';
+
+export type Unwrap<T> =
+  T extends LiveData<infer Z>
+    ? Unwrap<Z>
+    : T extends LiveData<infer A>[]
+      ? Unwrap<A>[]
+      : T;
+
+export type Flat<T> = T extends LiveData<infer P> ? LiveData<Unwrap<P>> : T;