// SPDX-License-Identifier: AGPL-3.0
pragma solidity ^0.8.30;

import { Mimc } from "./Mimc.sol";

/**
 * @title Library to perform SparseMerkleProof actions using the MiMC hashing algorithm
 * @author ConsenSys Software Inc.
 * @custom:security-contact security-report@linea.build
 */
library SparseMerkleProof {
  using Mimc for *;

  /**
   * The Account struct represents the state of the account including the storage root, nonce, balance and codesize
   * @dev This is mapped directly to the output of the storage proof
   */
  struct Account {
    uint64 nonce;
    uint256 balance;
    bytes32 storageRoot;
    bytes32 mimcCodeHash;
    bytes32 keccakCodeHash;
    uint64 codeSize;
  }

  /**
   * Represents the leaf structure in both account and storage tries
   * @dev This is mapped directly to the output of the storage proof
   */
  struct Leaf {
    uint256 prev;
    uint256 next;
    bytes32 hKey;
    bytes32 hValue;
  }

  /**
   * Thrown when expected bytes length is incorrect
   */
  error WrongBytesLength(uint256 expectedLength, uint256 bytesLength);

  /**
   * Thrown when the length of bytes is not in exactly 32 byte chunks
   */
  error LengthNotMod32();

  /**
   * Thrown when the leaf index is higher than the tree depth
   */
  error MaxTreeLeafIndexExceed();

  /**
   * Thrown when the length of the unformatted proof is not provided exactly as expected (UNFORMATTED_PROOF_LENGTH)
   */
  error WrongProofLength(uint256 expectedLength, uint256 actualLength);

  uint256 internal constant TREE_DEPTH = 40;
  uint256 internal constant UNFORMATTED_PROOF_LENGTH = 42;
  bytes32 internal constant ZERO_HASH = 0x0;
  uint256 internal constant MAX_TREE_LEAF_INDEX = 2 ** TREE_DEPTH - 1;

  /**
   * @notice Formats input, computes root and returns true if it matches the provided merkle root
   * @param _rawProof Raw sparse merkle tree proof
   * @param _leafIndex Index of the leaf
   * @param _root Sparse merkle root
   * @return If the computed merkle root matches the provided one
   */
  function verifyProof(bytes[] calldata _rawProof, uint256 _leafIndex, bytes32 _root) external pure returns (bool) {
    if (_rawProof.length != UNFORMATTED_PROOF_LENGTH) {
      revert WrongProofLength(UNFORMATTED_PROOF_LENGTH, _rawProof.length);
    }

    (bytes32 nextFreeNode, bytes32 leafHash, bytes32[] memory proof) = _formatProof(_rawProof);
    return _verify(proof, leafHash, _leafIndex, _root, nextFreeNode);
  }

  /**
   * @notice Hash a value using MIMC hash
   * @param _input Value to hash
   * @return bytes32 Mimc hash
   */
  function mimcHash(bytes calldata _input) external pure returns (bytes32) {
    return Mimc.hash(_input);
  }

  /**
   * @notice Get leaf
   * @param _encodedLeaf Encoded leaf bytes (prev, next, hKey, hValue)
   * @return Leaf Formatted leaf struct
   */
  function getLeaf(bytes calldata _encodedLeaf) external pure returns (Leaf memory) {
    return _parseLeaf(_encodedLeaf);
  }

  /**
   * @notice Get account
   * @param _encodedAccountValue Encoded account value bytes (nonce, balance, storageRoot, mimcCodeHash, keccakCodeHash, codeSize)
   * @return Account Formatted account struct
   */
  function getAccount(bytes calldata _encodedAccountValue) external pure returns (Account memory) {
    return _parseAccount(_encodedAccountValue);
  }

  /**
   * @notice Hash account value
   * @param _value Encoded account value bytes (nonce, balance, storageRoot, mimcCodeHash, keccakCodeHash, codeSize)
   * @return bytes32 Account value hash
   */
  function hashAccountValue(bytes calldata _value) external pure returns (bytes32) {
    Account memory account = _parseAccount(_value);
    (bytes32 msb, bytes32 lsb) = _splitBytes32(account.keccakCodeHash);
    return
      Mimc.hash(
        abi.encode(
          account.nonce,
          account.balance,
          account.storageRoot,
          account.mimcCodeHash,
          lsb,
          msb,
          account.codeSize
        )
      );
  }

  /**
   * @notice Hash storage value
   * @param _value Encoded storage value bytes
   * @return bytes32 Storage value hash
   */
  function hashStorageValue(bytes32 _value) external pure returns (bytes32) {
    (bytes32 msb, bytes32 lsb) = _splitBytes32(_value);
    return Mimc.hash(abi.encodePacked(lsb, msb));
  }

  /**
   * @notice Parse leaf value
   * @param _encodedLeaf Encoded leaf bytes (prev, next, hKey, hValue)
   * @return {Leaf} Formatted leaf struct
   */
  function _parseLeaf(bytes calldata _encodedLeaf) private pure returns (Leaf memory) {
    if (_encodedLeaf.length != 128) {
      revert WrongBytesLength(128, _encodedLeaf.length);
    }
    return abi.decode(_encodedLeaf, (Leaf));
  }

  /**
   * @notice Parse account value
   * @param _value Encoded account value bytes (nonce, balance, storageRoot, mimcCodeHash, keccakCodeHash, codeSize)
   * @return {Account} Formatted account struct
   */
  function _parseAccount(bytes calldata _value) private pure returns (Account memory) {
    if (_value.length != 192) {
      revert WrongBytesLength(192, _value.length);
    }
    return abi.decode(_value, (Account));
  }

  /**
   * @notice Split bytes32 into two bytes32 taking most significant bits and least significant bits
   * @param _b bytes to split
   * @return msb Most significant bits
   * @return lsb Least significant bits
   */
  function _splitBytes32(bytes32 _b) private pure returns (bytes32 msb, bytes32 lsb) {
    assembly {
      msb := shr(128, _b)
      lsb := and(_b, 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF)
    }
  }

  /**
   * @notice Format proof
   * @param _rawProof Non formatted proof array
   * @return (bytes32, bytes32, bytes32[]) NextFreeNode, leafHash and formatted proof array
   */
  function _formatProof(bytes[] calldata _rawProof) private pure returns (bytes32, bytes32, bytes32[] memory) {
    uint256 rawProofLength = _rawProof.length;
    uint256 formattedProofLength = rawProofLength - 2;

    bytes32[] memory proof = new bytes32[](formattedProofLength);

    if (_rawProof[0].length != 0x40) {
      revert WrongBytesLength(0x40, _rawProof[0].length);
    }

    bytes32 nextFreeNode = bytes32(_rawProof[0][:32]);
    bytes32 leafHash = Mimc.hash(_rawProof[rawProofLength - 1]);

    for (uint256 i = 1; i < formattedProofLength; ) {
      proof[formattedProofLength - i] = Mimc.hash(_rawProof[i]);
      unchecked {
        ++i;
      }
    }

    // If the sibling leaf (_rawProof[formattedProofLength]) is equal to zero bytes we don't hash it
    if (_isZeroBytes(_rawProof[formattedProofLength])) {
      proof[0] = ZERO_HASH;
    } else {
      proof[0] = Mimc.hash(_rawProof[formattedProofLength]);
    }

    return (nextFreeNode, leafHash, proof);
  }

  /**
   * @notice Check if bytes contain only zero byte elements
   * @param _data Bytes to be checked
   * @return isZeroBytes true if bytes contain only zero byte elements, false otherwise
   */
  function _isZeroBytes(bytes calldata _data) private pure returns (bool isZeroBytes) {
    if (_data.length % 0x20 != 0) {
      revert LengthNotMod32();
    }

    isZeroBytes = true;
    assembly {
      let dataStart := _data.offset

      for {
        let currentPtr := dataStart
      } lt(currentPtr, add(dataStart, _data.length)) {
        currentPtr := add(currentPtr, 0x20)
      } {
        let dataWord := calldataload(currentPtr)

        if eq(iszero(dataWord), 0) {
          isZeroBytes := 0
          break
        }
      }
    }
  }

  /**
   * @notice Computes merkle root from proof and compares it to the provided root
   * @param _proof Sparse merkle tree proof
   * @param _leafHash Leaf hash
   * @param _leafIndex Index of the leaf
   * @param _root Sparse merkle root
   * @param _nextFreeNode Next free node
   * @return If the computed merkle root matches the provided one
   */
  function _verify(
    bytes32[] memory _proof,
    bytes32 _leafHash,
    uint256 _leafIndex,
    bytes32 _root,
    bytes32 _nextFreeNode
  ) private pure returns (bool) {
    bytes32 computedHash = _leafHash;
    uint256 currentIndex = _leafIndex;

    if (_leafIndex > MAX_TREE_LEAF_INDEX) {
      revert MaxTreeLeafIndexExceed();
    }

    for (uint256 height; height < TREE_DEPTH; ++height) {
      if ((currentIndex >> height) & 1 == 1) computedHash = Mimc.hash(abi.encodePacked(_proof[height], computedHash));
      else computedHash = Mimc.hash(abi.encodePacked(computedHash, _proof[height]));
    }

    return Mimc.hash(abi.encodePacked(_nextFreeNode, computedHash)) == _root;
  }
}
