Fix possible overflow with assignment operations

2 files changed, 21 insertions, 3 deletions

diff --git a/test b/test
index 2930e89..3256a12 100755
--- a/test
+++ b/test
@@ -115,6 +115,24 @@ public:
                 assert_equal(n, Zmod<12>(4));
         }
 },
+{
+        .name = "Multiplication overflow",
+        .f = []() {
+                Zmod<10> n = 8;
+                Zmod<10> m = 9;
+                auto prod = m * n;
+                assert_equal(prod.toint64(), 2);
+        }
+},
+{
+        .name = "Multiplication and assignment overflow",
+        .f = []() {
+                Zmod<10> n = 8;
+                Zmod<10> m = 9;
+                n *= m;
+                assert_equal(n.toint64(), 2);
+        }
+},
 };
 
 int main() {
diff --git a/zmodn.h b/zmodn.h
index 562aa22..f161ad1 100644
--- a/zmodn.h
+++ b/zmodn.h
@@ -17,9 +17,9 @@ public:
         Zmod operator+(const Zmod& z) const { return int64 + z.int64; }
         Zmod operator-(const Zmod& z) const { return int64 - z.int64; }
         Zmod operator*(const Zmod& z) const { return int64 * z.int64; }
-        Zmod operator+=(const Zmod& z) { return int64 += z.int64; }
-        Zmod operator-=(const Zmod& z) { return int64 -= z.int64; }
-        Zmod operator*=(const Zmod& z) { return int64 *= z.int64; }
+        Zmod operator+=(const Zmod& z) { return (*this) = int64 + z.int64; }
+        Zmod operator-=(const Zmod& z) { return (*this) = int64 - z.int64; }
+        Zmod operator*=(const Zmod& z) { return (*this) = int64 * z.int64; }
 
         bool operator==(const Zmod& z) const { return int64 == z.int64; }
         bool operator!=(const Zmod& z) const { return int64 != z.int64; }