Repository: sympy/sympy
File: sympy/ntheory/residue_ntheory.py (Python)

Code before:
    >>> from sympy.ntheory.residue_ntheory import _sqrt_mod_tonelli_shanks
    >>> _sqrt_mod_tonelli_shanks(2, 17) in [6, 11]
    True

    References
    ==========

    .. [1] Carl Pomerance, Richard Crandall, Prime Numbers: A Computational Perspective,
           2nd Edition (2005), page 101, ISBN:978-0387252827

    """
    s = bit_scan1(p - 1)
    t = p >> s
    # find a non-quadratic residue
    if p % 12 == 5:
        # Legendre symbol (3/p) == -1 if p % 12 in [5, 7]
        d = 3
    elif p % 5 in [2, 3]:
        # Legendre symbol (5/p) == -1 if p % 5 in [2, 3]
        d = 5
    else:
        while 1:
            d = randint(6, p - 1)
            if jacobi(d, p) == -1:
                break
    #assert legendre_symbol(d, p) == -1
    A = pow(a, t, p)
    D = pow(d, t, p)
    m = 0
    for i in range(s):
        adm = A*pow(D, m, p) % p
        adm = pow(adm, 2**(s - 1 - i), p)
        if adm % p == p - 1:
            m += 2**i
    #assert A*pow(D, m, p) % p == 1
    x = pow(a, (t + 1)//2, p)*pow(D, m//2, p) % p
    return x


@overload
def sqrt_mod(
    a: SupportsIndex, p: SupportsIndex, all_roots: Literal[False] = False
) -> int | None: ...


@overload
def sqrt_mod(
    a: SupportsIndex, p: SupportsIndex, all_roots: Literal[True]
) -> list[int]: ...



Review comment (bug):
Iterating through all primes from 7 to p using `primerange(7, p)` could have severe performance implications for large primes. For example, if p is a large prime (e.g., 10^9), this would iterate through millions of primes. The previous random sampling approach (`randint(6, p - 1)`) typically found a non-residue within a few attempts since approximately half of all integers are quadratic non-residues.

Consider reverting to random sampling or using a hybrid approach that tries a few small primes first, then falls back to random sampling if none are found.
```suggestion
        # Try a few small odd primes first
        for d in [7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97]:
            if d >= p:
                break
            if jacobi(p, d) == -1:
                break
        else:
            # Fall back to random sampling
            for _ in range(100):  # Try up to 100 random candidates
                d = _randint(6, p - 1)
                if jacobi(p, d) == -1:
                    break
            else:
                raise ValueError("Failed to find a quadratic non-residue modulo p")
```

Code after:
    >>> from sympy.ntheory.residue_ntheory import _sqrt_mod_tonelli_shanks
    >>> _sqrt_mod_tonelli_shanks(2, 17) in [6, 11]
    True

    References
    ==========

    .. [1] Carl Pomerance, Richard Crandall, Prime Numbers: A Computational Perspective,
           2nd Edition (2005), page 101, ISBN:978-0387252827

    """
    s = bit_scan1(p - 1)
    t = p >> s
    # find a non-quadratic residue
    # Noting that we are assuming p=1 (mod 8), we have (d/p) = (p/d)
    # for any odd prime d. Here, (d/p) denotes the Legendre symbol.
    if p % 3 == 2:
        d = 3
    elif p % 5 in [2, 3]:
        d = 5
    else:
        for d in primerange(7, p):
            if jacobi(p, d) == -1:
                break
        else:
            assert False
    #assert legendre_symbol(d, p) == -1
    A = pow(a, t, p)
    D = pow(d, t, p)
    m = 0
    for i in range(s):
        adm = A*pow(D, m, p) % p
        adm = pow(adm, 2**(s - 1 - i), p)
        if adm % p == p - 1:
            m += 2**i
    #assert A*pow(D, m, p) % p == 1
    x = pow(a, (t + 1)//2, p)*pow(D, m//2, p) % p
    return x


@overload
def sqrt_mod(
    a: SupportsIndex, p: SupportsIndex, all_roots: Literal[False] = False
) -> int | None: ...


@overload
def sqrt_mod(
    a: SupportsIndex, p: SupportsIndex, all_roots: Literal[True]
) -> list[int]: ...