Refactor abi loader to handle multiple abi in decoding

Author: Ferossgp

packages/transaction-decoder/src/abi-loader.ts

@@ -1,9 +1,10 @@
-import { Effect, Either, RequestResolver, Request, Array, pipe, Data, PrimaryKey, Schema, SchemaAST } from 'effect'
+import { Effect, Either, RequestResolver, Request, Array, Data, PrimaryKey, Schema, SchemaAST } from 'effect'
 import { ContractABI } from './abi-strategy/request-model.js'
 import { Abi } from 'viem'
 import * as AbiStore from './abi-store.js'
 import * as StrategyExecutorModule from './abi-strategy/strategy-executor.js'
 import { SAFE_MULTISEND_SIGNATURE, SAFE_MULTISEND_ABI, AA_ABIS } from './decoding/constants.js'
+import { errorFunctionSignatures, solidityError, solidityPanic } from './helpers/error.js'
 
 interface LoadParameters {
   readonly chainID: number
@@ -34,7 +35,7 @@ export class EmptyCalldataError extends Data.TaggedError('DecodeError')<
 class SchemaAbi extends Schema.make<Abi>(SchemaAST.objectKeyword) {}
 class AbiLoader extends Schema.TaggedRequest<AbiLoader>()('AbiLoader', {
   failure: Schema.instanceOf(MissingABIError),
-  success: SchemaAbi, // Abi
+  success: Schema.Array(Schema.Struct({ abi: SchemaAbi, id: Schema.optional(Schema.String) })),
   payload: {
     chainID: Schema.Number,
     address: Schema.String,
@@ -78,20 +79,19 @@ const setValue = (key: AbiLoader, abi: ContractABI | null) =>
         event: key.event,
         signature: key.signature,
       },
-      abi == null ? { status: 'not-found', result: null } : { status: 'success', result: abi },
+      abi == null
+        ? {
+            type: 'func' as const,
+            abi: '',
+            address: key.address,
+            chainID: key.chainID,
+            signature: key.signature || '',
+            status: 'not-found' as const,
+          }
+        : { ...abi, status: 'success' as const },
     )
   })
 
-const getBestMatch = (abi: ContractABI | null) => {
-  if (abi == null) return null
-
-  if (abi.type === 'address') {
-    return JSON.parse(abi.abi) as Abi
-  }
-
-  return JSON.parse(`[${abi.abi}]`) as Abi
-}
-
 /**
  * Data loader for contracts abi
  *
@@ -146,27 +146,41 @@ export const AbiLoaderRequestResolver = RequestResolver.makeBatched((requests: A
     const requestGroups = Array.groupBy(requests, makeRequestKey)
     const uniqueRequests = Object.values(requestGroups).map((group) => group[0])
 
-    const [remaining, cachedResults] = yield* pipe(
-      getMany(uniqueRequests),
-      Effect.map(
-        Array.partitionMap((resp, i) => {
-          return resp.status === 'empty'
-            ? Either.left(uniqueRequests[i])
-            : Either.right([uniqueRequests[i], resp.result] as const)
-        }),
-      ),
-      Effect.orElseSucceed(() => [uniqueRequests, []] as const),
-    )
+    const allCachedData = yield* getMany(uniqueRequests)
+
+    // Create a map of invalid sources for each request
+    const invalidSourcesMap = new Map<string, string[]>()
+    allCachedData.forEach((abis, i) => {
+      const request = uniqueRequests[i]
+      const invalid = abis
+        .filter((abi) => abi.status === 'invalid')
+        .map((abi) => abi.source)
+        .filter(Boolean) as string[]
+
+      invalidSourcesMap.set(makeRequestKey(request), invalid)
+    })
+
+    const [remaining, cachedResults] = Array.partitionMap(allCachedData, (abis, i) => {
+      // Filter out invalid/not-found ABIs and check if we have any valid ones
+      const validAbis = abis.filter((abi) => abi.status === 'success')
+      return validAbis.length === 0
+        ? Either.left(uniqueRequests[i])
+        : Either.right([uniqueRequests[i], validAbis] as const)
+    })
 
     //  Resolve ABI from the store
     yield* Effect.forEach(
       cachedResults,
-      ([request, abi]) => {
+      ([request, abis]) => {
         const group = requestGroups[makeRequestKey(request)]
-        const bestMatch = getBestMatch(abi)
-        const result = bestMatch ? Effect.succeed(bestMatch) : Effect.fail(new MissingABIError(request))
-
-        return Effect.forEach(group, (req) => Request.completeEffect(req, result), { discard: true })
+        const allMatches = abis.map((abi) => {
+          const parsedAbi = abi.type === 'address' ? (JSON.parse(abi.abi) as Abi) : (JSON.parse(`[${abi.abi}]`) as Abi)
+          return { abi: parsedAbi, id: abi.id }
+        })
+
+        return Effect.forEach(group, (req) => Request.completeEffect(req, Effect.succeed(allMatches)), {
+          discard: true,
+        })
       },
       {
         discard: true,
@@ -191,9 +205,14 @@ export const AbiLoaderRequestResolver = RequestResolver.makeBatched((requests: A
     const response = yield* Effect.forEach(
       remaining,
       (req) => {
-        const allAvailableStrategies = Array.prependAll(strategies.default, strategies[req.chainID] ?? []).filter(
-          (strategy) => strategy.type === 'address',
-        )
+        const invalidSources = invalidSourcesMap.get(makeRequestKey(req)) ?? []
+        const allAvailableStrategies = Array.prependAll(strategies.default, strategies[req.chainID] ?? [])
+          .filter((strategy) => strategy.type === 'address')
+          .filter((strategy) => !invalidSources.includes(strategy.id))
+
+        if (allAvailableStrategies.length === 0) {
+          return Effect.succeed(Either.right(req))
+        }
 
         return strategyExecutor
           .executeStrategiesSequentially(allAvailableStrategies, {
@@ -220,17 +239,22 @@ export const AbiLoaderRequestResolver = RequestResolver.makeBatched((requests: A
     // NOTE: Secondly we request strategies to fetch fragments
     const fragmentStrategyResults = yield* Effect.forEach(
       notFound,
-      ({ chainID, address, event, signature }) => {
-        const allAvailableStrategies = Array.prependAll(strategies.default, strategies[chainID] ?? []).filter(
-          (strategy) => strategy.type === 'fragment',
-        )
+      (req) => {
+        const invalidSources = invalidSourcesMap.get(makeRequestKey(req)) ?? []
+        const allAvailableStrategies = Array.prependAll(strategies.default, strategies[req.chainID] ?? [])
+          .filter((strategy) => strategy.type === 'fragment')
+          .filter((strategy) => !invalidSources.includes(strategy.id))
+
+        if (allAvailableStrategies.length === 0) {
+          return Effect.succeed(null)
+        }
 
         return strategyExecutor
           .executeStrategiesSequentially(allAvailableStrategies, {
-            address,
-            chainId: chainID,
-            event,
-            signature,
+            address: req.address,
+            chainId: req.chainID,
+            event: req.event,
+            signature: req.signature,
             strategyId: 'fragment-batch',
           })
           .pipe(
@@ -244,20 +268,44 @@ export const AbiLoaderRequestResolver = RequestResolver.makeBatched((requests: A
       },
     )
 
-    const strategyResults = Array.appendAll(addressStrategyResults, fragmentStrategyResults)
+    // Create a map to track which requests got results from address strategies
+    const addressResultsMap = new Map<string, ContractABI[]>()
+    addressStrategyResults.forEach((abis, i) => {
+      const request = remaining[i]
+      if (abis && abis.length > 0) {
+        addressResultsMap.set(makeRequestKey(request), abis)
+      }
+    })
+
+    // Create a map to track which requests got results from fragment strategies
+    const fragmentResultsMap = new Map<string, ContractABI[]>()
+    fragmentStrategyResults.forEach((abis, i) => {
+      const request = notFound[i]
+      if (abis && abis.length > 0) {
+        fragmentResultsMap.set(makeRequestKey(request), abis)
+      }
+    })
 
-    // Store results and resolve pending requests
+    // Resolve all remaining requests
     yield* Effect.forEach(
-      strategyResults,
-      (abis, i) => {
-        const request = remaining[i]
-        const abi = abis?.[0] ?? null
-        const bestMatch = getBestMatch(abi)
-        const result = bestMatch ? Effect.succeed(bestMatch) : Effect.fail(new MissingABIError(request))
+      remaining,
+      (request) => {
         const group = requestGroups[makeRequestKey(request)]
+        const addressAbis = addressResultsMap.get(makeRequestKey(request)) || []
+        const fragmentAbis = fragmentResultsMap.get(makeRequestKey(request)) || []
+        const allAbis = [...addressAbis, ...fragmentAbis]
+        const firstAbi = allAbis[0] || null
+
+        const allMatches = allAbis.map((abi) => {
+          const parsedAbi = abi.type === 'address' ? (JSON.parse(abi.abi) as Abi) : (JSON.parse(`[${abi.abi}]`) as Abi)
+          // TODO: We should figure out how to handle the db ID here, maybe we need to start providing the ids before inserting
+          return { abi: parsedAbi, id: undefined }
+        })
+
+        const result = allMatches.length > 0 ? Effect.succeed(allMatches) : Effect.fail(new MissingABIError(request))
 
         return Effect.zipRight(
-          setValue(request, abi),
+          setValue(request, firstAbi),
           Effect.forEach(group, (req) => Request.completeEffect(req, result), { discard: true }),
         )
       },
@@ -283,12 +331,17 @@ export const getAndCacheAbi = (params: AbiStore.AbiParams) =>
       return yield* Effect.fail(new EmptyCalldataError(params))
     }
 
+    if (params.signature && errorFunctionSignatures.includes(params.signature)) {
+      const errorAbis: Abi = [...solidityPanic, ...solidityError]
+      return [{ abi: errorAbis, id: undefined }]
+    }
+
     if (params.signature && params.signature === SAFE_MULTISEND_SIGNATURE) {
-      return yield* Effect.succeed(SAFE_MULTISEND_ABI)
+      return [{ abi: SAFE_MULTISEND_ABI, id: undefined }]
     }
 
     if (params.signature && AA_ABIS[params.signature]) {
-      return yield* Effect.succeed(AA_ABIS[params.signature])
+      return [{ abi: AA_ABIS[params.signature], id: undefined }]
     }
 
     return yield* Effect.request(new AbiLoader(params), AbiLoaderRequestResolver)

packages/transaction-decoder/src/decoding/abi-decode.ts

@@ -1,8 +1,18 @@
 import { formatAbiItem } from 'viem/utils'
 import type { DecodeResult, MostTypes, TreeNode } from '../types.js'
-import { Hex, Abi, decodeFunctionData, AbiParameter, AbiFunction, getAbiItem } from 'viem'
+import {
+  Hex,
+  Abi,
+  decodeFunctionData,
+  decodeEventLog as viemDecodeEventLog,
+  AbiParameter,
+  AbiFunction,
+  getAbiItem,
+} from 'viem'
 import { Data, Effect } from 'effect'
 import { messageFromUnknown } from '../helpers/error.js'
+import * as AbiStore from '../abi-store.js'
+import { getAndCacheAbi } from '../abi-loader.js'
 
 export class DecodeError extends Data.TaggedError('DecodeError')<{ message: string }> {
   constructor(message: string, error?: unknown) {
@@ -93,3 +103,102 @@ export const decodeMethod = (data: Hex, abi: Abi): Effect.Effect<DecodeResult |
       }
     }
   })
+
+/**
+ * Validates and decodes data using multiple ABIs with fallback support.
+ * Tries each ABI in sequence and marks failed ABIs as invalid in the store.
+ * Uses Effect.firstSuccessOf to return the first successful decode result.
+ */
+export const validateAndDecodeWithABIs = (
+  data: Hex,
+  params: AbiStore.AbiParams,
+): Effect.Effect<DecodeResult, DecodeError, AbiStore.AbiStore> =>
+  Effect.gen(function* () {
+    const { updateStatus } = yield* AbiStore.AbiStore
+
+    // TODO: When abi is returned from external source, it does not have ids.
+    // We could do a new db selct, change API of Store to return ids, or provide
+    // ids instead of DB auto-generated ids.
+    // Now it will fail only on the second call
+    const abiWithIds = yield* getAndCacheAbi(params)
+
+    // Create validation effects for store ABIs
+    const storeValidationEffects = abiWithIds.map(({ abi, id }) =>
+      Effect.gen(function* () {
+        const result = yield* decodeMethod(data, abi)
+
+        if (result == null) {
+          return yield* Effect.fail(new DecodeError(`ABI ${abi} failed to decode`))
+        }
+
+        return result
+      }).pipe(
+        Effect.catchAll((error: DecodeError) => {
+          return Effect.gen(function* () {
+            if (updateStatus && id != null) {
+              // Mark this ABI as invalid when it fails
+              yield* updateStatus(id, 'invalid').pipe(Effect.catchAll(() => Effect.void))
+            }
+            return yield* Effect.fail(error)
+          })
+        }),
+      ),
+    )
+
+    return yield* Effect.firstSuccessOf(storeValidationEffects)
+  })
+
+/**
+ * Validates and decodes event logs using multiple ABIs with fallback support.
+ * Tries each ABI in sequence and returns the first successful decode result along with the ABI used.
+ * Similar to validateAndDecodeWithABIs but specifically for event logs.
+ */
+export const validateAndDecodeEventWithABIs = (
+  topics: readonly Hex[],
+  data: Hex,
+  params: AbiStore.AbiParams,
+): Effect.Effect<{ eventName: string; args: any; abiItem: Abi }, DecodeError, AbiStore.AbiStore> =>
+  Effect.gen(function* () {
+    const abiWithIds = yield* getAndCacheAbi(params)
+
+    const validationEffects = abiWithIds.map(({ abi, id }) =>
+      Effect.gen(function* () {
+        const result = yield* decodeEventLog(topics, data, abi as Abi)
+
+        if (result == null) {
+          return yield* Effect.fail(new DecodeError(`ABI failed to decode event`))
+        }
+
+        return { ...result, abiItem: abi as Abi }
+      }).pipe(
+        Effect.catchAll((error: DecodeError) => {
+          return Effect.gen(function* () {
+            // Note: We don't mark ABIs as invalid for event decoding failures
+            // as the same ABI might work for other events on the same contract
+            return yield* Effect.fail(error)
+          })
+        }),
+      ),
+    )
+
+    return yield* Effect.firstSuccessOf(validationEffects)
+  })
+
+export const decodeEventLog = (
+  topics: readonly Hex[],
+  data: Hex,
+  abi: Abi,
+): Effect.Effect<{ eventName: string; args: any } | undefined, DecodeError> =>
+  Effect.gen(function* () {
+    const { eventName, args = {} } = yield* Effect.try({
+      try: () =>
+        viemDecodeEventLog({ abi, topics: topics as [] | [`0x${string}`, ...`0x${string}`[]], data, strict: false }),
+      catch: (error) => new DecodeError(`Could not decode event log`, error),
+    })
+
+    if (eventName == null) {
+      return undefined
+    }
+
+    return { eventName, args }
+  })

packages/transaction-decoder/src/decoding/calldata-decode.ts

@@ -1,7 +1,8 @@
 import { Effect } from 'effect'
 import { Hex, Address, encodeFunctionData, isAddress, getAddress } from 'viem'
-import { getAndCacheAbi, MissingABIError } from '../abi-loader.js'
+import { MissingABIError } from '../abi-loader.js'
 import * as AbiDecoder from './abi-decode.js'
+import { validateAndDecodeWithABIs } from './abi-decode.js'
 import { TreeNode } from '../types.js'
 import { PublicClient, RPCFetchError, UnknownNetwork } from '../public-client.js'
 import { SAFE_MULTISEND_SIGNATURE, SAFE_MULTISEND_NESTED_ABI } from './constants.js'
@@ -166,16 +167,11 @@ export const decodeMethod = ({
       }
     }
 
-    const abi = yield* getAndCacheAbi({
+    const decoded = yield* validateAndDecodeWithABIs(data, {
       address: implementationAddress ?? contractAddress,
       signature,
       chainID,
     })
-    const decoded = yield* AbiDecoder.decodeMethod(data, abi)
-
-    if (decoded == null) {
-      return yield* new AbiDecoder.DecodeError(`Failed to decode method: ${data}`)
-    }
 
     //MULTISEND: decode nested params for the multisend of the safe smart account
     if (