HIGH severity: Effective price calculations can be affected by edge cases in `MatLib::sqrt512` and `MatLib::div512by256`. **Description:** `Math512Lib::sqrt512` implements a full-width integer Newton-Raphson square root. This hinges on the assumption that the initial guess is larger than the upper limb and fits within 256 bits such that the iteration is strictly monotonically decreasing, i.e. converges on the true...

Vulnerable code:
function sqrt512(uint256 x1, uint256 x0) internal pure returns (uint256 root) {
    if (x1 == 0) {
        return FixedPointMathLib.sqrt(x0);
    }
@>  root = 1 << (128 + (LibBit.fls(x1) / 2));
    uint256 last;
    do {
        last = root;
        // Because `floor(sqrt(UINT512_MAX)) = 2^256-1` and guesses converging towards the
        // correct result the result of all divisions is guaranteed to fit within 256 bits.
@>      (, root) = div512by256(x1, x0, root);
        root = (root + last) / 2;
    } while (root != last);
    return root;
}

/// @dev Computes `[x1 x0] / d`
function div512by256(uint256 x1, uint256 x0, uint256 d)
    internal
    pure
    returns (uint256 y1, uint256 y0)
{
    if (d == 0) revert DivisorZero();
    assembly {
        // Compute first "digit" of long division result
        y1 := div(x1, d)
        // We take the remainder to continue the long division
        let r1 := mod(x1, d)
        // We complete the long division by computing `y0 = [r1 x0] / d`. We use the "512 by
        // 256 division" logic from Solady's `fullMulDiv` (Credit under MIT license:
        // https://github.com/Vectorized/solady/blob/main/src/utils/FixedPointMathLib.sol)

        // We need to compute `[r1 x0] mod d = r1 * 2^256 + x0 = (r1 * 2^128) * 2^128 + x0`.
@>      let r := addmod(mulmod(shl(128, x1), shl(128, 1), d), x0, d)

        // Same math from Solady, reference `fullMulDiv` for explanation.
        let t := and(d, sub(0, d))
        d := div(d, t)
        let inv := xor(2, mul(3, d)) // inverse mod 2**4
        inv := mul(inv, sub(2, mul(d, inv))) // inverse mod 2**8
        inv := mul(inv, sub(2, mul(d, inv))) // inverse mod 2**16
        inv := mul(inv, sub(2, mul(d, inv))) // inverse mod 2**32
        inv := mul(inv, sub(2, mul(d, inv))) // inverse mod 2**64
        inv := mul(inv, sub(2, mul(d, inv))) // inverse mod 2**128
        // Edits vs Solady: `x0` replaces `z`, `r1` replaces `p1`, final 256-bit result stored in `y0`
        y0 :=
            mul(
@>              or(mul(sub(r1, gt(r, x0)), add(div(sub(0, t), t), 1)), div(sub(x0, r), t)),
                mul(sub(2, mul(d, inv)), inv) // inverse mod 2**256
            )
    }
}

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

import {console, Test} from "forge-std/Test.sol";
import {stdError} from "forge-std/StdError.sol";
import {LibBit} from "solady/src/utils/LibBit.sol";
import {FixedPointMathLib} from "solady/src/utils/FixedPointMathLib.sol";
import {Math512Lib} from "../../src/libraries/Math512Lib.sol";

contract Math512LibHarness {
    using Math512Lib for uint256;

    function fullMul(uint256 x, uint256 y) public pure returns (uint256 z1, uint256 z0) {
        return Math512Lib.fullMul(x, y);
    }

    function sqrt512(uint256 x1, uint256 x0) public pure returns (uint256) {
        return Math512Lib.sqrt512(x1, x0);
    }

    function div512by256(uint256 x1, uint256 x0, uint256 d) public pure returns (uint256 y1, uint256 y0) {
        return Math512Lib.div512by256(x1, x0, d);
    }
}

contract Math512LibTest is Math512LibHarness, Test {
    // BUG: when msb is odd in [129, 245], the quotient is too wide
    function test_oddMsbInitialRootArithmeticError(uint8 msb) public {
        uint256 x0;

        // Choose a msb that is odd and in [129, 245] such that LibBit.fls(x1) returns an odd index
        msb = uint8(129 + 2 * bound(msb, 0, (245 - 129) / 2));

        uint256 x1 = uint256(1) << msb;

        uint256 p = LibBit.fls(x1);
        assertEq(p, msb, "index not equal to msb");
        assertEq(p % 2, 1, "index not odd");

        // Rounds down when p is odd, so we have r0 = 2^(128 + 64) = 2^193
        uint256 r0 = uint256(1) << (128 + (p / 2));

        // If root <= x1, then floor(([x1 x0]) / root) >= 2^256, i.e. quotient too wide.
        if (r0 <= x1) console.log("initial root <= x1 when msb is odd -> unsafe: quotient too wide");

        vm.expectRevert(stdError.arithmeticError);
        uint256 root = this.sqrt512(x1, x0);
    }

    // BUG: when msb is odd in [247, 255], initial root computation overflows
    function test_oddMsbInitialRootOverflow(uint8 msb) public {
        uint256 x0;

        // ensure msb is odd and in [247, 255]
        msb = uint8(247 + 2 * bound(msb, 0, (255 - 247) / 2));
        uint256 x1 = uint256(1) << msb;

        uint256 r0 = uint256(1) << (128 + (LibBit.fls(x1) / 2));
        if (msb != 255) assertGt(r0, x1, "initial root not > x1");

        vm.expectRevert("DivisorZero()");
        uint256 root = this.sqrt512(x1, x0);
    }

    // BUG: when msb is 254 (max even < 256), initial root computation overflows
    function test_evenMsbInitialRootOverflow() public {
        uint256 x0;

        uint8 msb = 254; // max even < 256
        uint256 x1 = uint256(1) << msb;

        uint256 r0 = uint256(1) << (128 + (LibBit.fls(x1) / 2));
        assertGt(r0, x1, "initial root not > x1");

        vm.expectRevert(stdError.arithmeticError);
        uint256 root = this.sqrt512(x1, x0);
    }

    // NOTE: when msb is even in [0, 254), initial root computation is safe
    function test_evenMsbInitialRoot(uint8 msb) public {
        uint256 x0;

        vm.assume(msb % 2 == 0 && msb < 254);
        uint256 x1 = uint256(1) << msb;

        uint256 r0 = uint256(1) << (128 + (LibBit.fls(x1) / 2));
        assertGt(r0, x1, "initial root not > x1");

        uint256 root = this.sqrt512(x1, x0);
    }

    // BUG: when x1 <= d, the quotient fits in 256 bits and y1 should be zero; however,
    // there are cases in which the high digit is not zero, so it should not be ignored
    function test_discardHighDigit(uint256 x0, uint256 x1, uint256 d) public {
        // Only allow inputs that should result in a y1 of 0 (i.e. quotient fits in 256 bits)
        vm.assume(d != 0 && x1 <= d);

        (uint256 y1, ) = this.div512by256(x1, x0, d);

        assertEq(y1, 0, "y1 not zero -> quotient too wide");
    }

    // BUG: when x1 <= d, the quotient fits in 256 bits and y1 should be zero; however,
    // there are cases in which y0 does not collapse to FixedPointMathLib.fullMulDiv(r1, x0, d)
    function test_fullMulDivEquivalence(uint256 x0, uint256 x1, uint256 d) public {
        // Compute 512-bit product
        (uint256 p1, uint256 p0) = this.fullMul(x1, x0);

        // Ensure the high digit does not exceed d; otherwise, the quotient does not fit in 256 bits
        vm.assume(d != 0 && p1 <= d);

        (uint256 y1, uint256 y0) = this.div512by256(p1, p0, d);
        uint256 z = FixedPointMathLib.fullMulDiv(x1, x0, d);

        assertEq(y1, 0, "y1 not zero -> quotient too wide");
        assertEq(y0, z, "y0 not equal to z -> incorrect quotient");
    }
}