Add `Buffer::autoindent_for_row`

Antonio Scandurra created

Change summary

zed/src/editor/buffer/mod.rs | 132 ++++++++++++++++++++++++++++++++++++-
1 file changed, 126 insertions(+), 6 deletions(-)

Detailed changes

zed/src/editor/buffer/mod.rs 🔗

@@ -674,6 +674,72 @@ impl Buffer {
         })
     }
 
+    fn indent_for_row(&self, row: u32) -> (usize, bool) {
+        let mut is_whitespace = true;
+        let mut indent = 0;
+        for c in self.chars_at(Point::new(row, 0)) {
+            match c {
+                ' ' => indent += 1,
+                '\n' => break,
+                _ => {
+                    is_whitespace = false;
+                    break;
+                }
+            }
+        }
+        (indent, is_whitespace)
+    }
+
+    fn autoindent_for_row(&self, row: u32) -> usize {
+        let mut indent_parent = None;
+        let mut indent = 2;
+        if let Some((language, syntax_tree)) = self.language.as_ref().zip(self.syntax_tree()) {
+            let row_start = Point::new(row, 0).into();
+            let mut cursor = syntax_tree.walk();
+            loop {
+                let node = cursor.node();
+                if row_start >= node.end_position() {
+                    if !cursor.goto_next_sibling() {
+                        break;
+                    }
+                } else if node.start_position() > row_start {
+                    break;
+                } else {
+                    if node.start_position().row as u32 != row
+                        && language.config.indent_nodes.contains(node.kind())
+                    {
+                        let parent_ends_at_row = node.end_position().row as u32 == row;
+                        indent_parent =
+                            Some((node.start_position().row as u32, parent_ends_at_row));
+                    }
+
+                    if !cursor.goto_first_child() {
+                        break;
+                    }
+                }
+            }
+
+            indent = language.config.indent;
+        }
+
+        if let Some((parent_row, parent_ends_at_row)) = indent_parent {
+            let (parent_indent, _) = self.indent_for_row(parent_row);
+            if parent_ends_at_row {
+                parent_indent
+            } else {
+                parent_indent + indent
+            }
+        } else {
+            for prev_row in (0..row).rev() {
+                let (prev_indent, is_whitespace) = self.indent_for_row(prev_row);
+                if prev_indent != 0 || !is_whitespace {
+                    return prev_indent;
+                }
+            }
+            0
+        }
+    }
+
     fn diff(&self, new_text: Arc<str>, ctx: &AppContext) -> Task<Diff> {
         // TODO: it would be nice to not allocate here.
         let old_text = self.text();
@@ -765,12 +831,6 @@ impl Buffer {
         self.visible_text.max_point()
     }
 
-    pub fn line(&self, row: u32) -> String {
-        self.chars_at(Point::new(row, 0))
-            .take_while(|c| *c != '\n')
-            .collect()
-    }
-
     pub fn text(&self) -> String {
         self.text_for_range(0..self.len()).collect()
     }
@@ -2632,6 +2692,7 @@ impl ToPoint for usize {
 mod tests {
     use super::*;
     use crate::{
+        language::LanguageConfig,
         test::{build_app_state, temp_tree},
         util::RandomCharIter,
         worktree::{Worktree, WorktreeHandle},
@@ -2643,9 +2704,11 @@ mod tests {
         cell::RefCell,
         cmp::Ordering,
         fs,
+        iter::FromIterator as _,
         rc::Rc,
         sync::atomic::{self, AtomicUsize},
     };
+    use unindent::Unindent as _;
 
     #[gpui::test]
     fn test_edit(ctx: &mut gpui::MutableAppContext) {
@@ -3631,6 +3694,63 @@ mod tests {
         }
     }
 
+    #[gpui::test]
+    async fn test_indent(mut ctx: gpui::TestAppContext) {
+        let grammar = tree_sitter_rust::language();
+        let lang = Arc::new(Language {
+            config: LanguageConfig {
+                indent: 3,
+                indent_nodes: std::collections::HashSet::from_iter(vec!["block".to_string()]),
+                ..Default::default()
+            },
+            grammar: grammar.clone(),
+            highlight_query: tree_sitter::Query::new(grammar, "").unwrap(),
+            theme_mapping: Default::default(),
+        });
+
+        let buffer = ctx.add_model(|ctx| {
+            let text = "
+                fn a() {}
+
+                fn b() {
+                }
+
+                fn c() {
+                 let badly_indented_line;
+
+                }
+
+                struct D {
+                    x: 1,
+
+
+                }
+                "
+            .unindent();
+            Buffer::from_history(0, History::new(text.into()), None, Some(lang), ctx)
+        });
+
+        buffer.condition(&ctx, |buf, _| !buf.is_parsing()).await;
+        buffer.read_with(&ctx, |buf, _| {
+            assert_eq!(buf.autoindent_for_row(6), 3);
+            assert_eq!(buf.autoindent_for_row(7), 3);
+
+            // Don't autoindent rows that start or end on the same row as the indent node.
+            assert_eq!(buf.autoindent_for_row(0), 0);
+            assert_eq!(buf.autoindent_for_row(2), 0);
+            assert_eq!(buf.autoindent_for_row(3), 0);
+            assert_eq!(buf.autoindent_for_row(4), 0);
+            assert_eq!(buf.autoindent_for_row(8), 0);
+
+            // We didn't find any matching indent node in the language for the struct definition.
+            // Don't autoindent the first line inside of the struct. Instead, align the second and
+            // third line to the first non-whitespace line preceding them.
+            assert_eq!(buf.autoindent_for_row(11), 0);
+            assert_eq!(buf.autoindent_for_row(12), 4);
+            assert_eq!(buf.autoindent_for_row(13), 4);
+        });
+    }
+
     impl Buffer {
         fn random_byte_range(&mut self, start_offset: usize, rng: &mut impl Rng) -> Range<usize> {
             let end = self.clip_offset(rng.gen_range(start_offset..=self.len()), Bias::Right);