From 8d3bc361fc96555bbd8006b653f049959afaa921 Mon Sep 17 00:00:00 2001
From: DittoETH <136378658+ditto-eth@users.noreply.github.com>
Date: Thu, 28 Sep 2023 15:30:00 -0700
Subject: [PATCH] feat: auto-import erc20 tokens after fork

---
 src/hooks/useGetLogs.ts              | 49 ++++++++++++++++++++++++++++
 src/screens/account-details.tsx      | 42 ++++++++++++++++++++++++
 src/screens/onboarding/configure.tsx |  1 +
 src/zustand/network.ts               | 13 ++++++++
 4 files changed, 105 insertions(+)
 create mode 100644 src/hooks/useGetLogs.ts

diff --git a/src/hooks/useGetLogs.ts b/src/hooks/useGetLogs.ts
new file mode 100644
index 0000000..5c380ec
--- /dev/null
+++ b/src/hooks/useGetLogs.ts
@@ -0,0 +1,49 @@
+import { queryOptions, useQuery } from '@tanstack/react-query'
+import {
+  type BlockNumber,
+  type BlockTag,
+  type Client,
+  type GetLogsParameters,
+  stringify,
+} from 'viem'
+
+import { createQueryKey } from '~/react-query'
+
+import type { AbiEvent } from 'abitype'
+import { useClient } from './useClient'
+
+type UseLogsParameters<
+  TAbiEvent extends AbiEvent,
+  TFromBlock extends BlockNumber | BlockTag,
+  TToBlock extends BlockNumber | BlockTag,
+> = GetLogsParameters<TAbiEvent, TAbiEvent[], undefined, TFromBlock, TToBlock>
+
+export const getLogsQueryKey = createQueryKey<
+  'get-logs',
+  [key: Client['key'], args: string]
+>('get-logs')
+
+export function useGetLogsQueryOptions<
+  TAbiEvent extends AbiEvent,
+  TFromBlock extends BlockNumber | BlockTag,
+  TToBlock extends BlockNumber | BlockTag,
+>(parameters: UseLogsParameters<TAbiEvent, TFromBlock, TToBlock>) {
+  const client = useClient()
+
+  return queryOptions({
+    enabled: Boolean(parameters),
+    queryKey: getLogsQueryKey([client.key, stringify(parameters)]),
+    async queryFn() {
+      return await client.getLogs(parameters)
+    },
+  })
+}
+
+export function useGetLogs<
+  TAbiEvent extends AbiEvent,
+  TFromBlock extends BlockNumber | BlockTag,
+  TToBlock extends BlockNumber | BlockTag,
+>(parameters: UseLogsParameters<TAbiEvent, TFromBlock, TToBlock>) {
+  const queryOptions = useGetLogsQueryOptions(parameters)
+  return useQuery(queryOptions)
+}
diff --git a/src/screens/account-details.tsx b/src/screens/account-details.tsx
index e50ad7c..2c61c0d 100644
--- a/src/screens/account-details.tsx
+++ b/src/screens/account-details.tsx
@@ -8,6 +8,7 @@ import {
   BaseError,
   formatUnits,
   isAddress,
+  parseAbiItem,
   parseUnits,
 } from 'viem'
 
@@ -31,7 +32,9 @@ import {
 } from '~/design-system'
 import { useErc20Balance } from '~/hooks/useErc20Balance'
 import { useErc20Metadata } from '~/hooks/useErc20Metadata'
+import { useGetLogs } from '~/hooks/useGetLogs'
 import { useSetErc20Balance } from '~/hooks/useSetErc20Balance'
+import { useNetworkStore } from '~/zustand'
 import { useTokensStore } from '~/zustand/tokens'
 
 export default function AccountDetails() {
@@ -88,8 +91,47 @@ export default function AccountDetails() {
   )
 }
 
+function useImportTransferredTokens() {
+  const { address: accountAddress } = useParams()
+  const { addToken } = useTokensStore()
+  const { network } = useNetworkStore()
+
+  const { data: transfersFrom } = useGetLogs({
+    event: parseAbiItem(
+      'event Transfer(address indexed from, address indexed to, uint256)',
+    ),
+    args: {
+      from: accountAddress as Address,
+    },
+    fromBlock: network.forkBlockNumber,
+    toBlock: 'latest',
+  })
+  const { data: transfersTo } = useGetLogs({
+    event: parseAbiItem(
+      'event Transfer(address indexed from, address indexed to, uint256)',
+    ),
+    args: {
+      to: accountAddress as Address,
+    },
+    fromBlock: network.forkBlockNumber,
+    toBlock: 'latest',
+  })
+
+  if (accountAddress) {
+    ;[
+      ...new Set([
+        ...(transfersFrom?.map((t) => t.address) || []),
+        ...(transfersTo?.map((t) => t.address) || []),
+      ]),
+    ].forEach((addr) => {
+      addToken(addr, accountAddress as Address)
+    })
+  }
+}
+
 function Tokens({ accountAddress }: { accountAddress: Address }) {
   const { tokens } = useTokensStore()
+  useImportTransferredTokens()
 
   if (!accountAddress) return null
   return (
diff --git a/src/screens/onboarding/configure.tsx b/src/screens/onboarding/configure.tsx
index bc0148a..cff1aec 100644
--- a/src/screens/onboarding/configure.tsx
+++ b/src/screens/onboarding/configure.tsx
@@ -73,6 +73,7 @@ export default function OnboardingConfigure() {
         rpcUrl: defaultNetwork.rpcUrl,
         network: {
           blockTime: Number(values.blockTime),
+          forkBlockNumber: BigInt(values.forkBlockNumber),
           chainId: Number(values.chainId),
           name: values.networkName,
           rpcUrl,
diff --git a/src/zustand/network.ts b/src/zustand/network.ts
index d8d759c..0317337 100644
--- a/src/zustand/network.ts
+++ b/src/zustand/network.ts
@@ -9,6 +9,7 @@ import { createStore } from './utils'
 type RpcUrl = string
 export type Network = {
   blockTime: number
+  forkBlockNumber: bigint
   chainId: number
   name: string
   rpcUrl: RpcUrl
@@ -35,6 +36,7 @@ export type NetworkStore = NetworkState & NetworkActions
 
 export const defaultNetwork = {
   blockTime: 0,
+  forkBlockNumber: 0n,
   chainId: -1,
   name: '',
   rpcUrl: '',
@@ -88,6 +90,16 @@ export const networkStore = createStore<NetworkStore>(
           return defaultNetwork.chainId
         }
       })()
+      const forkBlockNumber = await (async () => {
+        if (network_.forkBlockNumber) return network_.forkBlockNumber
+        try {
+          return await getClient({
+            rpcUrl,
+          }).getBlockNumber()
+        } catch {
+          return defaultNetwork.forkBlockNumber
+        }
+      })()
       const name = (() => {
         if (network_.name) return network_.name
         const chain = Object.values(chains).find(
@@ -106,6 +118,7 @@ export const networkStore = createStore<NetworkStore>(
           ...(index >= 0 ? networks[index] : defaultNetwork),
           ...network_,
           chainId,
+          forkBlockNumber,
           name,
         }
         if (index >= 0) networks[index] = network