From c3806cae5ec152ed4c10393bc60c26b42f7476d8 Mon Sep 17 00:00:00 2001
From: Evan Mesterhazy <emesterhazy@google.com>
Date: Wed, 14 Feb 2024 13:28:11 -0500
Subject: [PATCH] Add support for generics to #[derive(ContentHash)]

#3054
---
 lib/proc-macros/src/content_hash.rs | 13 ++++++++++++-
 lib/proc-macros/src/lib.rs          |  9 +++++++--
 lib/src/content_hash.rs             | 20 ++++++++++++++++++--
 3 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/lib/proc-macros/src/content_hash.rs b/lib/proc-macros/src/content_hash.rs
index 1e3a55e86b..1f286f4348 100644
--- a/lib/proc-macros/src/content_hash.rs
+++ b/lib/proc-macros/src/content_hash.rs
@@ -1,7 +1,18 @@
 use proc_macro2::TokenStream;
 use quote::{quote, quote_spanned};
 use syn::spanned::Spanned;
-use syn::{Data, Fields, Index};
+use syn::{parse_quote, Data, Fields, GenericParam, Generics, Index};
+
+pub fn add_trait_bounds(mut generics: Generics) -> Generics {
+    for param in &mut generics.params {
+        if let GenericParam::Type(ref mut type_param) = *param {
+            type_param
+                .bounds
+                .push(parse_quote!(crate::content_hash::ContentHash));
+        }
+    }
+    generics
+}
 
 pub fn generate_hash_impl(data: &Data) -> TokenStream {
     match *data {
diff --git a/lib/proc-macros/src/lib.rs b/lib/proc-macros/src/lib.rs
index 2dc47c0163..5aae404f1d 100644
--- a/lib/proc-macros/src/lib.rs
+++ b/lib/proc-macros/src/lib.rs
@@ -18,9 +18,14 @@ pub fn derive_content_hash(input: proc_macro::TokenStream) -> proc_macro::TokenS
     // Generate an expression to hash each of the fields in the struct.
     let hash_impl = content_hash::generate_hash_impl(&input.data);
 
+    // Handle structs and enums with generics.
+    let generics = content_hash::add_trait_bounds(input.generics);
+    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
+
     let expanded = quote! {
-        impl ::jj_lib::content_hash::ContentHash for #name {
-            fn hash(&self, state: &mut impl ::jj_lib::content_hash::DigestUpdate) {
+        impl #impl_generics ::jj_lib::content_hash::ContentHash for #name #ty_generics
+        #where_clause {
+            fn hash(&self, state: &mut impl digest::Update) {
                 #hash_impl
             }
         }
diff --git a/lib/src/content_hash.rs b/lib/src/content_hash.rs
index 735debdf29..d96c868750 100644
--- a/lib/src/content_hash.rs
+++ b/lib/src/content_hash.rs
@@ -225,15 +225,31 @@ mod tests {
 
     #[test]
     fn test_consistent_hashing() {
+        let expected_hash =
+            "14e42ea3d680bc815d0cea8ac20d3e872120014fb7bba8d82c3ffa7a8e6d63c41ef9631c60b73b150e3dd72efe50e8b0248321fe2b7eea09d879f3757b879372";
         content_hash! {
             struct Foo { x: Vec<Option<i32>>, y: i64 }
         }
-        insta::assert_snapshot!(
+        assert_eq!(
             hex::encode(hash(&Foo {
                 x: vec![None, Some(42)],
                 y: 17
             })),
-            @"14e42ea3d680bc815d0cea8ac20d3e872120014fb7bba8d82c3ffa7a8e6d63c41ef9631c60b73b150e3dd72efe50e8b0248321fe2b7eea09d879f3757b879372"
+            expected_hash
+        );
+
+        // Try again with an equivalent generic struct deriving ContentHash.
+        #[derive(ContentHash)]
+        struct GenericFoo<X, Y> {
+            x: X,
+            y: Y,
+        }
+        assert_eq!(
+            hex::encode(hash(&GenericFoo {
+                x: vec![None, Some(42)],
+                y: 17i64
+            })),
+            expected_hash
         );
     }