Remove remaining usages of old post thread query (#9184)

* Remove remaining usages of old post thread query

* Add PostThreadContext, cache mutator for threadgates on threads, pipe it through

* Replace getPostThread in threadgate query

* Replace in initQuote handling, which isn't even used rn...

* Missing import

* Revert ext change

src/state/queries/usePostThread/context.tsx

@@ -0,0 +1,43 @@
+import {createContext, useContext} from 'react'
+
+import {
+  type createPostThreadOtherQueryKey,
+  type createPostThreadQueryKey,
+} from '#/state/queries/usePostThread/types'
+
+/**
+ * Contains static metadata about the post thread query, suitable for
+ * context e.g. query keys and other things that don't update frequently.
+ *
+ * Be careful adding things here, as it could cause unnecessary re-renders.
+ */
+export type PostThreadContextType = {
+  postThreadQueryKey: ReturnType<typeof createPostThreadQueryKey>
+  postThreadOtherQueryKey: ReturnType<typeof createPostThreadOtherQueryKey>
+}
+
+const PostThreadContext = createContext<PostThreadContextType | undefined>(
+  undefined,
+)
+
+/**
+ * Use the current {@link PostThreadContext}, if one is available. If not,
+ * returns `undefined`.
+ */
+export function usePostThreadContext() {
+  return useContext(PostThreadContext)
+}
+
+export function PostThreadContextProvider({
+  children,
+  context,
+}: {
+  children: React.ReactNode
+  context?: PostThreadContextType
+}) {
+  return (
+    <PostThreadContext.Provider value={context}>
+      {children}
+    </PostThreadContext.Provider>
+  )
+}

src/state/queries/usePostThread/index.ts

@@ -11,6 +11,7 @@ import {
   TREE_VIEW_BELOW_DESKTOP,
   TREE_VIEW_BF,
 } from '#/state/queries/usePostThread/const'
+import {type PostThreadContextType} from '#/state/queries/usePostThread/context'
 import {
   createCacheMutator,
   getThreadPlaceholder,
@@ -31,6 +32,8 @@ import {useAgent, useSession} from '#/state/session'
 import {useMergeThreadgateHiddenReplies} from '#/state/threadgate-hidden-replies'
 import {useBreakpoints} from '#/alf'
 
+export * from '#/state/queries/usePostThread/context'
+export {useUpdatePostThreadThreadgateQueryCache} from '#/state/queries/usePostThread/queryCache'
 export * from '#/state/queries/usePostThread/types'
 
 export function usePostThread({anchor}: {anchor?: string}) {
@@ -277,8 +280,13 @@ export function usePostThread({anchor}: {anchor?: string}) {
     setOtherItemsVisible,
   ])
 
-  return useMemo(
-    () => ({
+  return useMemo(() => {
+    const context: PostThreadContextType = {
+      postThreadQueryKey,
+      postThreadOtherQueryKey,
+    }
+    return {
+      context,
       state: {
         /*
          * Copy in any query state that is useful
@@ -309,17 +317,18 @@ export function usePostThread({anchor}: {anchor?: string}) {
         setSort,
         setView,
       },
-    }),
-    [
-      query,
-      mutator.insertReplies,
-      otherItemsVisible,
-      sort,
-      view,
-      setSort,
-      setView,
-      threadgate,
-      items,
-    ],
-  )
+    }
+  }, [
+    query,
+    mutator.insertReplies,
+    otherItemsVisible,
+    sort,
+    view,
+    setSort,
+    setView,
+    threadgate,
+    items,
+    postThreadQueryKey,
+    postThreadOtherQueryKey,
+  ])
 }

src/state/queries/usePostThread/queryCache.ts

@@ -1,3 +1,4 @@
+import {useCallback} from 'react'
 import {
   type $Typed,
   type AppBskyActorDefs,
@@ -7,7 +8,7 @@ import {
   type AppBskyUnspeccedGetPostThreadV2,
   AtUri,
 } from '@atproto/api'
-import {type QueryClient} from '@tanstack/react-query'
+import {type QueryClient, useQueryClient} from '@tanstack/react-query'
 
 import {
   dangerousGetPostShadow,
@@ -18,6 +19,7 @@ import {findAllPostsInQueryData as findAllPostsInNotifsQueryData} from '#/state/
 import {findAllPostsInQueryData as findAllPostsInFeedQueryData} from '#/state/queries/post-feed'
 import {findAllPostsInQueryData as findAllPostsInQuoteQueryData} from '#/state/queries/post-quotes'
 import {findAllPostsInQueryData as findAllPostsInSearchQueryData} from '#/state/queries/search-posts'
+import {usePostThreadContext} from '#/state/queries/usePostThread'
 import {getBranch} from '#/state/queries/usePostThread/traversal'
 import {
   type ApiThreadItem,
@@ -322,3 +324,51 @@ export function* findAllProfilesInQueryData(
     }
   }
 }
+
+export function useUpdatePostThreadThreadgateQueryCache() {
+  const qc = useQueryClient()
+  const context = usePostThreadContext()
+
+  return useCallback(
+    (threadgate: AppBskyFeedDefs.ThreadgateView) => {
+      if (!context) return
+
+      function mutator<T>(thread: ApiThreadItem[]): T[] {
+        for (let i = 0; i < thread.length; i++) {
+          const item = thread[i]
+
+          if (!AppBskyUnspeccedDefs.isThreadItemPost(item.value)) continue
+
+          if (item.depth === 0) {
+            thread.splice(i, 1, {
+              ...item,
+              value: {
+                ...item.value,
+                post: {
+                  ...item.value.post,
+                  threadgate,
+                },
+              },
+            })
+          }
+        }
+
+        return thread as T[]
+      }
+
+      qc.setQueryData<AppBskyUnspeccedGetPostThreadV2.OutputSchema>(
+        context.postThreadQueryKey,
+        data => {
+          if (!data) return
+          return {
+            ...data,
+            thread: mutator<AppBskyUnspeccedGetPostThreadV2.ThreadItem>([
+              ...data.thread,
+            ]),
+          }
+        },
+      )
+    },
+    [qc, context],
+  )
+}