diff --git a/Makefile b/Makefile
index 1fdb623..4eb165f 100644
--- a/Makefile
+++ b/Makefile
@@ -1,5 +1,5 @@
 CC = /usr/bin/gcc
-CFLAGS = -Wall -g -O3 
+CFLAGS = -Wall -g -O3 -Wextra
 
 all: test/test_chacha \
 test/test_wots \
diff --git a/chacha.c b/chacha.c
index 036fa15..d97c1fa 100644
--- a/chacha.c
+++ b/chacha.c
@@ -1,7 +1,7 @@
 /*
- * This code is based on an OpenSSL implementation of chacha20. 
+ * This code is based on an OpenSSL implementation of chacha20.
  * Hence, the copyright below applies.
- * 
+ *
  */
 /* ====================================================================
  * Copyright (c) 2011-2013 The OpenSSL Project.  All rights reserved.
@@ -84,13 +84,12 @@ static const char sigma[16] = "expand 32-byte k";
  * |input| and writes the 64 output bytes to |output|. */
 static void chacha_core(unsigned char output[64], const uint32_t input[16],
 			int num_rounds)
-	{
+{
 	uint32_t x[16];
 	int i;
 
 	memcpy(x, input, sizeof(uint32_t) * 16);
-	for (i = 20; i > 0; i -= 2)
-		{
+	for (i = num_rounds; i > 0; i -= 2) {
 		QUARTERROUND( 0, 4, 8,12)
 		QUARTERROUND( 1, 5, 9,13)
 		QUARTERROUND( 2, 6,10,14)
@@ -99,20 +98,20 @@ static void chacha_core(unsigned char output[64], const uint32_t input[16],
 		QUARTERROUND( 1, 6,11,12)
 		QUARTERROUND( 2, 7, 8,13)
 		QUARTERROUND( 3, 4, 9,14)
-		}
+	}
 
 	for (i = 0; i < 16; ++i)
 		x[i] = PLUS(x[i], input[i]);
 	for (i = 0; i < 16; ++i)
 		U32TO8_LITTLE(output + 4 * i, x[i]);
-	}
+}
 
 void CRYPTO_chacha_20(unsigned char *out,
 		      const unsigned char *in, size_t in_len,
 		      const unsigned char key[32],
 		      const unsigned char nonce[12],
 		      uint32_t counter)
-	{
+{
 	uint32_t input[16];
 	unsigned char buf[64];
 	size_t todo, i;
@@ -137,8 +136,7 @@ void CRYPTO_chacha_20(unsigned char *out,
 	input[14] = U8TO32_LITTLE(nonce + 4);
 	input[15] = U8TO32_LITTLE(nonce + 8);
 
-	while (in_len > 0)
-		{
+	while (in_len > 0) {
 		todo = sizeof(buf);
 		if (in_len < todo)
 			todo = in_len;
@@ -154,15 +152,15 @@ void CRYPTO_chacha_20(unsigned char *out,
 		input[12]++;
 		if (input[12] == 0)
 			input[13]++;
-		}
 	}
-	
+}
+
 void CRYPTO_chacha_20_keystream(unsigned char *out,
 		      size_t out_len,
 		      const unsigned char key[32],
 		      const unsigned char nonce[12],
 		      uint32_t counter)
-	{
+{
 	uint32_t input[16];
 	unsigned char buf[64];
 	size_t todo, i;
@@ -187,8 +185,7 @@ void CRYPTO_chacha_20_keystream(unsigned char *out,
 	input[14] = U8TO32_LITTLE(nonce + 4);
 	input[15] = U8TO32_LITTLE(nonce + 8);
 
-	while (out_len > 0)
-		{
+	while (out_len > 0) {
 		todo = sizeof(buf);
 		if (out_len < todo)
 			todo = out_len;
@@ -203,6 +200,5 @@ void CRYPTO_chacha_20_keystream(unsigned char *out,
 		input[12]++;
 		if (input[12] == 0)
 			input[13]++;
-		}
 	}
-  
\ No newline at end of file
+}
diff --git a/hash.c b/hash.c
index 4f4ff0b..1acc1d2 100644
--- a/hash.c
+++ b/hash.c
@@ -7,14 +7,15 @@ Public domain.
 #include "prg.h"
 
 #include <stddef.h>
-#include "stdio.h"
+#include <stdio.h>
+#include <string.h>
 #include <openssl/sha.h>
 #include <openssl/hmac.h>
 #include <openssl/evp.h>
 
 
-#define SET_KEY_BIT(a,b) (a[15] = (a[15] & 253) | (b << 1))
-#define SET_BLOCK_BIT(a,b) (a[15] = (a[15] & 254) | b)
+#define SET_KEY_BIT(a, b) (a[15] = (a[15] & 253) | ((b << 1) & 2))
+#define SET_BLOCK_BIT(a, b) (a[15] = (a[15] & 254) | (b & 1))
 
 #define WOTS_SELECT_KEY(a) (a[15] = (a[15] & 254) | 1)
 #define WOTS_SELECT_BLOCK(a) (a[15] = (a[15] & 254) | 0)
@@ -22,25 +23,21 @@ Public domain.
 /**
  * Implements PRF_m
  */
-int prf_m(unsigned char *out, const unsigned char *in, size_t inlen, const unsigned char *key, int keylen)
+int prf_m(unsigned char *out, const unsigned char *in, size_t inlen, const unsigned char *key, unsigned int keylen)
 {
   unsigned int length;
-  if (keylen == 32){
+  if (keylen == 32) {
     HMAC(EVP_sha256(), key, keylen, in, inlen, out, &length);
-    if(length != 32)
-    {
-      fprintf(stderr, "HMAC outputs %d bytes... That should not happen...",length); 
+    if (length != 32) {
+      fprintf(stderr, "HMAC outputs %d bytes... That should not happen...", length);
     }
     return 0;
   }
-  else
-  {
-    if(keylen == 64)
-    {
+  else {
+    if (keylen == 64) {
       HMAC(EVP_sha512(), key, keylen, in, inlen, out, &length);
-      if(length != 64)
-      {
-	fprintf(stderr, "HMAC outputs %d bytes... That should not happen...",length); 
+      if (length != 64) {
+        fprintf(stderr, "HMAC outputs %d bytes... That should not happen...", length);
       }
       return 0;
     }
@@ -51,37 +48,32 @@ int prf_m(unsigned char *out, const unsigned char *in, size_t inlen, const unsig
 /*
  * Implemts H_m
  */
-int hash_m(unsigned char *out,const unsigned char *in,unsigned long long inlen,const unsigned char *key, const int keylen, const int m)
+int hash_m(unsigned char *out, const unsigned char *in, unsigned long long inlen, const unsigned char *key, const unsigned int keylen, const unsigned int m)
 {
-  if(keylen != m){
-    fprintf(stderr, "H_m takes m-bit keys, we got m=%d but a keylength of %d.\n",m,keylen); 
+  unsigned int i;
+  unsigned char buf[inlen + keylen + m];
+
+  if (keylen != m){
+    fprintf(stderr, "H_m takes m-bit keys, we got m=%d but a keylength of %d.\n", m, keylen);
     return 1;
   }
-  unsigned long long i;
-  unsigned char buf[inlen +keylen+m];
-  for(i=0;i<m;i++)
-  {
+  for (i=0; i < m; i++) {
     buf[i] = 0x00;
   }
-  for(i=0;i <keylen;i++)
-  {
-    buf[m+i] = key[i];
+  for (i=0; i < keylen; i++) {
+    buf[m + i] = key[i];
   }
-  for(i=0;i <inlen;i++)
-  {
-    buf[m+keylen+i] = in[i];
+  for (i=0; i < inlen; i++) {
+    buf[m + keylen + i] = in[i];
   }
-  
-  if(m == 32)
-  {  
-    SHA256(buf,inlen +keylen+m,out);
+
+  if (m == 32) {
+    SHA256(buf, inlen + keylen + m, out);
     return 0;
-  } 
-  else
-  {
-    if(m == 64)
-    {
-      SHA512(buf,inlen +keylen+m,out);
+  }
+  else {
+    if (m == 64) {
+      SHA512(buf, inlen + keylen + m, out);
       return 0;
     }
   }
@@ -91,68 +83,70 @@ int hash_m(unsigned char *out,const unsigned char *in,unsigned long long inlen,c
 /**
  * We assume the left half is in in[0]...in[n-1]
  */
-int hash_2n_n(unsigned char *out,const unsigned char *in, const unsigned char *pub_seed, unsigned char addr[16], const int n)
+int hash_2n_n(unsigned char *out, const unsigned char *in, const unsigned char *pub_seed, unsigned char addr[16], const unsigned int n)
 {
-  
+
   unsigned char buf[4*n];
   unsigned char key[n];
   unsigned char bitmask[2*n];
-  int i;
-  
-  SET_KEY_BIT(addr,1);
-  SET_BLOCK_BIT(addr,0);
+  unsigned int i;
+
+  SET_KEY_BIT(addr, 1);
+  SET_BLOCK_BIT(addr, 0);
   prg_with_counter(key, pub_seed, n, addr);
-  SET_KEY_BIT(addr,0);
+  SET_KEY_BIT(addr, 0);
   // Use MSB order
   prg_with_counter(bitmask, pub_seed, n, addr);
-  SET_BLOCK_BIT(addr,1);
+  SET_BLOCK_BIT(addr, 1);
   prg_with_counter(bitmask+n, pub_seed, n, addr);
-  for(i=0;i<n;i++)
-  {
+  for (i = 0; i < n; i++) {
     buf[i] = 0x00;
     buf[n+i] = key[i];
     buf[2*n+i] = in[i] ^ bitmask[i];
     buf[3*n+i] = in[n+i] ^ bitmask[n+i];
   }
-  if(n==32){
-    SHA256(buf,4*n,out);
+  if (n == 32) {
+    SHA256(buf, 4*n, out);
     return 0;
-  } else {
-    if(n==64){
-      SHA512(buf,4*n,out);
+  }
+  else {
+    if (n == 64) {
+      SHA512(buf, 4*n, out);
       return 0;
-    } else {
+    }
+    else {
       fprintf(stderr, "Hash.c:hash_2n_n: Code only supports n=32 or n=64");
       return -1;
     }
   }
 }
 
-int hash_n_n(unsigned char *out,const unsigned char *in, const unsigned char *pub_seed, unsigned char addr[16], const int n)
-{ 
+int hash_n_n(unsigned char *out, const unsigned char *in, const unsigned char *pub_seed, unsigned char addr[16], const unsigned int n)
+{
   unsigned char buf[3*n];
   unsigned char key[n];
   unsigned char bitmask[n];
-  int i;
-  
+  unsigned int i;
+
   WOTS_SELECT_KEY(addr);
   prg_with_counter(key, pub_seed, n, addr);
   WOTS_SELECT_BLOCK(addr);
   prg_with_counter(bitmask, pub_seed, n, addr);
-  for(i=0;i<n;i++)
-  {
+  for (i = 0; i < n; i++) {
     buf[i] = 0x00;
     buf[n+i] = key[i];
     buf[2*n+i] = in[i] ^ bitmask[i];
   }
-  if(n==32){
-    SHA256(buf,3*n,out);
+  if (n == 32) {
+    SHA256(buf, 3*n, out);
     return 0;
-  } else {
-    if(n==64){
-      SHA512(buf,3*n,out);
+  }
+  else {
+    if (n == 64) {
+      SHA512(buf, 3*n, out);
       return 0;
-    } else {
+    }
+    else {
       fprintf(stderr, "Hash.c:hash_n_n: Code only supports n=32 or n=64");
       return -1;
     }
diff --git a/prg.c b/prg.c
index 4688e5c..5a14ec6 100644
--- a/prg.c
+++ b/prg.c
@@ -16,33 +16,28 @@ const unsigned char zero_nonce[12] = {0};
  * Generates rlen output bytes using ChaCha20 with a zero nonce and counter = 0
  */
 void prg(unsigned char *r, unsigned long long rlen, const unsigned char *key, unsigned int key_len)
-{  
-  if(key_len == 32){
+{
+  if (key_len == 32) {
     CRYPTO_chacha_20_keystream(r, rlen, key, zero_nonce, 0);
   }
-  else
-  {
-    if(key_len == 64)
-    {
+  else {
+    if (key_len == 64) {
       unsigned long long left = rlen;
       u_int32_t counter = 0;
       unsigned char *c = (unsigned char*)&counter;
       unsigned int length;
-      unsigned int i = 0; 
+      unsigned int i = 0;
       unsigned char tmp[64];
-      while(left > 0)
-      {
-	HMAC(EVP_sha512(), key, key_len, c , 4, tmp, &length);
-	if(length != 64)
-	  {
-	    fprintf(stderr, "HMAC outputs %d bytes... That should not happen...",length); 
-	  }
-	for(i = 0; ((i < length) && (i < left));i++)
-	{
-	  r[rlen-left+i] = tmp[i];
-	}
-	left -=length;
-	counter++;
+      while (left > 0) {
+        HMAC(EVP_sha512(), key, key_len, c , 4, tmp, &length);
+        if (length != 64) {
+          fprintf(stderr, "HMAC outputs %d bytes... That should not happen...", length);
+        }
+        for (i = 0; ((i < length) && (i < left)); i++) {
+          r[rlen-left+i] = tmp[i];
+        }
+        left -= length;
+        counter++;
       }
     }
     else {
@@ -53,34 +48,31 @@ void prg(unsigned char *r, unsigned long long rlen, const unsigned char *key, un
 
 /**
  * Generates n output bytes using ChaCha20 (n=32) or HMAC-SHA2-512 (n=64).
- * 
+ *
  * For ChaCha, nonce and counter are set depending on the address addr. For HMAC, addr is used as message.
  */
 void prg_with_counter(unsigned char *r, const unsigned char *key, unsigned int n, const unsigned char addr[16])
 {
   int i;
   unsigned char nonce[12];
-  if(n == 32){
-    for(i = 0; i < 12; i++)
-    {
+  if (n == 32) {
+    for (i = 0; i < 12; i++) {
       nonce[i] = addr[i];
     }
     uint32_t counter;
-    counter = (((uint32_t)addr[12]) << 24)|(((uint32_t)addr[13]) << 16)|(((uint32_t)addr[14]) << 8)|addr[15];
+    counter = (((uint32_t)addr[12]) << 24) | (((uint32_t)addr[13]) << 16) | (((uint32_t)addr[14]) << 8) | addr[15];
     // TODO: Check address handling. Endianess?
     CRYPTO_chacha_20_keystream(r, n, key, nonce, counter);
-  } 
-  else 
-  {
-    if(n == 64)
-    {
+  }
+  else {
+    if (n == 64) {
       unsigned int length;
       HMAC(EVP_sha512(), key, n, addr, 16, r, &length);
-      if(length != 64)
-      {
-	fprintf(stderr, "HMAC outputs %d bytes... That should not happen...",length); 
+      if (length != 64) {
+        fprintf(stderr, "HMAC outputs %d bytes... That should not happen...", length);
       }
-    } else {
+    }
+    else {
       fprintf(stderr,"prg.c:: Code only supports 32 byte and 64 byte seeds");
     }
   }
diff --git a/randombytes.c b/randombytes.c
index 310d1ce..effebc1 100644
--- a/randombytes.c
+++ b/randombytes.c
@@ -11,13 +11,13 @@ This code was taken from the SPHINCS reference implementation and is public doma
 
 static int fd = -1;
 
-void randombytes(unsigned char *x,unsigned long long xlen)
+void randombytes(unsigned char *x, unsigned long long xlen)
 {
   int i;
 
   if (fd == -1) {
     for (;;) {
-      fd = open("/dev/urandom",O_RDONLY);
+      fd = open("/dev/urandom", O_RDONLY);
       if (fd != -1) break;
       sleep(1);
     }
@@ -26,7 +26,7 @@ void randombytes(unsigned char *x,unsigned long long xlen)
   while (xlen > 0) {
     if (xlen < 1048576) i = xlen; else i = 1048576;
 
-    i = read(fd,x,i);
+    i = read(fd, x, i);
     if (i < 1) {
       sleep(1);
       continue;
diff --git a/test/test_chacha.c b/test/test_chacha.c
index 5392fdd..6ec2ddb 100644
--- a/test/test_chacha.c
+++ b/test/test_chacha.c
@@ -14,7 +14,7 @@ int main()
   int n = 32;
   unsigned char seed[32] = {0};
 //   unsigned char seed[64] = {0,0};
-  
+
   unsigned char out[2*n];
   unsigned char addr[16] = {2};
 
@@ -24,7 +24,7 @@ int main()
   printf("\n");
   hexdump(out, 2*n);
   printf("\n");
-  
+
   printf("Case 2: key = 1\n");
   seed[31] = 1;
   prg_with_counter(out, seed, n, addr);
diff --git a/test/test_wots.c b/test/test_wots.c
index 794db0c..146db3d 100644
--- a/test/test_wots.c
+++ b/test/test_wots.c
@@ -36,9 +36,8 @@ int main()
   wots_sign(sig, msg, seed, &params, pub_seed, addr);
   wots_pkFromSig(pk2, sig, msg, &params, pub_seed, addr);
 
-  for(i=0;i<sig_len;i++)
-    if(pk1[i] != pk2[i])
-    {
+  for (i = 0; i < sig_len; i++)
+    if (pk1[i] != pk2[i]) {
       printf("pk1 != pk2 %d\n",i);
       return -1;
     }
@@ -49,6 +48,6 @@ int main()
   printf("\nsig: ");
   hexdump(sig, sig_len);
   printf("\n");
-  
+
   return 0;
 }
diff --git a/test/test_xmss.c b/test/test_xmss.c
index 1407b6a..3895867 100644
--- a/test/test_xmss.c
+++ b/test/test_xmss.c
@@ -6,8 +6,6 @@
 #define MLEN 3491
 #define SIGNATURES 50
 
-
-
 unsigned char mi[MLEN];
 unsigned long long smlen;
 unsigned long long mlen;
@@ -16,13 +14,13 @@ int main()
 {
   int r;
   unsigned long long i;
-  int m = 64;
-  int n = 64;
-  int h = 8;
-  int w = 16;
-  
+  unsigned int m = 64;
+  unsigned int n = 64;
+  unsigned int h = 8;
+  unsigned int w = 16;
+
   unsigned long errors = 0;
-  
+
   unsigned char sk[3*n+4];
   unsigned char pk[2*n];
 
@@ -34,21 +32,20 @@ int main()
   unsigned char sm[MLEN+signature_length];
 
   FILE *urandom = fopen("/dev/urandom", "r");
-  for(i=0;i<MLEN;i++) mi[i] = fgetc(urandom);
+  for (i = 0; i < MLEN; i++) mi[i] = fgetc(urandom);
 
   printf("keypair\n");
   xmss_keypair(pk, sk, params);
   // check pub_seed in SK
-  for(i=0;i<n;i++)
-  {
-    if(pk[n+i] != sk[4+m+n+i]) printf("pk.pub_seed != sk.pub_seed %llu",i);
+  for (i = 0; i < n; i++) {
+    if (pk[n+i] != sk[4+m+n+i]) printf("pk.pub_seed != sk.pub_seed %llu",i);
   }
 
   // check index
   unsigned long idx = ((unsigned long)sk[0] << 24) | ((unsigned long)sk[1] << 16) | ((unsigned long)sk[2] << 8) | sk[3];
-  if(idx) printf("\nidx != 0 %lu\n",idx);
-  
-  for(i=0;i<(1<<h);i++){
+  if (idx) printf("\nidx != 0 %lu\n",idx);
+
+  for (i = 0; i < SIGNATURES; i++) {
     printf("sign\n");
     xmss_sign(sk, sm, &smlen, mi, MLEN, params);
     idx = ((unsigned long)sm[0] << 24) | ((unsigned long)sm[1] << 16) | ((unsigned long)sm[2] << 8) | sm[3];
@@ -57,11 +54,11 @@ int main()
     r = memcmp(mi, sm+signature_length,MLEN);
     printf("%d\n", r);
 
-        /* Test valid signature */
+    /* Test valid signature */
     printf("verify\n");
     r = xmss_sign_open(mo, &mlen, sm, smlen, pk, params);
     printf("%d\n", r);
-    if(r != 0) errors++;
+    if (r != 0) errors++;
     r = memcmp(mi,mo,MLEN);
     printf("%d\n", r);
     printf("%llu\n", MLEN-mlen);
@@ -70,7 +67,7 @@ int main()
     sm[signature_length+10] ^= 1;
     r = xmss_sign_open(mo, &mlen, sm, smlen, pk, params);
     printf("%d\n", r+1);
-    if(r == 0) errors++;
+    if (r == 0) errors++;
     r = memcmp(mi,mo,MLEN);
     printf("%d\n", (r!=0) - 1);
     printf("%llu\n", mlen+1);
@@ -81,37 +78,37 @@ int main()
     sm[2] ^= 1;
     r = xmss_sign_open(mo, &mlen, sm, smlen, pk, params);
     printf("%d\n", r+1);
-    if(r == 0) errors++;
+    if (r == 0) errors++;
     r = memcmp(mi,mo,MLEN);
     printf("%d\n", (r!=0) - 1);
     printf("%llu\n", mlen+1);
-    
+
     /* Modified R */
     sm[2] ^= 1;
     sm[5] ^= 1;
     r = xmss_sign_open(mo, &mlen, sm, smlen, pk, params);
     printf("%d\n", r+1);
-    if(r == 0) errors++;
+    if (r == 0) errors++;
     r = memcmp(mi,mo,MLEN);
     printf("%d\n", (r!=0) - 1);
     printf("%llu\n", mlen+1);
-    
+
     /* Modified OTS sig */
     sm[5] ^= 1;
     sm[240] ^= 1;
     r = xmss_sign_open(mo, &mlen, sm, smlen, pk, params);
     printf("%d\n", r+1);
-    if(r == 0) errors++;
+    if (r == 0) errors++;
     r = memcmp(mi,mo,MLEN);
     printf("%d\n", (r!=0) - 1);
     printf("%llu\n", mlen+1);
-    
+
     /* Modified AUTH */
     sm[240] ^= 1;
     sm[signature_length - 10] ^= 1;
     r = xmss_sign_open(mo, &mlen, sm, smlen, pk, params);
     printf("%d\n", r+1);
-    if(r == 0) errors++;
+    if (r == 0) errors++;
     r = memcmp(mi,mo,MLEN);
     printf("%d\n", (r!=0) - 1);
     printf("%llu\n", mlen+1);
diff --git a/test/test_xmss_fast.c b/test/test_xmss_fast.c
index 34b5977..bca8e10 100644
--- a/test/test_xmss_fast.c
+++ b/test/test_xmss_fast.c
@@ -7,8 +7,6 @@
 #define MLEN 3491
 #define SIGNATURES 256
 
-
-
 unsigned char mi[MLEN];
 unsigned long long smlen;
 unsigned long long mlen;
@@ -17,14 +15,14 @@ int main()
 {
   int r;
   unsigned long long i;
-  int m = 32;
-  int n = 32;
-  int h = 8;
-  int w = 16;
-  int k = 2;
-  
+  unsigned int m = 32;
+  unsigned int n = 32;
+  unsigned int h = 8;
+  unsigned int w = 16;
+  unsigned int k = 2;
+
   unsigned long errors = 0;
-  
+
   unsigned char sk[3*n+4];
   unsigned char pk[2*n];
 
@@ -43,7 +41,7 @@ int main()
   unsigned char retain[((1 << k) - k - 1)*n];
   bds_state s;
   bds_state *state = &s;
-  for(i=0;i<h-k;i++)
+  for (i = 0; i < h-k; i++)
     treehash[i].node = &th_nodes[n*i];
   xmss_set_bds_state(state, stack, stackoffset, stacklevels, auth, keep, treehash, retain, 0);
 
@@ -52,21 +50,20 @@ int main()
   unsigned char sm[MLEN+signature_length];
 
   FILE *urandom = fopen("/dev/urandom", "r");
-  for(i=0;i<MLEN;i++) mi[i] = fgetc(urandom);
+  for (i = 0; i < MLEN; i++) mi[i] = fgetc(urandom);
 
   printf("keypair\n");
   xmss_keypair(pk, sk, state, params);
   // check pub_seed in SK
-  for(i=0;i<n;i++)
-  {
-    if(pk[n+i] != sk[4+m+n+i]) printf("pk.pub_seed != sk.pub_seed %llu",i);
+  for (i = 0; i < n; i++) {
+    if (pk[n+i] != sk[4+m+n+i]) printf("pk.pub_seed != sk.pub_seed %llu",i);
   }
 
   // check index
   unsigned long idx = ((unsigned long)sk[0] << 24) | ((unsigned long)sk[1] << 16) | ((unsigned long)sk[2] << 8) | sk[3];
-  if(idx) printf("\nidx != 0 %lu\n",idx);
-  
-  for(i=0;i<((1<<h));i++){
+  if (idx) printf("\nidx != 0 %lu\n",idx);
+
+  for (i = 0; i < SIGNATURES; i++) {
     printf("sign\n");
     xmss_sign(sk, state, sm, &smlen, mi, MLEN, params);
     idx = ((unsigned long)sm[0] << 24) | ((unsigned long)sm[1] << 16) | ((unsigned long)sm[2] << 8) | sm[3];
@@ -79,7 +76,7 @@ int main()
     printf("verify\n");
     r = xmss_sign_open(mo, &mlen, sm, smlen, pk, params);
     printf("%d\n", r);
-    if(r != 0) errors++;
+    if (r != 0) errors++;
     r = memcmp(mi,mo,MLEN);
     printf("%d\n", r);
     printf("%llu\n", MLEN-mlen);
@@ -88,7 +85,7 @@ int main()
     sm[signature_length+10] ^= 1;
     r = xmss_sign_open(mo, &mlen, sm, smlen, pk, params);
     printf("%d\n", r+1);
-    if(r == 0) errors++;
+    if (r == 0) errors++;
     r = memcmp(mi,mo,MLEN);
     printf("%d\n", (r!=0) - 1);
     printf("%llu\n", mlen+1);
@@ -99,37 +96,37 @@ int main()
     sm[2] ^= 1;
     r = xmss_sign_open(mo, &mlen, sm, smlen, pk, params);
     printf("%d\n", r+1);
-    if(r == 0) errors++;
+    if (r == 0) errors++;
     r = memcmp(mi,mo,MLEN);
     printf("%d\n", (r!=0) - 1);
     printf("%llu\n", mlen+1);
-    
+
     /* Modified R */
     sm[2] ^= 1;
     sm[5] ^= 1;
     r = xmss_sign_open(mo, &mlen, sm, smlen, pk, params);
     printf("%d\n", r+1);
-    if(r == 0) errors++;
+    if (r == 0) errors++;
     r = memcmp(mi,mo,MLEN);
     printf("%d\n", (r!=0) - 1);
     printf("%llu\n", mlen+1);
-    
+
     /* Modified OTS sig */
     sm[5] ^= 1;
     sm[240] ^= 1;
     r = xmss_sign_open(mo, &mlen, sm, smlen, pk, params);
     printf("%d\n", r+1);
-    if(r == 0) errors++;
+    if (r == 0) errors++;
     r = memcmp(mi,mo,MLEN);
     printf("%d\n", (r!=0) - 1);
     printf("%llu\n", mlen+1);
-    
+
     /* Modified AUTH */
     sm[240] ^= 1;
     sm[signature_length - 10] ^= 1;
     r = xmss_sign_open(mo, &mlen, sm, smlen, pk, params);
     printf("%d\n", r+1);
-    if(r == 0) errors++;
+    if (r == 0) errors++;
     r = memcmp(mi,mo,MLEN);
     printf("%d\n", (r!=0) - 1);
     printf("%llu\n", mlen+1);
@@ -141,5 +138,3 @@ int main()
   printf("closed urandom\n");
   return 0;
 }
-
-
diff --git a/test/test_xmssmt.c b/test/test_xmssmt.c
index ad2a9a4..f11b275 100644
--- a/test/test_xmssmt.c
+++ b/test/test_xmssmt.c
@@ -6,7 +6,6 @@
 #define MLEN 3491
 #define SIGNATURES 1024
 
-
 unsigned char mi[MLEN];
 unsigned long long smlen;
 unsigned long long mlen;
@@ -15,55 +14,54 @@ int main()
 {
   int r;
   unsigned long long i,j;
-  int m = 32;
-  int n = 32;
-  int h = 20;
-  int d = 5;
-  int w = 16;
-  
+  unsigned int m = 32;
+  unsigned int n = 32;
+  unsigned int h = 20;
+  unsigned int d = 5;
+  unsigned int w = 16;
+
   xmssmt_params p;
   xmssmt_params *params = &p;
   xmssmt_set_params(params, m, n, h, d, w);
-  
+
   unsigned char sk[(params->index_len+2*n+m)];
   unsigned char pk[2*n];
-  
+
   unsigned long long signature_length = params->index_len + m + (d*params->xmss_par.wots_par.keysize) + h*n;
   unsigned char mo[MLEN+signature_length];
   unsigned char sm[MLEN+signature_length];
 
   FILE *urandom = fopen("/dev/urandom", "r");
-  for(i=0;i<MLEN;i++) mi[i] = fgetc(urandom);
+  for (i = 0; i < MLEN; i++) mi[i] = fgetc(urandom);
 
   printf("keypair\n");
   xmssmt_keypair(pk, sk, params);
   // check pub_seed in SK
-  for(i=0;i<n;i++)
-  {
-    if(pk[n+i] != sk[params->index_len+m+n+i]) printf("pk.pub_seed != sk.pub_seed %llu",i);
+  for (i = 0; i < n; i++) {
+    if (pk[n+i] != sk[params->index_len+m+n+i]) printf("pk.pub_seed != sk.pub_seed %llu",i);
   }
   printf("pk checked\n");
 
   unsigned int idx_len = params->index_len;
   // check index
   unsigned long long idx = 0;
-  for(i = 0; i < idx_len; i++){
+  for (i = 0; i < idx_len; i++) {
     idx |= ((unsigned long long)sk[i]) << 8*(idx_len - 1 - i);
   }
 
-  if(idx) printf("\nidx != 0: %llu\n",idx);
-  
-  for(i=0;i<(1<<h);i++){
+  if (idx) printf("\nidx != 0: %llu\n",idx);
+
+  for (i = 0; i < SIGNATURES; i++) {
     printf("sign\n");
     xmssmt_sign(sk, sm, &smlen, mi, MLEN, params);
     idx = 0;
-    for(j = 0; j < idx_len; j++){
+    for (j = 0; j < idx_len; j++) {
       idx += ((unsigned long long)sm[j]) << 8*(idx_len - 1 - j);
     }
     printf("\nidx = %llu\n",idx);
     r = memcmp(mi, sm+signature_length,MLEN);
     printf("%d\n", r);
-    
+
     /* Test valid signature */
     printf("verify\n");
     r = xmssmt_sign_open(mo, &mlen, sm, smlen, pk, params);
@@ -89,7 +87,6 @@ int main()
     r = memcmp(mi,mo,MLEN);
     printf("%d\n", (r!=0) - 1);
     printf("%llu\n", mlen+1);
-    
   }
   fclose(urandom);
   return 0;
diff --git a/test/test_xmssmt_fast.c b/test/test_xmssmt_fast.c
index 3c12a85..ce3cb1c 100644
--- a/test/test_xmssmt_fast.c
+++ b/test/test_xmssmt_fast.c
@@ -6,7 +6,6 @@
 #define MLEN 3491
 #define SIGNATURES 4096
 
-
 unsigned char mi[MLEN];
 unsigned long long smlen;
 unsigned long long mlen;
@@ -15,12 +14,12 @@ int main()
 {
   int r;
   unsigned long long i,j;
-  int m = 32;
-  int n = 32;
-  int h = 12;
-  int d = 2;
-  int w = 16;
-  int k = 2;
+  unsigned int m = 32;
+  unsigned int n = 32;
+  unsigned int h = 12;
+  unsigned int d = 3;
+  unsigned int w = 16;
+  unsigned int k = 2;
 
   xmssmt_params p;
   xmssmt_params *params = &p;
@@ -31,18 +30,19 @@ int main()
   unsigned int tree_h = h / d;
 
   // stack needs to be larger than regular (H-K-1), since we re-use for 'next'
-  unsigned char stack[2*d * (tree_h + 1)*n];
-  unsigned char stacklevels[2*d * (tree_h + 1)*n];
-  unsigned char auth[2*d * tree_h*n];
-  unsigned char keep[2*d * (tree_h >> 1)*n];
-  treehash_inst treehash[2*d * (tree_h-k)];
-  unsigned char th_nodes[2*d * (tree_h-k)*n];
-  unsigned char retain[2*d * ((1 << k) - k - 1)*n];
+  unsigned char stack[(2*d-1) * (tree_h + 1)*n];
+  unsigned char stacklevels[(2*d-1) * (tree_h + 1)*n];
+  unsigned char auth[(2*d-1) * tree_h*n];
+  unsigned char keep[(2*d-1) * (tree_h >> 1)*n];
+  treehash_inst treehash[(2*d-1) * (tree_h-k)];
+  unsigned char th_nodes[(2*d-1) * (tree_h-k)*n];
+  unsigned char retain[(2*d-1) * ((1 << k) - k - 1)*n];
   unsigned char wots_sigs[d * params->xmss_par.wots_par.keysize];
-  bds_state states[2*d]; // first d are 'regular' states, second d are 'next'
+  // first d are 'regular' states, second d are 'next'; top tree has no 'next'
+  bds_state states[2*d-1];
 
-  for (i = 0; i < 2*d; i++) {
-    for(j=0;j<tree_h-k;j++)
+  for (i = 0; i < 2*d-1; i++) {
+    for (j = 0; j < tree_h-k; j++)
       treehash[i*(tree_h-k) + j].node = th_nodes + (i*(tree_h-k) + j) * n;
     xmss_set_bds_state(states + i,
       stack + i*(tree_h + 1)*n, 0, stacklevels + i*(tree_h + 1),
@@ -62,32 +62,31 @@ int main()
   unsigned char sm[MLEN+signature_length];
 
   FILE *urandom = fopen("/dev/urandom", "r");
-  for(i=0;i<MLEN;i++) mi[i] = fgetc(urandom);
+  for (i = 0; i < MLEN; i++) mi[i] = fgetc(urandom);
 
   printf("keypair\n");
   xmssmt_keypair(pk, sk, states, wots_sigs, params);
   // check pub_seed in SK
-  for(i=0;i<n;i++)
-  {
-    if(pk[n+i] != sk[params->index_len+m+n+i]) printf("pk.pub_seed != sk.pub_seed %llu",i);
+  for (i = 0; i < n; i++) {
+    if (pk[n+i] != sk[params->index_len+m+n+i]) printf("pk.pub_seed != sk.pub_seed %llu",i);
   }
   printf("pk checked\n");
 
   unsigned int idx_len = params->index_len;
   // check index
   unsigned long long idx = 0;
-  for(i = 0; i < idx_len; i++){
+  for (i = 0; i < idx_len; i++) {
     idx |= ((unsigned long long)sk[i]) << 8*(idx_len - 1 - i);
   }
 
-  if(idx) printf("\nidx != 0: %llu\n",idx);
+  if (idx) printf("\nidx != 0: %llu\n",idx);
 
-  for(i=0;i<SIGNATURES;i++){
+  for (i = 0; i < SIGNATURES; i++) {
     printf("sign\n");
     xmssmt_sign(sk, states, wots_sigs, sm, &smlen, mi, MLEN, params);
 
     idx = 0;
-    for(j = 0; j < idx_len; j++){
+    for (j = 0; j < idx_len; j++) {
       idx += ((unsigned long long)sm[j]) << 8*(idx_len - 1 - j);
     }
     printf("\nidx = %llu\n",idx);
diff --git a/wots.c b/wots.c
index 1280c1b..80ab3f8 100644
--- a/wots.c
+++ b/wots.c
@@ -17,11 +17,11 @@ Public domain.
  * in the 16byte hash address
  */
 #define SET_HASH_ADDRESS(a, v) {\
-  a[15] = (a[15] & 1) | ((v << 1) & 255);\
+  a[15] = (a[15] & 1) | ((v << 1) & 254);\
   a[14] = (a[14] & 254) | ((v >> 7) & 1);}
-  
+
 #define SET_CHAIN_ADDRESS(a, v) {\
-  a[14] = (a[14] & 1) | ((v << 1) & 255);\
+  a[14] = (a[14] & 1) | ((v << 1) & 254);\
   a[13] = (v >> 7) & 255;\
   a[12] = (a[12] & 254) | ((v >> 15) & 1);}
 
@@ -40,7 +40,7 @@ void wots_set_params(wots_params *params, int m, int n, int w)
 
 /**
  * Helper method for pseudorandom key generation
- * Expands a 32 byte array into a len*n byte array 
+ * Expands an n-byte array into a len*n byte array
  * this is done using chacha20 with nonce 0 and counter 0
  */
 static void expand_seed(unsigned char *outseeds, const unsigned char *inseed, const wots_params *params)
@@ -51,29 +51,26 @@ static void expand_seed(unsigned char *outseeds, const unsigned char *inseed, co
 /**
  * Computes the chaining function.
  * out and in have to be n-byte arrays
- * 
+ *
  * interpretes in as start-th value of the chain
  * addr has to contain the address of the chain
  */
-static void gen_chain(unsigned char *out, const unsigned char *in, int start, int steps, const wots_params *params, const unsigned char *pub_seed, unsigned char addr[16])
+static void gen_chain(unsigned char *out, const unsigned char *in, unsigned int start, unsigned int steps, const wots_params *params, const unsigned char *pub_seed, unsigned char addr[16])
 {
-  unsigned int i,j;
-  for(j=0;j<params->n;j++) 
+  unsigned int i, j;
+  for (j = 0; j < params->n; j++)
     out[j] = in[j];
 
-  for(i=start;i<(start+steps) && i<params->w;i++){
-    SET_HASH_ADDRESS(addr,i);
-//     printf("Hash %d:",i);
-//     hexdump(addr,16);
-//     printf("\n");
-     hash_n_n(out,out, pub_seed, addr,params->n);
+  for (i = start; i < (start+steps) && i < params->w; i++) {
+    SET_HASH_ADDRESS(addr, i);
+    hash_n_n(out, out, pub_seed, addr, params->n);
   }
 }
 
 /**
  * base_w algorithm as described in draft.
- * 
- * 
+ *
+ *
  */
 static void base_w(int *output, const unsigned char *input, int in_len, const wots_params *params)
 {
@@ -82,10 +79,9 @@ static void base_w(int *output, const unsigned char *input, int in_len, const wo
   int total = 0;
   int bits = 0;
   int consumed = 0;
-  
-  for(consumed = 0; consumed < 8 * in_len; consumed += params->log_w)
-  {
-    if(bits == 0){
+
+  for (consumed = 0; consumed < 8 * in_len; consumed += params->log_w) {
+    if (bits == 0) {
       total = input[in_len - 1 - in];
       in++;
       bits += 8;
@@ -96,29 +92,13 @@ static void base_w(int *output, const unsigned char *input, int in_len, const wo
   }
 }
 
-/**
- * Alternative base w algorithm for w = 16 to check...
- */
-static void base_w_alternative(int *output, unsigned char *input, int in_len, const wots_params *params)
-{
-  unsigned int i = 0;
-  for(i = 0; i < in_len; i += 2)
-  {
-    output[i] = input[in_len - 1 - (i / 2)] >> 4;
-    output[i+1] = input[in_len - 1 - (i / 2)] & 0xf;
-  }
-}
-
 void wots_pkgen(unsigned char *pk, const unsigned char *sk, const wots_params *params, const unsigned char *pub_seed, unsigned char addr[16])
 {
   unsigned int i;
   expand_seed(pk, sk, params);
-  for(i=0;i<params->len;i++){
-    SET_CHAIN_ADDRESS(addr,i);
-//     printf("Chain: %d\n",i);
-//     hexdump(addr,16);
-//     printf("\n");
-     gen_chain(pk+i*params->n, pk+i*params->n, 0, params->w-1, params, pub_seed, addr);
+  for (i=0; i < params->len; i++) {
+    SET_CHAIN_ADDRESS(addr, i);
+    gen_chain(pk+i*params->n, pk+i*params->n, 0, params->w-1, params, pub_seed, addr);
   }
 }
 
@@ -127,38 +107,33 @@ void wots_sign(unsigned char *sig, const unsigned char *msg, const unsigned char
 {
   int basew[params->len];
   int csum = 0;
-  unsigned int i=0;
+  unsigned int i = 0;
 
   base_w(basew, msg, params->m, params);
-  
-  for(i=0;i<params->len_1;i++)
-  {
+
+  for (i=0; i < params->len_1; i++) {
     csum += params->w - 1 - basew[i];
   }
 
-  csum = csum << ( 8 - ( ( params->len_2 * params->log_w ) % 8 ));
-  
+  csum = csum << (8 - ((params->len_2 * params->log_w) % 8));
+
   int len_2_bytes = ((params->len_2 * params->log_w) + 7) / 8;
-  
+
   unsigned char csum_bytes[len_2_bytes];
   to_byte(csum_bytes, csum, len_2_bytes);
-  
+
   int csum_basew[len_2_bytes / params->log_w];
   base_w(csum_basew, csum_bytes, len_2_bytes, params);
-  
-  for(i = 0; i < params->len_2; i++)
-  {
+
+  for (i = 0; i < params->len_2; i++) {
     basew[params->len_1 + i] = csum_basew[i];
   }
-  
+
   expand_seed(sig, sk, params);
-  
-  for(i=0;i<params->len;i++){
-    SET_CHAIN_ADDRESS(addr,i);
-//     printf("Chain: %d\n",i);
-//     hexdump(addr,16);
-//     printf("\n");
-     gen_chain(sig+i*params->n, sig+i*params->n, 0, basew[i], params, pub_seed, addr);
+
+  for (i = 0; i < params->len; i++) {
+    SET_CHAIN_ADDRESS(addr, i);
+    gen_chain(sig+i*params->n, sig+i*params->n, 0, basew[i], params, pub_seed, addr);
   }
 }
 
@@ -166,35 +141,29 @@ void wots_pkFromSig(unsigned char *pk, const unsigned char *sig, const unsigned
 {
   int basew[params->len];
   int csum = 0;
-  unsigned int i=0;
+  unsigned int i = 0;
 
   base_w(basew, msg, params->m, params);
-  
-  for(i=0;i<params->len_1;i++)
-  {
+
+  for (i=0; i < params->len_1; i++) {
     csum += params->w - 1 - basew[i];
   }
 
-  csum = csum << ( 8 - ( ( params->len_2 * params->log_w ) % 8 ));
-  
+  csum = csum << (8 - ((params->len_2 * params->log_w) % 8));
+
   int len_2_bytes = ((params->len_2 * params->log_w) + 7) / 8;
-  
+
   unsigned char csum_bytes[len_2_bytes];
   to_byte(csum_bytes, csum, len_2_bytes);
-  
+
   int csum_basew[len_2_bytes / params->log_w];
   base_w(csum_basew, csum_bytes, len_2_bytes, params);
-  
-  for(i = 0; i < params->len_2; i++)
-  {
+
+  for (i = 0; i < params->len_2; i++) {
     basew[params->len_1 + i] = csum_basew[i];
   }
-  
-  for(i=0;i<params->len;i++){
-    SET_CHAIN_ADDRESS(addr,i);
-//     printf("Chain: %d\n",i);
-//     hexdump(addr,16);
-//     printf("\n");
-     gen_chain(pk+i*params->n, sig+i*params->n, basew[i], params->w-1-basew[i], params, pub_seed, addr);
+  for (i=0; i < params->len; i++) {
+    SET_CHAIN_ADDRESS(addr, i);
+    gen_chain(pk+i*params->n, sig+i*params->n, basew[i], params->w-1-basew[i], params, pub_seed, addr);
   }
 }
diff --git a/wots.h b/wots.h
index 83b4558..b9585eb 100644
--- a/wots.h
+++ b/wots.h
@@ -12,15 +12,15 @@ Public domain.
  * 
  * Meaning as defined in draft-irtf-cfrg-xmss-hash-based-signatures-02
  */
-typedef struct{
-  int len_1;
-  int len_2;
-  int len;
-  int m;
-  int n;
-  int w;
-  int log_w;
-  int keysize;
+typedef struct {
+  unsigned int len_1;
+  unsigned int len_2;
+  unsigned int len;
+  unsigned int m;
+  unsigned int n;
+  unsigned int w;
+  unsigned int log_w;
+  unsigned int keysize;
 } wots_params;
 
 /**
diff --git a/xmss.c b/xmss.c
index 7f8631f..e626166 100644
--- a/xmss.c
+++ b/xmss.c
@@ -24,32 +24,32 @@ Public domain.
  * in the 16byte hash address
  */
 #define SET_LAYER_ADDRESS(a, v) {\
-  a[6] = (a[6] & 3) | ((v << 2) & 255);\
-  a[5] = (a[5] & 252) | ((v >> 6) & 255);}  
+  a[6] = (a[6] & 3) | ((v << 2) & 252);\
+  a[5] = (a[5] & 252) | ((v >> 6) & 3);}
 
 #define SET_TREE_ADDRESS(a, v) {\
-  a[9] = (a[9] & 3) | ((v << 2) & 255);\
+  a[9] = (a[9] & 3) | ((v << 2) & 252);\
   a[8] = (v >> 6) & 255;\
   a[7] = (v >> 14) & 255;\
-  a[6] = (a[6] & 252) | ((v >> 22) & 255);}  
-  
+  a[6] = (a[6] & 252) | ((v >> 22) & 3);}
+
 #define SET_OTS_BIT(a, b) {\
-  a[9] = (a[9] & 253) | (b << 1);}
+  a[9] = (a[9] & 253) | ((b << 1) & 2);}
 
 #define SET_OTS_ADDRESS(a, v) {\
-  a[12] = (a[12] & 1) | ((v << 1) & 255);\
+  a[12] = (a[12] & 1) | ((v << 1) & 254);\
   a[11] = (v >> 7) & 255;\
   a[10] = (v >> 15) & 255;\
-  a[9] = (a[9] & 254) | ((v >> 23) & 1);}  
-  
+  a[9] = (a[9] & 254) | ((v >> 23) & 1);}
+
 #define ZEROISE_OTS_ADDR(a) {\
   a[12] = (a[12] & 254);\
   a[13] = 0;\
   a[14] = 0;\
   a[15] = 0;}
-  
+
 #define SET_LTREE_BIT(a, b) {\
-  a[9] = (a[9] & 254) | b;}
+  a[9] = (a[9] & 254) | (b & 1);}
 
 #define SET_LTREE_ADDRESS(a, v) {\
   a[12] = v & 255;\
@@ -57,23 +57,23 @@ Public domain.
   a[10] = (v >> 16) & 255;}
 
 #define SET_LTREE_TREE_HEIGHT(a, v) {\
-  a[13] = (a[13] & 3) | ((v << 2) & 255);}
+  a[13] = (a[13] & 3) | ((v << 2) & 252);}
 
 #define SET_LTREE_TREE_INDEX(a, v) {\
-  a[15] = (a[15] & 3) | ((v << 2) & 255);\
+  a[15] = (a[15] & 3) | ((v << 2) & 252);\
   a[14] = (v >> 6) & 255;\
   a[13] = (a[13] & 252) | ((v >> 14) & 3);}
-  
+
 #define SET_NODE_PADDING(a) {\
   a[10] = 0;\
-  a[11] = a[11] & 3;}  
+  a[11] = a[11] & 3;}
 
 #define SET_NODE_TREE_HEIGHT(a, v) {\
-  a[12] = (a[12] & 3) | ((v << 2) & 255);\
+  a[12] = (a[12] & 3) | ((v << 2) & 252);\
   a[11] = (a[11] & 252) | ((v >> 6) & 3);}
 
 #define SET_NODE_TREE_INDEX(a, v) {\
-  a[15] = (a[15] & 3) | ((v << 2) & 255);\
+  a[15] = (a[15] & 3) | ((v << 2) & 252);\
   a[14] = (v >> 6) & 255;\
   a[13] = (v >> 14) & 255;\
   a[12] = (a[12] & 252) | ((v >> 22) & 3);}
@@ -82,7 +82,7 @@ Public domain.
 /**
  * Used for pseudorandom keygeneration,
  * generates the seed for the WOTS keypair at address addr
- * 
+ *
  * takes n byte sk_seed and returns n byte seed using 16 byte address addr.
  */
 static void get_seed(unsigned char *seed, const unsigned char *sk_seed, int n, unsigned char addr[16])
@@ -110,12 +110,12 @@ void xmss_set_params(xmss_params *params, int m, int n, int h, int w)
 /**
  * Initialize xmssmt_params struct
  * parameter names are the same as in the draft
- * 
+ *
  * Especially h is the total tree height, i.e. the XMSS trees have height h/d
  */
 void xmssmt_set_params(xmssmt_params *params, int m, int n, int h, int d, int w)
 {
-  if(h % d){
+  if (h % d) {
     fprintf(stderr, "d must devide h without remainder!\n");
     return;
   }
@@ -133,43 +133,40 @@ void xmssmt_set_params(xmssmt_params *params, int m, int n, int h, int d, int w)
  * Computes a leaf from a WOTS public key using an L-tree.
  */
 static void l_tree(unsigned char *leaf, unsigned char *wots_pk, const xmss_params *params, const unsigned char *pub_seed, unsigned char addr[16])
-{ 
+{
   unsigned int l = params->wots_par.len;
   unsigned int n = params->n;
   unsigned long i = 0;
   unsigned int height = 0;
-  
+
   //ADRS.setTreeHeight(0);
-  SET_LTREE_TREE_HEIGHT(addr,height);
+  SET_LTREE_TREE_HEIGHT(addr, height);
   unsigned long bound;
-  while ( l > 1 ) 
-  {
-     bound = l >> 1; //floor(l / 2); 
-     for ( i = 0; i < bound; i = i + 1 ) {
+  while (l > 1) {
+     bound = l >> 1; //floor(l / 2);
+     for (i = 0; i < bound; i++) {
        //ADRS.setTreeIndex(i);
-       SET_LTREE_TREE_INDEX(addr,i);
+       SET_LTREE_TREE_INDEX(addr, i);
        //wots_pk[i] = RAND_HASH(pk[2i], pk[2i + 1], SEED, ADRS);
-       hash_2n_n(wots_pk+i*n,wots_pk+i*2*n, pub_seed, addr, n);
+       hash_2n_n(wots_pk+i*n, wots_pk+i*2*n, pub_seed, addr, n);
      }
      //if ( l % 2 == 1 ) {
-     if(l&1)
-     {
+     if (l & 1) {
        //pk[floor(l / 2) + 1] = pk[l];
-       memcpy(wots_pk+(l>>1)*n,wots_pk+(l-1)*n, n);
+       memcpy(wots_pk+(l>>1)*n, wots_pk+(l-1)*n, n);
        //l = ceil(l / 2);
        l=(l>>1)+1;
      }
-     else
-     {
+     else {
        //l = ceil(l / 2);
        l=(l>>1);
-     }     
+     }
      //ADRS.setTreeHeight(ADRS.getTreeHeight() + 1);
      height++;
-     SET_LTREE_TREE_HEIGHT(addr,height);
+     SET_LTREE_TREE_HEIGHT(addr, height);
    }
    //return pk[0];
-   memcpy(leaf,wots_pk,n);
+   memcpy(leaf, wots_pk, n);
 }
 
 /**
@@ -183,13 +180,13 @@ static void gen_leaf_wots(unsigned char *leaf, const unsigned char *sk_seed, con
   get_seed(seed, sk_seed, params->n, ots_addr);
   wots_pkgen(pk, seed, &(params->wots_par), pub_seed, ots_addr);
 
-  l_tree(leaf, pk, params, pub_seed, ltree_addr); 
+  l_tree(leaf, pk, params, pub_seed, ltree_addr);
 }
 
 /**
  * Merkle's TreeHash algorithm. The address only needs to initialize the first 78 bits of addr. Everything else will be set by treehash.
  * Currently only used for key generation.
- * 
+ *
  */
 static void treehash(unsigned char *node, int height, int index, const unsigned char *sk_seed, const xmss_params *params, const unsigned char *pub_seed, const unsigned char addr[16])
 {
@@ -208,32 +205,30 @@ static void treehash(unsigned char *node, int height, int index, const unsigned
   memcpy(node_addr, ltree_addr, 10);
   SET_LTREE_BIT(node_addr, 0);
   SET_NODE_PADDING(node_addr);
-  
-  int lastnode,i;
+
+  unsigned int lastnode, i;
   unsigned char stack[(height+1)*n];
-  unsigned int  stacklevels[height+1];
-  unsigned int  stackoffset=0;
-  
-  lastnode = idx+(1<<height);
-
-  for(;idx<lastnode;idx++) 
-  {
-    SET_LTREE_ADDRESS(ltree_addr,idx);
-    SET_OTS_ADDRESS(ots_addr,idx);
-    gen_leaf_wots(stack+stackoffset*n,sk_seed,params, pub_seed, ltree_addr, ots_addr);
+  unsigned int stacklevels[height+1];
+  unsigned int stackoffset=0;
+
+  lastnode = idx+(1 << height);
+
+  for (; idx < lastnode; idx++) {
+    SET_LTREE_ADDRESS(ltree_addr, idx);
+    SET_OTS_ADDRESS(ots_addr, idx);
+    gen_leaf_wots(stack+stackoffset*n, sk_seed, params, pub_seed, ltree_addr, ots_addr);
     stacklevels[stackoffset] = 0;
     stackoffset++;
-    while(stackoffset>1 && stacklevels[stackoffset-1] == stacklevels[stackoffset-2])
-    {
-      SET_NODE_TREE_HEIGHT(node_addr,stacklevels[stackoffset-1]);
+    while (stackoffset>1 && stacklevels[stackoffset-1] == stacklevels[stackoffset-2]) {
+      SET_NODE_TREE_HEIGHT(node_addr, stacklevels[stackoffset-1]);
       SET_NODE_TREE_INDEX(node_addr, (idx >> (stacklevels[stackoffset-1]+1)));
-      hash_2n_n(stack+(stackoffset-2)*n,stack+(stackoffset-2)*n, pub_seed,
+      hash_2n_n(stack+(stackoffset-2)*n, stack+(stackoffset-2)*n, pub_seed,
           node_addr, n);
       stacklevels[stackoffset-2]++;
       stackoffset--;
     }
   }
-  for(i=0;i<n;i++)
+  for (i=0; i < n; i++)
     node[i] = stack[i];
 }
 
@@ -243,43 +238,38 @@ static void treehash(unsigned char *node, int height, int index, const unsigned
 static void validate_authpath(unsigned char *root, const unsigned char *leaf, unsigned long leafidx, const unsigned char *authpath, const xmss_params *params, const unsigned char *pub_seed, unsigned char addr[16])
 {
   unsigned int n = params->n;
-  
-  int i,j;
+
+  unsigned int i, j;
   unsigned char buffer[2*n];
 
   // If leafidx is odd (last bit = 1), current path element is a right child and authpath has to go to the left.
   // Otherwise, it is the other way around
-  if(leafidx&1)
-  {
-    for(j=0;j<n;j++)
+  if (leafidx & 1) {
+    for (j = 0; j < n; j++)
       buffer[n+j] = leaf[j];
-    for(j=0;j<n;j++)
+    for (j = 0; j < n; j++)
       buffer[j] = authpath[j];
   }
-  else
-  {
-    for(j=0;j<n;j++)
+  else {
+    for (j = 0; j < n; j++)
       buffer[j] = leaf[j];
-    for(j=0;j<n;j++)
+    for (j = 0; j < n; j++)
       buffer[n+j] = authpath[j];
   }
   authpath += n;
 
-  for(i=0;i<params->h-1;i++)
-  {
-    SET_NODE_TREE_HEIGHT(addr,i);
+  for (i=0; i < params->h-1; i++) {
+    SET_NODE_TREE_HEIGHT(addr, i);
     leafidx >>= 1;
     SET_NODE_TREE_INDEX(addr, leafidx);
-    if(leafidx&1)
-    {
-      hash_2n_n(buffer+n,buffer,pub_seed, addr, n);
-      for(j=0;j<n;j++)
+    if (leafidx&1) {
+      hash_2n_n(buffer+n, buffer, pub_seed, addr, n);
+      for (j = 0; j < n; j++)
         buffer[j] = authpath[j];
     }
-    else
-    {
-      hash_2n_n(buffer,buffer,pub_seed, addr, n);
-      for(j=0;j<n;j++)
+    else {
+      hash_2n_n(buffer, buffer, pub_seed, addr, n);
+      for (j = 0; j < n; j++)
         buffer[j+n] = authpath[j];
     }
     authpath += n;
@@ -287,26 +277,26 @@ static void validate_authpath(unsigned char *root, const unsigned char *leaf, un
   SET_NODE_TREE_HEIGHT(addr, (params->h-1));
   leafidx >>= 1;
   SET_NODE_TREE_INDEX(addr, leafidx);
-  hash_2n_n(root,buffer,pub_seed,addr,n);
+  hash_2n_n(root, buffer, pub_seed, addr, n);
 }
 
 /**
  * Computes the authpath and the root. This method is using a lot of space as we build the whole tree and then select the authpath nodes.
  * For more efficient algorithms see e.g. the chapter on hash-based signatures in Bernstein, Buchmann, Dahmen. "Post-quantum Cryptography", Springer 2009.
- * It returns the authpath in "authpath" with the node on level 0 at index 0.  
+ * It returns the authpath in "authpath" with the node on level 0 at index 0.
  */
 static void compute_authpath_wots(unsigned char *root, unsigned char *authpath, unsigned long leaf_idx, const unsigned char *sk_seed, const xmss_params *params, unsigned char *pub_seed, unsigned char addr[16])
 {
   unsigned int i, j, level;
-  int n = params->n;
-  int h = params->h;
-  
+  unsigned int n = params->n;
+  unsigned int h = params->h;
+
   unsigned char tree[2*(1<<h)*n];
 
   unsigned char ots_addr[16];
   unsigned char ltree_addr[16];
   unsigned char node_addr[16];
-  
+
   memcpy(ots_addr, addr, 10);
   SET_OTS_BIT(ots_addr, 1);
   memcpy(ltree_addr, addr, 10);
@@ -316,25 +306,22 @@ static void compute_authpath_wots(unsigned char *root, unsigned char *authpath,
   SET_LTREE_BIT(node_addr, 0);
   SET_NODE_PADDING(node_addr);
 
-  
+
   // Compute all leaves
-  for(i = 0; i < (1<<h); i++)
-  {
-    SET_LTREE_ADDRESS(ltree_addr,i);
-    SET_OTS_ADDRESS(ots_addr,i);
+  for (i = 0; i < (1U << h); i++) {
+    SET_LTREE_ADDRESS(ltree_addr, i);
+    SET_OTS_ADDRESS(ots_addr, i);
     gen_leaf_wots(tree+((1<<h)*n + i*n), sk_seed, params, pub_seed, ltree_addr, ots_addr);
   }
-  
-  
+
+
   level = 0;
   // Compute tree:
-  // Outer loop: For each inner layer 
-  for (i = (1<<h); i > 1; i>>=1)
-  {
+  // Outer loop: For each inner layer
+  for (i = (1<<h); i > 1; i>>=1) {
     SET_NODE_TREE_HEIGHT(node_addr, level);
     // Inner loop: for each pair of sibling nodes
-    for (j = 0; j < i; j+=2)
-    {
+    for (j = 0; j < i; j+=2) {
       SET_NODE_TREE_INDEX(node_addr, j>>1);
       hash_2n_n(tree + (i>>1)*n + (j>>1) * n, tree + i*n + j*n, pub_seed, node_addr, n);
     }
@@ -342,9 +329,9 @@ static void compute_authpath_wots(unsigned char *root, unsigned char *authpath,
   }
 
   // copy authpath
-  for(i=0;i<h;i++)
+  for (i=0; i < h; i++)
     memcpy(authpath + i*n, tree + ((1<<h)>>i)*n + ((leaf_idx >> i) ^ 1) * n, n);
-  
+
   // copy root
   memcpy(root, tree+n, n);
 }
@@ -365,11 +352,11 @@ int xmss_keypair(unsigned char *pk, unsigned char *sk, xmss_params *params)
   sk[2] = 0;
   sk[3] = 0;
   // Init SK_SEED (n byte), SK_PRF (m byte), and PUB_SEED (n byte)
-  randombytes(sk+4,2*n+m);
+  randombytes(sk+4, 2*n+m);
   // Copy PUB_SEED to public key
-  memcpy(pk+n, sk+4+n+m,n);
+  memcpy(pk+n, sk+4+n+m, n);
 
-  unsigned char addr[16] = {0,0,0,0};
+  unsigned char addr[16] = {0, 0, 0, 0};
   // Compute root
   treehash(pk, params->h, 0, sk+4, params, sk+4+n+m, addr);
   return 0;
@@ -377,51 +364,51 @@ int xmss_keypair(unsigned char *pk, unsigned char *sk, xmss_params *params)
 
 /**
  * Signs a message.
- * Returns 
+ * Returns
  * 1. an array containing the signature followed by the message AND
  * 2. an updated secret key!
- * 
+ *
  */
 int xmss_sign(unsigned char *sk, unsigned char *sig_msg, unsigned long long *sig_msg_len, const unsigned char *msg, unsigned long long msglen, const xmss_params *params)
 {
   unsigned int n = params->n;
   unsigned int m = params->m;
-  
+
   // Extract SK
   unsigned long idx = ((unsigned long)sk[0] << 24) | ((unsigned long)sk[1] << 16) | ((unsigned long)sk[2] << 8) | sk[3];
   unsigned char sk_seed[n];
-  memcpy(sk_seed,sk+4,n);
+  memcpy(sk_seed, sk+4, n);
   unsigned char sk_prf[m];
-  memcpy(sk_prf,sk+4+n,m);
+  memcpy(sk_prf, sk+4+n, m);
   unsigned char pub_seed[n];
-  memcpy(pub_seed,sk+4+n+m,n);  
-  
+  memcpy(pub_seed, sk+4+n+m, n);
+
   // Update SK
   sk[0] = ((idx + 1) >> 24) & 255;
   sk[1] = ((idx + 1) >> 16) & 255;
   sk[2] = ((idx + 1) >> 8) & 255;
   sk[3] = (idx + 1) & 255;
-  // -- Secret key for this non-forward-secure version is now updated. 
-  // -- A productive implementation should use a file handle instead and write the updated secret key at this point! 
-  
+  // -- Secret key for this non-forward-secure version is now updated.
+  // -- A productive implementation should use a file handle instead and write the updated secret key at this point!
+
   // Init working params
   unsigned long long i;
   unsigned char R[m];
   unsigned char msg_h[m];
   unsigned char root[n];
   unsigned char ots_seed[n];
-  unsigned char ots_addr[16] = {0,0,0,0};
-  
+  unsigned char ots_addr[16] = {0, 0, 0, 0};
+
   // ---------------------------------
   // Message Hashing
   // ---------------------------------
-  
-  // Message Hash: 
+
+  // Message Hash:
   // First compute pseudorandom key
-  prf_m(R, msg, msglen, sk_prf, m); 
+  prf_m(R, msg, msglen, sk_prf, m);
   // Then use it for message digest
   hash_m(msg_h, msg, msglen, R, m, m);
-  
+
   // Start collecting signature
   *sig_msg_len = 0;
 
@@ -430,42 +417,42 @@ int xmss_sign(unsigned char *sk, unsigned char *sig_msg, unsigned long long *sig
   sig_msg[1] = (idx >> 16) & 255;
   sig_msg[2] = (idx >> 8) & 255;
   sig_msg[3] = idx & 255;
-  
+
   sig_msg += 4;
   *sig_msg_len += 4;
-  
+
   // Copy R to signature
-  for(i=0; i<m; i++)
+  for (i = 0; i < m; i++)
     sig_msg[i] = R[i];
 
   sig_msg += m;
   *sig_msg_len += m;
-  
+
   // ----------------------------------
-  // Now we start to "really sign" 
+  // Now we start to "really sign"
   // ----------------------------------
-  
+
   // Prepare Address
-  SET_OTS_BIT(ots_addr,1);
-  SET_OTS_ADDRESS(ots_addr,idx);
-  
+  SET_OTS_BIT(ots_addr, 1);
+  SET_OTS_ADDRESS(ots_addr, idx);
+
   // Compute seed for OTS key pair
   get_seed(ots_seed, sk_seed, n, ots_addr);
-     
+
   // Compute WOTS signature
   wots_sign(sig_msg, msg_h, ots_seed, &(params->wots_par), pub_seed, ots_addr);
-  
+
   sig_msg += params->wots_par.keysize;
   *sig_msg_len += params->wots_par.keysize;
 
   compute_authpath_wots(root, sig_msg, idx, sk_seed, params, pub_seed, ots_addr);
   sig_msg += params->h*n;
   *sig_msg_len += params->h*n;
-  
-  //Whipe secret elements?  
+
+  //Whipe secret elements?
   //zerobytes(tsk, CRYPTO_SECRETKEYBYTES);
 
-  memcpy(sig_msg,msg,msglen);
+  memcpy(sig_msg, msg, msglen);
   *sig_msg_len += msglen;
 
   return 0;
@@ -478,38 +465,38 @@ int xmss_sign_open(unsigned char *msg, unsigned long long *msglen, const unsigne
 {
   unsigned int n = params->n;
   unsigned int m = params->m;
-    
+
   unsigned long long i, m_len;
   unsigned long idx=0;
   unsigned char wots_pk[params->wots_par.keysize];
   unsigned char pkhash[n];
   unsigned char root[n];
   unsigned char msg_h[m];
-  
+
   unsigned char pub_seed[n];
-  memcpy(pub_seed,pk+n,n);  
-  
+  memcpy(pub_seed, pk+n, n);
+
   // Init addresses
-  unsigned char ots_addr[16] = {0,0,0,0};
+  unsigned char ots_addr[16] = {0, 0, 0, 0};
   unsigned char ltree_addr[16];
   unsigned char node_addr[16];
-  
+
   SET_OTS_BIT(ots_addr, 1);
-  
+
   memcpy(ltree_addr, ots_addr, 10);
   SET_OTS_BIT(ltree_addr, 0);
   SET_LTREE_BIT(ltree_addr, 1);
-  
+
   memcpy(node_addr, ltree_addr, 10);
   SET_LTREE_BIT(node_addr, 0);
-  SET_NODE_PADDING(node_addr);  
-  
+  SET_NODE_PADDING(node_addr);
+
   // Extract index
   idx = ((unsigned long)sig_msg[0] << 24) | ((unsigned long)sig_msg[1] << 16) | ((unsigned long)sig_msg[2] << 8) | sig_msg[3];
-  printf("verify:: idx = %lu\n",idx);
+  printf("verify:: idx = %lu\n", idx);
   sig_msg += 4;
   sig_msg_len -= 4;
-  
+
   // hash message (recall, R is now on pole position at sig_msg
   unsigned long long tmp_sig_len = m+params->wots_par.keysize+params->h*n;
   m_len = sig_msg_len - tmp_sig_len;
@@ -517,43 +504,43 @@ int xmss_sign_open(unsigned char *msg, unsigned long long *msglen, const unsigne
 
   sig_msg += m;
   sig_msg_len -= m;
-  
+
   //-----------------------
   // Verify signature
   //-----------------------
-  
+
   // Prepare Address
-  SET_OTS_ADDRESS(ots_addr,idx);
-  // Check WOTS signature 
+  SET_OTS_ADDRESS(ots_addr, idx);
+  // Check WOTS signature
   wots_pkFromSig(wots_pk, sig_msg, msg_h, &(params->wots_par), pub_seed, ots_addr);
 
   sig_msg += params->wots_par.keysize;
   sig_msg_len -= params->wots_par.keysize;
-  
+
   // Compute Ltree
-  SET_LTREE_ADDRESS(ltree_addr, idx); 
+  SET_LTREE_ADDRESS(ltree_addr, idx);
   l_tree(pkhash, wots_pk, params, pub_seed, ltree_addr);
-  
+
   // Compute root
-  validate_authpath(root, pkhash, idx, sig_msg, params, pub_seed, node_addr);  
+  validate_authpath(root, pkhash, idx, sig_msg, params, pub_seed, node_addr);
 
   sig_msg += params->h*n;
   sig_msg_len -= params->h*n;
-  
-  for(i=0;i<n;i++)
-    if(root[i] != pk[i])
+
+  for (i=0; i < n; i++)
+    if (root[i] != pk[i])
       goto fail;
-  
+
   *msglen = sig_msg_len;
-  for(i=0;i<*msglen;i++)
+  for (i=0; i < *msglen; i++)
     msg[i] = sig_msg[i];
 
   return 0;
-  
-  
+
+
 fail:
   *msglen = sig_msg_len;
-  for(i=0;i<*msglen;i++)
+  for (i=0; i < *msglen; i++)
     msg[i] = 0;
   *msglen = -1;
   return -1;
@@ -570,18 +557,18 @@ int xmssmt_keypair(unsigned char *pk, unsigned char *sk, xmssmt_params *params)
   unsigned int m = params->m;
   unsigned int i;
   // Set idx = 0
-  for (i = 0; i < params->index_len; i++){
+  for (i = 0; i < params->index_len; i++) {
     sk[i] = 0;
   }
   // Init SK_SEED (n byte), SK_PRF (m byte), and PUB_SEED (n byte)
-  randombytes(sk+params->index_len,2*n+m);
+  randombytes(sk+params->index_len, 2*n+m);
   // Copy PUB_SEED to public key
-  memcpy(pk+n, sk+params->index_len+n+m,n);
+  memcpy(pk+n, sk+params->index_len+n+m, n);
 
   // Set address to point on the single tree on layer d-1
-  unsigned char addr[16] = {0,0,0,0};
+  unsigned char addr[16] = {0, 0, 0, 0};
   SET_LAYER_ADDRESS(addr, (params->d-1));
-  
+
   // Compute root
   treehash(pk, params->xmss_par.h, 0, sk+params->index_len, &(params->xmss_par), pk+n, addr);
   return 0;
@@ -589,10 +576,10 @@ int xmssmt_keypair(unsigned char *pk, unsigned char *sk, xmssmt_params *params)
 
 /**
  * Signs a message.
- * Returns 
+ * Returns
  * 1. an array containing the signature followed by the message AND
  * 2. an updated secret key!
- * 
+ *
  */
 int xmssmt_sign(unsigned char *sk, unsigned char *sig_msg, unsigned long long *sig_msg_len, const unsigned char *msg, unsigned long long msglen, const xmssmt_params *params)
 {
@@ -603,7 +590,7 @@ int xmssmt_sign(unsigned char *sk, unsigned char *sig_msg, unsigned long long *s
   unsigned long long idx_tree;
   unsigned long long idx_leaf;
   unsigned long long i;
-  
+
   unsigned char sk_seed[n];
   unsigned char sk_prf[m];
   unsigned char pub_seed[n];
@@ -612,109 +599,109 @@ int xmssmt_sign(unsigned char *sk, unsigned char *sig_msg, unsigned long long *s
   unsigned char msg_h[m];
   unsigned char root[n];
   unsigned char ots_seed[n];
-  unsigned char ots_addr[16] = {0,0,0,0};
-  
+  unsigned char ots_addr[16] = {0, 0, 0, 0};
+
   // Extract SK
   unsigned long long idx = 0;
-  for(i = 0; i < idx_len; i++){
+  for (i = 0; i < idx_len; i++) {
     idx |= ((unsigned long long)sk[i]) << 8*(idx_len - 1 - i);
   }
-  
-  memcpy(sk_seed,sk+idx_len,n);
-  memcpy(sk_prf,sk+idx_len+n,m);
-  memcpy(pub_seed,sk+idx_len+n+m,n);  
-  
+
+  memcpy(sk_seed, sk+idx_len, n);
+  memcpy(sk_prf, sk+idx_len+n, m);
+  memcpy(pub_seed, sk+idx_len+n+m, n);
+
   // Update SK
-  for(i = 0; i < idx_len; i++){
+  for (i = 0; i < idx_len; i++) {
     sk[i] = ((idx + 1) >> 8*(idx_len - 1 - i)) & 255;
   }
-  // -- Secret key for this non-forward-secure version is now updated. 
-  // -- A productive implementation should use a file handle instead and write the updated secret key at this point! 
-  
-  
+  // -- Secret key for this non-forward-secure version is now updated.
+  // -- A productive implementation should use a file handle instead and write the updated secret key at this point!
+
+
   // ---------------------------------
   // Message Hashing
   // ---------------------------------
-  
-  // Message Hash: 
+
+  // Message Hash:
   // First compute pseudorandom key
-  prf_m(R, msg, msglen, sk_prf, m); 
+  prf_m(R, msg, msglen, sk_prf, m);
   // Then use it for message digest
   hash_m(msg_h, msg, msglen, R, m, m);
-  
+
   // Start collecting signature
   *sig_msg_len = 0;
 
   // Copy index to signature
-  for(i = 0; i < idx_len; i++){
+  for (i = 0; i < idx_len; i++) {
     sig_msg[i] = (idx >> 8*(idx_len - 1 - i)) & 255;
   }
-  
+
   sig_msg += idx_len;
   *sig_msg_len += idx_len;
-  
+
   // Copy R to signature
-  for(i=0; i<m; i++)
+  for (i=0; i < m; i++)
     sig_msg[i] = R[i];
 
   sig_msg += m;
   *sig_msg_len += m;
-  
+
   // ----------------------------------
-  // Now we start to "really sign" 
+  // Now we start to "really sign"
   // ----------------------------------
-  
+
   // Handle lowest layer separately as it is slightly different...
-  
+
   // Prepare Address
-  SET_OTS_BIT(ots_addr,1);
+  SET_OTS_BIT(ots_addr, 1);
   idx_tree = idx >> tree_h;
   idx_leaf = (idx & ((1 << tree_h)-1));
-  SET_LAYER_ADDRESS(ots_addr,0);
+  SET_LAYER_ADDRESS(ots_addr, 0);
   SET_TREE_ADDRESS(ots_addr, idx_tree);
   SET_OTS_ADDRESS(ots_addr, idx_leaf);
-  
+
   // Compute seed for OTS key pair
   get_seed(ots_seed, sk_seed, n, ots_addr);
-     
+
   // Compute WOTS signature
   wots_sign(sig_msg, msg_h, ots_seed, &(params->xmss_par.wots_par), pub_seed, ots_addr);
-  
+
   sig_msg += params->xmss_par.wots_par.keysize;
   *sig_msg_len += params->xmss_par.wots_par.keysize;
 
   compute_authpath_wots(root, sig_msg, idx_leaf, sk_seed, &(params->xmss_par), pub_seed, ots_addr);
   sig_msg += tree_h*n;
   *sig_msg_len += tree_h*n;
-  
+
   // Now loop over remaining layers...
   unsigned int j;
-  for(j = 1; j < params->d; j++){
+  for (j = 1; j < params->d; j++) {
     // Prepare Address
     idx_leaf = (idx_tree & ((1 << tree_h)-1));
     idx_tree = idx_tree >> tree_h;
-    SET_LAYER_ADDRESS(ots_addr,j);
+    SET_LAYER_ADDRESS(ots_addr, j);
     SET_TREE_ADDRESS(ots_addr, idx_tree);
     SET_OTS_ADDRESS(ots_addr, idx_leaf);
-    
+
     // Compute seed for OTS key pair
     get_seed(ots_seed, sk_seed, n, ots_addr);
-      
+
     // Compute WOTS signature
     wots_sign(sig_msg, root, ots_seed, &(params->xmss_par.wots_par), pub_seed, ots_addr);
-    
+
     sig_msg += params->xmss_par.wots_par.keysize;
     *sig_msg_len += params->xmss_par.wots_par.keysize;
 
     compute_authpath_wots(root, sig_msg, idx_leaf, sk_seed, &(params->xmss_par), pub_seed, ots_addr);
     sig_msg += tree_h*n;
-    *sig_msg_len += tree_h*n;   
+    *sig_msg_len += tree_h*n;
   }
-  
-  //Whipe secret elements?  
+
+  //Whipe secret elements?
   //zerobytes(tsk, CRYPTO_SECRETKEYBYTES);
 
-  memcpy(sig_msg,msg,msglen);
+  memcpy(sig_msg, msg, msglen);
   *sig_msg_len += msglen;
 
   return 0;
@@ -727,35 +714,35 @@ int xmssmt_sign_open(unsigned char *msg, unsigned long long *msglen, const unsig
 {
   unsigned int n = params->n;
   unsigned int m = params->m;
-  
+
   unsigned int tree_h = params->xmss_par.h;
   unsigned int idx_len = params->index_len;
   unsigned long long idx_tree;
   unsigned long long idx_leaf;
-  
+
   unsigned long long i, m_len;
   unsigned long long idx=0;
   unsigned char wots_pk[params->xmss_par.wots_par.keysize];
   unsigned char pkhash[n];
   unsigned char root[n];
   unsigned char msg_h[m];
-  
+
   unsigned char pub_seed[n];
-  memcpy(pub_seed,pk+n,n);  
-  
+  memcpy(pub_seed, pk+n, n);
+
   // Init addresses
-  unsigned char ots_addr[16] = {0,0,0,0};
+  unsigned char ots_addr[16] = {0, 0, 0, 0};
   unsigned char ltree_addr[16];
   unsigned char node_addr[16];
-  
+
   // Extract index
-  for(i = 0; i < idx_len; i++){
+  for (i = 0; i < idx_len; i++) {
     idx |= ((unsigned long long)sig_msg[i]) << (8*(idx_len - 1 - i));
   }
-  printf("verify:: idx = %llu\n",idx);
+  printf("verify:: idx = %llu\n", idx);
   sig_msg += idx_len;
   sig_msg_len -= idx_len;
-  
+
   // hash message (recall, R is now on pole position at sig_msg
   unsigned long long tmp_sig_len = m+ (params->d * params->xmss_par.wots_par.keysize) + (params->h * n);
   m_len = sig_msg_len - tmp_sig_len;
@@ -763,95 +750,95 @@ int xmssmt_sign_open(unsigned char *msg, unsigned long long *msglen, const unsig
 
   sig_msg += m;
   sig_msg_len -= m;
-  
+
   //-----------------------
   // Verify signature
   //-----------------------
-  
+
   // Prepare Address
   idx_tree = idx >> tree_h;
   idx_leaf = (idx & ((1 << tree_h)-1));
-  SET_LAYER_ADDRESS(ots_addr,0);
+  SET_LAYER_ADDRESS(ots_addr, 0);
   SET_TREE_ADDRESS(ots_addr, idx_tree);
   SET_OTS_BIT(ots_addr, 1);
-  
+
   memcpy(ltree_addr, ots_addr, 10);
   SET_OTS_BIT(ltree_addr, 0);
   SET_LTREE_BIT(ltree_addr, 1);
-  
+
   memcpy(node_addr, ltree_addr, 10);
   SET_LTREE_BIT(node_addr, 0);
-  SET_NODE_PADDING(node_addr);  
-  
-  SET_OTS_ADDRESS(ots_addr,idx_leaf);
-  
-  // Check WOTS signature 
+  SET_NODE_PADDING(node_addr);
+
+  SET_OTS_ADDRESS(ots_addr, idx_leaf);
+
+  // Check WOTS signature
   wots_pkFromSig(wots_pk, sig_msg, msg_h, &(params->xmss_par.wots_par), pub_seed, ots_addr);
 
   sig_msg += params->xmss_par.wots_par.keysize;
   sig_msg_len -= params->xmss_par.wots_par.keysize;
-  
+
   // Compute Ltree
-  SET_LTREE_ADDRESS(ltree_addr, idx_leaf); 
+  SET_LTREE_ADDRESS(ltree_addr, idx_leaf);
   l_tree(pkhash, wots_pk, &(params->xmss_par), pub_seed, ltree_addr);
-  
+
   // Compute root
-  validate_authpath(root, pkhash, idx_leaf, sig_msg, &(params->xmss_par), pub_seed, node_addr);  
+  validate_authpath(root, pkhash, idx_leaf, sig_msg, &(params->xmss_par), pub_seed, node_addr);
 
   sig_msg += tree_h*n;
   sig_msg_len -= tree_h*n;
-  
-  for(i = 1; i < params->d; i++){
+
+  for (i = 1; i < params->d; i++) {
     // Prepare Address
     idx_leaf = (idx_tree & ((1 << tree_h)-1));
     idx_tree = idx_tree >> tree_h;
-    
-    SET_LAYER_ADDRESS(ots_addr,i);
+
+    SET_LAYER_ADDRESS(ots_addr, i);
     SET_TREE_ADDRESS(ots_addr, idx_tree);
     SET_OTS_BIT(ots_addr, 1);
-    
+
     memcpy(ltree_addr, ots_addr, 10);
     SET_OTS_BIT(ltree_addr, 0);
     SET_LTREE_BIT(ltree_addr, 1);
-    
+
     memcpy(node_addr, ltree_addr, 10);
     SET_LTREE_BIT(node_addr, 0);
-    SET_NODE_PADDING(node_addr);  
-    
-    SET_OTS_ADDRESS(ots_addr,idx_leaf);
-    
-    // Check WOTS signature 
+    SET_NODE_PADDING(node_addr);
+
+    SET_OTS_ADDRESS(ots_addr, idx_leaf);
+
+    // Check WOTS signature
     wots_pkFromSig(wots_pk, sig_msg, root, &(params->xmss_par.wots_par), pub_seed, ots_addr);
 
     sig_msg += params->xmss_par.wots_par.keysize;
     sig_msg_len -= params->xmss_par.wots_par.keysize;
-    
+
     // Compute Ltree
-    SET_LTREE_ADDRESS(ltree_addr, idx_leaf); 
+    SET_LTREE_ADDRESS(ltree_addr, idx_leaf);
     l_tree(pkhash, wots_pk, &(params->xmss_par), pub_seed, ltree_addr);
-    
+
     // Compute root
-    validate_authpath(root, pkhash, idx_leaf, sig_msg, &(params->xmss_par), pub_seed, node_addr);  
+    validate_authpath(root, pkhash, idx_leaf, sig_msg, &(params->xmss_par), pub_seed, node_addr);
 
     sig_msg += tree_h*n;
     sig_msg_len -= tree_h*n;
-    
+
   }
-  
-  for(i=0;i<n;i++)
-    if(root[i] != pk[i])
+
+  for (i=0; i < n; i++)
+    if (root[i] != pk[i])
       goto fail;
-  
+
   *msglen = sig_msg_len;
-  for(i=0;i<*msglen;i++)
+  for (i=0; i < *msglen; i++)
     msg[i] = sig_msg[i];
 
   return 0;
-  
-  
+
+
 fail:
   *msglen = sig_msg_len;
-  for(i=0;i<*msglen;i++)
+  for (i=0; i < *msglen; i++)
     msg[i] = 0;
   *msglen = -1;
   return -1;
diff --git a/xmss.h b/xmss.h
index 0fb3b3d..d00403e 100644
--- a/xmss.h
+++ b/xmss.h
@@ -9,25 +9,25 @@ Public domain.
 #ifndef XMSS_H
 #define XMSS_H
 typedef struct{
-  int level;
+  unsigned int level;
   unsigned long long subtree;
-  int subleaf;
+  unsigned int subleaf;
 } leafaddr;
 
 typedef struct{
   wots_params wots_par;
-  int n;
-  int m;
-  int h;
+  unsigned int n;
+  unsigned int m;
+  unsigned int h;
 } xmss_params;
 
 typedef struct{
   xmss_params xmss_par;
-  int n;
-  int m;
-  int h;
-  int d;
-  int index_len;
+  unsigned int n;
+  unsigned int m;
+  unsigned int h;
+  unsigned int d;
+  unsigned int index_len;
 } xmssmt_params;
 /**
  * Initializes parameter set.
diff --git a/xmss_commons.c b/xmss_commons.c
index 16b8bae..f80bb73 100644
--- a/xmss_commons.c
+++ b/xmss_commons.c
@@ -11,7 +11,7 @@ Public domain.
 void to_byte(unsigned char *out, unsigned int in, int bytes)
 {
   int i;
-  for(i = 0; i < bytes; i++){
+  for (i = 0; i < bytes; i++) {
     out[i] = in & 0xff;
     in = in >> 8;
   }
diff --git a/xmss_fast.c b/xmss_fast.c
index dd80f05..17074f9 100644
--- a/xmss_fast.c
+++ b/xmss_fast.c
@@ -24,32 +24,32 @@ Public domain.
  * in the 16byte hash address
  */
 #define SET_LAYER_ADDRESS(a, v) {\
-  a[6] = (a[6] & 3) | ((v << 2) & 255);\
-  a[5] = (a[5] & 252) | ((v >> 6) & 255);}  
+  a[6] = (a[6] & 3) | ((v << 2) & 252);\
+  a[5] = (a[5] & 252) | ((v >> 6) & 3);}
 
 #define SET_TREE_ADDRESS(a, v) {\
-  a[9] = (a[9] & 3) | ((v << 2) & 255);\
+  a[9] = (a[9] & 3) | ((v << 2) & 252);\
   a[8] = (v >> 6) & 255;\
   a[7] = (v >> 14) & 255;\
-  a[6] = (a[6] & 252) | ((v >> 22) & 255);}  
-  
+  a[6] = (a[6] & 252) | ((v >> 22) & 3);}
+
 #define SET_OTS_BIT(a, b) {\
-  a[9] = (a[9] & 253) | (b << 1);}
+  a[9] = (a[9] & 253) | ((b << 1) & 2);}
 
 #define SET_OTS_ADDRESS(a, v) {\
-  a[12] = (a[12] & 1) | ((v << 1) & 255);\
+  a[12] = (a[12] & 1) | ((v << 1) & 254);\
   a[11] = (v >> 7) & 255;\
   a[10] = (v >> 15) & 255;\
-  a[9] = (a[9] & 254) | ((v >> 23) & 1);}  
-  
+  a[9] = (a[9] & 254) | ((v >> 23) & 1);}
+
 #define ZEROISE_OTS_ADDR(a) {\
   a[12] = (a[12] & 254);\
   a[13] = 0;\
   a[14] = 0;\
   a[15] = 0;}
-  
+
 #define SET_LTREE_BIT(a, b) {\
-  a[9] = (a[9] & 254) | b;}
+  a[9] = (a[9] & 254) | (b & 1);}
 
 #define SET_LTREE_ADDRESS(a, v) {\
   a[12] = v & 255;\
@@ -57,23 +57,23 @@ Public domain.
   a[10] = (v >> 16) & 255;}
 
 #define SET_LTREE_TREE_HEIGHT(a, v) {\
-  a[13] = (a[13] & 3) | ((v << 2) & 255);}
+  a[13] = (a[13] & 3) | ((v << 2) & 252);}
 
 #define SET_LTREE_TREE_INDEX(a, v) {\
-  a[15] = (a[15] & 3) | ((v << 2) & 255);\
+  a[15] = (a[15] & 3) | ((v << 2) & 252);\
   a[14] = (v >> 6) & 255;\
   a[13] = (a[13] & 252) | ((v >> 14) & 3);}
-  
+
 #define SET_NODE_PADDING(a) {\
   a[10] = 0;\
-  a[11] = a[11] & 3;}  
+  a[11] = a[11] & 3;}
 
 #define SET_NODE_TREE_HEIGHT(a, v) {\
-  a[12] = (a[12] & 3) | ((v << 2) & 255);\
+  a[12] = (a[12] & 3) | ((v << 2) & 252);\
   a[11] = (a[11] & 252) | ((v >> 6) & 3);}
 
 #define SET_NODE_TREE_INDEX(a, v) {\
-  a[15] = (a[15] & 3) | ((v << 2) & 255);\
+  a[15] = (a[15] & 3) | ((v << 2) & 252);\
   a[14] = (v >> 6) & 255;\
   a[13] = (v >> 14) & 255;\
   a[12] = (a[12] & 252) | ((v >> 22) & 3);}
@@ -81,7 +81,7 @@ Public domain.
 /**
  * Used for pseudorandom keygeneration,
  * generates the seed for the WOTS keypair at address addr
- * 
+ *
  * takes n byte sk_seed and returns n byte seed using 16 byte address addr.
  */
 static void get_seed(unsigned char *seed, const unsigned char *sk_seed, int n, unsigned char addr[16])
@@ -99,7 +99,7 @@ static void get_seed(unsigned char *seed, const unsigned char *sk_seed, int n, u
  */
 int xmss_set_params(xmss_params *params, int m, int n, int h, int w, int k)
 {
-  if(k >= h || k < 2 || (h - k) % 2){
+  if (k >= h || k < 2 || (h - k) % 2) {
     fprintf(stderr, "For BDS traversal, H - K must be even, with H > K >= 2!\n");
     return 1;
   }
@@ -132,12 +132,12 @@ void xmss_set_bds_state(bds_state *state, unsigned char *stack, int stackoffset,
 /**
  * Initialize xmssmt_params struct
  * parameter names are the same as in the draft
- * 
+ *
  * Especially h is the total tree height, i.e. the XMSS trees have height h/d
  */
 int xmssmt_set_params(xmssmt_params *params, int m, int n, int h, int d, int w, int k)
 {
-  if(h % d){
+  if (h % d) {
     fprintf(stderr, "d must divide h without remainder!\n");
     return 1;
   }
@@ -158,29 +158,27 @@ int xmssmt_set_params(xmssmt_params *params, int m, int n, int h, int d, int w,
  * Computes a leaf from a WOTS public key using an L-tree.
  */
 static void l_tree(unsigned char *leaf, unsigned char *wots_pk, const xmss_params *params, const unsigned char *pub_seed, unsigned char addr[16])
-{ 
+{
   unsigned int l = params->wots_par.len;
   unsigned int n = params->n;
   unsigned long i = 0;
   unsigned int height = 0;
-  
+
   //ADRS.setTreeHeight(0);
-  SET_LTREE_TREE_HEIGHT(addr,height);
+  SET_LTREE_TREE_HEIGHT(addr, height);
   unsigned long bound;
-  while ( l > 1 ) 
-  {
-     bound = l >> 1; //floor(l / 2); 
-     for ( i = 0; i < bound; i = i + 1 ) {
+  while (l > 1) {
+     bound = l >> 1; //floor(l / 2);
+     for (i = 0; i < bound; i = i + 1) {
        //ADRS.setTreeIndex(i);
-       SET_LTREE_TREE_INDEX(addr,i);
+       SET_LTREE_TREE_INDEX(addr, i);
        //wots_pk[i] = RAND_HASH(pk[2i], pk[2i + 1], SEED, ADRS);
-       hash_2n_n(wots_pk+i*n,wots_pk+i*2*n, pub_seed, addr, n);
+       hash_2n_n(wots_pk+i*n, wots_pk+i*2*n, pub_seed, addr, n);
      }
      //if ( l % 2 == 1 ) {
-     if(l&1)
-     {
+     if (l & 1) {
        //pk[floor(l / 2) + 1] = pk[l];
-       memcpy(wots_pk+(l>>1)*n,wots_pk+(l-1)*n, n);
+       memcpy(wots_pk+(l>>1)*n, wots_pk+(l-1)*n, n);
        //l = ceil(l / 2);
        l=(l>>1)+1;
      }
@@ -188,13 +186,13 @@ static void l_tree(unsigned char *leaf, unsigned char *wots_pk, const xmss_param
      {
        //l = ceil(l / 2);
        l=(l>>1);
-     }     
+     }
      //ADRS.setTreeHeight(ADRS.getTreeHeight() + 1);
      height++;
-     SET_LTREE_TREE_HEIGHT(addr,height);
+     SET_LTREE_TREE_HEIGHT(addr, height);
    }
    //return pk[0];
-   memcpy(leaf,wots_pk,n);
+   memcpy(leaf, wots_pk, n);
 }
 
 /**
@@ -208,11 +206,11 @@ static void gen_leaf_wots(unsigned char *leaf, const unsigned char *sk_seed, con
   get_seed(seed, sk_seed, params->n, ots_addr);
   wots_pkgen(pk, seed, &(params->wots_par), pub_seed, ots_addr);
 
-  l_tree(leaf, pk, params, pub_seed, ltree_addr); 
+  l_tree(leaf, pk, params, pub_seed, ltree_addr);
 }
 
 static int treehash_minheight_on_stack(bds_state* state, const xmss_params *params, const treehash_inst *treehash) {
-  int r = params->h, i;
+  unsigned int r = params->h, i;
   for (i = 0; i < treehash->stackusage; i++) {
     if (state->stacklevels[state->stackoffset - i - 1] < r) {
       r = state->stacklevels[state->stackoffset - i - 1];
@@ -224,7 +222,7 @@ static int treehash_minheight_on_stack(bds_state* state, const xmss_params *para
 /**
  * Merkle's TreeHash algorithm. The address only needs to initialize the first 78 bits of addr. Everything else will be set by treehash.
  * Currently only used for key generation.
- * 
+ *
  */
 static void treehash_setup(unsigned char *node, int height, int index, bds_state *state, const unsigned char *sk_seed, const xmss_params *params, const unsigned char *pub_seed, const unsigned char addr[16])
 {
@@ -244,34 +242,32 @@ static void treehash_setup(unsigned char *node, int height, int index, bds_state
   memcpy(node_addr, ltree_addr, 10);
   SET_LTREE_BIT(node_addr, 0);
   SET_NODE_PADDING(node_addr);
-  
-  int lastnode,i;
+
+  unsigned int lastnode, i;
   unsigned char stack[(height+1)*n];
-  unsigned int  stacklevels[height+1];
-  unsigned int  stackoffset=0;
-  
-  int nodeh;
+  unsigned int stacklevels[height+1];
+  unsigned int stackoffset=0;
+  unsigned int nodeh;
 
   lastnode = idx+(1<<height);
 
-  for(i = 0; i < h-k; i++) {
+  for (i = 0; i < h-k; i++) {
     state->treehash[i].h = i;
     state->treehash[i].completed = 1;
     state->treehash[i].stackusage = 0;
   }
 
   i = 0;
-  for(;idx<lastnode;idx++) 
-  {
-    SET_LTREE_ADDRESS(ltree_addr,idx);
-    SET_OTS_ADDRESS(ots_addr,idx);
-    gen_leaf_wots(stack+stackoffset*n,sk_seed,params, pub_seed, ltree_addr, ots_addr);
+  for (; idx < lastnode; idx++) {
+    SET_LTREE_ADDRESS(ltree_addr, idx);
+    SET_OTS_ADDRESS(ots_addr, idx);
+    gen_leaf_wots(stack+stackoffset*n, sk_seed, params, pub_seed, ltree_addr, ots_addr);
     stacklevels[stackoffset] = 0;
     stackoffset++;
     if (h - k > 0 && i == 3) {
       memcpy(state->treehash[0].node, stack+stackoffset*n, n);
     }
-    while(stackoffset>1 && stacklevels[stackoffset-1] == stacklevels[stackoffset-2])
+    while (stackoffset>1 && stacklevels[stackoffset-1] == stacklevels[stackoffset-2])
     {
       nodeh = stacklevels[stackoffset-1];
       if (i >> nodeh == 1) {
@@ -285,9 +281,9 @@ static void treehash_setup(unsigned char *node, int height, int index, bds_state
           memcpy(state->retain + ((1 << (h - 1 - nodeh)) + nodeh - h + (((i >> nodeh) - 3) >> 1)) * n, stack+(stackoffset-1)*n, n);
         }
       }
-      SET_NODE_TREE_HEIGHT(node_addr,stacklevels[stackoffset-1]);
+      SET_NODE_TREE_HEIGHT(node_addr, stacklevels[stackoffset-1]);
       SET_NODE_TREE_INDEX(node_addr, (idx >> (stacklevels[stackoffset-1]+1)));
-      hash_2n_n(stack+(stackoffset-2)*n,stack+(stackoffset-2)*n, pub_seed,
+      hash_2n_n(stack+(stackoffset-2)*n, stack+(stackoffset-2)*n, pub_seed,
           node_addr, n);
       stacklevels[stackoffset-2]++;
       stackoffset--;
@@ -295,7 +291,7 @@ static void treehash_setup(unsigned char *node, int height, int index, bds_state
     i++;
   }
 
-  for(i=0;i<n;i++)
+  for (i = 0; i < n; i++)
     node[i] = stack[i];
 }
 
@@ -350,43 +346,38 @@ static void treehash_update(treehash_inst *treehash, bds_state *state, const uns
 static void validate_authpath(unsigned char *root, const unsigned char *leaf, unsigned long leafidx, const unsigned char *authpath, const xmss_params *params, const unsigned char *pub_seed, unsigned char addr[16])
 {
   unsigned int n = params->n;
-  
-  int i,j;
+
+  unsigned int i, j;
   unsigned char buffer[2*n];
 
   // If leafidx is odd (last bit = 1), current path element is a right child and authpath has to go to the left.
   // Otherwise, it is the other way around
-  if(leafidx&1)
-  {
-    for(j=0;j<n;j++)
+  if (leafidx & 1) {
+    for (j = 0; j < n; j++)
       buffer[n+j] = leaf[j];
-    for(j=0;j<n;j++)
+    for (j = 0; j < n; j++)
       buffer[j] = authpath[j];
   }
-  else
-  {
-    for(j=0;j<n;j++)
+  else {
+    for (j = 0; j < n; j++)
       buffer[j] = leaf[j];
-    for(j=0;j<n;j++)
+    for (j = 0; j < n; j++)
       buffer[n+j] = authpath[j];
   }
   authpath += n;
 
-  for(i=0;i<params->h-1;i++)
-  {
-    SET_NODE_TREE_HEIGHT(addr,i);
+  for (i = 0; i < params->h-1; i++) {
+    SET_NODE_TREE_HEIGHT(addr, i);
     leafidx >>= 1;
     SET_NODE_TREE_INDEX(addr, leafidx);
-    if(leafidx&1)
-    {
-      hash_2n_n(buffer+n,buffer,pub_seed, addr, n);
-      for(j=0;j<n;j++)
+    if (leafidx & 1) {
+      hash_2n_n(buffer+n, buffer, pub_seed, addr, n);
+      for (j = 0; j < n; j++)
         buffer[j] = authpath[j];
     }
-    else
-    {
-      hash_2n_n(buffer,buffer,pub_seed, addr, n);
-      for(j=0;j<n;j++)
+    else {
+      hash_2n_n(buffer, buffer, pub_seed, addr, n);
+      for (j = 0; j < n; j++)
         buffer[j+n] = authpath[j];
     }
     authpath += n;
@@ -394,41 +385,45 @@ static void validate_authpath(unsigned char *root, const unsigned char *leaf, un
   SET_NODE_TREE_HEIGHT(addr, (params->h-1));
   leafidx >>= 1;
   SET_NODE_TREE_INDEX(addr, leafidx);
-  hash_2n_n(root,buffer,pub_seed,addr,n);
+  hash_2n_n(root, buffer, pub_seed, addr, n);
 }
 
 /**
  * Performs one treehash update on the instance that needs it the most.
  * Returns 1 if such an instance was not found
  **/
-static char bds_treehash_update(bds_state *state, const unsigned char *sk_seed, const xmss_params *params, unsigned char *pub_seed, const unsigned char addr[16]) {
-  unsigned int i;
-  int level, l_min, low;
-  int h = params->h;
-  int k = params->k;
-
-  l_min = h;
-  level = h - k;
-  for (i = 0; i < h - k; i++) {
-    if (state->treehash[i].completed) {
-      low = h;
-    }
-    else if (state->treehash[i].stackusage == 0) {
-      low = i;
-    }
-    else {
-      low = treehash_minheight_on_stack(state, params, &(state->treehash[i]));
+static char bds_treehash_update(bds_state *state, unsigned int updates, const unsigned char *sk_seed, const xmss_params *params, unsigned char *pub_seed, const unsigned char addr[16]) {
+  unsigned int i, j;
+  unsigned int level, l_min, low;
+  unsigned int h = params->h;
+  unsigned int k = params->k;
+  unsigned int used = 0;
+
+  for (j = 0; j < updates; j++) {
+    l_min = h;
+    level = h - k;
+    for (i = 0; i < h - k; i++) {
+      if (state->treehash[i].completed) {
+        low = h;
+      }
+      else if (state->treehash[i].stackusage == 0) {
+        low = i;
+      }
+      else {
+        low = treehash_minheight_on_stack(state, params, &(state->treehash[i]));
+      }
+      if (low < l_min) {
+        level = i;
+        l_min = low;
+      }
     }
-    if (low < l_min) {
-      level = i;
-      l_min = low;
+    if (level == h - k) {
+      break;
     }
+    treehash_update(&(state->treehash[level]), state, sk_seed, params, pub_seed, addr);
+    used++;
   }
-  if (level == h - k) {
-    return 1;
-  }
-  treehash_update(&(state->treehash[level]), state, sk_seed, params, pub_seed, addr);
-  return 0;
+  return updates - used;
 }
 
 /**
@@ -464,15 +459,14 @@ static char bds_state_update(bds_state *state, const unsigned char *sk_seed, con
   SET_OTS_BIT(node_addr, 0);
   SET_NODE_PADDING(node_addr);
 
-  gen_leaf_wots(state->stack+state->stackoffset*n,sk_seed,params, pub_seed, ltree_addr, ots_addr);
+  gen_leaf_wots(state->stack+state->stackoffset*n, sk_seed, params, pub_seed, ltree_addr, ots_addr);
 
   state->stacklevels[state->stackoffset] = 0;
   state->stackoffset++;
   if (h - k > 0 && idx == 3) {
     memcpy(state->treehash[0].node, state->stack+state->stackoffset*n, n);
   }
-  while(state->stackoffset>1 && state->stacklevels[state->stackoffset-1] == state->stacklevels[state->stackoffset-2])
-  {
+  while (state->stackoffset>1 && state->stacklevels[state->stackoffset-1] == state->stacklevels[state->stackoffset-2]) {
     nodeh = state->stacklevels[state->stackoffset-1];
     if (idx >> nodeh == 1) {
       memcpy(state->auth + nodeh*n, state->stack+(state->stackoffset-1)*n, n);
@@ -487,7 +481,7 @@ static char bds_state_update(bds_state *state, const unsigned char *sk_seed, con
     }
     SET_NODE_TREE_HEIGHT(node_addr, state->stacklevels[state->stackoffset-1]);
     SET_NODE_TREE_INDEX(node_addr, (idx >> (state->stacklevels[state->stackoffset-1]+1)));
-    hash_2n_n(state->stack+(state->stackoffset-2)*n,state->stack+(state->stackoffset-2)*n, pub_seed, node_addr, n);
+    hash_2n_n(state->stack+(state->stackoffset-2)*n, state->stack+(state->stackoffset-2)*n, pub_seed, node_addr, n);
 
     state->stacklevels[state->stackoffset-2]++;
     state->stackoffset--;
@@ -501,22 +495,22 @@ static char bds_state_update(bds_state *state, const unsigned char *sk_seed, con
  * next leaf node, using the algorithm described by Buchmann, Dahmen and Szydlo
  * in "Post Quantum Cryptography", Springer 2009.
  */
-static int bds_round(bds_state *state, const unsigned long leaf_idx, const unsigned char *sk_seed, const xmss_params *params, const int updates, unsigned char *pub_seed, unsigned char addr[16])
+static void bds_round(bds_state *state, const unsigned long leaf_idx, const unsigned char *sk_seed, const xmss_params *params, unsigned char *pub_seed, unsigned char addr[16])
 {
   unsigned int i;
-  int n = params->n;
-  int h = params->h;
-  int k = params->k;
+  unsigned int n = params->n;
+  unsigned int h = params->h;
+  unsigned int k = params->k;
 
-  int tau = h;
-  int startidx;
-  int offset, rowidx;
+  unsigned int tau = h;
+  unsigned int startidx;
+  unsigned int offset, rowidx;
   unsigned char buf[2 * n];
 
   unsigned char ots_addr[16];
   unsigned char ltree_addr[16];
   unsigned char node_addr[16];
-  
+
   memcpy(ots_addr, addr, 10);
   SET_OTS_BIT(ots_addr, 1);
   memcpy(ltree_addr, addr, 10);
@@ -542,8 +536,8 @@ static int bds_round(bds_state *state, const unsigned long leaf_idx, const unsig
     memcpy(state->keep + (tau >> 1)*n, state->auth + tau*n, n);
   }
   if (tau == 0) {
-    SET_LTREE_ADDRESS(ltree_addr,leaf_idx);
-    SET_OTS_ADDRESS(ots_addr,leaf_idx);
+    SET_LTREE_ADDRESS(ltree_addr, leaf_idx);
+    SET_OTS_ADDRESS(ots_addr, leaf_idx);
     gen_leaf_wots(state->auth, sk_seed, params, pub_seed, ltree_addr, ots_addr);
   }
   else {
@@ -563,7 +557,7 @@ static int bds_round(bds_state *state, const unsigned long leaf_idx, const unsig
 
     for (i = 0; i < ((tau < h - k) ? tau : (h - k)); i++) {
       startidx = leaf_idx + 1 + 3 * (1 << i);
-      if (startidx < 1 << h) {
+      if (startidx < 1U << h) {
         state->treehash[i].h = i;
         state->treehash[i].next_idx = startidx;
         state->treehash[i].completed = 0;
@@ -571,13 +565,6 @@ static int bds_round(bds_state *state, const unsigned long leaf_idx, const unsig
       }
     }
   }
-  int remaining = 0;
-  for (i = 0; i < updates; i++) {
-    if (bds_treehash_update(state, sk_seed, params, pub_seed, addr)) {
-      remaining++;
-    }
-  }
-  return remaining;
 }
 
 /*
@@ -595,11 +582,11 @@ int xmss_keypair(unsigned char *pk, unsigned char *sk, bds_state *state, xmss_pa
   sk[2] = 0;
   sk[3] = 0;
   // Init SK_SEED (n byte), SK_PRF (m byte), and PUB_SEED (n byte)
-  randombytes(sk+4,2*n+m);
+  randombytes(sk+4, 2*n+m);
   // Copy PUB_SEED to public key
-  memcpy(pk+n, sk+4+n+m,n);
+  memcpy(pk+n, sk+4+n+m, n);
 
-  unsigned char addr[16] = {0,0,0,0};
+  unsigned char addr[16] = {0, 0, 0, 0};
   // Compute root
   treehash_setup(pk, params->h, 0, state, sk+4, params, sk+4+n+m, addr);
   return 0;
@@ -607,10 +594,10 @@ int xmss_keypair(unsigned char *pk, unsigned char *sk, bds_state *state, xmss_pa
 
 /**
  * Signs a message.
- * Returns 
+ * Returns
  * 1. an array containing the signature followed by the message AND
  * 2. an updated secret key!
- * 
+ *
  */
 int xmss_sign(unsigned char *sk, bds_state *state, unsigned char *sig_msg, unsigned long long *sig_msg_len, const unsigned char *msg, unsigned long long msglen, const xmss_params *params)
 {
@@ -618,41 +605,41 @@ int xmss_sign(unsigned char *sk, bds_state *state, unsigned char *sig_msg, unsig
   unsigned int n = params->n;
   unsigned int m = params->m;
   unsigned int k = params->k;
-  
+
   // Extract SK
   unsigned long idx = ((unsigned long)sk[0] << 24) | ((unsigned long)sk[1] << 16) | ((unsigned long)sk[2] << 8) | sk[3];
   unsigned char sk_seed[n];
-  memcpy(sk_seed,sk+4,n);
+  memcpy(sk_seed, sk+4, n);
   unsigned char sk_prf[m];
-  memcpy(sk_prf,sk+4+n,m);
+  memcpy(sk_prf, sk+4+n, m);
   unsigned char pub_seed[n];
-  memcpy(pub_seed,sk+4+n+m,n);  
-  
+  memcpy(pub_seed, sk+4+n+m, n);
+
   // Update SK
   sk[0] = ((idx + 1) >> 24) & 255;
   sk[1] = ((idx + 1) >> 16) & 255;
   sk[2] = ((idx + 1) >> 8) & 255;
   sk[3] = (idx + 1) & 255;
-  // -- Secret key for this non-forward-secure version is now updated. 
-  // -- A productive implementation should use a file handle instead and write the updated secret key at this point! 
-  
+  // -- Secret key for this non-forward-secure version is now updated.
+  // -- A productive implementation should use a file handle instead and write the updated secret key at this point!
+
   // Init working params
   unsigned long long i;
   unsigned char R[m];
   unsigned char msg_h[m];
   unsigned char ots_seed[n];
-  unsigned char ots_addr[16] = {0,0,0,0};
-  
+  unsigned char ots_addr[16] = {0, 0, 0, 0};
+
   // ---------------------------------
   // Message Hashing
   // ---------------------------------
-  
-  // Message Hash: 
+
+  // Message Hash:
   // First compute pseudorandom key
-  prf_m(R, msg, msglen, sk_prf, m); 
+  prf_m(R, msg, msglen, sk_prf, m);
   // Then use it for message digest
   hash_m(msg_h, msg, msglen, R, m, m);
-  
+
   // Start collecting signature
   *sig_msg_len = 0;
 
@@ -661,48 +648,49 @@ int xmss_sign(unsigned char *sk, bds_state *state, unsigned char *sig_msg, unsig
   sig_msg[1] = (idx >> 16) & 255;
   sig_msg[2] = (idx >> 8) & 255;
   sig_msg[3] = idx & 255;
-  
+
   sig_msg += 4;
   *sig_msg_len += 4;
-  
+
   // Copy R to signature
-  for(i=0; i<m; i++)
+  for (i = 0; i < m; i++)
     sig_msg[i] = R[i];
 
   sig_msg += m;
   *sig_msg_len += m;
-  
+
   // ----------------------------------
-  // Now we start to "really sign" 
+  // Now we start to "really sign"
   // ----------------------------------
-  
+
   // Prepare Address
-  SET_OTS_BIT(ots_addr,1);
-  SET_OTS_ADDRESS(ots_addr,idx);
-  
+  SET_OTS_BIT(ots_addr, 1);
+  SET_OTS_ADDRESS(ots_addr, idx);
+
   // Compute seed for OTS key pair
   get_seed(ots_seed, sk_seed, n, ots_addr);
-     
+
   // Compute WOTS signature
   wots_sign(sig_msg, msg_h, ots_seed, &(params->wots_par), pub_seed, ots_addr);
-  
+
   sig_msg += params->wots_par.keysize;
   *sig_msg_len += params->wots_par.keysize;
 
   // the auth path was already computed during the previous round
   memcpy(sig_msg, state->auth, h*n);
 
-  if (idx < (1 << h) - 1) {
-    bds_round(state, idx, sk_seed, params, (h - k) >> 1, pub_seed, ots_addr);
+  if (idx < (1U << h) - 1) {
+    bds_round(state, idx, sk_seed, params, pub_seed, ots_addr);
+    bds_treehash_update(state, (h - k) >> 1, sk_seed, params, pub_seed, ots_addr);
   }
 
   sig_msg += params->h*n;
   *sig_msg_len += params->h*n;
-  
-  //Whipe secret elements?  
+
+  //Whipe secret elements?
   //zerobytes(tsk, CRYPTO_SECRETKEYBYTES);
 
-  memcpy(sig_msg,msg,msglen);
+  memcpy(sig_msg, msg, msglen);
   *sig_msg_len += msglen;
 
   return 0;
@@ -715,38 +703,38 @@ int xmss_sign_open(unsigned char *msg, unsigned long long *msglen, const unsigne
 {
   unsigned int n = params->n;
   unsigned int m = params->m;
-    
+
   unsigned long long i, m_len;
   unsigned long idx=0;
   unsigned char wots_pk[params->wots_par.keysize];
   unsigned char pkhash[n];
   unsigned char root[n];
   unsigned char msg_h[m];
-  
+
   unsigned char pub_seed[n];
-  memcpy(pub_seed,pk+n,n);  
-  
+  memcpy(pub_seed, pk+n, n);
+
   // Init addresses
-  unsigned char ots_addr[16] = {0,0,0,0};
+  unsigned char ots_addr[16] = {0, 0, 0, 0};
   unsigned char ltree_addr[16];
   unsigned char node_addr[16];
-  
+
   SET_OTS_BIT(ots_addr, 1);
-  
+
   memcpy(ltree_addr, ots_addr, 10);
   SET_OTS_BIT(ltree_addr, 0);
   SET_LTREE_BIT(ltree_addr, 1);
-  
+
   memcpy(node_addr, ltree_addr, 10);
   SET_LTREE_BIT(node_addr, 0);
-  SET_NODE_PADDING(node_addr);  
-  
+  SET_NODE_PADDING(node_addr);
+
   // Extract index
   idx = ((unsigned long)sig_msg[0] << 24) | ((unsigned long)sig_msg[1] << 16) | ((unsigned long)sig_msg[2] << 8) | sig_msg[3];
-  printf("verify:: idx = %lu\n",idx);
+  printf("verify:: idx = %lu\n", idx);
   sig_msg += 4;
   sig_msg_len -= 4;
-  
+
   // hash message (recall, R is now on pole position at sig_msg
   unsigned long long tmp_sig_len = m+params->wots_par.keysize+params->h*n;
   m_len = sig_msg_len - tmp_sig_len;
@@ -754,43 +742,43 @@ int xmss_sign_open(unsigned char *msg, unsigned long long *msglen, const unsigne
 
   sig_msg += m;
   sig_msg_len -= m;
-  
+
   //-----------------------
   // Verify signature
   //-----------------------
-  
+
   // Prepare Address
-  SET_OTS_ADDRESS(ots_addr,idx);
-  // Check WOTS signature 
+  SET_OTS_ADDRESS(ots_addr, idx);
+  // Check WOTS signature
   wots_pkFromSig(wots_pk, sig_msg, msg_h, &(params->wots_par), pub_seed, ots_addr);
 
   sig_msg += params->wots_par.keysize;
   sig_msg_len -= params->wots_par.keysize;
-  
+
   // Compute Ltree
-  SET_LTREE_ADDRESS(ltree_addr, idx); 
+  SET_LTREE_ADDRESS(ltree_addr, idx);
   l_tree(pkhash, wots_pk, params, pub_seed, ltree_addr);
-  
+
   // Compute root
-  validate_authpath(root, pkhash, idx, sig_msg, params, pub_seed, node_addr);  
+  validate_authpath(root, pkhash, idx, sig_msg, params, pub_seed, node_addr);
 
   sig_msg += params->h*n;
   sig_msg_len -= params->h*n;
-  
-  for(i=0;i<n;i++)
-    if(root[i] != pk[i])
+
+  for (i = 0; i < n; i++)
+    if (root[i] != pk[i])
       goto fail;
-  
+
   *msglen = sig_msg_len;
-  for(i=0;i<*msglen;i++)
+  for (i = 0; i < *msglen; i++)
     msg[i] = sig_msg[i];
 
   return 0;
-  
-  
+
+
 fail:
   *msglen = sig_msg_len;
-  for(i=0;i<*msglen;i++)
+  for (i = 0; i < *msglen; i++)
     msg[i] = 0;
   *msglen = -1;
   return -1;
@@ -808,25 +796,22 @@ int xmssmt_keypair(unsigned char *pk, unsigned char *sk, bds_state *states, unsi
   unsigned int i;
   unsigned char ots_seed[params->n];
   // Set idx = 0
-  for (i = 0; i < params->index_len; i++){
+  for (i = 0; i < params->index_len; i++) {
     sk[i] = 0;
   }
   // Init SK_SEED (n byte), SK_PRF (m byte), and PUB_SEED (n byte)
-  randombytes(sk+params->index_len,2*n+m);
+  randombytes(sk+params->index_len, 2*n+m);
   // Copy PUB_SEED to public key
-  memcpy(pk+n, sk+params->index_len+n+m,n);
+  memcpy(pk+n, sk+params->index_len+n+m, n);
 
   // Set address to point on the single tree on layer d-1
-  unsigned char addr[16] = {0,0,0,0};
-  SET_OTS_BIT(addr,1);
+  unsigned char addr[16] = {0, 0, 0, 0};
+  SET_OTS_BIT(addr, 1);
   SET_TREE_ADDRESS(addr, 0);
   SET_OTS_ADDRESS(addr, 0);
   SET_LAYER_ADDRESS(addr, 0);
   // Set up state and compute wots signatures for all but topmost tree root
   for (i = 0; i < params->d - 1; i++) {
-    SET_LAYER_ADDRESS(addr, i);
-    SET_TREE_ADDRESS(addr, 0);
-    SET_OTS_ADDRESS(addr, 0);
     // Compute seed for OTS key pair
     treehash_setup(pk, params->xmss_par.h, 0, states + i, sk+params->index_len, &(params->xmss_par), pk+n, addr);
     SET_LAYER_ADDRESS(addr, (i+1));
@@ -839,10 +824,10 @@ int xmssmt_keypair(unsigned char *pk, unsigned char *sk, bds_state *states, unsi
 
 /**
  * Signs a message.
- * Returns 
+ * Returns
  * 1. an array containing the signature followed by the message AND
  * 2. an updated secret key!
- * 
+ *
  */
 int xmssmt_sign(unsigned char *sk, bds_state *states, unsigned char *wots_sigs, unsigned char *sig_msg, unsigned long long *sig_msg_len, const unsigned char *msg, unsigned long long msglen, const xmssmt_params *params)
 {
@@ -855,8 +840,8 @@ int xmssmt_sign(unsigned char *sk, bds_state *states, unsigned char *wots_sigs,
   unsigned long long idx_tree;
   unsigned long long idx_leaf;
   unsigned long long i, j;
+  int needswap_upto = -1;
   unsigned int updates;
-  unsigned int first_nonwots;
 
   unsigned char sk_seed[n];
   unsigned char sk_prf[m];
@@ -865,72 +850,73 @@ int xmssmt_sign(unsigned char *sk, bds_state *states, unsigned char *wots_sigs,
   unsigned char R[m];
   unsigned char msg_h[m];
   unsigned char ots_seed[n];
-  unsigned char ots_addr[16] = {0,0,0,0};
+  unsigned char addr[16] = {0, 0, 0, 0};
+  unsigned char ots_addr[16] = {0, 0, 0, 0};
   bds_state tmp;
 
   // Extract SK
   unsigned long long idx = 0;
-  for(i = 0; i < idx_len; i++){
+  for (i = 0; i < idx_len; i++) {
     idx |= ((unsigned long long)sk[i]) << 8*(idx_len - 1 - i);
   }
-  
-  memcpy(sk_seed,sk+idx_len,n);
-  memcpy(sk_prf,sk+idx_len+n,m);
-  memcpy(pub_seed,sk+idx_len+n+m,n);  
-  
+
+  memcpy(sk_seed, sk+idx_len, n);
+  memcpy(sk_prf, sk+idx_len+n, m);
+  memcpy(pub_seed, sk+idx_len+n+m, n);
+
   // Update SK
-  for(i = 0; i < idx_len; i++){
+  for (i = 0; i < idx_len; i++) {
     sk[i] = ((idx + 1) >> 8*(idx_len - 1 - i)) & 255;
   }
-  // -- Secret key for this non-forward-secure version is now updated. 
-  // -- A productive implementation should use a file handle instead and write the updated secret key at this point! 
-  
-  
+  // -- Secret key for this non-forward-secure version is now updated.
+  // -- A productive implementation should use a file handle instead and write the updated secret key at this point!
+
+
   // ---------------------------------
   // Message Hashing
   // ---------------------------------
-  
-  // Message Hash: 
+
+  // Message Hash:
   // First compute pseudorandom key
-  prf_m(R, msg, msglen, sk_prf, m); 
+  prf_m(R, msg, msglen, sk_prf, m);
   // Then use it for message digest
   hash_m(msg_h, msg, msglen, R, m, m);
-  
+
   // Start collecting signature
   *sig_msg_len = 0;
 
   // Copy index to signature
-  for(i = 0; i < idx_len; i++){
+  for (i = 0; i < idx_len; i++) {
     sig_msg[i] = (idx >> 8*(idx_len - 1 - i)) & 255;
   }
-  
+
   sig_msg += idx_len;
   *sig_msg_len += idx_len;
-  
+
   // Copy R to signature
-  for(i=0; i<m; i++)
+  for (i = 0; i < m; i++)
     sig_msg[i] = R[i];
 
   sig_msg += m;
   *sig_msg_len += m;
-  
+
   // ----------------------------------
-  // Now we start to "really sign" 
+  // Now we start to "really sign"
   // ----------------------------------
-  
+
   // Handle lowest layer separately as it is slightly different...
-  
+
   // Prepare Address
-  SET_OTS_BIT(ots_addr,1);
+  SET_OTS_BIT(ots_addr, 1);
   idx_tree = idx >> tree_h;
   idx_leaf = (idx & ((1 << tree_h)-1));
-  SET_LAYER_ADDRESS(ots_addr,0);
+  SET_LAYER_ADDRESS(ots_addr, 0);
   SET_TREE_ADDRESS(ots_addr, idx_tree);
   SET_OTS_ADDRESS(ots_addr, idx_leaf);
-  
+
   // Compute seed for OTS key pair
   get_seed(ots_seed, sk_seed, n, ots_addr);
-     
+
   // Compute WOTS signature
   wots_sign(sig_msg, msg_h, ots_seed, &(params->xmss_par.wots_par), pub_seed, ots_addr);
 
@@ -942,7 +928,7 @@ int xmssmt_sign(unsigned char *sk, bds_state *states, unsigned char *wots_sigs,
   *sig_msg_len += tree_h*n;
 
   // prepare signature of remaining layers
-  for(i = 1; i < params->d; i++){
+  for (i = 1; i < params->d; i++) {
     // put WOTS signature in place
     memcpy(sig_msg, wots_sigs + (i-1)*params->xmss_par.wots_par.keysize, params->xmss_par.wots_par.keysize);
 
@@ -952,82 +938,64 @@ int xmssmt_sign(unsigned char *sk, bds_state *states, unsigned char *wots_sigs,
     // put AUTH nodes in place
     memcpy(sig_msg, states[i].auth, tree_h*n);
     sig_msg += tree_h*n;
-    *sig_msg_len += tree_h*n;   
+    *sig_msg_len += tree_h*n;
   }
 
-  SET_LAYER_ADDRESS(ots_addr, 0);
-  SET_TREE_ADDRESS(ots_addr, (idx_tree + 1));
-
-  updates = tree_h - k;
-
-  // if a NEXT-tree exists within the hypertree
-  if ((1 + idx_tree) * (1 << tree_h) + idx_leaf < (1 << h)) {
-    // mandatory update for NEXT_0 (does not count towards h-k)
-    bds_state_update(&states[params->d], sk_seed, &(params->xmss_par), pub_seed, ots_addr);
-
-    // check if we're at the end of a tree
-    for (i = 0; i < params->d; i++) {
-      if (((idx + 1) & ((1 << ((i+1)*tree_h)) - 1)) == 0) {
-        memcpy(&tmp, states+params->d + i, sizeof(bds_state));
-        memcpy(states+params->d + i, states + i, sizeof(bds_state));
-        memcpy(states + i, &tmp, sizeof(bds_state));
+  updates = (tree_h - k) >> 1;
 
-        SET_TREE_ADDRESS(ots_addr, ((idx + 1) >> ((i+2) * tree_h)));
-        SET_OTS_ADDRESS(ots_addr, (((idx >> ((i+1) * tree_h)) + 1) & ((1 << tree_h)-1)));
-        SET_LAYER_ADDRESS(ots_addr, (i+1));
-
-        get_seed(ots_seed, sk+params->index_len, n, ots_addr);
-        wots_sign(wots_sigs + i*params->xmss_par.wots_par.keysize, states[i].stack, ots_seed, &(params->xmss_par.wots_par), pub_seed, ots_addr);
-
-        states[params->d + i].stackoffset = 0;
-        states[params->d + i].next_leaf = 0;
-
-        updates--; // WOTS-signing counts as one update
+  SET_TREE_ADDRESS(addr, (idx_tree + 1));
+  // mandatory update for NEXT_0 (does not count towards h-k/2) if NEXT_0 exists
+  if ((1 + idx_tree) * (1 << tree_h) + idx_leaf < (1ULL << h)) {
+    bds_state_update(&states[params->d], sk_seed, &(params->xmss_par), pub_seed, addr);
+  }
 
-        // this bds_round is needed to initialise the state, but should not perform updates
-        // note that one should still pass the (reduced) current idx, as bds_round sets up for idx+1
-        bds_round(&states[i+1], ((idx >> ((i+1)*tree_h))) & ((1 << tree_h)-1), sk_seed, &(params->xmss_par), 0, pub_seed, ots_addr);
-        for (j = 0; j < tree_h-k; j++) {
-          states[i].treehash[j].completed = 1;
+  for (i = 0; i < params->d; i++) {
+    // check if we're not at the end of a tree
+    if (! (((idx + 1) & ((1ULL << ((i+1)*tree_h)) - 1)) == 0)) {
+      idx_leaf = (idx >> (tree_h * i)) & ((1 << tree_h)-1);
+      idx_tree = (idx >> (tree_h * (i+1)));
+      SET_LAYER_ADDRESS(addr, i);
+      SET_TREE_ADDRESS(addr, idx_tree);
+      if (i == (unsigned int) (needswap_upto + 1)) {
+        bds_round(&states[i], idx_leaf, sk_seed, &(params->xmss_par), pub_seed, addr);
+      }
+      updates = bds_treehash_update(&states[i], updates, sk_seed, &(params->xmss_par), pub_seed, addr);
+      SET_TREE_ADDRESS(addr, (idx_tree + 1));
+      // if a NEXT-tree exists for this level;
+      if ((1 + idx_tree) * (1 << tree_h) + idx_leaf < (1ULL << (h - tree_h * i))) {
+        if (i > 0 && updates > 0 && states[params->d + i].next_leaf < (1ULL << h)) {
+          bds_state_update(&states[params->d + i], sk_seed, &(params->xmss_par), pub_seed, addr);
+          updates--;
         }
       }
     }
-  }
+    else if (idx < (1ULL << h) - 1) {
+      memcpy(&tmp, states+params->d + i, sizeof(bds_state));
+      memcpy(states+params->d + i, states + i, sizeof(bds_state));
+      memcpy(states + i, &tmp, sizeof(bds_state));
 
-  SET_LAYER_ADDRESS(ots_addr, 0);
-  SET_TREE_ADDRESS(ots_addr, idx_tree);
+      SET_LAYER_ADDRESS(ots_addr, (i+1));
+      SET_TREE_ADDRESS(ots_addr, ((idx + 1) >> ((i+2) * tree_h)));
+      SET_OTS_ADDRESS(ots_addr, (((idx >> ((i+1) * tree_h)) + 1) & ((1 << tree_h)-1)));
 
-  first_nonwots = (tree_h - k) - updates;
+      get_seed(ots_seed, sk+params->index_len, n, ots_addr);
+      wots_sign(wots_sigs + i*params->xmss_par.wots_par.keysize, states[i].stack, ots_seed, &(params->xmss_par.wots_par), pub_seed, ots_addr);
 
-  if (first_nonwots == 0) {
-    updates = bds_round(&states[0], idx_leaf, sk_seed, &(params->xmss_par), (tree_h - k) >> 1, pub_seed, ots_addr);
-  }
+      states[params->d + i].stackoffset = 0;
+      states[params->d + i].next_leaf = 0;
 
-  for (i = 1; updates > 0 && i < params->d; i++) {
-    idx_leaf = (idx_tree & ((1 << tree_h)-1));
-    idx_tree = idx_tree >> tree_h;
-    if (first_nonwots > i) {
-      continue;
-    }
-    SET_LAYER_ADDRESS(ots_addr, i);
-    SET_TREE_ADDRESS(ots_addr, idx_tree);
-    SET_OTS_ADDRESS(ots_addr, idx_leaf);
-    while (updates > 0 && !bds_treehash_update(&states[i], sk_seed, &(params->xmss_par), pub_seed, ots_addr)) {
-      updates--;
-    }
-    SET_TREE_ADDRESS(ots_addr, (idx_tree + 1));
-    // if a NEXT-tree exists for this level;
-    if ((1 + idx_tree) * (1 << tree_h) + idx_leaf < (1 << (h - tree_h * i))) {
-      while (updates > 0 && !bds_state_update(&states[params->d + i], sk_seed, &(params->xmss_par), pub_seed, ots_addr)) {
-        updates--;
+      updates--; // WOTS-signing counts as one update
+      needswap_upto = i;
+      for (j = 0; j < tree_h-k; j++) {
+        states[i].treehash[j].completed = 1;
       }
     }
   }
-  
-  //Whipe secret elements?  
+
+  //Whipe secret elements?
   //zerobytes(tsk, CRYPTO_SECRETKEYBYTES);
 
-  memcpy(sig_msg,msg,msglen);
+  memcpy(sig_msg, msg, msglen);
   *sig_msg_len += msglen;
 
   return 0;
@@ -1040,35 +1008,35 @@ int xmssmt_sign_open(unsigned char *msg, unsigned long long *msglen, const unsig
 {
   unsigned int n = params->n;
   unsigned int m = params->m;
-  
+
   unsigned int tree_h = params->xmss_par.h;
   unsigned int idx_len = params->index_len;
   unsigned long long idx_tree;
   unsigned long long idx_leaf;
-  
+
   unsigned long long i, m_len;
   unsigned long long idx=0;
   unsigned char wots_pk[params->xmss_par.wots_par.keysize];
   unsigned char pkhash[n];
   unsigned char root[n];
   unsigned char msg_h[m];
-  
+
   unsigned char pub_seed[n];
-  memcpy(pub_seed,pk+n,n);  
-  
+  memcpy(pub_seed, pk+n, n);
+
   // Init addresses
-  unsigned char ots_addr[16] = {0,0,0,0};
+  unsigned char ots_addr[16] = {0, 0, 0, 0};
   unsigned char ltree_addr[16];
   unsigned char node_addr[16];
-  
+
   // Extract index
-  for(i = 0; i < idx_len; i++){
+  for (i = 0; i < idx_len; i++) {
     idx |= ((unsigned long long)sig_msg[i]) << (8*(idx_len - 1 - i));
   }
-  printf("verify:: idx = %llu\n",idx);
+  printf("verify:: idx = %llu\n", idx);
   sig_msg += idx_len;
   sig_msg_len -= idx_len;
-  
+
   // hash message (recall, R is now on pole position at sig_msg
   unsigned long long tmp_sig_len = m+ (params->d * params->xmss_par.wots_par.keysize) + (params->h * n);
   m_len = sig_msg_len - tmp_sig_len;
@@ -1076,95 +1044,95 @@ int xmssmt_sign_open(unsigned char *msg, unsigned long long *msglen, const unsig
 
   sig_msg += m;
   sig_msg_len -= m;
-  
+
   //-----------------------
   // Verify signature
   //-----------------------
-  
+
   // Prepare Address
   idx_tree = idx >> tree_h;
   idx_leaf = (idx & ((1 << tree_h)-1));
-  SET_LAYER_ADDRESS(ots_addr,0);
+  SET_LAYER_ADDRESS(ots_addr, 0);
   SET_TREE_ADDRESS(ots_addr, idx_tree);
   SET_OTS_BIT(ots_addr, 1);
-  
+
   memcpy(ltree_addr, ots_addr, 10);
   SET_OTS_BIT(ltree_addr, 0);
   SET_LTREE_BIT(ltree_addr, 1);
-  
+
   memcpy(node_addr, ltree_addr, 10);
   SET_LTREE_BIT(node_addr, 0);
-  SET_NODE_PADDING(node_addr);  
-  
-  SET_OTS_ADDRESS(ots_addr,idx_leaf);
-  
-  // Check WOTS signature 
+  SET_NODE_PADDING(node_addr);
+
+  SET_OTS_ADDRESS(ots_addr, idx_leaf);
+
+  // Check WOTS signature
   wots_pkFromSig(wots_pk, sig_msg, msg_h, &(params->xmss_par.wots_par), pub_seed, ots_addr);
 
   sig_msg += params->xmss_par.wots_par.keysize;
   sig_msg_len -= params->xmss_par.wots_par.keysize;
-  
+
   // Compute Ltree
-  SET_LTREE_ADDRESS(ltree_addr, idx_leaf); 
+  SET_LTREE_ADDRESS(ltree_addr, idx_leaf);
   l_tree(pkhash, wots_pk, &(params->xmss_par), pub_seed, ltree_addr);
-  
+
   // Compute root
-  validate_authpath(root, pkhash, idx_leaf, sig_msg, &(params->xmss_par), pub_seed, node_addr);  
+  validate_authpath(root, pkhash, idx_leaf, sig_msg, &(params->xmss_par), pub_seed, node_addr);
 
   sig_msg += tree_h*n;
   sig_msg_len -= tree_h*n;
-  
-  for(i = 1; i < params->d; i++){
+
+  for (i = 1; i < params->d; i++) {
     // Prepare Address
     idx_leaf = (idx_tree & ((1 << tree_h)-1));
     idx_tree = idx_tree >> tree_h;
-    
-    SET_LAYER_ADDRESS(ots_addr,i);
+
+    SET_LAYER_ADDRESS(ots_addr, i);
     SET_TREE_ADDRESS(ots_addr, idx_tree);
     SET_OTS_BIT(ots_addr, 1);
-    
+
     memcpy(ltree_addr, ots_addr, 10);
     SET_OTS_BIT(ltree_addr, 0);
     SET_LTREE_BIT(ltree_addr, 1);
-    
+
     memcpy(node_addr, ltree_addr, 10);
     SET_LTREE_BIT(node_addr, 0);
-    SET_NODE_PADDING(node_addr);  
-    
-    SET_OTS_ADDRESS(ots_addr,idx_leaf);
-    
-    // Check WOTS signature 
+    SET_NODE_PADDING(node_addr);
+
+    SET_OTS_ADDRESS(ots_addr, idx_leaf);
+
+    // Check WOTS signature
     wots_pkFromSig(wots_pk, sig_msg, root, &(params->xmss_par.wots_par), pub_seed, ots_addr);
 
     sig_msg += params->xmss_par.wots_par.keysize;
     sig_msg_len -= params->xmss_par.wots_par.keysize;
-    
+
     // Compute Ltree
-    SET_LTREE_ADDRESS(ltree_addr, idx_leaf); 
+    SET_LTREE_ADDRESS(ltree_addr, idx_leaf);
     l_tree(pkhash, wots_pk, &(params->xmss_par), pub_seed, ltree_addr);
-    
+
     // Compute root
-    validate_authpath(root, pkhash, idx_leaf, sig_msg, &(params->xmss_par), pub_seed, node_addr);  
+    validate_authpath(root, pkhash, idx_leaf, sig_msg, &(params->xmss_par), pub_seed, node_addr);
 
     sig_msg += tree_h*n;
     sig_msg_len -= tree_h*n;
-    
+
   }
-  
-  for(i=0;i<n;i++)
-    if(root[i] != pk[i])
+
+  for (i = 0; i < n; i++)
+    if (root[i] != pk[i])
       goto fail;
-  
+
   *msglen = sig_msg_len;
-  for(i=0;i<*msglen;i++)
+  for (i = 0; i < *msglen; i++)
     msg[i] = sig_msg[i];
 
   return 0;
-  
-  
+
+
 fail:
   *msglen = sig_msg_len;
-  for(i=0;i<*msglen;i++)
+  for (i = 0; i < *msglen; i++)
     msg[i] = 0;
   *msglen = -1;
   return -1;
diff --git a/xmss_fast.h b/xmss_fast.h
index f5968ca..97d1310 100644
--- a/xmss_fast.h
+++ b/xmss_fast.h
@@ -9,45 +9,45 @@ Public domain.
 #ifndef XMSS_H
 #define XMSS_H
 typedef struct{
-  int level;
+  unsigned int level;
   unsigned long long subtree;
-  int subleaf;
+  unsigned int subleaf;
 } leafaddr;
 
 typedef struct{
   wots_params wots_par;
-  int n;
-  int m;
-  int h;
-  int k;
+  unsigned int n;
+  unsigned int m;
+  unsigned int h;
+  unsigned int k;
 } xmss_params;
 
 typedef struct{
   xmss_params xmss_par;
-  int n;
-  int m;
-  int h;
-  int d;
-  int index_len;
+  unsigned int n;
+  unsigned int m;
+  unsigned int h;
+  unsigned int d;
+  unsigned int index_len;
 } xmssmt_params;
 
 typedef struct{
-  int h;
-  int next_idx;
-  int stackusage;
+  unsigned int h;
+  unsigned int next_idx;
+  unsigned int stackusage;
   unsigned char completed;
   unsigned char *node;
 } treehash_inst;
 
 typedef struct {
   unsigned char *stack;
-  int stackoffset;
+  unsigned int stackoffset;
   unsigned char *stacklevels;
   unsigned char *auth;
   unsigned char *keep;
   treehash_inst *treehash;
   unsigned char *retain;
-  int next_leaf;
+  unsigned int next_leaf;
 } bds_state;
 
 /**
diff --git a/zerobytes.c b/zerobytes.c
index 04056d2..b9944ec 100644
--- a/zerobytes.c
+++ b/zerobytes.c
@@ -6,9 +6,9 @@ Public domain.
 #include "zerobytes.h"
 
 unsigned char *zerobytes(unsigned char *r,unsigned long long n)
-{ 
-  volatile unsigned char *p=r; 
-  while (n--) 
-    *(p++) = 0; 
-  return r; 
+{
+  volatile unsigned char *p=r;
+  while (n--)
+    *(p++) = 0;
+  return r;
 }