ghost-dao-contracts/src/types/Weaver.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

import {IWeaver} from "../interfaces/IWeaver.sol";

import {Hashes} from "../libraries/Hashes.sol";
import {Checkpoints} from "../libraries/Checkpoints.sol";

abstract contract Weaver is IWeaver {
    using Checkpoints for Checkpoints.Trace256;
    using Checkpoints for Checkpoints.Trace160;

    uint256 public constant DEPTH = 8;
    uint256 public constant SLOTS = 2 ** DEPTH;
    uint256 public constant ENTRIES = DEPTH * SLOTS;

    uint256 public override currentSession;
    uint256 public override startSession;
    address public previousWeaver;

    mapping(uint256 => mapping(uint256 => Checkpoints.Trace256)) internal _treeNodes;
    mapping(uint256 => mapping(uint256 => Checkpoints.Trace160)) internal _slotLenghts;
    mapping(uint256 => mapping(uint256 => bytes32[])) internal _slotValues;
    mapping(uint256 => uint256) internal _filledEntries;

    constructor(address _previousWeaver) {
        previousWeaver = _previousWeaver;
        if (_previousWeaver != address(0)) {
            uint256 newSession = IWeaver(_previousWeaver).currentSession() + 1;
            currentSession = newSession;
            startSession = newSession;
        }
    }

    function getSlotValues(
        uint256 globalIndex,
        uint256 session,
        uint256 atBlock
    ) external override view returns (bytes32[] memory) {
        if (session < startSession) {
            return IWeaver(previousWeaver).getSlotValues(globalIndex, session, atBlock);
        }

        uint256 slotIndex = globalIndex % SLOTS;
        // forge-lint: disable-next-line(unsafe-typecast)
        uint256 length = _slotLenghts[session][slotIndex].upperLookupRecent(uint96(atBlock));
        bytes32[] memory values = new bytes32[](length);

        uint256 i;
        for (; i < length; ) {
            values[i] = _slotValues[session][slotIndex][i];
            unchecked { ++i; }
        }

        return values;
    }

    function getRoot(uint256 session, uint256 atBlock) public override view returns (bytes32, uint256) {
        if (session < startSession) {
            return IWeaver(previousWeaver).getRoot(session, atBlock);
        }

        uint256 currentLevelCount = SLOTS >> 1;
        bytes32[] memory currentLevel = new bytes32[](currentLevelCount);

        uint256 i;
        for (; i < currentLevelCount;) {
            uint256 j = i << 1;
            bytes32 left = bytes32(_treeNodes[session][j].upperLookup(atBlock));
            bytes32 right = bytes32(_treeNodes[session][j | 1].upperLookup(atBlock));

            bytes32 leftHash = Hashes.efficientKeccak256(left);
            bytes32 rightHash = Hashes.efficientKeccak256(right);

            currentLevel[i] = Hashes.efficientKeccak256(leftHash, rightHash);
            unchecked { ++i; }
        }

        while (currentLevelCount > 1) {
            currentLevelCount >>= 1;
            i = 0;
            for (; i < currentLevelCount;) {
                uint256 j = i << 1;
                bytes32 left = currentLevel[j];
                bytes32 right = currentLevel[j | 1];

                currentLevel[i] = Hashes.efficientKeccak256(left, right);
                unchecked { ++i; }
            }
        }

        if (session < currentSession) {
            unchecked { ++session; }
        }

        return (currentLevel[0], session);
    }

    function getProof(
        uint256 globalIndex,
        uint256 session,
        uint256 atBlock
    ) external override view returns (bytes32[] memory) {
        if (session < startSession) {
            return IWeaver(previousWeaver).getProof(globalIndex, session, atBlock);
        }

        uint256 currentLevelCount = SLOTS;
        uint256 currentIndex = globalIndex % SLOTS;

        bytes32[] memory proof = new bytes32[](DEPTH);
        bytes32[] memory currentLevel = new bytes32[](currentLevelCount);
        uint256 i;

        for (; i < currentLevelCount;) {
            bytes32 leaf = bytes32(_treeNodes[session][i].upperLookup(atBlock));
            currentLevel[i] = Hashes.efficientKeccak256(leaf);
            unchecked { ++i; }
        }

        uint256 proofIndex;
        while (currentLevelCount > 1) {
            currentLevelCount >>= 1;
            uint256 siblingIndex = currentIndex ^ 1;

            proof[proofIndex] = currentLevel[siblingIndex];
            i = 0;

            for (; i < currentLevelCount;) {
                uint256 j = i << 1;
                bytes32 left = currentLevel[j];
                bytes32 right = currentLevel[j | 1];

                currentLevel[i] = Hashes.efficientKeccak256(left, right);
                unchecked { ++i; }
            }

            currentIndex >>= 1;
            unchecked { ++proofIndex; }
        }

        return proof;
    }

    function _insertTreeNode(bytes32 who, uint256 amount) internal returns (uint256 globalIndex) {
        uint256 session = currentSession;
        globalIndex = _filledEntries[session];
        if (globalIndex >= ENTRIES) {
            unchecked { ++session; }
            currentSession = session;
            globalIndex = 0;
        }

        uint256 slotIndex = globalIndex % SLOTS;
        uint256 valueIndex = globalIndex / SLOTS;
        uint160 length = uint160(_slotValues[session][slotIndex].length + 1);

        bytes32 preimage = _computeArgumentsHash(valueIndex, amount, who);
        bytes32 currHash = Hashes.efficientKeccak256(preimage);
        bytes32 slotHash = bytes32(_treeNodes[session][slotIndex].latest());

        slotHash = Hashes.efficientKeccak256(slotHash, currHash);

        _treeNodes[session][slotIndex].push(block.number, uint256(slotHash));
        _slotLenghts[session][slotIndex].push(uint96(block.number), length);
        _slotValues[session][slotIndex].push(preimage);

        unchecked { ++_filledEntries[session]; }
    }

    function _computeArgumentsHash(uint256 i, uint256 a, bytes32 r) internal pure returns (bytes32) {
        return Hashes.efficientKeccak256(Hashes.efficientKeccak256(bytes32(a), r), bytes32(i));
    }
}