Merkle.hs

{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}

module Tahoe.CHK.Merkle (
    MerkleTree (MerkleNode, MerkleLeaf),
    Direction (..),
    leaf,
    leafNumberToNodeNumber,
    breadthFirstList,
    merklePathLengthForSize,
    makeTree,
    makeTreePartial,
    merkleProof,
    neededHashes,
    firstLeafNum,
    rootHash,
    pairHash,
    emptyLeafHash,
    size,
    height,
    mapTree,
    merklePath,
    leafHashes,
    -- exported for testing in ghci
    treeFromRows,
    buildTreeOutOfAllTheNodes,
) where

import Data.Binary (Binary (get, put))
import Data.Binary.Get (getRemainingLazyByteString)
import Data.Binary.Put (putByteString)
import Data.TreeDiff.Class (ToExpr)
import GHC.Generics (Generic)

import Data.List.HT (
    padLeft,
 )
import Data.Tuple.HT (
    mapFst,
 )

import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as LBS

import Data.Text (
    pack,
 )
import qualified Data.Text as T
import Data.Text.Encoding (
    encodeUtf8,
 )

import Data.ByteString.Base32 (
    encodeBase32Unpadded,
 )

import Tahoe.CHK.Crypto (
    taggedHash,
    taggedPairHash,
 )

import Crypto.Hash (HashAlgorithm (hashDigestSize))
import Crypto.Hash.Algorithms (SHA256 (SHA256))
import Tahoe.Util (
    chunkedBy,
    nextPowerOf,
    toBinary,
 )

data MerkleTree
    = MerkleLeaf B.ByteString
    | MerkleNode B.ByteString MerkleTree MerkleTree
    deriving (Eq, Ord, Generic, ToExpr)

{- | A constructor for a MerkleLeaf that enforces correct byte string length
 (error on incorrect length).
-}
leaf :: B.ByteString -> MerkleTree
leaf bs
    | B.length bs == 32 = MerkleLeaf bs
    | otherwise = error $ "Constructed MerkleLeaf with hash of length " <> show (B.length bs)

-- | Count the number of nodes in a tree.
size :: MerkleTree -> Int
size = sum . mapTree (const 1)

-- | Measure the height of a tree.
height :: MerkleTree -> Int
height (MerkleLeaf _) = 1
height (MerkleNode _ left _) = 1 + height left

mapTree :: (MerkleTree -> a) -> MerkleTree -> [a]
mapTree f l@(MerkleLeaf _) = [f l]
mapTree f n@(MerkleNode _ left right) = f n : mapTree f left ++ mapTree f right

instance Show MerkleTree where
    show (MerkleLeaf value) =
        T.unpack $ T.concat ["MerkleLeaf ", encodeBase32Unpadded value]
    show (MerkleNode value left right) =
        T.unpack $
            T.concat
                [ "MerkleNode " :: T.Text
                , encodeBase32Unpadded value
                , " ("
                , T.pack $ show left
                , ")"
                , " ("
                , T.pack $ show right
                , ")"
                ]

emptyLeafHash :: Int -> B.ByteString
emptyLeafHash = taggedHash (hashDigestSize SHA256) "Merkle tree empty leaf" . encodeUtf8 . pack . show

pairHash :: B.ByteString -> B.ByteString -> B.ByteString
pairHash = taggedPairHash (hashDigestSize SHA256) "Merkle tree internal node"

rootHash :: MerkleTree -> B.ByteString
rootHash (MerkleLeaf value) = value
rootHash (MerkleNode value _ _) = value

-- Like makeTree but error on empty list
makeTreePartial :: [B.ByteString] -> MerkleTree
makeTreePartial = unJust . makeTree
  where
    unJust Nothing = error "Merkle.makeTreePartial failed to make a tree"
    unJust (Just t) = t

-- Make a merkle tree for the given values.  Extra values are generated to
-- fill the tree if necessary.  The given values are the values of the leaf
-- nodes.
makeTree :: [B.ByteString] -> Maybe MerkleTree
makeTree [] = Nothing
makeTree leaves =
    Just $ makeTree' (pad leaves)
  where
    -- Pad the leaves out to the next power of two so the tree is full.
    pad :: [B.ByteString] -> [B.ByteString]
    pad leaves' = leaves' ++ padding (length leaves')

    -- Create the padding for the pad function.  The number of leaves in the
    -- tree must be a power of 2 (a height zero tree has 2 ^ 0 leaves, a
    -- height one tree has 2 ^ 1 leaves, etc) so compute a number of empty
    -- leaves that when added to the non-empty leaves gives us a power of 2.
    -- This could be none if we happened to already have a number of leaves
    -- that is a power of 2.
    --
    -- This function assumes that the number of non-empty leaves is at least
    -- half the number of total leaves.  If it is fewer it will create less
    -- padding than necessary.  This should be reasonable since if there fewer
    -- leaves then a smaller tree could hold them all.
    padding :: Int -> [B.ByteString]
    padding numLeaves = emptyLeafHash <$> [numLeaves .. nextPowerOf 2 numLeaves - 1]

    -- Turn a length-of-power-of-2 list into a tree
    makeTree' :: [B.ByteString] -> MerkleTree
    makeTree' [x] = leaf x
    makeTree' xs =
        makeNode (makeTree' left) (makeTree' right)
      where
        (left, right) = splitAt (length xs `div` 2) xs

    -- Make a parent node referencing two given child nodes, calculating the
    -- parent node's hash in the process.
    makeNode :: MerkleTree -> MerkleTree -> MerkleTree
    makeNode left right =
        MerkleNode (pairHash (rootHash left) (rootHash right)) left right

-- | Represent a direction to take when walking down a binary tree.
data Direction = TurnLeft | TurnRight deriving (Show, Ord, Eq)

{- | Return a list of tuples of node numbers and corresponding merkle hashes.
 The node numbers correspond to a numbering of the nodes in the tree where the
 root node is numbered 1, each node's left child is the node's number times
 two, and the node's right child is the node's number times two plus one.
-}
merkleProof :: MerkleTree -> Int -> Maybe [(Int, B.ByteString)]
merkleProof tree targetLeaf = merkleProof' 1 tree $ merklePath (height tree) targetLeaf

{- | Compute the path to a leaf from the root of a merkle tree of a certain
 height.
-}
merklePath :: Int -> Int -> [Direction]
merklePath height' leafNum = padLeft TurnLeft (height' - 1) (toBinary TurnLeft TurnRight leafNum)

-- | Compute the length of a merkle path through a tree of the given height.
merklePathLengthForSize :: Int -> Int
merklePathLengthForSize size' = ceiling . logBase (2 :: Double) . fromIntegral $ nextPowerOf 2 size'

-- Convert a tree to a breadth-first list of its hash values.
breadthFirstList :: MerkleTree -> [B.ByteString]
breadthFirstList tree = traverse' [tree]
  where
    traverse' :: [MerkleTree] -> [B.ByteString]
    traverse' [] = []
    traverse' trees =
        [rootHash tree' | tree' <- trees] ++ traverse' (concat [children tree'' | tree'' <- trees])

    children (MerkleLeaf _) = []
    children (MerkleNode _ left right) = [left, right]

{- | Construct Just a merkle proof along the pre-computed path or Nothing if
 the path runs past the leaves of the tree.
-}
merkleProof' :: Int -> MerkleTree -> [Direction] -> Maybe [(Int, B.ByteString)]
merkleProof' _ _ [] = Just []
merkleProof' thisNodeNum (MerkleNode _ left right) (d : ds) =
    case d of
        TurnLeft ->
            ((rightChildNum, rootHash right) :) <$> merkleProof' leftChildNum left ds
        TurnRight ->
            ((leftChildNum, rootHash left) :) <$> merkleProof' rightChildNum right ds
  where
    leftChildNum = thisNodeNum * 2
    rightChildNum = thisNodeNum * 2 + 1
merkleProof' _ (MerkleLeaf _) ds = error $ show ds

{- | Translate a leaf number to a node number.  Leaf numbers are zero indexed
 and identify leaves of a tree from left to right.  Node numbers are one
 indexed and identify nodes of a tree from top to bottom, left to right.
-}
leafNumberToNodeNumber :: MerkleTree -> Int -> Int
leafNumberToNodeNumber tree leafNum = 1 + leafNum + firstLeafNum tree

{- | Get a merkle proof but re-number the node numbers to be zero-indexed
 instead of one-indexed.
-}
neededHashes :: MerkleTree -> Int -> Maybe [(Int, B.ByteString)]
neededHashes tree = fmap (map $ mapFst (subtract 1)) . merkleProof tree

{- | Determine the smallest index into the breadth first list for the given
 tree where a leaf may be found.
-}
firstLeafNum :: MerkleTree -> Int
firstLeafNum tree = size tree `div` 2

{- | Serialize a MerkleTree to bytes by concatenating all of the leaf hashes
 left to right.

 This serialization includes no framing so the only thing we can do is
 consume all available input.  Use this instance with `isolate` and bring
 your own framing mechanism to determine how many bytes to process.
-}
instance Binary MerkleTree where
    put = putByteString . B.concat . breadthFirstList
    get =
        getRemainingLazyByteString
            >>= maybe (fail "could not construct MerkleTree") pure
                . buildTreeOutOfAllTheNodes
                . chunkedBy (hashDigestSize SHA256)
                . LBS.toStrict

-- | Get a list of all of the leaf hashes of a tree from left to right.
leafHashes :: MerkleTree -> [B.ByteString]
leafHashes (MerkleLeaf h) = [h]
leafHashes (MerkleNode _ l r) = leafHashes l <> leafHashes r

{- | Make a merkle tree out of a flat list of all nodes (start from
 root, then first two children, etc .. [0, 1, 2] is a two-layer
 tree, [0, 1, 2, 3, 4, 5, 6] is three-layer, etc
-}
buildTreeOutOfAllTheNodes :: [B.ByteString] -> Maybe MerkleTree
buildTreeOutOfAllTheNodes nodes
    | validMerkleSize nodes = Just (head (treeFromRows [] (clumpRows powersOfTwo nodes)))
    | otherwise = Nothing

{- | Increasing consecutive powers of 2 from 2 ^ 0 to the maximum value
 representable in `Int`.
-}
powersOfTwo :: [Int]
powersOfTwo = (2 ^) <$> [0 :: Int .. 62]

{- | Determine whether a list of nodes is a possible representation of a
 merkle tree.

 It is possible if the number of elements in the list is one less than a
 positive power of 2.
-}
validMerkleSize :: [a] -> Bool
validMerkleSize nodes =
    head (dropWhile (< size') (tail powersOfTwo)) == size'
  where
    size' = length nodes + 1

{- | Reorganize a flat list of merkle tree node values into a list of lists of
 merkle tree node values.  Each inner list gives the values from left to right
 at a particular height in the tree.  The head of the outer list gives the
 leaves.
-}
clumpRows ::
    -- | The numbers of elements of the flat list to take to make this (the
    -- head) and subsequent (the tail) clumps.
    [Int] ->
    -- | The values of the nodes themselves.
    [B.ByteString] ->
    [[B.ByteString]]
clumpRows _ [] = []
clumpRows [] _ = error "Ran out of clump lengths (too many nodes!)"
clumpRows (p : ps) rows = clumpRows ps (drop p rows) ++ [take p rows]

-- | Given some children
treeFromRows ::
    -- | Some children to attach to a list of nodes representing the next
    -- shallowest level of the tree.
    [MerkleTree] ->
    -- | The values of the nodes to create at the next shallowest level of the
    -- tree.
    [[B.ByteString]] ->
    -- | The nodes forming the shallowest level of the tree.  If we built a
    -- full tree, there will be exactly one node here.
    [MerkleTree]
-- if we've processed nothing yet, we're on the "all leafs" children row
treeFromRows [] (children : rest) = treeFromRows (MerkleLeaf <$> children) rest
-- if we're out of other stuff then we're done
treeFromRows children [] = children
-- with only a single thing in the "rest", we're at the root
treeFromRows [left, right] [[root]] = [MerkleNode root left right]
-- this recursion is harder to think about: we want to "collect" done
-- stuff from the first argument and build it up into a tree. kind of.
treeFromRows (left : right : children) (row : rest) = treeFromRows (mTree (left : right : children) row) rest
treeFromRows x y = error $ "treeFromRows not sure what to do with " <> show x <> " " <> show y

-- this does the "second recursion"; see above -- building out a row
-- of parents from children + parent node content
mTree :: [MerkleTree] -> [B.ByteString] -> [MerkleTree]
mTree [left, right] [head'] = [MerkleNode head' left right]
mTree (left : right : more) row = MerkleNode (head row) left right : mTree more (tail row)
mTree x y = error $ "mTree not sure what to do with " <> show x <> " " <> show y