Compound-like voting extension for ERC1155.

MockERC1155votes.sol

// SPDX-License-Identifier: AGPL-3.0-only
pragma solidity >=0.8.0;

import "https://github.com/Rari-Capital/solmate/blob/main/src/tokens/ERC1155.sol";
import "https://github.com/Rari-Capital/solmate/blob/main/src/utils/SafeCastLib.sol";

/// @notice Compound-like voting extension for ERC1155.
/// @author z0r0z.eth
abstract contract ERC1155votes is ERC1155 {
    using SafeCastLib for uint256;

    /// -----------------------------------------------------------------------
    /// Events
    /// -----------------------------------------------------------------------

    event DelegateChanged(
        address indexed delegator,
        uint256 id,
        address indexed fromDelegate,
        address indexed toDelegate
    );

    event DelegateVotesChanged(
        address indexed delegate,
        uint256 id,
        uint256 previousBalance,
        uint256 newBalance
    );

    /// -----------------------------------------------------------------------
    /// Checkpoint Storage
    /// -----------------------------------------------------------------------
     
    mapping(address => mapping(uint256 => address)) private _delegates;

    mapping(address => mapping(uint256 => uint256)) public numCheckpoints;

    mapping(address => mapping(uint256 => mapping(uint256 => Checkpoint))) public checkpoints;
    
    struct Checkpoint {
        uint64 fromTimestamp;
        uint192 votes;
    }

    /// -----------------------------------------------------------------------
    /// Delegation Logic
    /// -----------------------------------------------------------------------

    function delegates(address account, uint256 id) public view returns (address) {
        address current = _delegates[account][id];

        return current == address(0) ? account : current;
    }

    function getCurrentVotes(address account, uint256 id) external view returns (uint256) {
        // this is safe from underflow because decrement only occurs if `nCheckpoints` is positive
        unchecked {
            uint256 nCheckpoints = numCheckpoints[account][id];

            return
                nCheckpoints != 0
                    ? checkpoints[account][nCheckpoints - 1][id].votes
                    : 0;
        }
    }

    function getPriorVotes(
        address account, 
        uint256 id,
        uint256 timestamp
    )
        external
        view
        returns (uint256)
    {
        require(block.timestamp > timestamp, "UNDETERMINED");

        uint256 nCheckpoints = numCheckpoints[account][id];

        if (nCheckpoints == 0) return 0;

        // this is safe from underflow because decrement only occurs if `nCheckpoints` is positive
        unchecked {
            if (
                checkpoints[account][nCheckpoints - 1][id].fromTimestamp <=
                timestamp
            ) return checkpoints[account][nCheckpoints - 1][id].votes;

            if (checkpoints[account][0][id].fromTimestamp > timestamp) return 0;

            uint256 lower;

            uint256 upper = nCheckpoints - 1;

            while (upper > lower) {
                uint256 center = upper - (upper - lower) / 2;

                Checkpoint memory cp = checkpoints[account][center][id];

                if (cp.fromTimestamp == timestamp) {
                    return cp.votes;
                } else if (cp.fromTimestamp < timestamp) {
                    lower = center;
                } else {
                    upper = center - 1;
                }
            }

            return checkpoints[account][lower][id].votes;
        }
    }

    function delegate(address account, uint256 id) external payable {
        address currentDelegate = delegates(msg.sender, id);

        _delegates[msg.sender][id] = account;

        _moveDelegates(currentDelegate, account, id, balanceOf[msg.sender][id]);

        emit DelegateChanged(msg.sender, id, currentDelegate, account);
    }

    function _moveDelegates(
        address srcRep,
        address dstRep,
        uint256 id,
        uint256 amount
    ) internal {
        if (srcRep != dstRep && amount != 0) {
            if (srcRep != address(0)) {
                uint256 srcRepNum = numCheckpoints[srcRep][id];

                uint256 srcRepOld = srcRepNum != 0
                    ? checkpoints[srcRep][srcRepNum - 1][id].votes
                    : 0;

                uint256 srcRepNew = srcRepOld - amount;

                _writeCheckpoint(srcRep, id, srcRepNum, srcRepOld, srcRepNew);
            }

            if (dstRep != address(0)) {
                uint256 dstRepNum = numCheckpoints[dstRep][id];

                uint256 dstRepOld = dstRepNum != 0
                    ? checkpoints[dstRep][dstRepNum - 1][id].votes
                    : 0;

                uint256 dstRepNew = dstRepOld + amount;

                _writeCheckpoint(dstRep, id, dstRepNum, dstRepOld, dstRepNew);
            }
        }
    }

    function _writeCheckpoint(
        address delegatee,
        uint256 id,
        uint256 nCheckpoints,
        uint256 oldVotes,
        uint256 newVotes
    ) internal {
        unchecked {
            // this is safe from underflow because decrement only occurs if `nCheckpoints` is positive
            if (
                nCheckpoints != 0 &&
                checkpoints[delegatee][nCheckpoints - 1][id].fromTimestamp ==
                block.timestamp
            ) {
                checkpoints[delegatee][nCheckpoints - 1][id].votes = newVotes.safeCastTo192();
            } else {
                checkpoints[delegatee][nCheckpoints][id] = Checkpoint(
                    block.timestamp.safeCastTo64(),
                    newVotes.safeCastTo192()
                );

                // cannot realistically overflow
                numCheckpoints[delegatee][id] = nCheckpoints + 1;
            }
        }

        emit DelegateVotesChanged(delegatee, id, oldVotes, newVotes);
    }
}

contract MockERC1155votes is ERC1155votes {
    function uri(uint256) public pure virtual override returns (string memory) {}

    function mint(
        address to,
        uint256 id,
        uint256 amount,
        bytes memory data
    ) public virtual {
        _mint(to, id, amount, data);

        _moveDelegates(address(0), delegates(to, id), id, amount);
    }

    function burn(
        address from,
        uint256 id,
        uint256 amount
    ) public virtual {
        _burn(from, id, amount);

        _moveDelegates(delegates(from, id), address(0), id, amount);
    }
}