// SPDX-License-Identifier: MIT
pragma solidity 0.8.23;

/**
 * @title SSZ — the minimal Simple-Serialize merkleization primitives the Ethereum
 *        light client needs: SHA-256 chunk hashing, fixed-size merkleization,
 *        Merkle branch verification, and BeaconBlockHeader hash_tree_root.
 *
 * @dev   All hashing is SHA-256 (the beacon-chain hash), via the 0x02 precompile —
 *        NOT keccak. Little-endian integer serialization matches the consensus spec.
 */
library SSZ {
    /// @notice sha256(a || b) via the precompile.
    function hashPair(bytes32 a, bytes32 b) internal view returns (bytes32 r) {
        bytes memory input = new bytes(64);
        assembly {
            mstore(add(input, 32), a)
            mstore(add(input, 64), b)
        }
        r = _sha(input);
    }

    function _sha(bytes memory input) private view returns (bytes32 r) {
        bool ok;
        assembly {
            let out := mload(0x40)
            ok := staticcall(gas(), 0x02, add(input, 32), mload(input), out, 32)
            r := mload(out)
        }
        require(ok, "SSZ:sha256");
    }

    /// @notice uint64 -> a 32-byte SSZ chunk (little-endian, right-zero-padded).
    function uint64Chunk(uint64 v) internal pure returns (bytes32 out) {
        // little-endian bytes of v in the LOW 8 bytes of the (big-endian) word,
        // laid out at the start of the 32-byte chunk.
        uint256 le;
        uint256 vv = uint256(v);
        for (uint256 i = 0; i < 8; i++) {
            le |= ((vv >> (8 * i)) & 0xff) << (8 * (31 - i));
        }
        out = bytes32(le);
    }

    /// @notice hash_tree_root of a BeaconBlockHeader (5 fields, merkleized as 8 leaves).
    function beaconBlockHeaderRoot(
        uint64 slot,
        uint64 proposerIndex,
        bytes32 parentRoot,
        bytes32 stateRoot,
        bytes32 bodyRoot
    ) internal view returns (bytes32) {
        bytes32 zero = bytes32(0);
        bytes32 l01 = hashPair(uint64Chunk(slot), uint64Chunk(proposerIndex));
        bytes32 l23 = hashPair(parentRoot, stateRoot);
        bytes32 l45 = hashPair(bodyRoot, zero);
        bytes32 l67 = hashPair(zero, zero);
        bytes32 a = hashPair(l01, l23);
        bytes32 b = hashPair(l45, l67);
        return hashPair(a, b);
    }

    /// @notice is_valid_merkle_branch (consensus-spec). Folds `leaf` up `depth`
    ///         levels using `branch`, choosing sibling order by the bits of `index`,
    ///         and checks the result equals `root`.
    function isValidMerkleBranch(
        bytes32 leaf,
        bytes32[] memory branch,
        uint256 depth,
        uint256 index,
        bytes32 root
    ) internal view returns (bool) {
        require(branch.length >= depth, "SSZ:branch len");
        bytes32 value = leaf;
        for (uint256 i = 0; i < depth; i++) {
            if ((index >> i) & 1 == 1) {
                value = hashPair(branch[i], value);
            } else {
                value = hashPair(value, branch[i]);
            }
        }
        return value == root;
    }

    /// @notice Merkleize `leaves` (length must be a power of two) into a single root.
    function merkleize(bytes32[] memory leaves) internal view returns (bytes32) {
        uint256 n = leaves.length;
        require(n != 0 && (n & (n - 1)) == 0, "SSZ:not pow2");
        while (n > 1) {
            for (uint256 i = 0; i < n; i += 2) {
                leaves[i / 2] = hashPair(leaves[i], leaves[i + 1]);
            }
            n /= 2;
        }
        return leaves[0];
    }
}
