Continue comments when opening a new line

helix-core/src/comment.rs

@@ -1,8 +1,14 @@
-//! This module contains the functionality toggle comments on lines over the selection
-//! using the comment character defined in the user's `languages.toml`
+//! This module contains the some comment-related features
+//! using the comment character defined in the user's `languages.toml`:
+//! * toggle comments on lines over the selection.
+//! * continue comment when opening a new line.
+
+use tree_sitter::QueryCursor;
 
 use crate::{
-    find_first_non_whitespace_char, Change, Rope, RopeSlice, Selection, Tendril, Transaction,
+    find_first_non_whitespace_char,
+    syntax::{CapturedNode, LanguageConfiguration},
+    Change, Range, Rope, RopeSlice, Selection, Syntax, Tendril, Transaction,
 };
 use std::borrow::Cow;
 
@@ -94,6 +100,107 @@ pub fn toggle_line_comments(doc: &Rope, selection: &Selection, token: Option<&st
     Transaction::change(doc, changes.into_iter())
 }
 
+/// Return token if the current line is commented.
+/// Otherwise, return None.
+pub fn continue_comment<'a>(doc: &Rope, line: usize, tokens: &'a [String]) -> Option<&'a str> {
+    // TODO: don't continue shebangs.
+    if tokens.is_empty() {
+        return None;
+    }
+
+    let mut result = None;
+    let line_slice = doc.line(line);
+    if let Some(pos) = find_first_non_whitespace_char(line_slice) {
+        let len = line_slice.len_chars();
+        for token in tokens {
+            // line can be shorter than pos + token len
+            let fragment = Cow::from(line_slice.slice(pos..std::cmp::min(pos + token.len(), len)));
+            if fragment == *token {
+                // Purposefully not break here to overwrite the result when a longer comment token
+                // matches.
+                result = Some(token.as_str());
+            }
+        }
+    }
+
+    result
+}
+
+pub fn continue_block_comment<'a>(
+    doc: &Rope,
+    syntax: Option<&Syntax>,
+    lang_config: &'a LanguageConfiguration,
+    range: &Range,
+    open_below: bool,
+) -> Option<&'a str> {
+    if let Some((doc_syntax, block_comment_tokens)) =
+        syntax.zip(lang_config.block_comment_tokens.as_ref())
+    {
+        let slice_tree = doc_syntax.tree().root_node();
+        let slice = doc.slice(..);
+        let line_pos = slice.char_to_line(range.cursor(slice));
+        let mut cursor = QueryCursor::new();
+        let mut found_block_comments = false;
+
+        let should_insert_comment_middle = |node: CapturedNode| {
+            let node_start = doc.byte_to_line(node.start_byte());
+            let node_end = doc.byte_to_line(node.end_byte());
+            // NOTE: we use line comparison to allow opening a comment on the newline after the
+            // end of the block comment.
+            // We also do not want to continue the comment if opening below when the cursor is
+            // on the last line of the block comment.
+            line_pos >= node_start && (line_pos < node_end || (line_pos == node_end && !open_below))
+        };
+
+        {
+            let nodes = lang_config.textobject_query().and_then(|query| {
+                query.capture_nodes_any(&["comment.block.around"], slice_tree, slice, &mut cursor)
+            });
+            if let Some(nodes) = nodes {
+                for node in nodes {
+                    found_block_comments = true;
+                    if should_insert_comment_middle(node) {
+                        return Some(&block_comment_tokens.middle);
+                    }
+                }
+            }
+        }
+
+        // Many tree-sitter grammars don't contain a block comment token, so search the comment
+        // token and check that it starts and ends with the correct block comment tokens.
+        // FIXME: this doesn't take into account that a line comment followed by a block comment
+        // counts as only one comment object.
+        // TODO: Maybe that's caused by the + in the query?
+        if !found_block_comments {
+            let nodes = lang_config.textobject_query().and_then(|query| {
+                query.capture_nodes_any(&["comment.around"], slice_tree, slice, &mut cursor)
+            });
+            if let Some(nodes) = nodes {
+                for node in nodes {
+                    let node_start = doc.byte_to_char(node.start_byte());
+                    let node_end = doc.byte_to_char(node.end_byte());
+                    let start_token_len = block_comment_tokens.start.len();
+                    let end_token_len = block_comment_tokens.end.len();
+                    if doc.len_chars() > node_start + start_token_len && doc.len_chars() > node_end
+                    {
+                        let comment_start = doc.slice(node_start..node_start + start_token_len);
+                        let comment_end = doc.slice(node_end - end_token_len..node_end);
+                        if comment_start == block_comment_tokens.start
+                            && comment_end == block_comment_tokens.end
+                        {
+                            // We're in a block comment.
+                            if should_insert_comment_middle(node) {
+                                return Some(&block_comment_tokens.middle);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+    None
+}
+
 #[cfg(test)]
 mod test {
     use super::*;

helix-core/src/syntax.rs

@@ -55,6 +55,13 @@ pub struct Configuration {
     pub language: Vec<LanguageConfiguration>,
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+pub struct BlockCommentTokens {
+    pub start: String,
+    pub middle: String,
+    pub end: String,
+}
+
 // largely based on tree-sitter/cli/src/loader.rs
 #[derive(Debug, Serialize, Deserialize)]
 #[serde(rename_all = "kebab-case", deny_unknown_fields)]
@@ -66,7 +73,10 @@ pub struct LanguageConfiguration {
     #[serde(default)]
     pub shebangs: Vec<String>, // interpreter(s) associated with language
     pub roots: Vec<String>,      // these indicate project roots <.git, Cargo.toml>
-    pub comment_token: Option<String>,
+    pub comment_token: Option<String>, // TODO: remove.
+    #[serde(default)]
+    pub comment_tokens: Vec<String>,
+    pub block_comment_tokens: Option<BlockCommentTokens>,
 
     #[serde(default, skip_serializing, deserialize_with = "deserialize_lsp_config")]
     pub config: Option<serde_json::Value>,

helix-term/src/commands.rs

@@ -2241,6 +2241,7 @@ async fn make_format_callback(
     Ok(call)
 }
 
+#[derive(PartialEq)]
 enum Open {
     Below,
     Above,
@@ -2293,6 +2294,9 @@ fn open(cx: &mut Context, open: Open) {
         let mut text = String::with_capacity(1 + indent_len);
         text.push_str(doc.line_ending.as_str());
         text.push_str(&indent);
+
+        handle_comment_continue(doc, range, &mut text, cursor_line, open == Open::Below);
+
         let text = text.repeat(count);
 
         // calculate new selection ranges
@@ -2312,6 +2316,51 @@ fn open(cx: &mut Context, open: Open) {
     transaction = transaction.with_selection(Selection::new(ranges, selection.primary_index()));
 
     doc.apply(&transaction, view.id);
+
+    // Since we might have added a comment token, move to the end of the line.
+    goto_line_end_newline(cx);
+}
+
+fn handle_comment_continue(
+    doc: &Document,
+    range: &Range,
+    text: &mut String,
+    cursor_line: usize,
+    open_below: bool,
+) {
+    if let Some(lang_config) = doc.language_config() {
+        let line_comment_tokens = &lang_config.comment_tokens;
+        if let Some(token) = comment::continue_block_comment(
+            doc.text(),
+            doc.syntax(),
+            lang_config,
+            range,
+            open_below,
+        ) {
+            text.push(' ');
+            text.push_str(token);
+            text.push(' ');
+        } else if let Some(token) =
+            comment::continue_comment(doc.text(), cursor_line, line_comment_tokens)
+        {
+            text.push_str(token);
+            text.push(' ');
+        } else if let Some(ref block_comment_tokens) = lang_config.block_comment_tokens {
+            // FIXME: this doesn't work for the following lines of a comment start.
+            // FIXME: this is too indented in some cases.
+            if comment::continue_comment(
+                doc.text(),
+                cursor_line,
+                &[block_comment_tokens.start.clone()],
+            )
+            .is_some()
+            {
+                text.push(' ');
+                text.push_str(&block_comment_tokens.middle);
+                text.push(' ');
+            }
+        }
+    }
 }
 
 // o inserts a new line after each line with a selection
@@ -2774,6 +2823,9 @@ pub mod insert {
                 text.reserve_exact(1 + indent.len());
                 text.push_str(doc.line_ending.as_str());
                 text.push_str(&indent);
+
+                handle_comment_continue(doc, range, &mut text, current_line, false);
+
                 pos + offs + text.chars().count()
             };
 
@@ -3560,7 +3612,7 @@ fn toggle_comments(cx: &mut Context) {
     let (view, doc) = current!(cx.editor);
     let token = doc
         .language_config()
-        .and_then(|lc| lc.comment_token.as_ref())
+        .and_then(|lc| lc.comment_tokens.get(0))
         .map(|tc| tc.as_ref());
     let transaction = comment::toggle_line_comments(doc.text(), doc.selection(view.id), token);