diff --git a/programs/psa/hmac_psa.c b/programs/psa/hmac_psa.c
index 019b7ff..7bf69ab 100644
--- a/programs/psa/hmac_psa.c
+++ b/programs/psa/hmac_psa.c
@@ -1,9 +1,24 @@
-/*
- * This is a simple example of multi-part HMAC computation using the PSA
- * Crypto API. It comes with a companion program hmac_non_psa.c, which does
- * the same operations with the legacy MD API. The goal is that comparing the
- * two programs will help people migrating to the PSA Crypto API.
+/**
+ * PSA API multi-part HMAC demonstration.
  *
+ * This programs computes the HMAC of two messages using the multi-part API.
+ *
+ * It comes with a companion program hmac_non_psa.c, which does the same
+ * operations with the legacy MD API. The goal is that comparing the two
+ * programs will help people migrating to the PSA Crypto API.
+ *
+ * When it comes to multi-part HMAC operations, the `mbedtls_md_context`
+ * serves a dual purpose (1) hold the key, and (2) save progress information
+ * for the current operation. With PSA those roles are held by two disinct
+ * objects: (1) a psa_key_id_t to hold the key, and (2) a psa_operation_t for
+ * multi-part progress.
+ *
+ * This program and its companion hmac_non_psa.c illustrate this by doing the
+ * same sequence of multi-part HMAC computation with both APIs; looking at the
+ * two side by side should make the differences and similarities clear.
+ */
+
+/*
  *  Copyright The Mbed TLS Contributors
  *  SPDX-License-Identifier: Apache-2.0
  *
@@ -20,22 +35,17 @@
  *  limitations under the License.
  */
 
-/*
- * When in comes to multi-part HMAC operations, the `mbedtls_md_context`
- * serves a dual purpose (1) hold the key, and (2) save progress information
- * for the current operation. With PSA those roles are held by two disinct
- * objects: (1) a psa_key_id_t to hold the key, and (2) a psa_operation_t for
- * multi-part progress.
- *
- * This program and its companion hmac_non_psa.c illustrate this by doing the
- * same sequence of multi-part HMAC computation with both APIs; looking at the
- * two side by side should make the differences and similarities clear.
- */
-
-#include <stdio.h>
-
+/* First include Mbed TLS headers to get the Mbed TLS configuration and
+ * platform definitions that we'll use in this program. Also include
+ * standard C headers for functions we'll use here. */
 #include "mbedtls/build_info.h"
 
+#include "psa/crypto.h"
+
+#include <stdlib.h>
+#include <stdio.h>
+
+/* If the build options we need are not enabled, compile a placeholder. */
 #if !defined(MBEDTLS_PSA_CRYPTO_C) || \
     defined(MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER)
 int main( void )
@@ -46,20 +56,22 @@
 }
 #else
 
-#include "psa/crypto.h"
+/* The real program starts here. */
 
-/*
- * Dummy inputs for HMAC
- */
+/* Dummy inputs for HMAC */
 const unsigned char msg1_part1[] = { 0x01, 0x02 };
 const unsigned char msg1_part2[] = { 0x03, 0x04 };
 const unsigned char msg2_part1[] = { 0x05, 0x05 };
 const unsigned char msg2_part2[] = { 0x06, 0x06 };
 
+/* Dummy key material - never do this in production!
+ * This example program uses SHA-256, so a 32-byte key makes sense. */
 const unsigned char key_bytes[32] = { 0 };
 
+/* Buffer for the output - using SHA-256, so 32-byte output */
 unsigned char out[32];
 
+/* Print the contents of the output buffer in hex */
 void print_out( const char *title )
 {
     printf( "%s:", title );
@@ -68,13 +80,26 @@
     printf( "\n" );
 }
 
-#define CHK( code )     \
-    do {                \
-        status = code;     \
-        if( status != PSA_SUCCESS )  \
-            goto exit;  \
-    } while( 0 )
+/* Run a PSA function and bail out if it fails. */
+#define PSA_CHECK( expr )                                       \
+    do                                                          \
+    {                                                           \
+        status = ( expr );                                      \
+        if( status != PSA_SUCCESS )                             \
+        {                                                       \
+            printf( "Error %d at line %d: %s\n",                \
+                    (int) status,                               \
+                    __LINE__,                                   \
+                    #expr );                                    \
+            goto exit;                                          \
+        }                                                       \
+    }                                                           \
+    while( 0 )
 
+/*
+ * This function demonstrates computation of the HMAC of two messages using
+ * the multipart API.
+ */
 psa_status_t hmac_demo(void)
 {
     psa_status_t status;
@@ -86,7 +111,7 @@
     psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_SIGN_MESSAGE );
     psa_set_key_algorithm( &attributes, alg );
     psa_set_key_type( &attributes, PSA_KEY_TYPE_HMAC );
-    psa_set_key_bits( &attributes, 8 * sizeof( key_bytes ) );
+    psa_set_key_bits( &attributes, 8 * sizeof( key_bytes ) ); // optional
 
     status = psa_import_key( &attributes, key_bytes, sizeof( key_bytes ), &key );
     if( status != PSA_SUCCESS )
@@ -97,21 +122,21 @@
     size_t out_len = 0;
 
     /* compute HMAC(key, msg1_part1 | msg1_part2) */
-    CHK( psa_mac_sign_setup( &op, key, alg ) );
-    CHK( psa_mac_update( &op, msg1_part1, sizeof( msg1_part1 ) ) );
-    CHK( psa_mac_update( &op, msg1_part2, sizeof( msg1_part2 ) ) );
-    CHK( psa_mac_sign_finish( &op, out, sizeof( out ), &out_len ) );
+    PSA_CHECK( psa_mac_sign_setup( &op, key, alg ) );
+    PSA_CHECK( psa_mac_update( &op, msg1_part1, sizeof( msg1_part1 ) ) );
+    PSA_CHECK( psa_mac_update( &op, msg1_part2, sizeof( msg1_part2 ) ) );
+    PSA_CHECK( psa_mac_sign_finish( &op, out, sizeof( out ), &out_len ) );
     print_out( "msg1" );
 
     /* compute HMAC(key, msg2_part1 | msg2_part2) */
-    CHK( psa_mac_sign_setup( &op, key, alg ) );
-    CHK( psa_mac_update( &op, msg2_part1, sizeof( msg2_part1 ) ) );
-    CHK( psa_mac_update( &op, msg2_part2, sizeof( msg2_part2 ) ) );
-    CHK( psa_mac_sign_finish( &op, out, sizeof( out ), &out_len ) );
+    PSA_CHECK( psa_mac_sign_setup( &op, key, alg ) );
+    PSA_CHECK( psa_mac_update( &op, msg2_part1, sizeof( msg2_part1 ) ) );
+    PSA_CHECK( psa_mac_update( &op, msg2_part2, sizeof( msg2_part2 ) ) );
+    PSA_CHECK( psa_mac_sign_finish( &op, out, sizeof( out ), &out_len ) );
     print_out( "msg2" );
 
 exit:
-    psa_mac_abort( &op );
+    psa_mac_abort( &op ); // needed on error, harmless on success
     psa_destroy_key( key );
 
     return( status );
@@ -119,13 +144,19 @@
 
 int main(void)
 {
-    psa_status_t status = psa_crypto_init();
-    if( status != PSA_SUCCESS )
-        printf( "psa init: %d\n", status );
+    psa_status_t status = PSA_SUCCESS;
 
-    status = hmac_demo();
-    if( status != PSA_SUCCESS )
-        printf( "hmac_demo: %d\n", status );
+    /* Initialize the PSA crypto library. */
+    PSA_CHECK( psa_crypto_init( ) );
+
+    /* Run the demo */
+    PSA_CHECK( hmac_demo() );
+
+    /* Deinitialize the PSA crypto library. */
+    mbedtls_psa_crypto_free( );
+
+exit:
+    return( status == PSA_SUCCESS ? EXIT_SUCCESS : EXIT_FAILURE );
 }
 
 #endif