module dcrypt.pqc.sphincs.treeutil;

import dcrypt.pqc.sphincs.common: num_digits, is_hash_2n_n, hash_2n_n_mask;


package template TreeUtil(alias hash_2n_n, H, M)
if (is_hash_2n_n!(hash_2n_n, H) && 2*H.length == M.length) {

	alias hash_2n_n_mask!(hash_2n_n, H, M) hash_nodes;

	/// Compute the root hash of a tree given a leaf and its authentication path.
	/// 
	/// Params:
	/// leaf	=	A leaf of the tree.
	/// leafidx	=	The index of the leaf.
	/// authpath	=	Authentication path as generated by `gen_subtree_authpath()`.
	/// masks	=	Bitmasks for the tree.
	@safe @nogc pure nothrow
	public H validate_authpath(in ref H leaf, in uint leafidx, in H[] authpath, in M[] masks)
	in {
		assert(masks.length >= authpath.length, "Got to few bitmasks.");
	} body {

		uint idx = leafidx;
		H p = leaf;
		
		foreach(i;0..authpath.length) {
			if(idx % 2 == 0) {
				p = hash_nodes(p, authpath[i], masks[i]);
			} else {
				p = hash_nodes(authpath[i], p, masks[i]);
			}
			idx >>= 1;
		}
		
		return p;
	}

	/// Calculate root hash of a L-tree or a binary tree.
	/// The number of leaves is not required to be a power of 2.
	/// 
	/// Params:
	/// len	=	Number of leaves. Must be larger than 0.
	/// leaves	=	The leaves of the tree. Exactly l.
	/// masks	=	One bitmask for each layer of the tree.
	@safe @nogc pure nothrow
	public H hash_tree(uint len)(in H[] leaves, in M[] masks) if(len > 0)
	in {
		assert(leaves.length == len);
		assert(1<<masks.length >= leaves.length, "Not enough bitmasks.");
	} body {
		enum height = num_digits(len-1, 2); // Get minimal height of the tree given the number of leaves.
		static assert(2*len > (1<<height), "Tree height is too large.");

		hash_stack!height stack;
		
		/// The algorithm in a nutshell:
		/// After pushing a leaf, reduce the stack size by merging the top two
		/// elements as long as they belong to the same level in the tree.
		/// The number of trailing zeros of the current leaf index +1 tells us how
		/// many times we can merge the top two stack elements.

		size_t i;
		uint maskLevel;
		for(i = 0; i < len; ++i) {

			stack.push(leaves[i]);
			
			//			if(addr.subleaf == laddr.subleaf) {
			//				// That's the leaf we want to generate the authpath for.
			//				stack.start_authpath();
			//			}
			
			auto zeromap = i+1; // Number of trailing zeros tells us how many times to call stack.reduce().

			maskLevel = 0;
			while((zeromap & 1) == 0) {
				stack.reduce(masks[maskLevel]);
				
				++maskLevel;
				zeromap >>= 1;
			}
			
			//++addr.subleaf;
		}

		/// If the tree is a L-tree (number of leaves not a power of 2),
		/// then there is still something to do.
		static if(len < (1<<height)) {
			static assert(len > 0, "Hash tree is not defined for 0 leaves.");
			i = i-1;
			i >>= maskLevel;
			for(; maskLevel < height; ++maskLevel) {
				if(i & 1) {
					stack.reduce(masks[maskLevel]);
				}
				i >>= 1;
			}
		}
		
		H root = stack.pop();
		assert(stack.empty);
		
		return root;
	}

	/// L-tree hash sanity test.
	private unittest {
		import dcrypt.pqc.sphincs.sphincs256: hash_2n_n;

		enum l = 3;
		enum hash_bytes = 32;
		alias ubyte[hash_bytes] hash256;
		alias TreeUtil!(hash_2n_n, ubyte[hash_bytes], ubyte[2*hash_bytes]) Tree;
		hash256[l] leaves;
		for(uint i = 0; i < l; ++i) { leaves[i][] = cast(ubyte) (i+1); } // Make leaves distinct.
		
		ubyte[64][2] masks;
		masks[0][] = 1;
		masks[1][] = 2;
		
		hash256 root = Tree.hash_tree!l(leaves, masks);
		
		leaves[0][] ^= masks[0][0..hash_bytes];
		leaves[1][] ^= masks[0][hash_bytes..$];
		
		hash256 root2 = hash_2n_n(leaves[0], leaves[1]);
		
		root2[] ^= masks[1][0..hash_bytes];
		leaves[2][] ^= masks[1][hash_bytes..$];
		
		root2 = hash_2n_n(root2, leaves[2]);
		assert(root == root2);
	}

	/// Test hash_ltree against result of reference implementation (l_tree()).
	private unittest {

		import dcrypt.pqc.sphincs.sphincs256: hash_2n_n, hash256;

		enum l = 67;
		enum hash_bytes = 32;
		alias TreeUtil!(hash_2n_n, ubyte[hash_bytes], ubyte[2*hash_bytes]) Tree;
		
		hash256[l] leaves;
		for(uint i = 0; i < leaves.length; ++i) {
			leaves[i] = cast(ubyte) i;
		}
		
		ubyte[2*hash_bytes][7] masks;
		
		for(uint i = 0; i < masks.length; ++i) { 
			masks[i][0..hash_bytes] = cast(ubyte) (1+2*i);
			masks[i][hash_bytes..$] = cast(ubyte) (2+2*i);
		}
		
		hash256 root = Tree.hash_tree!l(leaves, masks);
		
		assert(root == x"59641ed4970735d4e1d84ec00e4780d1ab211ebd9339b9962de2a15ead43e1e4", "hash_ltree() failed.");
	}

	/// Helper struct for treehash algorithm.
	struct hash_stack(uint height)
	{
		
		@safe @nogc:
		
		private {
			H[height+1] stack;
			uint stackptr = 0;
			int authpath_marker = -1;	/// Points to a element belonging to the authpath. -1 means there is no such element.

			H[height] authpath;
			uint authpath_ptr = 0;
		}
		
		invariant {
			assert(stackptr <= stack.length, "Stack grew higher than allowed.");
			assert(authpath_marker >= -1 && authpath_marker < cast(int) stack.length);
			
			// TODO: Why does this fail?
			//assert(authpath_marker >= -1 && authpath_marker < stack.length);
			
		}
		
		@property
		bool empty() nothrow {
			return stackptr == 0;
		}

		void push(in ref H h) nothrow {
			stack[stackptr] = h;
			++stackptr;
		}
		
		/// Start creating the authpath for the node on the top.
		void start_authpath() nothrow {
			assert(!empty, "Can't start creating authpath in empty stack.");
			assert(authpath_marker == -1, "Authpath already started.");
			
			authpath_marker = stackptr-1;
		}
		
		H pop() nothrow {
			assert(!empty, "Stack is empty can't pop().");
			--stackptr;
			return stack[stackptr];
		}
		
		H[height] get_authpath() nothrow {
			assert(authpath_ptr == authpath.length, "Authpath is not yet constructed.");
			return authpath;
		}
		
		/// Merge the top two hashes into one.
		/// 
		/// Params:
		/// mask_lower	=	Bitmask for lower hash.
		/// mask_top	=	Bitmask for the hash on the top.
		void reduce(in ref M mask) nothrow {
			assert(stackptr >= 2, "Less than two hashes on the stack. Can't merge.");
			
			--stackptr;
			
			// Check if one of the top two nodes belongs to the authpath.
			// If yes, add it to the authpath.
			if(authpath_marker == stackptr) {
				authpath[authpath_ptr] = stack[stackptr-1];
				++authpath_ptr;
				--authpath_marker;
			} else if(authpath_marker == stackptr-1) {
				authpath[authpath_ptr] = stack[stackptr-0];
				++authpath_ptr;
			}
			
			stack[stackptr-1] = hash_nodes(stack[stackptr-1], stack[stackptr-0], mask);
		}
	}

	/// Sanity test for hash_stack.
	private unittest {
		import dcrypt.pqc.sphincs.sphincs256: hash_2n_n;

		enum hash_bytes = 32;
		alias ubyte[hash_bytes] hash_t;
		alias TreeUtil!(hash_2n_n, ubyte[hash_bytes], ubyte[2*hash_bytes]) Tree;
		Tree.hash_stack!4 stack;

		hash_t mask1 = 111;
		hash_t mask2 = 222;
		
		ubyte[2*hash_bytes] mask = mask1~mask2;
		
		hash_t[4] l;
		for(uint i = 0; i < l.length; ++i) l[i] = cast(ubyte) (i+1);
		
		assert(stack.empty);
		stack.push(l[0]);
		stack.push(l[1]);
		assert(l[1] == stack.pop());
		assert(l[0] == stack.pop());
		
		stack.push(l[0]);
		stack.push(l[1]);
		
		stack.reduce(mask);
		hash_t hash1 = stack.pop();
		assert(stack.empty);
		
		hash_t hash2 = hash_nodes(l[0], l[1], mask);
		assert(hash1 == hash2);
	}
}