Address comments

quantumlib · Oct 16, 2024 · d8ba6d3 · d8ba6d3
1 parent 5828eb7
commit d8ba6d3
Show file tree

Hide file tree

Showing 4 changed files with 67 additions and 90 deletions.
diff --git a/qualtran/_infra/data_types.py b/qualtran/_infra/data_types.py
@@ -779,6 +779,7 @@ class QMontgomeryUInt(QDType):
         We follow Montgomery form as described in the above paper; namely, r = 2^bitsize.
     """
 
+    # TODO(https://github.com/quantumlib/Qualtran/issues/1471): Add modulus p as a class member.
     bitsize: SymbolicInt
 
     @property

diff --git a/qualtran/_infra/data_types_test.py b/qualtran/_infra/data_types_test.py
@@ -139,16 +139,26 @@ def test_qmontgomeryuint():
 @pytest.mark.parametrize('val', [1, 5, 7, 9])
 def test_qmontgomeryuint_operations(val, p):
     qmontgomeryuint_8 = QMontgomeryUInt(8)
+    # Convert value to montgomery form and get the modular inverse.
     val_m = qmontgomeryuint_8.uint_to_montgomery(val, p)
     mod_inv = qmontgomeryuint_8.montgomery_inverse(val_m, p)
-    assert qmontgomeryuint_8.montgomery_product(val_m, mod_inv, p) == 1
+
+    # Calculate the product in montgomery form and convert back to normal form for assertion.
+    assert (
+        qmontgomeryuint_8.montgomery_to_uint(
+            qmontgomeryuint_8.montgomery_product(val_m, mod_inv, p), p
+        )
+        == 1
+    )
 
 
 @pytest.mark.parametrize('p', [13, 17, 29])
 @pytest.mark.parametrize('val', [1, 5, 7, 9])
 def test_qmontgomeryuint_conversions(val, p):
     qmontgomeryuint_8 = QMontgomeryUInt(8)
-    assert val == qmontgomeryuint_8.montgomery_to_uint(qmontgomeryuint_8.uint_to_montgomery(val, p), p)
+    assert val == qmontgomeryuint_8.montgomery_to_uint(
+        qmontgomeryuint_8.uint_to_montgomery(val, p), p
+    )
 
 
 def test_qgf():

diff --git a/qualtran/bloqs/factoring/ecc/ec_add.py b/qualtran/bloqs/factoring/ecc/ec_add.py
@@ -251,10 +251,7 @@ def on_classical_vals(
                 f1 = 0
             else:
                 lam = QMontgomeryUInt(self.n).montgomery_product(
-                    y,
-                    QMontgomeryUInt(self.n).montgomery_inverse(x, self.n, self.mod),
-                    self.n,
-                    self.mod,
+                    int(y), QMontgomeryUInt(self.n).montgomery_inverse(int(x), self.mod), self.mod
                 )
                 # TODO(https://github.com/quantumlib/Qualtran/issues/1461): Fix bug in circuit
                 # which flips f1 when lam and lam_r are equal.
@@ -540,9 +537,11 @@ def signature(self) -> 'Signature':
     def on_classical_vals(
         self, x: 'ClassicalValT', y: 'ClassicalValT', lam: 'ClassicalValT'
     ) -> Dict[str, 'ClassicalValT']:
-        x = (x - QMontgomeryUInt(self.n).montgomery_product(lam, lam, self.n, self.mod)) % self.mod
+        x = (
+            x - QMontgomeryUInt(self.n).montgomery_product(int(lam), int(lam), self.mod)
+        ) % self.mod
         if lam > 0:
-            y = QMontgomeryUInt(self.n).montgomery_product(x, lam, self.n, self.mod)
+            y = QMontgomeryUInt(self.n).montgomery_product(int(x), int(lam), self.mod)
         return {'x': x, 'y': y, 'lam': lam}
 
     def build_composite_bloq(
@@ -1071,29 +1070,29 @@ def build_composite_bloq(
 
     def on_classical_vals(self, a, b, x, y, lam_r) -> Dict[str, Union['ClassicalValT', sympy.Expr]]:
         curve_a = (
-            QMontgomeryUInt(self.n).montgomery_to_uint(lam_r, self.n, self.mod)
+            QMontgomeryUInt(self.n).montgomery_to_uint(lam_r, self.mod)
             * 2
-            * QMontgomeryUInt(self.n).montgomery_to_uint(b, self.n, self.mod)
-            - (3 * QMontgomeryUInt(self.n).montgomery_to_uint(a, self.n, self.mod) ** 2)
+            * QMontgomeryUInt(self.n).montgomery_to_uint(b, self.mod)
+            - (3 * QMontgomeryUInt(self.n).montgomery_to_uint(a, self.mod) ** 2)
         ) % self.mod
         p1 = ECPoint(
-            QMontgomeryUInt(self.n).montgomery_to_uint(a, self.n, self.mod),
-            QMontgomeryUInt(self.n).montgomery_to_uint(b, self.n, self.mod),
+            QMontgomeryUInt(self.n).montgomery_to_uint(a, self.mod),
+            QMontgomeryUInt(self.n).montgomery_to_uint(b, self.mod),
             mod=self.mod,
             curve_a=curve_a,
         )
         p2 = ECPoint(
-            QMontgomeryUInt(self.n).montgomery_to_uint(x, self.n, self.mod),
-            QMontgomeryUInt(self.n).montgomery_to_uint(y, self.n, self.mod),
+            QMontgomeryUInt(self.n).montgomery_to_uint(x, self.mod),
+            QMontgomeryUInt(self.n).montgomery_to_uint(y, self.mod),
             mod=self.mod,
             curve_a=curve_a,
         )
         result = p1 + p2
         return {
             'a': a,
             'b': b,
-            'x': QMontgomeryUInt(self.n).uint_to_montgomery(result.x, self.n, self.mod),
-            'y': QMontgomeryUInt(self.n).uint_to_montgomery(result.y, self.n, self.mod),
+            'x': QMontgomeryUInt(self.n).uint_to_montgomery(result.x, self.mod),
+            'y': QMontgomeryUInt(self.n).uint_to_montgomery(result.y, self.mod),
             'lam_r': lam_r,
         }
 

diff --git a/qualtran/bloqs/factoring/ecc/ec_add_test.py b/qualtran/bloqs/factoring/ecc/ec_add_test.py
@@ -33,23 +33,22 @@
 from qualtran.resource_counting.generalizers import ignore_alloc_free, ignore_split_join
 
 
-@pytest.mark.slow
 @pytest.mark.parametrize(
-    ['n', 'm'], [(n, m) for n in range(7, 10) for m in range(1, n + 1) if n % m == 0]
+    ['n', 'm'], [(n, m) for n in range(7, 8) for m in range(1, n + 1) if n % m == 0]
 )
-@pytest.mark.parametrize('a,b', [(15, 13), (2, 10), (8, 3), (0, 0)])
-@pytest.mark.parametrize('x,y', [(15, 13), (2, 10), (8, 3), (0, 0)])
+@pytest.mark.parametrize('a,b', [(15, 13), (2, 10)])
+@pytest.mark.parametrize('x,y', [(15, 13), (0, 0)])
 def test_ec_add_steps_classical_fast(n, m, a, b, x, y):
     p = 17
     lam_num = (3 * a**2) % p
     lam_denom = (2 * b) % p
     lam_r = 0 if b == 0 else (lam_num * pow(lam_denom, -1, mod=p)) % p
 
-    a = QMontgomeryUInt(n).uint_to_montgomery(a, n, p)
-    b = QMontgomeryUInt(n).uint_to_montgomery(b, n, p)
-    x = QMontgomeryUInt(n).uint_to_montgomery(x, n, p)
-    y = QMontgomeryUInt(n).uint_to_montgomery(y, n, p)
-    lam_r = QMontgomeryUInt(n).uint_to_montgomery(lam_r, n, p) if lam_r != 0 else p
+    a = QMontgomeryUInt(n).uint_to_montgomery(a, p)
+    b = QMontgomeryUInt(n).uint_to_montgomery(b, p)
+    x = QMontgomeryUInt(n).uint_to_montgomery(x, p)
+    y = QMontgomeryUInt(n).uint_to_montgomery(y, p)
+    lam_r = QMontgomeryUInt(n).uint_to_montgomery(lam_r, p) if lam_r != 0 else p
 
     bloq = _ECAddStepOne(n=n, mod=p)
     ret1 = bloq.call_classically(a=a, b=b, x=x, y=y)
@@ -159,7 +158,7 @@ def test_ec_add_steps_classical_fast(n, m, a, b, x, y):
 
 @pytest.mark.slow
 @pytest.mark.parametrize(
-    ['n', 'm'], [(n, m) for n in range(7, 10) for m in range(1, n + 1) if n % m == 0]
+    ['n', 'm'], [(n, m) for n in range(7, 9) for m in range(1, n + 1) if n % m == 0]
 )
 @pytest.mark.parametrize(
     'a,b',
@@ -178,34 +177,18 @@ def test_ec_add_steps_classical_fast(n, m, a, b, x, y):
         (0, 0),
     ],
 )
-@pytest.mark.parametrize(
-    'x,y',
-    [
-        (15, 13),
-        (2, 10),
-        (8, 3),
-        (12, 1),
-        (6, 6),
-        (5, 8),
-        (10, 15),
-        (1, 12),
-        (3, 0),
-        (1, 5),
-        (10, 2),
-        (0, 0),
-    ],
-)
+@pytest.mark.parametrize('x,y', [(15, 13), (5, 8), (10, 15), (1, 12), (3, 0), (1, 5), (10, 2)])
 def test_ec_add_steps_classical(n, m, a, b, x, y):
     p = 17
     lam_num = (3 * a**2) % p
     lam_denom = (2 * b) % p
     lam_r = 0 if b == 0 else (lam_num * pow(lam_denom, -1, mod=p)) % p
 
-    a = QMontgomeryUInt(n).uint_to_montgomery(a, n, p)
-    b = QMontgomeryUInt(n).uint_to_montgomery(b, n, p)
-    x = QMontgomeryUInt(n).uint_to_montgomery(x, n, p)
-    y = QMontgomeryUInt(n).uint_to_montgomery(y, n, p)
-    lam_r = QMontgomeryUInt(n).uint_to_montgomery(lam_r, n, p) if lam_r != 0 else p
+    a = QMontgomeryUInt(n).uint_to_montgomery(a, p)
+    b = QMontgomeryUInt(n).uint_to_montgomery(b, p)
+    x = QMontgomeryUInt(n).uint_to_montgomery(x, p)
+    y = QMontgomeryUInt(n).uint_to_montgomery(y, p)
+    lam_r = QMontgomeryUInt(n).uint_to_montgomery(lam_r, p) if lam_r != 0 else p
 
     bloq = _ECAddStepOne(n=n, mod=p)
     ret1 = bloq.call_classically(a=a, b=b, x=x, y=y)
@@ -314,36 +297,36 @@ def test_ec_add_steps_classical(n, m, a, b, x, y):
 
 
 @pytest.mark.parametrize(
-    ['n', 'm'], [(n, m) for n in range(7, 10) for m in range(1, n + 1) if n % m == 0]
+    ['n', 'm'], [(n, m) for n in range(7, 8) for m in range(1, n + 1) if n % m == 0]
 )
-@pytest.mark.parametrize('a,b', [(15, 13), (2, 10), (8, 3), (0, 0)])
-@pytest.mark.parametrize('x,y', [(15, 13), (2, 10), (8, 3), (0, 0)])
+@pytest.mark.parametrize('a,b', [(15, 13), (2, 10)])
+@pytest.mark.parametrize('x,y', [(15, 13), (0, 0)])
 def test_ec_add_classical_fast(n, m, a, b, x, y):
     p = 17
     bloq = ECAdd(n=n, mod=p, window_size=m)
     lam_num = (3 * a**2) % p
     lam_denom = (2 * b) % p
     lam_r = p if b == 0 else (lam_num * pow(lam_denom, -1, mod=p)) % p
     ret1 = bloq.call_classically(
-        a=QMontgomeryUInt(n).uint_to_montgomery(a, n, p),
-        b=QMontgomeryUInt(n).uint_to_montgomery(b, n, p),
-        x=QMontgomeryUInt(n).uint_to_montgomery(x, n, p),
-        y=QMontgomeryUInt(n).uint_to_montgomery(y, n, p),
-        lam_r=QMontgomeryUInt(n).uint_to_montgomery(lam_r, n, p),
+        a=QMontgomeryUInt(n).uint_to_montgomery(a, p),
+        b=QMontgomeryUInt(n).uint_to_montgomery(b, p),
+        x=QMontgomeryUInt(n).uint_to_montgomery(x, p),
+        y=QMontgomeryUInt(n).uint_to_montgomery(y, p),
+        lam_r=QMontgomeryUInt(n).uint_to_montgomery(lam_r, p),
     )
     ret2 = bloq.decompose_bloq().call_classically(
-        a=QMontgomeryUInt(n).uint_to_montgomery(a, n, p),
-        b=QMontgomeryUInt(n).uint_to_montgomery(b, n, p),
-        x=QMontgomeryUInt(n).uint_to_montgomery(x, n, p),
-        y=QMontgomeryUInt(n).uint_to_montgomery(y, n, p),
-        lam_r=QMontgomeryUInt(n).uint_to_montgomery(lam_r, n, p),
+        a=QMontgomeryUInt(n).uint_to_montgomery(a, p),
+        b=QMontgomeryUInt(n).uint_to_montgomery(b, p),
+        x=QMontgomeryUInt(n).uint_to_montgomery(x, p),
+        y=QMontgomeryUInt(n).uint_to_montgomery(y, p),
+        lam_r=QMontgomeryUInt(n).uint_to_montgomery(lam_r, p),
     )
     assert ret1 == ret2
 
 
 @pytest.mark.slow
 @pytest.mark.parametrize(
-    ['n', 'm'], [(n, m) for n in range(7, 10) for m in range(1, n + 1) if n % m == 0]
+    ['n', 'm'], [(n, m) for n in range(7, 9) for m in range(1, n + 1) if n % m == 0]
 )
 @pytest.mark.parametrize(
     'a,b',
@@ -362,42 +345,26 @@ def test_ec_add_classical_fast(n, m, a, b, x, y):
         (0, 0),
     ],
 )
-@pytest.mark.parametrize(
-    'x,y',
-    [
-        (15, 13),
-        (2, 10),
-        (8, 3),
-        (12, 1),
-        (6, 6),
-        (5, 8),
-        (10, 15),
-        (1, 12),
-        (3, 0),
-        (1, 5),
-        (10, 2),
-        (0, 0),
-    ],
-)
+@pytest.mark.parametrize('x,y', [(15, 13), (5, 8), (10, 15), (1, 12), (3, 0), (1, 5), (10, 2)])
 def test_ec_add_classical(n, m, a, b, x, y):
     p = 17
     bloq = ECAdd(n=n, mod=p, window_size=m)
     lam_num = (3 * a**2) % p
     lam_denom = (2 * b) % p
     lam_r = p if b == 0 else (lam_num * pow(lam_denom, -1, mod=p)) % p
     ret1 = bloq.call_classically(
-        a=QMontgomeryUInt(n).uint_to_montgomery(a, n, p),
-        b=QMontgomeryUInt(n).uint_to_montgomery(b, n, p),
-        x=QMontgomeryUInt(n).uint_to_montgomery(x, n, p),
-        y=QMontgomeryUInt(n).uint_to_montgomery(y, n, p),
-        lam_r=QMontgomeryUInt(n).uint_to_montgomery(lam_r, n, p),
+        a=QMontgomeryUInt(n).uint_to_montgomery(a, p),
+        b=QMontgomeryUInt(n).uint_to_montgomery(b, p),
+        x=QMontgomeryUInt(n).uint_to_montgomery(x, p),
+        y=QMontgomeryUInt(n).uint_to_montgomery(y, p),
+        lam_r=QMontgomeryUInt(n).uint_to_montgomery(lam_r, p),
     )
     ret2 = bloq.decompose_bloq().call_classically(
-        a=QMontgomeryUInt(n).uint_to_montgomery(a, n, p),
-        b=QMontgomeryUInt(n).uint_to_montgomery(b, n, p),
-        x=QMontgomeryUInt(n).uint_to_montgomery(x, n, p),
-        y=QMontgomeryUInt(n).uint_to_montgomery(y, n, p),
-        lam_r=QMontgomeryUInt(n).uint_to_montgomery(lam_r, n, p),
+        a=QMontgomeryUInt(n).uint_to_montgomery(a, p),
+        b=QMontgomeryUInt(n).uint_to_montgomery(b, p),
+        x=QMontgomeryUInt(n).uint_to_montgomery(x, p),
+        y=QMontgomeryUInt(n).uint_to_montgomery(y, p),
+        lam_r=QMontgomeryUInt(n).uint_to_montgomery(lam_r, p),
     )
     assert ret1 == ret2