feat: diff improvements

src/reports/code-diff.spec.ts

@@ -2,8 +2,15 @@ import {describe, expect, it} from 'vitest';
 import pre32 from './mocks/pre3-2.json';
 import post32 from './mocks/post3-2.json';
 import {diffCode, downloadContract} from './code-diff';
+import {diffSlot} from './raw-storage-diff';
 
 describe.skip('code diffs', () => {
+  it('should diff slots', () => {
+    diffSlot(1, '0x0', {
+      previousValue: '0x4816b2C2895f97fB918f1aE7Da403750a0eE372e',
+      newValue: '0xE5e48Ad1F9D1A894188b483DcF91f4FaD6AbA43b',
+    });
+  });
   it(
     'should download contract',
     () => {

src/reports/diff-reports.ts

@@ -6,7 +6,7 @@ import {renderReserve, renderReserveDiff} from './reserve';
 import {AaveV3Reserve, type AaveV3Snapshot} from './snapshot-types';
 import {renderStrategy, renderStrategyDiff} from './strategy';
 import {diffCode, downloadContract} from './code-diff';
-import {bytes32ToAddress} from '../utils/storageSlots';
+import {diffRawStorage} from './raw-storage-diff';
 
 function hasDiff(input: Record<string, any>): boolean {
   if (!input) return false;
@@ -127,21 +127,7 @@ export async function diffReports<A extends AaveV3Snapshot, B extends AaveV3Snap
   }
 
   if (raw) {
-    // ERC1967 slot https://github.com/OpenZeppelin/openzeppelin-contracts/blob/master/contracts/proxy/ERC1967/ERC1967Utils.sol#L21C53-L21C119
-    const erc1967Slot = '0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc';
-    Object.keys(raw).map((contract) => {
-      if (raw[contract].stateDiff[erc1967Slot] && raw[contract].stateDiff[erc1967Slot].previousValue != '0x0000000000000000000000000000000000000000000000000000000000000000') {
-        const fromAddress = bytes32ToAddress(raw[contract].stateDiff[erc1967Slot].previousValue);
-        const toAddress = bytes32ToAddress(raw[contract].stateDiff[erc1967Slot].newValue);
-        const from = downloadContract(pre.chainId, fromAddress);
-        const to = downloadContract(pre.chainId, toAddress);
-        const result = diffCode(from, to);
-        writeFileSync(
-          `./diffs/${pre.chainId}_${contract}_${fromAddress}_${toAddress}.diff`,
-          result,
-        );
-      }
-    });
+    diffRawStorage(pre.chainId, raw);
   }
 
   try {

src/reports/raw-storage-diff.ts

@@ -0,0 +1,86 @@
+import {writeFileSync, mkdirSync, readFileSync} from 'fs';
+import {bytes32ToAddress} from '../utils/storageSlots';
+import {diffCode, downloadContract} from './code-diff';
+import {RawStorage, SlotDiff} from './snapshot-types';
+import {isKnownAddress} from '../govv3/utils/checkAddress';
+import {Address, getContract, zeroHash} from 'viem';
+import {getClient} from '@bgd-labs/rpc-env';
+import {IPool_ABI} from '@bgd-labs/aave-address-book/abis';
+
+export function diffSlot(chainId: number, address: Address, slot: SlotDiff) {
+  // pure new deployments cannot be diffed, we just download the code in that case
+  if (slot.previousValue == zeroHash) {
+    const toAddress = bytes32ToAddress(slot.newValue);
+    const to = downloadContract(chainId, toAddress);
+    mkdirSync('./diffs', {recursive: true});
+    writeFileSync(`./diffs/${chainId}_${address}_${toAddress}.diff`, readFileSync(to), {});
+  } else {
+    const fromAddress = bytes32ToAddress(slot.previousValue);
+    const toAddress = bytes32ToAddress(slot.newValue);
+    const from = downloadContract(chainId, fromAddress);
+    const to = downloadContract(chainId, toAddress);
+    const result = diffCode(from, to);
+    mkdirSync('./diffs', {recursive: true});
+    writeFileSync(`./diffs/${chainId}_${address}_${fromAddress}_${toAddress}.diff`, result, {});
+  }
+}
+
+export async function diffRawStorage(chainId: number, raw: RawStorage) {
+  // ERC1967 slot https://github.com/OpenZeppelin/openzeppelin-contracts/blob/master/contracts/proxy/ERC1967/ERC1967Utils.sol#L21C53-L21C119
+  const erc1967ImplSlot = '0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc';
+  const erc1967AdminSlot = '0xb53127684a568b3173ae13b9f8a6016e243e63b6e8ee1178d6a717850b5d6103';
+  await Promise.all(
+    (Object.keys(raw) as (keyof typeof raw)[]).map(async (contract) => {
+      if (raw[contract]) {
+        const contractName = isKnownAddress(contract, chainId);
+        // label known contracts if no label was set on foundry
+        if (!raw[contract].label) {
+          if (contractName) raw[contract].label = contractName.join(', ');
+        }
+
+        // create contract diff if storage changed
+        if (raw[contract].stateDiff[erc1967ImplSlot]) {
+          raw[contract].stateDiff[erc1967ImplSlot].label = 'Implementation slot';
+          diffSlot(chainId, contract, raw[contract].stateDiff[erc1967ImplSlot]);
+        }
+
+        // admin slot
+        if (raw[contract].stateDiff[erc1967AdminSlot]) {
+          raw[contract].stateDiff[erc1967ImplSlot].label = 'Admin slot';
+          // ... we might want to fetch the owner in that case
+        }
+
+        // Diff code logic libraries
+        if (contractName && contractName[contractName.length - 1] === 'POOL') {
+          const oldPool = getContract({
+            client: getClient(chainId, {}),
+            abi: IPool_ABI,
+            address: bytes32ToAddress(raw[contract].stateDiff[erc1967ImplSlot].previousValue),
+          });
+          const newPool = getContract({
+            client: getClient(chainId, {}),
+            abi: IPool_ABI,
+            address: bytes32ToAddress(raw[contract].stateDiff[erc1967ImplSlot].previousValue),
+          });
+          const addresses = await Promise.all([
+            oldPool.read.getSupplyLogic(),
+            newPool.read.getSupplyLogic(),
+            oldPool.read.getBorrowLogic(),
+            newPool.read.getBorrowLogic(),
+            oldPool.read.getLiquidationLogic(),
+            newPool.read.getLiquidationLogic(),
+            oldPool.read.getPoolLogic(),
+            newPool.read.getPoolLogic(),
+            oldPool.read.getFlashLoanLogic(),
+            newPool.read.getFlashLoanLogic(),
+            oldPool.read.getEModeLogic(),
+            newPool.read.getEModeLogic(),
+          ]);
+          for (let i = 0; i < addresses.length; i = i + 2) {
+            diffSlot(chainId, contract, {previousValue: addresses[i], newValue: addresses[i + 1]});
+          }
+        }
+      }
+    }),
+  );
+}

src/reports/snapshot-types.ts

@@ -1,4 +1,4 @@
-import {Address} from 'viem';
+import {Address, Hex} from 'viem';
 import {zksync} from 'viem/chains';
 import {z} from 'zod';
 
@@ -103,25 +103,32 @@ const zodChainId = z.nativeEnum(CHAIN_ID);
 
 export type CHAIN_ID = z.infer<typeof zodChainId>;
 
+export const slotDiff = z.object({
+  previousValue: z.string() as z.ZodType<Hex>,
+  newValue: z.string() as z.ZodType<Hex>,
+  label: z.string().optional(), // does not initially exist, but we might want to inject information here
+});
+
+export const rawStorageSchema = z.record(
+  z.string() as z.ZodType<Address>,
+  z.object({
+    label: z.string().nullable(),
+    balanceDiff: z.string().nullable(),
+    stateDiff: z.record(z.string(), slotDiff),
+  }),
+);
+
+export type RawStorage = z.infer<typeof rawStorageSchema>;
+
+export type SlotDiff = z.infer<typeof slotDiff>;
+
 export const aaveV3SnapshotSchema = z.object({
   reserves: z.record(aaveV3ReserveSchema),
   strategies: z.record(aaveV3StrategySchema),
   eModes: z.record(aaveV3EmodeSchema),
   poolConfig: aaveV3ConfigSchema,
   chainId: zodChainId,
-  raw: z
-    .record(
-      z.string(),
-      z.object({
-        label: z.string().nullable(),
-        balanceDiff: z.string().nullable(),
-        stateDiff: z.record(
-          z.string(),
-          z.object({previousValue: z.string(), newValue: z.string()}),
-        ),
-      }),
-    )
-    .optional(),
+  raw: rawStorageSchema.optional(),
 });
 
 export const aDIReceiverConfigSchema = z.object({