diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 39144a3..fccb800 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -749,8 +749,8 @@
     }
 
     /* On success, store the allocated export-formatted key. */
-    slot->data.key.data = output;
-    slot->data.key.bytes = data_length;
+    slot->key.data = output;
+    slot->key.bytes = data_length;
 
     return( PSA_SUCCESS );
 }
@@ -983,8 +983,8 @@
     }
 
     /* On success, store the allocated export-formatted key. */
-    slot->data.key.data = output;
-    slot->data.key.bytes = data_length;
+    slot->key.data = output;
+    slot->key.bytes = data_length;
 
     return( PSA_SUCCESS );
 }
@@ -1017,14 +1017,14 @@
 static psa_status_t psa_allocate_buffer_to_slot( psa_key_slot_t *slot,
                                                  size_t buffer_length )
 {
-    if( slot->data.key.data != NULL )
+    if( slot->key.data != NULL )
         return( PSA_ERROR_ALREADY_EXISTS );
 
-    slot->data.key.data = mbedtls_calloc( 1, buffer_length );
-    if( slot->data.key.data == NULL )
+    slot->key.data = mbedtls_calloc( 1, buffer_length );
+    if( slot->key.data == NULL )
         return( PSA_ERROR_INSUFFICIENT_MEMORY );
 
-    slot->data.key.bytes = buffer_length;
+    slot->key.bytes = buffer_length;
     return( PSA_SUCCESS );
 }
 
@@ -1037,7 +1037,7 @@
     if( status != PSA_SUCCESS )
         return( status );
 
-    memcpy( slot->data.key.data, data, data_length );
+    memcpy( slot->key.data, data, data_length );
     return( PSA_SUCCESS );
 }
 
@@ -1342,23 +1342,14 @@
 /** Wipe key data from a slot. Preserve metadata such as the policy. */
 static psa_status_t psa_remove_key_data_from_memory( psa_key_slot_t *slot )
 {
-#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-    if( psa_get_se_driver( slot->attr.lifetime, NULL, NULL ) &&
-        psa_key_slot_is_external( slot ) )
-    {
-        /* No key material to clean. */
-    }
-    else
-#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
-    {
-        /* Data pointer will always be either a valid pointer or NULL in an
-         * initialized slot, so we can just free it. */
-        if( slot->data.key.data != NULL )
-            mbedtls_platform_zeroize( slot->data.key.data, slot->data.key.bytes);
-        mbedtls_free( slot->data.key.data );
-        slot->data.key.data = NULL;
-        slot->data.key.bytes = 0;
-    }
+    /* Data pointer will always be either a valid pointer or NULL in an
+     * initialized slot, so we can just free it. */
+    if( slot->key.data != NULL )
+        mbedtls_platform_zeroize( slot->key.data, slot->key.bytes);
+
+    mbedtls_free( slot->key.data );
+    slot->key.data = NULL;
+    slot->key.bytes = 0;
 
     return( PSA_SUCCESS );
 }
@@ -1443,7 +1434,7 @@
          * three actions. */
         psa_crypto_prepare_transaction( PSA_CRYPTO_TRANSACTION_DESTROY_KEY );
         psa_crypto_transaction.key.lifetime = slot->attr.lifetime;
-        psa_crypto_transaction.key.slot = slot->data.se.slot_number;
+        psa_crypto_transaction.key.slot = psa_key_slot_get_slot_number( slot );
         psa_crypto_transaction.key.id = slot->attr.id;
         status = psa_crypto_save_transaction( );
         if( status != PSA_SUCCESS )
@@ -1460,7 +1451,8 @@
             goto exit;
         }
 
-        status = psa_destroy_se_key( driver, slot->data.se.slot_number );
+        status = psa_destroy_se_key( driver,
+                                     psa_key_slot_get_slot_number( slot ) );
         if( overall_status == PSA_SUCCESS )
             overall_status = status;
     }
@@ -1616,7 +1608,8 @@
 
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
     if( psa_key_slot_is_external( slot ) )
-        psa_set_key_slot_number( attributes, slot->data.se.slot_number );
+        psa_set_key_slot_number( attributes,
+                                 psa_key_slot_get_slot_number( slot ) );
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */
 
     switch( slot->attr.type )
@@ -1637,8 +1630,8 @@
                 mbedtls_rsa_context *rsa = NULL;
 
                 status = psa_load_rsa_representation( slot->attr.type,
-                                                      slot->data.key.data,
-                                                      slot->data.key.bytes,
+                                                      slot->key.data,
+                                                      slot->key.bytes,
                                                       &rsa );
                 if( status != PSA_SUCCESS )
                     break;
@@ -1684,12 +1677,12 @@
                                                     size_t data_size,
                                                     size_t *data_length )
 {
-    if( slot->data.key.bytes > data_size )
+    if( slot->key.bytes > data_size )
         return( PSA_ERROR_BUFFER_TOO_SMALL );
-    memcpy( data, slot->data.key.data, slot->data.key.bytes );
-    memset( data + slot->data.key.bytes, 0,
-            data_size - slot->data.key.bytes );
-    *data_length = slot->data.key.bytes;
+    memcpy( data, slot->key.data, slot->key.bytes );
+    memset( data + slot->key.bytes, 0,
+            data_size - slot->key.bytes );
+    *data_length = slot->key.bytes;
     return( PSA_SUCCESS );
 }
 
@@ -1727,7 +1720,7 @@
         if( method == NULL )
             return( PSA_ERROR_NOT_SUPPORTED );
         return( method( drv_context,
-                        slot->data.se.slot_number,
+                        psa_key_slot_get_slot_number( slot ),
                         data, data_size, data_length ) );
     }
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */
@@ -1768,8 +1761,8 @@
             mbedtls_rsa_context *rsa = NULL;
             status = psa_load_rsa_representation(
                                     slot->attr.type,
-                                    slot->data.key.data,
-                                    slot->data.key.bytes,
+                                    slot->key.data,
+                                    slot->key.bytes,
                                     &rsa );
             if( status != PSA_SUCCESS )
                 return( status );
@@ -1797,8 +1790,8 @@
             mbedtls_ecp_keypair *ecp = NULL;
             status = psa_load_ecp_representation(
                                     slot->attr.type,
-                                    slot->data.key.data,
-                                    slot->data.key.bytes,
+                                    slot->key.data,
+                                    slot->key.bytes,
                                     &ecp );
             if( status != PSA_SUCCESS )
                 return( status );
@@ -2071,8 +2064,9 @@
      * we can roll back to a state where the key doesn't exist. */
     if( *p_drv != NULL )
     {
+        psa_key_slot_number_t slot_number;
         status = psa_find_se_slot_for_key( attributes, method, *p_drv,
-                                           &slot->data.se.slot_number );
+                                           &slot_number );
         if( status != PSA_SUCCESS )
             return( status );
 
@@ -2080,7 +2074,7 @@
         {
             psa_crypto_prepare_transaction( PSA_CRYPTO_TRANSACTION_CREATE_KEY );
             psa_crypto_transaction.key.lifetime = slot->attr.lifetime;
-            psa_crypto_transaction.key.slot = slot->data.se.slot_number;
+            psa_crypto_transaction.key.slot = slot_number;
             psa_crypto_transaction.key.id = slot->attr.id;
             status = psa_crypto_save_transaction( );
             if( status != PSA_SUCCESS )
@@ -2089,6 +2083,9 @@
                 return( status );
             }
         }
+
+        status = psa_copy_key_material_into_slot(
+            slot, (uint8_t *)( &slot_number ), sizeof( slot_number ) );
     }
 
     if( *p_drv == NULL && method == PSA_KEY_CREATION_REGISTER )
@@ -2140,13 +2137,15 @@
         if( driver != NULL )
         {
             psa_se_key_data_storage_t data;
+            psa_key_slot_number_t slot_number =
+                psa_key_slot_get_slot_number( slot ) ;
+
 #if defined(static_assert)
-            static_assert( sizeof( slot->data.se.slot_number ) ==
+            static_assert( sizeof( slot_number ) ==
                            sizeof( data.slot_number ),
                            "Slot number size does not match psa_se_key_data_storage_t" );
 #endif
-            memcpy( &data.slot_number, &slot->data.se.slot_number,
-                    sizeof( slot->data.se.slot_number ) );
+            memcpy( &data.slot_number, &slot_number, sizeof( slot_number ) );
             status = psa_save_persistent_key( &slot->attr,
                                               (uint8_t*) &data,
                                               sizeof( data ) );
@@ -2157,8 +2156,8 @@
             /* Key material is saved in export representation in the slot, so
              * just pass the slot buffer for storage. */
             status = psa_save_persistent_key( &slot->attr,
-                                              slot->data.key.data,
-                                              slot->data.key.bytes );
+                                              slot->key.data,
+                                              slot->key.bytes );
         }
     }
 #endif /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
@@ -2264,8 +2263,8 @@
 
             psa_status_t status = psa_load_rsa_representation(
                                     slot->attr.type,
-                                    slot->data.key.data,
-                                    slot->data.key.bytes,
+                                    slot->key.data,
+                                    slot->key.bytes,
                                     &rsa );
             if( status != PSA_SUCCESS )
                 return( status );
@@ -2345,8 +2344,8 @@
         }
         status = drv->key_management->p_import(
             psa_get_se_driver_context( driver ),
-            slot->data.se.slot_number, attributes, data, data_length,
-            &bits );
+            psa_key_slot_get_slot_number( slot ),
+            attributes, data, data_length, &bits );
         if( status != PSA_SUCCESS )
             goto exit;
         if( bits > PSA_MAX_KEY_BITS )
@@ -2423,8 +2422,8 @@
                                            psa_key_slot_t *target )
 {
     psa_status_t status = psa_copy_key_material_into_slot( target,
-                                                           source->data.key.data,
-                                                           source->data.key.bytes );
+                                                           source->key.data,
+                                                           source->key.bytes );
     if( status != PSA_SUCCESS )
         return( status );
 
@@ -3236,7 +3235,7 @@
         return( ret );
 
     ret = mbedtls_cipher_cmac_starts( &operation->ctx.cmac,
-                                      slot->data.key.data,
+                                      slot->key.data,
                                       key_bits );
     return( ret );
 }
@@ -3382,8 +3381,8 @@
         }
 
         status = psa_hmac_setup_internal( &operation->ctx.hmac,
-                                          slot->data.key.data,
-                                          slot->data.key.bytes,
+                                          slot->key.data,
+                                          slot->key.bytes,
                                           hash_alg );
     }
     else
@@ -3969,8 +3968,8 @@
         mbedtls_rsa_context *rsa = NULL;
 
         status = psa_load_rsa_representation( slot->attr.type,
-                                              slot->data.key.data,
-                                              slot->data.key.bytes,
+                                              slot->key.data,
+                                              slot->key.bytes,
                                               &rsa );
         if( status != PSA_SUCCESS )
             goto exit;
@@ -4001,8 +4000,8 @@
         {
             mbedtls_ecp_keypair *ecp = NULL;
             status = psa_load_ecp_representation( slot->attr.type,
-                                                  slot->data.key.data,
-                                                  slot->data.key.bytes,
+                                                  slot->key.data,
+                                                  slot->key.bytes,
                                                   &ecp );
             if( status != PSA_SUCCESS )
                 goto exit;
@@ -4079,8 +4078,8 @@
         mbedtls_rsa_context *rsa = NULL;
 
         status = psa_load_rsa_representation( slot->attr.type,
-                                              slot->data.key.data,
-                                              slot->data.key.bytes,
+                                              slot->key.data,
+                                              slot->key.bytes,
                                               &rsa );
         if( status != PSA_SUCCESS )
             goto exit;
@@ -4104,8 +4103,8 @@
         {
             mbedtls_ecp_keypair *ecp = NULL;
             status = psa_load_ecp_representation( slot->attr.type,
-                                                  slot->data.key.data,
-                                                  slot->data.key.bytes,
+                                                  slot->key.data,
+                                                  slot->key.bytes,
                                                   &ecp );
             if( status != PSA_SUCCESS )
                 goto exit;
@@ -4188,8 +4187,8 @@
     {
         mbedtls_rsa_context *rsa = NULL;
         status = psa_load_rsa_representation( slot->attr.type,
-                                              slot->data.key.data,
-                                              slot->data.key.bytes,
+                                              slot->key.data,
+                                              slot->key.bytes,
                                               &rsa );
         if( status != PSA_SUCCESS )
             goto rsa_exit;
@@ -4294,8 +4293,8 @@
     {
         mbedtls_rsa_context *rsa = NULL;
         status = psa_load_rsa_representation( slot->attr.type,
-                                              slot->data.key.data,
-                                              slot->data.key.bytes,
+                                              slot->key.data,
+                                              slot->key.bytes,
                                               &rsa );
         if( status != PSA_SUCCESS )
             goto exit;
@@ -4455,8 +4454,8 @@
     {
         /* Two-key Triple-DES is 3-key Triple-DES with K1=K3 */
         uint8_t keys[24];
-        memcpy( keys, slot->data.key.data, 16 );
-        memcpy( keys + 16, slot->data.key.data, 8 );
+        memcpy( keys, slot->key.data, 16 );
+        memcpy( keys + 16, slot->key.data, 8 );
         ret = mbedtls_cipher_setkey( &operation->ctx.cipher,
                                      keys,
                                      192, cipher_operation );
@@ -4465,7 +4464,7 @@
 #endif
     {
         ret = mbedtls_cipher_setkey( &operation->ctx.cipher,
-                                     slot->data.key.data,
+                                     slot->key.data,
                                      (int) key_bits, cipher_operation );
     }
     if( ret != 0 )
@@ -4964,7 +4963,7 @@
             mbedtls_ccm_init( &operation->ctx.ccm );
             status = mbedtls_to_psa_error(
                 mbedtls_ccm_setkey( &operation->ctx.ccm, cipher_id,
-                                    operation->slot->data.key.data,
+                                    operation->slot->key.data,
                                     (unsigned int) key_bits ) );
             if( status != 0 )
                 goto cleanup;
@@ -4986,7 +4985,7 @@
             mbedtls_gcm_init( &operation->ctx.gcm );
             status = mbedtls_to_psa_error(
                 mbedtls_gcm_setkey( &operation->ctx.gcm, cipher_id,
-                                    operation->slot->data.key.data,
+                                    operation->slot->key.data,
                                     (unsigned int) key_bits ) );
             if( status != 0 )
                 goto cleanup;
@@ -5006,7 +5005,7 @@
             mbedtls_chachapoly_init( &operation->ctx.chachapoly );
             status = mbedtls_to_psa_error(
                 mbedtls_chachapoly_setkey( &operation->ctx.chachapoly,
-                                           operation->slot->data.key.data ) );
+                                           operation->slot->key.data ) );
             if( status != 0 )
                 goto cleanup;
             break;
@@ -6129,8 +6128,8 @@
 
     status = psa_key_derivation_input_internal( operation,
                                                 step, slot->attr.type,
-                                                slot->data.key.data,
-                                                slot->data.key.bytes );
+                                                slot->key.data,
+                                                slot->key.bytes );
 
     unlock_status = psa_unlock_key_slot( slot );
 
@@ -6215,8 +6214,8 @@
             mbedtls_ecp_keypair *ecp = NULL;
             psa_status_t status = psa_load_ecp_representation(
                                     private_key->attr.type,
-                                    private_key->data.key.data,
-                                    private_key->data.key.bytes,
+                                    private_key->key.data,
+                                    private_key->key.bytes,
                                     &ecp );
             if( status != PSA_SUCCESS )
                 return( status );
@@ -6558,16 +6557,16 @@
         if( status != PSA_SUCCESS )
             return( status );
 
-        status = psa_generate_random( slot->data.key.data,
-                                      slot->data.key.bytes );
+        status = psa_generate_random( slot->key.data,
+                                      slot->key.bytes );
         if( status != PSA_SUCCESS )
             return( status );
 
         slot->attr.bits = (psa_key_bits_t) bits;
 #if defined(MBEDTLS_DES_C)
         if( type == PSA_KEY_TYPE_DES )
-            psa_des_set_key_parity( slot->data.key.data,
-                                    slot->data.key.bytes );
+            psa_des_set_key_parity( slot->key.data,
+                                    slot->key.bytes );
 #endif /* MBEDTLS_DES_C */
     }
     else
@@ -6611,9 +6610,9 @@
 
         status = psa_export_rsa_key( type,
                                      &rsa,
-                                     slot->data.key.data,
+                                     slot->key.data,
                                      bytes,
-                                     &slot->data.key.bytes );
+                                     &slot->key.bytes );
         mbedtls_rsa_free( &rsa );
         if( status != PSA_SUCCESS )
             psa_remove_key_data_from_memory( slot );
@@ -6657,11 +6656,11 @@
         }
 
         status = mbedtls_to_psa_error(
-            mbedtls_ecp_write_key( &ecp, slot->data.key.data, bytes ) );
+            mbedtls_ecp_write_key( &ecp, slot->key.data, bytes ) );
 
         mbedtls_ecp_keypair_free( &ecp );
         if( status != PSA_SUCCESS ) {
-            memset( slot->data.key.data, 0, bytes );
+            memset( slot->key.data, 0, bytes );
             psa_remove_key_data_from_memory( slot );
         }
         return( status );