Fix hot reloading components with keys (#2886)

* Fix hot reloading components with keys

* include component formatted segment keys, but not dynamic component value

* Fix component_literal_dyn_idx index

* add a new test for hot reloading components with keys

* Even more tests

* make clippy happy

* fix typo

packages/rsx/src/assign_dyn_ids.rs

@@ -69,10 +69,13 @@ impl<'a> DynIdVisitor<'a> {
                         if let HotLiteral::Fmted(segments) = literal {
                             self.assign_formatted_segment(segments);
                         }
-                        component.component_literal_dyn_idx[index]
-                            .set(self.component_literal_index);
-                        self.component_literal_index += 1;
-                        index += 1;
+                        // Don't include keys in the component dynamic pool
+                        if !property.name.is_likely_key() {
+                            component.component_literal_dyn_idx[index]
+                                .set(self.component_literal_index);
+                            self.component_literal_index += 1;
+                            index += 1;
+                        }
                     }
                 }
             }

packages/rsx/src/attribute.rs

@@ -406,6 +406,10 @@ impl AttributeName {
         matches!(self, Self::BuiltIn(ident) if ident.to_string().starts_with("on"))
     }
 
+    pub fn is_likely_key(&self) -> bool {
+        matches!(self, Self::BuiltIn(ident) if ident == "key")
+    }
+
     pub fn span(&self) -> proc_macro2::Span {
         match self {
             Self::Custom(lit) => lit.span(),

packages/rsx/src/component.rs

@@ -171,10 +171,10 @@ impl Component {
     }
 
     pub fn get_key(&self) -> Option<&AttributeValue> {
-        self.fields.iter().find_map(|attr| match &attr.name {
-            AttributeName::BuiltIn(key) if key == "key" => Some(&attr.value),
-            _ => None,
-        })
+        self.fields
+            .iter()
+            .find(|attr| attr.name.is_likely_key())
+            .map(|attr| &attr.value)
     }
 
     /// Ensure there's no duplicate props - this will be a compile error but we can move it to a
@@ -252,48 +252,41 @@ impl Component {
         self.spreads.first().map(|spread| &spread.expr)
     }
 
-    fn make_field_idents(&self) -> Vec<(TokenStream2, TokenStream2)> {
-        let mut dynamic_literal_index = 0;
+    // Iterate over the props of the component (without spreads, key, and custom attributes)
+    pub(crate) fn component_props(&self) -> impl Iterator<Item = &Attribute> {
         self.fields
             .iter()
-            .filter_map(move |attr| {
-                let Attribute { name, value, .. } = attr;
+            .filter(move |attr| !attr.name.is_likely_key())
+    }
 
-                let attr = match name {
-                    AttributeName::BuiltIn(k) => {
-                        if k == "key" {
-                            return None;
+    fn make_field_idents(&self) -> Vec<(TokenStream2, TokenStream2)> {
+        let mut dynamic_literal_index = 0;
+        self.component_props()
+            .map(|attribute| {
+                let release_value = attribute.value.to_token_stream();
+
+            // In debug mode, we try to grab the value from the dynamic literal pool if possible
+            let value = if let AttributeValue::AttrLiteral(literal) = &attribute.value {
+                let idx = self.component_literal_dyn_idx[dynamic_literal_index].get();
+                dynamic_literal_index += 1;
+                let debug_value = quote! { __dynamic_literal_pool.component_property(#idx, &*__template_read, #literal) };
+                quote! {
+                    {
+                        #[cfg(debug_assertions)]
+                        {
+                            #debug_value
                         }
-                        quote! { #k }
-                    }
-                    AttributeName::Custom(_) => return None,
-                    AttributeName::Spread(_) => return None,
-                };
-
-                let release_value = value.to_token_stream();
-
-                // In debug mode, we try to grab the value from the dynamic literal pool if possible
-                let value = if let AttributeValue::AttrLiteral(literal) = &value {
-                    let idx = self.component_literal_dyn_idx[dynamic_literal_index].get();
-                    dynamic_literal_index += 1;
-                    let debug_value = quote! { __dynamic_literal_pool.component_property(#idx, &*__template_read, #literal) };
-                    quote! {
+                        #[cfg(not(debug_assertions))]
                         {
-                            #[cfg(debug_assertions)]
-                            {
-                                #debug_value
-                            }
-                            #[cfg(not(debug_assertions))]
-                            {
-                                #release_value
-                            }
+                            #release_value
                         }
                     }
-                } else {
-                    release_value
-                };
+                }
+            } else {
+                release_value
+            };
 
-                Some((attr, value))
+                (attribute.name.to_token_stream(), value)
             })
             .collect()
     }
@@ -345,6 +338,7 @@ fn normalize_path(name: &mut syn::Path) -> Option<AngleBracketedGenericArguments
 mod tests {
     use super::*;
     use prettier_please::PrettyUnparse;
+    use syn::parse_quote;
 
     /// Ensure we can parse a component
     #[test]
@@ -482,4 +476,23 @@ mod tests {
 
         let _parsed: syn::Path = syn::parse2(input).unwrap();
     }
+
+    #[test]
+    fn identifies_key() {
+        let input = quote! {
+            Link { key: "{value}", to: Route::List, class: "pure-button", "Go back" }
+        };
+
+        let component: Component = syn::parse2(input).unwrap();
+
+        // The key should exist
+        assert_eq!(component.get_key(), Some(&parse_quote!("{value}")));
+
+        // The key should not be included in the properties
+        let properties = component
+            .component_props()
+            .map(|attr| attr.name.to_string())
+            .collect::<Vec<_>>();
+        assert_eq!(properties, ["to", "class"]);
+    }
 }

packages/rsx/src/element.rs

@@ -237,7 +237,7 @@ impl Element {
         }
 
         for attr in attrs {
-            if attr.name.to_string() == "key" {
+            if attr.name.is_likely_key() {
                 continue;
             }
 
@@ -301,15 +301,10 @@ impl Element {
     }
 
     pub(crate) fn key(&self) -> Option<&AttributeValue> {
-        for attr in &self.raw_attributes {
-            if let AttributeName::BuiltIn(name) = &attr.name {
-                if name == "key" {
-                    return Some(&attr.value);
-                }
-            }
-        }
-
-        None
+        self.raw_attributes
+            .iter()
+            .find(|attr| attr.name.is_likely_key())
+            .map(|attr| &attr.value)
     }
 
     fn completion_hints(&self) -> TokenStream2 {

packages/rsx/src/hot_reload/diff.rs

@@ -406,13 +406,15 @@ impl HotReloadResult {
         }
 
         // Then check if the fields are the same
-        if new_component.fields.len() != old_component.fields.len() {
+        let new_non_key_fields: Vec<_> = new_component.component_props().collect();
+        let old_non_key_fields: Vec<_> = old_component.component_props().collect();
+        if new_non_key_fields.len() != old_non_key_fields.len() {
             return None;
         }
 
-        let mut new_fields = new_component.fields.clone();
+        let mut new_fields = new_non_key_fields.clone();
         new_fields.sort_by_key(|attribute| attribute.name.to_string());
-        let mut old_fields = old_component.fields.iter().enumerate().collect::<Vec<_>>();
+        let mut old_fields = old_non_key_fields.iter().enumerate().collect::<Vec<_>>();
         old_fields.sort_by_key(|(_, attribute)| attribute.name.to_string());
 
         // The literal component properties for the component in same the order as the original component property literals

packages/rsx/src/template_body.rs

@@ -320,7 +320,7 @@ impl TemplateBody {
                 }
             })
             .flat_map(|component| {
-                component.fields.iter().filter_map(|field| {
+                component.component_props().filter_map(|field| {
                     if let AttributeValue::AttrLiteral(literal) = &field.value {
                         Some(literal)
                     } else {

packages/rsx/tests/hotreload_pattern.rs

@@ -1085,6 +1085,109 @@ fn component_with_handlers() {
     );
 }
 
+#[test]
+fn component_remove_key() {
+    let a = quote! {
+        Component {
+            key: "{key}",
+            class: 123,
+            id: 456.789,
+            other: true,
+            dynamic1,
+            dynamic2,
+            blah: "hello {world}",
+            onclick: |e| { println!("clicked") },
+        }
+    };
+
+    // changing lit values
+    let b = quote! {
+        Component {
+            class: 456,
+            id: 789.456,
+            other: false,
+            dynamic1,
+            dynamic2,
+            blah: "goodbye {world}",
+            onclick: |e| { println!("clicked") },
+        }
+    };
+
+    let hot_reload = hot_reload_from_tokens(a, b).unwrap();
+    let template = hot_reload.get(&0).unwrap();
+    assert_eq!(
+        template.component_values,
+        &[
+            HotReloadLiteral::Int(456),
+            HotReloadLiteral::Float(789.456),
+            HotReloadLiteral::Bool(false),
+            HotReloadLiteral::Fmted(FmtedSegments::new(vec![
+                FmtSegment::Literal { value: "goodbye " },
+                FmtSegment::Dynamic { id: 1 }
+            ]))
+        ]
+    );
+}
+
+#[test]
+fn component_modify_key() {
+    let a = quote! {
+        Component {
+            key: "{key}",
+            class: 123,
+            id: 456.789,
+            other: true,
+            dynamic1,
+            dynamic2,
+            blah1: "hello {world123}",
+            blah2: "hello {world}",
+            onclick: |e| { println!("clicked") },
+        }
+    };
+
+    // changing lit values
+    let b = quote! {
+        Component {
+            key: "{key}-{world}",
+            class: 456,
+            id: 789.456,
+            other: false,
+            dynamic1,
+            dynamic2,
+            blah1: "hello {world123}",
+            blah2: "hello {world}",
+            onclick: |e| { println!("clicked") },
+        }
+    };
+
+    let hot_reload = hot_reload_from_tokens(a, b).unwrap();
+    let template = hot_reload.get(&0).unwrap();
+    assert_eq!(
+        template.key,
+        Some(FmtedSegments::new(vec![
+            FmtSegment::Dynamic { id: 0 },
+            FmtSegment::Literal { value: "-" },
+            FmtSegment::Dynamic { id: 2 },
+        ]))
+    );
+    assert_eq!(
+        template.component_values,
+        &[
+            HotReloadLiteral::Int(456),
+            HotReloadLiteral::Float(789.456),
+            HotReloadLiteral::Bool(false),
+            HotReloadLiteral::Fmted(FmtedSegments::new(vec![
+                FmtSegment::Literal { value: "hello " },
+                FmtSegment::Dynamic { id: 1 }
+            ])),
+            HotReloadLiteral::Fmted(FmtedSegments::new(vec![
+                FmtSegment::Literal { value: "hello " },
+                FmtSegment::Dynamic { id: 2 }
+            ]))
+        ]
+    );
+}
+
 #[test]
 fn duplicating_dynamic_nodes() {
     let a = quote! {