# HG changeset patch
# User Seth Troisi <sethtroisi@google.com>
# Date 1780469697 25200
#      Tue Jun 02 23:54:57 2026 -0700
# Node ID 92708a163d2f4c009c0b55cf5a738a6dcf89229b
# Parent  cf985f9efb5d57a8029bfe4eadf5c652b8d918c8
Added mpz_perfect_square

Extract initial checks into mpn_perfect_square_trivial_p

diff -r cf985f9efb5d -r 92708a163d2f ChangeLog
--- a/ChangeLog	Wed Jun 03 00:07:28 2026 -0700
+++ b/ChangeLog	Tue Jun 02 23:54:57 2026 -0700
@@ -1,3 +1,8 @@
+2026-06-01  Seth Troisi <sethtroisi@google.com>
+	* mpz/perfpow.c (mpz_perfect_square): New function with root return.
+	* mpn/perfpow.c (mpn_perfect_square_trivial_p): Checks used before sqrt.
+	* tests/mpz/t-perfsqr.c: Add tests for mpz_perfect_square
+
 2026-04-04  Marc Glisse <marc.glisse@inria.fr>
 
 	* gmpxx.h (operator""): Remove deprecated space.
diff -r cf985f9efb5d -r 92708a163d2f doc/gmp.texi
--- a/doc/gmp.texi	Wed Jun 03 00:07:28 2026 -0700
+++ b/doc/gmp.texi	Tue Jun 02 23:54:57 2026 -0700
@@ -3602,6 +3602,16 @@
 be perfect squares.
 @end deftypefun
 
+@deftypefun int mpz_perfect_square (mpz_t @var{rop}, const mpz_t @var{op})
+@cindex Perfect square functions
+@cindex Root testing functions
+Return non-zero and sets @var{rop} to the square root if @var{op} is a perfect
+square, i.e., if the square root of @var{op} is an integer.  Under this
+definition both 0 and 1 are considered to be perfect squares and @var{rop} is
+set to @var{op} in both cases. If @var{op} is not a perfect square the value of
+@var{rop} may be clobbered.
+@end deftypefun
+
 
 @need 2000
 @node Number Theoretic Functions, Integer Comparisons, Integer Roots, Integer Functions
@@ -5635,8 +5645,7 @@
 case the return value is zero or non-zero according to whether the remainder
 would have been zero or non-zero.
 
-A return value of zero indicates a perfect square.  See also
-@code{mpn_perfect_square_p}.
+A return value of zero indicates a perfect square.  See also @code{mpn_perfect_square_p}.
 @end deftypefun
 
 @deftypefun size_t mpn_sizeinbase (const mp_limb_t *@var{xp}, mp_size_t @var{n}, int @var{base})
@@ -5719,8 +5728,12 @@
 
 @deftypefun int mpn_perfect_square_p (const mp_limb_t *@var{s1p}, mp_size_t @var{n})
 Return non-zero iff @{@var{s1p}, @var{n}@} is a perfect square.
-The most significant limb of the input @{@var{s1p}, @var{n}@} must be
-non-zero.
+The most significant limb of the input @{@var{s1p}, @var{n}@} must be non-zero.
+@end deftypefun
+
+@deftypefun int mpn_perfect_square_trivial_p (const mp_limb_t *@var{s1p}, mp_size_t @var{n})
+Return non-zero iff @{@var{s1p}, @var{n}@} passes preliminary perfect square checks.
+The most significant limb of the input @{@var{s1p}, @var{n}@} must be non-zero.
 @end deftypefun
 
 @deftypefun void mpn_and_n (mp_limb_t *@var{rp}, const mp_limb_t *@var{s1p}, const mp_limb_t *@var{s2p}, mp_size_t @var{n})
@@ -9250,9 +9263,10 @@
 A significant fraction of non-squares can be quickly identified by checking
 whether the input is a quadratic residue modulo small integers.
 
-@code{mpz_perfect_square_p} first tests the input mod 256, which means just
-examining the low byte.  Only 44 different values occur for squares mod 256,
-so 82.8% of inputs can be immediately identified as non-squares.
+@code{mpz_perfect_square}/@code{mpz_perfect_square_p} first tests the input
+mod 256, which means just examining the low byte.  Only 44 different values
+occur for squares mod 256, so 82.8% of inputs can be immediately identified
+as non-squares.
 
 On a 32-bit system similar tests are done mod 9, 5, 7, 13 and 17, for a total
 99.25% of inputs identified as non-squares.  On a 64-bit system 97 is tested
diff -r cf985f9efb5d -r 92708a163d2f doc/tasks.html
--- a/doc/tasks.html	Wed Jun 03 00:07:28 2026 -0700
+++ b/doc/tasks.html	Tue Jun 02 23:54:57 2026 -0700
@@ -579,10 +579,6 @@
      probably simplest to regard this as a compiler compatibility issue, and
      leave it to users or sysadmins to ensure application and library code is
      built the same.
-<li> <code>mpz_sqrt_if_perfect_square</code>: When
-     <code>mpz_perfect_square_p</code> does its tests it calculates a square
-     root and then discards it.  For some applications it might be useful to
-     return that root.  Suggested by Jason Moxham.
 <li> <code>mpz_get_ull</code>, <code>mpz_set_ull</code>,
      <code>mpz_get_sll</code>, <code>mpz_get_sll</code>: Conversions for
      <code>long long</code>.  These would aid interoperability, though a
diff -r cf985f9efb5d -r 92708a163d2f gmp-h.in
--- a/gmp-h.in	Wed Jun 03 00:07:28 2026 -0700
+++ b/gmp-h.in	Tue Jun 02 23:54:57 2026 -0700
@@ -972,6 +972,9 @@
 __GMP_DECLSPEC int mpz_perfect_square_p (mpz_srcptr) __GMP_ATTRIBUTE_PURE;
 #endif
 
+#define mpz_perfect_square __gmpz_perfect_square
+__GMP_DECLSPEC int mpz_perfect_square (mpz_ptr, mpz_srcptr);
+
 #define mpz_popcount __gmpz_popcount
 #if __GMP_INLINE_PROTOTYPES || defined (__GMP_FORCE_mpz_popcount)
 __GMP_DECLSPEC mp_bitcnt_t mpz_popcount (mpz_srcptr) __GMP_NOTHROW __GMP_ATTRIBUTE_PURE;
@@ -1570,6 +1573,9 @@
 #define mpn_perfect_square_p __MPN(perfect_square_p)
 __GMP_DECLSPEC int mpn_perfect_square_p (mp_srcptr, mp_size_t) __GMP_ATTRIBUTE_PURE;
 
+#define mpn_perfect_square_trivial_p __MPN(perfect_square_trivial_p)
+__GMP_DECLSPEC int mpn_perfect_square_trivial_p (mp_srcptr, mp_size_t) __GMP_ATTRIBUTE_PURE;
+
 #define mpn_perfect_power_p __MPN(perfect_power_p)
 __GMP_DECLSPEC int mpn_perfect_power_p (mp_srcptr, mp_size_t) __GMP_ATTRIBUTE_PURE;
 
diff -r cf985f9efb5d -r 92708a163d2f mpn/generic/perfsqr.c
--- a/mpn/generic/perfsqr.c	Wed Jun 03 00:07:28 2026 -0700
+++ b/mpn/generic/perfsqr.c	Tue Jun 02 23:54:57 2026 -0700
@@ -1,8 +1,12 @@
 /* mpn_perfect_square_p(u,usize) -- Return non-zero if U is a perfect square,
    zero otherwise.
 
-Copyright 1991, 1993, 1994, 1996, 1997, 2000-2002, 2005, 2012 Free Software
-Foundation, Inc.
+mpn_perfect_square_trivial_p(u, usize) -- Returns non-zero if U not trivially
+disqualified from being a perfect square checks. Returns zero if impossible
+to be a perfect square. 
+
+Copyright 1991, 1993, 1994, 1996, 1997, 2000-2002, 2005, 2012, 2026 Free
+Software Foundation, Inc.
 
 This file is part of the GNU MP Library.
 
@@ -174,14 +178,11 @@
       }								\
   } while (0)
 
-
 int
-mpn_perfect_square_p (mp_srcptr up, mp_size_t usize)
+mpn_perfect_square_trivial_p (mp_srcptr up, mp_size_t usize)
 {
   ASSERT (usize >= 1);
 
-  TRACE (gmp_printf ("mpn_perfect_square_p %Nd\n", up, usize));
-
   /* The first test excludes 212/256 (82.8%) of the perfect square candidates
      in O(1) time.  */
   {
@@ -217,6 +218,16 @@
      according to their residues modulo small primes (or powers of
      primes).  See perfsqr.h.  */
   PERFSQR_MOD_TEST (up, usize);
+  return 1;
+}
+
+int
+mpn_perfect_square_p (mp_srcptr up, mp_size_t usize)
+{
+  TRACE (gmp_printf ("mpn_perfect_square_p %Nd\n", up, usize));
+
+  if (! mpn_perfect_square_trivial_p (up, usize))
+    return 0;
 
 
   /* For the third and last test, we finally compute the square root,
diff -r cf985f9efb5d -r 92708a163d2f mpz/perfsqr.c
--- a/mpz/perfsqr.c	Wed Jun 03 00:07:28 2026 -0700
+++ b/mpz/perfsqr.c	Tue Jun 02 23:54:57 2026 -0700
@@ -1,7 +1,10 @@
 /* mpz_perfect_square_p(arg) -- Return non-zero if ARG is a perfect square,
    zero otherwise.
 
-Copyright 1991, 1993, 1994, 1996, 2000, 2001 Free Software Foundation, Inc.
+mpz_perfect_square(rop, arg) -- Return non-zero if ARG is a perfect square
+   sets rop to root, otherwise returns zero and leaves rop unchanged.
+
+Copyright 1991, 1993, 1994, 1996, 2000, 2001, 2026 Free Software Foundation, Inc.
 
 This file is part of the GNU MP Library.
 
@@ -32,3 +35,58 @@
 #define __GMP_FORCE_mpz_perfect_square_p 1
 
 #include "gmp-impl.h"
+
+int
+mpz_perfect_square (mpz_ptr rop, mpz_srcptr u)
+{
+  mp_size_t u_size = SIZ (u);
+
+  /* Handle <= 0 */
+  if (__GMP_UNLIKELY ( u_size <= 0))
+    {
+      /* Zero is a square with rop = 0 */
+      if (__GMP_UNLIKELY ( u_size == 0))
+      {
+        SIZ(rop) = 0;
+        return 1;
+      }
+      /* Negatives are never perfect squares. */
+      return 0;
+    }
+
+  if (! mpn_perfect_square_trivial_p (PTR (u), u_size))
+    return 0;
+
+  int rop_size = (SIZ (u) + 1) / 2;
+
+  /* mpn_sqrtrem doesn't allow sp == np */
+  if (rop == u)
+    {
+      TMP_DECL;
+      TMP_MARK;
+
+      mp_ptr rop_ptr = TMP_ALLOC_LIMBS (rop_size);
+      int res = ! mpn_sqrtrem (rop_ptr, NULL, PTR (u), u_size);
+
+      /* It might be better to always copy so the value in rop
+         is more determistic. */
+      if (res)
+        {
+          MPN_COPY (PTR(rop), rop_ptr, rop_size);
+          SIZ(rop) = rop_size;
+        }
+
+      TMP_FREE;
+      return res;
+    }
+
+  /* Make sure rop has space for the square root. */
+  MPZ_REALLOC(rop, rop_size);
+
+  int res = ! mpn_sqrtrem (PTR (rop), NULL, PTR (u), u_size);
+  if (res)
+    {
+      SIZ(rop) = rop_size;
+    }
+  return res;
+}
diff -r cf985f9efb5d -r 92708a163d2f tests/mpz/t-perfsqr.c
--- a/tests/mpz/t-perfsqr.c	Wed Jun 03 00:07:28 2026 -0700
+++ b/tests/mpz/t-perfsqr.c	Tue Jun 02 23:54:57 2026 -0700
@@ -1,4 +1,4 @@
-/* Test mpz_perfect_square_p.
+/* Test mpz_perfect_square and mpz_perfect_square_p.
 
 Copyright 2000-2002 Free Software Foundation, Inc.
 
@@ -26,8 +26,8 @@
 #include "mpn/perfsqr.h"
 
 
-/* check_modulo() exercises mpz_perfect_square_p on squares which cover each
-   possible quadratic residue to each divisor used within
+/* check_modulo() exercises mpz_perfect_square/mpz_perfect_square_p on squares
+   which cover each possible quadratic residue to each divisor used within
    mpn_perfect_square_p, ensuring those residues aren't incorrectly claimed
    to be non-residues.
 
@@ -45,11 +45,13 @@
   static const unsigned long  divisor[] = PERFSQR_DIVISORS;
   unsigned long  i, j;
 
-  mpz_t  alldiv, others, n;
+  mpz_t  alldiv, others, n, root, rop;
 
   mpz_init (alldiv);
   mpz_init (others);
+  mpz_init (root);
   mpz_init (n);
+  mpz_init (rop);
 
   /* product of all divisors */
   mpz_set_ui (alldiv, 1L);
@@ -67,71 +69,164 @@
       for (j = 1; j <= divisor[i]; j++)
         {
           /* square */
-          mpz_mul_ui (n, others, j);
-          mpz_mul (n, n, n);
+          mpz_mul_ui (root, others, j);
+          mpz_mul (n, root, root);
           if (! mpz_perfect_square_p (n))
             {
               printf ("mpz_perfect_square_p got 0, want 1\n");
               mpz_trace ("  n", n);
               abort ();
             }
+          if (! mpz_perfect_square (rop, n))
+            {
+              printf ("mpz_perfect_square got 0, want 1\n");
+              mpz_trace ("  n", n);
+              abort ();
+            }
+          if (! mpz_cmp (rop, n))
+            {
+              gmp_printf ("mpz_perfect_square rop %Zd, want %Zd\n", rop, n);
+              abort ();
+            }
         }
     }
 
   mpz_clear (alldiv);
   mpz_clear (others);
+  mpz_clear (root);
+  mpz_clear (n);
+  mpz_clear (rop);
+}
+
+/* Check negative, 0, and 1 */
+void
+check_edge_cases (void)
+{
+  mpz_t  n, root;
+  mpz_init (n);
+  mpz_init (root);
+
+  mpz_set_si(n, -4);
+  if (mpz_perfect_square_p (n))
+    {
+      printf ("mpz_perfect_square_p -4 not a square\n");
+      abort ();
+    }
+  if (mpz_perfect_square (root, n))
+    {
+      printf ("mpz_perfect_square -4 not a square\n");
+      abort ();
+    }
+
+  // 0 and 1 are both perfect squares with respective root 0 and 1.
+  for (unsigned int m = 0; m < 2; m++)
+    {
+      mpz_set_ui(n, m);
+      if (! mpz_perfect_square_p (n))
+        {
+          printf ("mpz_perfect_square_p %u should be a square\n", m);
+          abort ();
+        }
+      if (! mpz_perfect_square (root, n) || mpz_cmp_ui (root, m) != 0)
+        {
+          printf ("mpz_perfect_square %u should be a square with root = %u\n", m, m);
+          abort ();
+        }
+    }
+
+  mpz_clear (root);
   mpz_clear (n);
 }
 
-
 /* Exercise mpz_perfect_square_p compared to what mpz_sqrt says. */
 void
 check_sqrt (int reps)
 {
-  mpz_t x2, x2t, x;
+  mpz_t x2, x2t, x, rop;
   mp_size_t x2n;
   int res;
   int i;
-  /* int cnt = 0; */
+  int want;
+  int cnt = 0;
   gmp_randstate_ptr rands = RANDS;
   mpz_t bs;
 
   mpz_init (bs);
 
   mpz_init (x2);
+  mpz_init (x2t);
   mpz_init (x);
-  mpz_init (x2t);
+  mpz_init (rop);
 
   for (i = 0; i < reps; i++)
     {
       mpz_urandomb (bs, rands, 9);
       x2n = mpz_get_ui (bs);
       mpz_rrandomb (x2, rands, x2n);
-      /* mpz_out_str (stdout, -16, x2); puts (""); */
+
+      mpz_sqrt (x, x2);
+      mpz_mul (x2t, x, x);
+      want = mpz_cmp(x2, x2t) == 0;
 
       res = mpz_perfect_square_p (x2);
-      mpz_sqrt (x, x2);
-      mpz_mul (x2t, x, x);
-
-      if (res != (mpz_cmp (x2, x2t) == 0))
+      if (res != want)
         {
           printf    ("mpz_perfect_square_p and mpz_sqrt differ\n");
           mpz_trace ("   x  ", x);
           mpz_trace ("   x2 ", x2);
           mpz_trace ("   x2t", x2t);
           printf    ("   mpz_perfect_square_p %d\n", res);
-          printf    ("   mpz_sqrt             %d\n", mpz_cmp (x2, x2t) == 0);
+          printf    ("   mpz_sqrt             %d\n", want);
           abort ();
         }
 
-      /* cnt += res != 0; */
+      res = mpz_perfect_square (rop, x2);
+      if (res != want)
+        {
+          printf    ("mpz_perfect_square and mpz_sqrt differ\n");
+          mpz_trace ("   x  ", x);
+          mpz_trace ("   x2 ", x2);
+          mpz_trace ("   x2t", x2t);
+          printf    ("   mpz_perfect_square %d\n", res);
+          printf    ("   mpz_sqrt           %d\n", want);
+          abort ();
+        }
+      if (res && mpz_cmp(x, rop) != 0)
+        {
+          printf    ("mpz_perfect_square and mpz_sqrt differ\n");
+          mpz_trace ("   x  ", x);
+          mpz_trace ("   x2 ", x2);
+          mpz_trace ("   x2t", x2t);
+          mpz_trace ("   mpz_perfect_square rop", rop);
+          mpz_trace ("   mpz_sqrt", x);
+          abort ();
+        }
+      /* Check that same variable as input and output works */
+      mpz_set(rop, x2);
+      res = mpz_perfect_square (rop, rop);
+      if (res != want || (res && mpz_cmp(x, rop) != 0))
+        {
+          printf    ("mpz_perfect_square differ when output == input\n");
+          mpz_trace ("   x2 ", x2);
+          printf    ("   mpz_perfect_square %d\n", res);
+          mpz_trace ("   mpz_perfect_square rop", rop);
+          abort ();
+        }
+
+      cnt += res != 0;
     }
+
+  if (reps > 1000 && cnt == 0) {
+    printf("No perfect squares found in %d reps", reps);
+    abort();
+  }
   /* printf ("%d/%d perfect squares\n", cnt, reps); */
 
   mpz_clear (bs);
   mpz_clear (x2);
   mpz_clear (x);
   mpz_clear (x2t);
+  mpz_clear (rop);
 }
 
 
@@ -146,6 +241,7 @@
   if (argc == 2)
      reps = atoi (argv[1]);
 
+  check_edge_cases ();
   check_modulo ();
   check_sqrt (reps);