Extract Spaces refactor / cleanup

Have extract_spaces method of ExtractSpaces trait take an arena so that:

1. We can remove a lot of `todo!()`s.
2. Recursively collapse several layers of SpaceBefore/SpaceAfter into a
   single Spaces.
3. Avoid some fuzzer panics (but not all).

Rename new snapshot

Author: gamebox

crates/compiler/can/src/annotation.rs

@@ -1086,7 +1086,7 @@ fn canonicalize_has_clause(
     } = clause;
     let region = *region;
 
-    let var_name = var.extract_spaces().item;
+    let var_name = var.extract_spaces(env.arena).item;
     debug_assert!(
         var_name.starts_with(char::is_lowercase),
         "Vars should have been parsed as lowercase"
@@ -1192,8 +1192,10 @@ fn can_extension_type(
             );
 
             if valid_extension_type(shallow_dealias_with_scope(scope, &ext_type)) {
-                if matches!(loc_ann.extract_spaces().item, TypeAnnotation::Wildcard)
-                    && matches!(ext_problem_kind, ExtensionTypeKind::TagUnion)
+                if matches!(
+                    loc_ann.extract_spaces(env.arena).item,
+                    TypeAnnotation::Wildcard
+                ) && matches!(ext_problem_kind, ExtensionTypeKind::TagUnion)
                     && pol == CanPolarity::Pos
                 {
                     // Wildcards are redundant in positive positions, since they will always be

crates/compiler/can/src/def.rs

@@ -542,7 +542,7 @@ fn canonicalize_claimed_ability_impl<'a>(
 ) -> Result<(Symbol, Symbol), ()> {
     let ability_home = ability.module_id();
 
-    match loc_impl.extract_spaces().item {
+    match loc_impl.extract_spaces(env.arena).item {
         AssignedField::LabelOnly(label) => {
             let label_str = label.value;
             let region = label.region;
@@ -788,7 +788,7 @@ fn canonicalize_opaque<'a>(
 
         for has_ability in has_abilities.value.items {
             let region = has_ability.region;
-            let (ability, opt_impls) = match has_ability.value.extract_spaces().item {
+            let (ability, opt_impls) = match has_ability.value.extract_spaces(env.arena).item {
                 ast::ImplementsAbility::ImplementsAbility { ability, impls } => (ability, impls),
                 _ => internal_error!("spaces not extracted"),
             };
@@ -835,7 +835,7 @@ fn canonicalize_opaque<'a>(
 
                 // First up canonicalize all the claimed implementations, building a map of ability
                 // member -> implementation.
-                for loc_impl in impls.extract_spaces().item.items {
+                for loc_impl in impls.extract_spaces(env.arena).item.items {
                     let (member, impl_symbol) =
                         match canonicalize_claimed_ability_impl(env, scope, ability, loc_impl) {
                             Ok((member, impl_symbol)) => (member, impl_symbol),
@@ -3005,7 +3005,7 @@ fn to_pending_type_def<'a>(
 
             for member in *members {
                 let name_region = member.name.region;
-                let member_name = member.name.extract_spaces().item;
+                let member_name = member.name.extract_spaces(env.arena).item;
 
                 let member_sym = match scope.introduce(member_name.into(), name_region) {
                     Ok(sym) => sym,

crates/compiler/can/src/pattern.rs

@@ -856,7 +856,7 @@ pub fn canonicalize_record_destructs<'a>(
     let mut opt_erroneous = None;
 
     for loc_pattern in patterns.iter() {
-        match loc_pattern.value.extract_spaces().item {
+        match loc_pattern.value.extract_spaces(env.arena).item {
             Identifier { ident: label } => {
                 match scope.introduce(label.into(), region) {
                     Ok(symbol) => {

crates/compiler/fmt/src/annotation.rs

@@ -13,15 +13,12 @@ use bumpalo::{
     collections::{String, Vec},
     Bump,
 };
-use roc_parse::{ast::Spaced, ident::UppercaseIdent};
-use roc_parse::{
-    ast::{
-        AbilityImpls, AssignedField, Collection, CommentOrNewline, Expr, ExtractSpaces,
-        FunctionArrow, ImplementsAbilities, ImplementsAbility, ImplementsClause, Spaceable, Spaces,
-        SpacesAfter, SpacesBefore, Tag, TypeAnnotation, TypeHeader,
-    },
-    expr::merge_spaces,
+use roc_parse::ast::{
+    merge_spaces, AbilityImpls, AssignedField, Collection, CommentOrNewline, Expr, ExtractSpaces,
+    FunctionArrow, ImplementsAbilities, ImplementsAbility, ImplementsClause, Spaceable, Spaces,
+    SpacesAfter, SpacesBefore, Tag, TypeAnnotation, TypeHeader,
 };
+use roc_parse::{ast::Spaced, ident::UppercaseIdent};
 use roc_region::all::Loc;
 
 /// Does an AST node need parens around it?
@@ -618,7 +615,7 @@ impl<'a> Formattable for ImplementsClause<'a> {
     }
 
     fn format_with_options(&self, buf: &mut Buf, parens: Parens, newlines: Newlines, indent: u16) {
-        buf.push_str(self.var.value.extract_spaces().item);
+        buf.push_str(self.var.value.extract_spaces(buf.bump()).item);
         buf.spaces(1);
         buf.push_str(roc_parse::keyword::IMPLEMENTS);
         buf.spaces(1);
@@ -1447,7 +1444,7 @@ pub struct NodeSpaces<'a, T> {
 impl<'a, T: Copy> ExtractSpaces<'a> for NodeSpaces<'a, T> {
     type Item = T;
 
-    fn extract_spaces(&self) -> Spaces<'a, T> {
+    fn extract_spaces(&self, _arena: &'a Bump) -> Spaces<'a, T> {
         Spaces {
             before: self.before,
             item: self.item,

crates/compiler/fmt/src/collection.rs

@@ -1,7 +1,4 @@
-use roc_parse::{
-    ast::{Collection, CommentOrNewline, ExtractSpaces},
-    expr::merge_spaces,
-};
+use roc_parse::ast::{merge_spaces, Collection, CommentOrNewline, ExtractSpaces};
 
 use crate::{
     annotation::{is_collection_multiline, Formattable, Newlines},
@@ -33,7 +30,7 @@ impl Braces {
     }
 }
 
-pub fn fmt_collection<'a, 'buf, T: ExtractSpaces<'a> + Formattable + std::fmt::Debug>(
+pub fn fmt_collection<'a, 'buf: 'a, T: ExtractSpaces<'a> + Formattable + std::fmt::Debug>(
     buf: &mut Buf<'buf>,
 
     indent: u16,
@@ -59,7 +56,7 @@ pub fn fmt_collection<'a, 'buf, T: ExtractSpaces<'a> + Formattable + std::fmt::D
 
         for (index, item) in items.iter().enumerate() {
             let is_first_item = index == 0;
-            let item = item.extract_spaces();
+            let item = item.extract_spaces(buf.arena);
             let is_only_newlines = item.before.iter().all(|s| s.is_newline());
             let last_after_was_only_newlines = last_after.iter().all(|s| s.is_newline());
 

crates/compiler/fmt/src/def.rs

@@ -17,12 +17,11 @@ use crate::{Buf, MigrationFlags};
 use bumpalo::Bump;
 use roc_error_macros::internal_error;
 use roc_parse::ast::{
-    AbilityMember, Defs, Expr, ExtractSpaces, ImportAlias, ImportAsKeyword, ImportExposingKeyword,
-    ImportedModuleName, IngestedFileAnnotation, IngestedFileImport, ModuleImport,
-    ModuleImportParams, Pattern, Spaces, SpacesBefore, StrLiteral, TypeAnnotation, TypeDef,
-    TypeHeader, ValueDef,
+    merge_spaces, AbilityMember, Defs, Expr, ExtractSpaces, ImportAlias, ImportAsKeyword,
+    ImportExposingKeyword, ImportedModuleName, IngestedFileAnnotation, IngestedFileImport,
+    ModuleImport, ModuleImportParams, Pattern, Spaces, SpacesBefore, StrLiteral, TypeAnnotation,
+    TypeDef, TypeHeader, ValueDef,
 };
-use roc_parse::expr::merge_spaces;
 use roc_parse::header::Keyword;
 use roc_region::all::Loc;
 
@@ -435,8 +434,10 @@ impl<'a> Formattable for TypeDef<'a> {
                 typ: ann,
                 derived: has_abilities,
             } => {
-                let ann_is_where_clause =
-                    matches!(ann.extract_spaces().item, TypeAnnotation::Where(..));
+                let ann_is_where_clause = matches!(
+                    ann.extract_spaces(buf.arena).item,
+                    TypeAnnotation::Where(..)
+                );
 
                 // Always put the has-abilities clause on a newline if the opaque annotation
                 // contains a where-has clause.
@@ -476,7 +477,7 @@ impl<'a> Formattable for TypeDef<'a> {
             } => {
                 let header_lifted = type_head_lift_spaces(buf.text.bump(), *header);
                 header_lifted.item.format(buf, indent);
-                let implements = loc_implements.value.extract_spaces();
+                let implements = loc_implements.value.extract_spaces(buf.arena);
                 let before_implements = merge_spaces_conservative(
                     buf.text.bump(),
                     header_lifted.after,
@@ -513,7 +514,7 @@ impl<'a> Formattable for TypeDef<'a> {
                             before,
                             item,
                             after,
-                        } = member.name.value.extract_spaces();
+                        } = member.name.value.extract_spaces(buf.arena);
                         fmt_spaces(buf, before.iter(), indent + INDENT);
                         buf.ensure_ends_with_newline();
 
@@ -833,7 +834,7 @@ impl<'a> Formattable for ValueDef<'a> {
         use roc_parse::ast::ValueDef::*;
         match self {
             Annotation(loc_pattern, loc_annotation) => {
-                let pat_parens = if ann_pattern_needs_parens(&loc_pattern.value) {
+                let pat_parens = if ann_pattern_needs_parens(buf.arena, &loc_pattern.value) {
                     Parens::InApply
                 } else {
                     Parens::NotNeeded
@@ -860,7 +861,7 @@ impl<'a> Formattable for ValueDef<'a> {
                 body_pattern,
                 body_expr,
             } => {
-                let pat_parens = if ann_pattern_needs_parens(&ann_pattern.value) {
+                let pat_parens = if ann_pattern_needs_parens(buf.arena, &ann_pattern.value) {
                     Parens::InApply
                 } else {
                     Parens::NotNeeded
@@ -888,10 +889,12 @@ impl<'a> Formattable for ValueDef<'a> {
     }
 }
 
-fn ann_pattern_needs_parens(value: &Pattern<'_>) -> bool {
-    match value.extract_spaces().item {
+fn ann_pattern_needs_parens(arena: &'_ Bump, value: &Pattern<'_>) -> bool {
+    match value.extract_spaces(arena).item {
         Pattern::Tag(_) => true,
-        Pattern::Apply(func, _args) if matches!(func.extract_spaces().item, Pattern::Tag(..)) => {
+        Pattern::Apply(func, _args)
+            if matches!(func.extract_spaces(arena).item, Pattern::Tag(..)) =>
+        {
             true
         }
         _ => false,
@@ -1016,17 +1019,17 @@ pub fn fmt_body<'a>(
     body: &'a Expr<'a>,
     indent: u16,
 ) {
-    let pattern_extracted = pattern.extract_spaces();
+    let pattern_extracted = pattern.extract_spaces(buf.arena);
     // Check if this is an assignment into the unit value
     let is_unit_assignment = if let Pattern::RecordDestructure(collection) = pattern_extracted.item
     {
         allow_simplify_empty_record_destructure
             && collection.is_empty()
             && pattern_extracted.before.iter().all(|s| s.is_newline())
             && pattern_extracted.after.iter().all(|s| s.is_newline())
-            && !matches!(body.extract_spaces().item, Expr::Defs(..))
-            && !matches!(body.extract_spaces().item, Expr::Return(..))
-            && !matches!(body.extract_spaces().item, Expr::DbgStmt { .. })
+            && !matches!(body.extract_spaces(buf.arena).item, Expr::Defs(..))
+            && !matches!(body.extract_spaces(buf.arena).item, Expr::Return(..))
+            && !matches!(body.extract_spaces(buf.arena).item, Expr::DbgStmt { .. })
             && !starts_with_expect_ident(body)
     } else {
         false
@@ -1186,7 +1189,7 @@ impl<'a> Formattable for AbilityMember<'a> {
         _newlines: Newlines,
         indent: u16,
     ) {
-        let Spaces { before, item, .. } = self.name.value.extract_spaces();
+        let Spaces { before, item, .. } = self.name.value.extract_spaces(buf.arena);
         fmt_spaces(buf, before.iter(), indent);
 
         buf.indent(indent);

crates/compiler/fmt/src/expr.rs

@@ -13,11 +13,9 @@ use bumpalo::collections::Vec;
 use bumpalo::Bump;
 use roc_module::called_via::{self, BinOp, UnaryOp};
 use roc_parse::ast::{
-    AssignedField, Base, Collection, CommentOrNewline, Expr, ExtractSpaces, Pattern, Spaceable,
-    Spaces, SpacesAfter, SpacesBefore, WhenBranch,
+    merge_spaces, AssignedField, Base, Collection, CommentOrNewline, Expr, ExtractSpaces, Pattern,
+    Spaceable, Spaces, SpacesAfter, SpacesBefore, StrLiteral, StrSegment, WhenBranch,
 };
-use roc_parse::ast::{StrLiteral, StrSegment};
-use roc_parse::expr::merge_spaces;
 use roc_parse::ident::Accessor;
 use roc_parse::keyword;
 use roc_region::all::Loc;
@@ -674,6 +672,7 @@ fn fmt_apply(
     indent: u16,
     buf: &mut Buf<'_>,
 ) {
+    let arena = buf.arena;
     // should_reflow_outdentable, aka should we transform this:
     //
     // ```
@@ -692,27 +691,28 @@ fn fmt_apply(
     //   2,
     // ]
     // ```
-    let should_reflow_outdentable = loc_expr.extract_spaces().after.is_empty()
+    let should_reflow_outdentable = loc_expr.extract_spaces(buf.arena).after.is_empty()
         && except_last(loc_args).all(|a| !a.is_multiline())
         && loc_args
             .last()
             .map(|a| {
-                a.extract_spaces().item.is_multiline()
-                    && is_outdentable_collection(&a.value.extract_spaces().item)
-                    && (a.extract_spaces().before == [CommentOrNewline::Newline]
-                        || a.extract_spaces().before.is_empty())
+                a.extract_spaces(arena).item.is_multiline()
+                    && is_outdentable_collection(&a.value.extract_spaces(arena).item)
+                    && (a.extract_spaces(arena).before == [CommentOrNewline::Newline]
+                        || a.extract_spaces(arena).before.is_empty())
             })
             .unwrap_or_default();
 
     let needs_indent = !should_reflow_outdentable
-        && (!loc_expr.extract_spaces().after.is_empty()
+        && (!loc_expr.extract_spaces(arena).after.is_empty()
             || except_last(loc_args).any(|a| a.is_multiline())
             || loc_expr.is_multiline()
             || loc_args
                 .last()
                 .map(|a| {
                     a.is_multiline()
-                        && (!a.extract_spaces().before.is_empty() || !is_outdentable(&a.value))
+                        && (!a.extract_spaces(arena).before.is_empty()
+                            || !is_outdentable(arena, &a.value))
                 })
                 .unwrap_or_default());
 
@@ -868,9 +868,9 @@ pub(crate) fn format_sq_literal(buf: &mut Buf, s: &str) {
     buf.push('\'');
 }
 
-fn is_outdentable(expr: &Expr) -> bool {
+fn is_outdentable(arena: &'_ Bump, expr: &Expr) -> bool {
     matches!(
-        expr.extract_spaces().item,
+        expr.extract_spaces(arena).item,
         Expr::Tuple(_) | Expr::List(_) | Expr::Record(_) | Expr::Closure(..)
     )
 }
@@ -1554,15 +1554,15 @@ pub fn format_spaces(buf: &mut Buf, spaces: &[CommentOrNewline], newlines: Newli
     }
 }
 
-fn is_when_patterns_multiline(when_branch: &WhenBranch) -> bool {
+fn is_when_patterns_multiline(arena: &'_ Bump, when_branch: &WhenBranch) -> bool {
     let patterns = when_branch.patterns;
     let (first_pattern, rest) = patterns.split_first().unwrap();
 
     let is_multiline_patterns = if let Some((last_pattern, inner_patterns)) = rest.split_last() {
-        !first_pattern.value.extract_spaces().after.is_empty()
-            || !last_pattern.value.extract_spaces().before.is_empty()
+        !first_pattern.value.extract_spaces(arena).after.is_empty()
+            || !last_pattern.value.extract_spaces(arena).before.is_empty()
             || inner_patterns.iter().any(|p| {
-                let spaces = p.value.extract_spaces();
+                let spaces = p.value.extract_spaces(arena);
                 !spaces.before.is_empty() || !spaces.after.is_empty()
             })
     } else {
@@ -1624,7 +1624,7 @@ fn fmt_when<'a>(
         let expr = &branch.value;
         let patterns = &branch.patterns;
         let is_multiline_expr = expr.is_multiline();
-        let is_multiline_patterns = is_when_patterns_multiline(branch);
+        let is_multiline_patterns = is_when_patterns_multiline(buf.bump(), branch);
 
         for (pattern_index, pattern) in patterns.iter().enumerate() {
             if pattern_index == 0 {

crates/compiler/fmt/src/lib.rs

@@ -21,6 +21,7 @@ pub struct Buf<'a> {
     beginning_of_line: bool,
     line_indent: u16,
     flags: MigrationFlags,
+    pub arena: &'a Bump,
 }
 
 #[derive(Debug, Copy, Clone)]
@@ -44,6 +45,7 @@ impl<'a> Buf<'a> {
             newlines_to_flush: 0,
             beginning_of_line: true,
             flags,
+            arena,
         }
     }
 
@@ -59,6 +61,10 @@ impl<'a> Buf<'a> {
         self.text.into_bump_str()
     }
 
+    pub fn bump(&'a self) -> &'a Bump {
+        self.text.bump()
+    }
+
     pub fn indent(&mut self, indent: u16) {
         if self.beginning_of_line {
             self.line_indent = indent;

crates/compiler/fmt/src/pattern.rs

@@ -7,9 +7,9 @@ use crate::spaces::{fmt_comments_only, fmt_spaces, NewlineAt, INDENT};
 use crate::Buf;
 use bumpalo::Bump;
 use roc_parse::ast::{
-    Base, CommentOrNewline, Pattern, PatternAs, Spaceable, Spaces, SpacesAfter, SpacesBefore,
+    merge_spaces, Base, CommentOrNewline, Pattern, PatternAs, Spaceable, Spaces, SpacesAfter,
+    SpacesBefore,
 };
-use roc_parse::expr::merge_spaces;
 use roc_region::all::Loc;
 
 pub fn fmt_pattern<'a>(buf: &mut Buf, pattern: &'a Pattern<'a>, indent: u16, parens: Parens) {