feat: snapshot/revert blocks

MimmyJau · Nov 16, 2023 · c345267 · c345267
1 parent 0edf9d9
commit c345267
Show file tree

Hide file tree

Showing 9 changed files with 220 additions and 65 deletions.
diff --git a/src/app.tsx b/src/app.tsx
@@ -35,6 +35,7 @@ import {
 } from '~/zustand'
 
 import { type AppMeta, AppMetaContext } from './contexts'
+import { useSnapshot } from './hooks/useSnapshot'
 import Layout from './screens/_layout'
 import AccountDetails from './screens/account-details'
 import BlockConfig from './screens/block-config'
@@ -229,7 +230,8 @@ function NetworkChangedEmitter() {
 
 /** Keeps block number in sync. */
 function SyncBlockNumber() {
-  usePendingBlock()
+  const { data: block } = usePendingBlock()
+  useSnapshot({ blockNumber: block?.number })
   return null
 }
 

diff --git a/src/components/Header.tsx b/src/components/Header.tsx
@@ -1,6 +1,6 @@
 import { type ReactNode, useCallback, useMemo } from 'react'
 import { Link } from 'react-router-dom'
-import { formatGwei } from 'viem'
+import { type Hex, formatGwei } from 'viem'
 
 import { Tooltip } from '~/components'
 import { BrandIcon } from '~/components/svgs'
@@ -27,6 +27,8 @@ import { usePendingBlock } from '~/hooks/usePendingBlock'
 import { getMessenger } from '~/messengers'
 import { useAccountStore, useNetworkStore, useSessionsStore } from '~/zustand'
 
+import { useRevert } from '../hooks/useRevert'
+import { useSnapshot } from '../hooks/useSnapshot'
 import * as styles from './Header.css'
 
 const contentMessenger = getMessenger('wallet:contentScript')
@@ -352,19 +354,24 @@ function Block() {
 
 function BlockNumber() {
   const { data: block } = usePendingBlock()
+  const { data: snapshot } = useSnapshot({
+    blockNumber: block?.number ? block?.number - 1n : undefined,
+    enabled: false,
+  })
   return (
     <Box position="relative">
-      <Inline wrap={false}>
-        <HeaderItem label="Block">
-          <Text size="12px" tabular>
-            {block?.number ? block?.number.toString() : ''}
-          </Text>
-        </HeaderItem>
-        {block && (
-          <Inset horizontal="2px">
-            <MineButton />
-          </Inset>
-        )}
+      <Inline gap="4px" wrap={false}>
+        <Box width="fit">
+          <HeaderItem label="Block">
+            <Text size="12px" tabular>
+              {block?.number ? block?.number.toString() : ''}
+            </Text>
+          </HeaderItem>
+        </Box>
+        <Inline wrap={false}>
+          {block && <MineButton />}
+          <RevertButton snapshot={snapshot} />
+        </Inline>
       </Inline>
     </Box>
   )
@@ -414,11 +421,7 @@ function MineButton() {
   const { mutateAsync: mine } = useMine()
 
   return (
-    <Box
-      key={block?.number?.toString()}
-      position="absolute"
-      style={{ marginTop: '8px' }}
-    >
+    <Box key={block?.number?.toString()} style={{ marginTop: '8px' }}>
       <Button.Symbol
         label="Mine Block"
         height="20px"
@@ -433,3 +436,23 @@ function MineButton() {
     </Box>
   )
 }
+
+function RevertButton({ snapshot }: { snapshot?: Hex }) {
+  const { mutateAsync: revert } = useRevert()
+
+  return (
+    <Box key={snapshot} style={{ marginTop: '8px' }}>
+      <Button.Symbol
+        disabled={!snapshot}
+        label="Revert Block"
+        height="20px"
+        onClick={(e) => {
+          e.preventDefault()
+          revert({ id: snapshot! })
+        }}
+        symbol="backward.fill"
+        variant="ghost primary"
+      />
+    </Box>
+  )
+}
diff --git a/src/design-system/symbols/generated/index.ts b/src/design-system/symbols/generated/index.ts
diff --git a/src/design-system/tokens.ts b/src/design-system/tokens.ts
@@ -655,6 +655,8 @@ export const symbolNames = [
   'plus',
   'gear',
   'minus',
+  'backward.fill',
+  'arrow.counterclockwise',
 ] as const
 export type SymbolName = typeof symbolNames[number]
 

diff --git a/src/hooks/usePendingBlock.ts b/src/hooks/usePendingBlock.ts
@@ -3,6 +3,7 @@ import { type Block, type Client, type Transaction } from 'viem'
 
 import {
   createQueryKey,
+  filterInfiniteQueryData,
   queryClient,
   updateInfiniteQueryData,
 } from '~/react-query'
@@ -42,55 +43,76 @@ export function usePendingBlockQueryOptions({
       const block = await client.getBlock({
         blockTag: 'pending',
       })
-      const prevBlock = queryClient.getQueryData([
-        'pending-block',
-        client.key,
-      ]) as Block
 
-      if (
-        block &&
-        prevBlock &&
-        prevBlock.number === block.number &&
-        prevBlock.transactions.length === block.transactions.length
-      )
-        return prevBlock || null
+      try {
+        const prevBlock = queryClient.getQueryData([
+          'pending-block',
+          client.key,
+        ]) as Block
 
-      queryClient.invalidateQueries({
-        queryKey: getAccountTokensQueryKey([client.key]),
-      })
-      queryClient.invalidateQueries({
-        queryKey: getBalanceQueryKey([client.key]),
-      })
-      queryClient.invalidateQueries({
-        queryKey: getContractsQueryKey([client.key]),
-      })
-      queryClient.invalidateQueries({
-        queryKey: getErc20BalanceQueryKey([client.key]),
-      })
-      queryClient.invalidateQueries({
-        queryKey: getBlockQueryKey([client.key, 'latest']),
-      })
-      queryClient.invalidateQueries({
-        queryKey: getNonceQueryKey([client.key]),
-      })
-      queryClient.invalidateQueries({
-        queryKey: getPendingTransactionsQueryKey([client.key]),
-      })
-      queryClient.invalidateQueries({
-        queryKey: getTxpoolQueryKey([client.key]),
-      })
+        if (
+          block &&
+          prevBlock &&
+          prevBlock.number === block.number &&
+          prevBlock.transactions.length === block.transactions.length
+        )
+          return prevBlock || null
 
-      const latestBlock = await client.getBlock({
-        includeTransactions: true,
-      })
-      queryClient.setQueryData<InfiniteData<Block[]>>(
-        getInfiniteBlocksQueryKey([client.key]),
-        updateInfiniteQueryData<Block[]>([latestBlock]),
-      )
-      queryClient.setQueryData<InfiniteData<Transaction[]>>(
-        getInfiniteBlockTransactionsQueryKey([client.key]),
-        updateInfiniteQueryData<Transaction[]>([...latestBlock.transactions]),
-      )
+        queryClient.invalidateQueries({
+          queryKey: getAccountTokensQueryKey([client.key]),
+        })
+        queryClient.invalidateQueries({
+          queryKey: getBalanceQueryKey([client.key]),
+        })
+        queryClient.invalidateQueries({
+          queryKey: getContractsQueryKey([client.key]),
+        })
+        queryClient.invalidateQueries({
+          queryKey: getErc20BalanceQueryKey([client.key]),
+        })
+        queryClient.invalidateQueries({
+          queryKey: getBlockQueryKey([client.key, 'latest']),
+        })
+        queryClient.invalidateQueries({
+          queryKey: getNonceQueryKey([client.key]),
+        })
+        queryClient.invalidateQueries({
+          queryKey: getPendingTransactionsQueryKey([client.key]),
+        })
+        queryClient.invalidateQueries({
+          queryKey: getTxpoolQueryKey([client.key]),
+        })
+
+        const latestBlock = await client.getBlock({
+          includeTransactions: true,
+        })
+
+        if (prevBlock.number && latestBlock.number < prevBlock.number) {
+          queryClient.setQueryData<InfiniteData<Block[]>>(
+            getInfiniteBlocksQueryKey([client.key]),
+            filterInfiniteQueryData<Block[]>(
+              (block) => block.number! <= latestBlock.number,
+            ),
+          )
+          queryClient.setQueryData<InfiniteData<Transaction[]>>(
+            getInfiniteBlockTransactionsQueryKey([client.key]),
+            filterInfiniteQueryData<Transaction[]>(
+              (tx) => !tx.blockNumber || tx.blockNumber! <= latestBlock.number,
+            ),
+          )
+        } else {
+          queryClient.setQueryData<InfiniteData<Block[]>>(
+            getInfiniteBlocksQueryKey([client.key]),
+            updateInfiniteQueryData<Block[]>([latestBlock]),
+          )
+          queryClient.setQueryData<InfiniteData<Transaction[]>>(
+            getInfiniteBlockTransactionsQueryKey([client.key]),
+            updateInfiniteQueryData<Transaction[]>([
+              ...latestBlock.transactions,
+            ]),
+          )
+        }
+      } catch {}
 
       return block || null
     },

diff --git a/src/hooks/useRevert.ts b/src/hooks/useRevert.ts
@@ -0,0 +1,21 @@
+import { useMutation } from '@tanstack/react-query'
+import type { RevertParameters } from 'viem'
+
+import { queryClient } from '~/react-query'
+
+import { useClient } from './useClient'
+import { usePendingBlockQueryOptions } from './usePendingBlock'
+
+export function useRevert() {
+  const client = useClient()
+  const pendingBlockQueryOptions = usePendingBlockQueryOptions()
+
+  return useMutation({
+    mutationFn: async ({ id }: RevertParameters) => {
+      await client.revert({
+        id,
+      })
+      queryClient.invalidateQueries(pendingBlockQueryOptions)
+    },
+  })
+}
diff --git a/src/hooks/useSnapshot.ts b/src/hooks/useSnapshot.ts
@@ -0,0 +1,37 @@
+import { queryOptions, useQuery } from '@tanstack/react-query'
+import type { Client } from 'viem'
+
+import { createQueryKey } from '~/react-query'
+
+import { useClient } from './useClient'
+
+type UseSnapshotParameters = {
+  blockNumber?: bigint | null
+  enabled?: boolean
+}
+
+export const getSnapshotQueryKey = createQueryKey<
+  'snapshot',
+  [key: Client['key'], blockNumber: string]
+>('snapshot')
+
+export function useSnapshotQueryOptions({
+  blockNumber,
+  enabled = true,
+}: UseSnapshotParameters) {
+  const client = useClient()
+  return queryOptions({
+    gcTime: Infinity,
+    staleTime: 0,
+    enabled: Boolean(enabled && blockNumber),
+    queryKey: getSnapshotQueryKey([client.key, (blockNumber || '').toString()]),
+    async queryFn() {
+      return (await client.snapshot()) || null
+    },
+  })
+}
+
+export function useSnapshot({ blockNumber, enabled }: UseSnapshotParameters) {
+  const queryOptions = useSnapshotQueryOptions({ blockNumber, enabled })
+  return useQuery(queryOptions)
+}
diff --git a/src/react-query.tsx b/src/react-query.tsx
@@ -68,6 +68,21 @@ export function createQueryKey<
     [id, ...(deps ? deps : [])] as unknown as [key, ...deps]
 }
 
+export function filterInfiniteQueryData<data extends unknown[]>(
+  predicate: (page: data[number]) => boolean,
+) {
+  return (prev: InfiniteData<data> | undefined) => {
+    const pages = prev?.pages
+      .map((page) => page.filter(predicate))
+      .filter((page) => page.length > 0)! as data[]
+    return {
+      ...prev,
+      pageParams: Array.from({ length: pages!.length }, (_, i) => i),
+      pages,
+    }
+  }
+}
+
 export function updateInfiniteQueryData<data extends unknown[]>(data: data) {
   return (prev: InfiniteData<data> | undefined) => {
     if (!prev) return
@@ -78,6 +93,7 @@ export function updateInfiniteQueryData<data extends unknown[]>(data: data) {
         pageParams: [...prev.pageParams, prev.pageParams.length],
         pages: [data, first, ...rest],
       }
+
     return {
       ...prev,
       pages: [[...data, ...first], ...rest],

diff --git a/src/screens/index.tsx b/src/screens/index.tsx
@@ -5,6 +5,7 @@ import { Link, useSearchParams } from 'react-router-dom'
 import { toast } from 'sonner'
 import {
   type Address,
+  type Block,
   type Hex,
   type Transaction,
   formatEther,
@@ -60,6 +61,8 @@ import {
 } from '~/zustand'
 import type { Account } from '~/zustand/account'
 
+import { useRevert } from '../hooks/useRevert'
+import { useSnapshot } from '../hooks/useSnapshot'
 import OnboardingStart from './onboarding/start'
 
 export default function Index() {
@@ -519,6 +522,13 @@ function Blocks() {
                           {block.transactions.length || '0'}
                         </Text>
                       </Column>
+                      <Column width="content">
+                        <Box style={{ width: '20px' }}>
+                          {status !== 'pending' && (
+                            <RevertButton block={block} />
+                          )}
+                        </Box>
+                      </Column>
                     </Columns>
                   </Box>
                 </VirtualList.Link>
@@ -531,6 +541,28 @@ function Blocks() {
   )
 }
 
+function RevertButton({ block }: { block?: Block }) {
+  const { data: snapshot } = useSnapshot({
+    blockNumber: block?.number,
+    enabled: false,
+  })
+  const { mutateAsync: revert } = useRevert()
+
+  if (!snapshot || !block?.timestamp) return null
+  return (
+    <Button.Symbol
+      label="Revert Block"
+      height="20px"
+      onClick={(e) => {
+        e.preventDefault()
+        revert({ id: snapshot! })
+      }}
+      symbol="arrow.counterclockwise"
+      variant="ghost primary"
+    />
+  )
+}
+
 ////////////////////////////////////////////////////////////////////////
 // Transactions