Use only one limb parameter for assign

Signed-off-by: Gabor Mezei <gabor.mezei@arm.com>
diff --git a/library/bignum_core.c b/library/bignum_core.c
index 7074a09..83c115e 100644
--- a/library/bignum_core.c
+++ b/library/bignum_core.c
@@ -163,18 +163,11 @@
 }
 
 void mbedtls_mpi_core_cond_assign( mbedtls_mpi_uint *X,
-                                   size_t X_limbs,
                                    const mbedtls_mpi_uint *Y,
-                                   size_t Y_limbs,
+                                   size_t limbs,
                                    unsigned char assign )
 {
-    /* all-bits 1 if assign is 1, all-bits 0 if assign is 0 */
-    mbedtls_mpi_uint limb_mask = mbedtls_ct_mpi_uint_mask( assign );
-
-    mbedtls_ct_mpi_uint_cond_assign( Y_limbs, X, Y, assign );
-
-    for( size_t i = Y_limbs; i < X_limbs; i++ )
-        X[i] &= ~limb_mask;
+    mbedtls_ct_mpi_uint_cond_assign( limbs, X, Y, assign );
 }
 
 void mbedtls_mpi_core_cond_swap( mbedtls_mpi_uint *X,
diff --git a/library/bignum_core.h b/library/bignum_core.h
index d895148..cf7688d 100644
--- a/library/bignum_core.h
+++ b/library/bignum_core.h
@@ -79,9 +79,9 @@
  *          the condition was true or not.
  *
  * \param[OUT] X        The address of the first MPI. This must be initialized.
- * \param      X_limbs  The number of limbs of \p X.
+ *                      It must have at least \p limbs limbs.
  * \param[IN]  Y        The address of the second MPI. This must be initialized.
- * \param      Y_limbs  The number of limbs of \p Y.
+ * \param      limbs    The number of limbs of \p Y.
  * \param      assign   The condition deciding whether to perform the
  *                      assignment or not. Must be either 0 or 1:
  *                      * \c 1: Perform the assignment `X = Y`.
@@ -95,9 +95,8 @@
  *                 neither its original value nor the value in \p Y.
  */
 void mbedtls_mpi_core_cond_assign( mbedtls_mpi_uint *X,
-                                   size_t X_limbs,
                                    const mbedtls_mpi_uint *Y,
-                                   size_t Y_limbs,
+                                   size_t limbs,
                                    unsigned char assign );
 
 /**
diff --git a/library/bignum_mod_raw.c b/library/bignum_mod_raw.c
index 98994ab..3c7f889 100644
--- a/library/bignum_mod_raw.c
+++ b/library/bignum_mod_raw.c
@@ -46,8 +46,7 @@
                                       const mbedtls_mpi_mod_modulus *m,
                                       unsigned char assign )
 {
-    mbedtls_mpi_core_cond_assign( X, m->limbs,
-                                  Y, m->limbs, assign );
+    mbedtls_mpi_core_cond_assign( X, Y, m->limbs, assign );
 }
 
 void mbedtls_mpi_mod_raw_cond_swap( mbedtls_mpi_uint *X,
diff --git a/library/constant_time.c b/library/constant_time.c
index 7bf67f4..d01998b 100644
--- a/library/constant_time.c
+++ b/library/constant_time.c
@@ -682,11 +682,17 @@
     MPI_VALIDATE_RET( X != NULL );
     MPI_VALIDATE_RET( Y != NULL );
 
+    /* all-bits 1 if assign is 1, all-bits 0 if assign is 0 */
+    mbedtls_mpi_uint limb_mask = mbedtls_ct_mpi_uint_mask( assign );
+
     MBEDTLS_MPI_CHK( mbedtls_mpi_grow( X, Y->n ) );
 
     X->s = mbedtls_ct_cond_select_sign( assign, Y->s, X->s );
 
-    mbedtls_mpi_core_cond_assign( X->p, X->n, Y->p, Y->n, assign );
+    mbedtls_mpi_core_cond_assign( X->p, Y->p, Y->n, assign );
+
+    for( size_t i = Y->n; i < X->n; i++ )
+        X->p[i] &= ~limb_mask;
 
 cleanup:
     return( ret );