diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index aef18ac..7bce391 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3709,35 +3709,34 @@
     return( status );
 }
 
-static psa_status_t psa_validate_tag_length( psa_aead_operation_t *operation,
-                                             psa_algorithm_t alg ) {
+static psa_status_t psa_validate_tag_length( psa_algorithm_t alg ) {
     const uint8_t tag_len = PSA_ALG_AEAD_GET_TAG_LENGTH( alg );
 
     switch( PSA_ALG_AEAD_WITH_SHORTENED_TAG( alg, 0 ) )
     {
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_CCM)
+#if defined(PSA_WANT_ALG_CCM)
         case PSA_ALG_AEAD_WITH_SHORTENED_TAG( PSA_ALG_CCM, 0 ):
             /* CCM allows the following tag lengths: 4, 6, 8, 10, 12, 14, 16.*/
             if( tag_len < 4 || tag_len > 16 || tag_len % 2 )
                 return( PSA_ERROR_INVALID_ARGUMENT );
             break;
-#endif /* MBEDTLS_PSA_BUILTIN_ALG_CCM */
+#endif /* PSA_WANT_ALG_CCM */
 
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_GCM)
+#if defined(PSA_WANT_ALG_GCM)
         case PSA_ALG_AEAD_WITH_SHORTENED_TAG( PSA_ALG_GCM, 0 ):
             /* GCM allows the following tag lengths: 4, 8, 12, 13, 14, 15, 16. */
             if( tag_len != 4 && tag_len != 8 && ( tag_len < 12 || tag_len > 16 ) )
                 return( PSA_ERROR_INVALID_ARGUMENT );
             break;
-#endif /* MBEDTLS_PSA_BUILTIN_ALG_GCM */
+#endif /* PSA_WANT_ALG_GCM */
 
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305)
+#if defined(PSA_WANT_ALG_CHACHA20_POLY1305)
         case PSA_ALG_AEAD_WITH_SHORTENED_TAG( PSA_ALG_CHACHA20_POLY1305, 0 ):
             /* We only support the default tag length. */
             if( tag_len != 16 )
                 return( PSA_ERROR_INVALID_ARGUMENT );
             break;
-#endif /* MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305 */
+#endif /* PSA_WANT_ALG_CHACHA20_POLY1305 */
 
         default:
             (void) tag_len;
@@ -3803,7 +3802,7 @@
     if( status != PSA_SUCCESS )
         goto exit;
 
-    if( ( status = psa_validate_tag_length( operation, alg ) ) != PSA_SUCCESS )
+    if( ( status = psa_validate_tag_length( alg ) ) != PSA_SUCCESS )
         goto exit;
 
     operation->key_type = psa_get_key_type( &attributes );