diff --git a/library/bignum_mod.c b/library/bignum_mod.c
index acf45e9..84f3896 100644
--- a/library/bignum_mod.c
+++ b/library/bignum_mod.c
@@ -88,7 +88,7 @@
             N->rep.mont.mm = 0;
             break;
         case MBEDTLS_MPI_MOD_REP_OPT_RED:
-            mbedtls_free(N->rep.ored);
+            N->rep.ored.modp = NULL;
             break;
         case MBEDTLS_MPI_MOD_REP_INVALID:
             break;
@@ -136,33 +136,25 @@
     return ret;
 }
 
-int mbedtls_mpi_mod_modulus_setup(mbedtls_mpi_mod_modulus *N,
-                                  const mbedtls_mpi_uint *p,
-                                  size_t p_limbs,
-                                  mbedtls_mpi_mod_rep_selector int_rep)
+static inline void standard_modulus_setup(mbedtls_mpi_mod_modulus *N,
+                                          const mbedtls_mpi_uint *p,
+                                          size_t p_limbs,
+                                          mbedtls_mpi_mod_rep_selector int_rep)
 {
-    int ret = 0;
-
     N->p = p;
     N->limbs = p_limbs;
     N->bits = mbedtls_mpi_core_bitlen(p, p_limbs);
+    N->int_rep = int_rep;
+}
 
-    switch (int_rep) {
-        case MBEDTLS_MPI_MOD_REP_MONTGOMERY:
-            N->int_rep = int_rep;
-            N->rep.mont.mm = mbedtls_mpi_core_montmul_init(N->p);
-            ret = set_mont_const_square(&N->rep.mont.rr, N->p, N->limbs);
-            break;
-        case MBEDTLS_MPI_MOD_REP_OPT_RED:
-            N->int_rep = int_rep;
-            N->rep.ored = NULL;
-            break;
-        default:
-            ret = MBEDTLS_ERR_MPI_BAD_INPUT_DATA;
-            goto exit;
-    }
-
-exit:
+int mbedtls_mpi_mod_modulus_setup(mbedtls_mpi_mod_modulus *N,
+                                  const mbedtls_mpi_uint *p,
+                                  size_t p_limbs)
+{
+    int ret = 0;
+    standard_modulus_setup(N, p, p_limbs, MBEDTLS_MPI_MOD_REP_MONTGOMERY);
+    N->rep.mont.mm = mbedtls_mpi_core_montmul_init(N->p);
+    ret = set_mont_const_square(&N->rep.mont.rr, N->p, N->limbs);
 
     if (ret != 0) {
         mbedtls_mpi_mod_modulus_free(N);
@@ -171,6 +163,16 @@
     return ret;
 }
 
+int mbedtls_mpi_mod_optred_modulus_setup(mbedtls_mpi_mod_modulus *N,
+                                         const mbedtls_mpi_uint *p,
+                                         size_t p_limbs,
+                                         mbedtls_mpi_modp_fn modp)
+{
+    standard_modulus_setup(N, p, p_limbs, MBEDTLS_MPI_MOD_REP_OPT_RED);
+    N->rep.ored.modp = modp;
+    return 0;
+}
+
 int mbedtls_mpi_mod_mul(mbedtls_mpi_mod_residue *X,
                         const mbedtls_mpi_mod_residue *A,
                         const mbedtls_mpi_mod_residue *B,
@@ -235,8 +237,7 @@
     mbedtls_mpi_mod_modulus Nmont;
     mbedtls_mpi_mod_modulus_init(&Nmont);
 
-    MBEDTLS_MPI_CHK(mbedtls_mpi_mod_modulus_setup(&Nmont, N->p, N->limbs,
-                                                  MBEDTLS_MPI_MOD_REP_MONTGOMERY));
+    MBEDTLS_MPI_CHK(mbedtls_mpi_mod_modulus_setup(&Nmont, N->p, N->limbs));
 
     /* We'll use X->p to hold the Montgomery form of the input A->p */
     mbedtls_mpi_core_to_mont_rep(X->p, A->p, Nmont.p, Nmont.limbs,
diff --git a/library/bignum_mod.h b/library/bignum_mod.h
index db177ed..39e8fd2 100644
--- a/library/bignum_mod.h
+++ b/library/bignum_mod.h
@@ -98,10 +98,11 @@
     /* Skip 1 as it is slightly easier to accidentally pass to functions. */
     /** Montgomery representation. */
     MBEDTLS_MPI_MOD_REP_MONTGOMERY = 2,
-    /** TODO: document this.
-     *
-     * Residues are in canonical representation.
-     */
+    /* Optimised reduction available. This indicates a coordinate modulus (P)
+     * and one or more of the following have been configured:
+     * - A nist curve (MBEDTLS_ECP_DP_SECPXXXR1_ENABLED) & MBEDTLS_ECP_NIST_OPTIM.
+     * - A Kobliz Curve.
+     * - A Fast Reduction Curve CURVE25519 or CURVE448. */
     MBEDTLS_MPI_MOD_REP_OPT_RED,
 } mbedtls_mpi_mod_rep_selector;
 
@@ -123,7 +124,11 @@
     mbedtls_mpi_uint mm;         /* Montgomery const for -N^{-1} mod 2^{ciL} */
 } mbedtls_mpi_mont_struct;
 
-typedef void *mbedtls_mpi_opt_red_struct;
+typedef int (*mbedtls_mpi_modp_fn)(mbedtls_mpi_uint *X, size_t X_limbs);
+
+typedef struct {
+    mbedtls_mpi_modp_fn modp;    /* The optimised reduction function pointer */
+} mbedtls_mpi_opt_red_struct;
 
 typedef struct {
     const mbedtls_mpi_uint *p;
@@ -197,16 +202,29 @@
  *                  not be modified in any way until after
  *                  mbedtls_mpi_mod_modulus_free() is called.
  * \param p_limbs   The number of limbs of \p p.
- * \param int_rep   The internal representation to be used for residues
- *                  associated with \p N (see #mbedtls_mpi_mod_rep_selector).
  *
  * \return      \c 0 if successful.
- * \return      #MBEDTLS_ERR_MPI_BAD_INPUT_DATA if \p int_rep is invalid.
  */
 int mbedtls_mpi_mod_modulus_setup(mbedtls_mpi_mod_modulus *N,
                                   const mbedtls_mpi_uint *p,
-                                  size_t p_limbs,
-                                  mbedtls_mpi_mod_rep_selector int_rep);
+                                  size_t p_limbs);
+
+/** Setup an optimised-reduction compatible modulus structure.
+ *
+ * \param[out] N    The address of the modulus structure to populate.
+ * \param[in] p     The address of the limb array storing the value of \p N.
+ *                  The memory pointed to by \p p will be used by \p N and must
+ *                  not be modified in any way until after
+ *                  mbedtls_mpi_mod_modulus_free() is called.
+ * \param p_limbs   The number of limbs of \p p.
+ * \param modp      A pointer to the optimised reduction function to use. \p p.
+ *
+ * \return      \c 0 if successful.
+ */
+int mbedtls_mpi_mod_optred_modulus_setup(mbedtls_mpi_mod_modulus *N,
+                                         const mbedtls_mpi_uint *p,
+                                         size_t p_limbs,
+                                         mbedtls_mpi_modp_fn modp);
 
 /** Free elements of a modulus structure.
  *
diff --git a/library/bignum_mod_raw.c b/library/bignum_mod_raw.c
index ef8c2b3..c539197 100644
--- a/library/bignum_mod_raw.c
+++ b/library/bignum_mod_raw.c
@@ -114,8 +114,6 @@
     (void) mbedtls_mpi_core_add_if(X, N->p, N->limbs, (unsigned) c);
 }
 
-#if defined(MBEDTLS_TEST_HOOKS)
-
 MBEDTLS_STATIC_TESTABLE
 void mbedtls_mpi_mod_raw_fix_quasi_reduction(mbedtls_mpi_uint *X,
                                              const mbedtls_mpi_mod_modulus *N)
@@ -125,7 +123,6 @@
     (void) mbedtls_mpi_core_add_if(X, N->p, N->limbs, (unsigned) c);
 }
 
-#endif /* MBEDTLS_TEST_HOOKS */
 
 void mbedtls_mpi_mod_raw_mul(mbedtls_mpi_uint *X,
                              const mbedtls_mpi_uint *A,
@@ -133,8 +130,31 @@
                              const mbedtls_mpi_mod_modulus *N,
                              mbedtls_mpi_uint *T)
 {
-    mbedtls_mpi_core_montmul(X, A, B, N->limbs, N->p, N->limbs,
-                             N->rep.mont.mm, T);
+    /* Standard (A * B) multiplication stored into pre-allocated T
+     * buffer of fixed limb size of (2N + 1).
+     *
+     * The space may not not fully filled by when
+     * MBEDTLS_MPI_MOD_REP_OPT_RED is used. */
+    const size_t T_limbs = BITS_TO_LIMBS(N->bits) * 2;
+    switch (N->int_rep) {
+        case MBEDTLS_MPI_MOD_REP_MONTGOMERY:
+            mbedtls_mpi_core_montmul(X, A, B, N->limbs, N->p, N->limbs,
+                                     N->rep.mont.mm, T);
+            break;
+        case MBEDTLS_MPI_MOD_REP_OPT_RED:
+            mbedtls_mpi_core_mul(T, A, N->limbs, B, N->limbs);
+
+            /* Optimised Reduction */
+            (*N->rep.ored.modp)(T, T_limbs);
+
+            /* Convert back to canonical representation */
+            mbedtls_mpi_mod_raw_fix_quasi_reduction(T, N);
+            memcpy(X, T, N->limbs * sizeof(mbedtls_mpi_uint));
+            break;
+        default:
+            break;
+    }
+
 }
 
 size_t mbedtls_mpi_mod_raw_inv_prime_working_limbs(size_t AN_limbs)
diff --git a/library/bn_mul.h b/library/bn_mul.h
index c5994f7..43dd5c2 100644
--- a/library/bn_mul.h
+++ b/library/bn_mul.h
@@ -248,27 +248,39 @@
 
 #endif /* AMD64 */
 
-#if defined(__aarch64__)
+// The following assembly code assumes that a pointer will fit in a 64-bit register
+// (including ILP32 __aarch64__ ABIs such as on watchOS, hence the 2^32 - 1)
+#if defined(__aarch64__) && (UINTPTR_MAX == 0xfffffffful || UINTPTR_MAX == 0xfffffffffffffffful)
 
+/*
+ * There are some issues around different compilers requiring different constraint
+ * syntax for updating pointers from assembly code (see notes for
+ * MBEDTLS_ASM_AARCH64_PTR_CONSTRAINT in common.h), especially on aarch64_32 (aka ILP32).
+ *
+ * For this reason we cast the pointers to/from uintptr_t here.
+ */
 #define MULADDC_X1_INIT             \
-    asm(
+    do { uintptr_t muladdc_d = (uintptr_t) d, muladdc_s = (uintptr_t) s; asm(
 
 #define MULADDC_X1_CORE             \
-        "ldr x4, [%2], #8   \n\t"   \
-        "ldr x5, [%1]       \n\t"   \
+        "ldr x4, [%x2], #8  \n\t"   \
+        "ldr x5, [%x1]      \n\t"   \
         "mul x6, x4, %4     \n\t"   \
         "umulh x7, x4, %4   \n\t"   \
         "adds x5, x5, x6    \n\t"   \
         "adc x7, x7, xzr    \n\t"   \
         "adds x5, x5, %0    \n\t"   \
         "adc %0, x7, xzr    \n\t"   \
-        "str x5, [%1], #8   \n\t"
+        "str x5, [%x1], #8  \n\t"
 
 #define MULADDC_X1_STOP                                                 \
-         : "+r" (c),  "+r" (d), "+r" (s), "+m" (*(uint64_t (*)[16]) d)  \
+         : "+r" (c),                                                    \
+           "+r" (muladdc_d),                                            \
+           "+r" (muladdc_s),                                            \
+           "+m" (*(uint64_t (*)[16]) d)                                 \
          : "r" (b), "m" (*(const uint64_t (*)[16]) s)                   \
          : "x4", "x5", "x6", "x7", "cc"                                 \
-    );
+    ); d = (mbedtls_mpi_uint *)muladdc_d; s = (mbedtls_mpi_uint *)muladdc_s; } while (0);
 
 #endif /* Aarch64 */
 
diff --git a/library/common.h b/library/common.h
index b48a1fc..97846c4 100644
--- a/library/common.h
+++ b/library/common.h
@@ -69,6 +69,44 @@
 #define MBEDTLS_TEST_HOOK_TEST_ASSERT(TEST)
 #endif /* defined(MBEDTLS_TEST_HOOKS) */
 
+/** \def ARRAY_LENGTH
+ * Return the number of elements of a static or stack array.
+ *
+ * \param array         A value of array (not pointer) type.
+ *
+ * \return The number of elements of the array.
+ */
+/* A correct implementation of ARRAY_LENGTH, but which silently gives
+ * a nonsensical result if called with a pointer rather than an array. */
+#define ARRAY_LENGTH_UNSAFE(array)            \
+    (sizeof(array) / sizeof(*(array)))
+
+#if defined(__GNUC__)
+/* Test if arg and &(arg)[0] have the same type. This is true if arg is
+ * an array but not if it's a pointer. */
+#define IS_ARRAY_NOT_POINTER(arg)                                     \
+    (!__builtin_types_compatible_p(__typeof__(arg),                \
+                                   __typeof__(&(arg)[0])))
+/* A compile-time constant with the value 0. If `const_expr` is not a
+ * compile-time constant with a nonzero value, cause a compile-time error. */
+#define STATIC_ASSERT_EXPR(const_expr)                                \
+    (0 && sizeof(struct { unsigned int STATIC_ASSERT : 1 - 2 * !(const_expr); }))
+
+/* Return the scalar value `value` (possibly promoted). This is a compile-time
+ * constant if `value` is. `condition` must be a compile-time constant.
+ * If `condition` is false, arrange to cause a compile-time error. */
+#define STATIC_ASSERT_THEN_RETURN(condition, value)   \
+    (STATIC_ASSERT_EXPR(condition) ? 0 : (value))
+
+#define ARRAY_LENGTH(array)                                           \
+    (STATIC_ASSERT_THEN_RETURN(IS_ARRAY_NOT_POINTER(array),         \
+                               ARRAY_LENGTH_UNSAFE(array)))
+
+#else
+/* If we aren't sure the compiler supports our non-standard tricks,
+ * fall back to the unsafe implementation. */
+#define ARRAY_LENGTH(array) ARRAY_LENGTH_UNSAFE(array)
+#endif
 /** Allow library to access its structs' private members.
  *
  * Although structs defined in header files are publicly available,
@@ -169,6 +207,34 @@
 #endif
 /* *INDENT-ON* */
 
+/*
+ * Define the constraint used for read-only pointer operands to aarch64 asm.
+ *
+ * This is normally the usual "r", but for aarch64_32 (aka ILP32,
+ * as found in watchos), "p" is required to avoid warnings from clang.
+ *
+ * Note that clang does not recognise '+p' or '=p', and armclang
+ * does not recognise 'p' at all. Therefore, to update a pointer from
+ * aarch64 assembly, it is necessary to use something like:
+ *
+ * uintptr_t uptr = (uintptr_t) ptr;
+ * asm( "ldr x4, [%x0], #8" ... : "+r" (uptr) : : )
+ * ptr = (void*) uptr;
+ *
+ * Note that the "x" in "%x0" is neccessary; writing "%0" will cause warnings.
+ */
+#if defined(__aarch64__) && defined(MBEDTLS_HAVE_ASM)
+#if UINTPTR_MAX == 0xfffffffful
+/* ILP32: Specify the pointer operand slightly differently, as per #7787. */
+#define MBEDTLS_ASM_AARCH64_PTR_CONSTRAINT "p"
+#elif UINTPTR_MAX == 0xfffffffffffffffful
+/* Normal case (64-bit pointers): use "r" as the constraint for pointer operands to asm */
+#define MBEDTLS_ASM_AARCH64_PTR_CONSTRAINT "r"
+#else
+#error Unrecognised pointer size for aarch64
+#endif
+#endif
+
 /* Always provide a static assert macro, so it can be used unconditionally.
  * It will expand to nothing on some systems.
  * Can be used outside functions (but don't add a trailing ';' in that case:
diff --git a/library/constant_time.c b/library/constant_time.c
index 68b9bdb..f2cdddf 100644
--- a/library/constant_time.c
+++ b/library/constant_time.c
@@ -31,10 +31,18 @@
 #include "mbedtls/platform_util.h"
 
 #include <string.h>
-#if defined(MBEDTLS_USE_PSA_CRYPTO)
-#define PSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,    \
-                                                           psa_to_ssl_errors,              \
-                                                           psa_generic_status_to_mbedtls)
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO) && defined(MBEDTLS_SSL_SOME_SUITES_USE_MAC)
+#include "psa/crypto.h"
+/* Define a local translating function to save code size by not using too many
+ * arguments in each translating place. */
+static int local_err_translation(psa_status_t status)
+{
+    return psa_status_to_mbedtls(status, psa_to_ssl_errors,
+                                 ARRAY_LENGTH(psa_to_ssl_errors),
+                                 psa_generic_status_to_mbedtls);
+}
+#define PSA_TO_MBEDTLS_ERR(status) local_err_translation(status)
 #endif
 
 #if !defined(MBEDTLS_CT_ASM)
@@ -57,8 +65,9 @@
  * only used here.
  */
 #if defined(MBEDTLS_EFFICIENT_UNALIGNED_ACCESS) && \
-    (defined(MBEDTLS_CT_ARM_ASM) || defined(MBEDTLS_CT_AARCH64_ASM))
-
+    ((defined(MBEDTLS_CT_ARM_ASM) && (UINTPTR_MAX == 0xfffffffful)) || \
+    defined(MBEDTLS_CT_AARCH64_ASM))
+/* We check pointer sizes to avoid issues with them not matching register size requirements */
 #define MBEDTLS_EFFICIENT_UNALIGNED_VOLATILE_ACCESS
 
 static inline uint32_t mbedtls_get_unaligned_volatile_uint32(volatile const unsigned char *p)
@@ -71,7 +80,7 @@
 #if defined(MBEDTLS_CT_ARM_ASM)
     asm volatile ("ldr %0, [%1]" : "=r" (r) : "r" (p) :);
 #elif defined(MBEDTLS_CT_AARCH64_ASM)
-    asm volatile ("ldr %w0, [%1]" : "=r" (r) : "r" (p) :);
+    asm volatile ("ldr %w0, [%1]" : "=r" (r) : MBEDTLS_ASM_AARCH64_PTR_CONSTRAINT(p) :);
 #else
 #error No assembly defined for mbedtls_get_unaligned_volatile_uint32
 #endif
diff --git a/library/ecp_curves.c b/library/ecp_curves.c
index af649a2..a4fa663 100644
--- a/library/ecp_curves.c
+++ b/library/ecp_curves.c
@@ -4922,7 +4922,7 @@
 static int ecp_mod_p192(mbedtls_mpi *N)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    size_t expected_width = 2 * ((192 + biL - 1) / biL);
+    size_t expected_width = BITS_TO_LIMBS(192) * 2;
     MBEDTLS_MPI_CHK(mbedtls_mpi_grow(N, expected_width));
     ret = mbedtls_ecp_mod_p192_raw(N->p, expected_width);
 
@@ -4936,7 +4936,7 @@
     mbedtls_mpi_uint c = 0, last_carry[WIDTH] = { 0 };
     mbedtls_mpi_uint *p, *end;
 
-    if (Nn != 2*((192 + biL - 1)/biL)) {
+    if (Nn != BITS_TO_LIMBS(192) * 2) {
         return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
     }
 
@@ -5082,7 +5082,7 @@
 static int ecp_mod_p224(mbedtls_mpi *N)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    size_t expected_width =  2 * 224 / biL;
+    size_t expected_width = BITS_TO_LIMBS(224) * 2;
     MBEDTLS_MPI_CHK(mbedtls_mpi_grow(N, expected_width));
     ret = mbedtls_ecp_mod_p224_raw(N->p, expected_width);
 cleanup:
@@ -5092,7 +5092,7 @@
 MBEDTLS_STATIC_TESTABLE
 int mbedtls_ecp_mod_p224_raw(mbedtls_mpi_uint *X, size_t X_limbs)
 {
-    if (X_limbs != 2 * 224 / biL) {
+    if (X_limbs != BITS_TO_LIMBS(224) * 2) {
         return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
     }
 
@@ -5135,7 +5135,7 @@
 static int ecp_mod_p256(mbedtls_mpi *N)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    size_t expected_width = 2 * 256 / biL;
+    size_t expected_width = BITS_TO_LIMBS(256) * 2;
     MBEDTLS_MPI_CHK(mbedtls_mpi_grow(N, expected_width));
     ret = mbedtls_ecp_mod_p256_raw(N->p, expected_width);
 cleanup:
@@ -5145,7 +5145,7 @@
 MBEDTLS_STATIC_TESTABLE
 int mbedtls_ecp_mod_p256_raw(mbedtls_mpi_uint *X, size_t X_limbs)
 {
-    if (X_limbs != 2 * 256 / biL) {
+    if (X_limbs != BITS_TO_LIMBS(256) * 2) {
         return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
     }
 
@@ -5215,7 +5215,7 @@
 static int ecp_mod_p384(mbedtls_mpi *N)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    size_t expected_width = 2 * ((384 + biL - 1) / biL);
+    size_t expected_width = BITS_TO_LIMBS(384) * 2;
     MBEDTLS_MPI_CHK(mbedtls_mpi_grow(N, expected_width));
     ret = mbedtls_ecp_mod_p384_raw(N->p, expected_width);
 cleanup:
@@ -5225,7 +5225,7 @@
 MBEDTLS_STATIC_TESTABLE
 int mbedtls_ecp_mod_p384_raw(mbedtls_mpi_uint *X, size_t X_limbs)
 {
-    if (X_limbs != 2*((384 + biL - 1)/biL)) {
+    if (X_limbs != BITS_TO_LIMBS(384) * 2) {
         return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
     }
 
@@ -5337,7 +5337,7 @@
 static int ecp_mod_p521(mbedtls_mpi *N)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    size_t expected_width = 2 * P521_WIDTH;
+    size_t expected_width = BITS_TO_LIMBS(521) * 2;
     MBEDTLS_MPI_CHK(mbedtls_mpi_grow(N, expected_width));
     ret = mbedtls_ecp_mod_p521_raw(N->p, expected_width);
 cleanup:
@@ -5349,7 +5349,7 @@
 {
     mbedtls_mpi_uint carry = 0;
 
-    if (X_limbs != 2 * P521_WIDTH || X[2 * P521_WIDTH - 1] != 0) {
+    if (X_limbs != BITS_TO_LIMBS(521) * 2) {
         return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
     }
 
@@ -5423,7 +5423,7 @@
 static int ecp_mod_p255(mbedtls_mpi *N)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    size_t expected_width = 2 * P255_WIDTH;
+    size_t expected_width = BITS_TO_LIMBS(255) * 2;
     MBEDTLS_MPI_CHK(mbedtls_mpi_grow(N, expected_width));
     ret = mbedtls_ecp_mod_p255_raw(N->p, expected_width);
 cleanup:
@@ -5434,7 +5434,7 @@
 int mbedtls_ecp_mod_p255_raw(mbedtls_mpi_uint *X, size_t X_Limbs)
 {
 
-    if (X_Limbs != 2 * P255_WIDTH) {
+    if (X_Limbs != BITS_TO_LIMBS(255) * 2) {
         return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
     }
 
@@ -5492,7 +5492,7 @@
 static int ecp_mod_p448(mbedtls_mpi *N)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    size_t expected_width = 2 * ((448 + biL - 1) / biL);
+    size_t expected_width = BITS_TO_LIMBS(448) * 2;
 
     /* This is required as some tests and use cases do not pass in a Bignum of
      * the correct size, and expect the growth to be done automatically, which
@@ -5522,7 +5522,7 @@
     size_t round;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
-    if (X_limbs <= P448_WIDTH) {
+    if (X_limbs != BITS_TO_LIMBS(448) * 2) {
         return 0;
     }
 
@@ -5577,9 +5577,9 @@
     (void) mbedtls_mpi_core_add(X, X, Q, Q_limbs);
 
     /* M = B0 */
-    if (ciL > 4) {
-        M[P224_WIDTH_MIN] &= ((mbedtls_mpi_uint)-1) >> (P224_UNUSED_BITS);
-    }
+#ifdef MBEDTLS_HAVE_INT64
+    M[P224_WIDTH_MIN] &= ((mbedtls_mpi_uint)-1) >> (P224_UNUSED_BITS);
+ #endif
     memset(M + P224_WIDTH_MAX, 0, ((M_limbs - P224_WIDTH_MAX) * ciL));
 
     /* M = M + Q = B0 + B1 */
@@ -5734,7 +5734,7 @@
 static int ecp_mod_p192k1(mbedtls_mpi *N)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    size_t expected_width = 2 * ((192 + biL - 1) / biL);
+    size_t expected_width = BITS_TO_LIMBS(192) * 2;
     MBEDTLS_MPI_CHK(mbedtls_mpi_grow(N, expected_width));
     ret = mbedtls_ecp_mod_p192k1_raw(N->p, expected_width);
 
@@ -5750,7 +5750,7 @@
                                   0x01, 0x00, 0x00, 0x00)
     };
 
-    if (X_limbs != 2 * ((192 + biL - 1) / biL)) {
+    if (X_limbs != BITS_TO_LIMBS(192) * 2) {
         return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
     }
 
@@ -5768,7 +5768,7 @@
 static int ecp_mod_p224k1(mbedtls_mpi *N)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    size_t expected_width =  2 * 224 / biL;
+    size_t expected_width = BITS_TO_LIMBS(224) * 2;
     MBEDTLS_MPI_CHK(mbedtls_mpi_grow(N, expected_width));
     ret = mbedtls_ecp_mod_p224k1_raw(N->p, expected_width);
 
@@ -5784,7 +5784,7 @@
                                   0x01, 0x00, 0x00, 0x00)
     };
 
-    if (X_limbs != 2 * 224 / biL) {
+    if (X_limbs !=  BITS_TO_LIMBS(224) * 2) {
         return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
     }
 
@@ -5802,7 +5802,7 @@
 static int ecp_mod_p256k1(mbedtls_mpi *N)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    size_t expected_width = 2 * ((256 + biL - 1) / biL);
+    size_t expected_width = BITS_TO_LIMBS(256) * 2;
     MBEDTLS_MPI_CHK(mbedtls_mpi_grow(N, expected_width));
     ret = mbedtls_ecp_mod_p256k1_raw(N->p, expected_width);
 
@@ -5818,7 +5818,7 @@
                                   0x01, 0x00, 0x00, 0x00)
     };
 
-    if (X_limbs != 2 * ((256 + biL - 1) / biL)) {
+    if (X_limbs != BITS_TO_LIMBS(256) * 2) {
         return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
     }
 
@@ -5831,20 +5831,24 @@
 MBEDTLS_STATIC_TESTABLE
 int mbedtls_ecp_modulus_setup(mbedtls_mpi_mod_modulus *N,
                               const mbedtls_ecp_group_id id,
-                              const mbedtls_ecp_curve_type ctype)
+                              const mbedtls_ecp_modulus_type ctype)
 {
+    mbedtls_mpi_modp_fn modp = NULL;
     mbedtls_mpi_uint *p = NULL;
     size_t p_limbs;
 
-    if (!(ctype == (mbedtls_ecp_curve_type) MBEDTLS_ECP_MOD_COORDINATE || \
-          ctype == (mbedtls_ecp_curve_type) MBEDTLS_ECP_MOD_SCALAR)) {
+    if (!(ctype == (mbedtls_ecp_modulus_type) MBEDTLS_ECP_MOD_COORDINATE || \
+          ctype == (mbedtls_ecp_modulus_type) MBEDTLS_ECP_MOD_SCALAR)) {
         return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
     }
 
     switch (id) {
 #if defined(MBEDTLS_ECP_DP_SECP192R1_ENABLED)
         case MBEDTLS_ECP_DP_SECP192R1:
-            if (ctype == (mbedtls_ecp_curve_type) MBEDTLS_ECP_MOD_COORDINATE) {
+            if (ctype == (mbedtls_ecp_modulus_type) MBEDTLS_ECP_MOD_COORDINATE) {
+#if defined(MBEDTLS_ECP_NIST_OPTIM)
+                modp = &mbedtls_ecp_mod_p192_raw;
+#endif
                 p = (mbedtls_mpi_uint *) secp192r1_p;
                 p_limbs = CHARS_TO_LIMBS(sizeof(secp192r1_p));
             } else {
@@ -5856,7 +5860,10 @@
 
 #if defined(MBEDTLS_ECP_DP_SECP224R1_ENABLED)
         case MBEDTLS_ECP_DP_SECP224R1:
-            if (ctype == (mbedtls_ecp_curve_type) MBEDTLS_ECP_MOD_COORDINATE) {
+            if (ctype == (mbedtls_ecp_modulus_type) MBEDTLS_ECP_MOD_COORDINATE) {
+#if defined(MBEDTLS_ECP_NIST_OPTIM)
+                modp = &mbedtls_ecp_mod_p224_raw;
+#endif
                 p = (mbedtls_mpi_uint *) secp224r1_p;
                 p_limbs = CHARS_TO_LIMBS(sizeof(secp224r1_p));
             } else {
@@ -5868,7 +5875,10 @@
 
 #if defined(MBEDTLS_ECP_DP_SECP256R1_ENABLED)
         case MBEDTLS_ECP_DP_SECP256R1:
-            if (ctype == (mbedtls_ecp_curve_type) MBEDTLS_ECP_MOD_COORDINATE) {
+            if (ctype == (mbedtls_ecp_modulus_type) MBEDTLS_ECP_MOD_COORDINATE) {
+#if defined(MBEDTLS_ECP_NIST_OPTIM)
+                modp = &mbedtls_ecp_mod_p256_raw;
+#endif
                 p = (mbedtls_mpi_uint *) secp256r1_p;
                 p_limbs = CHARS_TO_LIMBS(sizeof(secp256r1_p));
             } else {
@@ -5880,7 +5890,10 @@
 
 #if defined(MBEDTLS_ECP_DP_SECP384R1_ENABLED)
         case MBEDTLS_ECP_DP_SECP384R1:
-            if (ctype == (mbedtls_ecp_curve_type) MBEDTLS_ECP_MOD_COORDINATE) {
+            if (ctype == (mbedtls_ecp_modulus_type) MBEDTLS_ECP_MOD_COORDINATE) {
+#if defined(MBEDTLS_ECP_NIST_OPTIM)
+                modp = &mbedtls_ecp_mod_p384_raw;
+#endif
                 p = (mbedtls_mpi_uint *) secp384r1_p;
                 p_limbs = CHARS_TO_LIMBS(sizeof(secp384r1_p));
             } else {
@@ -5892,7 +5905,10 @@
 
 #if defined(MBEDTLS_ECP_DP_SECP521R1_ENABLED)
         case MBEDTLS_ECP_DP_SECP521R1:
-            if (ctype == (mbedtls_ecp_curve_type) MBEDTLS_ECP_MOD_COORDINATE) {
+            if (ctype == (mbedtls_ecp_modulus_type) MBEDTLS_ECP_MOD_COORDINATE) {
+#if defined(MBEDTLS_ECP_NIST_OPTIM)
+                modp = &mbedtls_ecp_mod_p521_raw;
+#endif
                 p = (mbedtls_mpi_uint *) secp521r1_p;
                 p_limbs = CHARS_TO_LIMBS(sizeof(secp521r1_p));
             } else {
@@ -5904,7 +5920,7 @@
 
 #if defined(MBEDTLS_ECP_DP_BP256R1_ENABLED)
         case MBEDTLS_ECP_DP_BP256R1:
-            if (ctype == (mbedtls_ecp_curve_type) MBEDTLS_ECP_MOD_COORDINATE) {
+            if (ctype == (mbedtls_ecp_modulus_type) MBEDTLS_ECP_MOD_COORDINATE) {
                 p = (mbedtls_mpi_uint *) brainpoolP256r1_p;
                 p_limbs = CHARS_TO_LIMBS(sizeof(brainpoolP256r1_p));
             } else {
@@ -5916,7 +5932,7 @@
 
 #if defined(MBEDTLS_ECP_DP_BP384R1_ENABLED)
         case MBEDTLS_ECP_DP_BP384R1:
-            if (ctype == (mbedtls_ecp_curve_type) MBEDTLS_ECP_MOD_COORDINATE) {
+            if (ctype == (mbedtls_ecp_modulus_type) MBEDTLS_ECP_MOD_COORDINATE) {
                 p = (mbedtls_mpi_uint *) brainpoolP384r1_p;
                 p_limbs = CHARS_TO_LIMBS(sizeof(brainpoolP384r1_p));
             } else {
@@ -5928,7 +5944,7 @@
 
 #if defined(MBEDTLS_ECP_DP_BP512R1_ENABLED)
         case MBEDTLS_ECP_DP_BP512R1:
-            if (ctype == (mbedtls_ecp_curve_type) MBEDTLS_ECP_MOD_COORDINATE) {
+            if (ctype == (mbedtls_ecp_modulus_type) MBEDTLS_ECP_MOD_COORDINATE) {
                 p = (mbedtls_mpi_uint *) brainpoolP512r1_p;
                 p_limbs = CHARS_TO_LIMBS(sizeof(brainpoolP512r1_p));
             } else {
@@ -5940,7 +5956,8 @@
 
 #if defined(MBEDTLS_ECP_DP_CURVE25519_ENABLED)
         case MBEDTLS_ECP_DP_CURVE25519:
-            if (ctype == (mbedtls_ecp_curve_type) MBEDTLS_ECP_MOD_COORDINATE) {
+            if (ctype == (mbedtls_ecp_modulus_type) MBEDTLS_ECP_MOD_COORDINATE) {
+                modp = &mbedtls_ecp_mod_p255_raw;
                 p = (mbedtls_mpi_uint *) curve25519_p;
                 p_limbs = CHARS_TO_LIMBS(sizeof(curve25519_p));
             } else {
@@ -5952,7 +5969,8 @@
 
 #if defined(MBEDTLS_ECP_DP_SECP192K1_ENABLED)
         case MBEDTLS_ECP_DP_SECP192K1:
-            if (ctype == (mbedtls_ecp_curve_type) MBEDTLS_ECP_MOD_COORDINATE) {
+            if (ctype == (mbedtls_ecp_modulus_type) MBEDTLS_ECP_MOD_COORDINATE) {
+                modp = &mbedtls_ecp_mod_p192k1_raw;
                 p = (mbedtls_mpi_uint *) secp192k1_p;
                 p_limbs = CHARS_TO_LIMBS(sizeof(secp192k1_p));
             } else {
@@ -5964,7 +5982,8 @@
 
 #if defined(MBEDTLS_ECP_DP_SECP224K1_ENABLED)
         case MBEDTLS_ECP_DP_SECP224K1:
-            if (ctype == (mbedtls_ecp_curve_type) MBEDTLS_ECP_MOD_COORDINATE) {
+            if (ctype == (mbedtls_ecp_modulus_type) MBEDTLS_ECP_MOD_COORDINATE) {
+                modp = &mbedtls_ecp_mod_p224k1_raw;
                 p = (mbedtls_mpi_uint *) secp224k1_p;
                 p_limbs = CHARS_TO_LIMBS(sizeof(secp224k1_p));
             } else {
@@ -5976,7 +5995,8 @@
 
 #if defined(MBEDTLS_ECP_DP_SECP256K1_ENABLED)
         case MBEDTLS_ECP_DP_SECP256K1:
-            if (ctype == (mbedtls_ecp_curve_type) MBEDTLS_ECP_MOD_COORDINATE) {
+            if (ctype == (mbedtls_ecp_modulus_type) MBEDTLS_ECP_MOD_COORDINATE) {
+                modp = &mbedtls_ecp_mod_p256k1_raw;
                 p = (mbedtls_mpi_uint *) secp256k1_p;
                 p_limbs = CHARS_TO_LIMBS(sizeof(secp256k1_p));
             } else {
@@ -5988,7 +6008,8 @@
 
 #if defined(MBEDTLS_ECP_DP_CURVE448_ENABLED)
         case MBEDTLS_ECP_DP_CURVE448:
-            if (ctype == (mbedtls_ecp_curve_type) MBEDTLS_ECP_MOD_COORDINATE) {
+            if (ctype == (mbedtls_ecp_modulus_type) MBEDTLS_ECP_MOD_COORDINATE) {
+                modp = &mbedtls_ecp_mod_p448_raw;
                 p = (mbedtls_mpi_uint *) curve448_p;
                 p_limbs = CHARS_TO_LIMBS(sizeof(curve448_p));
             } else {
@@ -6003,9 +6024,14 @@
             return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
     }
 
-    if (mbedtls_mpi_mod_modulus_setup(N, p, p_limbs,
-                                      MBEDTLS_MPI_MOD_REP_MONTGOMERY)) {
-        return MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    if (modp != NULL) {
+        if (mbedtls_mpi_mod_optred_modulus_setup(N, p, p_limbs, modp)) {
+            return MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+        }
+    } else {
+        if (mbedtls_mpi_mod_modulus_setup(N, p, p_limbs)) {
+            return MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+        }
     }
     return 0;
 }
diff --git a/library/ecp_invasive.h b/library/ecp_invasive.h
index 1dc5567..94867b9 100644
--- a/library/ecp_invasive.h
+++ b/library/ecp_invasive.h
@@ -306,7 +306,7 @@
  * \param[in,out] N The address of the modulus structure to populate.
  *                  Must be initialized.
  * \param[in] id    The mbedtls_ecp_group_id for which to initialise the modulus.
- * \param[in] ctype The mbedtls_ecp_curve_type identifier for a coordinate modulus (P)
+ * \param[in] ctype The mbedtls_ecp_modulus_type identifier for a coordinate modulus (P)
  *                  or a scalar modulus (N).
  *
  * \return          \c 0 if successful.
@@ -317,7 +317,7 @@
 MBEDTLS_STATIC_TESTABLE
 int mbedtls_ecp_modulus_setup(mbedtls_mpi_mod_modulus *N,
                               const mbedtls_ecp_group_id id,
-                              const mbedtls_ecp_curve_type ctype);
+                              const mbedtls_ecp_modulus_type ctype);
 
 #endif /* MBEDTLS_TEST_HOOKS && MBEDTLS_ECP_C */
 
diff --git a/library/lmots.c b/library/lmots.c
index 4061edd..4ef2c51 100644
--- a/library/lmots.c
+++ b/library/lmots.c
@@ -45,9 +45,15 @@
 
 #include "psa/crypto.h"
 
-#define PSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,   \
-                                                           psa_to_lms_errors,             \
-                                                           psa_generic_status_to_mbedtls)
+/* Define a local translating function to save code size by not using too many
+ * arguments in each translating place. */
+static int local_err_translation(psa_status_t status)
+{
+    return psa_status_to_mbedtls(status, psa_to_lms_errors,
+                                 ARRAY_LENGTH(psa_to_lms_errors),
+                                 psa_generic_status_to_mbedtls);
+}
+#define PSA_TO_MBEDTLS_ERR(status) local_err_translation(status)
 
 #define PUBLIC_KEY_TYPE_OFFSET     (0)
 #define PUBLIC_KEY_I_KEY_ID_OFFSET (PUBLIC_KEY_TYPE_OFFSET + \
diff --git a/library/lms.c b/library/lms.c
index acc3523..823ce09 100644
--- a/library/lms.c
+++ b/library/lms.c
@@ -46,9 +46,15 @@
 
 #include "mbedtls/platform.h"
 
-#define PSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,   \
-                                                           psa_to_lms_errors,             \
-                                                           psa_generic_status_to_mbedtls)
+/* Define a local translating function to save code size by not using too many
+ * arguments in each translating place. */
+static int local_err_translation(psa_status_t status)
+{
+    return psa_status_to_mbedtls(status, psa_to_lms_errors,
+                                 ARRAY_LENGTH(psa_to_lms_errors),
+                                 psa_generic_status_to_mbedtls);
+}
+#define PSA_TO_MBEDTLS_ERR(status) local_err_translation(status)
 
 #define SIG_Q_LEAF_ID_OFFSET     (0)
 #define SIG_OTS_SIG_OFFSET       (SIG_Q_LEAF_ID_OFFSET + \
diff --git a/library/oid.c b/library/oid.c
index a580992..47a311b 100644
--- a/library/oid.c
+++ b/library/oid.c
@@ -543,7 +543,7 @@
                         mbedtls_pk_type_t,
                         pk_alg)
 
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
 /*
  * For elliptic curves that use namedCurve inside ECParams (RFC 5480)
  */
@@ -674,7 +674,7 @@
                         oid_ecp_grp_algid,
                         mbedtls_ecp_group_id,
                         grp_id)
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 
 #if defined(MBEDTLS_CIPHER_C)
 /*
diff --git a/library/pk.c b/library/pk.c
index 91796de..52eb0d5 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -31,7 +31,7 @@
 #if defined(MBEDTLS_RSA_C)
 #include "mbedtls/rsa.h"
 #endif
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
 #include "mbedtls/ecp.h"
 #endif
 #if defined(MBEDTLS_ECDSA_C)
@@ -125,12 +125,12 @@
         case MBEDTLS_PK_RSA:
             return &mbedtls_rsa_info;
 #endif /* MBEDTLS_RSA_C */
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
         case MBEDTLS_PK_ECKEY:
             return &mbedtls_eckey_info;
         case MBEDTLS_PK_ECKEY_DH:
             return &mbedtls_eckeydh_info;
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 #if defined(MBEDTLS_PK_CAN_ECDSA_SOME)
         case MBEDTLS_PK_ECDSA:
             return &mbedtls_ecdsa_info;
@@ -196,42 +196,6 @@
 }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
 
-#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
-int mbedtls_pk_update_public_key_from_keypair(mbedtls_pk_context *pk,
-                                              mbedtls_ecp_keypair *ecp_keypair)
-{
-    int ret = MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
-
-    if (pk == NULL) {
-        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
-    }
-    /* The raw public key storing mechanism is only supported for EC keys so
-     * we fail silently for other ones. */
-    if ((pk->pk_info->type != MBEDTLS_PK_ECKEY) &&
-        (pk->pk_info->type != MBEDTLS_PK_ECKEY_DH) &&
-        (pk->pk_info->type != MBEDTLS_PK_ECDSA)) {
-        return 0;
-    }
-
-    ret = mbedtls_ecp_point_write_binary(&ecp_keypair->grp, &ecp_keypair->Q,
-                                         MBEDTLS_ECP_PF_UNCOMPRESSED,
-                                         &pk->pub_raw_len,
-                                         pk->pub_raw,
-                                         MBEDTLS_PK_MAX_EC_PUBKEY_RAW_LEN);
-    if (ret != 0) {
-        return ret;
-    }
-
-    pk->ec_family = mbedtls_ecc_group_to_psa(ecp_keypair->grp.id,
-                                             &pk->ec_bits);
-    if (pk->ec_family == 0) {
-        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
-    }
-
-    return 0;
-}
-#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
-
 #if defined(MBEDTLS_PK_RSA_ALT_SUPPORT)
 /*
  * Initialize an RSA-alt context
@@ -903,14 +867,14 @@
                               psa_key_usage_t usage,
                               psa_algorithm_t alg2)
 {
-#if !defined(MBEDTLS_ECP_LIGHT) && !defined(MBEDTLS_RSA_C)
+#if !defined(MBEDTLS_PK_HAVE_ECC_KEYS) && !defined(MBEDTLS_RSA_C)
     ((void) pk);
     ((void) key);
     ((void) alg);
     ((void) usage);
     ((void) alg2);
-#else /* !MBEDTLS_ECP_LIGHT && !MBEDTLS_RSA_C */
-#if defined(MBEDTLS_ECP_LIGHT)
+#else /* !MBEDTLS_PK_HAVE_ECC_KEYS && !MBEDTLS_RSA_C */
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     if (mbedtls_pk_get_type(pk) == MBEDTLS_PK_ECKEY) {
         size_t d_len;
         psa_ecc_family_t curve_id;
@@ -965,7 +929,7 @@
 
         return mbedtls_pk_setup_opaque(pk, *key);
     } else
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 #if defined(MBEDTLS_RSA_C)
     if (mbedtls_pk_get_type(pk) == MBEDTLS_PK_RSA) {
         unsigned char buf[MBEDTLS_PK_RSA_PRV_DER_MAX_BYTES];
@@ -1006,7 +970,7 @@
         return mbedtls_pk_setup_opaque(pk, *key);
     } else
 #endif /* MBEDTLS_RSA_C */
-#endif /* !MBEDTLS_ECP_LIGHT && !MBEDTLS_RSA_C */
+#endif /* !MBEDTLS_PK_HAVE_ECC_KEYS && !MBEDTLS_RSA_C */
     return MBEDTLS_ERR_PK_TYPE_MISMATCH;
 }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
diff --git a/library/pk_internal.h b/library/pk_internal.h
index 388f94a..3d05f57 100644
--- a/library/pk_internal.h
+++ b/library/pk_internal.h
@@ -25,7 +25,7 @@
 
 #include "mbedtls/pk.h"
 
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
 #include "mbedtls/ecp.h"
 #endif
 
@@ -44,7 +44,7 @@
                                                                     psa_pk_status_to_mbedtls)
 #endif
 
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
 /**
  * Public function mbedtls_pk_ec() can be used to get direct access to the
  * wrapped ecp_keypair structure pointed to the pk_ctx. However this is not
@@ -115,21 +115,7 @@
 #if defined(MBEDTLS_ECP_DP_CURVE25519_ENABLED) || defined(MBEDTLS_ECP_DP_CURVE448_ENABLED)
 #define MBEDTLS_PK_HAVE_RFC8410_CURVES
 #endif /* MBEDTLS_ECP_DP_CURVE25519_ENABLED || MBEDTLS_ECP_DP_CURVE448_ENABLED */
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 
-#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
-/**
- * \brief   Copy the public key content in raw format from "ctx->pk_ctx"
- *          (which is an ecp_keypair) into the internal "ctx->pub_raw" buffer.
- *
- * \note    This is a temporary function that can be removed as soon as the pk
- *          module is free from ECP_C
- *
- * \param pk   It is the pk_context which is going to be updated. It acts both
- *             as input and output.
- */
-int mbedtls_pk_update_public_key_from_keypair(mbedtls_pk_context *pk,
-                                              mbedtls_ecp_keypair *ecp_keypair);
-#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 
 #endif /* MBEDTLS_PK_INTERNAL_H */
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 0cadab2..54a4d5d 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -634,7 +634,7 @@
 };
 #endif /* MBEDTLS_RSA_C */
 
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
 /*
  * Generic EC key
  */
@@ -1335,7 +1335,7 @@
 #endif
     eckey_debug,            /* Same underlying key structure */
 };
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 
 #if defined(MBEDTLS_PK_CAN_ECDSA_SOME)
 static int ecdsa_can_do(mbedtls_pk_type_t type)
diff --git a/library/pk_wrap.h b/library/pk_wrap.h
index b4b974f..1436d78 100644
--- a/library/pk_wrap.h
+++ b/library/pk_wrap.h
@@ -120,7 +120,7 @@
 extern const mbedtls_pk_info_t mbedtls_rsa_info;
 #endif
 
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
 extern const mbedtls_pk_info_t mbedtls_eckey_info;
 extern const mbedtls_pk_info_t mbedtls_eckeydh_info;
 #endif
diff --git a/library/pkparse.c b/library/pkparse.c
index 4c55d34..483176a 100644
--- a/library/pkparse.c
+++ b/library/pkparse.c
@@ -37,7 +37,7 @@
 #if defined(MBEDTLS_RSA_C) || defined(MBEDTLS_ECP_C)
 #include "pkwrite.h"
 #endif
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
 #include "pk_internal.h"
 #endif
 #if defined(MBEDTLS_ECDSA_C)
@@ -64,10 +64,10 @@
 #include "mbedtls/platform.h"
 
 /* Helper for Montgomery curves */
-#if defined(MBEDTLS_ECP_LIGHT) && defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS) && defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
 #define MBEDTLS_PK_IS_RFC8410_GROUP_ID(id)  \
     ((id == MBEDTLS_ECP_DP_CURVE25519) || (id == MBEDTLS_ECP_DP_CURVE448))
-#endif /* MBEDTLS_ECP_LIGHT && MBEDTLS_PK_HAVE_RFC8410_CURVES */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS && MBEDTLS_PK_HAVE_RFC8410_CURVES */
 
 #if defined(MBEDTLS_FS_IO)
 /*
@@ -174,7 +174,7 @@
 }
 #endif /* MBEDTLS_FS_IO */
 
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
 /* Minimally parse an ECParameters buffer to and mbedtls_asn1_buf
  *
  * ECParameters ::= CHOICE {
@@ -655,7 +655,6 @@
     mbedtls_ecp_keypair *eck = mbedtls_pk_ec_rw(*pk);
 
     if ((ret = mbedtls_mpi_read_binary_le(&eck->d, key, len)) != 0) {
-        mbedtls_ecp_keypair_free(eck);
         return MBEDTLS_ERROR_ADD(MBEDTLS_ERR_PK_KEY_INVALID_FORMAT, ret);
     }
 #endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
@@ -664,9 +663,6 @@
      * which never contain a public key. As such, derive the public key
      * unconditionally. */
     if ((ret = pk_derive_public_key(pk, key, len, f_rng, p_rng)) != 0) {
-#if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
-        mbedtls_ecp_keypair_free(eck);
-#endif /* !MBEDTLS_PK_USE_PSA_EC_DATA */
         return ret;
     }
 
@@ -674,7 +670,6 @@
      * into PSA. */
 #if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
     if ((ret = mbedtls_ecp_check_privkey(&eck->grp, &eck->d)) != 0) {
-        mbedtls_ecp_keypair_free(eck);
         return ret;
     }
 #endif /* !MBEDTLS_PK_USE_PSA_EC_DATA */
@@ -793,7 +788,7 @@
 
     return ret;
 }
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 
 #if defined(MBEDTLS_RSA_C)
 /*
@@ -878,7 +873,7 @@
     }
 
     ret = mbedtls_oid_get_pk_alg(&alg_oid, pk_alg);
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     if (ret == MBEDTLS_ERR_OID_NOT_FOUND) {
         ret = mbedtls_oid_get_ec_grp_algid(&alg_oid, ec_grp_id);
         if (ret == 0) {
@@ -952,7 +947,7 @@
         ret = pk_get_rsapubkey(p, end, mbedtls_pk_rsa(*pk));
     } else
 #endif /* MBEDTLS_RSA_C */
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     if (pk_alg == MBEDTLS_PK_ECKEY_DH || pk_alg == MBEDTLS_PK_ECKEY) {
 #if defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
         if (MBEDTLS_PK_IS_RFC8410_GROUP_ID(ec_grp_id)) {
@@ -966,7 +961,7 @@
             ret = pk_get_ecpubkey(p, end, pk);
         }
     } else
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
     ret = MBEDTLS_ERR_PK_UNKNOWN_PK_ALG;
 
     if (ret == 0 && *p != end) {
@@ -1170,7 +1165,7 @@
 }
 #endif /* MBEDTLS_RSA_C */
 
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
 /*
  * Parse a SEC1 encoded private EC key
  */
@@ -1186,10 +1181,11 @@
     unsigned char *d;
     unsigned char *end = p + keylen;
     unsigned char *end2;
-    mbedtls_ecp_keypair *eck = mbedtls_pk_ec_rw(*pk);
 #if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_status_t status;
+#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
+    mbedtls_ecp_keypair *eck = mbedtls_pk_ec_rw(*pk);
 #endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 
     /*
@@ -1226,7 +1222,6 @@
 
 #if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
     if ((ret = mbedtls_mpi_read_binary(&eck->d, p, len)) != 0) {
-        mbedtls_ecp_keypair_free(eck);
         return MBEDTLS_ERROR_ADD(MBEDTLS_ERR_PK_KEY_INVALID_FORMAT, ret);
     }
 #endif
@@ -1243,11 +1238,9 @@
                                         0)) == 0) {
             if ((ret = pk_get_ecparams(&p, p + len, &params)) != 0 ||
                 (ret = pk_use_ecparams(&params, pk)) != 0) {
-                mbedtls_ecp_keypair_free(eck);
                 return ret;
             }
         } else if (ret != MBEDTLS_ERR_ASN1_UNEXPECTED_TAG) {
-            mbedtls_ecp_keypair_free(eck);
             return MBEDTLS_ERROR_ADD(MBEDTLS_ERR_PK_KEY_INVALID_FORMAT, ret);
         }
     }
@@ -1283,7 +1276,6 @@
                 }
             }
         } else if (ret != MBEDTLS_ERR_ASN1_UNEXPECTED_TAG) {
-            mbedtls_ecp_keypair_free(eck);
             return MBEDTLS_ERROR_ADD(MBEDTLS_ERR_PK_KEY_INVALID_FORMAT, ret);
         }
     }
@@ -1311,21 +1303,19 @@
 
     if (!pubkey_done) {
         if ((ret = pk_derive_public_key(pk, d, d_len, f_rng, p_rng)) != 0) {
-            mbedtls_ecp_keypair_free(eck);
             return ret;
         }
     }
 
 #if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
     if ((ret = mbedtls_ecp_check_privkey(&eck->grp, &eck->d)) != 0) {
-        mbedtls_ecp_keypair_free(eck);
         return ret;
     }
 #endif /* !MBEDTLS_PK_USE_PSA_EC_DATA */
 
     return 0;
 }
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 
 /*
  * Parse an unencrypted PKCS#8 encoded private key
@@ -1354,7 +1344,7 @@
     mbedtls_ecp_group_id ec_grp_id = MBEDTLS_ECP_DP_NONE;
     const mbedtls_pk_info_t *pk_info;
 
-#if !defined(MBEDTLS_ECP_LIGHT)
+#if !defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     (void) f_rng;
     (void) p_rng;
 #endif
@@ -1419,7 +1409,7 @@
         }
     } else
 #endif /* MBEDTLS_RSA_C */
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     if (pk_alg == MBEDTLS_PK_ECKEY || pk_alg == MBEDTLS_PK_ECKEY_DH) {
 #if defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
         if (MBEDTLS_PK_IS_RFC8410_GROUP_ID(ec_grp_id)) {
@@ -1441,7 +1431,7 @@
             }
         }
     } else
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
     return MBEDTLS_ERR_PK_UNKNOWN_PK_ALG;
 
     return 0;
@@ -1608,7 +1598,7 @@
     }
 #endif /* MBEDTLS_RSA_C */
 
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     /* Avoid calling mbedtls_pem_read_buffer() on non-null-terminated string */
     if (key[keylen - 1] != '\0') {
         ret = MBEDTLS_ERR_PEM_NO_HEADER_FOOTER_PRESENT;
@@ -1637,7 +1627,7 @@
     } else if (ret != MBEDTLS_ERR_PEM_NO_HEADER_FOOTER_PRESENT) {
         return ret;
     }
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 
     /* Avoid calling mbedtls_pem_read_buffer() on non-null-terminated string */
     if (key[keylen - 1] != '\0') {
@@ -1743,7 +1733,7 @@
     mbedtls_pk_init(pk);
 #endif /* MBEDTLS_RSA_C */
 
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     pk_info = mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY);
     if (mbedtls_pk_setup(pk, pk_info) == 0 &&
         pk_parse_key_sec1_der(pk,
@@ -1751,13 +1741,13 @@
         return 0;
     }
     mbedtls_pk_free(pk);
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 
-    /* If MBEDTLS_RSA_C is defined but MBEDTLS_ECP_LIGHT isn't,
+    /* If MBEDTLS_RSA_C is defined but MBEDTLS_PK_HAVE_ECC_KEYS isn't,
      * it is ok to leave the PK context initialized but not
      * freed: It is the caller's responsibility to call pk_init()
      * before calling this function, and to call pk_free()
-     * when it fails. If MBEDTLS_ECP_LIGHT is defined but MBEDTLS_RSA_C
+     * when it fails. If MBEDTLS_PK_HAVE_ECC_KEYS is defined but MBEDTLS_RSA_C
      * isn't, this leads to mbedtls_pk_free() being called
      * twice, once here and once by the caller, but this is
      * also ok and in line with the mbedtls_pk_free() calls
diff --git a/library/pkwrite.c b/library/pkwrite.c
index 218d0c1..5f801e2 100644
--- a/library/pkwrite.c
+++ b/library/pkwrite.c
@@ -38,10 +38,10 @@
 #include "mbedtls/ecp.h"
 #include "mbedtls/platform_util.h"
 #endif
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
 #include "pk_internal.h"
 #endif
-#if defined(MBEDTLS_RSA_C) || defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_RSA_C) || defined(MBEDTLS_PK_HAVE_ECC_KEYS)
 #include "pkwrite.h"
 #endif
 #if defined(MBEDTLS_ECDSA_C)
@@ -58,7 +58,7 @@
 #include "mbedtls/platform.h"
 
 /* Helper for Montgomery curves */
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
 #if defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
 static inline int mbedtls_pk_is_rfc8410(const mbedtls_pk_context *pk)
 {
@@ -76,6 +76,7 @@
 #endif
     return 0;
 }
+
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 /* It is assumed that the input key is opaque */
 static psa_ecc_family_t pk_get_opaque_ec_family(const mbedtls_pk_context *pk)
@@ -93,7 +94,7 @@
 }
 #endif /* MBETLS_USE_PSA_CRYPTO */
 #endif /* MBEDTLS_PK_HAVE_RFC8410_CURVES */
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 /* It is assumed that the input key is opaque */
@@ -158,7 +159,7 @@
 }
 #endif /* MBEDTLS_RSA_C */
 
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
 #if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
 static int pk_write_ec_pubkey(unsigned char **p, unsigned char *start,
                               const mbedtls_pk_context *pk)
@@ -316,7 +317,7 @@
     return ret;
 }
 #endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 static int pk_write_opaque_pubkey(unsigned char **p, unsigned char *start,
@@ -353,7 +354,7 @@
         MBEDTLS_ASN1_CHK_ADD(len, pk_write_rsa_pubkey(p, start, key));
     } else
 #endif
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     if (mbedtls_pk_get_type(key) == MBEDTLS_PK_ECKEY) {
         MBEDTLS_ASN1_CHK_ADD(len, pk_write_ec_pubkey(p, start, key));
     } else
@@ -375,7 +376,7 @@
     int has_par = 1;
     size_t len = 0, par_len = 0, oid_len = 0;
     mbedtls_pk_type_t pk_type;
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     mbedtls_ecp_group_id ec_grp_id = MBEDTLS_ECP_DP_NONE;
 #endif
     const char *oid;
@@ -404,20 +405,20 @@
     MBEDTLS_ASN1_CHK_ADD(len, mbedtls_asn1_write_tag(&c, buf, MBEDTLS_ASN1_BIT_STRING));
 
     pk_type = mbedtls_pk_get_type(key);
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     if (pk_type == MBEDTLS_PK_ECKEY) {
         ec_grp_id = mbedtls_pk_get_group_id(key);
     }
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     if (pk_type == MBEDTLS_PK_OPAQUE) {
         psa_key_type_t opaque_key_type = pk_get_opaque_key_type(key);
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
         if (PSA_KEY_TYPE_IS_ECC(opaque_key_type)) {
             pk_type = MBEDTLS_PK_ECKEY;
             ec_grp_id = mbedtls_pk_get_group_id(key);
         } else
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
         if (PSA_KEY_TYPE_IS_RSA(opaque_key_type)) {
             /* The rest of the function works as for legacy RSA contexts. */
             pk_type = MBEDTLS_PK_RSA;
@@ -429,7 +430,7 @@
     }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
 
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     if (pk_type == MBEDTLS_PK_ECKEY) {
         /* Some groups have their own AlgorithmIdentifier OID, others are handled
          * by mbedtls_oid_get_oid_by_pk_alg() below */
@@ -445,7 +446,7 @@
             return ret;
         }
     }
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 
     if (oid_len == 0) {
         if ((ret = mbedtls_oid_get_oid_by_pk_alg(pk_type, &oid,
@@ -464,7 +465,7 @@
     return (int) len;
 }
 
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
 #if defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
 /*
  * RFC8410 section 7
@@ -572,7 +573,7 @@
 
     return (int) len;
 }
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 
 #if defined(MBEDTLS_RSA_C)
 static int pk_write_rsa_der(unsigned char **p, unsigned char *buf,
@@ -691,9 +692,9 @@
 #if defined(MBEDTLS_RSA_C)
     int is_rsa_opaque = 0;
 #endif /* MBEDTLS_RSA_C */
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     int is_ec_opaque = 0;
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     psa_key_type_t opaque_key_type;
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
@@ -710,9 +711,9 @@
 #if defined(MBEDTLS_RSA_C)
         is_rsa_opaque = PSA_KEY_TYPE_IS_RSA(opaque_key_type);
 #endif /* MBEDTLS_RSA_C */
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
         is_ec_opaque = PSA_KEY_TYPE_IS_ECC(opaque_key_type);
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
     }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
 
@@ -721,7 +722,7 @@
         return pk_write_rsa_der(&c, buf, key);
     } else
 #endif /* MBEDTLS_RSA_C */
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     if ((mbedtls_pk_get_type(key) == MBEDTLS_PK_ECKEY) || is_ec_opaque) {
 #if defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
         if (mbedtls_pk_is_rfc8410(key)) {
@@ -730,7 +731,7 @@
 #endif /* MBEDTLS_PK_HAVE_RFC8410_CURVES */
         return pk_write_ec_der(&c, buf, key);
     } else
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
     return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
 
     return (int) len;
@@ -781,12 +782,12 @@
     unsigned char output_buf[PRV_DER_MAX_BYTES];
     const char *begin, *end;
     size_t olen = 0;
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     int is_ec_opaque = 0;
 #if defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
     int is_montgomery_opaque = 0;
 #endif /* MBEDTLS_PK_HAVE_RFC8410_CURVES */
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 #if defined(MBEDTLS_RSA_C)
     int is_rsa_opaque = 0;
 #endif
@@ -802,14 +803,14 @@
 #if defined(MBEDTLS_RSA_C)
         is_rsa_opaque = PSA_KEY_TYPE_IS_RSA(opaque_key_type);
 #endif
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
         is_ec_opaque = PSA_KEY_TYPE_IS_ECC(opaque_key_type);
 #if defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
         if (pk_get_opaque_ec_family(key) == PSA_ECC_FAMILY_MONTGOMERY) {
             is_montgomery_opaque = 1;
         }
 #endif /* MBEDTLS_PK_HAVE_RFC8410_CURVES */
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
     }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
 
@@ -819,7 +820,7 @@
         end = PEM_END_PRIVATE_KEY_RSA;
     } else
 #endif
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     if ((mbedtls_pk_get_type(key) == MBEDTLS_PK_ECKEY) || is_ec_opaque) {
 #if defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
         if (is_montgomery_opaque ||
@@ -828,13 +829,13 @@
             begin = PEM_BEGIN_PRIVATE_KEY_PKCS8;
             end = PEM_END_PRIVATE_KEY_PKCS8;
         } else
-#endif
+#endif /* MBEDTLS_PK_HAVE_RFC8410_CURVES */
         {
             begin = PEM_BEGIN_PRIVATE_KEY_EC;
             end = PEM_END_PRIVATE_KEY_EC;
         }
     } else
-#endif
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
     return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
 
     if ((ret = mbedtls_pem_write_buffer(begin, end,
diff --git a/library/pkwrite.h b/library/pkwrite.h
index 8db2333..aa2f17b 100644
--- a/library/pkwrite.h
+++ b/library/pkwrite.h
@@ -73,7 +73,7 @@
 
 #endif /* MBEDTLS_RSA_C */
 
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
 /*
  * EC public keys:
  *  SubjectPublicKeyInfo  ::=  SEQUENCE  {      1 + 2
@@ -98,10 +98,10 @@
  */
 #define MBEDTLS_PK_ECP_PRV_DER_MAX_BYTES    (29 + 3 * MBEDTLS_ECP_MAX_BYTES)
 
-#else /* MBEDTLS_ECP_LIGHT */
+#else /* MBEDTLS_PK_HAVE_ECC_KEYS */
 
 #define MBEDTLS_PK_ECP_PUB_DER_MAX_BYTES   0
 #define MBEDTLS_PK_ECP_PRV_DER_MAX_BYTES   0
 
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 #endif /* MBEDTLS_PK_WRITE_H */
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 5e38c3a..ac6bd5b 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -84,8 +84,6 @@
 #include "mbedtls/sha512.h"
 #include "md_psa.h"
 
-#define ARRAY_LENGTH(array) (sizeof(array) / sizeof(*(array)))
-
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF) ||          \
     defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT) ||  \
     defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXPAND)
@@ -390,7 +388,56 @@
 /* Key management */
 /****************************************************************/
 
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
+psa_ecc_family_t mbedtls_ecc_group_to_psa(mbedtls_ecp_group_id grpid,
+                                          size_t *bits)
+{
+    switch (grpid) {
+        case MBEDTLS_ECP_DP_SECP192R1:
+            *bits = 192;
+            return PSA_ECC_FAMILY_SECP_R1;
+        case MBEDTLS_ECP_DP_SECP224R1:
+            *bits = 224;
+            return PSA_ECC_FAMILY_SECP_R1;
+        case MBEDTLS_ECP_DP_SECP256R1:
+            *bits = 256;
+            return PSA_ECC_FAMILY_SECP_R1;
+        case MBEDTLS_ECP_DP_SECP384R1:
+            *bits = 384;
+            return PSA_ECC_FAMILY_SECP_R1;
+        case MBEDTLS_ECP_DP_SECP521R1:
+            *bits = 521;
+            return PSA_ECC_FAMILY_SECP_R1;
+        case MBEDTLS_ECP_DP_BP256R1:
+            *bits = 256;
+            return PSA_ECC_FAMILY_BRAINPOOL_P_R1;
+        case MBEDTLS_ECP_DP_BP384R1:
+            *bits = 384;
+            return PSA_ECC_FAMILY_BRAINPOOL_P_R1;
+        case MBEDTLS_ECP_DP_BP512R1:
+            *bits = 512;
+            return PSA_ECC_FAMILY_BRAINPOOL_P_R1;
+        case MBEDTLS_ECP_DP_CURVE25519:
+            *bits = 255;
+            return PSA_ECC_FAMILY_MONTGOMERY;
+        case MBEDTLS_ECP_DP_SECP192K1:
+            *bits = 192;
+            return PSA_ECC_FAMILY_SECP_K1;
+        case MBEDTLS_ECP_DP_SECP224K1:
+            *bits = 224;
+            return PSA_ECC_FAMILY_SECP_K1;
+        case MBEDTLS_ECP_DP_SECP256K1:
+            *bits = 256;
+            return PSA_ECC_FAMILY_SECP_K1;
+        case MBEDTLS_ECP_DP_CURVE448:
+            *bits = 448;
+            return PSA_ECC_FAMILY_MONTGOMERY;
+        default:
+            *bits = 0;
+            return 0;
+    }
+}
+
 mbedtls_ecp_group_id mbedtls_ecc_group_of_psa(psa_ecc_family_t curve,
                                               size_t bits,
                                               int bits_is_sloppy)
@@ -482,7 +529,7 @@
     (void) bits_is_sloppy;
     return MBEDTLS_ECP_DP_NONE;
 }
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
 
 psa_status_t psa_validate_unstructured_key_bit_size(psa_key_type_t type,
                                                     size_t bits)
@@ -7718,10 +7765,8 @@
         psa_jpake_computation_stage_t *computation_stage =
             &operation->computation_stage.jpake;
 
-        computation_stage->state = PSA_PAKE_STATE_SETUP;
-        computation_stage->sequence = PSA_PAKE_SEQ_INVALID;
-        computation_stage->input_step = PSA_PAKE_STEP_X1_X2;
-        computation_stage->output_step = PSA_PAKE_STEP_X1_X2;
+        memset(computation_stage, 0, sizeof(*computation_stage));
+        computation_stage->step = PSA_PAKE_STEP_KEY_SHARE;
     } else
 #endif /* PSA_WANT_ALG_JPAKE */
     {
@@ -7890,59 +7935,32 @@
     return status;
 }
 
-/* Auxiliary function to convert core computation stage(step, sequence, state) to single driver step. */
+/* Auxiliary function to convert core computation stage to single driver step. */
 #if defined(PSA_WANT_ALG_JPAKE)
 static psa_crypto_driver_pake_step_t convert_jpake_computation_stage_to_driver_step(
     psa_jpake_computation_stage_t *stage)
 {
-    switch (stage->state) {
-        case PSA_PAKE_OUTPUT_X1_X2:
-        case PSA_PAKE_INPUT_X1_X2:
-            switch (stage->sequence) {
-                case PSA_PAKE_X1_STEP_KEY_SHARE:
-                    return PSA_JPAKE_X1_STEP_KEY_SHARE;
-                case PSA_PAKE_X1_STEP_ZK_PUBLIC:
-                    return PSA_JPAKE_X1_STEP_ZK_PUBLIC;
-                case PSA_PAKE_X1_STEP_ZK_PROOF:
-                    return PSA_JPAKE_X1_STEP_ZK_PROOF;
-                case PSA_PAKE_X2_STEP_KEY_SHARE:
-                    return PSA_JPAKE_X2_STEP_KEY_SHARE;
-                case PSA_PAKE_X2_STEP_ZK_PUBLIC:
-                    return PSA_JPAKE_X2_STEP_ZK_PUBLIC;
-                case PSA_PAKE_X2_STEP_ZK_PROOF:
-                    return PSA_JPAKE_X2_STEP_ZK_PROOF;
-                default:
-                    return PSA_JPAKE_STEP_INVALID;
-            }
-            break;
-        case PSA_PAKE_OUTPUT_X2S:
-            switch (stage->sequence) {
-                case PSA_PAKE_X1_STEP_KEY_SHARE:
-                    return PSA_JPAKE_X2S_STEP_KEY_SHARE;
-                case PSA_PAKE_X1_STEP_ZK_PUBLIC:
-                    return PSA_JPAKE_X2S_STEP_ZK_PUBLIC;
-                case PSA_PAKE_X1_STEP_ZK_PROOF:
-                    return PSA_JPAKE_X2S_STEP_ZK_PROOF;
-                default:
-                    return PSA_JPAKE_STEP_INVALID;
-            }
-            break;
-        case PSA_PAKE_INPUT_X4S:
-            switch (stage->sequence) {
-                case PSA_PAKE_X1_STEP_KEY_SHARE:
-                    return PSA_JPAKE_X4S_STEP_KEY_SHARE;
-                case PSA_PAKE_X1_STEP_ZK_PUBLIC:
-                    return PSA_JPAKE_X4S_STEP_ZK_PUBLIC;
-                case PSA_PAKE_X1_STEP_ZK_PROOF:
-                    return PSA_JPAKE_X4S_STEP_ZK_PROOF;
-                default:
-                    return PSA_JPAKE_STEP_INVALID;
-            }
-            break;
-        default:
-            return PSA_JPAKE_STEP_INVALID;
+    psa_crypto_driver_pake_step_t key_share_step;
+    if (stage->round == PSA_JPAKE_FIRST) {
+        int is_x1;
+
+        if (stage->io_mode == PSA_JPAKE_OUTPUT) {
+            is_x1 = (stage->outputs < 1);
+        } else {
+            is_x1 = (stage->inputs < 1);
+        }
+
+        key_share_step = is_x1 ?
+                         PSA_JPAKE_X1_STEP_KEY_SHARE :
+                         PSA_JPAKE_X2_STEP_KEY_SHARE;
+    } else if (stage->round == PSA_JPAKE_SECOND) {
+        key_share_step = (stage->io_mode == PSA_JPAKE_OUTPUT) ?
+                         PSA_JPAKE_X2S_STEP_KEY_SHARE :
+                         PSA_JPAKE_X4S_STEP_KEY_SHARE;
+    } else {
+        return PSA_JPAKE_STEP_INVALID;
     }
-    return PSA_JPAKE_STEP_INVALID;
+    return key_share_step + stage->step - PSA_PAKE_STEP_KEY_SHARE;
 }
 #endif /* PSA_WANT_ALG_JPAKE */
 
@@ -7981,12 +7999,6 @@
 #if defined(PSA_WANT_ALG_JPAKE)
         if (operation->alg == PSA_ALG_JPAKE) {
             operation->stage = PSA_PAKE_OPERATION_STAGE_COMPUTATION;
-            psa_jpake_computation_stage_t *computation_stage =
-                &operation->computation_stage.jpake;
-            computation_stage->state = PSA_PAKE_STATE_READY;
-            computation_stage->sequence = PSA_PAKE_SEQ_INVALID;
-            computation_stage->input_step = PSA_PAKE_STEP_X1_X2;
-            computation_stage->output_step = PSA_PAKE_STEP_X1_X2;
         } else
 #endif /* PSA_WANT_ALG_JPAKE */
         {
@@ -7997,9 +8009,10 @@
 }
 
 #if defined(PSA_WANT_ALG_JPAKE)
-static psa_status_t psa_jpake_output_prologue(
+static psa_status_t psa_jpake_prologue(
     psa_pake_operation_t *operation,
-    psa_pake_step_t step)
+    psa_pake_step_t step,
+    psa_jpake_io_mode_t io_mode)
 {
     if (step != PSA_PAKE_STEP_KEY_SHARE &&
         step != PSA_PAKE_STEP_ZK_PUBLIC &&
@@ -8010,84 +8023,66 @@
     psa_jpake_computation_stage_t *computation_stage =
         &operation->computation_stage.jpake;
 
-    if (computation_stage->state == PSA_PAKE_STATE_INVALID) {
+    if (computation_stage->round != PSA_JPAKE_FIRST &&
+        computation_stage->round != PSA_JPAKE_SECOND) {
         return PSA_ERROR_BAD_STATE;
     }
 
-    if (computation_stage->state != PSA_PAKE_STATE_READY &&
-        computation_stage->state != PSA_PAKE_OUTPUT_X1_X2 &&
-        computation_stage->state != PSA_PAKE_OUTPUT_X2S) {
+    /* Check that the step we are given is the one we were expecting */
+    if (step != computation_stage->step) {
         return PSA_ERROR_BAD_STATE;
     }
 
-    if (computation_stage->state == PSA_PAKE_STATE_READY) {
-        if (step != PSA_PAKE_STEP_KEY_SHARE) {
-            return PSA_ERROR_BAD_STATE;
-        }
-
-        switch (computation_stage->output_step) {
-            case PSA_PAKE_STEP_X1_X2:
-                computation_stage->state = PSA_PAKE_OUTPUT_X1_X2;
-                break;
-            case PSA_PAKE_STEP_X2S:
-                computation_stage->state = PSA_PAKE_OUTPUT_X2S;
-                break;
-            default:
-                return PSA_ERROR_BAD_STATE;
-        }
-
-        computation_stage->sequence = PSA_PAKE_X1_STEP_KEY_SHARE;
-    }
-
-    /* Check if step matches current sequence */
-    switch (computation_stage->sequence) {
-        case PSA_PAKE_X1_STEP_KEY_SHARE:
-        case PSA_PAKE_X2_STEP_KEY_SHARE:
-            if (step != PSA_PAKE_STEP_KEY_SHARE) {
-                return PSA_ERROR_BAD_STATE;
-            }
-            break;
-
-        case PSA_PAKE_X1_STEP_ZK_PUBLIC:
-        case PSA_PAKE_X2_STEP_ZK_PUBLIC:
-            if (step != PSA_PAKE_STEP_ZK_PUBLIC) {
-                return PSA_ERROR_BAD_STATE;
-            }
-            break;
-
-        case PSA_PAKE_X1_STEP_ZK_PROOF:
-        case PSA_PAKE_X2_STEP_ZK_PROOF:
-            if (step != PSA_PAKE_STEP_ZK_PROOF) {
-                return PSA_ERROR_BAD_STATE;
-            }
-            break;
-
-        default:
-            return PSA_ERROR_BAD_STATE;
+    if (step == PSA_PAKE_STEP_KEY_SHARE &&
+        computation_stage->inputs == 0 &&
+        computation_stage->outputs == 0) {
+        /* Start of the round, so function decides whether we are inputting
+         * or outputting */
+        computation_stage->io_mode = io_mode;
+    } else if (computation_stage->io_mode != io_mode) {
+        /* Middle of the round so the mode we are in must match the function
+         * called by the user */
+        return PSA_ERROR_BAD_STATE;
     }
 
     return PSA_SUCCESS;
 }
 
-static psa_status_t psa_jpake_output_epilogue(
-    psa_pake_operation_t *operation)
+static psa_status_t psa_jpake_epilogue(
+    psa_pake_operation_t *operation,
+    psa_jpake_io_mode_t io_mode)
 {
-    psa_jpake_computation_stage_t *computation_stage =
+    psa_jpake_computation_stage_t *stage =
         &operation->computation_stage.jpake;
 
-    if ((computation_stage->state == PSA_PAKE_OUTPUT_X1_X2 &&
-         computation_stage->sequence == PSA_PAKE_X2_STEP_ZK_PROOF) ||
-        (computation_stage->state == PSA_PAKE_OUTPUT_X2S &&
-         computation_stage->sequence == PSA_PAKE_X1_STEP_ZK_PROOF)) {
-        computation_stage->state = PSA_PAKE_STATE_READY;
-        computation_stage->output_step++;
-        computation_stage->sequence = PSA_PAKE_SEQ_INVALID;
+    if (stage->step == PSA_PAKE_STEP_ZK_PROOF) {
+        /* End of an input/output */
+        if (io_mode == PSA_JPAKE_INPUT) {
+            stage->inputs++;
+            if (stage->inputs == PSA_JPAKE_EXPECTED_INPUTS(stage->round)) {
+                stage->io_mode = PSA_JPAKE_OUTPUT;
+            }
+        }
+        if (io_mode == PSA_JPAKE_OUTPUT) {
+            stage->outputs++;
+            if (stage->outputs == PSA_JPAKE_EXPECTED_OUTPUTS(stage->round)) {
+                stage->io_mode = PSA_JPAKE_INPUT;
+            }
+        }
+        if (stage->inputs == PSA_JPAKE_EXPECTED_INPUTS(stage->round) &&
+            stage->outputs == PSA_JPAKE_EXPECTED_OUTPUTS(stage->round)) {
+            /* End of a round, move to the next round */
+            stage->inputs = 0;
+            stage->outputs = 0;
+            stage->round++;
+        }
+        stage->step = PSA_PAKE_STEP_KEY_SHARE;
     } else {
-        computation_stage->sequence++;
+        stage->step++;
     }
-
     return PSA_SUCCESS;
 }
+
 #endif /* PSA_WANT_ALG_JPAKE */
 
 psa_status_t psa_pake_output(
@@ -8121,7 +8116,7 @@
     switch (operation->alg) {
 #if defined(PSA_WANT_ALG_JPAKE)
         case PSA_ALG_JPAKE:
-            status = psa_jpake_output_prologue(operation, step);
+            status = psa_jpake_prologue(operation, step, PSA_JPAKE_OUTPUT);
             if (status != PSA_SUCCESS) {
                 goto exit;
             }
@@ -8145,7 +8140,7 @@
     switch (operation->alg) {
 #if defined(PSA_WANT_ALG_JPAKE)
         case PSA_ALG_JPAKE:
-            status = psa_jpake_output_epilogue(operation);
+            status = psa_jpake_epilogue(operation, PSA_JPAKE_OUTPUT);
             if (status != PSA_SUCCESS) {
                 goto exit;
             }
@@ -8162,100 +8157,6 @@
     return status;
 }
 
-#if defined(PSA_WANT_ALG_JPAKE)
-static psa_status_t psa_jpake_input_prologue(
-    psa_pake_operation_t *operation,
-    psa_pake_step_t step)
-{
-    if (step != PSA_PAKE_STEP_KEY_SHARE &&
-        step != PSA_PAKE_STEP_ZK_PUBLIC &&
-        step != PSA_PAKE_STEP_ZK_PROOF) {
-        return PSA_ERROR_INVALID_ARGUMENT;
-    }
-
-    psa_jpake_computation_stage_t *computation_stage =
-        &operation->computation_stage.jpake;
-
-    if (computation_stage->state == PSA_PAKE_STATE_INVALID) {
-        return PSA_ERROR_BAD_STATE;
-    }
-
-    if (computation_stage->state != PSA_PAKE_STATE_READY &&
-        computation_stage->state != PSA_PAKE_INPUT_X1_X2 &&
-        computation_stage->state != PSA_PAKE_INPUT_X4S) {
-        return PSA_ERROR_BAD_STATE;
-    }
-
-    if (computation_stage->state == PSA_PAKE_STATE_READY) {
-        if (step != PSA_PAKE_STEP_KEY_SHARE) {
-            return PSA_ERROR_BAD_STATE;
-        }
-
-        switch (computation_stage->input_step) {
-            case PSA_PAKE_STEP_X1_X2:
-                computation_stage->state = PSA_PAKE_INPUT_X1_X2;
-                break;
-            case PSA_PAKE_STEP_X2S:
-                computation_stage->state = PSA_PAKE_INPUT_X4S;
-                break;
-            default:
-                return PSA_ERROR_BAD_STATE;
-        }
-
-        computation_stage->sequence = PSA_PAKE_X1_STEP_KEY_SHARE;
-    }
-
-    /* Check if step matches current sequence */
-    switch (computation_stage->sequence) {
-        case PSA_PAKE_X1_STEP_KEY_SHARE:
-        case PSA_PAKE_X2_STEP_KEY_SHARE:
-            if (step != PSA_PAKE_STEP_KEY_SHARE) {
-                return PSA_ERROR_BAD_STATE;
-            }
-            break;
-
-        case PSA_PAKE_X1_STEP_ZK_PUBLIC:
-        case PSA_PAKE_X2_STEP_ZK_PUBLIC:
-            if (step != PSA_PAKE_STEP_ZK_PUBLIC) {
-                return PSA_ERROR_BAD_STATE;
-            }
-            break;
-
-        case PSA_PAKE_X1_STEP_ZK_PROOF:
-        case PSA_PAKE_X2_STEP_ZK_PROOF:
-            if (step != PSA_PAKE_STEP_ZK_PROOF) {
-                return PSA_ERROR_BAD_STATE;
-            }
-            break;
-
-        default:
-            return PSA_ERROR_BAD_STATE;
-    }
-
-    return PSA_SUCCESS;
-}
-
-static psa_status_t psa_jpake_input_epilogue(
-    psa_pake_operation_t *operation)
-{
-    psa_jpake_computation_stage_t *computation_stage =
-        &operation->computation_stage.jpake;
-
-    if ((computation_stage->state == PSA_PAKE_INPUT_X1_X2 &&
-         computation_stage->sequence == PSA_PAKE_X2_STEP_ZK_PROOF) ||
-        (computation_stage->state == PSA_PAKE_INPUT_X4S &&
-         computation_stage->sequence == PSA_PAKE_X1_STEP_ZK_PROOF)) {
-        computation_stage->state = PSA_PAKE_STATE_READY;
-        computation_stage->input_step++;
-        computation_stage->sequence = PSA_PAKE_SEQ_INVALID;
-    } else {
-        computation_stage->sequence++;
-    }
-
-    return PSA_SUCCESS;
-}
-#endif /* PSA_WANT_ALG_JPAKE */
-
 psa_status_t psa_pake_input(
     psa_pake_operation_t *operation,
     psa_pake_step_t step,
@@ -8288,7 +8189,7 @@
     switch (operation->alg) {
 #if defined(PSA_WANT_ALG_JPAKE)
         case PSA_ALG_JPAKE:
-            status = psa_jpake_input_prologue(operation, step);
+            status = psa_jpake_prologue(operation, step, PSA_JPAKE_INPUT);
             if (status != PSA_SUCCESS) {
                 goto exit;
             }
@@ -8312,7 +8213,7 @@
     switch (operation->alg) {
 #if defined(PSA_WANT_ALG_JPAKE)
         case PSA_ALG_JPAKE:
-            status = psa_jpake_input_epilogue(operation);
+            status = psa_jpake_epilogue(operation, PSA_JPAKE_INPUT);
             if (status != PSA_SUCCESS) {
                 goto exit;
             }
@@ -8347,8 +8248,7 @@
     if (operation->alg == PSA_ALG_JPAKE) {
         psa_jpake_computation_stage_t *computation_stage =
             &operation->computation_stage.jpake;
-        if (computation_stage->input_step != PSA_PAKE_STEP_DERIVE ||
-            computation_stage->output_step != PSA_PAKE_STEP_DERIVE) {
+        if (computation_stage->round != PSA_JPAKE_FINISHED) {
             status = PSA_ERROR_BAD_STATE;
             goto exit;
         }
diff --git a/library/psa_crypto_ffdh.h b/library/psa_crypto_ffdh.h
index 5d7d951..5298f5a 100644
--- a/library/psa_crypto_ffdh.h
+++ b/library/psa_crypto_ffdh.h
@@ -47,8 +47,8 @@
  * \retval #PSA_ERROR_INVALID_ARGUMENT
  *         \p key_buffer_size, \p peer_key_length, \p shared_secret_size
  *         do not match
- * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY   \emptydescription
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED   \emptydescription
  */
 psa_status_t mbedtls_psa_key_agreement_ffdh(
     const psa_key_attributes_t *attributes,
@@ -73,9 +73,9 @@
  * \retval #PSA_SUCCESS  The public key was exported successfully.
  * \retval #PSA_ERROR_BUFFER_TOO_SMALL
  *         The size of \p key_buffer is too small.
- * \retval #PSA_ERROR_NOT_PERMITTED
- * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ * \retval #PSA_ERROR_NOT_PERMITTED         \emptydescription
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY   \emptydescription
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED   \emptydescription
  */
 psa_status_t mbedtls_psa_export_ffdh_public_key(
     const psa_key_attributes_t *attributes,
@@ -103,8 +103,8 @@
  *         Key size in bits is invalid.
  * \retval #PSA_ERROR_BUFFER_TOO_SMALL
  *         The size of \p key_buffer is too small.
- * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY   \emptydescription
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED   \emptydescription
  */
 psa_status_t mbedtls_psa_ffdh_generate_key(
     const psa_key_attributes_t *attributes,
diff --git a/library/psa_crypto_pake.c b/library/psa_crypto_pake.c
index 4136614..e22bcf8 100644
--- a/library/psa_crypto_pake.c
+++ b/library/psa_crypto_pake.c
@@ -80,65 +80,37 @@
  */
 
 /*
- * The first PAKE step shares the same sequences of the second PAKE step
- * but with a second set of KEY_SHARE/ZK_PUBLIC/ZK_PROOF outputs/inputs.
- * It's simpler to share the same sequences numbers of the first
- * set of KEY_SHARE/ZK_PUBLIC/ZK_PROOF outputs/inputs in both PAKE steps.
+ * Possible sequence of calls to implementation:
  *
- * State sequence with step, state & sequence enums:
- *   => Input & Output Step = PSA_PAKE_STEP_INVALID
- *   => state = PSA_PAKE_STATE_INVALID
- *   psa_pake_setup()
- *   => Input & Output Step = PSA_PAKE_STEP_X1_X2
- *   => state = PSA_PAKE_STATE_SETUP
- *   => sequence = PSA_PAKE_SEQ_INVALID
- *   |
- *   |--- In any order: (First round input before or after first round output)
- *   |   | First call of psa_pake_output() or psa_pake_input() sets
- *   |   | state = PSA_PAKE_STATE_READY
- *   |   |
- *   |   |------ In Order: => state = PSA_PAKE_OUTPUT_X1_X2
- *   |   |       | psa_pake_output() => sequence = PSA_PAKE_X1_STEP_KEY_SHARE
- *   |   |       | psa_pake_output() => sequence = PSA_PAKE_X1_STEP_ZK_PUBLIC
- *   |   |       | psa_pake_output() => sequence = PSA_PAKE_X1_STEP_ZK_PROOF
- *   |   |       | psa_pake_output() => sequence = PSA_PAKE_X2_STEP_KEY_SHARE
- *   |   |       | psa_pake_output() => sequence = PSA_PAKE_X2_STEP_ZK_PUBLIC
- *   |   |       | psa_pake_output() => sequence = PSA_PAKE_X2_STEP_ZK_PROOF
- *   |   |       | => state = PSA_PAKE_STATE_READY
- *   |   |       | => sequence = PSA_PAKE_SEQ_INVALID
- *   |   |       | => Output Step = PSA_PAKE_STEP_X2S
- *   |   |
- *   |   |------ In Order: => state = PSA_PAKE_INPUT_X1_X2
- *   |   |       | psa_pake_input() => sequence = PSA_PAKE_X1_STEP_KEY_SHARE
- *   |   |       | psa_pake_input() => sequence = PSA_PAKE_X1_STEP_ZK_PUBLIC
- *   |   |       | psa_pake_input() => sequence = PSA_PAKE_X1_STEP_ZK_PROOF
- *   |   |       | psa_pake_input() => sequence = PSA_PAKE_X2_STEP_KEY_SHARE
- *   |   |       | psa_pake_input() => sequence = PSA_PAKE_X2_STEP_ZK_PUBLIC
- *   |   |       | psa_pake_input() => sequence = PSA_PAKE_X2_STEP_ZK_PROOF
- *   |   |       | => state = PSA_PAKE_STATE_READY
- *   |   |       | => sequence = PSA_PAKE_SEQ_INVALID
- *   |   |       | => Output Step = PSA_PAKE_INPUT_X4S
- *   |
- *   |--- In any order: (Second round input before or after second round output)
- *   |   |
- *   |   |------ In Order: => state = PSA_PAKE_OUTPUT_X2S
- *   |   |       | psa_pake_output() => sequence = PSA_PAKE_X1_STEP_KEY_SHARE
- *   |   |       | psa_pake_output() => sequence = PSA_PAKE_X1_STEP_ZK_PUBLIC
- *   |   |       | psa_pake_output() => sequence = PSA_PAKE_X1_STEP_ZK_PROOF
- *   |   |       | => state = PSA_PAKE_STATE_READY
- *   |   |       | => sequence = PSA_PAKE_SEQ_INVALID
- *   |   |       | => Output Step = PSA_PAKE_STEP_DERIVE
- *   |   |
- *   |   |------ In Order: => state = PSA_PAKE_INPUT_X4S
- *   |   |       | psa_pake_input() => sequence = PSA_PAKE_X1_STEP_KEY_SHARE
- *   |   |       | psa_pake_input() => sequence = PSA_PAKE_X1_STEP_ZK_PUBLIC
- *   |   |       | psa_pake_input() => sequence = PSA_PAKE_X1_STEP_ZK_PROOF
- *   |   |       | => state = PSA_PAKE_STATE_READY
- *   |   |       | => sequence = PSA_PAKE_SEQ_INVALID
- *   |   |       | => Output Step = PSA_PAKE_STEP_DERIVE
- *   |
- *   psa_pake_get_implicit_key()
- *   => Input & Output Step = PSA_PAKE_STEP_INVALID
+ * |--- In any order:
+ * |   |
+ * |   |------ In Order
+ * |   |       | mbedtls_psa_pake_output(PSA_JPAKE_X1_STEP_KEY_SHARE)
+ * |   |       | mbedtls_psa_pake_output(PSA_JPAKE_X1_STEP_ZK_PUBLIC)
+ * |   |       | mbedtls_psa_pake_output(PSA_JPAKE_X1_STEP_ZK_PROOF)
+ * |   |       | mbedtls_psa_pake_output(PSA_JPAKE_X2_STEP_KEY_SHARE)
+ * |   |       | mbedtls_psa_pake_output(PSA_JPAKE_X2_STEP_ZK_PUBLIC)
+ * |   |       | mbedtls_psa_pake_output(PSA_JPAKE_X2_STEP_ZK_PROOF)
+ * |   |
+ * |   |------ In Order:
+ * |           | mbedtls_psa_pake_input(PSA_JPAKE_X1_STEP_KEY_SHARE)
+ * |           | mbedtls_psa_pake_input(PSA_JPAKE_X1_STEP_ZK_PUBLIC)
+ * |           | mbedtls_psa_pake_input(PSA_JPAKE_X1_STEP_ZK_PROOF)
+ * |           | mbedtls_psa_pake_input(PSA_JPAKE_X2_STEP_KEY_SHARE)
+ * |           | mbedtls_psa_pake_input(PSA_JPAKE_X2_STEP_ZK_PUBLIC)
+ * |           | mbedtls_psa_pake_input(PSA_JPAKE_X2_STEP_ZK_PROOF)
+ * |
+ * |--- In any order:
+ * |   |
+ * |   |------ In Order
+ * |   |       | mbedtls_psa_pake_output(PSA_JPAKE_X2S_STEP_KEY_SHARE)
+ * |   |       | mbedtls_psa_pake_output(PSA_JPAKE_X2S_STEP_ZK_PUBLIC)
+ * |   |       | mbedtls_psa_pake_output(PSA_JPAKE_X2S_STEP_ZK_PROOF)
+ * |   |
+ * |   |------ In Order:
+ * |           | mbedtls_psa_pake_input(PSA_JPAKE_X4S_STEP_KEY_SHARE)
+ * |           | mbedtls_psa_pake_input(PSA_JPAKE_X4S_STEP_ZK_PUBLIC)
+ * |           | mbedtls_psa_pake_input(PSA_JPAKE_X4S_STEP_ZK_PROOF)
  */
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
diff --git a/library/psa_crypto_pake.h b/library/psa_crypto_pake.h
index 001c987..f21b0e6 100644
--- a/library/psa_crypto_pake.h
+++ b/library/psa_crypto_pake.h
@@ -43,8 +43,8 @@
  *         compatible with the PAKE algorithm, or the hash algorithm in
  *         \p cipher_suite is not supported or not compatible with the PAKE
  *         algorithm and primitive.
- * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY   \emptydescription
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED   \emptydescription
  */
 psa_status_t mbedtls_psa_pake_setup(mbedtls_psa_pake_operation_t *operation,
                                     const psa_crypto_driver_pake_inputs_t *inputs);
@@ -78,10 +78,10 @@
  *         Success.
  * \retval #PSA_ERROR_BUFFER_TOO_SMALL
  *         The size of the \p output buffer is too small.
- * \retval #PSA_ERROR_INSUFFICIENT_ENTROPY
- * \retval #PSA_ERROR_CORRUPTION_DETECTED
- * \retval #PSA_ERROR_DATA_CORRUPT
- * \retval #PSA_ERROR_DATA_INVALID
+ * \retval #PSA_ERROR_INSUFFICIENT_ENTROPY  \emptydescription
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED   \emptydescription
+ * \retval #PSA_ERROR_DATA_CORRUPT          \emptydescription
+ * \retval #PSA_ERROR_DATA_INVALID          \emptydescription
  */
 psa_status_t mbedtls_psa_pake_output(mbedtls_psa_pake_operation_t *operation,
                                      psa_crypto_driver_pake_step_t step,
@@ -116,10 +116,10 @@
  * \retval #PSA_ERROR_NOT_SUPPORTED
  *         the \p input is not supported for the \p operation's algorithm, cipher
  *         suite or \p step.
- * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- * \retval #PSA_ERROR_CORRUPTION_DETECTED
- * \retval #PSA_ERROR_DATA_CORRUPT
- * \retval #PSA_ERROR_DATA_INVALID
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY   \emptydescription
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED   \emptydescription
+ * \retval #PSA_ERROR_DATA_CORRUPT          \emptydescription
+ * \retval #PSA_ERROR_DATA_INVALID          \emptydescription
  */
 psa_status_t mbedtls_psa_pake_input(mbedtls_psa_pake_operation_t *operation,
                                     psa_crypto_driver_pake_step_t step,
@@ -143,10 +143,10 @@
  * \retval #PSA_ERROR_NOT_SUPPORTED
  *         Input from a PAKE is not supported by the algorithm in the \p output
  *         key derivation operation.
- * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- * \retval #PSA_ERROR_CORRUPTION_DETECTED
- * \retval #PSA_ERROR_DATA_CORRUPT
- * \retval #PSA_ERROR_DATA_INVALID
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY   \emptydescription
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED   \emptydescription
+ * \retval #PSA_ERROR_DATA_CORRUPT          \emptydescription
+ * \retval #PSA_ERROR_DATA_INVALID          \emptydescription
  */
 psa_status_t mbedtls_psa_pake_get_implicit_key(
     mbedtls_psa_pake_operation_t *operation,
@@ -164,7 +164,7 @@
  *
  * \retval #PSA_SUCCESS
  *         Success.
- * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED   \emptydescription
  */
 psa_status_t mbedtls_psa_pake_abort(mbedtls_psa_pake_operation_t *operation);
 
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index a7cb9b5..a10cb2b 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -36,8 +36,6 @@
 #include <string.h>
 #include "mbedtls/platform.h"
 
-#define ARRAY_LENGTH(array) (sizeof(array) / sizeof(*(array)))
-
 typedef struct {
     psa_key_slot_t key_slots[MBEDTLS_PSA_KEY_SLOT_COUNT];
     unsigned key_slots_initialized : 1;
diff --git a/library/sha512.c b/library/sha512.c
index b8b2485..ff92a1b 100644
--- a/library/sha512.c
+++ b/library/sha512.c
@@ -1001,8 +1001,6 @@
 };
 #endif /* MBEDTLS_SHA512_C */
 
-#define ARRAY_LENGTH(a)   (sizeof(a) / sizeof((a)[0]))
-
 static int mbedtls_sha512_common_self_test(int verbose, int is384)
 {
     int i, buflen, ret = 0;
diff --git a/library/ssl_cookie.c b/library/ssl_cookie.c
index ae7a420..098aced 100644
--- a/library/ssl_cookie.c
+++ b/library/ssl_cookie.c
@@ -37,9 +37,15 @@
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 #include "md_psa.h"
-#define PSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,   \
-                                                           psa_to_ssl_errors,             \
-                                                           psa_generic_status_to_mbedtls)
+/* Define a local translating function to save code size by not using too many
+ * arguments in each translating place. */
+static int local_err_translation(psa_status_t status)
+{
+    return psa_status_to_mbedtls(status, psa_to_ssl_errors,
+                                 ARRAY_LENGTH(psa_to_ssl_errors),
+                                 psa_generic_status_to_mbedtls);
+}
+#define PSA_TO_MBEDTLS_ERR(status) local_err_translation(status)
 #endif
 
 /*
diff --git a/library/ssl_msg.c b/library/ssl_msg.c
index b5019ef..e07be05 100644
--- a/library/ssl_msg.c
+++ b/library/ssl_msg.c
@@ -49,9 +49,15 @@
 #endif
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
-#define PSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,   \
-                                                           psa_to_ssl_errors,             \
-                                                           psa_generic_status_to_mbedtls)
+/* Define a local translating function to save code size by not using too many
+ * arguments in each translating place. */
+static int local_err_translation(psa_status_t status)
+{
+    return psa_status_to_mbedtls(status, psa_to_ssl_errors,
+                                 ARRAY_LENGTH(psa_to_ssl_errors),
+                                 psa_generic_status_to_mbedtls);
+}
+#define PSA_TO_MBEDTLS_ERR(status) local_err_translation(status)
 #endif
 
 #if defined(MBEDTLS_SSL_SOME_SUITES_USE_MAC)
diff --git a/library/ssl_ticket.c b/library/ssl_ticket.c
index 7d07d19..1adaa07 100644
--- a/library/ssl_ticket.c
+++ b/library/ssl_ticket.c
@@ -31,9 +31,15 @@
 #include <string.h>
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
-#define PSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,   \
-                                                           psa_to_ssl_errors,             \
-                                                           psa_generic_status_to_mbedtls)
+/* Define a local translating function to save code size by not using too many
+ * arguments in each translating place. */
+static int local_err_translation(psa_status_t status)
+{
+    return psa_status_to_mbedtls(status, psa_to_ssl_errors,
+                                 ARRAY_LENGTH(psa_to_ssl_errors),
+                                 psa_generic_status_to_mbedtls);
+}
+#define PSA_TO_MBEDTLS_ERR(status) local_err_translation(status)
 #endif
 
 /*
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index f0067f4..bc9f4f8 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -51,12 +51,15 @@
 #endif
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
-#define PSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status, \
-                                                           psa_to_ssl_errors, \
-                                                           psa_generic_status_to_mbedtls)
-#define PSA_TO_MD_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status, \
-                                                      psa_to_md_errors, \
-                                                      psa_generic_status_to_mbedtls)
+/* Define local translating functions to save code size by not using too many
+ * arguments in each translating place. */
+static int local_err_translation(psa_status_t status)
+{
+    return psa_status_to_mbedtls(status, psa_to_ssl_errors,
+                                 ARRAY_LENGTH(psa_to_ssl_errors),
+                                 psa_generic_status_to_mbedtls);
+}
+#define PSA_TO_MBEDTLS_ERR(status) local_err_translation(status)
 #endif
 
 #if defined(MBEDTLS_TEST_HOOKS)
@@ -748,8 +751,6 @@
 }
 
 #if defined(MBEDTLS_SSL_PROTO_TLS1_3) && defined(MBEDTLS_SSL_SESSION_TICKETS)
-#define ARRAY_LENGTH(a) (sizeof(a) / sizeof(*(a)))
-
 static const char *ticket_flag_name_table[] =
 {
     [0] = "ALLOW_PSK_RESUMPTION",
diff --git a/library/ssl_tls12_client.c b/library/ssl_tls12_client.c
index fc96dae..28f9cdb 100644
--- a/library/ssl_tls12_client.c
+++ b/library/ssl_tls12_client.c
@@ -33,9 +33,17 @@
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 #include "mbedtls/psa_util.h"
 #include "psa/crypto.h"
-#define PSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,   \
-                                                           psa_to_ssl_errors,             \
-                                                           psa_generic_status_to_mbedtls)
+#if defined(MBEDTLS_KEY_EXCHANGE_ECDHE_PSK_ENABLED)
+/* Define a local translating function to save code size by not using too many
+ * arguments in each translating place. */
+static int local_err_translation(psa_status_t status)
+{
+    return psa_status_to_mbedtls(status, psa_to_ssl_errors,
+                                 ARRAY_LENGTH(psa_to_ssl_errors),
+                                 psa_generic_status_to_mbedtls);
+}
+#define PSA_TO_MBEDTLS_ERR(status) local_err_translation(status)
+#endif /* MBEDTLS_KEY_EXCHANGE_ECDHE_PSK_ENABLED */
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
 
 #include <string.h>
diff --git a/library/ssl_tls12_server.c b/library/ssl_tls12_server.c
index 5f5a433..006ce17 100644
--- a/library/ssl_tls12_server.c
+++ b/library/ssl_tls12_server.c
@@ -34,9 +34,18 @@
 #include <string.h>
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
-#define PSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,   \
-                                                           psa_to_ssl_errors,             \
-                                                           psa_generic_status_to_mbedtls)
+/* Define a local translating function to save code size by not using too many
+ * arguments in each translating place. */
+#if defined(MBEDTLS_KEY_EXCHANGE_SOME_ECDH_ENABLED) || \
+    defined(MBEDTLS_KEY_EXCHANGE_SOME_ECDHE_ENABLED)
+static int local_err_translation(psa_status_t status)
+{
+    return psa_status_to_mbedtls(status, psa_to_ssl_errors,
+                                 ARRAY_LENGTH(psa_to_ssl_errors),
+                                 psa_generic_status_to_mbedtls);
+}
+#define PSA_TO_MBEDTLS_ERR(status) local_err_translation(status)
+#endif
 #endif
 
 #if defined(MBEDTLS_ECP_C)
@@ -2589,14 +2598,17 @@
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-    unsigned char buf[
-        PSA_KEY_EXPORT_ECC_KEY_PAIR_MAX_SIZE(PSA_VENDOR_ECC_MAX_CURVE_BITS)];
+    mbedtls_pk_context *pk;
+    mbedtls_pk_type_t pk_type;
     psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;
+#if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
     uint16_t tls_id = 0;
     psa_ecc_family_t ecc_family;
     size_t key_len;
-    mbedtls_pk_context *pk;
     mbedtls_ecp_group_id grp_id;
+    unsigned char buf[PSA_KEY_EXPORT_ECC_KEY_PAIR_MAX_SIZE(PSA_VENDOR_ECC_MAX_CURVE_BITS)];
+    mbedtls_ecp_keypair *key;
+#endif /* !MBEDTLS_PK_USE_PSA_EC_DATA */
 
     pk = mbedtls_ssl_own_key(ssl);
 
@@ -2604,18 +2616,20 @@
         return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
     }
 
-#if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
-    mbedtls_ecp_keypair *key = mbedtls_pk_ec_rw(*pk);
-#endif /* !MBEDTLS_PK_USE_PSA_EC_DATA */
+    pk_type = mbedtls_pk_get_type(pk);
 
-    switch (mbedtls_pk_get_type(pk)) {
+    switch (pk_type) {
         case MBEDTLS_PK_OPAQUE:
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+        case MBEDTLS_PK_ECKEY:
+        case MBEDTLS_PK_ECKEY_DH:
+        case MBEDTLS_PK_ECDSA:
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
             if (!mbedtls_pk_can_do(pk, MBEDTLS_PK_ECKEY)) {
                 return MBEDTLS_ERR_SSL_PK_TYPE_MISMATCH;
             }
 
             ssl->handshake->ecdh_psa_privkey = pk->priv_id;
-
             /* Key should not be destroyed in the TLS library */
             ssl->handshake->ecdh_psa_privkey_is_external = 1;
 
@@ -2633,9 +2647,11 @@
 
             ret = 0;
             break;
+#if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
         case MBEDTLS_PK_ECKEY:
         case MBEDTLS_PK_ECKEY_DH:
         case MBEDTLS_PK_ECDSA:
+            key = mbedtls_pk_ec_rw(*pk);
             grp_id = mbedtls_pk_get_group_id(pk);
             if (grp_id == MBEDTLS_ECP_DP_NONE) {
                 return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
@@ -2660,36 +2676,29 @@
                              PSA_KEY_TYPE_ECC_KEY_PAIR(ssl->handshake->ecdh_psa_type));
             psa_set_key_bits(&key_attributes, ssl->handshake->ecdh_bits);
 
-#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
-            status = psa_export_key(pk->priv_id, buf, sizeof(buf), &key_len);
-            if (status != PSA_SUCCESS) {
-                ret = PSA_TO_MBEDTLS_ERR(status);
-                goto cleanup;
-            }
-#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
             key_len = PSA_BITS_TO_BYTES(key->grp.pbits);
             ret = mbedtls_ecp_write_key(key, buf, key_len);
             if (ret != 0) {
-                goto cleanup;
+                mbedtls_platform_zeroize(buf, sizeof(buf));
+                break;
             }
-#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 
             status = psa_import_key(&key_attributes, buf, key_len,
                                     &ssl->handshake->ecdh_psa_privkey);
             if (status != PSA_SUCCESS) {
                 ret = PSA_TO_MBEDTLS_ERR(status);
-                goto cleanup;
+                mbedtls_platform_zeroize(buf, sizeof(buf));
+                break;
             }
 
+            mbedtls_platform_zeroize(buf, sizeof(buf));
             ret = 0;
             break;
+#endif /* !MBEDTLS_PK_USE_PSA_EC_DATA */
         default:
             ret = MBEDTLS_ERR_SSL_PK_TYPE_MISMATCH;
     }
 
-cleanup:
-    mbedtls_platform_zeroize(buf, sizeof(buf));
-
     return ret;
 }
 #elif defined(MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED) || \
diff --git a/library/ssl_tls13_client.c b/library/ssl_tls13_client.c
index 3dffc1d..6ec3170 100644
--- a/library/ssl_tls13_client.c
+++ b/library/ssl_tls13_client.c
@@ -35,9 +35,17 @@
 #include "ssl_debug_helpers.h"
 #include "md_psa.h"
 
-#define PSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,   \
-                                                           psa_to_ssl_errors,             \
-                                                           psa_generic_status_to_mbedtls)
+#if defined(PSA_WANT_ALG_ECDH)
+/* Define a local translating function to save code size by not using too many
+ * arguments in each translating place. */
+static int local_err_translation(psa_status_t status)
+{
+    return psa_status_to_mbedtls(status, psa_to_ssl_errors,
+                                 ARRAY_LENGTH(psa_to_ssl_errors),
+                                 psa_generic_status_to_mbedtls);
+}
+#define PSA_TO_MBEDTLS_ERR(status) local_err_translation(status)
+#endif
 
 /* Write extensions */
 
diff --git a/library/ssl_tls13_generic.c b/library/ssl_tls13_generic.c
index a59f01c..fa193ff 100644
--- a/library/ssl_tls13_generic.c
+++ b/library/ssl_tls13_generic.c
@@ -39,9 +39,18 @@
 #include "psa/crypto.h"
 #include "mbedtls/psa_util.h"
 
-#define PSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,   \
-                                                           psa_to_ssl_errors,             \
-                                                           psa_generic_status_to_mbedtls)
+#if defined(MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_EPHEMERAL_ENABLED) || \
+    defined(PSA_WANT_ALG_ECDH)
+/* Define a local translating function to save code size by not using too many
+ * arguments in each translating place. */
+static int local_err_translation(psa_status_t status)
+{
+    return psa_status_to_mbedtls(status, psa_to_ssl_errors,
+                                 ARRAY_LENGTH(psa_to_ssl_errors),
+                                 psa_generic_status_to_mbedtls);
+}
+#define PSA_TO_MBEDTLS_ERR(status) local_err_translation(status)
+#endif
 
 const uint8_t mbedtls_ssl_tls13_hello_retry_request_magic[
     MBEDTLS_SERVER_HELLO_RANDOM_LEN] =
diff --git a/library/ssl_tls13_keys.c b/library/ssl_tls13_keys.c
index 540f854..81daf0a 100644
--- a/library/ssl_tls13_keys.c
+++ b/library/ssl_tls13_keys.c
@@ -36,9 +36,15 @@
 #include "psa/crypto.h"
 #include "md_psa.h"
 
-#define PSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,   \
-                                                           psa_to_ssl_errors,             \
-                                                           psa_generic_status_to_mbedtls)
+/* Define a local translating function to save code size by not using too many
+ * arguments in each translating place. */
+static int local_err_translation(psa_status_t status)
+{
+    return psa_status_to_mbedtls(status, psa_to_ssl_errors,
+                                 ARRAY_LENGTH(psa_to_ssl_errors),
+                                 psa_generic_status_to_mbedtls);
+}
+#define PSA_TO_MBEDTLS_ERR(status) local_err_translation(status)
 
 #define MBEDTLS_SSL_TLS1_3_LABEL(name, string)       \
     .name = string,
diff --git a/library/x509.c b/library/x509.c
index 6e16c4c..b600f45 100644
--- a/library/x509.c
+++ b/library/x509.c
@@ -53,13 +53,17 @@
 #include <time.h>
 #endif
 
-#define CHECK(code) if ((ret = (code)) != 0) { return ret; }
+#define CHECK(code)                                     \
+    do {                                                \
+        if ((ret = (code)) != 0) {                      \
+            return ret;                                 \
+        }                                               \
+    } while (0)
+
 #define CHECK_RANGE(min, max, val)                      \
-    do                                                  \
-    {                                                   \
-        if ((val) < (min) || (val) > (max))    \
-        {                                               \
-            return ret;                              \
+    do {                                                \
+        if ((val) < (min) || (val) > (max)) {           \
+            return ret;                                 \
         }                                               \
     } while (0)
 
@@ -1700,16 +1704,19 @@
     return 0;
 }
 
-#define PRINT_ITEM(i)                           \
-    {                                           \
-        ret = mbedtls_snprintf(p, n, "%s" i, sep);    \
-        MBEDTLS_X509_SAFE_SNPRINTF;                        \
-        sep = ", ";                             \
-    }
+#define PRINT_ITEM(i)                                   \
+    do {                                                \
+        ret = mbedtls_snprintf(p, n, "%s" i, sep);      \
+        MBEDTLS_X509_SAFE_SNPRINTF;                     \
+        sep = ", ";                                     \
+    } while (0)
 
-#define CERT_TYPE(type, name)                    \
-    if (ns_cert_type & (type))                 \
-    PRINT_ITEM(name);
+#define CERT_TYPE(type, name)                           \
+    do {                                                \
+        if (ns_cert_type & (type)) {                    \
+            PRINT_ITEM(name);                           \
+        }                                               \
+    } while (0)
 
 int mbedtls_x509_info_cert_type(char **buf, size_t *size,
                                 unsigned char ns_cert_type)
@@ -1734,9 +1741,12 @@
     return 0;
 }
 
-#define KEY_USAGE(code, name)    \
-    if (key_usage & (code))    \
-    PRINT_ITEM(name);
+#define KEY_USAGE(code, name)       \
+    do {                            \
+        if ((key_usage) & (code)) { \
+            PRINT_ITEM(name);       \
+        }                           \
+    } while (0)
 
 int mbedtls_x509_info_key_usage(char **buf, size_t *size,
                                 unsigned int key_usage)
diff --git a/library/x509_crt.c b/library/x509_crt.c
index 9b3414a..4508e50 100644
--- a/library/x509_crt.c
+++ b/library/x509_crt.c
@@ -49,7 +49,6 @@
 #include "mbedtls/psa_util.h"
 #include "md_psa.h"
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
-#include "x509_invasive.h"
 #include "pk_internal.h"
 
 #include "mbedtls/platform.h"
@@ -107,7 +106,7 @@
     MBEDTLS_X509_ID_FLAG(MBEDTLS_MD_SHA384) |
     MBEDTLS_X509_ID_FLAG(MBEDTLS_MD_SHA512),
     0xFFFFFFF, /* Any PK alg    */
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     /* Curves at or above 128-bit security level. Note that this selection
      * should be aligned with ssl_preset_default_curves in ssl_tls.c. */
     MBEDTLS_X509_ID_FLAG(MBEDTLS_ECP_DP_SECP256R1) |
@@ -117,9 +116,9 @@
     MBEDTLS_X509_ID_FLAG(MBEDTLS_ECP_DP_BP384R1) |
     MBEDTLS_X509_ID_FLAG(MBEDTLS_ECP_DP_BP512R1) |
     0,
-#else /* MBEDTLS_ECP_LIGHT */
+#else /* MBEDTLS_PK_HAVE_ECC_KEYS */
     0,
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
     2048,
 };
 
@@ -158,13 +157,13 @@
     /* Only ECDSA */
     MBEDTLS_X509_ID_FLAG(MBEDTLS_PK_ECDSA) |
     MBEDTLS_X509_ID_FLAG(MBEDTLS_PK_ECKEY),
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     /* Only NIST P-256 and P-384 */
     MBEDTLS_X509_ID_FLAG(MBEDTLS_ECP_DP_SECP256R1) |
     MBEDTLS_X509_ID_FLAG(MBEDTLS_ECP_DP_SECP384R1),
-#else /* MBEDTLS_ECP_LIGHT */
+#else /* MBEDTLS_PK_HAVE_ECC_KEYS */
     0,
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
     0,
 };
 
@@ -234,7 +233,7 @@
     }
 #endif /* MBEDTLS_RSA_C */
 
-#if defined(MBEDTLS_ECP_LIGHT)
+#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
     if (pk_alg == MBEDTLS_PK_ECDSA ||
         pk_alg == MBEDTLS_PK_ECKEY ||
         pk_alg == MBEDTLS_PK_ECKEY_DH) {
@@ -250,7 +249,7 @@
 
         return -1;
     }
-#endif /* MBEDTLS_ECP_LIGHT */
+#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
 
     return -1;
 }
@@ -2862,7 +2861,6 @@
 
 #endif /* !AF_INET6 || MBEDTLS_TEST_SW_INET_PTON */ //no-check-names
 
-MBEDTLS_STATIC_TESTABLE
 size_t mbedtls_x509_crt_parse_cn_inet_pton(const char *cn, void *dst)
 {
     return strchr(cn, ':') == NULL
diff --git a/library/x509_invasive.h b/library/x509_invasive.h
deleted file mode 100644
index d8fd74b..0000000
--- a/library/x509_invasive.h
+++ /dev/null
@@ -1,53 +0,0 @@
-/**
- * \file x509_invasive.h
- *
- * \brief x509 module: interfaces for invasive testing only.
- *
- * The interfaces in this file are intended for testing purposes only.
- * They SHOULD NOT be made available in library integrations except when
- * building the library for testing.
- */
-/*
- *  Copyright The Mbed TLS Contributors
- *  SPDX-License-Identifier: Apache-2.0
- *
- *  Licensed under the Apache License, Version 2.0 (the "License"); you may
- *  not use this file except in compliance with the License.
- *  You may obtain a copy of the License at
- *
- *  http://www.apache.org/licenses/LICENSE-2.0
- *
- *  Unless required by applicable law or agreed to in writing, software
- *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
- *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- *  See the License for the specific language governing permissions and
- *  limitations under the License.
- */
-
-#ifndef MBEDTLS_X509_INVASIVE_H
-#define MBEDTLS_X509_INVASIVE_H
-
-#include "common.h"
-
-#if defined(MBEDTLS_TEST_HOOKS)
-
-/**
- * \brief          This function parses a CN string as an IP address.
- *
- * \param cn       The CN string to parse. CN string MUST be NUL-terminated.
- * \param dst      The target buffer to populate with the binary IP address.
- *                 The buffer MUST be 16 bytes to save IPv6, and should be
- *                 4-byte aligned if the result will be used as struct in_addr.
- *                 e.g. uint32_t dst[4]
- *
- * \note           \cn is parsed as an IPv6 address if string contains ':',
- *                 else \cn is parsed as an IPv4 address.
- *
- * \return         Length of binary IP address; num bytes written to target.
- * \return         \c 0 on failure to parse CN string as an IP address.
- */
-size_t mbedtls_x509_crt_parse_cn_inet_pton(const char *cn, void *dst);
-
-#endif /* MBEDTLS_TEST_HOOKS */
-
-#endif /* MBEDTLS_X509_INVASIVE_H */