MultiHash.hs

-- | Module      : MultiFormats.MultiHash
--   Description : Associates multicodec values with their hash functions.
--   Copyright   : Zoey McBride (c) 2026
--   License     : AGPL-3.0-or-later
--   Maintainer  : zoeymcbride@mailbox.org
--   Stability   : experimental
--
-- Provides a lookup table to verify multicodec values as multihashes and a way
-- to resolve hash functions from multicodecs or their multicodec name. See:
-- https://github.com/multiformats/multihash
module MultiFormats.MultiHash
  ( MultiHash (MultiHash, multiHashCodec, multiHashFunction),
    resolveMultiHash,
    resolveHashCodec,
  )
where

import Control.DeepSeq (NFData (..))
import Control.Monad (join, (<$!>))
import Crypto.Hash (HashAlgorithm)
import Crypto.Hash qualified as Hash
import Crypto.Hash.Algorithms qualified as Algo
import Data.BaseSystems (base16lower, encoder)
import Data.ByteArray qualified as Memory
import Data.ByteString.Builder (Builder)
import Data.ByteString.Builder qualified as Builder
import Data.ByteString.Internal (ByteString (BS))
import Data.ByteString.Lazy (LazyByteString)
import Data.ByteString.Lazy qualified as LBS
import Data.Function (on)
import Data.Map (Map)
import Data.Map qualified as LazyMap
import Data.Maybe (fromMaybe, mapMaybe)
import Data.Text qualified as Text
import Data.XCodec.BinaryTranscoder (packValueBE)
import Foreign.ForeignPtr (newForeignPtr)
import Foreign.Marshal.Alloc (finalizerFree)
import Foreign.Marshal.Array (copyArray, mallocArray)
import GHC.IO (unsafePerformIO)
import MultiFormats.MultiCodec (MultiCodec (..))
import MultiFormats.MultiCodec qualified as MC

-- | Alias for the type of function signture we want from a hash function.
type HashFunction = Builder -> LazyByteString

-- | Forms the MultiHash implementaton.
data MultiHash = MultiHash
  { -- | Stores the MultiCodec associated with multiHashCodec. Used for Eq, Ord.
    multiHashCodec :: MultiCodec,
    -- | Lazy loads the hash function when called.
    multiHashFunction :: ~HashFunction
  }

-- | Evaluates MultiHash into strict normal form.
instance NFData MultiHash where
  rnf (MultiHash codec hashfunc) = rnf codec `seq` rnf hashfunc

instance Show MultiHash where
  show (MultiHash codec@(MultiCodec value) _) =
    let name = fromMaybe "<invalid codec>" $ MC.resolveCodecName codec
        hex = "0x" ++ Text.unpack (encoder base16lower $ packValueBE value)
     in "MultiHash(" ++ name ++ "=" ++ hex ++ ")"

instance Eq MultiHash where
  (==) = (==) `on` multiHashCodec

instance Ord MultiHash where
  compare = compare `on` multiHashCodec

-- | Resolves a MultiHash from a MultiCodec.
resolveHashCodec :: MultiCodec -> Maybe MultiHash
resolveHashCodec codec =
  -- Strictly contructs the MultiHash with the codec and the hash lookup
  (MultiHash $! codec) <$!> join (LazyMap.lookup codec hashFunctionTable)

-- | Resolves a MultiHash from MultiCodec name.
resolveMultiHash :: String -> Maybe MultiHash
resolveMultiHash name = resolveHashCodec =<< MC.resolveMultiCodec name

-- | Runs the hash function and copies the result into a new buffer.
hashBytesCopy :: (HashAlgorithm a) => a -> HashFunction
hashBytesCopy algo builder =
  -- This was done as a pure function with pack/unpack, but we're going to use
  -- IO to improve the performance of the copy, which needs to happen to
  -- instantiate a new bytestring type from Crypto.Hash function results.
  unsafePerformIO $!
    -- In order to work with the data produced by the hash via Memory, we have
    -- pass an IO transformation on the pointer to the data. The data at the
    -- pointer gets GC'd after this, so we can't copy by referenece and have to
    -- malloc to retain the data (and not double free and crash) unfortunately.
    Memory.withByteArray digest $ \srcptr -> do
      -- Copy the array into newly allocated space
      destptr <- mallocArray len
      copyArray destptr srcptr len
      -- Associate the finalizer with free for the pointer to be used in the GC
      hashptr <- newForeignPtr finalizerFree destptr
      -- Construct the bytestring from the pointer, then O(1) convert to lazy
      return $ LBS.fromStrict (BS hashptr len)
  where
    -- Loads the entire input into memory. hashWith only accepts instances of
    -- ByteArrayAccess, but LazyByteString doesn't implement this.
    access = LBS.toStrict (Builder.toLazyByteString builder)
    -- Runs the hash algorithm on the accessed data.
    digest = Hash.hashWith algo access
    -- Gets the length (in bytes) of the pointed to memory.
    len = Memory.length digest

-- | Resolves all the String names from MultiCodec table and assigns an
-- implementation of the hash function if it exists
hashFunctionTable :: Map MultiCodec (Maybe HashFunction)
hashFunctionTable =
  -- We want a lazy map so we can skip evaluting hash functions that never get
  -- looked up, but we have to evaluate all the hash codecs that exist.
  LazyMap.fromList . mapMaybe pairEntries $
    [ ("identity", Just Builder.toLazyByteString),
      ("sha1", hashSingle Algo.SHA1),
      ("sha2-256", hashSingle Algo.SHA256),
      ("sha2-512", hashSingle Algo.SHA512),
      ("sha3-256", hashSingle Algo.SHA3_256),
      ("sha3-512", hashSingle Algo.SHA3_512),
      ("sha3-224", hashSingle Algo.SHA3_224),
      ("sha3-384", hashSingle Algo.SHA3_384),
      ("sha3-512", hashSingle Algo.SHA3_512),
      ("keccak-224", hashSingle Algo.Keccak_224),
      ("keccak-256", hashSingle Algo.Keccak_256),
      ("keccak-384", hashSingle Algo.Keccak_384),
      ("keccak-512", hashSingle Algo.Keccak_512),
      ("md5", hashSingle Algo.MD5),
      ("dbl-sha2-256", hashDouble Algo.SHA256),
      -- These have HashAlgorithms in Algo, but idk how to use them
      ("blake3", Nothing),
      ("shake-128", Nothing),
      ("shake-256", Nothing)
    ]
  where
    -- Applies the hash algorithm once
    hashSingle algo = Just $ hashBytesCopy algo
    -- Applies the hash algorithm twice
    hashDouble algo = do
      hash <- hashSingle algo
      return $ hash . Builder.lazyByteString . hash
    -- Returns a map entry if the multicodec exists
    pairEntries (name, hfunc) =
      case MC.resolveMultiCodec name of
        Just multicodec -> Just (multicodec, hfunc)
        Nothing -> Nothing