fix: slot-overwrite bug

package.json

@@ -74,7 +74,6 @@
     "@typechain/ethers-v5": "7.0.1",
     "@typechain/hardhat": "2.0.2",
     "@types/chai": "4.2.18",
-    "@types/chai-as-promised": "7.1.5",
     "@types/diff": "^5.0.1",
     "@types/lodash": "4.14.170",
     "@types/lodash.isequal": "^4.5.5",

src/logic/editable-storage-logic.ts

@@ -1,5 +1,5 @@
 import { SetVariablesType, SmockVMManager } from '../types';
-import { fromHexString, toFancyAddress } from '../utils';
+import { fromHexString, remove0x, toFancyAddress, toHexString } from '../utils';
 import { computeStorageSlots, SolidityStorageLayout } from '../utils/storage';
 
 export class EditableStorageLogic {
@@ -19,7 +19,14 @@ export class EditableStorageLogic {
     const slots = computeStorageSlots(this.storageLayout, { [variableName]: value });
 
     for (const slot of slots) {
-      await this.vmManager.putContractStorage(toFancyAddress(this.contractAddress), fromHexString(slot.key), fromHexString(slot.val));
+      let prevStorageValue = await this.vmManager.getContractStorage(toFancyAddress(this.contractAddress), fromHexString(slot.key));
+      let stringValue = remove0x(toHexString(prevStorageValue));
+      if (stringValue && (slot.type === 'address' || slot.type === 'bool' || slot.type.startsWith('contract'))) {
+        let lastResult = slot.val.slice(0, slot.val.length - stringValue.length).concat(stringValue);
+        await this.vmManager.putContractStorage(toFancyAddress(this.contractAddress), fromHexString(slot.key), fromHexString(lastResult));
+      } else {
+        await this.vmManager.putContractStorage(toFancyAddress(this.contractAddress), fromHexString(slot.key), fromHexString(slot.val));
+      }
     }
   }
 

src/utils/storage.ts

@@ -40,6 +40,7 @@ export interface SolidityStorageLayout {
 interface StorageSlotPair {
   key: string;
   val: string;
+  type: string;
 }
 
 /**
@@ -139,6 +140,7 @@ export function computeStorageSlots(storageLayout: SolidityStorageLayout, variab
       prevSlots.push({
         key: slot.key,
         val: mergedVal,
+        type: slot.type,
       });
     }
 
@@ -228,6 +230,7 @@ function encodeVariable(
         {
           key: slotKey,
           val: padNumHexSlotValue(variable, storageObj.offset),
+          type: variableType.label,
         },
       ];
     } else if (variableType.label === 'bool') {
@@ -249,6 +252,7 @@ function encodeVariable(
         {
           key: slotKey,
           val: padNumHexSlotValue(variable ? '1' : '0', storageObj.offset),
+          type: variableType.label,
         },
       ];
     } else if (variableType.label.startsWith('bytes')) {
@@ -260,6 +264,7 @@ function encodeVariable(
         {
           key: slotKey,
           val: padBytesHexSlotValue(remove0x(variable).padEnd(parseInt(variableType.numberOfBytes, 10) * 2, '0'), storageObj.offset),
+          type: variableType.label,
         },
       ];
     } else if (variableType.label.startsWith('uint') || variableType.label.startsWith('int')) {
@@ -271,6 +276,7 @@ function encodeVariable(
         {
           key: slotKey,
           val: padNumHexSlotValue(variable, storageObj.offset),
+          type: variableType.label,
         },
       ];
     } else if (variableType.label.startsWith('struct')) {
@@ -314,6 +320,7 @@ function encodeVariable(
               ethers.BigNumber.from(bytes.length * 2).toHexString(),
             ])
           ),
+          type: variableType.label,
         },
       ];
     } else {
@@ -324,6 +331,7 @@ function encodeVariable(
       slots = slots.concat({
         key: slotKey,
         val: padNumHexSlotValue(bytes.length * 2 + 1, 0),
+        type: variableType.label,
       });
 
       // Each storage slot has 32 bytes so we make sure to slice the large bytes into 32bytes chunks
@@ -335,6 +343,7 @@ function encodeVariable(
         slots = slots.concat({
           key: key,
           val: ethers.utils.hexlify(ethers.utils.concat([bytes.slice(i * 32, i * 32 + 32), ethers.constants.HashZero]).slice(0, 32)),
+          type: variableType.label,
         });
       }
 

test/contracts/mock/Bug.sol

@@ -0,0 +1,11 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+contract Bug {
+  address public myAddress;
+  bool public myBool;
+  address public newAddress;
+  address public newAddress2;
+  bool public newBool;
+  bool public newBool2;
+}

test/unit/fake/initialization.spec.ts

@@ -3,11 +3,10 @@ import { makeRandomAddress } from '@src/utils';
 import { Receiver, Receiver__factory, Returner } from '@typechained';
 import receiverArtifact from 'artifacts/test/contracts/watchable-function-logic/Receiver.sol/Receiver.json';
 import chai, { expect } from 'chai';
-import chaiAsPromised from 'chai-as-promised';
 import { ethers, network } from 'hardhat';
 import storageArtifact from 'test/unit/fake/testdata/Storage.json';
 
-chai.use(chaiAsPromised);
+chai.use(smock.matchers);
 
 describe('Fake: Initialization', () => {
   it('should work with the contract name', async () => {
@@ -91,21 +90,4 @@ describe('Fake: Initialization', () => {
     const fake = await smock.fake(storageArtifact);
     expect(fake.store._watchable).not.to.be.undefined;
   });
-
-  it('should throw error for invalid json string abi', async () => {
-    await expect(smock.fake(`{invalid}`)).to.be.rejectedWith(
-      Error,
-      `unable to generate smock spec from abi string.\nUnexpected token i in JSON at position 1`
-    );
-  });
-
-  it('should throw error for non existent contract', async () => {
-    let regex: RegExp = /unable to generate smock spec from contract name*/;
-    await expect(smock.fake('NonExistentContract')).to.be.rejectedWith(Error, regex);
-  });
-
-  it('should throw error if contract name is ambiguous', async () => {
-    let regex: RegExp = /unable to generate smock spec from contract name*/;
-    await expect(smock.fake('Storage')).to.be.rejectedWith(Error, regex);
-  });
 });

test/unit/mock/editable-storage-logic.spec.ts

@@ -14,7 +14,7 @@ describe('Mock: Editable storage logic', () => {
   });
 
   beforeEach(async () => {
-    mock = await storageGetterFactory.deploy(1);
+    mock = await storageGetterFactory.deploy(4);
   });
 
   describe('setVariable', () => {
@@ -64,7 +64,6 @@ describe('Mock: Editable storage logic', () => {
 
     it('should be able to set an address in a packed storage slot', async () => {
       await mock.setVariable('_packedB', ADDRESS_EXAMPLE);
-
       expect(await mock.getPackedAddress()).to.equal(ADDRESS_EXAMPLE);
     });
 

test/unit/mock/mytest.spec.ts

@@ -0,0 +1,40 @@
+import { MockContract, MockContractFactory, smock } from '@src';
+import { ADDRESS_EXAMPLE, ADDRESS_EXAMPLE2, ADDRESS_EXAMPLE3 } from '@test-utils';
+import { Bug, Bug__factory } from '@typechained';
+import { expect } from 'chai';
+
+describe('Bug Fix', () => {
+  let bugFactory: MockContractFactory<Bug__factory>;
+  let mock: MockContract<Bug>;
+
+  before(async () => {
+    bugFactory = await smock.mock('Bug');
+  });
+
+  beforeEach(async () => {
+    mock = await bugFactory.deploy();
+  });
+
+  describe('setVariable', () => {
+    it('should be able to fix slot overwrite on packed variables', async () => {
+      await mock.setVariable('myAddress', ADDRESS_EXAMPLE);
+      await mock.setVariable('myBool', true);
+      await mock.setVariable('newAddress', ADDRESS_EXAMPLE2);
+      await mock.setVariable('newAddress2', ADDRESS_EXAMPLE3);
+      await mock.setVariable('newBool', false);
+      await mock.setVariable('newBool2', true);
+      console.log('____myAddress____', await mock.myAddress());
+      console.log('____myBool____', await mock.myBool());
+      console.log('____newAddress____', await mock.newAddress());
+      console.log('____newAddress2____', await mock.newAddress2());
+      console.log('____newBool____', await mock.newBool());
+      console.log('____newBool2____', await mock.newBool2());
+      expect(await mock.myAddress()).to.equal(ADDRESS_EXAMPLE);
+      expect(await mock.myBool()).to.equal(true);
+      expect(await mock.newAddress()).to.equal(ADDRESS_EXAMPLE2);
+      expect(await mock.newAddress2()).to.equal(ADDRESS_EXAMPLE3);
+      expect(await mock.newBool()).to.equal(false);
+      expect(await mock.newBool2()).to.equal(true);
+    });
+  });
+});

test/unit/mock/mytest.spec2.ts

@@ -0,0 +1,25 @@
+import { MockContract, MockContractFactory, smock } from '@src';
+import { ADDRESS_EXAMPLE } from '@test-utils';
+import { StorageGetter, StorageGetter__factory } from '@typechained';
+import { expect } from 'chai';
+
+describe('Mock: Editable storage logic', () => {
+  let storageGetterFactory: MockContractFactory<StorageGetter__factory>;
+  let mock: MockContract<StorageGetter>;
+
+  before(async () => {
+    storageGetterFactory = await smock.mock('StorageGetter');
+  });
+
+  beforeEach(async () => {
+    mock = await storageGetterFactory.deploy(1);
+  });
+
+  describe('setVariable', () => {
+    it('should be able to fix bug', async () => {
+      await mock.setVariable('_address', ADDRESS_EXAMPLE);
+      await mock.setVariable('_bool', true);
+      expect(await mock.getAddress()).to.equal(ADDRESS_EXAMPLE);
+    });
+  });
+});

test/utils/constants.ts

@@ -2,6 +2,10 @@ import { utils } from 'ethers';
 
 export const ADDRESS_EXAMPLE = '0x558ba9b8d78713fbf768c1f8a584485B4003f43F';
 
+export const ADDRESS_EXAMPLE2 = '0x4fd96172C40913c8100a47337f21DBf25C9e2fC6';
+
+export const ADDRESS_EXAMPLE3 = '0x62eA762d9f2672f82d26b94f56849881831B9eaB';
+
 export const BYTES32_EXAMPLE = '0x1234123412341234123412341234123412341234123412341234123412341234';
 
 export const BYTES_EXAMPLE =

yarn.lock

@@ -914,13 +914,6 @@
   dependencies:
     "@types/node" "*"
 
-"@types/[email protected]":
-  version "7.1.5"
-  resolved "https://registry.yarnpkg.com/@types/chai-as-promised/-/chai-as-promised-7.1.5.tgz#6e016811f6c7a64f2eed823191c3a6955094e255"
-  integrity sha512-jStwss93SITGBwt/niYrkf2C+/1KTeZCZl1LaeezTlqppAKeoQC7jxyqYuP72sxBGKCIbw7oHgbYssIRzT5FCQ==
-  dependencies:
-    "@types/chai" "*"
-
 "@types/chai@*":
   version "4.2.21"
   resolved "https://registry.yarnpkg.com/@types/chai/-/chai-4.2.21.tgz#9f35a5643129df132cf3b5c1ec64046ea1af0650"