From dba471cf8bbeaed338a6b4af3aeb31d655461863 Mon Sep 17 00:00:00 2001
From: pedrobonamin <pedrobonamin@gmail.com>
Date: Tue, 1 Oct 2024 08:44:57 +0200
Subject: [PATCH] chore(core): update editState to use createSWR cache

---
 .../document/document-pair/editState.ts       | 16 ++++++++------
 .../store/_legacy/document/document-store.ts  |  4 ++--
 .../_legacy/document/documentsStorage.ts      | 21 -------------------
 3 files changed, 12 insertions(+), 29 deletions(-)
 delete mode 100644 packages/sanity/src/core/store/_legacy/document/documentsStorage.ts

diff --git a/packages/sanity/src/core/store/_legacy/document/document-pair/editState.ts b/packages/sanity/src/core/store/_legacy/document/document-pair/editState.ts
index 3affaf92b48..6565d4fac96 100644
--- a/packages/sanity/src/core/store/_legacy/document/document-pair/editState.ts
+++ b/packages/sanity/src/core/store/_legacy/document/document-pair/editState.ts
@@ -1,10 +1,10 @@
 import {type SanityClient} from '@sanity/client'
 import {type SanityDocument, type Schema} from '@sanity/types'
-import {combineLatest, defer, type Observable, of} from 'rxjs'
-import {map, publishReplay, refCount, startWith, switchMap, tap} from 'rxjs/operators'
+import {combineLatest, type Observable} from 'rxjs'
+import {map, publishReplay, refCount, startWith, switchMap} from 'rxjs/operators'
 
+import {createSWR} from '../../../../util/rxSwr'
 import {type PairListenerOptions} from '../getPairListener'
-import {type DocumentsStorage} from '../documentsStorage'
 import {type IdPair, type PendingMutationsEvent} from '../types'
 import {memoize} from '../utils/createMemoizer'
 import {memoizeKeyGen} from './memoizeKeyGen'
@@ -15,6 +15,8 @@ interface TransactionSyncLockState {
   enabled: boolean
 }
 
+const swr = createSWR<[SanityDocument, SanityDocument, TransactionSyncLockState]>({maxSize: 50})
+
 /**
  * @hidden
  * @beta */
@@ -43,6 +45,7 @@ export const editState = memoize(
     typeName: string,
   ): Observable<EditStateFor> => {
     const liveEdit = isLiveEditEnabled(ctx.schema, typeName)
+
     return snapshotPair(
       ctx.client,
       idPair,
@@ -60,14 +63,15 @@ export const editState = memoize(
           ),
         ]),
       ),
-      map(([draftSnapshot, publishedSnapshot, transactionSyncLock]) => ({
+      swr(`${idPair.publishedId}-${idPair.draftId}`),
+      map(({value: [draftSnapshot, publishedSnapshot, transactionSyncLock], fromCache}) => ({
         id: idPair.publishedId,
         type: typeName,
         draft: draftSnapshot,
         published: publishedSnapshot,
         liveEdit,
-        ready: true,
-        transactionSyncLock,
+        ready: !fromCache,
+        transactionSyncLock: fromCache ? null : transactionSyncLock,
       })),
       startWith({
         id: idPair.publishedId,
diff --git a/packages/sanity/src/core/store/_legacy/document/document-store.ts b/packages/sanity/src/core/store/_legacy/document/document-store.ts
index 977960afe9c..813962d3fdc 100644
--- a/packages/sanity/src/core/store/_legacy/document/document-store.ts
+++ b/packages/sanity/src/core/store/_legacy/document/document-store.ts
@@ -30,6 +30,7 @@ import {getInitialValueStream, type InitialValueMsg, type InitialValueOptions} f
 import {listenQuery, type ListenQueryOptions} from './listenQuery'
 import {resolveTypeForDocument} from './resolveTypeForDocument'
 import {type IdPair} from './types'
+
 /**
  * @hidden
  * @beta */
@@ -108,7 +109,7 @@ export function createDocumentStore({
   // internal operations, and a `getClient` method that we expose to user-land
   // for things like validations
   const client = getClient(DEFAULT_STUDIO_CLIENT_OPTIONS)
-  const storage = createDocumentsStorage()
+
   const ctx = {
     client,
     getClient,
@@ -118,7 +119,6 @@ export function createDocumentStore({
     i18n,
     serverActionsEnabled,
     pairListenerOptions,
-    storage,
   }
 
   return {
diff --git a/packages/sanity/src/core/store/_legacy/document/documentsStorage.ts b/packages/sanity/src/core/store/_legacy/document/documentsStorage.ts
deleted file mode 100644
index 5da6aa16c74..00000000000
--- a/packages/sanity/src/core/store/_legacy/document/documentsStorage.ts
+++ /dev/null
@@ -1,21 +0,0 @@
-import {type SanityDocument} from '@sanity/types'
-import QuickLRU from 'quick-lru'
-
-export interface DocumentsStorage {
-  save: (id: string, doc: SanityDocument) => void
-  get: (id: string) => SanityDocument | null
-}
-
-export function createDocumentsStorage(): DocumentsStorage {
-  const documentsCache = new QuickLRU<string, SanityDocument | null>({
-    maxSize: 50,
-  })
-  return {
-    save(id, doc) {
-      documentsCache.set(id, doc)
-    },
-    get(id) {
-      return documentsCache.get(id) || null
-    },
-  }
-}