// Copyright (c) 2019-2022 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package standalone

import (
	"github.com/decred/dcrd/chaincfg/chainhash"
	"github.com/decred/dcrd/wire"
)

// CalcMerkleRootInPlace is an in-place version of CalcMerkleRoot that reuses
// the backing array of the provided slice to perform the calculation thereby
// preventing extra allocations.  It is the caller's responsibility to ensure it
// is safe to mutate the entries in the provided slice.
//
// The function internally appends an additional entry in the case the number of
// provided leaves is odd, so the caller may wish to pre-allocate space for one
// additional element in the backing array in that case to ensure it doesn't
// need to be reallocated to expand it.
//
// For example:
//
//	allocLen := len(leaves) + len(leaves)&1
//	leaves := make([]chainhash.Hash, len(leaves), allocLen)
//	// populate the leaves
//
// See CalcMerkleRoot for more details on how the merkle root is calculated.
func CalcMerkleRootInPlace(leaves []chainhash.Hash) chainhash.Hash {
	if len(leaves) == 0 {
		// All zero.
		return chainhash.Hash{}
	}

	// Create a buffer to reuse for hashing the branches and some long lived
	// slices into it to avoid reslicing.
	var buf [2 * chainhash.HashSize]byte
	var left = buf[:chainhash.HashSize]
	var right = buf[chainhash.HashSize:]
	var both = buf[:]

	// The following algorithm works by replacing the leftmost entries in the
	// slice with the concatenations of each subsequent set of 2 hashes and
	// shrinking the slice by half to account for the fact that each level of
	// the tree is half the size of the previous one.  In the case a level is
	// unbalanced (there is no final right child), the final node is duplicated
	// so it ultimately is concatenated with itself.
	//
	// For example, the following illustrates calculating a tree with 5 leaves:
	//
	// [0 1 2 3 4]                              (5 entries)
	// 1st iteration: [h(0||1) h(2||3) h(4||4)] (3 entries)
	// 2nd iteration: [h(h01||h23) h(h44||h44)] (2 entries)
	// 3rd iteration: [h(h0123||h4444)]         (1 entry)
	for len(leaves) > 1 {
		// When there is no right child, the parent is generated by hashing the
		// concatenation of the left child with itself.
		if len(leaves)&1 != 0 {
			leaves = append(leaves, leaves[len(leaves)-1])
		}

		// Set the parent node to the hash of the concatenation of the left and
		// right children.
		for i := 0; i < len(leaves)/2; i++ {
			copy(left, leaves[i*2][:])
			copy(right, leaves[i*2+1][:])
			leaves[i] = chainhash.HashH(both)
		}
		leaves = leaves[:len(leaves)/2]
	}
	return leaves[0]
}

// CalcMerkleRoot treats the provided slice of hashes as leaves of a merkle tree
// and returns the resulting merkle root.
//
// A merkle tree is a tree in which every non-leaf node is the hash of its
// children nodes.  A diagram depicting how this works for Decred transactions
// where h(x) is a blake256 hash follows:
//
//	         root = h1234 = h(h12 + h34)
//	        /                           \
//	  h12 = h(h1 + h2)            h34 = h(h3 + h4)
//	   /            \              /            \
//	h1 = h(tx1)  h2 = h(tx2)    h3 = h(tx3)  h4 = h(tx4)
//
// The number of inputs is not always a power of two which results in a
// balanced tree structure as above.  In that case, parent nodes with no
// children are also zero and parent nodes with only a single left node
// are calculated by concatenating the left node with itself before hashing.
func CalcMerkleRoot(leaves []chainhash.Hash) chainhash.Hash {
	if len(leaves) == 0 {
		// All zero.
		return chainhash.Hash{}
	}

	// Copy the leaves so they can be safely mutated by the in-place merkle root
	// calculation.  Note that the backing array is provided with space for one
	// additional item when the number of leaves is odd as an optimization for
	// the in-place calculation to avoid the need grow the backing array.
	allocLen := len(leaves) + len(leaves)&1
	dupLeaves := make([]chainhash.Hash, len(leaves), allocLen)
	copy(dupLeaves, leaves)
	return CalcMerkleRootInPlace(dupLeaves)
}

// CalcTxTreeMerkleRoot calculates and returns the merkle root for the provided
// transactions.  The full (including witness data) hashes for the transactions
// are used as required for merkle roots.
//
// See CalcMerkleRoot for more details on how the merkle root is calculated.
func CalcTxTreeMerkleRoot(transactions []*wire.MsgTx) chainhash.Hash {
	if len(transactions) == 0 {
		// All zero.
		return chainhash.Hash{}
	}

	// Note that the backing array is provided with space for one additional
	// item when the number of leaves is odd as an optimization for the in-place
	// calculation to avoid the need grow the backing array.
	allocLen := len(transactions) + len(transactions)&1
	leaves := make([]chainhash.Hash, 0, allocLen)
	for _, tx := range transactions {
		leaves = append(leaves, tx.TxHashFull())
	}
	return CalcMerkleRootInPlace(leaves)
}

// CalcCombinedTxTreeMerkleRoot calculates and returns the combined merkle root
// for the provided regular and stake transaction trees in accordance with
// DCP0005.
//
// In particular, the final merkle root is the result of a merkle tree that
// itself has the individual merkle roots of the two transaction trees as
// leaves.  The full (including witness data) hashes for the transactions are
// used as required for merkle roots.
//
// A diagram depicting this follows:
//
//	root = blake256(regularTreeRoot || stakeTreeRoot)
//	          /                           \
//	   regularTreeRoot               stakeTreeRoot
//
// It is also worth noting that it also happens to be exactly equivalent to the
// blake256 hash of the concatenation of the two individual merkle roots due to
// the way two leaf merkle trees are calculated:
//
//	blake256(regularTreeRoot || stakeTreeRoot)
func CalcCombinedTxTreeMerkleRoot(regularTxns, stakeTxns []*wire.MsgTx) chainhash.Hash {
	regularRoot := CalcTxTreeMerkleRoot(regularTxns)
	stakeRoot := CalcTxTreeMerkleRoot(stakeTxns)

	// Return the hash of the concatenation of the individual merkle roots.
	//
	// This is a slightly faster equivalent of calling CalcMerkleRoot with the
	// two individual merkle roots.
	var both [chainhash.HashSize * 2]byte
	copy(both[:chainhash.HashSize], regularRoot[:])
	copy(both[chainhash.HashSize:], stakeRoot[:])
	return chainhash.HashH(both[:])
}
