KCL: Fewer allocations when recasting array expressions (#7854)

Using a `.flat_map` required me to allocate a vector inside the closure. Turning that whole part into a `for` loop that pushed to a vec was faster.

Author: adamchalmers

rust/kcl-lib/src/parsing/ast/types/mod.rs

@@ -1944,6 +1944,10 @@ impl CallExpressionKw {
             .chain(self.arguments.iter().map(|arg| (arg.label.as_ref(), &arg.arg)))
     }
 
+    pub fn num_arguments(&self) -> usize {
+        self.arguments.len() + if self.unlabeled.is_some() { 1 } else { 0 }
+    }
+
     pub fn replace_value(&mut self, source_range: SourceRange, new_value: Expr) {
         if let Some(unlabeled) = &mut self.unlabeled {
             unlabeled.replace_value(source_range, new_value.clone());

rust/kcl-lib/src/unparser.rs

@@ -527,32 +527,31 @@ impl ArrayExpression {
         let num_items = self.elements.len() + self.non_code_meta.non_code_nodes_len();
         let mut elems = self.elements.iter();
         let mut found_line_comment = false;
-        let mut format_items: Vec<_> = (0..num_items)
-            .flat_map(|i| {
-                if let Some(noncode) = self.non_code_meta.non_code_nodes.get(&i) {
-                    noncode
-                        .iter()
-                        .map(|nc| {
-                            found_line_comment |= nc.value.should_cause_array_newline();
-                            nc.recast(options, 0)
-                        })
-                        .collect::<Vec<_>>()
-                } else {
-                    let el = elems.next().unwrap();
-                    let s = format!("{}, ", el.recast(options, 0, ExprContext::Other));
-                    vec![s]
-                }
-            })
-            .collect();
+        let mut format_items: Vec<_> = Vec::with_capacity(num_items);
+        for i in 0..num_items {
+            if let Some(noncode) = self.non_code_meta.non_code_nodes.get(&i) {
+                format_items.extend(noncode.iter().map(|nc| {
+                    found_line_comment |= nc.value.should_cause_array_newline();
+                    nc.recast(options, 0)
+                }));
+            } else {
+                let el = elems.next().unwrap();
+                let s = format!("{}, ", el.recast(options, 0, ExprContext::Other));
+                format_items.push(s);
+            }
+        }
 
         // Format these items into a one-line array.
         if let Some(item) = format_items.last_mut() {
             if let Some(norm) = item.strip_suffix(", ") {
                 *item = norm.to_owned();
             }
         }
-        let format_items = format_items; // Remove mutability
-        let flat_recast = format!("[{}]", format_items.join(""));
+        let mut flat_recast = "[".to_owned();
+        for fi in &format_items {
+            flat_recast.push_str(fi)
+        }
+        flat_recast.push(']');
 
         // We might keep the one-line representation, if it's short enough.
         let max_array_length = 40;
@@ -562,29 +561,31 @@ impl ArrayExpression {
         }
 
         // Otherwise, we format a multi-line representation.
+        let mut output = "[\n".to_owned();
         let inner_indentation = if ctxt == ExprContext::Pipe {
             options.get_indentation_offset_pipe(indentation_level + 1)
         } else {
             options.get_indentation(indentation_level + 1)
         };
-        let formatted_array_lines = format_items
-            .iter()
-            .map(|s| {
-                format!(
-                    "{inner_indentation}{}{}",
-                    if let Some(x) = s.strip_suffix(" ") { x } else { s },
-                    if s.ends_with('\n') { "" } else { "\n" }
-                )
-            })
-            .collect::<Vec<String>>()
-            .join("")
-            .to_owned();
+        for format_item in format_items {
+            output.push_str(&inner_indentation);
+            output.push_str(if let Some(x) = format_item.strip_suffix(" ") {
+                x
+            } else {
+                &format_item
+            });
+            if !format_item.ends_with('\n') {
+                output.push('\n')
+            }
+        }
         let end_indent = if ctxt == ExprContext::Pipe {
             options.get_indentation_offset_pipe(indentation_level)
         } else {
             options.get_indentation(indentation_level)
         };
-        format!("[\n{formatted_array_lines}{end_indent}]")
+        output.push_str(&end_indent);
+        output.push(']');
+        output
     }
 }