From 162dca87c332921e53e16fe026207da5a3aa25d3 Mon Sep 17 00:00:00 2001
From: Sean Parkinson <sean@wolfssl.com>
Date: Fri, 23 Dec 2022 09:21:13 +1000
Subject: [PATCH] SP math: rework mod 3

Simplification when only calculating mod and modulus is 3.
2^(2*n) * x mod 3 = x mod 3.
Add all digits and calculate mod 3 of sum.
---
 wolfcrypt/src/sp_int.c | 50 +++++++++++++++++++++++++++---------------
 wolfcrypt/test/test.c  | 26 +++++++++++++++++-----
 2 files changed, 52 insertions(+), 24 deletions(-)

diff --git a/wolfcrypt/src/sp_int.c b/wolfcrypt/src/sp_int.c
index 29efb617ab..df7bcf0c2c 100644
--- a/wolfcrypt/src/sp_int.c
+++ b/wolfcrypt/src/sp_int.c
@@ -6095,27 +6095,41 @@ static void _sp_div_3(const sp_int* a, sp_int* r, sp_int_digit* rem)
 
     /* Check whether only mod value needed. */
     if (r == NULL) {
-        /* Divide starting at most significant word down to least. */
-        for (i = a->used - 1; i >= 0; i--) {
+        /*    2^2 mod 3 = 4 mod 3 = 1.
+         * => 2^(2*n) mod 3 = (2^2 mod 3)^n mod 3 = 1^n mod 3 = 1
+         * => (2^(2*n) * x) mod 3 = (2^(2*n) mod 3) * (x mod 3) = x mod 3
+         *
+         * Calculate mod 3 on sum of digits as SP_WORD_SIZE is a multiple of 2.
+         */
     #ifndef SQR_MUL_ASM
-            /* Combine remainder from last operation with this word. */
-            t = ((sp_int_word)tr << SP_WORD_SIZE) | a->dp[i];
-            /* Get top digit after multipling by (2^SP_WORD_SIZE) / 3. */
-            tt = (t * SP_DIV_3_CONST) >> SP_WORD_SIZE;
-            /* Subtract trail division. */
-            tr = (sp_int_digit)(t - (sp_int_word)tt * 3);
+        t = 0;
+        /* Sum the digits. */
+        for (i = 0; i < a->used; i++) {
+            t += a->dp[i];
+        }
+        /* Sum digits of sum. */
+        t = (t >> SP_WORD_SIZE) + (t & SP_MASK);
+        /* Get top digit after multipling by (2^SP_WORD_SIZE) / 3. */
+        tt = (t * SP_DIV_3_CONST) >> SP_WORD_SIZE;
+        /* Subtract trail division. */
+        tr = (sp_int_digit)(t - (sp_int_word)tt * 3);
     #else
-            /* Multiply digit by (2^SP_WORD_SIZE) / 3. */
-            SP_ASM_MUL(l, tt, a->dp[i], t);
-            /* Add remainder multiplied by (2^SP_WORD_SIZE) / 3 to top digit. */
-            tt += tr * SP_DIV_3_CONST;
-            /* Subtract trail division from digit. */
-            tr = a->dp[i] - (tt * 3);
+        /* Sum the digits. */
+        for (i = 0; i < a->used; i++) {
+            SP_ASM_ADDC_REG(l, tr, a->dp[i]);
+        }
+        /* Sum digits of sum - can get carry. */
+        SP_ASM_ADDC_REG(l, tt, tr);
+        /* Multiply digit by (2^SP_WORD_SIZE) / 3. */
+        SP_ASM_MUL(t, tr, l, t);
+        /* Add remainder multiplied by (2^SP_WORD_SIZE) / 3 to top digit. */
+        tr += tt * SP_DIV_3_CONST;
+        /* Subtract trail division from digit. */
+        tr = l - (tr * 3);
     #endif
-            /* tr is 0..5 but need 0..2 */
-            /* Fix up remainder. */
-            tr = sp_rem6[tr];
-        }
+        /* tr is 0..5 but need 0..2 */
+        /* Fix up remainder. */
+        tr = sp_rem6[tr];
         *rem = tr;
     }
     /* At least result needed - remainder is calculated anyway. */
diff --git a/wolfcrypt/test/test.c b/wolfcrypt/test/test.c
index d29d9aed9f..21bcee165c 100644
--- a/wolfcrypt/test/test.c
+++ b/wolfcrypt/test/test.c
@@ -41281,12 +41281,13 @@ static int mp_test_mod_2d(mp_int* a, mp_int* r, mp_int* t, WC_RNG* rng)
 }
 #endif
 
-#if (defined(HAVE_ECC) && defined(HAVE_COMP_KEY)) || \
-                            (defined(OPENSSL_EXTRA) && defined(WOLFSSL_KEY_GEN))
-static int mp_test_mod_d(mp_int* a)
+#if (defined(HAVE_ECC) && defined(HAVE_COMP_KEY)) || defined(WOLFSSL_KEY_GEN)
+static int mp_test_mod_d(mp_int* a, WC_RNG* rng)
 {
     int ret;
     mp_digit r;
+    mp_digit rem;
+    int i;
 
     if (mp_set(a, 1) != MP_OKAY)
         return -13130;
@@ -41305,6 +41306,20 @@ static int mp_test_mod_d(mp_int* a)
     if (ret != MP_OKAY)
         return -13134;
 
+    for (i = MP_MAX_TEST_BYTE_LEN - 16; i <= MP_MAX_TEST_BYTE_LEN; i++) {
+        ret = randNum(a, i, rng, NULL);
+        if (ret != MP_OKAY)
+            return -13135;
+        ret = mp_mod_d(a, 3, &r);
+        if (ret != MP_OKAY)
+            return -13136;
+        ret = mp_div_d(a, 3, a, &rem);
+        if (ret != MP_OKAY)
+            return -13137;
+        if (r != rem)
+            return -13138;
+    }
+
     return 0;
 }
 #endif
@@ -41992,9 +42007,8 @@ WOLFSSL_TEST_SUBROUTINE int mp_test(void)
     if ((ret = mp_test_mod_2d(&a, &r1, &p, &rng)) != 0)
         return ret;
 #endif
-#if (defined(HAVE_ECC) && defined(HAVE_COMP_KEY)) || \
-                            (defined(OPENSSL_EXTRA) && defined(WOLFSSL_KEY_GEN))
-    if ((ret = mp_test_mod_d(&a)) != 0)
+#if (defined(HAVE_ECC) && defined(HAVE_COMP_KEY)) || defined(WOLFSSL_KEY_GEN)
+    if ((ret = mp_test_mod_d(&a, &rng)) != 0)
         return ret;
 #endif
     if ((ret = mp_test_mul_sqr(&a, &b, &r1, &r2, &rng)) != 0)