From e6e87ae1971e47a1b0c15cd44b376200902043c0 Mon Sep 17 00:00:00 2001
From: smartcontracts <kelvinfichter@gmail.com>
Date: Thu, 27 May 2021 01:29:49 -0400
Subject: [PATCH] feat[smock]: add support for overloaded functions (#966)

* feat[smock]: add support for overloaded functions

* chore: add changeset
---
 .changeset/neat-melons-lie.md                 |  5 +++++
 packages/smock/src/smockit/smockit.ts         |  2 +-
 .../TestHelpers_BasicReturnContract.sol       | 19 +++++++++++++++++++
 .../smockit/function-manipulation.spec.ts     | 17 +++++++++++++++++
 4 files changed, 42 insertions(+), 1 deletion(-)
 create mode 100644 .changeset/neat-melons-lie.md

diff --git a/.changeset/neat-melons-lie.md b/.changeset/neat-melons-lie.md
new file mode 100644
index 000000000000..833a9b34d248
--- /dev/null
+++ b/.changeset/neat-melons-lie.md
@@ -0,0 +1,5 @@
+---
+'@eth-optimism/smock': patch
+---
+
+Fix a bug where overloaded functions would not be handled correctly
diff --git a/packages/smock/src/smockit/smockit.ts b/packages/smock/src/smockit/smockit.ts
index 0a844dba640b..9ca98ba8c087 100644
--- a/packages/smock/src/smockit/smockit.ts
+++ b/packages/smock/src/smockit/smockit.ts
@@ -211,7 +211,7 @@ export const smockit = async (
     let mockFn: any
     if (fn !== null) {
       params = this.interface.decodeFunctionData(fn, toHexString(data))
-      mockFn = this.smocked[fn.name]
+      mockFn = this.smocked[fn.name] || this.smocked[fn.format()]
     } else {
       params = toHexString(data)
       mockFn = this.smocked.fallback
diff --git a/packages/smock/test/contracts/TestHelpers_BasicReturnContract.sol b/packages/smock/test/contracts/TestHelpers_BasicReturnContract.sol
index 1a56f11e127f..e575412ee1c6 100644
--- a/packages/smock/test/contracts/TestHelpers_BasicReturnContract.sol
+++ b/packages/smock/test/contracts/TestHelpers_BasicReturnContract.sol
@@ -131,4 +131,23 @@ contract TestHelpers_BasicReturnContract {
             uint256[] memory _out
         )
     {}
+
+    function overloadedFunction(
+        uint256 _paramA,
+        uint256 _paramB
+    )
+        public
+        returns (
+            uint256
+        )
+    {}
+
+    function overloadedFunction(
+        uint256
+    )
+        public
+        returns (
+            uint256
+        )
+    {}
 }
diff --git a/packages/smock/test/smockit/function-manipulation.spec.ts b/packages/smock/test/smockit/function-manipulation.spec.ts
index de94a9bf2576..cc481e5b01dc 100644
--- a/packages/smock/test/smockit/function-manipulation.spec.ts
+++ b/packages/smock/test/smockit/function-manipulation.spec.ts
@@ -133,6 +133,23 @@ describe('[smock]: function manipulation tests', () => {
       // TODO
     })
 
+    describe('overloaded functions', () => {
+      it('should be able to modify both versions of an overloaded function', async () => {
+        const expected1 = 1234
+        const expected2 = 5678
+        mock.smocked['overloadedFunction(uint256)'].will.return.with(expected1)
+        mock.smocked['overloadedFunction(uint256,uint256)'].will.return.with(
+          expected2
+        )
+        expect(
+          await mock.callStatic['overloadedFunction(uint256)'](0)
+        ).to.equal(expected1)
+        expect(
+          await mock.callStatic['overloadedFunction(uint256,uint256)'](0, 0)
+        ).to.equal(expected2)
+      })
+    })
+
     describe('returning with data', () => {
       describe('fixed data types', () => {
         describe('default behaviors', () => {