SmtLib updates

contracts/identitytreestore/IdentityTreeStore.sol

@@ -2,11 +2,12 @@
 pragma solidity 0.8.27;
 
 import {ReverseHashLib} from "../lib/ReverseHashLib.sol";
 import {PoseidonUnit2L, PoseidonUnit3L} from "../lib/Poseidon.sol";
 import {IState} from "../interfaces/IState.sol";
 import {IOnchainCredentialStatusResolver} from "../interfaces/IOnchainCredentialStatusResolver.sol";
 import {IRHSStorage} from "../interfaces/IRHSStorage.sol";
 import {Initializable} from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
+import {IHasher} from "../interfaces/IHasher.sol";
 
 error NodeNotFound();
 error InvalidStateNode();
@@ -32,7 +33,7 @@
     /**
      * @dev Version of contract
      */
-    string public constant VERSION = "1.1.0";
+    string public constant VERSION = "2.0.0";
 
     /**
      * @dev Max SMT depth for the CredentialStatus proof
@@ -59,6 +60,7 @@
     /// @custom:storage-location erc7201:iden3.storage.IdentityTreeStore.Main
     struct IdentityTreeStoreMainStorage {
         IState _state;
+        IHasher _hasher;
     }
 
     // keccak256(abi.encode(uint256(keccak256("iden3.storage.IdentityTreeStore.Main")) - 1)) & ~bytes32(uint256(0xff));
@@ -81,13 +83,22 @@
     /**
      * @dev Function to call first time for initialization of the proxy.
      * @param state The state contract address to be used to check state of the identities
+     * @param hasher The hasher to use in hashFunction
      **/
-    function initialize(address state) public initializer {
-        IdentityTreeStoreMainStorage storage $its = _getIdentityTreeStoreMainStorage();
-        ReverseHashLib.Data storage $rhl = _getReverseHashLibDataStorage();
+    function initialize(address state, IHasher hasher) public initializer {
+        _getIdentityTreeStoreMainStorage()._state = IState(state);
+        _initializeHasher(hasher);
+    }
 
-        $its._state = IState(state);
-        $rhl.hashFunction = _hashFunc;
+    /**
+     * @dev Initialize needed data
+     * @param hasher Hasher for SmtLib
+     */
+    function initializeHasher(IHasher hasher) external reinitializer(2) {
+        // Initialize in case the hasher has not been set yet
+        if (address(_getIdentityTreeStoreMainStorage()._hasher) == address(0)) {
+            _initializeHasher(hasher);
+        }
     }
 
     /**
@@ -98,6 +109,10 @@
         return _getReverseHashLibDataStorage().savePreimages(nodes);
     }
 
+    function getStateAddress() external view returns (IState) {
+        return _getIdentityTreeStoreMainStorage()._state;
+    }
+
     /**
      * @dev Returns a node by its key. Note that a node contains an array.
      * @param key The key of the node
@@ -241,13 +256,19 @@
         return NodeType.Unknown;
     }
 
-    function _hashFunc(uint256[] memory preimage) internal pure returns (uint256) {
+    function _hashFunc(uint256[] memory preimage) internal view returns (uint256) {
+        IdentityTreeStoreMainStorage storage $its = _getIdentityTreeStoreMainStorage();
         if (preimage.length == 2) {
-            return PoseidonUnit2L.poseidon([preimage[0], preimage[1]]);
+            return $its._hasher.hash2([preimage[0], preimage[1]]);
         }
         if (preimage.length == 3) {
-            return PoseidonUnit3L.poseidon([preimage[0], preimage[1], preimage[2]]);
+            return $its._hasher.hash3([preimage[0], preimage[1], preimage[2]]);
         }
         revert UnsupportedLength();
     }
+
+    function _initializeHasher(IHasher hasher) internal {
+        _getIdentityTreeStoreMainStorage()._hasher = hasher;
+        _getReverseHashLibDataStorage().hashFunction = _hashFunc;
+    }
 }

contracts/interfaces/IHasher.sol

@@ -0,0 +1,28 @@
+// SPDX-License-Identifier: GPL-3.0
+pragma solidity 0.8.27;
+
+/**
+ * @dev IHasher. Interface for generating hashes. Specifically used for Merkle Tree hashing.
+ */
+interface IHasher {
+    /**
+     * @dev hash1. hashes one uint256 parameter and returns the resulting hash as uint256.
+     * @param params The parameters array of size 1 to be hashed.
+     * @return The resulting hash as uint256.
+     */
+    function hash1(uint256[1] memory params) external pure returns (uint256);
+
+    /**
+     * @dev hash2. hashes two uint256 parameters and returns the resulting hash as uint256.
+     * @param params The parameters array of size 2 to be hashed.
+     * @return The resulting hash as uint256.
+     */
+    function hash2(uint256[2] memory params) external pure returns (uint256);
+
+    /**
+     * @dev hash3. hashes three uint256 parameters and returns the resulting hash as uint256.
+     * @param params The parameters array of size 3 to be hashed.
+     * @return The resulting hash as uint256.
+     */
+    function hash3(uint256[3] memory params) external pure returns (uint256);
+}

contracts/lib/IdentityBase.sol

@@ -6,6 +6,7 @@ import {IOnchainCredentialStatusResolver} from "../interfaces/IOnchainCredential
 import {IdentityLib} from "../lib/IdentityLib.sol";
 import {SmtLib} from "../lib/SmtLib.sol";
 import {ERC165} from "@openzeppelin/contracts/utils/introspection/ERC165.sol";
+import {IHasher} from "../interfaces/IHasher.sol";
 
 error IdentityIdMismatch();
 
@@ -46,12 +47,13 @@ abstract contract IdentityBase is IIdentifiable, IOnchainCredentialStatusResolve
      * @dev Initialization of IdentityLib library
      * @param _stateContractAddr - address of the State contract
      */
-    function initialize(address _stateContractAddr, bytes2 idType) public virtual {
+    function initialize(address _stateContractAddr, bytes2 idType, IHasher hasher) public virtual {
         _getIdentityBaseStorage().identity.initialize(
             _stateContractAddr,
             address(this),
             getSmtDepth(),
-            idType
+            idType,
+            hasher
         );
     }
 

contracts/lib/IdentityLib.sol

@@ -5,6 +5,7 @@ import {IState} from "../interfaces/IState.sol";
 import {SmtLib} from "../lib/SmtLib.sol";
 import {PoseidonUnit3L, PoseidonUnit4L} from "../lib/Poseidon.sol";
 import {GenesisUtils} from "../lib/GenesisUtils.sol";
+import {IHasher} from "../interfaces/IHasher.sol";
 
 error SMTDepthIsGreaterThanMaxAllowed();
 error IdTypeNotSupported();
@@ -86,7 +87,8 @@ library IdentityLib {
         address _stateContractAddr,
         address _identityAddr,
         uint256 depth,
-        bytes2 idType
+        bytes2 idType,
+        IHasher hasher
     ) external {
         if (depth > IDENTITY_MAX_SMT_DEPTH) {
             revert SMTDepthIsGreaterThanMaxAllowed();
@@ -96,9 +98,9 @@ library IdentityLib {
             revert IdTypeNotSupported();
         }
         self.isOldStateGenesis = true;
-        self.trees.claimsTree.initialize(depth);
-        self.trees.revocationsTree.initialize(depth);
-        self.trees.rootsTree.initialize(depth);
+        self.trees.claimsTree.initialize(depth, hasher);
+        self.trees.revocationsTree.initialize(depth, hasher);
+        self.trees.rootsTree.initialize(depth, hasher);
         self.id = GenesisUtils.calcIdFromEthAddress(idType, _identityAddr);
     }
 

contracts/lib/ReverseHashLib.sol

@@ -4,7 +4,7 @@ pragma solidity 0.8.27;
 library ReverseHashLib {
     struct Data {
         mapping(uint256 => uint256[]) hashesToPreimages;
-        function(uint256[] memory) pure returns (uint256) hashFunction;
+        function(uint256[] memory) view returns (uint256) hashFunction;
     }
 
     /**