From f115dfca2e9f089e86d962ceb8004e1ac74439b5 Mon Sep 17 00:00:00 2001
From: Mukundan314 <mukundan314@gmail.com>
Date: Sat, 2 Nov 2024 22:35:20 +0530
Subject: [PATCH] Fix mod_sqrt for a = 1, p = 2

---
 pyrival/algebra/mod_sqrt.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyrival/algebra/mod_sqrt.py b/pyrival/algebra/mod_sqrt.py
index 38bc175..8c1840f 100644
--- a/pyrival/algebra/mod_sqrt.py
+++ b/pyrival/algebra/mod_sqrt.py
@@ -1,8 +1,8 @@
 def mod_sqrt(a, p):
     """returns x s.t. x**2 == a (mod p)"""
     a %= p
-    if a == 0:
-        return 0
+    if a < 2:
+        return a
     assert pow(a, (p - 1) // 2, p) == 1
     if p & 3 == 3:
         return pow(a, (p + 1) // 4, p)