fix(core,pg): refactor functions that get build info and storage layout

.changeset/silly-jars-think.md

@@ -0,0 +1,6 @@
+---
+'@chugsplash/plugins': patch
+'@chugsplash/core': patch
+---
+
+Refactor functions that get build info and storage layout

packages/core/src/actions/artifacts.ts

@@ -10,6 +10,7 @@ import {
 } from '../languages/solidity/types'
 import { Integration } from '../constants'
 import {
+  addEnumMembersToStorageLayout,
   createDeploymentFolderForNetwork,
   readBuildInfo,
   readContractArtifact,
@@ -104,24 +105,26 @@ export const getConstructorArgs = (
  * Reads the storageLayout portion of the compiler artifact for a given contract. Reads the
  * artifact from the local file system.
  *
- * @param fullyQualifiedName Fully qualified name of the contract.
+ * @param contractFullyQualifiedName Fully qualified name of the contract.
  * @param artifactFolder Relative path to the folder where artifacts are stored.
  * @return Storage layout object from the compiler output.
  */
-export const readStorageLayout = (
-  fullyQualifiedName: string,
-  artifactPaths: ArtifactPaths,
-  integration: Integration
+export const getStorageLayout = (
+  buildInfoPath: string,
+  contractFullyQualifiedName: string
 ): SolidityStorageLayout => {
-  const { sourceName, contractName } = readContractArtifact(
-    artifactPaths,
-    fullyQualifiedName,
-    integration
+  const buildInfo = readBuildInfo(buildInfoPath)
+  const [sourceName, contractName] = contractFullyQualifiedName.split(':')
+  const contractOutput = buildInfo.output.contracts[sourceName][contractName]
+  const outputSource = buildInfo.output.sources[sourceName]
+
+  addEnumMembersToStorageLayout(
+    contractOutput.storageLayout,
+    contractName,
+    outputSource.ast.nodes
   )
-  const buildInfo = readBuildInfo(artifactPaths, fullyQualifiedName)
-  const output = buildInfo.output.contracts[sourceName][contractName]
 
-  return (output as any).storageLayout
+  return contractOutput.storageLayout
 }
 
 export const getDeployedBytecode = async (
@@ -156,13 +159,12 @@ export const createDeploymentArtifacts = async (
     parsedConfig.contracts
   )) {
     const artifact = readContractArtifact(
-      artifactPaths,
-      contractConfig.contract,
+      artifactPaths[referenceName].contractArtifactPath,
       integration
     )
     const { sourceName, contractName, bytecode, abi } = artifact
 
-    const buildInfo = readBuildInfo(artifactPaths, contractConfig.contract)
+    const buildInfo = readBuildInfo(artifactPaths[referenceName].buildInfoPath)
 
     const { constructorArgValues } = getConstructorArgs(
       parsedConfig,
@@ -203,10 +205,9 @@ export const createDeploymentArtifacts = async (
       deployedBytecode: await provider.getCode(contractConfig.proxy),
       devdoc,
       userdoc,
-      storageLayout: readStorageLayout(
-        contractConfig.contract,
-        artifactPaths,
-        integration
+      storageLayout: getStorageLayout(
+        artifactPaths[referenceName].buildInfoPath,
+        contractConfig.contract
       ),
     }
 

packages/core/src/actions/bundle.ts

@@ -11,7 +11,7 @@ import {
 } from '../languages/solidity/types'
 import { readContractArtifact } from '../utils'
 import {
-  readStorageLayout,
+  getStorageLayout,
   getCreationCodeWithConstructorArgs,
 } from './artifacts'
 import {
@@ -231,15 +231,13 @@ export const bundleLocal = async (
   for (const [referenceName, contractConfig] of Object.entries(
     parsedConfig.contracts
   )) {
-    const storageLayout = readStorageLayout(
-      contractConfig.contract,
-      artifactPaths,
-      integration
+    const storageLayout = getStorageLayout(
+      artifactPaths[referenceName].buildInfoPath,
+      contractConfig.contract
     )
 
     const { abi, bytecode } = readContractArtifact(
-      artifactPaths,
-      contractConfig.contract,
+      artifactPaths[referenceName].contractArtifactPath,
       integration
     )
     const creationCode = getCreationCodeWithConstructorArgs(

packages/core/src/actions/types.ts

@@ -1,15 +1,4 @@
 import { BigNumber } from 'ethers'
-import { Fragment } from 'ethers/lib/utils'
-
-// TODO
-export type BuildInfo = any
-export type ContractArtifact = {
-  abi: Array<Fragment>
-  sourceName: string
-  contractName: string
-  bytecode: string
-}
-export type ContractASTNode = any
 
 /**
  * Possible action types.

packages/core/src/config/parse.ts

@@ -173,20 +173,14 @@ export const parseChugSplashConfig = async (
       )
     }
 
-    const {
-      contract,
-      externalProxy,
-      externalProxyType,
-      variables,
-      constructorArgs,
-    } = userContractConfig
+    const { externalProxy, externalProxyType, variables, constructorArgs } =
+      userContractConfig
 
     // Change the `contract` fields to be a fully qualified name. This ensures that it's easy for the
     // executor to create the `CanonicalConfigArtifacts` when it eventually compiles the canonical
     // config.
     const { sourceName, contractName } = readContractArtifact(
-      artifactPaths,
-      contract,
+      artifactPaths[referenceName].contractArtifactPath,
       integration
     )
     const contractFullyQualifiedName = `${sourceName}:${contractName}`

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

@@ -1,3 +1,5 @@
+import { Fragment } from 'ethers/lib/utils'
+
 /**
  * Represents the JSON objects outputted by the Solidity compiler that describe the structure of
  * state within the contract. See
@@ -39,11 +41,10 @@ export interface SolidityStorageLayout {
 }
 
 /**
- * Mapping from a contract to its build info file path and artifact path. Note that each
- * contract matches the `contract` field in the `UserChugSplashConfig`.
+ * Mapping from a contract's reference name to its build info file path and artifact path.
  */
 export type ArtifactPaths = {
-  [contract: string]: {
+  [referenceName: string]: {
     buildInfoPath: string
     contractArtifactPath: string
   }
@@ -75,6 +76,17 @@ export interface CompilerInput {
   }
 }
 
+// TODO
+export type BuildInfo = any
+export type ContractASTNode = any
+
+export type ContractArtifact = {
+  abi: Array<Fragment>
+  sourceName: string
+  contractName: string
+  bytecode: string
+}
+
 export interface CompilerOutputMetadata {
   sources: {
     [sourceName: string]: {

packages/core/src/tasks/index.ts

@@ -283,8 +283,10 @@ export const chugsplashCommitAbstractSubtask = async (
   }
 
   const chugsplashInputs: Array<ChugSplashInput> = []
-  for (const contractConfig of Object.values(parsedConfig.contracts)) {
-    const buildInfo = readBuildInfo(artifactPaths, contractConfig.contract)
+  for (const [referenceName, contractConfig] of Object.entries(
+    parsedConfig.contracts
+  )) {
+    const buildInfo = readBuildInfo(artifactPaths[referenceName].buildInfoPath)
 
     const prevChugSplashInput = chugsplashInputs.find(
       (input) => input.solcLongVersion === buildInfo.solcLongVersion

packages/core/src/utils.ts

@@ -41,19 +41,19 @@ import {
   UserConfigVariable,
 } from './config/types'
 import {
-  BuildInfo,
   ChugSplashActionBundle,
   ChugSplashActionType,
-  ContractArtifact,
-  ContractASTNode,
-  readStorageLayout,
+  getStorageLayout,
 } from './actions'
 import { Integration, keywords } from './constants'
 import 'core-js/features/array/at'
 import { getLatestDeployedStorageLayout } from './deployed'
 import { FoundryContractArtifact } from './types'
 import {
   ArtifactPaths,
+  BuildInfo,
+  ContractArtifact,
+  ContractASTNode,
   CompilerOutputSources,
   SolidityStorageLayout,
 } from './languages/solidity/types'
@@ -772,10 +772,9 @@ permission to call the 'upgradeTo' function on each of them.
           remoteExecution,
           canonicalConfigFolderPath
         )
-        const newStorageLayout = readStorageLayout(
-          contractConfig.contract,
-          artifactPaths,
-          integration
+        const newStorageLayout = getStorageLayout(
+          artifactPaths[referenceName].buildInfoPath,
+          contractConfig.contract
         )
 
         assertStorageCompatiblePreserveKeywords(
@@ -798,8 +797,7 @@ permission to call the 'upgradeTo' function on each of them.
       // user will be able to upgrade the proxy in the future.
       if (contractConfig.proxyType === 'oz-uups') {
         const artifact = readContractArtifact(
-          artifactPaths,
-          contractConfig.contract,
+          artifactPaths[referenceName].contractArtifactPath,
           integration
         )
         const containsPublicUpgradeTo = artifact.abi.some(
@@ -964,13 +962,17 @@ export const assertValidContracts = (
   parsedConfig: ParsedChugSplashConfig,
   artifactPaths: ArtifactPaths
 ) => {
-  for (const contractConfig of Object.values(parsedConfig.contracts)) {
-    // Split the contract's fully qualified name into its source name and contract name.
+  for (const [referenceName, contractConfig] of Object.entries(
+    parsedConfig.contracts
+  )) {
+    // Get the source name and contract name from its fully qualified name
     const [sourceName, contractName] = contractConfig.contract.split(':')
-    const buildInfo = readBuildInfo(artifactPaths, contractConfig.contract)
 
-    const sourceOutput = buildInfo.output.sources[sourceName]
-    const childContractNode = sourceOutput.ast.nodes.find(
+    const buildInfoPath = artifactPaths[referenceName].buildInfoPath
+    const buildInfo = readBuildInfo(buildInfoPath)
+    const outputSource = buildInfo.output.sources[sourceName]
+
+    const childContractNode = outputSource.ast.nodes.find(
       (node) =>
         node.nodeType === 'ContractDefinition' && node.name === contractName
     )
@@ -1003,7 +1005,7 @@ export const assertValidContracts = (
               `User attempted to assign a value to a non-immutable state variable '${node.name}' in\n` +
                 `the contract: ${contractNode.name}. This is not allowed because the value will not exist in\n` +
                 `the upgradeable contract. Please remove the value in the contract and define it in your ChugSplash\n` +
-                `file instead. You can also set '${node.name} to be a constant or immutable variable instead.`
+                `file instead. Alternatively, can also set '${node.name} to be a constant or immutable variable.`
             )
           }
 
@@ -1136,28 +1138,11 @@ export const getBundleCompletionTxnHash = async (
 
 /**
  * Retrieves an artifact by name from the local file system.
- *
- * @param name Contract name or fully qualified name.
- * @returns Artifact.
  */
 export const readContractArtifact = (
-  artifactPaths: ArtifactPaths,
-  contract: string,
+  contractArtifactPath: string,
   integration: Integration
 ): ContractArtifact => {
-  let contractArtifactPath: string
-  if (artifactPaths[contract]) {
-    contractArtifactPath = artifactPaths[contract].contractArtifactPath
-  } else {
-    // The contract must be a fully qualified name.
-    const contractName = contract.split(':').at(-1)
-    if (contractName === undefined) {
-      throw new Error('Could not use contract name to get build info')
-    } else {
-      contractArtifactPath = artifactPaths[contractName].contractArtifactPath
-    }
-  }
-
   const artifact: ContractArtifact = JSON.parse(
     fs.readFileSync(contractArtifactPath, 'utf8')
   )
@@ -1174,44 +1159,35 @@ export const readContractArtifact = (
 /**
  * Reads the build info from the local file system.
  *
- * @param artifactPaths ArtifactPaths object.
- * @param fullyQualifiedName Fully qualified name of the contract.
+ * @param buildInfoPath Path to the build info file.
  * @returns BuildInfo object.
  */
-export const readBuildInfo = (
-  artifactPaths: ArtifactPaths,
-  fullyQualifiedName: string
-): BuildInfo => {
-  const [sourceName, contractName] = fullyQualifiedName.split(':')
-  const { buildInfoPath } =
-    artifactPaths[fullyQualifiedName] ?? artifactPaths[contractName]
+export const readBuildInfo = (buildInfoPath: string): BuildInfo => {
   const buildInfo: BuildInfo = JSON.parse(
     fs.readFileSync(buildInfoPath, 'utf8')
   )
 
-  const contractOutput = buildInfo.output.contracts[sourceName][contractName]
-  const sourceNodes = buildInfo.output.sources[sourceName].ast.nodes
-
   if (!semver.satisfies(buildInfo.solcVersion, '>=0.4.x <0.9.x')) {
     throw new Error(
       `Storage layout for Solidity version ${buildInfo.solcVersion} not yet supported. Sorry!`
     )
   }
 
-  if (!('storageLayout' in contractOutput)) {
+  if (
+    !buildInfo.input.settings.outputSelection['*']['*'].includes(
+      'storageLayout'
+    )
+  ) {
     throw new Error(
-      `Storage layout for ${fullyQualifiedName} not found. Did you forget to set the storage layout
-compiler option in your hardhat config? Read more:
-https://github.com/ethereum-optimism/smock#note-on-using-smoddit`
+      `Storage layout not found. Did you forget to set the "storageLayout" compiler option in your\n` +
+        `Hardhat/Foundry config file?\n\n` +
+        `If you're using Hardhat, see how to configure your project here:\n` +
+        `https://github.com/chugsplash/chugsplash/blob/develop/docs/hardhat/setup-project.md#setup-chugsplash-using-typescript\n\n` +
+        `If you're using Foundry, see how to configure your project here:\n` +
+        `https://github.com/chugsplash/chugsplash/blob/develop/docs/foundry/getting-started.md#3-configure-your-foundrytoml-file`
     )
   }
 
-  addEnumMembersToStorageLayout(
-    contractOutput.storageLayout,
-    contractName,
-    sourceNodes
-  )
-
   return buildInfo
 }
 

packages/plugins/src/foundry/utils/index.ts

@@ -68,9 +68,11 @@ export const getArtifactPaths = async (
 ): Promise<ArtifactPaths> => {
   const artifactPaths: ArtifactPaths = {}
 
-  for (const { contract } of Object.values(contractConfigs)) {
+  for (const [referenceName, contractConfig] of Object.entries(
+    contractConfigs
+  )) {
     const { sourceName, contractName } = getContractArtifact(
-      contract,
+      contractConfig.contract,
       artifactFolder
     )
     const buildInfo = await getBuildInfo(buildInfoFolder, sourceName)
@@ -79,7 +81,7 @@ export const getArtifactPaths = async (
     const fileName = `${contractName}.json`
     const contractArtifactPath = path.join(artifactFolder, folderName, fileName)
 
-    artifactPaths[contract] = {
+    artifactPaths[referenceName] = {
       buildInfoPath: path.join(buildInfoFolder, `${buildInfo.id}.json`),
       contractArtifactPath,
     }