Update LMS and LMOTS api

Fix function names and parameters. Move macros to be more private.
Update implementation.

Signed-off-by: Raef Coles <raef.coles@arm.com>
diff --git a/include/mbedtls/lms.h b/include/mbedtls/lms.h
index 8430309..c463b2a 100644
--- a/include/mbedtls/lms.h
+++ b/include/mbedtls/lms.h
@@ -32,33 +32,24 @@
 
 #include "lmots.h"
 
-#include "mbedtls/private_access.h"
+#include "mbedtls/build_info.h"
 
 #define MBEDTLS_ERR_LMS_BAD_INPUT_DATA   -0x0011 /**< Bad data has been input to an LMS function */
-#define MBEDTLS_ERR_LMS_OUT_OF_PRIV_KEYS -0x0013 /**< Specified LMS key has utilised all of its private keys */
+#define MBEDTLS_ERR_LMS_OUT_OF_PRIVATE_KEYS -0x0013 /**< Specified LMS key has utilised all of its private keys */
 #define MBEDTLS_ERR_LMS_VERIFY_FAILED    -0x0015 /**< LMS signature verification failed */
 #define MBEDTLS_ERR_LMS_ALLOC_FAILED     -0x0017 /**< LMS failed to allocate space for a private key */
 #define MBEDTLS_ERR_LMS_BUFFER_TOO_SMALL -0x0019 /**< Input/output buffer is too small to contain requited data */
 
-#define MBEDTLS_LMS_TYPE_LEN            (4)
-#define MBEDTLS_LMS_H_TREE_HEIGHT       (10)
 #define MBEDTLS_LMS_M_NODE_BYTES        (32) /* The length of a hash output, 32 for SHA256 */
+#define MBEDTLS_LMS_TYPE_LEN            (4)
+#define MBEDTLS_LMS_H_TREE_HEIGHT       (10u)
 
 #define MBEDTLS_LMS_SIG_LEN (MBEDTLS_LMOTS_Q_LEAF_ID_LEN + MBEDTLS_LMOTS_SIG_LEN + \
                              MBEDTLS_LMS_TYPE_LEN + MBEDTLS_LMS_H_TREE_HEIGHT * MBEDTLS_LMS_M_NODE_BYTES)
 
-#define MBEDTLS_LMS_PUBKEY_LEN (MBEDTLS_LMS_TYPE_LEN + MBEDTLS_LMOTS_TYPE_LEN + \
+#define MBEDTLS_LMS_PUBLIC_KEY_LEN (MBEDTLS_LMS_TYPE_LEN + MBEDTLS_LMOTS_TYPE_LEN + \
                                 MBEDTLS_LMOTS_I_KEY_ID_LEN + MBEDTLS_LMS_M_NODE_BYTES)
 
-#define MBEDTLS_LMS_SIG_Q_LEAF_ID_OFFSET    (0)
-#define MBEDTLS_LMS_SIG_OTS_SIG_OFFSET      (MBEDTLS_LMS_SIG_Q_LEAF_ID_OFFSET + MBEDTLS_LMOTS_Q_LEAF_ID_LEN)
-#define MBEDTLS_LMS_SIG_TYPE_OFFSET         (MBEDTLS_LMS_SIG_OTS_SIG_OFFSET   + MBEDTLS_LMOTS_SIG_LEN)
-#define MBEDTLS_LMS_SIG_PATH_OFFSET         (MBEDTLS_LMS_SIG_TYPE_OFFSET      + MBEDTLS_LMS_TYPE_LEN)
-
-#define MBEDTLS_LMS_PUBKEY_TYPE_OFFSET      (0)
-#define MBEDTLS_LMS_PUBKEY_OTSTYPE_OFFSET   (MBEDTLS_LMS_PUBKEY_TYPE_OFFSET     + MBEDTLS_LMS_TYPE_LEN)
-#define MBEDTLS_LMS_PUBKEY_I_KEY_ID_OFFSET  (MBEDTLS_LMS_PUBKEY_OTSTYPE_OFFSET  + MBEDTLS_LMOTS_TYPE_LEN)
-#define MBEDTLS_LMS_PUBKEY_ROOT_NODE_OFFSET (MBEDTLS_LMS_PUBKEY_I_KEY_ID_OFFSET + MBEDTLS_LMOTS_I_KEY_ID_LEN)
 
 #ifdef __cplusplus
 extern "C" {
@@ -72,85 +63,234 @@
 } mbedtls_lms_algorithm_type_t;
 
 
-/** LMS context structure.
+/** LMS parameters structure.
+ *
+ * This contains the metadata associated with an LMS key, detailing the
+ * algorithm type, the type of the underlying OTS algorithm, and the key ID.
+ */
+typedef struct {
+    unsigned char MBEDTLS_PRIVATE(I_key_identifier[MBEDTLS_LMOTS_I_KEY_ID_LEN]); /*!< The key
+                                                     identifier. */
+    mbedtls_lmots_algorithm_type_t MBEDTLS_PRIVATE(otstype); /*!< The LM-OTS key type identifier as
+                                                               per IANA. Only SHA256_N32_W8 is
+                                                               currently supported. */
+    mbedtls_lms_algorithm_type_t MBEDTLS_PRIVATE(type); /*!< The LMS key type identifier as per
+                                                             IANA. Only SHA256_M32_H10 is currently
+                                                             supported. */
+} mbedtls_lms_parameters_t;
+
+/** LMS public context structure.
+ *
+ *A LMS public key is the hash output that is the root of the merkle tree, and
+ * the applicable parameter set
  *
  * The context must be initialized before it is used. A public key must either
- * be imported, or an algorithm type set, a private key generated and the public
- * key calculated from it. A context that does not contain a public key cannot
- * verify, and a context that does not contain a private key cannot sign.
+ * be imported or generated from a private context.
  *
  * \dot
- * digraph lmots {
+ * digraph lms_public_t {
  *   UNINITIALIZED -> INIT [label="init"];
- *   TYPE_SET -> INIT [label="free"];
- *   PRIVATE -> INIT [label="free"];
- *   PUBLIC -> INIT [label="free"];
- *   "PRIVATE+PUBLIC" -> INIT [label="free"];
- *   INIT -> TYPE_SET [label="set_algorithm_type"];
- *   INIT -> PUBLIC [label="import_public"];
- *   PUBLIC -> PUBLIC [label="export_pubkey"];
- *   "PRIVATE+PUBLIC" -> "PRIVATE+PUBLIC" [label="export_pubkey"];
- *   PRIVATE -> "PRIVATE+PUBLIC" [label="gen_pubkey"];
- *   TYPE_SET -> PRIVATE [label="gen_privkey"];
+ *   HAVE_PUBLIC_KEY -> INIT [label="free"];
+ *   INIT -> HAVE_PUBLIC_KEY [label="import_public_key"];
+ *   INIT -> HAVE_PUBLIC_KEY [label="calculate_public_key from private key"];
+ *   HAVE_PUBLIC_KEY -> HAVE_PUBLIC_KEY [label="export_public_key"];
  * }
  * \enddot
  */
 typedef struct {
-    unsigned char MBEDTLS_PRIVATE(have_privkey); /*!< Whether the context contains a private key.
-                                                     Boolean values only. */
-    unsigned char MBEDTLS_PRIVATE(have_pubkey); /*!< Whether the context contains a public key.
-                                                     Boolean values only. */
-    unsigned char MBEDTLS_PRIVATE(I_key_identifier)[MBEDTLS_LMOTS_I_KEY_ID_LEN]; /*!< The key
-                                                     identifier. */
-    mbedtls_lms_algorithm_type_t MBEDTLS_PRIVATE(type); /*!< The LMS key type identifier as per
-                                                     IANA. Only SHA256_M32_H10 is currently
-                                                     supported. */
-    mbedtls_lmots_algorithm_type_t MBEDTLS_PRIVATE(otstype); /*!< The LM-OTS key type identifier as
-                                                     per IANA. Only SHA256_N32_W8 is currently
-                                                     supported. */
-    unsigned int MBEDTLS_PRIVATE(q_next_usable_key); /*!< The index of the next OTS key that has not
-                                                     been used. */
-    mbedtls_lmots_context *MBEDTLS_PRIVATE(priv_keys); /*!< The private key material. One OTS key
-                                                     for each leaf node in the merkle tree. */
+    mbedtls_lms_parameters_t MBEDTLS_PRIVATE(params);
     unsigned char MBEDTLS_PRIVATE(T_1_pub_key)[MBEDTLS_LMS_M_NODE_BYTES]; /*!< The public key, in
                                                      the form of the merkle tree root node. */
-} mbedtls_lms_context;
+    unsigned char MBEDTLS_PRIVATE(have_public_key); /*!< Whether the context contains a public key.
+                                                     Boolean values only. */
+} mbedtls_lms_public_t;
 
 
+/** LMS private context structure.
+ *
+ * A LMS private key is a set of LMOTS private keys, an index to the next usable
+ * key, and the applicable parameter set.
+ *
+ * The context must be initialized before it is used. A public key must either
+ * be imported or generated from a private context.
+ *
+ * \dot
+ * digraph lms_public_t {
+ *   UNINITIALIZED -> INIT [label="init"];
+ *   HAVE_PRIVATE_KEY -> INIT [label="free"];
+ *   INIT -> HAVE_PRIVATE_KEY [label="generate_private_key"];
+ * }
+ * \enddot
+ */
+typedef struct {
+    mbedtls_lms_parameters_t MBEDTLS_PRIVATE(params);
+    uint32_t MBEDTLS_PRIVATE(q_next_usable_key); /*!< The index of the next OTS key that has not
+                                                      been used. */
+    mbedtls_lmots_private_t *MBEDTLS_PRIVATE(ots_private_keys); /*!< The private key material. One OTS key
+                                                              for each leaf node in the merkle tree. */
+    mbedtls_lmots_public_t *MBEDTLS_PRIVATE(ots_public_keys); /*!< The OTS key public keys, used to
+                                                                   build the merkle tree. */
+    unsigned char MBEDTLS_PRIVATE(have_private_key); /*!< Whether the context contains a private key.
+                                                     Boolean values only. */
+} mbedtls_lms_private_t;
+
 /**
- * \brief                    This function initializes an LMS context
+ * \brief                    This function initializes an LMS public context
  *
  * \param ctx                The uninitialized LMS context that will then be
  *                           initialized.
  */
-void mbedtls_lms_init( mbedtls_lms_context *ctx );
+void mbedtls_lms_init_public( mbedtls_lms_public_t *ctx );
 
 /**
- * \brief                    This function uninitializes an LMS context
+ * \brief                    This function uninitializes an LMS public context
  *
  * \param ctx                The initialized LMS context that will then be
  *                           uninitialized.
  */
-void mbedtls_lms_free( mbedtls_lms_context *ctx );
+void mbedtls_lms_free_public( mbedtls_lms_public_t *ctx );
 
 /**
- * \brief                    This function sets the type of an LMS context
+ * \brief                    This function imports an LMS public key into a
+ *                           public LMS context.
  *
- * \note                     The parameter set in the context will then be used
- *                           for keygen operations etc.
+ * \note                     Before this function is called, the context must
+ *                           have been initialized.
  *
- * \param ctx                The initialized LMS context.
- * \param type               The type that will be set in the context.
- * \param otstype            The type of the LMOTS implementation used by this
- *                           context.
+ * \note                     See IETF RFC8554 for details of the encoding of
+ *                           this public key.
+ *
+ * \param ctx                The initialized LMS context store the key in.
+ * \param key                The buffer from which the key will be read.
+ *                           #MBEDTLS_LMS_PUBLIC_KEY_LEN bytes will be read from
+ *                           this.
+ * \param key_size           The size of the key being imported.
+ *
+ * \return         \c 0 on success.
+ * \return         A non-zero error code on failure.
  */
-int mbedtls_lms_set_algorithm_type( mbedtls_lms_context *ctx,
-                                    mbedtls_lms_algorithm_type_t type,
-                                    mbedtls_lmots_algorithm_type_t otstype);
+int mbedtls_lms_import_public_key( mbedtls_lms_public_t *ctx,
+                                   const unsigned char *key, size_t key_size );
+
+/**
+ * \brief                    This function verifies a LMS signature, using a
+ *                           LMS context that contains a public key.
+ *
+ * \note                     Before this function is called, the context must
+ *                           have been initialized and must contain a public key
+ *                           (either by import or generation).
+ *
+ * \param ctx                The initialized LMS public context from which the
+ *                           public key will be read.
+ * \param msg                The buffer from which the message will be read.
+ * \param msg_size           The size of the message that will be read.
+ * \param sig                The buf from which the signature will be read.
+ *                           #MBEDTLS_LMS_SIG_LEN bytes will be read from
+ *                           this.
+ * \param sig_size           The size of the signature to be verified.
+ *
+ * \return         \c 0 on successful verification.
+ * \return         A non-zero error code on failure.
+ */
+int mbedtls_lms_verify( const mbedtls_lms_public_t *ctx,
+                        const unsigned char *msg, size_t msg_size,
+                        const unsigned char *sig, size_t sig_size );
+
+/**
+ * \brief                    This function initializes an LMS private context
+ *
+ * \param ctx                The uninitialized LMS private context that will
+ *                           then be initialized. */
+void mbedtls_lms_init_private( mbedtls_lms_private_t *ctx );
+
+/**
+ * \brief                    This function uninitializes an LMS private context
+ *
+ * \param ctx                The initialized LMS private context that will then
+ *                           be uninitialized.
+ */
+void mbedtls_lms_free_private( mbedtls_lms_private_t *ctx );
+
+/**
+ * \brief                    This function generates an LMS private key, and
+ *                           stores in into an LMS private context.
+ *
+ * \warning                  This function is **not intended for use in
+ *                           production**, due to as-yet unsolved problems with
+ *                           handling stateful keys.
+ *
+ * \note                     The seed must have at least 256 bits of entropy.
+ *
+ * \param ctx                The initialized LMOTS context to generate the key
+ *                           into.
+ * \param type               The LMS parameter set identifier.
+ * \param otstype            The LMOTS parameter set identifier.
+ * \param f_rng              The RNG function to be used to generate the key ID.
+ * \param p_rng              The RNG context to be passed to f_rng
+ * \param seed               The seed used to deterministically generate the
+ *                           key.
+ * \param seed_size          The length of the seed.
+ *
+ * \return         \c 0 on success.
+ * \return         A non-zero error code on failure.
+ */
+int mbedtls_lms_generate_private_key( mbedtls_lms_private_t *ctx,
+                                      mbedtls_lms_algorithm_type_t type,
+                                      mbedtls_lmots_algorithm_type_t otstype,
+                                      int (*f_rng)(void *, unsigned char *, size_t),
+                                      void* p_rng, unsigned char *seed,
+                                      size_t seed_size );
+
+/**
+ * \brief                    This function generates an LMS public key from a
+ *                           LMS context that already contains a private key.
+ *
+ * \note                     Before this function is called, the context must
+ *                           have been initialized and the context must contain
+ *                           a private key.
+ *
+ * \param ctx                The initialized LMS public context to generate the key
+ *                           from and store it into.
+ *
+ * \param ctx                The LMS private context to read the private key
+ *                           from. This must have been initialized and contain a
+ *                           private key.
+ *
+ * \return         \c 0 on success.
+ * \return         A non-zero error code on failure.
+ */
+int mbedtls_lms_calculate_public_key( mbedtls_lms_public_t *ctx,
+                                      mbedtls_lms_private_t *priv_ctx );
+
+/**
+ * \brief                    This function exports an LMS public key from a
+ *                           LMS public context that already contains a public
+ *                           key.
+ *
+ * \note                     Before this function is called, the context must
+ *                           have been initialized and the context must contain
+ *                           a public key.
+ *
+ * \note                     See IETF RFC8554 for details of the encoding of
+ *                           this public key.
+ *
+ * \param ctx                The initialized LMS public context that contains
+ *                           the public key.
+ * \param key                The buffer into which the key will be output. Must
+ *                           be at least #MBEDTLS_LMS_PUBLIC_KEY_LEN in size.
+ * \param key_size           The size of the key buffer.
+ * \param key_len            If not NULL, will be written with the size of the
+ *                           key.
+ *
+ * \return         \c 0 on success.
+ * \return         A non-zero error code on failure.
+ */
+int mbedtls_lms_export_public_key( mbedtls_lms_public_t *ctx, unsigned char *key,
+                                   size_t key_size, size_t *key_len );
 
 /**
  * \brief                    This function creates a LMS signature, using a
- *                           LMOTS context that contains a private key.
+ *                           LMS context that contains unused private keys.
  *
  * \warning                  This function is **not intended for use in
  *                           production**, due to as-yet unsolved problems with
@@ -167,135 +307,27 @@
  *                           important to not perform copy operations on LMS
  *                           contexts that contain private key material.
  *
- * \param ctx                The initialized LMS context from which the
+ * \param ctx                The initialized LMS private context from which the
  *                           private key will be read.
  * \param f_rng              The RNG function to be used for signature
  *                           generation.
  * \param p_rng              The RNG context to be passed to f_rng
  * \param msg                The buffer from which the message will be read.
- * \param msg_len            The size of the message that will be read.
+ * \param msg_size           The size of the message that will be read.
  * \param sig                The buf into which the signature will be stored.
- *                           Must be at least #MBEDTLS_LMOTS_SIG_LEN in size.
+ *                           Must be at least #MBEDTLS_LMS_SIG_LEN in size.
+ * \param sig_size           The size of the buffer the signature will be
+ *                           written into.
+ * \param sig_len            If not NULL, will be written with the size of the
+ *                           signature.
  *
  * \return         \c 0 on success.
  * \return         A non-zero error code on failure.
  */
-int mbedtls_lms_sign( mbedtls_lms_context *ctx,
+int mbedtls_lms_sign( mbedtls_lms_private_t *ctx,
                       int (*f_rng)(void *, unsigned char *, size_t),
-                      void* p_rng, unsigned char *msg, unsigned int msg_len,
-                      unsigned char *sig);
-
-/**
- * \brief                    This function verifies a LMS signature, using a
- *                           LMS context that contains a public key.
- *
- * \note                     Before this function is called, the context must
- *                           have been initialized and must contain a public key
- *                           (either by import or generation).
- *
- * \param ctx                The initialized LMS context from which the public
- *                           key will be read.
- * \param msg                The buffer from which the message will be read.
- * \param msg_len            The size of the message that will be read.
- * \param sig                The buf from which the signature will be read.
- *                           #MBEDTLS_LMS_SIG_LEN bytes will be read from
- *                           this.
- *
- * \return         \c 0 on successful verification.
- * \return         A non-zero error code on failure.
- */
-int mbedtls_lms_verify( const mbedtls_lms_context *ctx,
-                        const unsigned char *msg, unsigned int msg_len,
-                        const unsigned char *sig );
-
-/**
- * \brief                    This function imports an LMOTS public key into a
- *                           LMS context.
- *
- * \note                     Before this function is called, the context must
- *                           have been initialized.
- *
- * \note                     See IETF RFC8554 for details of the encoding of
- *                           this public key.
- *
- * \param ctx                The initialized LMS context store the key in.
- * \param key                The buffer from which the key will be read.
- *                           #MBEDTLS_LMS_PUBKEY_LEN bytes will be read from
- *                           this.
- *
- * \return         \c 0 on success.
- * \return         A non-zero error code on failure.
- */
-int mbedtls_lms_import_pubkey( mbedtls_lms_context *ctx,
-                               const unsigned char *key );
-
-/**
- * \brief                    This function exports an LMOTS public key from a
- *                           LMS context that already contains a public key.
- *
- * \note                     Before this function is called, the context must
- *                           have been initialized and the context must contain
- *                           a public key.
- *
- * \note                     See IETF RFC8554 for details of the encoding of
- *                           this public key.
- *
- * \param ctx                The initialized LMS context that contains the
- *                           publc key.
- * \param key                The buffer into which the key will be output. Must
- *                           be at least #MBEDTLS_LMS_PUBKEY_LEN in size.
- *
- * \return         \c 0 on success.
- * \return         A non-zero error code on failure.
- */
-int mbedtls_lms_export_pubkey( mbedtls_lms_context *ctx,
-                               unsigned char *key );
-
-/**
- * \brief                    This function generates an LMS public key from a
- *                           LMS context that already contains a private key.
- *
- * \note                     Before this function is called, the context must
- *                           have been initialized and the context must contain
- *                           a private key.
- *
- * \param ctx                The initialized LMS context to generate the key
- *                           from and store it into.
- *
- * \return         \c 0 on success.
- * \return         A non-zero error code on failure.
- */
-int mbedtls_lms_gen_pubkey( mbedtls_lms_context *ctx );
-
-/**
- * \brief                    This function generates an LMS private key, and
- *                           stores in into an LMS context.
- *
- * \warning                  This function is **not intended for use in
- *                           production**, due to as-yet unsolved problems with
- *                           handling stateful keys.
- *
- * \note                     Before this function is called, the context must
- *                           have been initialized and the type of the LMS
- *                           context set using mbedtls_lmots_set_algorithm_type
- *
- * \note                     The seed must have at least 256 bits of entropy.
- *
- * \param ctx                The initialized LMOTS context to generate the key
- *                           into.
- * \param f_rng              The RNG function to be used to generate the key ID.
- * \param p_rng              The RNG context to be passed to f_rng
- * \param seed               The seed used to deterministically generate the
- *                           key.
- * \param seed_len           The length of the seed.
- *
- * \return         \c 0 on success.
- * \return         A non-zero error code on failure.
- */
-int mbedtls_lms_gen_privkey( mbedtls_lms_context *ctx,
-                             int (*f_rng)(void *, unsigned char *, size_t),
-                             void* p_rng, unsigned char *seed,
-                             size_t seed_len );
+                      void* p_rng, unsigned char *msg, unsigned int msg_size,
+                      unsigned char *sig, size_t sig_size, size_t *sig_len);
 
 #ifdef __cplusplus
 }
diff --git a/library/lmots.c b/library/lmots.c
index 915291c..e2f86e6 100644
--- a/library/lmots.c
+++ b/library/lmots.c
@@ -44,18 +44,29 @@
 
 #include "psa/crypto.h"
 
-#define W_SYMBOL_BIT_LEN      (8)
-#define CHECKSUM_LEN          (2)
-#define I_SYMBOL_IDX_LEN      (2)
-#define J_HASH_IDX_LEN        (1)
+#define MBEDTLS_LMOTS_SIG_C_RANDOM_OFFSET  (MBEDTLS_LMOTS_SIG_TYPE_OFFSET     + MBEDTLS_LMOTS_TYPE_LEN)
+#define MBEDTLS_LMOTS_SIG_SIGNATURE_OFFSET (MBEDTLS_LMOTS_SIG_C_RANDOM_OFFSET + MBEDTLS_LMOTS_C_RANDOM_VALUE_LEN)
+
+#define MBEDTLS_LMOTS_PUBLIC_KEY_TYPE_OFFSET      (0)
+#define MBEDTLS_LMOTS_PUBLIC_KEY_I_KEY_ID_OFFSET  (MBEDTLS_LMOTS_PUBLIC_KEY_TYPE_OFFSET      + MBEDTLS_LMOTS_TYPE_LEN)
+#define MBEDTLS_LMOTS_PUBLIC_KEY_Q_LEAF_ID_OFFSET (MBEDTLS_LMOTS_PUBLIC_KEY_I_KEY_ID_OFFSET  + MBEDTLS_LMOTS_I_KEY_ID_LEN)
+#define MBEDTLS_LMOTS_PUBLIC_KEY_KEY_HASH_OFFSET  (MBEDTLS_LMOTS_PUBLIC_KEY_Q_LEAF_ID_OFFSET + MBEDTLS_LMOTS_Q_LEAF_ID_LEN)
+
+/* We only support parameter sets that use 8-bit digits, as it does not require
+ * translation logic between digits and bytes */
+#define W_WINTERNITZ_PARAMETER (8u)
+#define CHECKSUM_LEN           (2)
+#define I_DIGIT_IDX_LEN        (2)
+#define J_HASH_IDX_LEN         (1)
+#define D_CONST_LEN            (2)
+
+#define DIGIT_MAX_VALUE          ((1u << W_WINTERNITZ_PARAMETER) - 1u)
+
 #define D_CONST_LEN           (2)
+static const unsigned char D_PUBLIC_CONSTANT_BYTES[D_CONST_LEN] = {0x80, 0x80};
+static const unsigned char D_MESSAGE_CONSTANT_BYTES[D_CONST_LEN] = {0x81, 0x81};
 
-#define SYMBOL_MAX_VAL        ((1 << W_SYMBOL_BIT_LEN) - 1)
-
-#define D_PBLC_CONSTANT       (0x8080)
-#define D_MESG_CONSTANT       (0x8181)
-
-static void val_to_network_bytes(unsigned int val, size_t len, unsigned char *bytes)
+void unsigned_int_to_network_bytes(unsigned int val, size_t len, unsigned char *bytes)
 {
     size_t idx;
 
@@ -64,7 +75,7 @@
     }
 }
 
-static unsigned int network_bytes_to_val(size_t len, const unsigned char *bytes)
+unsigned int network_bytes_to_unsigned_int(size_t len, const unsigned char *bytes)
 {
     size_t idx;
     unsigned int val = 0;
@@ -76,30 +87,28 @@
     return val;
 }
 
-static unsigned short lmots_checksum_generate( const unsigned char* digest )
+static unsigned short lmots_checksum_calculate( const unsigned char* digest )
 {
     size_t idx;
-    unsigned short sum = 0;
+    unsigned sum = 0;
 
     for ( idx = 0; idx < MBEDTLS_LMOTS_N_HASH_LEN; idx++ )
     {
-        sum += ( 1 << W_SYMBOL_BIT_LEN ) - 1 - digest[idx];
+        sum += DIGIT_MAX_VALUE - digest[idx];
     }
 
     return sum;
 }
 
-static int create_symbol_array( const unsigned char I_key_identifier[MBEDTLS_LMOTS_I_KEY_ID_LEN],
-                                const unsigned char q_leaf_identifier[MBEDTLS_LMOTS_Q_LEAF_ID_LEN],
-                                const unsigned char *msg,
-                                size_t msg_len,
-                                const unsigned char C_random_value[MBEDTLS_LMOTS_C_RANDOM_VALUE_LEN],
-                                unsigned char out[MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN] )
+static int create_digit_array_with_checksum( const mbedtls_lmots_parameters_t *params,
+                                             const unsigned char *msg,
+                                             size_t msg_len,
+                                             const unsigned char C_random_value[MBEDTLS_LMOTS_C_RANDOM_VALUE_LEN],
+                                             unsigned char out[MBEDTLS_LMOTS_P_SIG_DIGIT_COUNT] )
 {
     psa_hash_operation_t op;
     psa_status_t status;
     size_t output_hash_len;
-    unsigned char D_MESG_BYTES[D_CONST_LEN];
     unsigned short checksum;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
@@ -107,166 +116,146 @@
     status = psa_hash_setup( &op, PSA_ALG_SHA_256 );
     ret = mbedtls_lms_error_from_psa( status );
     if ( ret != 0 )
-    {
-        goto out;
-    }
+        goto exit;
 
-    status = psa_hash_update( &op, I_key_identifier, MBEDTLS_LMOTS_I_KEY_ID_LEN );
+    status = psa_hash_update( &op, params->MBEDTLS_PRIVATE(I_key_identifier),
+                              MBEDTLS_LMOTS_I_KEY_ID_LEN );
     ret = mbedtls_lms_error_from_psa( status );
     if ( ret != 0 )
-    {
-        goto out;
-    }
+        goto exit;
 
-    status = psa_hash_update( &op, q_leaf_identifier, MBEDTLS_LMOTS_Q_LEAF_ID_LEN );
+    status = psa_hash_update( &op, params->MBEDTLS_PRIVATE(q_leaf_identifier),
+                              MBEDTLS_LMOTS_Q_LEAF_ID_LEN );
     ret = mbedtls_lms_error_from_psa( status );
     if ( ret != 0 )
-    {
-        goto out;
-    }
+        goto exit;
 
-    val_to_network_bytes( D_MESG_CONSTANT, D_CONST_LEN, D_MESG_BYTES );
-    status = psa_hash_update( &op, D_MESG_BYTES, sizeof( D_MESG_BYTES ) );
+    status = psa_hash_update( &op, D_MESSAGE_CONSTANT_BYTES, D_CONST_LEN );
     ret = mbedtls_lms_error_from_psa( status );
     if ( ret != 0 )
-    {
-        goto out;
-    }
+        goto exit;
 
     status = psa_hash_update( &op, C_random_value, MBEDTLS_LMOTS_C_RANDOM_VALUE_LEN );
     ret = mbedtls_lms_error_from_psa( status );
     if ( ret != 0 )
-    {
-        goto out;
-    }
+        goto exit;
 
     status = psa_hash_update( &op, msg, msg_len );
     ret = mbedtls_lms_error_from_psa( status );
     if ( ret != 0 )
-    {
-        goto out;
-    }
+        goto exit;
 
-    status = psa_hash_finish( &op, out, MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN,
+    status = psa_hash_finish( &op, out, MBEDTLS_LMOTS_P_SIG_DIGIT_COUNT,
                               &output_hash_len );
     ret = mbedtls_lms_error_from_psa( status );
     if ( ret != 0 )
-    {
-        goto out;
-    }
+        goto exit;
 
-    checksum = lmots_checksum_generate( out );
-    val_to_network_bytes( checksum, CHECKSUM_LEN, out + MBEDTLS_LMOTS_N_HASH_LEN );
+    checksum = lmots_checksum_calculate( out );
+    unsigned_int_to_network_bytes( checksum, CHECKSUM_LEN, out + MBEDTLS_LMOTS_N_HASH_LEN );
 
-out:
+exit:
     psa_hash_abort( &op );
 
     return( ret );
 }
 
-static int hash_symbol_array( const unsigned char I_key_identifier[MBEDTLS_LMOTS_I_KEY_ID_LEN],
-                              const unsigned char q_leaf_identifier[MBEDTLS_LMOTS_Q_LEAF_ID_LEN],
-                              const unsigned char x_symbol_array[MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN][32],
-                              const unsigned char hash_idx_min_values[MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN],
-                              const unsigned char hash_idx_max_values[MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN],
-                              unsigned char output[MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN][32] )
+static int hash_digit_array( const mbedtls_lmots_parameters_t *params,
+                             const unsigned char x_digit_array[MBEDTLS_LMOTS_P_SIG_DIGIT_COUNT][MBEDTLS_LMOTS_N_HASH_LEN],
+                             const unsigned char *hash_idx_min_values,
+                             const unsigned char *hash_idx_max_values,
+                             unsigned char output[MBEDTLS_LMOTS_P_SIG_DIGIT_COUNT][MBEDTLS_LMOTS_N_HASH_LEN] )
 {
-    unsigned char i_symbol_idx;
+    unsigned char i_digit_idx;
     unsigned char j_hash_idx;
-    unsigned char i_symbol_idx_bytes[I_SYMBOL_IDX_LEN];
+    unsigned char i_digit_idx_bytes[I_DIGIT_IDX_LEN];
     unsigned char j_hash_idx_bytes[1];
-    unsigned short j_hash_idx_min;
-    unsigned short j_hash_idx_max;
+    /* These can't be unsigned chars, because they are sometimes set to
+     * #DIGIT_MAX_VALUE, which has a value of 256
+     */
+    unsigned int j_hash_idx_min;
+    unsigned int j_hash_idx_max;
     psa_hash_operation_t op;
     psa_status_t status;
     size_t output_hash_len;
-    unsigned char tmp_hash[32];
+    unsigned char tmp_hash[MBEDTLS_LMOTS_N_HASH_LEN];
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
-    for ( i_symbol_idx = 0; i_symbol_idx < MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN; i_symbol_idx++ )
+    op = psa_hash_operation_init();
+
+    for ( i_digit_idx = 0; i_digit_idx < MBEDTLS_LMOTS_P_SIG_DIGIT_COUNT; i_digit_idx++ )
     {
 
-        memcpy( tmp_hash, &x_symbol_array[i_symbol_idx], MBEDTLS_LMOTS_N_HASH_LEN );
+        memcpy( tmp_hash, &x_digit_array[i_digit_idx], MBEDTLS_LMOTS_N_HASH_LEN );
 
-        j_hash_idx_min = hash_idx_min_values != NULL ? hash_idx_min_values[i_symbol_idx] : 0;
-        j_hash_idx_max = hash_idx_max_values != NULL ? hash_idx_max_values[i_symbol_idx] : SYMBOL_MAX_VAL;
+        j_hash_idx_min = hash_idx_min_values != NULL ? hash_idx_min_values[i_digit_idx] : 0;
+        j_hash_idx_max = hash_idx_max_values != NULL ? hash_idx_max_values[i_digit_idx] : DIGIT_MAX_VALUE;
 
         for ( j_hash_idx = (unsigned char)j_hash_idx_min; j_hash_idx < j_hash_idx_max; j_hash_idx++ )
         {
-            op = psa_hash_operation_init();
             status = psa_hash_setup( &op, PSA_ALG_SHA_256 );
             ret = mbedtls_lms_error_from_psa( status );
             if ( ret != 0 )
-            {
-                goto out;
-            }
+                goto exit;
 
-            status = psa_hash_update( &op, I_key_identifier, MBEDTLS_LMOTS_I_KEY_ID_LEN );
+            status = psa_hash_update( &op,
+                                      params->MBEDTLS_PRIVATE(I_key_identifier),
+                                      MBEDTLS_LMOTS_I_KEY_ID_LEN );
             ret = mbedtls_lms_error_from_psa( status );
             if ( ret != 0 )
-            {
-                goto out;
-            }
+                goto exit;
 
-            status = psa_hash_update( &op, q_leaf_identifier, MBEDTLS_LMOTS_Q_LEAF_ID_LEN );
+            status = psa_hash_update( &op,
+                                      params->MBEDTLS_PRIVATE(q_leaf_identifier),
+                                      MBEDTLS_LMOTS_Q_LEAF_ID_LEN );
             ret = mbedtls_lms_error_from_psa( status );
             if ( ret != 0 )
-            {
-                goto out;
-            }
+                goto exit;
 
-            val_to_network_bytes( i_symbol_idx, I_SYMBOL_IDX_LEN, i_symbol_idx_bytes );
-            status = psa_hash_update( &op, i_symbol_idx_bytes, I_SYMBOL_IDX_LEN );
+            unsigned_int_to_network_bytes( i_digit_idx, I_DIGIT_IDX_LEN, i_digit_idx_bytes );
+            status = psa_hash_update( &op, i_digit_idx_bytes, I_DIGIT_IDX_LEN );
             ret = mbedtls_lms_error_from_psa( status );
             if ( ret != 0 )
-            {
-                goto out;
-            }
+                goto exit;
 
-            val_to_network_bytes( j_hash_idx, J_HASH_IDX_LEN, j_hash_idx_bytes );
+            unsigned_int_to_network_bytes( j_hash_idx, J_HASH_IDX_LEN, j_hash_idx_bytes );
             status = psa_hash_update( &op, j_hash_idx_bytes, J_HASH_IDX_LEN );
             ret = mbedtls_lms_error_from_psa( status );
             if ( ret != 0 )
-            {
-                goto out;
-            }
+                goto exit;
 
             status = psa_hash_update( &op, tmp_hash, MBEDTLS_LMOTS_N_HASH_LEN );
             ret = mbedtls_lms_error_from_psa( status );
             if ( ret != 0 )
-            {
-                goto out;
-            }
+                goto exit;
 
             status = psa_hash_finish( &op, tmp_hash, sizeof( tmp_hash ), &output_hash_len );
             ret = mbedtls_lms_error_from_psa( status );
             if ( ret != 0 )
-            {
-                goto out;
-            }
+                goto exit;
 
             psa_hash_abort( &op );
         }
 
-        memcpy( &output[i_symbol_idx], tmp_hash, MBEDTLS_LMOTS_N_HASH_LEN );
+        memcpy( &output[i_digit_idx], tmp_hash, MBEDTLS_LMOTS_N_HASH_LEN );
     }
 
-out:
+exit:
     if( ret )
     {
         psa_hash_abort( &op );
         return( ret );
     }
 
+    mbedtls_platform_zeroize( tmp_hash, sizeof( tmp_hash ) );
+
     return ret;
 }
 
-static int public_key_from_hashed_symbol_array( const unsigned char I_key_identifier[MBEDTLS_LMOTS_I_KEY_ID_LEN],
-                                                const unsigned char q_leaf_identifier[MBEDTLS_LMOTS_Q_LEAF_ID_LEN],
-                                                const unsigned char  y_hashed_symbols[MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN][32],
-                                                unsigned char *pub_key )
+static int public_key_from_hashed_digit_array( const mbedtls_lmots_parameters_t *params,
+                                               const unsigned char y_hashed_digits[MBEDTLS_LMOTS_P_SIG_DIGIT_COUNT][MBEDTLS_LMOTS_N_HASH_LEN],
+                                               unsigned char *pub_key )
 {
-    unsigned char D_PBLC_bytes[D_CONST_LEN];
     psa_hash_operation_t op;
     psa_status_t status;
     size_t output_hash_len;
@@ -276,44 +265,36 @@
     status = psa_hash_setup( &op, PSA_ALG_SHA_256 );
     ret = mbedtls_lms_error_from_psa( status );
     if ( ret != 0 )
-    {
-        goto out;
-    }
+        goto exit;
 
-    status = psa_hash_update( &op, I_key_identifier, MBEDTLS_LMOTS_I_KEY_ID_LEN );
+    status = psa_hash_update( &op,
+                              params->MBEDTLS_PRIVATE(I_key_identifier),
+                              MBEDTLS_LMOTS_I_KEY_ID_LEN );
     ret = mbedtls_lms_error_from_psa( status );
     if ( ret != 0 )
-    {
-        goto out;
-    }
+        goto exit;
 
-    status = psa_hash_update( &op, q_leaf_identifier, MBEDTLS_LMOTS_Q_LEAF_ID_LEN );
+    status = psa_hash_update( &op, params->MBEDTLS_PRIVATE(q_leaf_identifier),
+                              MBEDTLS_LMOTS_Q_LEAF_ID_LEN );
     ret = mbedtls_lms_error_from_psa( status );
     if ( ret != 0 )
-    {
-        goto out;
-    }
+        goto exit;
 
-    val_to_network_bytes( D_PBLC_CONSTANT, D_CONST_LEN, D_PBLC_bytes );
-    status = psa_hash_update( &op, D_PBLC_bytes, D_CONST_LEN );
+    status = psa_hash_update( &op, D_PUBLIC_CONSTANT_BYTES, D_CONST_LEN );
     ret = mbedtls_lms_error_from_psa( status );
     if ( ret != 0 )
-    {
-        goto out;
-    }
+        goto exit;
 
-    status = psa_hash_update( &op, ( unsigned char * )y_hashed_symbols,
-                              MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN * MBEDTLS_LMOTS_N_HASH_LEN );
+    status = psa_hash_update( &op, ( unsigned char * )y_hashed_digits,
+                              MBEDTLS_LMOTS_P_SIG_DIGIT_COUNT * MBEDTLS_LMOTS_N_HASH_LEN );
     ret = mbedtls_lms_error_from_psa( status );
     if ( ret != 0 )
-    {
-        goto out;
-    }
+        goto exit;
 
     status = psa_hash_finish( &op, pub_key, 32, &output_hash_len );
     ret = mbedtls_lms_error_from_psa( status );
 
-out:
+exit:
     psa_hash_abort( &op );
     return( ret );
 }
@@ -336,96 +317,352 @@
     }
 }
 
-void mbedtls_lmots_init( mbedtls_lmots_context *ctx )
+void mbedtls_lmots_init_public( mbedtls_lmots_public_t *ctx )
 {
-    if( ctx == NULL ) {
-        return;
-    }
-
-    mbedtls_platform_zeroize( ctx, sizeof( mbedtls_lmots_context  ) ) ;
+    mbedtls_platform_zeroize( ctx, sizeof( mbedtls_lmots_public_t  ) ) ;
 }
 
-void mbedtls_lmots_free( mbedtls_lmots_context *ctx )
+void mbedtls_lmots_free_public( mbedtls_lmots_public_t *ctx )
 {
-    if( ctx == NULL )
-    {
-        return;
-    }
-
-    mbedtls_platform_zeroize( ctx, sizeof( mbedtls_lmots_context  ) ) ;
+    mbedtls_platform_zeroize( ctx, sizeof( mbedtls_lmots_public_t  ) ) ;
 }
 
-int mbedtls_lmots_set_algorithm_type( mbedtls_lmots_context *ctx,
-                                      mbedtls_lmots_algorithm_type_t type )
+int mbedtls_lmots_import_public_key( mbedtls_lmots_public_t *ctx,
+                                 const unsigned char *key, size_t key_len )
 {
-    if( ctx == NULL )
+    if ( key_len < MBEDTLS_LMOTS_PUBLIC_KEY_LEN )
     {
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
 
-    ctx->MBEDTLS_PRIVATE(type) = type;
+    ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(type) =
+        network_bytes_to_unsigned_int( MBEDTLS_LMOTS_TYPE_LEN,
+                                       key + MBEDTLS_LMOTS_SIG_TYPE_OFFSET );
+
+    memcpy( ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(I_key_identifier),
+            key + MBEDTLS_LMOTS_PUBLIC_KEY_I_KEY_ID_OFFSET, MBEDTLS_LMOTS_I_KEY_ID_LEN );
+
+    memcpy( ctx->MBEDTLS_PRIVATE(MBEDTLS_PRIVATE(params).q_leaf_identifier),
+            key + MBEDTLS_LMOTS_PUBLIC_KEY_Q_LEAF_ID_OFFSET, MBEDTLS_LMOTS_Q_LEAF_ID_LEN );
+
+    memcpy( ctx->MBEDTLS_PRIVATE(public_key),
+            key + MBEDTLS_LMOTS_PUBLIC_KEY_KEY_HASH_OFFSET,
+            MBEDTLS_LMOTS_N_HASH_LEN );
+
+    ctx->MBEDTLS_PRIVATE(have_public_key) = 1;
 
     return( 0 );
 }
 
-int mbedtls_lmots_generate_pub_key_candidate( const unsigned char I_key_identifier[MBEDTLS_LMOTS_I_KEY_ID_LEN],
-                                              const unsigned char q_leaf_identifier[MBEDTLS_LMOTS_Q_LEAF_ID_LEN],
-                                              const unsigned char  *msg,
-                                              size_t msg_len,
-                                              const unsigned char *sig,
-                                              unsigned char *out )
+int mbedtls_lmots_calculate_public_key_candidate( const mbedtls_lmots_parameters_t *params,
+                                                  const unsigned char  *msg,
+                                                  size_t msg_size,
+                                                  const unsigned char *sig,
+                                                  size_t sig_size,
+                                                  unsigned char *out,
+                                                  size_t out_size,
+                                                  size_t *out_len)
 {
-    unsigned char tmp_symbol_array[MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN];
-    unsigned char y_hashed_symbols[MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN][32];
+    unsigned char tmp_digit_array[MBEDTLS_LMOTS_P_SIG_DIGIT_COUNT];
+    unsigned char y_hashed_digits[MBEDTLS_LMOTS_P_SIG_DIGIT_COUNT][MBEDTLS_LMOTS_N_HASH_LEN];
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
-    if (I_key_identifier == NULL || msg == NULL || sig == NULL || out == NULL)
+    if ( msg == NULL && msg_size != 0 )
+    {
+        return ( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
+    }
+
+    if ( sig_size != MBEDTLS_LMOTS_SIG_LEN || out_size < MBEDTLS_LMOTS_N_HASH_LEN )
     {
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
 
-    ret = create_symbol_array( I_key_identifier, q_leaf_identifier, msg, msg_len,
-                               sig + MBEDTLS_LMOTS_SIG_C_RANDOM_OFFSET, tmp_symbol_array );
+    ret = create_digit_array_with_checksum( params, msg, msg_size,
+                                            sig + MBEDTLS_LMOTS_SIG_C_RANDOM_OFFSET,
+                                            tmp_digit_array );
     if ( ret )
     {
         return ( ret );
     }
 
-    ret = hash_symbol_array( I_key_identifier, q_leaf_identifier,
-                             ( const unsigned char( *)[32] )(sig + MBEDTLS_LMOTS_SIG_SIGNATURE_OFFSET),
-                             tmp_symbol_array, NULL, y_hashed_symbols );
+    ret = hash_digit_array( params,
+                             ( const unsigned char( *)[MBEDTLS_LMOTS_N_HASH_LEN] )(sig + MBEDTLS_LMOTS_SIG_SIGNATURE_OFFSET),
+                             tmp_digit_array, NULL, y_hashed_digits );
     if ( ret )
     {
         return ( ret );
     }
 
-    ret = public_key_from_hashed_symbol_array( I_key_identifier, q_leaf_identifier,
-                                               ( const unsigned char( *)[32] )y_hashed_symbols,
+    ret = public_key_from_hashed_digit_array( params,
+                                               ( const unsigned char( *)[MBEDTLS_LMOTS_N_HASH_LEN] )y_hashed_digits,
                                                out );
     if ( ret )
     {
         return ( ret );
     }
 
+    if ( out_len != NULL )
+    {
+        *out_len = MBEDTLS_LMOTS_N_HASH_LEN;
+    }
+
     return( 0 );
 }
 
-int mbedtls_lmots_sign( mbedtls_lmots_context *ctx,
-                        int (*f_rng)(void *, unsigned char *, size_t),
-                        void *p_rng, const unsigned char *msg, size_t msg_len,
-                        unsigned char *sig )
+int mbedtls_lmots_verify( mbedtls_lmots_public_t *ctx, const unsigned char *msg,
+                          size_t msg_size, const unsigned char *sig,
+                          size_t sig_size )
 {
-    unsigned char tmp_symbol_array[MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN];
-    unsigned char tmp_sig[MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN][MBEDTLS_LMOTS_N_HASH_LEN];
+    unsigned char Kc_public_key_candidate[MBEDTLS_LMOTS_N_HASH_LEN];
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
-    if( ctx == NULL || f_rng == NULL || p_rng == NULL || msg == NULL || sig == NULL)
+    if ( msg == NULL && msg_size != 0 )
+    {
+        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
+    }
+
+    if ( !ctx->MBEDTLS_PRIVATE(have_public_key) )
+    {
+        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
+    }
+
+    if( ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE( type )
+        != MBEDTLS_LMOTS_SHA256_N32_W8 )
+    {
+        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
+    }
+
+    if ( network_bytes_to_unsigned_int( MBEDTLS_LMOTS_TYPE_LEN,
+                               sig + MBEDTLS_LMOTS_SIG_TYPE_OFFSET ) != MBEDTLS_LMOTS_SHA256_N32_W8 )
+    {
+        return( MBEDTLS_ERR_LMS_VERIFY_FAILED );
+    }
+
+    ret = mbedtls_lmots_calculate_public_key_candidate( &ctx->MBEDTLS_PRIVATE(params),
+                                                        msg, msg_size, sig, sig_size,
+                                                        Kc_public_key_candidate,
+                                                        MBEDTLS_LMOTS_N_HASH_LEN,
+                                                        NULL);
+    if ( ret )
+    {
+        return( ret );
+    }
+
+    if ( memcmp( &Kc_public_key_candidate, ctx->MBEDTLS_PRIVATE(public_key),
+                 sizeof( ctx->MBEDTLS_PRIVATE(public_key) ) ) )
+    {
+        return( MBEDTLS_ERR_LMS_VERIFY_FAILED );
+    }
+
+    return( 0 );
+}
+
+void mbedtls_lmots_init_private( mbedtls_lmots_private_t *ctx )
+{
+    mbedtls_platform_zeroize( ctx, sizeof( mbedtls_lmots_private_t  ) ) ;
+}
+
+void mbedtls_lmots_free_private( mbedtls_lmots_private_t *ctx )
+{
+    mbedtls_platform_zeroize( ctx, sizeof( mbedtls_lmots_private_t  ) ) ;
+}
+
+int mbedtls_lmots_generate_private_key( mbedtls_lmots_private_t *ctx,
+                                        mbedtls_lmots_algorithm_type_t type,
+                                        const unsigned char I_key_identifier[MBEDTLS_LMOTS_I_KEY_ID_LEN],
+                                        uint32_t q_leaf_identifier,
+                                        const unsigned char *seed,
+                                        size_t seed_size )
+{
+    psa_hash_operation_t op;
+    psa_status_t status;
+    size_t output_hash_len;
+    unsigned int i_digit_idx;
+    unsigned char i_digit_idx_bytes[2];
+    unsigned char const_bytes[1];
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+
+    if ( ctx->MBEDTLS_PRIVATE(have_private_key) )
+    {
+        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
+    }
+
+    if ( type != MBEDTLS_LMOTS_SHA256_N32_W8 ) {
+        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
+    }
+
+    ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(type) = type;
+
+    memcpy( ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(I_key_identifier),
+            I_key_identifier,
+            sizeof( ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(I_key_identifier) ) );
+
+    unsigned_int_to_network_bytes(q_leaf_identifier,
+                                  MBEDTLS_LMOTS_Q_LEAF_ID_LEN,
+                                  ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(q_leaf_identifier) );
+
+    unsigned_int_to_network_bytes( 0xFF, sizeof( const_bytes ), const_bytes );
+
+    for ( i_digit_idx = 0; i_digit_idx < MBEDTLS_LMOTS_P_SIG_DIGIT_COUNT; i_digit_idx++ )
+    {
+        op = psa_hash_operation_init( );
+        status = psa_hash_setup( &op, PSA_ALG_SHA_256 );
+        ret = mbedtls_lms_error_from_psa( status );
+        if ( ret != 0 )
+            goto exit;
+
+        ret = psa_hash_update( &op,
+                               ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(I_key_identifier),
+                               sizeof( ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(I_key_identifier) ) );
+        ret = mbedtls_lms_error_from_psa( status );
+        if ( ret )
+            goto exit;
+
+        status = psa_hash_update( &op,
+                                  ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(q_leaf_identifier),
+                                  MBEDTLS_LMOTS_Q_LEAF_ID_LEN );
+        ret = mbedtls_lms_error_from_psa( status );
+        if ( ret )
+            goto exit;
+
+        unsigned_int_to_network_bytes( i_digit_idx, I_DIGIT_IDX_LEN, i_digit_idx_bytes );
+        status = psa_hash_update( &op, i_digit_idx_bytes, I_DIGIT_IDX_LEN );
+        ret = mbedtls_lms_error_from_psa( status );
+        if ( ret )
+            goto exit;
+
+        status = psa_hash_update( &op, const_bytes, sizeof( const_bytes) );
+        ret = mbedtls_lms_error_from_psa( status );
+        if ( ret )
+            goto exit;
+
+        status = psa_hash_update( &op, seed, seed_size );
+        ret = mbedtls_lms_error_from_psa( status );
+        if ( ret )
+            goto exit;
+
+        status = psa_hash_finish( &op,
+                                  ctx->MBEDTLS_PRIVATE(private_key)[i_digit_idx],
+                                  32, &output_hash_len );
+        ret = mbedtls_lms_error_from_psa( status );
+        if ( ret )
+            goto exit;
+
+        psa_hash_abort( &op );
+    }
+
+    ctx->MBEDTLS_PRIVATE(have_private_key) = 1;
+
+exit:
+    if( ret )
+    {
+        psa_hash_abort( &op );
+        return( ret );
+    }
+
+    return ret;
+}
+
+int mbedtls_lmots_calculate_public_key( mbedtls_lmots_public_t *ctx,
+                                        mbedtls_lmots_private_t *priv_ctx)
+{
+    unsigned char y_hashed_digits[MBEDTLS_LMOTS_P_SIG_DIGIT_COUNT][MBEDTLS_LMOTS_N_HASH_LEN];
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+
+    if( ctx == NULL )
     {
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
 
     /* Check that a private key is loaded */
-    if ( !ctx->MBEDTLS_PRIVATE(have_privkey) )
+    if ( !priv_ctx->MBEDTLS_PRIVATE(have_private_key) )
+    {
+        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
+    }
+
+    ret = hash_digit_array( &priv_ctx->MBEDTLS_PRIVATE(params),
+                            ( const unsigned char( *)[MBEDTLS_LMOTS_N_HASH_LEN] )(priv_ctx->MBEDTLS_PRIVATE(private_key)),
+                            NULL, NULL, y_hashed_digits );
+    if ( ret )
+    {
+        return( ret );
+    }
+
+    ret = public_key_from_hashed_digit_array( &priv_ctx->MBEDTLS_PRIVATE(params),
+                                               ( const unsigned char( *)[MBEDTLS_LMOTS_N_HASH_LEN] )y_hashed_digits,
+                                               ctx->MBEDTLS_PRIVATE(public_key) );
+    if ( ret )
+    {
+        return( ret );
+    }
+
+    memcpy( &ctx->MBEDTLS_PRIVATE(params), &priv_ctx->MBEDTLS_PRIVATE(params),
+            sizeof( ctx->MBEDTLS_PRIVATE(params) ) );
+
+    ctx->MBEDTLS_PRIVATE(have_public_key = 1);
+
+    return( ret );
+}
+
+
+int mbedtls_lmots_export_public_key( mbedtls_lmots_public_t *ctx,
+                                     unsigned char *key, size_t key_size,
+                                     size_t *key_len )
+{
+    if( key_size < MBEDTLS_LMS_PUBLIC_KEY_LEN )
+    {
+        return( MBEDTLS_ERR_LMS_BUFFER_TOO_SMALL );
+    }
+
+    if( ! ctx->MBEDTLS_PRIVATE(have_public_key) )
+    {
+        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
+    }
+
+    unsigned_int_to_network_bytes( ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(type),
+                                   MBEDTLS_LMOTS_TYPE_LEN,
+                                   key + MBEDTLS_LMOTS_SIG_TYPE_OFFSET );
+
+    memcpy( key + MBEDTLS_LMOTS_PUBLIC_KEY_I_KEY_ID_OFFSET,
+            ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(I_key_identifier),
+            MBEDTLS_LMOTS_I_KEY_ID_LEN );
+
+    memcpy(key + MBEDTLS_LMOTS_PUBLIC_KEY_Q_LEAF_ID_OFFSET,
+           ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(q_leaf_identifier),
+           MBEDTLS_LMOTS_Q_LEAF_ID_LEN);
+
+    memcpy( key + MBEDTLS_LMOTS_PUBLIC_KEY_KEY_HASH_OFFSET, ctx->MBEDTLS_PRIVATE(public_key),
+            MBEDTLS_LMOTS_N_HASH_LEN );
+
+    if( key_len != NULL )
+    {
+        *key_len = MBEDTLS_LMS_PUBLIC_KEY_LEN;
+    }
+
+    return( 0 );
+}
+
+int mbedtls_lmots_sign( mbedtls_lmots_private_t *ctx,
+                        int (*f_rng)(void *, unsigned char *, size_t),
+                        void *p_rng, const unsigned char *msg, size_t msg_size,
+                        unsigned char *sig, size_t sig_size, size_t* sig_len )
+{
+    unsigned char tmp_digit_array[MBEDTLS_LMOTS_P_SIG_DIGIT_COUNT];
+    unsigned char tmp_sig[MBEDTLS_LMOTS_P_SIG_DIGIT_COUNT][MBEDTLS_LMOTS_N_HASH_LEN];
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+
+    if ( msg == NULL && msg_size != 0 )
+    {
+        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
+    }
+
+    if( sig_size < MBEDTLS_LMOTS_SIG_LEN )
+    {
+        return( MBEDTLS_ERR_LMS_BUFFER_TOO_SMALL );
+    }
+
+    /* Check that a private key is loaded */
+    if ( !ctx->MBEDTLS_PRIVATE(have_private_key) )
     {
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
@@ -436,285 +673,43 @@
         return( ret );
     }
 
-    ret = create_symbol_array( ctx->MBEDTLS_PRIVATE(I_key_identifier),
-                               ctx->MBEDTLS_PRIVATE(q_leaf_identifier_bytes),
-                               msg, msg_len, sig + MBEDTLS_LMOTS_SIG_C_RANDOM_OFFSET,
-                               tmp_symbol_array );
+    ret = create_digit_array_with_checksum( &ctx->MBEDTLS_PRIVATE(params),
+                                            msg, msg_size,
+                                            sig + MBEDTLS_LMOTS_SIG_C_RANDOM_OFFSET,
+                                            tmp_digit_array );
     if ( ret )
     {
         return( ret );
     }
 
-    ret = hash_symbol_array( ctx->MBEDTLS_PRIVATE(I_key_identifier),
-                             ctx->MBEDTLS_PRIVATE(q_leaf_identifier_bytes),
-                             ( const unsigned char( *)[32] )(ctx->MBEDTLS_PRIVATE(priv_key)),
-                             NULL, tmp_symbol_array, tmp_sig );
+    ret = hash_digit_array( &ctx->MBEDTLS_PRIVATE(params),
+                             ( const unsigned char( *)[MBEDTLS_LMOTS_N_HASH_LEN] )(ctx->MBEDTLS_PRIVATE(private_key)),
+                             NULL, tmp_digit_array, tmp_sig );
     if ( ret )
     {
         return( ret );
     }
 
-    val_to_network_bytes( ctx->MBEDTLS_PRIVATE(type), MBEDTLS_LMOTS_TYPE_LEN,
-                          sig + MBEDTLS_LMOTS_SIG_TYPE_OFFSET );
+    unsigned_int_to_network_bytes( ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(type),
+                                   MBEDTLS_LMOTS_TYPE_LEN,
+                                   sig + MBEDTLS_LMOTS_SIG_TYPE_OFFSET );
 
     /* We've got a valid signature now, so it's time to make sure the private
      * key can't be reused.
      */
-    ctx->MBEDTLS_PRIVATE(have_privkey) = 0;
-    mbedtls_platform_zeroize(ctx->MBEDTLS_PRIVATE(priv_key),
-                             sizeof(ctx->MBEDTLS_PRIVATE(priv_key)));
+    ctx->MBEDTLS_PRIVATE(have_private_key) = 0;
+    mbedtls_platform_zeroize(ctx->MBEDTLS_PRIVATE(private_key),
+                             sizeof(ctx->MBEDTLS_PRIVATE(private_key)));
 
     memcpy(sig + MBEDTLS_LMOTS_SIG_SIGNATURE_OFFSET, tmp_sig,
-           MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN * MBEDTLS_LMOTS_N_HASH_LEN);
+           MBEDTLS_LMOTS_P_SIG_DIGIT_COUNT * MBEDTLS_LMOTS_N_HASH_LEN);
 
-    return( 0 );
-}
-
-int mbedtls_lmots_verify( mbedtls_lmots_context *ctx, const unsigned char *msg,
-                          size_t msg_len, const unsigned char *sig )
-{
-    unsigned char Kc_public_key_candidate[32];
-    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-
-    if( ctx == NULL || msg == NULL || sig == NULL)
+    if( sig_len != NULL )
     {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    if ( !ctx->MBEDTLS_PRIVATE(have_pubkey) )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    if( ctx->MBEDTLS_PRIVATE(type ) != MBEDTLS_LMOTS_SHA256_N32_W8 )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    if ( network_bytes_to_val( MBEDTLS_LMOTS_TYPE_LEN,
-                               sig + MBEDTLS_LMOTS_SIG_TYPE_OFFSET ) != MBEDTLS_LMOTS_SHA256_N32_W8 )
-    {
-        return( MBEDTLS_ERR_LMS_VERIFY_FAILED );
-    }
-
-    ret = mbedtls_lmots_generate_pub_key_candidate( ctx->MBEDTLS_PRIVATE(I_key_identifier),
-                                                    ctx->MBEDTLS_PRIVATE(q_leaf_identifier_bytes),
-                                                    msg, msg_len, sig,
-                                                    Kc_public_key_candidate );
-    if ( ret )
-    {
-        return( ret );
-    }
-
-    if ( memcmp( &Kc_public_key_candidate, ctx->MBEDTLS_PRIVATE(pub_key),
-                 sizeof( ctx->MBEDTLS_PRIVATE(pub_key) ) ) )
-    {
-        return( MBEDTLS_ERR_LMS_VERIFY_FAILED );
+        *sig_len = MBEDTLS_LMS_SIG_LEN;
     }
 
     return( 0 );
 }
 
-int mbedtls_lmots_import_pubkey( mbedtls_lmots_context *ctx,
-                                 const unsigned char *key )
-{
-    if ( ctx == NULL || key == NULL)
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    ctx->MBEDTLS_PRIVATE(type) = network_bytes_to_val( MBEDTLS_LMOTS_TYPE_LEN,
-                                                       key + MBEDTLS_LMOTS_SIG_TYPE_OFFSET );
-
-    memcpy( ctx->MBEDTLS_PRIVATE(I_key_identifier), key + MBEDTLS_LMOTS_PUBKEY_I_KEY_ID_OFFSET,
-            MBEDTLS_LMOTS_I_KEY_ID_LEN );
-
-    memcpy( ctx->MBEDTLS_PRIVATE(q_leaf_identifier_bytes), key + MBEDTLS_LMOTS_PUBKEY_Q_LEAF_ID_OFFSET,
-            MBEDTLS_LMOTS_Q_LEAF_ID_LEN );
-    ctx->MBEDTLS_PRIVATE(q_leaf_identifier) = network_bytes_to_val( MBEDTLS_LMOTS_Q_LEAF_ID_LEN,
-                                                                    ctx->MBEDTLS_PRIVATE(q_leaf_identifier_bytes) );
-
-    memcpy( ctx->MBEDTLS_PRIVATE(pub_key), key + MBEDTLS_LMOTS_PUBKEY_KEY_HASH_OFFSET, MBEDTLS_LMOTS_N_HASH_LEN );
-
-    ctx->MBEDTLS_PRIVATE(have_pubkey) = 1;
-
-    return( 0 );
-}
-
-int mbedtls_lmots_export_pubkey( mbedtls_lmots_context *ctx,
-                                 unsigned char *key )
-{
-    if ( ctx == NULL || key == NULL)
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    if ( ! ctx->MBEDTLS_PRIVATE(have_pubkey) )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    val_to_network_bytes( ctx->MBEDTLS_PRIVATE(type), MBEDTLS_LMOTS_TYPE_LEN,
-                          key + MBEDTLS_LMOTS_SIG_TYPE_OFFSET );
-
-    memcpy( key + MBEDTLS_LMOTS_PUBKEY_I_KEY_ID_OFFSET, ctx->MBEDTLS_PRIVATE(I_key_identifier),
-            MBEDTLS_LMOTS_I_KEY_ID_LEN );
-
-    memcpy( key + MBEDTLS_LMOTS_PUBKEY_Q_LEAF_ID_OFFSET, ctx->MBEDTLS_PRIVATE(q_leaf_identifier_bytes),
-            MBEDTLS_LMOTS_Q_LEAF_ID_LEN );
-
-    memcpy( key + MBEDTLS_LMOTS_PUBKEY_KEY_HASH_OFFSET, ctx->MBEDTLS_PRIVATE(pub_key),
-            MBEDTLS_LMOTS_N_HASH_LEN );
-
-    return( 0 );
-}
-
-
-int mbedtls_lmots_gen_pubkey( mbedtls_lmots_context *ctx )
-{
-    unsigned char y_hashed_symbols[MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN][32];
-    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-
-    if( ctx == NULL )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    /* Check that a private key is loaded */
-    if ( !ctx->MBEDTLS_PRIVATE(have_privkey) )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    ret = hash_symbol_array( ctx->MBEDTLS_PRIVATE(I_key_identifier),
-                             ctx->MBEDTLS_PRIVATE(q_leaf_identifier_bytes),
-                             ( const unsigned char( *)[32] )(ctx->MBEDTLS_PRIVATE(priv_key)),
-                             NULL, NULL, y_hashed_symbols );
-    if ( ret )
-    {
-        return( ret );
-    }
-
-    ret = public_key_from_hashed_symbol_array( ctx->MBEDTLS_PRIVATE(I_key_identifier),
-                                               ctx->MBEDTLS_PRIVATE(q_leaf_identifier_bytes),
-                                               ( const unsigned char( *)[32] )y_hashed_symbols,
-                                               ctx->MBEDTLS_PRIVATE(pub_key) );
-    if ( ret )
-    {
-        return( ret );
-    }
-
-    ctx->MBEDTLS_PRIVATE(have_pubkey = 1);
-
-    return( ret );
-}
-
-int mbedtls_lmots_gen_privkey( mbedtls_lmots_context *ctx,
-                               const unsigned char I_key_identifier[MBEDTLS_LMOTS_I_KEY_ID_LEN],
-                               unsigned int q_leaf_identifier,
-                               const unsigned char *seed,
-                               size_t seed_len )
-{
-    psa_hash_operation_t op;
-    psa_status_t status;
-    size_t output_hash_len;
-    unsigned int i_symbol_idx;
-    unsigned char i_symbol_idx_bytes[2];
-    unsigned char const_bytes[1];
-    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-
-    if( ctx == NULL || I_key_identifier == NULL || seed == NULL)
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    if ( ctx->MBEDTLS_PRIVATE(have_privkey) )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    if ( ctx->MBEDTLS_PRIVATE(type) != MBEDTLS_LMOTS_SHA256_N32_W8 ) {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    memcpy( ctx->MBEDTLS_PRIVATE(I_key_identifier), I_key_identifier,
-            sizeof( ctx->MBEDTLS_PRIVATE(I_key_identifier) ) );
-
-    ctx->MBEDTLS_PRIVATE(q_leaf_identifier) = q_leaf_identifier;
-
-    val_to_network_bytes( ctx->MBEDTLS_PRIVATE(q_leaf_identifier), MBEDTLS_LMOTS_Q_LEAF_ID_LEN,
-                          ctx->MBEDTLS_PRIVATE(q_leaf_identifier_bytes) );
-
-    val_to_network_bytes( 0xFF, sizeof( const_bytes ), const_bytes );
-
-    for ( i_symbol_idx = 0; i_symbol_idx < MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN; i_symbol_idx++ )
-    {
-        op = psa_hash_operation_init( );
-        status = psa_hash_setup( &op, PSA_ALG_SHA_256 );
-        ret = mbedtls_lms_error_from_psa( status );
-        if ( ret != 0 )
-        {
-            goto out;
-        }
-
-        ret = psa_hash_update( &op, ctx->MBEDTLS_PRIVATE(I_key_identifier),
-                               sizeof( ctx->MBEDTLS_PRIVATE(I_key_identifier) ) );
-        ret = mbedtls_lms_error_from_psa( status );
-        if ( ret ) {
-            goto out;
-        }
-
-        status = psa_hash_update( &op, ctx->MBEDTLS_PRIVATE(q_leaf_identifier_bytes),
-                               sizeof( ctx->MBEDTLS_PRIVATE(q_leaf_identifier_bytes) ) );
-        ret = mbedtls_lms_error_from_psa( status );
-        if ( ret )
-        {
-            goto out;
-        }
-
-        val_to_network_bytes( i_symbol_idx, I_SYMBOL_IDX_LEN, i_symbol_idx_bytes );
-        status = psa_hash_update( &op, i_symbol_idx_bytes, I_SYMBOL_IDX_LEN );
-        ret = mbedtls_lms_error_from_psa( status );
-        if ( ret )
-        {
-            goto out;
-        }
-
-        status = psa_hash_update( &op, const_bytes, sizeof( const_bytes) );
-        ret = mbedtls_lms_error_from_psa( status );
-        if ( ret )
-        {
-            goto out;
-        }
-
-        status = psa_hash_update( &op, seed, seed_len );
-        ret = mbedtls_lms_error_from_psa( status );
-        if ( ret )
-        {
-            goto out;
-        }
-
-        status = psa_hash_finish( &op, ctx->MBEDTLS_PRIVATE(priv_key)[i_symbol_idx],
-                                  32, &output_hash_len );
-        ret = mbedtls_lms_error_from_psa( status );
-        if ( ret )
-        {
-            goto out;
-        }
-
-        psa_hash_abort( &op );
-    }
-
-    ctx->MBEDTLS_PRIVATE(have_privkey) = 1;
-
-out:
-    if( ret )
-    {
-        psa_hash_abort( &op );
-        return( ret );
-    }
-
-    return ret;
-}
-
 #endif /* MBEDTLS_LMS_C */
diff --git a/library/lmots.h b/library/lmots.h
index ec68967..ca7d4bf 100644
--- a/library/lmots.h
+++ b/library/lmots.h
@@ -26,7 +26,7 @@
 #ifndef MBEDTLS_LMOTS_H
 #define MBEDTLS_LMOTS_H
 
-#include "mbedtls/private_access.h"
+#include "mbedtls/build_info.h"
 
 #include "psa/crypto.h"
 
@@ -34,27 +34,19 @@
 #include <stddef.h>
 
 #define MBEDTLS_LMOTS_N_HASH_LEN            (32)
-#define MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN      (34)
+#define MBEDTLS_LMOTS_P_SIG_DIGIT_COUNT     (34)
 #define MBEDTLS_LMOTS_TYPE_LEN              (4)
 #define MBEDTLS_LMOTS_C_RANDOM_VALUE_LEN    (MBEDTLS_LMOTS_N_HASH_LEN)
 #define MBEDTLS_LMOTS_I_KEY_ID_LEN          (16)
 #define MBEDTLS_LMOTS_Q_LEAF_ID_LEN         (4)
 
 #define MBEDTLS_LMOTS_SIG_LEN (MBEDTLS_LMOTS_TYPE_LEN + MBEDTLS_LMOTS_C_RANDOM_VALUE_LEN + \
-                               (MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN * MBEDTLS_LMOTS_N_HASH_LEN))
+                               (MBEDTLS_LMOTS_P_SIG_DIGIT_COUNT * MBEDTLS_LMOTS_N_HASH_LEN))
 
-#define MBEDTLS_LMOTS_PUBKEY_LEN (MBEDTLS_LMOTS_TYPE_LEN + MBEDTLS_LMOTS_I_KEY_ID_LEN + \
+#define MBEDTLS_LMOTS_PUBLIC_KEY_LEN (MBEDTLS_LMOTS_TYPE_LEN + MBEDTLS_LMOTS_I_KEY_ID_LEN + \
                                   MBEDTLS_LMOTS_Q_LEAF_ID_LEN + MBEDTLS_LMOTS_N_HASH_LEN)
 
-#define MBEDTLS_LMOTS_SIG_TYPE_OFFSET      (0)
-#define MBEDTLS_LMOTS_SIG_C_RANDOM_OFFSET  (MBEDTLS_LMOTS_SIG_TYPE_OFFSET     + MBEDTLS_LMOTS_TYPE_LEN)
-#define MBEDTLS_LMOTS_SIG_SIGNATURE_OFFSET (MBEDTLS_LMOTS_SIG_C_RANDOM_OFFSET + MBEDTLS_LMOTS_C_RANDOM_VALUE_LEN)
-
-#define MBEDTLS_LMOTS_PUBKEY_TYPE_OFFSET      (0)
-#define MBEDTLS_LMOTS_PUBKEY_I_KEY_ID_OFFSET  (MBEDTLS_LMOTS_PUBKEY_TYPE_OFFSET      + MBEDTLS_LMOTS_TYPE_LEN)
-#define MBEDTLS_LMOTS_PUBKEY_Q_LEAF_ID_OFFSET (MBEDTLS_LMOTS_PUBKEY_I_KEY_ID_OFFSET  + MBEDTLS_LMOTS_I_KEY_ID_LEN)
-#define MBEDTLS_LMOTS_PUBKEY_KEY_HASH_OFFSET  (MBEDTLS_LMOTS_PUBKEY_Q_LEAF_ID_OFFSET + MBEDTLS_LMOTS_Q_LEAF_ID_LEN)
-
+#define MBEDTLS_LMOTS_SIG_TYPE_OFFSET       (0)
 
 #ifdef __cplusplus
 extern "C" {
@@ -68,54 +60,93 @@
 } mbedtls_lmots_algorithm_type_t;
 
 
-/** LMOTS context structure.
+/** LMOTS parameters structure.
+ *
+ * This contains the metadata associated with an LMOTS key, detailing the
+ * algorithm type, the key ID, and the leaf identifier should be key be part of
+ * a LMS key.
+ */
+typedef struct {
+    unsigned char MBEDTLS_PRIVATE(I_key_identifier[MBEDTLS_LMOTS_I_KEY_ID_LEN]); /*!< The key
+                                                     identifier. */
+    unsigned char MBEDTLS_PRIVATE(q_leaf_identifier[MBEDTLS_LMOTS_Q_LEAF_ID_LEN]); /*!< Which
+                                                      leaf of the LMS key this is.
+                                                      0 if the key is not part of an LMS key. */
+    mbedtls_lmots_algorithm_type_t MBEDTLS_PRIVATE(type); /*!< The LM-OTS key type identifier as
+                                                               per IANA. Only SHA256_N32_W8 is
+                                                               currently supported. */
+} mbedtls_lmots_parameters_t;
+
+/** LMOTS public context structure.
+ *
+ * A LMOTS public key is a hash output, and the applicable parameter set.
  *
  * The context must be initialized before it is used. A public key must either
- * be imported, or an algorithm type set, a private key generated and the public
- * key calculated from it. A context that does not contain a public key cannot
- * verify, and a context that does not contain a private key cannot sign.
- * Signing a message will remove the private key from the context, as private
- * keys can only be used a single time.
+ * be imported or generated from a private context.
  *
  * \dot
- * digraph lmots {
+ * digraph lmots_public_t {
  *   UNINITIALIZED -> INIT [label="init"];
- *   TYPE_SET -> INIT [label="free"];
- *   PRIVATE -> INIT [label="free"];
- *   PUBLIC -> INIT [label="free"];
- *   "PRIVATE+PUBLIC" -> INIT [label="free"];
- *   INIT -> TYPE_SET [label="set_algorithm_type"];
- *   PRIVATE -> TYPE_SET [label="sign"];
- *   "PRIVATE+PUBLIC" -> PUBLIC [label="sign"];
- *   INIT -> PUBLIC [label="import_public"];
- *   PUBLIC -> PUBLIC [label="export_pubkey"];
- *   "PRIVATE+PUBLIC" -> "PRIVATE+PUBLIC" [label="export_pubkey"];
- *   PRIVATE -> "PRIVATE+PUBLIC" [label="gen_pubkey"];
- *   TYPE_SET -> PRIVATE [label="gen_privkey"];
+ *   HAVE_PUBLIC_KEY -> INIT [label="free"];
+ *   INIT -> HAVE_PUBLIC_KEY [label="import_public_key"];
+ *   INIT -> HAVE_PUBLIC_KEY [label="calculate_public_key from private key"];
+ *   HAVE_PUBLIC_KEY -> HAVE_PUBLIC_KEY [label="export_public_key"];
  * }
  * \enddot
  */
 typedef struct {
-    unsigned char MBEDTLS_PRIVATE(have_privkey); /*!< Whether the context contains a private key.
+    mbedtls_lmots_parameters_t MBEDTLS_PRIVATE(params);
+    unsigned char MBEDTLS_PRIVATE(public_key)[32];
+    unsigned char MBEDTLS_PRIVATE(have_public_key); /*!< Whether the context contains a public key.
                                                      Boolean values only. */
-    unsigned char MBEDTLS_PRIVATE(have_pubkey); /*!< Whether the context contains a public key.
+} mbedtls_lmots_public_t;
+
+/** LMOTS private context structure.
+ *
+ * A LMOTS private key is one hash output for each of digit of the digest +
+ * checksum, and the applicable parameter set.
+ *
+ * The context must be initialized before it is used. A public key must either
+ * be imported or generated from a private context.
+ *
+ * \dot
+ * digraph lmots_public_t {
+ *   UNINITIALIZED -> INIT [label="init"];
+ *   HAVE_PRIVATE_KEY -> INIT [label="free"];
+ *   INIT -> HAVE_PRIVATE_KEY [label="generate_private_key"];
+ *   HAVE_PRIVATE_KEY -> INIT [label="sign"];
+ * }
+ * \enddot
+ */
+typedef struct {
+    mbedtls_lmots_parameters_t MBEDTLS_PRIVATE(params);
+    unsigned char MBEDTLS_PRIVATE(private_key)[MBEDTLS_LMOTS_P_SIG_DIGIT_COUNT][32];
+    unsigned char MBEDTLS_PRIVATE(have_private_key); /*!< Whether the context contains a private key.
                                                      Boolean values only. */
-    unsigned char MBEDTLS_PRIVATE(I_key_identifier[MBEDTLS_LMOTS_I_KEY_ID_LEN]); /*!< The key
-                                                     identifier. */
-    unsigned int MBEDTLS_PRIVATE(q_leaf_identifier); /*!< Which leaf of the LMS key this is.
-                                                     0 if the key is not part of an LMS key. */
-    unsigned char MBEDTLS_PRIVATE(q_leaf_identifier_bytes)[MBEDTLS_LMOTS_Q_LEAF_ID_LEN];/*!< The
-                                                     leaf identifier in network bytes form. */
-    mbedtls_lmots_algorithm_type_t MBEDTLS_PRIVATE(type); /*!< The LM-OTS key type identifier as
-                                                     per IANA. Only SHA256_N32_W8 is currently
-                                                     supported. */
-    unsigned char MBEDTLS_PRIVATE(priv_key[MBEDTLS_LMOTS_P_SIG_SYMBOL_LEN][32]); /*!< The private
-                                                     key, one hash output per byte of the encoded
-                                                     symbol string P (32 bytes of hash output +
-                                                     2 bytes of checksum). */
-    unsigned char MBEDTLS_PRIVATE(pub_key[32]); /*!< The public key, in the form of a SHA256
-                                                     output. */
-} mbedtls_lmots_context;
+} mbedtls_lmots_private_t;
+
+/**
+ * \brief                    This function converts an unsigned int into a
+ *                           network-byte-order (big endian) string.
+ *
+ * \param val                The unsigned integer value
+ * \param len                The length of the string.
+ * \param bytes              The string to output into.
+ *
+ * \return                   The corresponding LMS error code.
+ */
+void unsigned_int_to_network_bytes(unsigned int val, size_t len, unsigned char *bytes);
+
+/**
+ * \brief                    This function converts a network-byte-order
+ *                           (big endian) string into an unsigned integer.
+ *
+ * \param len                The length of the string.
+ * \param bytes              The string.
+ *
+ * \return                   The corresponding LMS error code.
+ */
+unsigned int network_bytes_to_unsigned_int(size_t len, const unsigned char *bytes);
 
 /**
  * \brief                    This function converts a \ref psa_status_t to a
@@ -129,35 +160,41 @@
 
 
 /**
- * \brief                    This function initializes an LMOTS context
+ * \brief                    This function initializes a public LMOTS context
  *
  * \param ctx                The uninitialized LMOTS context that will then be
  *                           initialized.
  */
-void mbedtls_lmots_init( mbedtls_lmots_context *ctx );
+void mbedtls_lmots_init_public( mbedtls_lmots_public_t *ctx );
 
 /**
- * \brief                    This function uninitializes an LMOTS context
+ * \brief                    This function uninitializes a public LMOTS context
  *
  * \param ctx                The initialized LMOTS context that will then be
  *                           uninitialized.
  */
-void mbedtls_lmots_free( mbedtls_lmots_context *ctx );
+void mbedtls_lmots_free_public( mbedtls_lmots_public_t *ctx );
 
 /**
- * \brief                    This function sets the type of an LMOTS context
+ * \brief                    This function imports an LMOTS public key into a
+ *                           LMOTS context.
  *
- * \note                     The parameter set in the context will then be used
- *                           for keygen operations etc.
+ * \note                     Before this function is called, the context must
+ *                           have been initialized.
  *
- * \param ctx                The initialized LMOTS context.
- * \param type               The type that will be set in the context.
+ * \note                     See IETF RFC8554 for details of the encoding of
+ *                           this public key.
+ *
+ * \param ctx                The initialized LMOTS context store the key in.
+ * \param key                The buffer from which the key will be read.
+ *                           #MBEDTLS_LMOTS_PUBLIC_KEY_LEN bytes will be read from
+ *                           this.
  *
  * \return         \c 0 on success.
  * \return         A non-zero error code on failure.
  */
-int mbedtls_lmots_set_algorithm_type( mbedtls_lmots_context *ctx,
-                                      mbedtls_lmots_algorithm_type_t type );
+int mbedtls_lmots_import_public_key( mbedtls_lmots_public_t *ctx,
+                                     const unsigned char *key, size_t key_size );
 
 /**
  * \brief                    This function creates a candidate public key from
@@ -170,12 +207,10 @@
  *                           mbedtls_lmots_verify will be used for LMOTS
  *                           signature verification.
  *
- * \param I_key_identifier   The key identifier of the key, as a 16-byte string.
- * \param q_leaf_identifier  The leaf identifier of key. If this LMOTS key is
- *                           not being used as part of an LMS key, this should
- *                           be set to 0.
+ * \param params             The LMOTS parameter set, q and I values as an
+ *                           mbedtls_lmots_parameters_t struct.
  * \param msg                The buffer from which the message will be read.
- * \param msg_len            The size of the message that will be read.
+ * \param msg_size           The size of the message that will be read.
  * \param sig                The buffer from which the signature will be read.
  *                           #MBEDTLS_LMOTS_SIG_LEN bytes will be read from this.
  * \param out                The buffer where the candidate public key will be
@@ -185,14 +220,129 @@
  * \return         \c 0 on success.
  * \return         A non-zero error code on failure.
  */
-int mbedtls_lmots_generate_pub_key_candidate( const unsigned char I_key_identifier[MBEDTLS_LMOTS_I_KEY_ID_LEN],
-                                              const unsigned char q_leaf_identifier[MBEDTLS_LMOTS_Q_LEAF_ID_LEN],
-                                              const unsigned char  *msg,
-                                              size_t msg_len,
-                                              const unsigned char *sig,
-                                              unsigned char *out );
+int mbedtls_lmots_calculate_public_key_candidate( const mbedtls_lmots_parameters_t *params,
+                                                  const unsigned char *msg,
+                                                  size_t msg_size,
+                                                  const unsigned char *sig,
+                                                  size_t sig_size,
+                                                  unsigned char *out,
+                                                  size_t out_size,
+                                                  size_t *out_len);
 
 /**
+ * \brief                    This function verifies a LMOTS signature, using a
+ *                           LMOTS context that contains a public key.
+ *
+ * \warning                  This function is **not intended for use in
+ *                           production**, due to as-yet unsolved problems with
+ *                           handling stateful keys.
+ *
+ * \note                     Before this function is called, the context must
+ *                           have been initialized and must contain a public key
+ *                           (either by import or calculation from a private key).
+ *
+ * \param ctx                The initialized LMOTS context from which the public
+ *                           key will be read.
+ * \param msg                The buffer from which the message will be read.
+ * \param msg_size           The size of the message that will be read.
+ * \param sig                The buf from which the signature will be read.
+ *                           #MBEDTLS_LMOTS_SIG_LEN bytes will be read from
+ *                           this.
+ *
+ * \return         \c 0 on successful verification.
+ * \return         A non-zero error code on failure.
+ */
+int mbedtls_lmots_verify( mbedtls_lmots_public_t *ctx, const unsigned char *msg,
+                          size_t msg_size, const unsigned char *sig,
+                          size_t sig_size );
+
+/**
+ * \brief                    This function initializes a private LMOTS context
+ *
+ * \param ctx                The uninitialized LMOTS context that will then be
+ *                           initialized.
+ */
+void mbedtls_lmots_init_private( mbedtls_lmots_private_t *ctx );
+
+/**
+ * \brief                    This function uninitializes a private LMOTS context
+ *
+ * \param ctx                The initialized LMOTS context that will then be
+ *                           uninitialized.
+ */
+void mbedtls_lmots_free_private( mbedtls_lmots_private_t *ctx );
+
+/**
+ * \brief                    This function generates an LMOTS private key, and
+ *                           stores in into an LMOTS context.
+ *
+ * \warning                  This function is **not intended for use in
+ *                           production**, due to as-yet unsolved problems with
+ *                           handling stateful keys.
+ *
+ * \note                     The seed must have at least 256 bits of entropy.
+ *
+ * \param ctx                The initialized LMOTS context to generate the key
+ *                           into.
+ * \param I_key_identifier   The key identifier of the key, as a 16-byte string.
+ * \param q_leaf_identifier  The leaf identifier of key. If this LMOTS key is
+ *                           not being used as part of an LMS key, this should
+ *                           be set to 0.
+ * \param seed               The seed used to deterministically generate the
+ *                           key.
+ * \param seed_size          The length of the seed.
+ *
+ * \return         \c 0 on success.
+ * \return         A non-zero error code on failure.
+ */
+int mbedtls_lmots_generate_private_key( mbedtls_lmots_private_t *ctx,
+                                        mbedtls_lmots_algorithm_type_t type,
+                                        const unsigned char I_key_identifier[MBEDTLS_LMOTS_I_KEY_ID_LEN],
+                                        uint32_t q_leaf_identifier,
+                                        const unsigned char *seed,
+                                        size_t seed_size );
+
+/**
+ * \brief                    This function generates an LMOTS public key from a
+ *                           LMOTS context that already contains a private key.
+ *
+ * \note                     Before this function is called, the context must
+ *                           have been initialized and the context must contain
+ *                           a private key.
+ *
+ * \param ctx                The initialized LMOTS context to generate the key
+ *                           from and store it into.
+ *
+ * \return         \c 0 on success.
+ * \return         A non-zero error code on failure.
+ */
+int mbedtls_lmots_calculate_public_key( mbedtls_lmots_public_t *ctx,
+                                        mbedtls_lmots_private_t *priv_ctx);
+
+
+/**
+ * \brief                    This function exports an LMOTS public key from a
+ *                           LMOTS context that already contains a public key.
+ *
+ * \note                     Before this function is called, the context must
+ *                           have been initialized and the context must contain
+ *                           a public key.
+ *
+ * \note                     See IETF RFC8554 for details of the encoding of
+ *                           this public key.
+ *
+ * \param ctx                The initialized LMOTS context that contains the
+ *                           publc key.
+ * \param key                The buffer into which the key will be output. Must
+ *                           be at least #MBEDTLS_LMOTS_PUBLIC_KEY_LEN in size.
+ *
+ * \return         \c 0 on success.
+ * \return         A non-zero error code on failure.
+ */
+int mbedtls_lmots_export_public_key( mbedtls_lmots_public_t *ctx,
+                                     unsigned char *key, size_t key_size,
+                                     size_t *key_len );
+/**
  * \brief                    This function creates a LMOTS signature, using a
  *                           LMOTS context that contains a private key.
  *
@@ -213,135 +363,18 @@
  *                           generation.
  * \param p_rng              The RNG context to be passed to f_rng
  * \param msg                The buffer from which the message will be read.
- * \param msg_len            The size of the message that will be read.
+ * \param msg_size           The size of the message that will be read.
  * \param sig                The buf into which the signature will be stored.
  *                           Must be at least #MBEDTLS_LMOTS_SIG_LEN in size.
  *
  * \return         \c 0 on success.
  * \return         A non-zero error code on failure.
  */
-int mbedtls_lmots_sign( mbedtls_lmots_context *ctx,
+int mbedtls_lmots_sign( mbedtls_lmots_private_t *ctx,
                         int (*f_rng)(void *, unsigned char *, size_t),
-                        void *p_rng, const unsigned char *msg, size_t msg_len,
-                        unsigned char *sig );
+                        void *p_rng, const unsigned char *msg, size_t msg_size,
+                        unsigned char *sig, size_t sig_size, size_t* sig_len );
 
-/**
- * \brief                    This function verifies a LMOTS signature, using a
- *                           LMOTS context that contains a public key.
- *
- * \warning                  This function is **not intended for use in
- *                           production**, due to as-yet unsolved problems with
- *                           handling stateful keys.
- *
- * \note                     Before this function is called, the context must
- *                           have been initialized and must contain a public key
- *                           (either by import or generation).
- *
- * \param ctx                The initialized LMOTS context from which the public
- *                           key will be read.
- * \param msg                The buffer from which the message will be read.
- * \param msg_len            The size of the message that will be read.
- * \param sig                The buf from which the signature will be read.
- *                           #MBEDTLS_LMOTS_SIG_LEN bytes will be read from
- *                           this.
- *
- * \return         \c 0 on successful verification.
- * \return         A non-zero error code on failure.
- */
-int mbedtls_lmots_verify( mbedtls_lmots_context *ctx, const unsigned char *msg,
-                          size_t msg_len, const unsigned char *sig );
-
-/**
- * \brief                    This function imports an LMOTS public key into a
- *                           LMOTS context.
- *
- * \note                     Before this function is called, the context must
- *                           have been initialized.
- *
- * \note                     See IETF RFC8554 for details of the encoding of
- *                           this public key.
- *
- * \param ctx                The initialized LMOTS context store the key in.
- * \param key                The buffer from which the key will be read.
- *                           #MBEDTLS_LMOTS_PUBKEY_LEN bytes will be read from
- *                           this.
- *
- * \return         \c 0 on success.
- * \return         A non-zero error code on failure.
- */
-int mbedtls_lmots_import_pubkey( mbedtls_lmots_context *ctx,
-                                 const unsigned char *key );
-
-/**
- * \brief                    This function exports an LMOTS public key from a
- *                           LMOTS context that already contains a public key.
- *
- * \note                     Before this function is called, the context must
- *                           have been initialized and the context must contain
- *                           a public key.
- *
- * \note                     See IETF RFC8554 for details of the encoding of
- *                           this public key.
- *
- * \param ctx                The initialized LMOTS context that contains the
- *                           publc key.
- * \param key                The buffer into which the key will be output. Must
- *                           be at least #MBEDTLS_LMOTS_PUBKEY_LEN in size.
- *
- * \return         \c 0 on success.
- * \return         A non-zero error code on failure.
- */
-int mbedtls_lmots_export_pubkey( mbedtls_lmots_context *ctx,
-                                 unsigned char *key );
-
-/**
- * \brief                    This function generates an LMOTS public key from a
- *                           LMOTS context that already contains a private key.
- *
- * \note                     Before this function is called, the context must
- *                           have been initialized and the context must contain
- *                           a private key.
- *
- * \param ctx                The initialized LMOTS context to generate the key
- *                           from and store it into.
- *
- * \return         \c 0 on success.
- * \return         A non-zero error code on failure.
- */
-int mbedtls_lmots_gen_pubkey( mbedtls_lmots_context *ctx );
-
-/**
- * \brief                    This function generates an LMOTS private key, and
- *                           stores in into an LMOTS context.
- *
- * \warning                  This function is **not intended for use in
- *                           production**, due to as-yet unsolved problems with
- *                           handling stateful keys.
- *
- * \note                     Before this function is called, the context must
- *                           have been initialized and the type of the LMOTS
- *                           context set using mbedtls_lmots_set_algorithm_type
- *
- * \note                     The seed must have at least 256 bits of entropy.
- *
- * \param ctx                The initialized LMOTS context to generate the key
- *                           into.
- * \param I_key_identifier   The key identifier of the key, as a 16-byte string.
- * \param q_leaf_identifier  The leaf identifier of key. If this LMOTS key is
- *                           not being used as part of an LMS key, this should
- *                           be set to 0.
- * \param seed               The seed used to deterministically generate the
- *                           key.
- * \param seed_len           The length of the seed.
- *
- * \return         \c 0 on success.
- * \return         A non-zero error code on failure.
- */
-int mbedtls_lmots_gen_privkey( mbedtls_lmots_context *ctx,
-                               const unsigned char I_key_identifier[MBEDTLS_LMOTS_I_KEY_ID_LEN],
-                               unsigned int q_leaf_identifier,
-                               const unsigned char *seed,
-                               size_t seed_len );
 
 #ifdef __cplusplus
 }
diff --git a/library/lms.c b/library/lms.c
index 4b4f151..b58aeea 100644
--- a/library/lms.c
+++ b/library/lms.c
@@ -54,45 +54,33 @@
 #define mbedtls_free   free
 #endif
 
-#define MERKLE_TREE_NODE_AM (1 << (MBEDTLS_LMS_H_TREE_HEIGHT + 1))
-#define MERKLE_TREE_LEAF_AM (1 << MBEDTLS_LMS_H_TREE_HEIGHT)
-#define MERKLE_TREE_INTR_AM (1 << MBEDTLS_LMS_H_TREE_HEIGHT)
+#define MBEDTLS_LMS_SIG_Q_LEAF_ID_OFFSET    (0)
+#define MBEDTLS_LMS_SIG_OTS_SIG_OFFSET      (MBEDTLS_LMS_SIG_Q_LEAF_ID_OFFSET + MBEDTLS_LMOTS_Q_LEAF_ID_LEN)
+#define MBEDTLS_LMS_SIG_TYPE_OFFSET         (MBEDTLS_LMS_SIG_OTS_SIG_OFFSET   + MBEDTLS_LMOTS_SIG_LEN)
+#define MBEDTLS_LMS_SIG_PATH_OFFSET         (MBEDTLS_LMS_SIG_TYPE_OFFSET      + MBEDTLS_LMS_TYPE_LEN)
+
+#define MBEDTLS_LMS_PUBLIC_KEY_TYPE_OFFSET      (0)
+#define MBEDTLS_LMS_PUBLIC_KEY_OTSTYPE_OFFSET   (MBEDTLS_LMS_PUBLIC_KEY_TYPE_OFFSET     + MBEDTLS_LMS_TYPE_LEN)
+#define MBEDTLS_LMS_PUBLIC_KEY_I_KEY_ID_OFFSET  (MBEDTLS_LMS_PUBLIC_KEY_OTSTYPE_OFFSET  + MBEDTLS_LMOTS_TYPE_LEN)
+#define MBEDTLS_LMS_PUBLIC_KEY_ROOT_NODE_OFFSET (MBEDTLS_LMS_PUBLIC_KEY_I_KEY_ID_OFFSET + MBEDTLS_LMOTS_I_KEY_ID_LEN)
+
+
+#define MERKLE_TREE_NODE_AM (1u << (MBEDTLS_LMS_H_TREE_HEIGHT + 1u))
+#define MERKLE_TREE_LEAF_NODE_AM (1u << MBEDTLS_LMS_H_TREE_HEIGHT)
+#define MERKLE_TREE_INTERNAL_NODE_AM (1u << MBEDTLS_LMS_H_TREE_HEIGHT)
 
 #define D_CONST_LEN           (2)
+static const unsigned char D_LEAF_CONSTANT_BYTES[D_CONST_LEN] = {0x82, 0x82};
+static const unsigned char D_INTERNAL_CONSTANT_BYTES[D_CONST_LEN] = {0x83, 0x83};
 
-#define D_LEAF_CONSTANT     (0x8282)
-#define D_INTR_CONSTANT     (0x8383)
-
-static void val_to_network_bytes(unsigned int val, size_t len, unsigned char *bytes)
-{
-    size_t idx;
-
-    for (idx = 0; idx < len; idx++) {
-        bytes[idx] = (val >> ((len - 1 - idx) * 8)) & 0xFF;
-    }
-}
-
-static unsigned int network_bytes_to_val(size_t len, const unsigned char *bytes)
-{
-    size_t idx;
-    unsigned int val = 0;
-
-    for (idx = 0; idx < len; idx++) {
-        val |= ((unsigned int)bytes[idx]) << (8 * (len - 1 - idx));
-    }
-
-    return val;
-}
-
-static int create_merkle_leaf_node( const mbedtls_lms_context *ctx,
+static int create_merkle_leaf_node( const unsigned char I_key_identifier[MBEDTLS_LMOTS_I_KEY_ID_LEN],
                                     unsigned char pub_key[MBEDTLS_LMOTS_N_HASH_LEN],
                                     unsigned int r_node_idx,
-                                    unsigned char out[32] )
+                                    unsigned char out[MBEDTLS_LMS_M_NODE_BYTES] )
 {
     psa_hash_operation_t op;
     psa_status_t status;
     size_t output_hash_len;
-    unsigned char D_LEAF_bytes[D_CONST_LEN];
     unsigned char r_node_idx_bytes[4];
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
@@ -100,64 +88,49 @@
     status = psa_hash_setup( &op, PSA_ALG_SHA_256 );
     ret = mbedtls_lms_error_from_psa( status );
     if ( ret != 0 )
-    {
-        goto out;
-    }
+        goto exit;
 
-    status = psa_hash_update( &op, ctx->MBEDTLS_PRIVATE(I_key_identifier),
-                              MBEDTLS_LMOTS_I_KEY_ID_LEN );
+    status = psa_hash_update( &op, I_key_identifier, MBEDTLS_LMOTS_I_KEY_ID_LEN );
     ret = mbedtls_lms_error_from_psa( status );
     if( ret )
-    {
-        goto out;
-    }
+        goto exit;
 
-    val_to_network_bytes( r_node_idx, 4, r_node_idx_bytes );
+    unsigned_int_to_network_bytes( r_node_idx, 4, r_node_idx_bytes );
     status = psa_hash_update( &op, r_node_idx_bytes, 4 );
     ret = mbedtls_lms_error_from_psa( status );
     if( ret )
-    {
-        goto out;
-    }
+        goto exit;
 
-    val_to_network_bytes( D_LEAF_CONSTANT, D_CONST_LEN, D_LEAF_bytes );
-    status = psa_hash_update( &op, D_LEAF_bytes, D_CONST_LEN );
+    status = psa_hash_update( &op, D_LEAF_CONSTANT_BYTES, D_CONST_LEN );
     ret = mbedtls_lms_error_from_psa( status );
     if( ret )
-    {
-        goto out;
-    }
+        goto exit;
 
     status = psa_hash_update( &op, pub_key, MBEDTLS_LMOTS_N_HASH_LEN );
     ret = mbedtls_lms_error_from_psa( status );
     if( ret )
-    {
-        goto out;
-    }
+        goto exit;
 
-    status = psa_hash_finish( &op, out, 32, &output_hash_len);
+    status = psa_hash_finish( &op, out, MBEDTLS_LMS_M_NODE_BYTES, &output_hash_len);
     ret = mbedtls_lms_error_from_psa( status );
     if( ret )
-    {
-        goto out;
-    }
+        goto exit;
 
-out:
+exit:
     psa_hash_abort( &op );
 
     return( ret );
 }
 
-static int create_merkle_intr_node( const mbedtls_lms_context *ctx,
-                                    const unsigned char left_node[32],
-                                    const unsigned char rght_node[32],
-                                    unsigned int r_node_idx,
-                                    unsigned char out[32] )
+static int create_merkle_internal_node( const unsigned char I_key_identifier[MBEDTLS_LMOTS_I_KEY_ID_LEN],
+                                        const unsigned char left_node[MBEDTLS_LMS_M_NODE_BYTES],
+                                        const unsigned char right_node[MBEDTLS_LMS_M_NODE_BYTES],
+                                        unsigned int r_node_idx,
+                                        unsigned char out[MBEDTLS_LMS_M_NODE_BYTES] )
 {
     psa_hash_operation_t op;
     psa_status_t status;
     size_t output_hash_len;
-    unsigned char D_INTR_bytes[D_CONST_LEN];
     unsigned char r_node_idx_bytes[4];
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
@@ -165,75 +138,62 @@
     status = psa_hash_setup( &op, PSA_ALG_SHA_256 );
     ret = mbedtls_lms_error_from_psa( status );
     if ( ret != 0 )
-    {
-        goto out;
-    }
+        goto exit;
 
-    status = psa_hash_update( &op, ctx->MBEDTLS_PRIVATE(I_key_identifier),
-                              MBEDTLS_LMOTS_I_KEY_ID_LEN );
+    status = psa_hash_update( &op, I_key_identifier, MBEDTLS_LMOTS_I_KEY_ID_LEN );
     ret = mbedtls_lms_error_from_psa( status );
     if( ret )
-    {
-        goto out;
-    }
+        goto exit;
 
-    val_to_network_bytes( r_node_idx, 4, r_node_idx_bytes );
+    unsigned_int_to_network_bytes( r_node_idx, 4, r_node_idx_bytes );
     status = psa_hash_update( &op, r_node_idx_bytes, 4 );
     ret = mbedtls_lms_error_from_psa( status );
     if( ret )
-    {
-        goto out;
-    }
+        goto exit;
 
-    val_to_network_bytes( D_INTR_CONSTANT, D_CONST_LEN, D_INTR_bytes );
-    status = psa_hash_update( &op, D_INTR_bytes, D_CONST_LEN );
+    status = psa_hash_update( &op, D_INTERNAL_CONSTANT_BYTES, D_CONST_LEN );
     ret = mbedtls_lms_error_from_psa( status );
     if( ret )
-    {
-        goto out;
-    }
+        goto exit;
 
     status = psa_hash_update( &op, left_node, MBEDTLS_LMOTS_N_HASH_LEN );
     ret = mbedtls_lms_error_from_psa( status );
     if( ret )
-    {
-        goto out;
-    }
+        goto exit;
 
-    status = psa_hash_update( &op, rght_node, MBEDTLS_LMOTS_N_HASH_LEN );
+    status = psa_hash_update( &op, right_node, MBEDTLS_LMOTS_N_HASH_LEN );
     ret = mbedtls_lms_error_from_psa( status );
     if( ret )
-    {
-        goto out;
-    }
+        goto exit;
 
-    ret = psa_hash_finish( &op, out, 32, &output_hash_len);
+    ret = psa_hash_finish( &op, out, MBEDTLS_LMS_M_NODE_BYTES, &output_hash_len);
     ret = mbedtls_lms_error_from_psa( status );
     if( ret )
-    {
-        goto out;
-    }
+        goto exit;
 
-out:
+exit:
     psa_hash_abort( &op );
 
     return ret;
 }
 
-static int generate_merkle_tree( mbedtls_lms_context *ctx,
-                                 unsigned char tree[MERKLE_TREE_NODE_AM][32] )
+static int calculate_merkle_tree( mbedtls_lms_private_t *ctx,
+                                 unsigned char tree[MERKLE_TREE_NODE_AM][MBEDTLS_LMS_M_NODE_BYTES] )
 {
     unsigned int priv_key_idx;
     unsigned int r_node_idx;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
     /* First create the leaf nodes, in ascending order */
-    for( priv_key_idx = 0; priv_key_idx < MERKLE_TREE_INTR_AM; priv_key_idx++ )
+    for( priv_key_idx = 0; priv_key_idx < MERKLE_TREE_INTERNAL_NODE_AM;
+         priv_key_idx++ )
     {
-        r_node_idx = MERKLE_TREE_INTR_AM + priv_key_idx;
+        r_node_idx = MERKLE_TREE_INTERNAL_NODE_AM + priv_key_idx;
 
-        ret = create_merkle_leaf_node( ctx, ctx->MBEDTLS_PRIVATE(priv_keys)[priv_key_idx].pub_key,
-                                       r_node_idx, tree[r_node_idx] );
+        ret = create_merkle_leaf_node(
+            ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(I_key_identifier),
+            ctx->MBEDTLS_PRIVATE(ots_public_keys)[priv_key_idx].MBEDTLS_PRIVATE(public_key),
+            r_node_idx, tree[r_node_idx] );
         if( ret )
         {
             return( ret );
@@ -242,11 +202,12 @@
 
     /* Then the internal nodes, in reverse order so that we can guarantee the
      * parent has been created */
-    for( r_node_idx = MERKLE_TREE_INTR_AM - 1; r_node_idx > 0; r_node_idx-- )
+    for( r_node_idx = MERKLE_TREE_INTERNAL_NODE_AM - 1; r_node_idx > 0;
+         r_node_idx-- )
     {
-        ret = create_merkle_intr_node( ctx, tree[(r_node_idx * 2)],
-                                       tree[(r_node_idx * 2 + 1)],
-                                       r_node_idx, tree[r_node_idx] );
+        ret = create_merkle_internal_node(
+            ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(I_key_identifier),
+            tree[(r_node_idx * 2)], tree[(r_node_idx * 2 + 1)], r_node_idx, tree[r_node_idx] );
         if( ret )
         {
             return( ret );
@@ -256,18 +217,17 @@
     return( 0 );
 }
 
-static int get_merkle_path( mbedtls_lms_context *ctx,
-                            unsigned int leaf_node_id, unsigned char path[MBEDTLS_LMS_H_TREE_HEIGHT][32] )
+static int get_merkle_path( mbedtls_lms_private_t *ctx,
+                            unsigned int leaf_node_id,
+                            unsigned char path[MBEDTLS_LMS_H_TREE_HEIGHT][MBEDTLS_LMS_M_NODE_BYTES] )
 {
-    unsigned char tree[MERKLE_TREE_NODE_AM][32];
+    unsigned char tree[MERKLE_TREE_NODE_AM][MBEDTLS_LMS_M_NODE_BYTES];
     unsigned int curr_node_id = leaf_node_id;
-    unsigned int parent_node_id;
-    unsigned char sibling_relative_id;
     unsigned int adjacent_node_id;
     unsigned int height;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
-    ret = generate_merkle_tree( ctx, tree);
+    ret = calculate_merkle_tree( ctx, tree);
     if( ret )
     {
         return( ret );
@@ -275,195 +235,107 @@
 
     for( height = 0; height < MBEDTLS_LMS_H_TREE_HEIGHT; height++ )
     {
-        parent_node_id = ( curr_node_id / 2 );
-
-        /* 0 if the node is a left child, 1 if the node is a right child */
-        sibling_relative_id = curr_node_id & 1;
-
-        adjacent_node_id = ( parent_node_id * 2 ) + ( 1 - sibling_relative_id );
+        adjacent_node_id = curr_node_id ^ 1;
 
         memcpy( &path[height], &tree[adjacent_node_id], MBEDTLS_LMOTS_N_HASH_LEN );
 
-        curr_node_id = parent_node_id;
+        curr_node_id >>=1;
     }
 
     return( 0 );
 }
 
-void mbedtls_lms_init( mbedtls_lms_context *ctx )
+void mbedtls_lms_init_public( mbedtls_lms_public_t *ctx )
 {
-    if( ctx == NULL )
-    {
-        return;
-    }
-
-    mbedtls_platform_zeroize( ctx, sizeof( mbedtls_lms_context ) ) ;
+    mbedtls_platform_zeroize( ctx, sizeof( mbedtls_lms_public_t ) ) ;
 }
 
-void mbedtls_lms_free( mbedtls_lms_context *ctx )
+void mbedtls_lms_free_public( mbedtls_lms_public_t *ctx )
 {
-    unsigned int idx;
-
-    if( ctx == NULL )
-    {
-        return;
-    }
-
-    if( ctx->MBEDTLS_PRIVATE(have_privkey) )
-    {
-        for( idx = 0; idx < MERKLE_TREE_LEAF_AM; idx++ )
-        {
-            mbedtls_lmots_free( &ctx->MBEDTLS_PRIVATE(priv_keys)[idx] );
-        }
-
-        mbedtls_free( ctx->MBEDTLS_PRIVATE(priv_keys) );
-    }
-
-    mbedtls_platform_zeroize( ctx, sizeof( mbedtls_lms_context ) );
+    mbedtls_platform_zeroize( ctx, sizeof( mbedtls_lms_public_t ) );
 }
 
-int mbedtls_lms_set_algorithm_type( mbedtls_lms_context *ctx,
-                                    mbedtls_lms_algorithm_type_t type,
-                                    mbedtls_lmots_algorithm_type_t otstype )
+int mbedtls_lms_import_public_key( mbedtls_lms_public_t *ctx,
+                               const unsigned char *key, size_t key_size )
 {
-    if( ctx == NULL )
+    mbedtls_lms_algorithm_type_t type;
+    mbedtls_lmots_algorithm_type_t otstype;
+
+    if( key_size < MBEDTLS_LMS_PUBLIC_KEY_LEN )
+    {
+        return( MBEDTLS_ERR_LMS_BUFFER_TOO_SMALL );
+    }
+
+    type = network_bytes_to_unsigned_int( MBEDTLS_LMS_TYPE_LEN, key + MBEDTLS_LMS_PUBLIC_KEY_TYPE_OFFSET );
+    if( type != MBEDTLS_LMS_SHA256_M32_H10 )
     {
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
+    ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(type) = type;
 
-    ctx->MBEDTLS_PRIVATE(type) = type;
-    ctx->MBEDTLS_PRIVATE(otstype) = otstype;
+    otstype = network_bytes_to_unsigned_int( MBEDTLS_LMOTS_TYPE_LEN,
+                                    key + MBEDTLS_LMS_PUBLIC_KEY_OTSTYPE_OFFSET );
+    if( otstype != MBEDTLS_LMOTS_SHA256_N32_W8 )
+    {
+        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
+    }
+    ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(otstype) = otstype;
+
+    memcpy( ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(I_key_identifier),
+            key + MBEDTLS_LMS_PUBLIC_KEY_I_KEY_ID_OFFSET,
+            MBEDTLS_LMOTS_I_KEY_ID_LEN );
+    memcpy( ctx->MBEDTLS_PRIVATE(T_1_pub_key), key + MBEDTLS_LMS_PUBLIC_KEY_ROOT_NODE_OFFSET,
+            MBEDTLS_LMOTS_N_HASH_LEN );
+
+    ctx->MBEDTLS_PRIVATE(have_public_key) = 1;
 
     return( 0 );
 }
 
-int mbedtls_lms_sign( mbedtls_lms_context *ctx,
-                      int ( *f_rng)(void *, unsigned char *, size_t),
-                      void* p_rng, unsigned char *msg, unsigned int msg_len,
-                      unsigned char *sig )
-{
-    unsigned int q_leaf_identifier;
-    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-
-    if( ctx == NULL )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    if( ! ctx->MBEDTLS_PRIVATE(have_privkey) )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    if( msg == NULL )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    if( sig == NULL )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-
-    if( ctx->MBEDTLS_PRIVATE(type) != MBEDTLS_LMS_SHA256_M32_H10 )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    if( ctx->MBEDTLS_PRIVATE(otstype) != MBEDTLS_LMOTS_SHA256_N32_W8 )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-
-    if( ctx->MBEDTLS_PRIVATE(q_next_usable_key) >= MERKLE_TREE_LEAF_AM )
-    {
-        return( MBEDTLS_ERR_LMS_OUT_OF_PRIV_KEYS );
-    }
-
-
-    q_leaf_identifier = ctx->MBEDTLS_PRIVATE(q_next_usable_key);
-    /* This new value must _always_ be written back to the disk before the
-     * signature is returned.
-     */
-    ctx->MBEDTLS_PRIVATE(q_next_usable_key) += 1;
-
-    ret = mbedtls_lmots_sign( &ctx->MBEDTLS_PRIVATE(priv_keys)[q_leaf_identifier],
-                              f_rng, p_rng, msg, msg_len,
-                              sig + MBEDTLS_LMS_SIG_OTS_SIG_OFFSET );
-    if( ret )
-    {
-        return( ret );
-    }
-
-    val_to_network_bytes( ctx->MBEDTLS_PRIVATE(type), MBEDTLS_LMS_TYPE_LEN,
-                          sig + MBEDTLS_LMS_SIG_TYPE_OFFSET );
-    val_to_network_bytes( q_leaf_identifier, MBEDTLS_LMOTS_Q_LEAF_ID_LEN,
-                          sig + MBEDTLS_LMS_SIG_Q_LEAF_ID_OFFSET);
-
-    ret = get_merkle_path( ctx, MERKLE_TREE_INTR_AM + q_leaf_identifier,
-                           ( unsigned char( * )[32] )( sig + MBEDTLS_LMS_SIG_PATH_OFFSET ) );
-    if( ret )
-    {
-        return( ret );
-    }
-
-    return( 0 );
-}
-
-int mbedtls_lms_verify( const mbedtls_lms_context *ctx,
-                        const unsigned char *msg, unsigned int msg_len,
-                        const unsigned char *sig )
+int mbedtls_lms_verify( const mbedtls_lms_public_t *ctx,
+                        const unsigned char *msg, size_t msg_size,
+                        const unsigned char *sig, size_t sig_size )
 {
     unsigned int q_leaf_identifier;
     unsigned char Kc_candidate_ots_pub_key[MBEDTLS_LMOTS_N_HASH_LEN];
-    unsigned char Tc_candidate_root_node[32];
+    unsigned char Tc_candidate_root_node[MBEDTLS_LMS_M_NODE_BYTES];
     unsigned int height;
     unsigned int curr_node_id;
     unsigned int parent_node_id;
     const unsigned char* left_node;
-    const unsigned char* rght_node;
+    const unsigned char* right_node;
+    mbedtls_lmots_parameters_t ots_params;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
-    if( ctx == NULL )
+    if( ! ctx->MBEDTLS_PRIVATE(have_public_key) )
     {
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
 
-    if( ! ctx->MBEDTLS_PRIVATE(have_pubkey) )
+    if( sig_size != MBEDTLS_LMS_SIG_LEN )
     {
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
 
-    if( msg == NULL)
+    if( ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(type)
+        != MBEDTLS_LMS_SHA256_M32_H10 )
     {
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
 
-    if( sig == NULL)
+    if( ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(otstype)
+        != MBEDTLS_LMOTS_SHA256_N32_W8 )
     {
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
 
-    if( ctx->MBEDTLS_PRIVATE(type) != MBEDTLS_LMS_SHA256_M32_H10 )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    if( ctx->MBEDTLS_PRIVATE(otstype) != MBEDTLS_LMOTS_SHA256_N32_W8 )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-
-    if( network_bytes_to_val( MBEDTLS_LMS_TYPE_LEN,
+    if( network_bytes_to_unsigned_int( MBEDTLS_LMS_TYPE_LEN,
                               sig + MBEDTLS_LMS_SIG_TYPE_OFFSET) != MBEDTLS_LMS_SHA256_M32_H10 )
     {
         return( MBEDTLS_ERR_LMS_VERIFY_FAILED );
     }
 
-    if( network_bytes_to_val( MBEDTLS_LMOTS_TYPE_LEN,
+    if( network_bytes_to_unsigned_int( MBEDTLS_LMOTS_TYPE_LEN,
                               sig + MBEDTLS_LMS_SIG_OTS_SIG_OFFSET + MBEDTLS_LMOTS_SIG_TYPE_OFFSET)
         != MBEDTLS_LMOTS_SHA256_N32_W8 )
     {
@@ -471,29 +343,39 @@
     }
 
 
-    q_leaf_identifier = network_bytes_to_val( MBEDTLS_LMOTS_Q_LEAF_ID_LEN,
+    q_leaf_identifier = network_bytes_to_unsigned_int( MBEDTLS_LMOTS_Q_LEAF_ID_LEN,
                                               sig + MBEDTLS_LMS_SIG_Q_LEAF_ID_OFFSET );
 
-    if( q_leaf_identifier >= MERKLE_TREE_LEAF_AM )
+    if( q_leaf_identifier >= MERKLE_TREE_LEAF_NODE_AM )
     {
         return( MBEDTLS_ERR_LMS_VERIFY_FAILED );
     }
 
-    ret = mbedtls_lmots_generate_pub_key_candidate( ctx->MBEDTLS_PRIVATE(I_key_identifier),
-                                                    sig + MBEDTLS_LMS_SIG_Q_LEAF_ID_OFFSET,
-                                                    msg, msg_len,
-                                                    sig + MBEDTLS_LMS_SIG_OTS_SIG_OFFSET,
-                                                    Kc_candidate_ots_pub_key );
+    memcpy(ots_params.MBEDTLS_PRIVATE(I_key_identifier),
+           ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(I_key_identifier),
+           MBEDTLS_LMOTS_I_KEY_ID_LEN);
+    unsigned_int_to_network_bytes( q_leaf_identifier,
+                                   MBEDTLS_LMOTS_Q_LEAF_ID_LEN,
+                                   ots_params.MBEDTLS_PRIVATE(q_leaf_identifier) );
+    ots_params.type = ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(otstype);
+
+    ret = mbedtls_lmots_calculate_public_key_candidate( &ots_params, msg, msg_size,
+                                                        sig + MBEDTLS_LMS_SIG_OTS_SIG_OFFSET,
+                                                        MBEDTLS_LMOTS_SIG_LEN,
+                                                        Kc_candidate_ots_pub_key,
+                                                        sizeof(Kc_candidate_ots_pub_key),
+                                                        NULL );
     if( ret )
     {
         return( ret );
     }
 
-    create_merkle_leaf_node( ctx, Kc_candidate_ots_pub_key,
-                             MERKLE_TREE_INTR_AM + q_leaf_identifier,
-                             Tc_candidate_root_node );
+    create_merkle_leaf_node(
+            ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(I_key_identifier),
+            Kc_candidate_ots_pub_key, MERKLE_TREE_INTERNAL_NODE_AM + q_leaf_identifier,
+            Tc_candidate_root_node );
 
-    curr_node_id = MERKLE_TREE_INTR_AM + q_leaf_identifier;
+    curr_node_id = MERKLE_TREE_INTERNAL_NODE_AM + q_leaf_identifier;
 
     for( height = 0; height < MBEDTLS_LMS_H_TREE_HEIGHT; height++ )
     {
@@ -502,17 +384,18 @@
         /* Left/right node ordering matters for the hash */
         if( curr_node_id & 1 )
         {
-            left_node = ( ( const unsigned char( * )[32] )( sig + MBEDTLS_LMS_SIG_PATH_OFFSET ) )[height];
-            rght_node = Tc_candidate_root_node;
+            left_node = ( ( const unsigned char( * )[MBEDTLS_LMS_M_NODE_BYTES] )( sig + MBEDTLS_LMS_SIG_PATH_OFFSET ) )[height];
+            right_node = Tc_candidate_root_node;
         }
         else
         {
             left_node = Tc_candidate_root_node;
-            rght_node = ( ( const unsigned char( * )[32] )( sig + MBEDTLS_LMS_SIG_PATH_OFFSET ) )[height];
+            right_node = ( ( const unsigned char( * )[MBEDTLS_LMS_M_NODE_BYTES] )( sig + MBEDTLS_LMS_SIG_PATH_OFFSET ) )[height];
         }
 
-        create_merkle_intr_node( ctx, left_node, rght_node, parent_node_id,
-                                 Tc_candidate_root_node);
+        create_merkle_internal_node(
+            ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(I_key_identifier),
+            left_node, right_node, parent_node_id, Tc_candidate_root_node);
 
         curr_node_id /= 2;
     }
@@ -526,115 +409,148 @@
     return( 0 );
 }
 
-int mbedtls_lms_import_pubkey( mbedtls_lms_context *ctx,
-                               const unsigned char *key )
+void mbedtls_lms_init_private( mbedtls_lms_private_t *ctx )
 {
-    mbedtls_lms_algorithm_type_t type;
-    mbedtls_lmots_algorithm_type_t otstype;
+    mbedtls_platform_zeroize( ctx, sizeof( mbedtls_lms_public_t ) ) ;
+}
 
-    if( ctx == NULL )
+void mbedtls_lms_free_private( mbedtls_lms_private_t *ctx )
+{
+    unsigned int idx;
+
+    if( ctx->MBEDTLS_PRIVATE(have_private_key) )
     {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
+        for( idx = 0; idx < MERKLE_TREE_LEAF_NODE_AM; idx++ )
+        {
+            mbedtls_lmots_free_private( &ctx->MBEDTLS_PRIVATE(ots_private_keys)[idx] );
+            mbedtls_lmots_free_public( &ctx->MBEDTLS_PRIVATE(ots_public_keys)[idx] );
+        }
+
+        mbedtls_free( ctx->MBEDTLS_PRIVATE(ots_private_keys) );
+        mbedtls_free( ctx->MBEDTLS_PRIVATE(ots_public_keys) );
     }
 
-    if( key == NULL )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
+    mbedtls_platform_zeroize( ctx, sizeof( mbedtls_lms_public_t ) );
+}
 
-    type = network_bytes_to_val( MBEDTLS_LMS_TYPE_LEN, key + MBEDTLS_LMS_PUBKEY_TYPE_OFFSET );
+
+int mbedtls_lms_generate_private_key( mbedtls_lms_private_t *ctx,
+                                      mbedtls_lms_algorithm_type_t type,
+                                      mbedtls_lmots_algorithm_type_t otstype,
+                                      int (*f_rng)(void *, unsigned char *, size_t),
+                                      void* p_rng, unsigned char *seed,
+                                      size_t seed_size )
+{
+    unsigned int idx = 0;
+    unsigned int free_idx = 0;
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+
     if( type != MBEDTLS_LMS_SHA256_M32_H10 )
     {
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
-    ctx->MBEDTLS_PRIVATE(type) = type;
 
-    otstype = network_bytes_to_val( MBEDTLS_LMOTS_TYPE_LEN,
-                                    key + MBEDTLS_LMS_PUBKEY_OTSTYPE_OFFSET );
     if( otstype != MBEDTLS_LMOTS_SHA256_N32_W8 )
     {
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
-    ctx->MBEDTLS_PRIVATE(otstype) = otstype;
 
-    memcpy( ctx->MBEDTLS_PRIVATE(I_key_identifier), key + MBEDTLS_LMS_PUBKEY_I_KEY_ID_OFFSET,
-            MBEDTLS_LMOTS_I_KEY_ID_LEN );
-    memcpy( ctx->MBEDTLS_PRIVATE(T_1_pub_key), key + MBEDTLS_LMS_PUBKEY_ROOT_NODE_OFFSET,
-            MBEDTLS_LMOTS_N_HASH_LEN );
+    if( ctx->MBEDTLS_PRIVATE(have_private_key) )
+    {
+        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
+    }
 
-    ctx->MBEDTLS_PRIVATE(have_pubkey) = 1;
+    ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(type) = type;
+    ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(otstype) = otstype;
+
+    f_rng( p_rng,
+           ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(I_key_identifier),
+           MBEDTLS_LMOTS_I_KEY_ID_LEN );
+
+    ctx->MBEDTLS_PRIVATE(ots_private_keys) = mbedtls_calloc( MERKLE_TREE_LEAF_NODE_AM,
+                                                             sizeof( mbedtls_lmots_private_t));
+    if( ctx->MBEDTLS_PRIVATE(ots_private_keys) == NULL )
+    {
+        ret = MBEDTLS_ERR_LMS_ALLOC_FAILED;
+        goto exit;
+    }
+
+    ctx->MBEDTLS_PRIVATE(ots_public_keys) = mbedtls_calloc( MERKLE_TREE_LEAF_NODE_AM,
+                                                            sizeof( mbedtls_lmots_public_t));
+    if( ctx->MBEDTLS_PRIVATE(ots_public_keys) == NULL )
+    {
+        ret = MBEDTLS_ERR_LMS_ALLOC_FAILED;
+        goto exit;
+    }
+
+    for( idx = 0; idx < MERKLE_TREE_LEAF_NODE_AM; idx++ )
+    {
+        mbedtls_lmots_init_private( &ctx->MBEDTLS_PRIVATE(ots_private_keys)[idx] );
+        mbedtls_lmots_init_public( &ctx->MBEDTLS_PRIVATE(ots_public_keys)[idx] );
+    }
+
+
+    for( idx = 0; idx < MERKLE_TREE_LEAF_NODE_AM; idx++ )
+    {
+        ret = mbedtls_lmots_generate_private_key( &ctx->MBEDTLS_PRIVATE(ots_private_keys)[idx],
+                                                  otstype,
+                                                  ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(I_key_identifier),
+                                                  idx, seed, seed_size );
+        if( ret)
+            goto exit;
+
+        ret = mbedtls_lmots_calculate_public_key( &ctx->MBEDTLS_PRIVATE(ots_public_keys)[idx],
+                                                  &ctx->MBEDTLS_PRIVATE(ots_private_keys)[idx] );
+        if( ret)
+            goto exit;
+    }
+
+    ctx->MBEDTLS_PRIVATE(q_next_usable_key) = 0;
+    ctx->MBEDTLS_PRIVATE(have_private_key) = 1;
+
+exit:
+    if( ret )
+    {
+        for ( free_idx = 0; free_idx < idx; free_idx++ ) {
+            mbedtls_lmots_free_private( &ctx->MBEDTLS_PRIVATE(ots_private_keys)[free_idx] );
+            mbedtls_lmots_free_public( &ctx->MBEDTLS_PRIVATE(ots_public_keys)[free_idx] );
+        }
+
+        mbedtls_free( ctx->MBEDTLS_PRIVATE(ots_private_keys) );
+        mbedtls_free( ctx->MBEDTLS_PRIVATE(ots_public_keys) );
+        return( ret );
+    }
 
     return( 0 );
 }
 
-int mbedtls_lms_export_pubkey( mbedtls_lms_context *ctx,
-                               unsigned char *key )
+int mbedtls_lms_calculate_public_key( mbedtls_lms_public_t *ctx,
+                                      mbedtls_lms_private_t *priv_ctx )
 {
-    if( ctx == NULL )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    if( key == NULL )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    if( ! ctx->MBEDTLS_PRIVATE(have_pubkey) )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    val_to_network_bytes( ctx->MBEDTLS_PRIVATE(type),
-                          MBEDTLS_LMS_TYPE_LEN, key + MBEDTLS_LMS_PUBKEY_TYPE_OFFSET );
-    val_to_network_bytes( ctx->MBEDTLS_PRIVATE(otstype),
-                          MBEDTLS_LMOTS_TYPE_LEN, key + MBEDTLS_LMS_PUBKEY_OTSTYPE_OFFSET );
-    memcpy( key + MBEDTLS_LMS_PUBKEY_I_KEY_ID_OFFSET,
-            ctx->MBEDTLS_PRIVATE(I_key_identifier),
-            MBEDTLS_LMOTS_I_KEY_ID_LEN );
-    memcpy( key + MBEDTLS_LMS_PUBKEY_ROOT_NODE_OFFSET,
-            ctx->MBEDTLS_PRIVATE(T_1_pub_key),
-            MBEDTLS_LMOTS_N_HASH_LEN );
-
-    return( 0 );
-}
-
-int mbedtls_lms_gen_pubkey( mbedtls_lms_context *ctx )
-{
-    unsigned char tree[MERKLE_TREE_NODE_AM][32];
-    unsigned int idx;
+    unsigned char tree[MERKLE_TREE_NODE_AM][MBEDTLS_LMS_M_NODE_BYTES];
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
-    if( ctx == NULL )
+    if( ! priv_ctx->MBEDTLS_PRIVATE( have_private_key ) )
     {
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
 
-    if( ! ctx->MBEDTLS_PRIVATE( have_privkey ) )
+    if( priv_ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(type)
+        != MBEDTLS_LMS_SHA256_M32_H10 )
     {
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
 
-    if( ctx->MBEDTLS_PRIVATE(type) != MBEDTLS_LMS_SHA256_M32_H10 )
+    if( priv_ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(otstype)
+        != MBEDTLS_LMOTS_SHA256_N32_W8 )
     {
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
 
-    if( ctx->MBEDTLS_PRIVATE(otstype) != MBEDTLS_LMOTS_SHA256_N32_W8 )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
+    memcpy( &ctx->MBEDTLS_PRIVATE(params), &priv_ctx->MBEDTLS_PRIVATE(params),
+            sizeof(mbedtls_lmots_parameters_t) );
 
-    for( idx = 0; idx < MERKLE_TREE_LEAF_AM; idx++ )
-    {
-        ret = mbedtls_lmots_gen_pubkey( &ctx->MBEDTLS_PRIVATE(priv_keys)[idx] );
-        if( ret )
-        {
-            return( ret );
-        }
-    }
-
-    ret = generate_merkle_tree( ctx, tree);
+    ret = calculate_merkle_tree( priv_ctx, tree);
     if( ret )
     {
         return( ret );
@@ -643,83 +559,112 @@
     /* Root node is always at position 1, due to 1-based indexing */
     memcpy( ctx->MBEDTLS_PRIVATE(T_1_pub_key), &tree[1], MBEDTLS_LMOTS_N_HASH_LEN );
 
-    ctx->MBEDTLS_PRIVATE(have_pubkey) = 1;
+    ctx->MBEDTLS_PRIVATE(have_public_key) = 1;
 
     return( 0 );
 }
 
-int mbedtls_lms_gen_privkey( mbedtls_lms_context *ctx,
-                             int ( *f_rng)(void *, unsigned char *, size_t),
-                             void* p_rng, unsigned char *seed,
-                             size_t seed_len )
+
+int mbedtls_lms_export_public_key( mbedtls_lms_public_t *ctx, unsigned char *key,
+                                   size_t key_size, size_t *key_len )
 {
-    unsigned int idx;
+    if( key_size < MBEDTLS_LMS_PUBLIC_KEY_LEN ) {
+        return( MBEDTLS_ERR_LMS_BUFFER_TOO_SMALL );
+    }
+
+    if( ! ctx->MBEDTLS_PRIVATE(have_public_key) )
+    {
+        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
+    }
+
+    unsigned_int_to_network_bytes(
+            ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(type),
+            MBEDTLS_LMS_TYPE_LEN, key + MBEDTLS_LMS_PUBLIC_KEY_TYPE_OFFSET );
+    unsigned_int_to_network_bytes(
+            ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(otstype),
+            MBEDTLS_LMOTS_TYPE_LEN, key + MBEDTLS_LMS_PUBLIC_KEY_OTSTYPE_OFFSET );
+    memcpy( key + MBEDTLS_LMS_PUBLIC_KEY_I_KEY_ID_OFFSET,
+            ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(I_key_identifier),
+            MBEDTLS_LMOTS_I_KEY_ID_LEN );
+    memcpy( key + MBEDTLS_LMS_PUBLIC_KEY_ROOT_NODE_OFFSET,
+            ctx->MBEDTLS_PRIVATE(T_1_pub_key),
+            MBEDTLS_LMOTS_N_HASH_LEN );
+
+    if( key_len != NULL ) {
+        *key_len = MBEDTLS_LMS_PUBLIC_KEY_LEN;
+    }
+
+    return( 0 );
+}
+
+
+int mbedtls_lms_sign( mbedtls_lms_private_t *ctx,
+                      int (*f_rng)(void *, unsigned char *, size_t),
+                      void* p_rng, unsigned char *msg, unsigned int msg_size,
+                      unsigned char *sig, size_t sig_size, size_t *sig_len)
+{
+    uint32_t q_leaf_identifier;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
-    if( ctx == NULL )
+    if( ! ctx->MBEDTLS_PRIVATE(have_private_key) )
     {
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
 
-    if( ctx->MBEDTLS_PRIVATE(type) != MBEDTLS_LMS_SHA256_M32_H10 )
+    if( sig_size < MBEDTLS_LMS_SIG_LEN )
+    {
+        return( MBEDTLS_ERR_LMS_BUFFER_TOO_SMALL );
+    }
+
+    if( ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(type) != MBEDTLS_LMS_SHA256_M32_H10 )
     {
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
 
-    if( ctx->MBEDTLS_PRIVATE(otstype) != MBEDTLS_LMOTS_SHA256_N32_W8 )
+    if( ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(otstype)
+        != MBEDTLS_LMOTS_SHA256_N32_W8 )
     {
         return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
     }
 
-    if( ctx->MBEDTLS_PRIVATE(have_privkey) )
+    if( ctx->MBEDTLS_PRIVATE(q_next_usable_key) >= MERKLE_TREE_LEAF_NODE_AM )
     {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
-    f_rng( p_rng, ctx->MBEDTLS_PRIVATE(I_key_identifier),
-           sizeof( ctx->MBEDTLS_PRIVATE(I_key_identifier) ) );
-
-    ctx->MBEDTLS_PRIVATE(priv_keys) = mbedtls_calloc( MERKLE_TREE_LEAF_AM,
-                                                      sizeof( mbedtls_lmots_context));
-    if( ctx->MBEDTLS_PRIVATE(priv_keys) == NULL )
-    {
-        ret = MBEDTLS_ERR_LMS_ALLOC_FAILED;
-        goto out;
-    }
-
-    for( idx = 0; idx < MERKLE_TREE_LEAF_AM; idx++ )
-    {
-        mbedtls_lmots_init( &ctx->MBEDTLS_PRIVATE(priv_keys)[idx] );
-        ret = mbedtls_lmots_set_algorithm_type( &ctx->MBEDTLS_PRIVATE(priv_keys)[idx],
-                                                ctx->MBEDTLS_PRIVATE(otstype) );
-        if( ret)
-        {
-            goto out;
-        }
+        return( MBEDTLS_ERR_LMS_OUT_OF_PRIVATE_KEYS );
     }
 
 
-    for( idx = 0; idx < MERKLE_TREE_LEAF_AM; idx++ )
-    {
-        ret = mbedtls_lmots_gen_privkey( &ctx->MBEDTLS_PRIVATE(priv_keys)[idx],
-                                         ctx->MBEDTLS_PRIVATE(I_key_identifier),
-                                         idx, seed, seed_len );
-        if( ret)
-        {
-            goto out;
-        }
-    }
+    q_leaf_identifier = ctx->MBEDTLS_PRIVATE(q_next_usable_key);
+    /* This new value must _always_ be written back to the disk before the
+     * signature is returned.
+     */
+    ctx->MBEDTLS_PRIVATE(q_next_usable_key) += 1;
 
-    ctx->MBEDTLS_PRIVATE(q_next_usable_key) = 0;
-    ctx->MBEDTLS_PRIVATE(have_privkey) = 1;
-
-out:
+    ret = mbedtls_lmots_sign( &ctx->MBEDTLS_PRIVATE(ots_private_keys)[q_leaf_identifier],
+                              f_rng, p_rng, msg, msg_size,
+                              sig + MBEDTLS_LMS_SIG_OTS_SIG_OFFSET,
+                              MBEDTLS_LMS_SIG_LEN, NULL );
     if( ret )
     {
-        mbedtls_free( ctx->MBEDTLS_PRIVATE(priv_keys) );
         return( ret );
     }
 
+    unsigned_int_to_network_bytes( ctx->MBEDTLS_PRIVATE(params).MBEDTLS_PRIVATE(type),
+                                   MBEDTLS_LMS_TYPE_LEN, sig + MBEDTLS_LMS_SIG_TYPE_OFFSET );
+    unsigned_int_to_network_bytes( q_leaf_identifier, MBEDTLS_LMOTS_Q_LEAF_ID_LEN,
+                                   sig + MBEDTLS_LMS_SIG_Q_LEAF_ID_OFFSET);
+
+    ret = get_merkle_path( ctx, MERKLE_TREE_INTERNAL_NODE_AM + q_leaf_identifier,
+                           ( unsigned char( * )[MBEDTLS_LMS_M_NODE_BYTES] )( sig + MBEDTLS_LMS_SIG_PATH_OFFSET ) );
+    if( ret )
+    {
+        return( ret );
+    }
+
+    if( sig_len != NULL ) {
+        *sig_len = MBEDTLS_LMS_SIG_LEN;
+    }
+
+
     return( 0 );
 }
 
diff --git a/tests/suites/test_suite_lmots.function b/tests/suites/test_suite_lmots.function
index 82dbcba..4492daa 100644
--- a/tests/suites/test_suite_lmots.function
+++ b/tests/suites/test_suite_lmots.function
@@ -15,7 +15,8 @@
 /* BEGIN_CASE */
 void lmots_sign_verify_test ( data_t * msg )
 {
-    mbedtls_lmots_context ctx;
+    mbedtls_lmots_public_t pub_ctx;
+    mbedtls_lmots_private_t priv_ctx;
     unsigned char sig[MBEDTLS_LMOTS_SIG_LEN];
     mbedtls_entropy_context entropy_ctx;
     mbedtls_ctr_drbg_context drbg_ctx;
@@ -23,22 +24,25 @@
 
     mbedtls_entropy_init( &entropy_ctx );
     mbedtls_ctr_drbg_init( &drbg_ctx );
-    mbedtls_lmots_init( &ctx );
+    mbedtls_lmots_init_public( &pub_ctx );
+    mbedtls_lmots_init_private( &priv_ctx );
 
     TEST_ASSERT( mbedtls_ctr_drbg_seed( &drbg_ctx, mbedtls_entropy_func,
                  &entropy_ctx, (uint8_t*)"", 0 ) == 0 );
     TEST_ASSERT( mbedtls_ctr_drbg_random( &drbg_ctx, seed, sizeof( seed ) ) == 0 );
 
-    TEST_ASSERT( mbedtls_lmots_set_algorithm_type(&ctx, MBEDTLS_LMOTS_SHA256_N32_W8) == 0 );
-    TEST_ASSERT( mbedtls_lmots_gen_privkey(&ctx, (uint8_t[16]){0}, 0x12, seed, sizeof( seed ) ) == 0 );
-    TEST_ASSERT( mbedtls_lmots_gen_pubkey(&ctx) == 0 );
-    TEST_ASSERT( mbedtls_lmots_sign(&ctx, mbedtls_ctr_drbg_random, &drbg_ctx, msg->x, msg->len, sig ) == 0 );
-    TEST_ASSERT( mbedtls_lmots_verify(&ctx, msg->x, msg->len, sig) == 0 );
+    TEST_ASSERT( mbedtls_lmots_generate_private_key(&priv_ctx, MBEDTLS_LMOTS_SHA256_N32_W8,
+                 (uint8_t[16]){0}, 0x12, seed, sizeof( seed ) ) == 0 );
+    TEST_ASSERT( mbedtls_lmots_calculate_public_key(&pub_ctx, &priv_ctx) == 0 );
+    TEST_ASSERT( mbedtls_lmots_sign(&priv_ctx, mbedtls_ctr_drbg_random, &drbg_ctx,
+                 msg->x, msg->len, sig, sizeof(sig), NULL ) == 0 );
+    TEST_ASSERT( mbedtls_lmots_verify(&pub_ctx, msg->x, msg->len, sig, sizeof(sig)) == 0 );
 
 exit:
     mbedtls_entropy_free( &entropy_ctx );
     mbedtls_ctr_drbg_free( &drbg_ctx );
-    mbedtls_lmots_free( &ctx );
+    mbedtls_lmots_free_public( &pub_ctx );
+    mbedtls_lmots_free_private( &priv_ctx );
 }
 /* END_CASE */
 
@@ -46,40 +50,40 @@
 void lmots_verify_test ( data_t * msg, data_t * sig, data_t * pub_key,
                           int expected_rc )
 {
-    mbedtls_lmots_context ctx;
+    mbedtls_lmots_public_t ctx;
 
-    mbedtls_lmots_init( &ctx );
+    mbedtls_lmots_init_public( &ctx );
 
-    mbedtls_lmots_import_pubkey( &ctx, pub_key->x );
+    mbedtls_lmots_import_public_key( &ctx, pub_key->x, pub_key->len );
 
-    TEST_ASSERT(mbedtls_lmots_verify( &ctx, msg->x, msg->len, sig->x ) == expected_rc );
+    TEST_ASSERT(mbedtls_lmots_verify( &ctx, msg->x, msg->len, sig->x, sig->len ) == expected_rc );
 
 exit:
-    mbedtls_lmots_free( &ctx );
+    mbedtls_lmots_free_public( &ctx );
 }
 /* END_CASE */
 
 /* BEGIN_CASE */
 void lmots_import_export_test (  data_t * pub_key )
 {
-    mbedtls_lmots_context ctx;
-    uint8_t exported_pub_key[MBEDTLS_LMOTS_PUBKEY_LEN];
+    mbedtls_lmots_public_t ctx;
+    uint8_t exported_pub_key[MBEDTLS_LMOTS_PUBLIC_KEY_LEN];
 
-    mbedtls_lmots_init( &ctx );
-    TEST_ASSERT( mbedtls_lmots_import_pubkey( &ctx, pub_key->x ) == 0 );
-    TEST_ASSERT( mbedtls_lmots_export_pubkey( &ctx, exported_pub_key ) == 0 );
+    mbedtls_lmots_init_public( &ctx );
+    TEST_ASSERT( mbedtls_lmots_import_public_key( &ctx, pub_key->x, pub_key->len ) == 0 );
+    TEST_ASSERT( mbedtls_lmots_export_public_key( &ctx, exported_pub_key, sizeof( exported_pub_key ), NULL ) == 0 );
 
-    TEST_ASSERT( memcmp( pub_key->x, exported_pub_key, MBEDTLS_LMOTS_PUBKEY_LEN ) == 0 );
+    TEST_ASSERT( memcmp( pub_key->x, exported_pub_key, MBEDTLS_LMOTS_PUBLIC_KEY_LEN ) == 0 );
 
 exit:
-    mbedtls_lmots_free( &ctx );
+    mbedtls_lmots_free_public( &ctx );
 }
 /* END_CASE */
 
 /* BEGIN_CASE */
 void lmots_reuse_test ( data_t * msg )
 {
-    mbedtls_lmots_context ctx;
+    mbedtls_lmots_private_t ctx;
     unsigned char sig[MBEDTLS_LMOTS_SIG_LEN];
     mbedtls_entropy_context entropy_ctx;
     mbedtls_ctr_drbg_context drbg_ctx;
@@ -92,19 +96,21 @@
 
     mbedtls_ctr_drbg_random( &drbg_ctx, seed, sizeof( seed ) );
 
-    mbedtls_lmots_init( &ctx );
-    TEST_ASSERT( mbedtls_lmots_set_algorithm_type( &ctx, MBEDTLS_LMOTS_SHA256_N32_W8 ) == 0 );
-    TEST_ASSERT( mbedtls_lmots_gen_privkey(&ctx, (uint8_t[16]){0}, 0x12, seed, sizeof( seed ) ) == 0 );
-    TEST_ASSERT( mbedtls_lmots_sign(&ctx, mbedtls_ctr_drbg_random, &drbg_ctx, msg->x, msg->len, sig ) == 0 );
+    mbedtls_lmots_init_private( &ctx );
+    TEST_ASSERT( mbedtls_lmots_generate_private_key(&ctx, MBEDTLS_LMOTS_SHA256_N32_W8,
+                                                    (uint8_t[16]){0}, 0x12, seed, sizeof( seed ) ) == 0 );
+    TEST_ASSERT( mbedtls_lmots_sign(&ctx, mbedtls_ctr_drbg_random, &drbg_ctx,
+                                    msg->x, msg->len, sig, sizeof( sig ), NULL ) == 0 );
 
     /* Running another sign operation should fail, since the key should now have
      * been erased.
      */
-    TEST_ASSERT( mbedtls_lmots_sign(&ctx, mbedtls_ctr_drbg_random, &drbg_ctx, msg->x, msg->len, sig ) != 0 );
+    TEST_ASSERT( mbedtls_lmots_sign(&ctx, mbedtls_ctr_drbg_random, &drbg_ctx,
+                                    msg->x, msg->len, sig, sizeof( sig ), NULL ) != 0 );
 
 exit:
     mbedtls_entropy_free( &entropy_ctx );
     mbedtls_ctr_drbg_free( &drbg_ctx );
-    mbedtls_lmots_free( &ctx );
+    mbedtls_lmots_free_private( &ctx );
 }
 /* END_CASE */
diff --git a/tests/suites/test_suite_lms.function b/tests/suites/test_suite_lms.function
index c6c7061..64ea900 100644
--- a/tests/suites/test_suite_lms.function
+++ b/tests/suites/test_suite_lms.function
@@ -13,7 +13,8 @@
 /* BEGIN_CASE */
 void lms_sign_verify_test ( data_t * msg )
 {
-    mbedtls_lms_context ctx;
+    mbedtls_lms_public_t pub_ctx;
+    mbedtls_lms_private_t priv_ctx;
     unsigned char sig[MBEDTLS_LMS_SIG_LEN];
     mbedtls_entropy_context entropy_ctx;
     mbedtls_ctr_drbg_context drbg_ctx;
@@ -22,29 +23,35 @@
 
     mbedtls_entropy_init( &entropy_ctx );
     mbedtls_ctr_drbg_init( &drbg_ctx );
-    mbedtls_lms_init( &ctx );
+    mbedtls_lms_init_public( &pub_ctx );
+    mbedtls_lms_init_private( &priv_ctx );
 
     TEST_ASSERT( mbedtls_ctr_drbg_seed( &drbg_ctx, mbedtls_entropy_func,
                  &entropy_ctx, ( uint8_t* )"", 0 ) == 0 );
     TEST_ASSERT( mbedtls_ctr_drbg_random( &drbg_ctx, seed, sizeof( seed ) ) == 0 );
 
-    TEST_ASSERT( mbedtls_lms_set_algorithm_type( &ctx, MBEDTLS_LMS_SHA256_M32_H10, MBEDTLS_LMOTS_SHA256_N32_W8 ) == 0 );
-
     /* Allocation failure isn't a test failure, since it likely just means there's not enough memory to run the test */
-    rc = mbedtls_lms_gen_privkey( &ctx, mbedtls_ctr_drbg_random, &drbg_ctx, seed, sizeof( seed ) );
+    rc = mbedtls_lms_generate_private_key( &priv_ctx, MBEDTLS_LMS_SHA256_M32_H10,
+                                           MBEDTLS_LMOTS_SHA256_N32_W8,
+                                           mbedtls_ctr_drbg_random, &drbg_ctx, seed,
+                                           sizeof( seed ) );
     TEST_ASSUME( rc != MBEDTLS_ERR_LMS_ALLOC_FAILED );
     TEST_ASSERT( rc == 0 );
 
-    TEST_ASSERT( mbedtls_lms_gen_pubkey( &ctx) == 0 );
+    TEST_ASSERT( mbedtls_lms_calculate_public_key( &pub_ctx, &priv_ctx ) == 0 );
 
-    TEST_ASSERT( mbedtls_lms_sign( &ctx, mbedtls_ctr_drbg_random, &drbg_ctx, msg->x, msg->len, sig ) == 0 );
+    TEST_ASSERT( mbedtls_lms_sign( &priv_ctx, mbedtls_ctr_drbg_random,
+                                   &drbg_ctx, msg->x, msg->len, sig,
+                                   sizeof( sig ), NULL ) == 0 );
 
-    TEST_ASSERT( mbedtls_lms_verify( &ctx, msg->x, msg->len, sig) == 0 );
+    TEST_ASSERT( mbedtls_lms_verify( &pub_ctx, msg->x, msg->len, sig,
+                                     sizeof( sig ) ) == 0 );
 
 exit:
     mbedtls_entropy_free( &entropy_ctx );
     mbedtls_ctr_drbg_free( &drbg_ctx );
-    mbedtls_lms_free( &ctx );
+    mbedtls_lms_free_public( &pub_ctx );
+    mbedtls_lms_free_private( &priv_ctx );
 }
 /* END_CASE */
 
@@ -52,34 +59,35 @@
 void lms_verify_test ( data_t * msg, data_t * sig, data_t * pub_key,
                           int expected_rc )
 {
-    mbedtls_lms_context ctx;
+    mbedtls_lms_public_t ctx;
 
-    mbedtls_lms_init( &ctx);
+    mbedtls_lms_init_public( &ctx);
 
-    mbedtls_lms_import_pubkey( &ctx, pub_key->x );
+    mbedtls_lms_import_public_key( &ctx, pub_key->x, pub_key->len );
 
-    TEST_ASSERT( mbedtls_lms_verify( &ctx, msg->x, msg->len, sig->x ) == expected_rc );
+    TEST_ASSERT( mbedtls_lms_verify( &ctx, msg->x, msg->len, sig->x, sig->len ) == expected_rc );
 
 exit:
-    mbedtls_lms_free( &ctx );
+    mbedtls_lms_free_public( &ctx );
 }
 /* END_CASE */
 
 /* BEGIN_CASE */
 void lms_import_export_test (  data_t * pub_key )
 {
-    mbedtls_lms_context ctx;
-    uint8_t exported_pub_key[MBEDTLS_LMS_PUBKEY_LEN];
+    mbedtls_lms_public_t ctx;
+    uint8_t exported_pub_key[MBEDTLS_LMS_PUBLIC_KEY_LEN];
 
-    mbedtls_lms_init(&ctx);
-    TEST_ASSERT( mbedtls_lms_import_pubkey( &ctx, pub_key->x ) == 0 );
-    TEST_ASSERT( mbedtls_lms_export_pubkey( &ctx, exported_pub_key) == 0 );
+    mbedtls_lms_init_public(&ctx);
+    TEST_ASSERT( mbedtls_lms_import_public_key( &ctx, pub_key->x, pub_key->len ) == 0 );
+    TEST_ASSERT( mbedtls_lms_export_public_key( &ctx, exported_pub_key,
+                                                sizeof(exported_pub_key), NULL ) == 0 );
 
-    ASSERT_COMPARE( pub_key->x, MBEDTLS_LMS_PUBKEY_LEN,
-                    exported_pub_key, MBEDTLS_LMS_PUBKEY_LEN );
+    ASSERT_COMPARE( pub_key->x, MBEDTLS_LMS_PUBLIC_KEY_LEN,
+                    exported_pub_key, MBEDTLS_LMS_PUBLIC_KEY_LEN );
 
 exit:
-    mbedtls_lms_free( &ctx );
+    mbedtls_lms_free_public( &ctx );
 }
 /* END_CASE */