From d9eed8c9fb2c164ae12b6bd8f4b44f71e563aee5 Mon Sep 17 00:00:00 2001
From: drgorillamd <83670532+drgorillamd@users.noreply.github.com>
Date: Thu, 16 May 2024 15:02:16 +0200
Subject: [PATCH] fix: invariant + optim

---
 package.json                             |  2 +-
 test/invariants/Greeter.invariants.t.sol | 36 +++++++++++++++---------
 test/invariants/Greeter.symbolic.t.sol   | 34 +++++++++++-----------
 3 files changed, 41 insertions(+), 31 deletions(-)

diff --git a/package.json b/package.json
index 47caadac..b9de5ca0 100644
--- a/package.json
+++ b/package.json
@@ -22,7 +22,7 @@
     "lint:sol-tests": "solhint -c .solhint.tests.json 'test/**/*.sol'",
     "prepare": "husky install",
     "test": "forge test -vvv",
-    "test:fuzz": "echidna test/invariants/Greeter.invariants.t.sol --contract GreeterInvariant --corpus-dir test/invariants/echidna_coverage/",
+    "test:fuzz": "echidna test/invariants/Greeter.invariants.t.sol --contract GreeterInvariant --corpus-dir test/invariants/echidna_coverage/ --test-mode assertion",
     "test:integration": "forge test --match-contract Integration -vvv",
     "test:symbolic": "halmos",
     "test:unit": "forge test --match-contract Unit -vvv",
diff --git a/test/invariants/Greeter.invariants.t.sol b/test/invariants/Greeter.invariants.t.sol
index 5489d2d5..b1737420 100644
--- a/test/invariants/Greeter.invariants.t.sol
+++ b/test/invariants/Greeter.invariants.t.sol
@@ -3,22 +3,30 @@ pragma solidity 0.8.23;
 
 import {Greeter, IERC20} from '../../src/contracts/Greeter.sol';
 
-contract GreeterInvariant is Greeter {
-  constructor() Greeter('a', IERC20(address(1))) {}
+interface IHevm {
+  function prank(address) external;
+}
+
+contract GreeterInvariant {
+  address constant HEVM_ADDRESS = 0x7109709ECfa91a80626fF3989D68f67F5b1DD12D;
+  IHevm hevm = IHevm(HEVM_ADDRESS);
+  Greeter public targetContract;
 
-  function echidna_greeterNeverEmpty() public view returns (bool) {
-    return keccak256(bytes(greeting)) != keccak256('');
+  constructor() {
+    targetContract = new Greeter('a', IERC20(address(1)));
   }
 
-  function echidna_onlyOwnerSetsGreeting() public returns (bool) {
-    // new greeting set, is the sender the owner?
-    try this.setGreeting('hello') {
-      if (msg.sender != OWNER) return false;
-      return true;
-    } catch {
-      // new greeting failed, is the sender not the owner?
-      if (msg.sender == OWNER) return false;
-      return true;
-    }
+  function checkGreeterNeverEmpty(string memory newGreeting) public {
+    (bool success,) = address(targetContract).call(abi.encodeCall(Greeter.setGreeting, newGreeting));
+
+    assert((success && keccak256(bytes(targetContract.greeting())) != keccak256(bytes(''))) || !success);
+  }
+
+  function checkOnlyOwnerSetsGreeting(address caller) public {
+    hevm.prank(caller);
+
+    (bool success,) = address(this).call(abi.encodeCall(Greeter.setGreeting, 'hello'));
+
+    assert((success && msg.sender == targetContract.OWNER()) || (!success && msg.sender != targetContract.OWNER()));
   }
 }
diff --git a/test/invariants/Greeter.symbolic.t.sol b/test/invariants/Greeter.symbolic.t.sol
index 25fe7303..0527373d 100644
--- a/test/invariants/Greeter.symbolic.t.sol
+++ b/test/invariants/Greeter.symbolic.t.sol
@@ -10,22 +10,32 @@ contract GreeterSymbolic is SymTest, Test {
   Greeter public targetContract;
 
   function setUp() public {
-    string memory _initialGreeting = svm.createString(64, 'initial greeting');
-    address _token = svm.createAddress('token');
+    string memory initialGreeting = svm.createString(64, 'initial greeting');
+    address token = svm.createAddress('token');
 
-    targetContract = new Greeter(_initialGreeting, IERC20(_token));
+    targetContract = new Greeter(initialGreeting, IERC20(token));
   }
 
-  function check_validState_greeterNeverEmpty(address caller, bytes4 selector) public {
+  function check_validState_greeterNeverEmpty(address caller) public {
     // Input conditions: any caller
     vm.prank(caller);
 
-    // Execution
-    (bool success,) = address(targetContract).call(gen_calldata(selector));
+    // Execution: Halmos cannot use a dynamic-sized array, iterate over multiple string lengths
+    for (uint256 i = 1; i < 3; i++) {
+      string memory greeting = svm.createString(i, 'greeting');
+      (bool success,) = address(targetContract).call(abi.encodeCall(Greeter.setGreeting, (greeting)));
+
+      // Output condition check
+      vm.assume(success); // discard failing calls
+      assert(keccak256(bytes(targetContract.greeting())) != keccak256(bytes('')));
+    }
+
+    // Add the empty string (bypass the non-empty check of svm.createString)
+    (bool success,) = address(targetContract).call(abi.encodeCall(Greeter.setGreeting, ('')));
 
     // Output condition check
     vm.assume(success); // discard failing calls
-    assert(keccak256(bytes(targetContract.greeting())) != keccak256(''));
+    assert(keccak256(bytes(targetContract.greeting())) != keccak256(bytes('')));
   }
 
   function check_setGreeting_onlyOwnerSetsGreeting(address caller) public {
@@ -41,15 +51,7 @@ contract GreeterSymbolic is SymTest, Test {
       assert(caller == targetContract.OWNER());
       assert(keccak256(bytes(targetContract.greeting())) == keccak256(bytes(newGreeting)));
     } else {
-      assert(caller != targetContract.OWNER() || keccak256(bytes(newGreeting)) == keccak256(''));
-    }
-  }
-
-  // either return a valid call to setGreeting or nothing (avoid halmos panicking on unknown contract call)
-  function gen_calldata(bytes4 selector) public view returns (bytes memory newCalldata) {
-    if (selector == targetContract.setGreeting.selector) {
-      string memory greeting = svm.createString(64, 'greeting');
-      newCalldata = abi.encodeWithSelector(selector, greeting);
+      assert(caller != targetContract.OWNER() || keccak256(bytes(newGreeting)) == keccak256(bytes('')));
     }
   }
 }