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),
    resolveHashCodec,
    resolveHashName,
  )
where

import Control.Monad (join)
import Crypto.Hash (HashAlgorithm)
import Crypto.Hash qualified as HF
import Crypto.Hash.Algorithms qualified as Algo
import Data.ByteArray qualified as Memory
import Data.ByteString (ByteString)
import Data.ByteString qualified as Bytes
import Data.Function (on)
import Data.Map (Map)
import Data.Map qualified as LazyMap
import Data.Maybe (fromMaybe, mapMaybe)
import MultiFormats.MultiCodec (MultiCodec (..))
import MultiFormats.MultiCodec qualified as MC
import Text.Printf (printf)

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

-- | 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
  }

instance Show MultiHash where
  show =
    printf "MultiHash(%s)"
      . fromMaybe "<invalid codec>"
      . MC.resolveCodecName
      . multiHashCodec

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 =
  MultiHash codec <$> join (LazyMap.lookup codec hashFunctionTable)

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

-- Provides an implementation for the given HashAlgorithm.
hashImpl :: (HashAlgorithm a) => a -> Maybe HashFunction
hashImpl algo =
  -- This code is really slow. It copies each byte from Memory to
  -- a new ByteString. I wonder if I can take the pointer from Memory and
  -- assign it to a ByteString for O(1) conversion. Or if we can convince the
  -- compiler that that's all that needs done.
  return $ Bytes.pack . Memory.unpack . HF.hashWith algo

-- | 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 $
    -- This should leave the hfuncs unevaluted in the Map.
    mapMaybe
      (\(name, hfunc) -> (,hfunc) <$> MC.resolveCodecValue name)
      [ ("identity", Just id),
        ("sha1", hashImpl Algo.SHA1),
        ("sha2-256", hashImpl Algo.SHA256),
        ("sha2-512", hashImpl Algo.SHA512),
        ("sha3-256", hashImpl Algo.SHA3_256),
        ("sha3-512", hashImpl Algo.SHA3_512),
        ("sha3-224", hashImpl Algo.SHA3_224),
        ("sha3-384", hashImpl Algo.SHA3_384),
        ("sha3-512", hashImpl Algo.SHA3_512),
        ("keccak-224", hashImpl Algo.Keccak_224),
        ("keccak-256", hashImpl Algo.Keccak_256),
        ("keccak-384", hashImpl Algo.Keccak_384),
        ("keccak-512", hashImpl Algo.Keccak_512),
        ("blake3", Nothing),
        -- Shake has HashAlgorithms in Algo, but idk how to use them
        ("shake-128", Nothing),
        ("shake-256", Nothing),
        ("md5", hashImpl Algo.MD5),
        -- This can be implemented w/ composition
        ("dbl-sha2-256", Nothing)
      ]