Simplify psa_pake_computation_stage_s structure

Signed-off-by: Przemek Stekiel <przemyslaw.stekiel@mobica.com>
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index 83c7e04..75ce1a3 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -1976,7 +1976,7 @@
 struct psa_pake_computation_stage_s {
     union {
         unsigned dummy;
-        psa_jpake_computation_stage_t MBEDTLS_PRIVATE(jpake_computation_stage);
+        psa_jpake_computation_stage_t MBEDTLS_PRIVATE(jpake);
     } MBEDTLS_PRIVATE(data);
 };
 
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 09d46ed..46d62b0 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -7203,7 +7203,7 @@
 
     if (operation->alg == PSA_ALG_JPAKE) {
         psa_jpake_computation_stage_t *computation_stage =
-            &operation->computation_stage.data.jpake_computation_stage;
+            &operation->computation_stage.data.jpake;
 
         computation_stage->state = PSA_PAKE_STATE_SETUP;
         computation_stage->sequence = PSA_PAKE_SEQ_INVALID;
@@ -7336,10 +7336,10 @@
 static psa_pake_driver_step_t convert_jpake_computation_stage_to_driver_step(
     psa_pake_computation_stage_t *stage)
 {
-    switch (stage->data.jpake_computation_stage.state) {
+    switch (stage->data.jpake.state) {
         case PSA_PAKE_OUTPUT_X1_X2:
         case PSA_PAKE_INPUT_X1_X2:
-            switch (stage->data.jpake_computation_stage.sequence) {
+            switch (stage->data.jpake.sequence) {
                 case PSA_PAKE_X1_STEP_KEY_SHARE:
                     return PSA_JPAKE_X1_STEP_KEY_SHARE;
                     break;
@@ -7363,7 +7363,7 @@
             }
             break;
         case PSA_PAKE_OUTPUT_X2S:
-            switch (stage->data.jpake_computation_stage.sequence) {
+            switch (stage->data.jpake.sequence) {
                 case PSA_PAKE_X1_STEP_KEY_SHARE:
                     return PSA_JPAKE_X2S_STEP_KEY_SHARE;
                     break;
@@ -7377,7 +7377,7 @@
             }
             break;
         case PSA_PAKE_INPUT_X4S:
-            switch (stage->data.jpake_computation_stage.sequence) {
+            switch (stage->data.jpake.sequence) {
                 case PSA_PAKE_X1_STEP_KEY_SHARE:
                     return PSA_JPAKE_X4S_STEP_KEY_SHARE;
                     break;
@@ -7400,7 +7400,7 @@
     psa_pake_operation_t *operation)
 {
     psa_jpake_computation_stage_t *computation_stage =
-        &operation->computation_stage.data.jpake_computation_stage;
+        &operation->computation_stage.data.jpake;
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     uint8_t *password = operation->data.inputs.password;
     size_t password_len = operation->data.inputs.password_len;
@@ -7439,7 +7439,7 @@
     psa_pake_step_t step)
 {
     psa_jpake_computation_stage_t *computation_stage =
-        &operation->computation_stage.data.jpake_computation_stage;
+        &operation->computation_stage.data.jpake;
 
     if (computation_stage->state == PSA_PAKE_STATE_INVALID) {
         return PSA_ERROR_BAD_STATE;
@@ -7510,7 +7510,7 @@
     psa_pake_operation_t *operation)
 {
     psa_jpake_computation_stage_t *computation_stage =
-        &operation->computation_stage.data.jpake_computation_stage;
+        &operation->computation_stage.data.jpake;
 
     if ((computation_stage->state == PSA_PAKE_OUTPUT_X1_X2 &&
          computation_stage->sequence == PSA_PAKE_X2_STEP_ZK_PROOF) ||
@@ -7598,7 +7598,7 @@
     size_t input_length)
 {
     psa_jpake_computation_stage_t *computation_stage =
-        &operation->computation_stage.data.jpake_computation_stage;
+        &operation->computation_stage.data.jpake;
 
     if (computation_stage->state == PSA_PAKE_STATE_INVALID) {
         return PSA_ERROR_BAD_STATE;
@@ -7675,7 +7675,7 @@
     psa_pake_operation_t *operation)
 {
     psa_jpake_computation_stage_t *computation_stage =
-        &operation->computation_stage.data.jpake_computation_stage;
+        &operation->computation_stage.data.jpake;
 
     if ((computation_stage->state == PSA_PAKE_INPUT_X1_X2 &&
          computation_stage->sequence == PSA_PAKE_X2_STEP_ZK_PROOF) ||
@@ -7762,7 +7762,7 @@
     uint8_t shared_key[MBEDTLS_PSA_PAKE_BUFFER_SIZE];
     size_t shared_key_len = 0;
     psa_jpake_computation_stage_t *computation_stage =
-        &operation->computation_stage.data.jpake_computation_stage;
+        &operation->computation_stage.data.jpake;
 
     if (operation->id == 0) {
         return PSA_ERROR_BAD_STATE;
@@ -7821,7 +7821,7 @@
 
     if (operation->alg == PSA_ALG_JPAKE) {
         psa_jpake_computation_stage_t *computation_stage =
-            &operation->computation_stage.data.jpake_computation_stage;
+            &operation->computation_stage.data.jpake;
 
         computation_stage->input_step = PSA_PAKE_STEP_INVALID;
         computation_stage->output_step = PSA_PAKE_STEP_INVALID;
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index f718349..0c44227 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -3119,10 +3119,8 @@
                        PSA_SUCCESS);
 
             /* Simulate that we are ready to get implicit key. */
-            operation.computation_stage.data.jpake_computation_stage.input_step =
-                PSA_PAKE_STEP_DERIVE;
-            operation.computation_stage.data.jpake_computation_stage.output_step =
-                PSA_PAKE_STEP_DERIVE;
+            operation.computation_stage.data.jpake.input_step = PSA_PAKE_STEP_DERIVE;
+            operation.computation_stage.data.jpake.output_step = PSA_PAKE_STEP_DERIVE;
 
             /* --- psa_pake_get_implicit_key --- */
             mbedtls_test_driver_pake_hooks.forced_status = forced_status;