From 6002ccdc567b26c69f6b55173f5d11b95b4d6966 Mon Sep 17 00:00:00 2001
From: "Tobin C. Harding" <me@tobin.cc>
Date: Fri, 7 Feb 2025 07:27:29 +1100
Subject: [PATCH] Add a tagged sha256t hash engine

We would like it if two different pre-tagged engines were considered
different types so it is not possible to mix them up.

Add a new `sha256t::HashEngine<T>` where `T` is a tag the same as on
`sha256t::Hash<T>`.
---
 bitcoin/src/taproot/mod.rs    |  2 +-
 hashes/src/internal_macros.rs |  2 +-
 hashes/src/sha256t/mod.rs     | 41 ++++++++++++++++++++++++++++-------
 3 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/bitcoin/src/taproot/mod.rs b/bitcoin/src/taproot/mod.rs
index d9842ec28..775357f07 100644
--- a/bitcoin/src/taproot/mod.rs
+++ b/bitcoin/src/taproot/mod.rs
@@ -1600,7 +1600,7 @@ mod test {
         fn empty_hash(tag_name: &str) -> [u8; 32] {
             let mut e = tag_engine(tag_name);
             e.input(&[]);
-            Hash::<TapBranchTag>::from_engine(e).to_byte_array()
+            sha256::Hash::from_engine(e).to_byte_array()
         }
         assert_eq!(empty_hash("TapLeaf"), Hash::<TapLeafTag>::hash(&[]).to_byte_array());
         assert_eq!(empty_hash("TapBranch"), Hash::<TapBranchTag>::hash(&[]).to_byte_array());
diff --git a/hashes/src/internal_macros.rs b/hashes/src/internal_macros.rs
index 3c6818bd4..5ae79f119 100644
--- a/hashes/src/internal_macros.rs
+++ b/hashes/src/internal_macros.rs
@@ -27,7 +27,7 @@ macro_rules! hash_trait_impls {
         $crate::impl_debug_only!(Hash, { $bits / 8 }, $reverse $(, $gen: $gent)*);
 
         impl<$($gen: $gent),*> $crate::GeneralHash for Hash<$($gen),*> {
-            type Engine = HashEngine;
+            type Engine = HashEngine<$($gen),*>;
 
             fn from_engine(e: Self::Engine) -> Hash<$($gen),*> { Self::from_engine(e) }
         }
diff --git a/hashes/src/sha256t/mod.rs b/hashes/src/sha256t/mod.rs
index f355ac82e..5d28b9d93 100644
--- a/hashes/src/sha256t/mod.rs
+++ b/hashes/src/sha256t/mod.rs
@@ -5,12 +5,12 @@
 use core::cmp;
 use core::marker::PhantomData;
 
+#[cfg(doc)]
+use crate::sha256::Midstate;
 use crate::{sha256, FromSliceError, HashEngine as _};
 
-type HashEngine = sha256::HashEngine;
-
 /// Trait representing a tag that can be used as a context for SHA256t hashes.
-pub trait Tag {
+pub trait Tag: Clone {
     /// The [`Midstate`] after pre-tagging the hash engine.
     const MIDSTATE: sha256::Midstate;
 }
@@ -57,14 +57,12 @@ where
     }
 
     /// Produces a hash from the current state of a given engine.
-    pub fn from_engine(e: HashEngine) -> Hash<T> {
-        Hash::from_byte_array(sha256::Hash::from_engine(e).to_byte_array())
+    pub fn from_engine(e: HashEngine<T>) -> Hash<T> {
+        Hash::from_byte_array(sha256::Hash::from_engine(e.0).to_byte_array())
     }
 
     /// Constructs a new engine.
-    pub fn engine() -> HashEngine {
-        sha256::HashEngine::from_midstate(T::MIDSTATE)
-    }
+    pub fn engine() -> HashEngine<T> { HashEngine::default() }
 
     /// Hashes some bytes.
     #[allow(clippy::self_named_constructors)] // Hash is a noun and a verb.
@@ -137,6 +135,33 @@ impl<T: Tag> core::hash::Hash for Hash<T> {
 
 crate::internal_macros::hash_trait_impls!(256, false, T: Tag);
 
+/// Engine to compute SHA256t hash function.
+#[derive(Debug, Clone)]
+pub struct HashEngine<T>(sha256::HashEngine, PhantomData<T>);
+
+impl<T: Tag> Default for HashEngine<T> {
+    fn default() -> Self {
+        let tagged = sha256::HashEngine::from_midstate(T::MIDSTATE);
+        HashEngine(tagged, PhantomData)
+    }
+}
+
+impl<T: Tag> crate::HashEngine for HashEngine<T> {
+    const BLOCK_SIZE: usize = 64; // Same as sha256::HashEngine::BLOCK_SIZE;
+    fn input(&mut self, data: &[u8]) { self.0.input(data) }
+    fn n_bytes_hashed(&self) -> u64 { self.0.n_bytes_hashed() }
+}
+
+crate::internal_macros::impl_io_write!(
+    HashEngine<T>,
+    |us: &mut HashEngine<T>, buf| {
+        us.input(buf);
+        Ok(buf.len())
+    },
+    |_us| { Ok(()) },
+    T: crate::sha256t::Tag
+);
+
 // Workaround macros being unavailable in attributes.
 #[doc(hidden)]
 #[macro_export]