feat(core): Add support for user defined types

.changeset/fuzzy-emus-stare.md

@@ -0,0 +1,6 @@
+---
+'@chugsplash/plugins': patch
+'@chugsplash/core': patch
+---
+
+Add support for user defined types

docs/variables.md

@@ -18,6 +18,7 @@ This is a reference that explains how to assign values to every variable type in
 - [Arrays](#arrays)
 - [Structs](#structs)
 - [Mappings](#mappings)
+- [User-Defined Value Types](#user-defined-types)
 
 
 ## Booleans
@@ -313,3 +314,17 @@ myMultiNestedMapping: {
   ...
 },
 ```
+
+## [User-Defined Value Types](https://docs.soliditylang.org/en/latest/types.html#user-defined-value-types)
+ChugSplash treats your user defined types as if they were their underlying types.
+
+Define your type in Solidity
+```solidity
+type UserDefinedType is uint256;
+UserDefinedType public userDefined;
+```
+
+In your ChugSplash config file:
+```ts
+userDefined: 1
+```

packages/core/src/actions/bundle.ts

@@ -1,6 +1,7 @@
 import { fromHexString, toHexString } from '@eth-optimism/core-utils'
 import { ethers, providers } from 'ethers'
 import MerkleTree from 'merkletreejs'
+import { astDereferencer } from 'solidity-ast/utils'
 
 import {
   CanonicalConfigArtifacts,
@@ -374,9 +375,18 @@ export const makeActionBundleFromConfig = async (
       })
     }
 
+    // Create an AST Dereferencer. We must convert the CompilerOutput type to `any` here because
+    // because a type error will be thrown otherwise. Coverting to `any` is harmless because we use
+    // Hardhat's default `CompilerOutput`, which is what OpenZeppelin expects.
+    const dereferencer = astDereferencer(compilerOutput as any)
+
     // Compute our storage slots.
     // TODO: One day we'll need to refactor this to support Vyper.
-    const slots = computeStorageSlots(storageLayout, contractConfig)
+    const slots = computeStorageSlots(
+      storageLayout,
+      contractConfig,
+      dereferencer
+    )
 
     // Add SET_STORAGE actions for each storage slot that we want to modify.
     for (const slot of slots) {

packages/core/src/languages/solidity/storage.ts

@@ -1,5 +1,6 @@
 import { add0x, fromHexString, remove0x } from '@eth-optimism/core-utils'
 import { BigNumber, ethers, utils } from 'ethers'
+import { ASTDereferencer } from 'solidity-ast/utils'
 
 import { isPreserveKeyword } from '../../utils'
 import { ParsedContractConfig } from '../../config/types'
@@ -9,6 +10,7 @@ import {
   SolidityStorageType,
   StorageSlotSegment,
 } from './types'
+
 import 'core-js/features/array/at'
 
 /**
@@ -67,7 +69,8 @@ export const encodeVariable = (
   storageTypes: {
     [name: string]: SolidityStorageType
   },
-  nestedSlotOffset: string
+  nestedSlotOffset: string,
+  dereferencer: ASTDereferencer
 ): Array<StorageSlotSegment> => {
   // The current slot key is the slot key of the current storage object plus the `nestedSlotOffset`.
   const slotKey = addStorageSlotKeys(storageObj.slot, nestedSlotOffset)
@@ -83,7 +86,11 @@ export const encodeVariable = (
     ]
   }
 
-  const variableType = storageTypes[storageObj.type]
+  const variableType = getStorageType(
+    storageObj.type,
+    storageTypes,
+    dereferencer
+  )
 
   // The Solidity compiler uses four encodings to encode state variables: "inplace", "mapping",
   // "dynamic_array", and "bytes". Each state variable is assigned an encoding depending on its
@@ -104,7 +111,8 @@ export const encodeVariable = (
         storageObj,
         storageTypes,
         elementSlotKey,
-        nestedSlotOffset
+        nestedSlotOffset,
+        dereferencer
       )
     } else if (
       variableType.label === 'address' ||
@@ -245,7 +253,13 @@ export const encodeVariable = (
           )
         }
         slots = slots.concat(
-          encodeVariable(varVal, memberStorageObj, storageTypes, slotKey)
+          encodeVariable(
+            varVal,
+            memberStorageObj,
+            storageTypes,
+            slotKey,
+            dereferencer
+          )
         )
       }
       return slots
@@ -319,7 +333,11 @@ export const encodeVariable = (
         )
       }
 
-      const mappingKeyStorageType = storageTypes[variableType.key]
+      const mappingKeyStorageType = getStorageType(
+        variableType.key,
+        storageTypes,
+        dereferencer
+      )
 
       // Encode the mapping key according to its Solidity compiler encoding. The encoding for the
       // mapping key is 'bytes' if the mapping key is a string or dynamic bytes. Otherwise, the
@@ -368,7 +386,13 @@ export const encodeVariable = (
       // `nestedSlotOffset` to '0' because it isn't used when calculating the storage slot
       // key (we already calculated the storage slot key above).
       slots = slots.concat(
-        encodeVariable(mappingVal, mappingValStorageObj, storageTypes, '0')
+        encodeVariable(
+          mappingVal,
+          mappingValStorageObj,
+          storageTypes,
+          '0',
+          dereferencer
+        )
       )
     }
     return slots
@@ -391,7 +415,8 @@ export const encodeVariable = (
         storageObj,
         storageTypes,
         utils.keccak256(slotKey), // The slot key of the array elements begins at the hash of the `slotKey`.
-        nestedSlotOffset
+        nestedSlotOffset,
+        dereferencer
       )
     )
     return slots
@@ -423,17 +448,27 @@ export const encodeArrayElements = (
     [name: string]: SolidityStorageType
   },
   elementSlotKey: string,
-  nestedSlotOffset: string
+  nestedSlotOffset: string,
+  dereferencer: ASTDereferencer
 ): Array<StorageSlotSegment> => {
-  const elementType = storageTypes[storageObj.type].base
+  const elementType = getStorageType(
+    storageObj.type,
+    storageTypes,
+    dereferencer
+  ).base
 
   if (elementType === undefined) {
     throw new Error(
       `Could not encode array elements for: ${storageObj.label}. Should never happen.`
     )
   }
 
-  const bytesPerElement = Number(storageTypes[elementType].numberOfBytes)
+  const elementStorageType = getStorageType(
+    elementType,
+    storageTypes,
+    dereferencer
+  )
+  const bytesPerElement = Number(elementStorageType.numberOfBytes)
 
   // Calculate the number of slots to increment when iterating over the array elements. This
   // number is only ever greater than one if `bytesPerElement` > 32, which could happen if the
@@ -460,7 +495,8 @@ export const encodeArrayElements = (
           type: elementType,
         },
         storageTypes,
-        nestedSlotOffset
+        nestedSlotOffset,
+        dereferencer
       )
     )
     // Increment the bytes offset every time we iterate over an element.
@@ -528,7 +564,8 @@ export const encodeBytesArrayElements = (
  */
 export const computeStorageSlots = (
   storageLayout: SolidityStorageLayout,
-  contractConfig: ParsedContractConfig
+  contractConfig: ParsedContractConfig,
+  dereferencer: ASTDereferencer
 ): Array<StorageSlotSegment> => {
   const storageEntries: { [storageObjLabel: string]: SolidityStorageObj } = {}
 
@@ -567,7 +604,13 @@ export const computeStorageSlots = (
 
     // Encode this variable as series of storage slot key/value pairs and save it.
     segments = segments.concat(
-      encodeVariable(variableValue, storageObj, storageLayout.types, '0')
+      encodeVariable(
+        variableValue,
+        storageObj,
+        storageLayout.types,
+        '0',
+        dereferencer
+      )
     )
   }
 
@@ -625,3 +668,43 @@ export const computeStorageSlots = (
 
   return segments
 }
+
+export const getStorageType = (
+  variableType: string,
+  storageTypes: {
+    [name: string]: SolidityStorageType
+  },
+  dereferencer: ASTDereferencer
+): SolidityStorageType => {
+  if (!variableType.startsWith('t_userDefinedValueType')) {
+    return storageTypes[variableType]
+  } else {
+    const userDefinedValueAstId = variableType.split(')').at(-1)
+
+    if (userDefinedValueAstId === undefined) {
+      throw new Error(
+        `Could not find AST ID for variable type: ${variableType}. Should never happen.`
+      )
+    }
+
+    const userDefinedValueNode = dereferencer(
+      ['UserDefinedValueTypeDefinition'],
+      parseInt(userDefinedValueAstId, 10)
+    )
+
+    const label =
+      userDefinedValueNode.underlyingType.typeDescriptions.typeString
+    if (label === undefined || label === null) {
+      throw new Error(
+        `Could not find label for user-defined value type: ${variableType}. Should never happen.`
+      )
+    }
+
+    const { encoding, numberOfBytes } = storageTypes[variableType]
+    return {
+      label,
+      encoding,
+      numberOfBytes,
+    }
+  }
+}

packages/core/src/utils.ts

@@ -1687,6 +1687,7 @@ export const getCanonicalConfigArtifacts = async (
     for (const { compilerInput, compilerOutput } of solcArray) {
       const contractOutput =
         compilerOutput.contracts?.[sourceName]?.[contractName]
+
       if (contractOutput !== undefined) {
         const creationCodeWithConstructorArgs =
           getCreationCodeWithConstructorArgs(

packages/demo/contracts/HelloChugSplash.sol

@@ -2,6 +2,7 @@
 pragma solidity ^0.8.15;
 
 contract HelloChugSplash {
+    type UFixed256x18 is uint256;
     uint8 public number;
     bool public stored;
     address public otherStorage;

packages/plugins/chugsplash/foundry/deploy.t.js

@@ -12,6 +12,29 @@ const variables = {
   bytes32Test: '0x' + '11'.repeat(32),
   longBytesTest:
     '0x123456789101112131415161718192021222324252627282930313233343536373839404142434445464',
+  userDefinedTypeTest: '1000000000000000000',
+  userDefinedBytesTest: '0x' + '11'.repeat(32),
+  userDefinedInt: ethers.constants.MinInt256.toString(),
+  userDefinedInt8: -128,
+  userDefinedUint8: 255,
+  userDefinedBool: true,
+  userDefinedFixedArray: ['1000000000000000000', '1000000000000000000'],
+  userDefinedFixedNestedArray: [
+    ['1000000000000000000', '1000000000000000000'],
+    ['1000000000000000000', '1000000000000000000'],
+  ],
+  userDefinedDynamicArray: [
+    '1000000000000000000',
+    '1000000000000000000',
+    '1000000000000000000',
+  ],
+  stringToUserDefinedMapping: {
+    testKey: '1000000000000000000',
+  },
+  userDefinedToStringMapping: {
+    // eslint-disable-next-line prettier/prettier
+    '1000000000000000000': 'testVal',
+  },
   contractTest: '0x' + '11'.repeat(20),
   enumTest: 1,
   simpleStruct: {
@@ -24,6 +47,7 @@ const variables = {
     b: {
       5: 'testVal',
     },
+    c: '1000000000000000000',
   },
   uint64FixedArray: [1, 10, 100, 1_000, 10_000],
   uint128FixedNestedArray: [

packages/plugins/contracts/Storage.sol

@@ -2,11 +2,19 @@
 pragma solidity ^0.8.9;
 
 contract Storage {
+    type UserDefinedType is uint256;
+    type UserDefinedBytes32 is bytes32;
+    type UserDefinedInt is int;
+    type UserDefinedInt8 is int8;
+    type UserDefinedUint8 is uint8;
+    type UserDefinedBool is bool;
+
     enum TestEnum { A, B, C }
     struct SimpleStruct { bytes32 a; uint128 b; uint128 c; }
     struct ComplexStruct {
         int32 a;
         mapping(uint32 => string) b;
+        UserDefinedType c;
     }
 
     int public minInt256;
@@ -18,6 +26,20 @@ contract Storage {
     bytes public bytesTest;
     bytes public longBytesTest;
     bytes32 public bytes32Test;
+
+    UserDefinedType public userDefinedTypeTest;
+    UserDefinedBytes32 public userDefinedBytesTest;
+    UserDefinedInt public userDefinedInt;
+    UserDefinedInt8 public userDefinedInt8;
+    UserDefinedUint8 public userDefinedUint8;
+    UserDefinedBool public userDefinedBool;
+    mapping(UserDefinedType => string) public userDefinedToStringMapping;
+    mapping(string => UserDefinedType) public stringToUserDefinedMapping;
+    UserDefinedType[2] public userDefinedFixedArray;
+    UserDefinedType[2][2] public userDefinedFixedNestedArray;
+    UserDefinedType[] public userDefinedDynamicArray;
+
+
     Storage public contractTest;
     TestEnum public enumTest;
     SimpleStruct public simpleStruct;