improve code

Author: capossele

contracts/src/RiscZeroVerifierFallbackRouter.sol

@@ -16,34 +16,32 @@
 
 pragma solidity ^0.8.9;
 
-import {Ownable, Ownable2Step} from "openzeppelin/contracts/access/Ownable2Step.sol";
-
-import {IRiscZeroVerifier, Receipt} from "./IRiscZeroVerifier.sol";
+import {IRiscZeroVerifier} from "./IRiscZeroVerifier.sol";
 import {RiscZeroVerifierRouter} from "./RiscZeroVerifierRouter.sol";
 
 /// @notice Router for IRiscZeroVerifier, allowing multiple implementations to be accessible behind a single address
 ///         and a fallback to the canonical router.
 /// @dev Extends RiscZeroVerifierRouter to add fallback behavior.
 contract RiscZeroVerifierFallbackRouter is RiscZeroVerifierRouter {
     /// @notice The canonical RISC Zero verifier router used as fallback.
-    IRiscZeroVerifier public FALLBACK_ROUTER;
+    RiscZeroVerifierRouter public fallbackRouter;
 
-    constructor(address owner, IRiscZeroVerifier fallbackRouter) RiscZeroVerifierRouter(owner) {
-        require(address(fallbackRouter) != address(0), "Fallback router address cannot be zero");
-        FALLBACK_ROUTER = fallbackRouter;
+    constructor(address owner, RiscZeroVerifierRouter canonicalRouter) RiscZeroVerifierRouter(owner) {
+        require(address(canonicalRouter) != address(0), "Fallback router address cannot be zero");
+        fallbackRouter = canonicalRouter;
     }
 
     /// @notice Sets the canonical RISC Zero verifier router.
-    function setFallbackRouter(IRiscZeroVerifier verifier) external onlyOwner {
+    function setFallbackRouter(RiscZeroVerifierRouter verifier) external onlyOwner {
         if (address(verifier) == address(0)) {
             revert VerifierAddressZero();
         }
-        FALLBACK_ROUTER = verifier;
+        fallbackRouter = verifier;
     }
 
     /// @notice Gets the canonical RISC Zero verifier router.
     function getFallbackRouter() external view returns (IRiscZeroVerifier) {
-        return FALLBACK_ROUTER;
+        return fallbackRouter;
     }
 
     /// @notice Get the associated verifier, falling back to the canonical router if unset.
@@ -52,7 +50,7 @@ contract RiscZeroVerifierFallbackRouter is RiscZeroVerifierRouter {
         IRiscZeroVerifier verifier = verifiers[selector];
         // If the verifier is unset, fall back to the canonical router.
         if (verifier == UNSET) {
-            return FALLBACK_ROUTER;
+            return fallbackRouter.getVerifier(selector);
         }
         if (verifier == TOMBSTONE) {
             revert SelectorRemoved({selector: selector});

contracts/test/RiscZeroVerifierFallbackRouter.t.sol

@@ -57,7 +57,7 @@ contract RiscZeroVerifierEmergencyStopTest is Test {
 
     function setUp() external {
         fallbackRouter = new RiscZeroVerifierRouter(address(this));
-        verifierRouter = new RiscZeroVerifierFallbackRouter(address(this), IRiscZeroVerifier(address(fallbackRouter)));
+        verifierRouter = new RiscZeroVerifierFallbackRouter(address(this), fallbackRouter);
 
         verifierMockA = new RiscZeroMockVerifier(bytes4(0xFFFFFFFF));
         verifierMockB = new RiscZeroMockVerifier(bytes4(uint32(1)));
@@ -78,7 +78,7 @@ contract RiscZeroVerifierEmergencyStopTest is Test {
         assertEq(address(verifierRouter.getFallbackRouter()), address(fallbackRouter));
 
         RiscZeroVerifierRouter newFallbackRouter = new RiscZeroVerifierRouter(address(this));
-        verifierRouter.setFallbackRouter(IRiscZeroVerifier(address(newFallbackRouter)));
+        verifierRouter.setFallbackRouter(newFallbackRouter);
 
         assertEq(address(verifierRouter.getFallbackRouter()), address(newFallbackRouter));
     }