CLI: Support Foundry contract naming conventions

packages/core/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 - CLI: Add `--requireReference` option. ([#900](https://github.com/OpenZeppelin/openzeppelin-upgrades/pull/900))
 - CLI: Simplify summary message when using `--contract`. ([#905](https://github.com/OpenZeppelin/openzeppelin-upgrades/pull/905))
+- CLI: Support Foundry contract naming conventions. ([#906](https://github.com/OpenZeppelin/openzeppelin-upgrades/pull/906))
 
 ## 1.30.1 (2023-10-11)
 

packages/core/src/cli/validate/find-contract.test.ts

@@ -0,0 +1,72 @@
+import test from 'ava';
+
+import { findContract } from './find-contract';
+
+test('fully qualified - match', async t => {
+  const sourceContract = {
+    name: 'Foo',
+    fullyQualifiedName: 'contracts/MyContract.sol:Foo',
+  };
+
+  t.deepEqual(findContract('MyContract.sol:Foo', sourceContract, [sourceContract]), sourceContract);
+});
+
+test('fully qualified without folder - match', async t => {
+  const sourceContract = {
+    name: 'Foo',
+    fullyQualifiedName: 'MyContract.sol:Foo',
+  };
+
+  t.deepEqual(findContract('MyContract.sol:Foo', sourceContract, [sourceContract]), sourceContract);
+});
+
+test('short name - match', async t => {
+  const sourceContract = {
+    name: 'Foo',
+    fullyQualifiedName: 'contracts/MyContract.sol:Foo',
+  };
+
+  t.deepEqual(findContract('Foo', sourceContract, [sourceContract]), sourceContract);
+});
+
+test('short name - match without folder', async t => {
+  const sourceContract = {
+    name: 'Foo',
+    fullyQualifiedName: 'MyContract.sol:Foo',
+  };
+
+  t.deepEqual(findContract('Foo', sourceContract, [sourceContract]), sourceContract);
+});
+
+test('short name with .sol - no match', async t => {
+  const sourceContract = {
+    name: 'Foo',
+    fullyQualifiedName: 'contracts/MyContract.sol:Foo',
+  };
+
+  for (const name of ['MyContract.sol', 'Foo.sol']) {
+    const error = t.throws(() => findContract(name, sourceContract, [sourceContract]));
+    t.assert(
+      error?.message.includes(`Could not find contract ${name} referenced in ${sourceContract.fullyQualifiedName}`),
+      error?.message,
+    );
+  }
+});
+
+test('short name with .sol - match', async t => {
+  const sourceContract = {
+    name: 'Foo',
+    fullyQualifiedName: 'contracts/Foo.sol:Foo',
+  };
+
+  t.deepEqual(findContract('Foo.sol', sourceContract, [sourceContract]), sourceContract);
+});
+
+test('short name with .sol - match without folder', async t => {
+  const sourceContract = {
+    name: 'Foo',
+    fullyQualifiedName: 'Foo.sol:Foo',
+  };
+
+  t.deepEqual(findContract('Foo.sol', sourceContract, [sourceContract]), sourceContract);
+});

packages/core/src/cli/validate/find-contract.ts

@@ -23,8 +23,14 @@ export class ReferenceContractNotFound extends Error {
   }
 }
 
-export function findContract(contractName: string, origin: SourceContract | undefined, allContracts: SourceContract[]) {
-  const foundContracts = allContracts.filter(c => c.fullyQualifiedName === contractName || c.name === contractName);
+type SourceContractIdentifier = Pick<SourceContract, 'name' | 'fullyQualifiedName'>;
+
+export function findContract<T extends SourceContractIdentifier>(
+  contractName: string,
+  origin: T | undefined,
+  allContracts: T[],
+): T {
+  const foundContracts = allContracts.filter(c => isMatch(contractName, c));
 
   if (foundContracts.length > 1) {
     const msg =
@@ -42,3 +48,38 @@ export function findContract(contractName: string, origin: SourceContract | unde
     throw new ReferenceContractNotFound(contractName, origin?.fullyQualifiedName);
   }
 }
+
+function isMatch(contractName: string, contract: SourceContractIdentifier) {
+  return (
+    contract.fullyQualifiedName === contractName || // contracts/MyContract.sol:MyContract
+    contract.name === contractName || // MyContract
+    matchesDotSolAndName(contractName, contract) || // MyContract.sol:MyContract
+    matchesDotSol(contractName, contract) // MyContract.sol
+  );
+}
+
+function matchesDotSolAndName(contractName: string, contract: SourceContractIdentifier) {
+  if (contractName.includes('.sol:')) {
+    const [fileWithoutExtension, name] = contractName.split('.sol:');
+    return matchesFullyQualifiedName(fileWithoutExtension, name, contract);
+  } else {
+    return false;
+  }
+}
+
+function matchesDotSol(contractName: string, contract: SourceContractIdentifier) {
+  if (contractName.endsWith('.sol')) {
+    const name = contractName.slice(0, contractName.length - 4);
+    return matchesFullyQualifiedName(name, name, contract);
+  } else {
+    return false;
+  }
+}
+
+function matchesFullyQualifiedName(fileNameWithoutExtension: string, name: string, contract: SourceContractIdentifier) {
+  const lastSlash = contract.fullyQualifiedName.lastIndexOf('/');
+  const fullyQualifiedWithoutPath =
+    lastSlash >= 0 ? contract.fullyQualifiedName.slice(lastSlash + 1) : contract.fullyQualifiedName;
+
+  return contract.name === name && fullyQualifiedWithoutPath === `${fileNameWithoutExtension}.sol:${name}`;
+}