Shrink Collection to make parse_expr_size test pass

compiler/parse/src/ast.rs

@@ -1,44 +1,10 @@
 use crate::header::{AppHeader, ImportsEntry, InterfaceHeader, PlatformHeader, TypedIdent};
 use crate::ident::Ident;
-use bumpalo::collections::String;
+use bumpalo::collections::{String, Vec};
 use bumpalo::Bump;
 use roc_module::operator::{BinOp, CalledVia, UnaryOp};
 use roc_region::all::{Loc, Position, Region};
 
-#[derive(Copy, Clone, Debug, PartialEq)]
-pub struct Collection<'a, T> {
-    pub items: &'a [T],
-    pub final_comments: &'a [CommentOrNewline<'a>],
-}
-
-impl<'a, T> Collection<'a, T> {
-    pub fn empty() -> Collection<'a, T> {
-        Collection {
-            items: &[],
-            final_comments: &[],
-        }
-    }
-
-    pub fn with_items(items: &'a [T]) -> Collection<'a, T> {
-        Collection {
-            items,
-            final_comments: &[],
-        }
-    }
-
-    pub fn iter(&self) -> impl Iterator<Item = &'a T> {
-        self.items.iter()
-    }
-
-    pub fn len(&self) -> usize {
-        self.items.len()
-    }
-
-    pub fn is_empty(&self) -> bool {
-        self.items.is_empty()
-    }
-}
-
 #[derive(Clone, Debug, PartialEq)]
 pub enum Module<'a> {
     Interface { header: InterfaceHeader<'a> },
@@ -523,6 +489,113 @@ impl<'a> Pattern<'a> {
         }
     }
 }
+#[derive(Copy, Clone, Debug)]
+pub struct Collection<'a, T> {
+    pub items: &'a [T],
+    // Use a pointer to a slice (rather than just a slice), in order to avoid bloating
+    // Ast variants. The final_comments field is rarely accessed in the hot path, so
+    // this shouldn't matter much for perf.
+    // Use an Option, so it's possible to initialize without allocating.
+    final_comments: Option<&'a &'a [CommentOrNewline<'a>]>,
+}
+
+impl<'a, T> Collection<'a, T> {
+    pub fn empty() -> Collection<'a, T> {
+        Collection {
+            items: &[],
+            final_comments: None,
+        }
+    }
+
+    pub fn with_items(items: &'a [T]) -> Collection<'a, T> {
+        Collection {
+            items,
+            final_comments: None,
+        }
+    }
+
+    pub fn with_items_and_comments(
+        arena: &'a Bump,
+        items: &'a [T],
+        comments: &'a [CommentOrNewline<'a>],
+    ) -> Collection<'a, T> {
+        Collection {
+            items,
+            final_comments: if comments.is_empty() {
+                None
+            } else {
+                Some(arena.alloc(comments))
+            },
+        }
+    }
+
+    pub fn replace_items<V>(&self, new_items: &'a [V]) -> Collection<'a, V> {
+        Collection {
+            items: new_items,
+            final_comments: self.final_comments,
+        }
+    }
+
+    pub fn ptrify_items(&self, arena: &'a Bump) -> Collection<'a, &'a T> {
+        let mut allocated = Vec::with_capacity_in(self.len(), arena);
+
+        for parsed_elem in self.items {
+            allocated.push(parsed_elem);
+        }
+
+        self.replace_items(allocated.into_bump_slice())
+    }
+
+    pub fn map_items<V: 'a>(&self, arena: &'a Bump, f: impl Fn(&'a T) -> V) -> Collection<'a, V> {
+        let mut allocated = Vec::with_capacity_in(self.len(), arena);
+
+        for parsed_elem in self.items {
+            allocated.push(f(parsed_elem));
+        }
+
+        self.replace_items(allocated.into_bump_slice())
+    }
+
+    pub fn map_items_result<V: 'a, E>(
+        &self,
+        arena: &'a Bump,
+        f: impl Fn(&T) -> Result<V, E>,
+    ) -> Result<Collection<'a, V>, E> {
+        let mut allocated = Vec::with_capacity_in(self.len(), arena);
+
+        for parsed_elem in self.items {
+            allocated.push(f(parsed_elem)?);
+        }
+
+        Ok(self.replace_items(allocated.into_bump_slice()))
+    }
+
+    pub fn final_comments(&self) -> &'a [CommentOrNewline<'a>] {
+        if let Some(final_comments) = self.final_comments {
+            *final_comments
+        } else {
+            &[]
+        }
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = &'a T> {
+        self.items.iter()
+    }
+
+    pub fn len(&self) -> usize {
+        self.items.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.items.is_empty()
+    }
+}
+
+impl<'a, T: PartialEq> PartialEq for Collection<'a, T> {
+    fn eq(&self, other: &Self) -> bool {
+        self.items == other.items && self.final_comments() == other.final_comments()
+    }
+}
 
 pub trait Spaceable<'a> {
     fn before(&'a self, _: &'a [CommentOrNewline<'a>]) -> Self;