From 994056e16d14706cae5a2546f78d350772ceb5e3 Mon Sep 17 00:00:00 2001
From: Jonathan Kelley <jkelleyrtp@gmail.com>
Date: Mon, 8 Apr 2024 22:31:58 -0700
Subject: [PATCH] Don't eat keys (#2280)

---
 packages/autofmt/src/component.rs          | 23 +++++++++++++++++-----
 packages/autofmt/src/writer.rs             |  9 ++++++++-
 packages/autofmt/tests/samples/complex.rsx | 15 ++++++++++++++
 3 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/packages/autofmt/src/component.rs b/packages/autofmt/src/component.rs
index 664a539bb..484810e4d 100644
--- a/packages/autofmt/src/component.rs
+++ b/packages/autofmt/src/component.rs
@@ -28,6 +28,7 @@ impl Writer<'_> {
             children,
             manual_props,
             prop_gen_args,
+            key,
             ..
         }: &Component,
     ) -> Result {
@@ -38,7 +39,7 @@ impl Writer<'_> {
         let mut opt_level = ShortOptimization::NoOpt;
 
         // check if we have a lot of attributes
-        let attr_len = self.field_len(fields, manual_props);
+        let attr_len = self.field_len(fields, manual_props) + self.key_len(key.as_ref());
         let is_short_attr_list = attr_len < 80;
         let is_small_children = self.is_short_children(children).is_some();
 
@@ -62,7 +63,7 @@ impl Writer<'_> {
         }
 
         // If there's nothing at all, empty optimization
-        if fields.is_empty() && children.is_empty() && manual_props.is_none() {
+        if fields.is_empty() && children.is_empty() && manual_props.is_none() && key.is_none() {
             opt_level = ShortOptimization::Empty;
         }
 
@@ -85,7 +86,7 @@ impl Writer<'_> {
             ShortOptimization::Oneliner => {
                 write!(self.out, " ")?;
 
-                self.write_component_fields(fields, manual_props, true)?;
+                self.write_component_fields(fields, key.as_ref(), manual_props, true)?;
 
                 if !children.is_empty() && !fields.is_empty() {
                     write!(self.out, ", ")?;
@@ -103,7 +104,7 @@ impl Writer<'_> {
 
             ShortOptimization::PropsOnTop => {
                 write!(self.out, " ")?;
-                self.write_component_fields(fields, manual_props, true)?;
+                self.write_component_fields(fields, key.as_ref(), manual_props, true)?;
 
                 if !children.is_empty() && !fields.is_empty() {
                     write!(self.out, ",")?;
@@ -114,7 +115,7 @@ impl Writer<'_> {
             }
 
             ShortOptimization::NoOpt => {
-                self.write_component_fields(fields, manual_props, false)?;
+                self.write_component_fields(fields, key.as_ref(), manual_props, false)?;
 
                 if !children.is_empty() && !fields.is_empty() {
                     write!(self.out, ",")?;
@@ -154,11 +155,23 @@ impl Writer<'_> {
     fn write_component_fields(
         &mut self,
         fields: &[ComponentField],
+        key: Option<&IfmtInput>,
         manual_props: &Option<syn::Expr>,
         sameline: bool,
     ) -> Result {
         let mut field_iter = fields.iter().peekable();
 
+        // write the key
+        if let Some(key) = key {
+            write!(self.out, "key: {}", ifmt_to_string(key))?;
+            if !fields.is_empty() {
+                write!(self.out, ",")?;
+                if sameline {
+                    write!(self.out, " ")?;
+                }
+            }
+        }
+
         while let Some(field) = field_iter.next() {
             if !sameline {
                 self.out.indented_tabbed_line().unwrap();
diff --git a/packages/autofmt/src/writer.rs b/packages/autofmt/src/writer.rs
index 73cc84a5e..7c840c763 100644
--- a/packages/autofmt/src/writer.rs
+++ b/packages/autofmt/src/writer.rs
@@ -1,4 +1,4 @@
-use dioxus_rsx::{AttributeType, BodyNode, ElementAttrValue, ForLoop, IfChain};
+use dioxus_rsx::{AttributeType, BodyNode, ElementAttrValue, ForLoop, IfChain, IfmtInput};
 use proc_macro2::{LineColumn, Span};
 use quote::ToTokens;
 use std::{
@@ -329,6 +329,13 @@ impl<'a> Writer<'a> {
 
         Ok(())
     }
+
+    pub(crate) fn key_len(&self, key: Option<&IfmtInput>) -> usize {
+        match key {
+            Some(key) => ifmt_to_string(key).len() + 5,
+            None => 0,
+        }
+    }
 }
 
 pub(crate) trait SpanLength {
diff --git a/packages/autofmt/tests/samples/complex.rsx b/packages/autofmt/tests/samples/complex.rsx
index 9031deea1..25e5f36c6 100644
--- a/packages/autofmt/tests/samples/complex.rsx
+++ b/packages/autofmt/tests/samples/complex.rsx
@@ -48,6 +48,21 @@ rsx! {
         }
     }
 
+    for i in 0..10 {
+        Component { key: "{i}", blah: 120 }
+    }
+    for i in 0..10 {
+        Component { key: "{i}" }
+    }
+
+    for i in 0..10 {
+        div { key: "{i}", blah: 120 }
+    }
+
+    for i in 0..10 {
+        div { key: "{i}" }
+    }
+
     div {
         "asdbascasdbasd"
         "asbdasbdabsd"