content: {-# LANGUAGE DeriveAnyClass #-}
  {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE DerivingStrategies #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE UndecidableInstances #-}
  
  -- | Merkle tree implementation.
  --
  -- See <https://tools.ietf.org/html/rfc6962>.
  module Cardano.Chain.Common.Merkle (
    -- * MerkleRoot
    MerkleRoot (..),
  
    -- * MerkleTree
    MerkleTree (..),
    mtRoot,
    mkMerkleTree,
    mkMerkleTreeDecoded,
  
    -- * MerkleNode
    MerkleNode (..),
    mkBranch,
    mkLeaf,
    mkLeafDecoded,
  ) where
  
  -- Cardano.Prelude has its own Rube Goldberg variant of 'Foldable' which we do not
  -- want. It would be great if we could write
  --   import           Cardano.Prelude hiding (toList, foldMap)
  -- but HLint insists that this is not OK because toList and foldMap are never
  -- used unqualified. The hiding in fact makes it clearer for the human reader
  -- what's going on.
  
  import Cardano.Crypto (Hash, hashDecoded, hashRaw, hashToBytes)
  import Cardano.Crypto.Raw (Raw)
  import Cardano.Ledger.Binary (
    Annotated (..),
    DecCBOR (..),
    EncCBOR (..),
    FromCBOR (..),
    ToCBOR (..),
    byronProtVer,
    fromByronCBOR,
    serializeBuilder,
    toByronCBOR,
   )
  import Cardano.Prelude
  import Data.Aeson (ToJSON)
  import Data.ByteString.Builder (Builder, byteString, word8)
  import qualified Data.ByteString.Builder.Extra as Builder
  import qualified Data.ByteString.Lazy as LBS
  import Data.Coerce (coerce)
  import qualified Data.Foldable as F
  import Formatting.Buildable (Buildable (..))
  import NoThunks.Class (NoThunks (..))
  import qualified Prelude
  
  --------------------------------------------------------------------------------
  -- MerkleRoot
  --------------------------------------------------------------------------------
  
  -- | Data type for root of Merkle tree
  newtype MerkleRoot a = MerkleRoot
    { getMerkleRoot :: Hash Raw
    -- ^ returns root 'Hash' of Merkle Tree
    }
    deriving (Show, Eq, Ord, Generic)
    deriving anyclass (NFData, NoThunks)
  
  instance Buildable (MerkleRoot a) where
    build (MerkleRoot h) = "MerkleRoot|" <> build h
  
  -- Used for debugging purposes only
  instance ToJSON a => ToJSON (MerkleRoot a)
  
  instance (EncCBOR a, Typeable a) => ToCBOR (MerkleRoot a) where
    toCBOR = toByronCBOR
    encodedSizeExpr size = encodedSizeExpr size . fmap getMerkleRoot
  
  instance DecCBOR a => FromCBOR (MerkleRoot a) where
    fromCBOR = fromByronCBOR
  
  instance EncCBOR a => EncCBOR (MerkleRoot a) where
    encCBOR = encCBOR . getMerkleRoot
  
  instance DecCBOR a => DecCBOR (MerkleRoot a) where
    decCBOR = MerkleRoot <$> decCBOR
  
  merkleRootToBuilder :: MerkleRoot a -> Builder
  merkleRootToBuilder (MerkleRoot h) = byteString (hashToBytes h)
  
  mkRoot :: MerkleRoot a -> MerkleRoot a -> MerkleRoot a
  mkRoot a b =
    MerkleRoot
      . hashRaw
      . toLazyByteString
      $ mconcat
        [word8 1, merkleRootToBuilder a, merkleRootToBuilder b]
  
  emptyHash :: MerkleRoot a
  emptyHash = MerkleRoot (hashRaw mempty)
  
  --------------------------------------------------------------------------------
  -- MerkleTree
  --------------------------------------------------------------------------------
  
  data MerkleTree a
    = MerkleEmpty
    | MerkleTree !Word32 !(MerkleNode a)
    deriving (Eq, Generic)
    deriving anyclass (NFData)
  
  instance Foldable MerkleTree where
    foldMap _ MerkleEmpty = mempty
    foldMap f (MerkleTree _ n) = F.foldMap f n
  
    null MerkleEmpty = True
    null _ = False
  
    length MerkleEmpty = 0
    length (MerkleTree s _) = fromIntegral s
  
  instance Show a => Show (MerkleTree a) where
    show tree = "Merkle tree: " <> show (F.toList tree)
  
  instance (EncCBOR a, Typeable a) => ToCBOR (MerkleTree a) where
    toCBOR = toByronCBOR
  
  instance (DecCBOR a, EncCBOR a) => FromCBOR (MerkleTree a) where
    fromCBOR = fromByronCBOR
  
  -- This instance is both faster and more space-efficient (as confirmed by a
  -- benchmark). Hashing turns out to be faster than decoding extra data.
  instance EncCBOR a => EncCBOR (MerkleTree a) where
    encCBOR = encCBOR . F.toList
  
  instance (DecCBOR a, EncCBOR a) => DecCBOR (MerkleTree a) where
    decCBOR = mkMerkleTree <$> decCBOR
  
  -- | Smart constructor for 'MerkleTree'
  mkMerkleTree :: EncCBOR a => [a] -> MerkleTree a
  mkMerkleTree = mkMerkleTree' (mkLeaf . getConst) . fmap Const
  
  -- | Reconstruct a 'MerkleTree' from a decoded list of items
  mkMerkleTreeDecoded :: [Annotated a ByteString] -> MerkleTree a
  mkMerkleTreeDecoded = mkMerkleTree' mkLeafDecoded
  
  mkMerkleTree' ::
    forall f a b. (f a b -> MerkleNode a) -> [f a b] -> MerkleTree a
  mkMerkleTree' _ [] = MerkleEmpty
  mkMerkleTree' leafBuilder ls = MerkleTree (fromIntegral lsLen) (go lsLen ls)
    where
      lsLen = length ls
      go :: Int -> [f a b] -> MerkleNode a
      go _ [x] = leafBuilder x
      go len xs = mkBranch (go i l) (go (len - i) r)
        where
          i = powerOfTwo len
          (l, r) = splitAt i xs
  
  -- | Return the largest power of two such that it's smaller than X.
  --
  -- >>> powerOfTwo 64
  -- 32
  -- >>> powerOfTwo 65
  -- 64
  powerOfTwo :: forall a. (Bits a, Num a) => a -> a
  powerOfTwo n
    | n .&. (n - 1) == 0 = n `shiftR` 1
    | otherwise = go n
    where
      {- “x .&. (x - 1)” clears the least significant bit:
  
             ↓
         01101000     x
         01100111     x - 1
         --------
         01100000     x .&. (x - 1)
  
         I could've used something like “until (\x -> x*2 > w) (*2) 1”,
         but bit tricks are fun. -}
      go :: a -> a
      go w = if w .&. (w - 1) == 0 then w else go (w .&. (w - 1))
  
  -- | Returns root of Merkle tree
  mtRoot :: MerkleTree a -> MerkleRoot a
  mtRoot MerkleEmpty = emptyHash
  mtRoot (MerkleTree _ n) = nodeRoot n
  
  --------------------------------------------------------------------------------
  -- MerkleNode
  --------------------------------------------------------------------------------
  
  data MerkleNode a
    = -- | MerkleBranch mRoot mLeft mRight
      MerkleBranch !(MerkleRoot a) !(MerkleNode a) !(MerkleNode a)
    | -- | MerkleLeaf mRoot mVal
      MerkleLeaf !(MerkleRoot a) a
    deriving (Eq, Show, Generic)
    deriving anyclass (NFData)
  
  instance Foldable MerkleNode where
    foldMap f x = case x of
      MerkleLeaf _ mVal -> f mVal
      MerkleBranch _ mLeft mRight ->
        F.foldMap f mLeft `mappend` F.foldMap f mRight
  
  toLazyByteString :: Builder -> LBS.ByteString
  toLazyByteString =
    Builder.toLazyByteStringWith (Builder.safeStrategy 1024 4096) mempty
  
  nodeRoot :: MerkleNode a -> MerkleRoot a
  nodeRoot (MerkleLeaf root _) = root
  nodeRoot (MerkleBranch root _ _) = root
  
  mkLeaf :: forall a. EncCBOR a => a -> MerkleNode a
  mkLeaf a = MerkleLeaf mRoot a
    where
      mRoot :: MerkleRoot a
      mRoot =
        MerkleRoot
          $ hashRaw
            (toLazyByteString (word8 0 <> serializeBuilder byronProtVer a))
  
  mkLeafDecoded :: Annotated a ByteString -> MerkleNode a
  mkLeafDecoded a = MerkleLeaf mRoot (unAnnotated a)
    where
      mRoot :: MerkleRoot a
      mRoot = MerkleRoot . coerce . hashDecoded $ prependTag <$> a
  
      prependTag = (LBS.toStrict (toLazyByteString (word8 0)) <>)
  
  mkBranch :: MerkleNode a -> MerkleNode a -> MerkleNode a
  mkBranch nodeA nodeB = MerkleBranch root nodeA nodeB
    where
      root = mkRoot (nodeRoot nodeA) (nodeRoot nodeB)
  