Use the correct guards on the context structures for MAC/HKDF/PRF

Signed-off-by: Steven Cooreman <steven.cooreman@silabs.com>
diff --git a/include/psa/crypto_builtin_composites.h b/include/psa/crypto_builtin_composites.h
index 81a8ec7..6979dec 100644
--- a/include/psa/crypto_builtin_composites.h
+++ b/include/psa/crypto_builtin_composites.h
@@ -43,7 +43,7 @@
 #define MBEDTLS_PSA_BUILTIN_MAC
 #endif
 
-#if defined(PSA_WANT_ALG_HMAC)
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_HMAC) || defined(PSA_CRYPTO_DRIVER_TEST)
 typedef struct
 {
     /** The HMAC algorithm in use */
@@ -55,7 +55,7 @@
 } mbedtls_psa_hmac_operation_t;
 
 #define MBEDTLS_PSA_HMAC_OPERATION_INIT {0, PSA_HASH_OPERATION_INIT, {0}}
-#endif /* PSA_WANT_ALG_HMAC */
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_HMAC */
 
 #include "mbedtls/cmac.h"
 
@@ -68,12 +68,12 @@
     union
     {
         unsigned dummy; /* Make the union non-empty even with no supported algorithms. */
-#if defined(PSA_WANT_ALG_HMAC)
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_HMAC) || defined(PSA_CRYPTO_DRIVER_TEST)
         mbedtls_psa_hmac_operation_t hmac;
-#endif
-#if defined(MBEDTLS_CMAC_C)
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_HMAC */
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_CMAC) || defined(PSA_CRYPTO_DRIVER_TEST)
         mbedtls_cipher_context_t cmac;
-#endif
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_CMAC */
     } ctx;
 } mbedtls_psa_mac_operation_t;
 
diff --git a/include/psa/crypto_struct.h b/include/psa/crypto_struct.h
index 79c4285..22b5fc2 100644
--- a/include/psa/crypto_struct.h
+++ b/include/psa/crypto_struct.h
@@ -168,7 +168,7 @@
     return( v );
 }
 
-#if defined(MBEDTLS_MD_C)
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF)
 typedef struct
 {
     uint8_t *info;
@@ -184,9 +184,10 @@
     unsigned int state : 2;
     unsigned int info_set : 1;
 } psa_hkdf_key_derivation_t;
-#endif /* MBEDTLS_MD_C */
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_HKDF */
 
-#if defined(MBEDTLS_MD_C)
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PRF) || \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS)
 typedef enum
 {
     PSA_TLS12_PRF_STATE_INIT,       /* no input provided */
@@ -221,7 +222,8 @@
     /* `HMAC_hash( prk, A(i) + seed )` in the notation of RFC 5246, Sect. 5. */
     uint8_t output_block[PSA_HASH_MAX_SIZE];
 } psa_tls12_prf_key_derivation_t;
-#endif /* MBEDTLS_MD_C */
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_TLS12_PRF) ||
+        * MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS */
 
 struct psa_key_derivation_s
 {
@@ -232,8 +234,11 @@
     {
         /* Make the union non-empty even with no supported algorithms. */
         uint8_t dummy;
-#if defined(MBEDTLS_MD_C)
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF)
         psa_hkdf_key_derivation_t hkdf;
+#endif
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PRF) || \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS)
         psa_tls12_prf_key_derivation_t tls12_prf;
 #endif
     } ctx;