3 files changed, 67 insertions, 41 deletions

diff --git a/helix-core/src/comment.rs b/helix-core/src/comment.rs
index 4fcece57..6fc1234d 100644
--- a/helix-core/src/comment.rs
+++ b/helix-core/src/comment.rs
@@ -1,17 +1,27 @@
 use crate::{
     find_first_non_whitespace_char, Change, Rope, RopeSlice, Selection, Tendril, Transaction,
 };
-use core::ops::Range;
 use std::borrow::Cow;
 
+/// Given text, a comment token, and a set of line indices, returns the following:
+/// - Whether the given lines should be considered commented
+///     - If any of the lines are uncommented, all lines are considered as such.
+/// - The lines to change for toggling comments
+///     - This is all provided lines excluding blanks lines.
+/// - The column of the comment tokens
+///     - Column of existing tokens, if the lines are commented; column to place tokens at otherwise.
+/// - The margin to the right of the comment tokens
+///     - Defaults to `1`. If any existing comment token is not followed by a space, changes to `0`.
 fn find_line_comment(
     token: &str,
     text: RopeSlice,
-    lines: Range<usize>,
-) -> (bool, Vec<usize>, usize) {
+    lines: impl IntoIterator<Item = usize>,
+) -> (bool, Vec<usize>, usize, usize) {
     let mut commented = true;
-    let mut skipped = Vec::new();
+    let mut to_change = Vec::new();
     let mut min = usize::MAX; // minimum col for find_first_non_whitespace_char
+    let mut margin = 1;
+    let token_len = token.chars().count();
     for line in lines {
         let line_slice = text.line(line);
         if let Some(pos) = find_first_non_whitespace_char(line_slice) {
@@ -29,46 +39,55 @@ fn find_line_comment(
                 // considered uncommented.
                 commented = false;
             }
-        } else {
-            // blank line
-            skipped.push(line);
+
+            // determine margin of 0 or 1 for uncommenting; if any comment token is not followed by a space,
+            // a margin of 0 is used for all lines.
+            if matches!(line_slice.get_char(pos + token_len), Some(c) if c != ' ') {
+                margin = 0;
+            }
+
+            // blank lines don't get pushed.
+            to_change.push(line);
         }
     }
-    (commented, skipped, min)
+    (commented, to_change, min, margin)
 }
 
 #[must_use]
 pub fn toggle_line_comments(doc: &Rope, selection: &Selection, token: Option<&str>) -> Transaction {
     let text = doc.slice(..);
-    let mut changes: Vec<Change> = Vec::new();
 
     let token = token.unwrap_or("//");
     let comment = Tendril::from(format!("{} ", token));
 
+    let mut lines: Vec<usize> = Vec::new();
+
+    let mut min_next_line = 0;
     for selection in selection {
         let (start, end) = selection.line_range(text);
-        let lines = start..end + 1;
-        let (commented, skipped, min) = find_line_comment(&token, text, lines.clone());
+        let start = start.max(min_next_line).min(text.len_lines());
+        let end = (end + 1).min(text.len_lines());
 
-        changes.reserve((end - start).saturating_sub(skipped.len()));
+        lines.extend(start..end);
+        min_next_line = end + 1;
+    }
 
-        for line in lines {
-            if skipped.contains(&line) {
-                continue;
-            }
+    let (commented, to_change, min, margin) = find_line_comment(&token, text, lines);
 
-            let pos = text.line_to_char(line) + min;
+    let mut changes: Vec<Change> = Vec::with_capacity(to_change.len());
 
-            if !commented {
-                // comment line
-                changes.push((pos, pos, Some(comment.clone())))
-            } else {
-                // uncomment line
-                let margin = 1; // TODO: margin is hardcoded 1 but could easily be 0
-                changes.push((pos, pos + token.len() + margin, None))
-            }
+    for line in to_change {
+        let pos = text.line_to_char(line) + min;
+
+        if !commented {
+            // comment line
+            changes.push((pos, pos, Some(comment.clone())));
+        } else {
+            // uncomment line
+            changes.push((pos, pos + token.len() + margin, None));
         }
     }
+
     Transaction::change(doc, changes.into_iter())
 }
 
@@ -90,23 +109,32 @@ mod test {
         let text = state.doc.slice(..);
 
         let res = find_line_comment("//", text, 0..3);
-        // (commented = true, skipped = [line 1], min = col 2)
-        assert_eq!(res, (false, vec![1], 2));
+        // (commented = true, to_change = [line 0, line 2], min = col 2, margin = 1)
+        assert_eq!(res, (false, vec![0, 2], 2, 1));
 
         // comment
         let transaction = toggle_line_comments(&state.doc, &state.selection, None);
         transaction.apply(&mut state.doc);
-        state.selection = state.selection.clone().map(transaction.changes());
+        state.selection = state.selection.map(transaction.changes());
 
         assert_eq!(state.doc, "  // 1\n\n  // 2\n  // 3");
 
         // uncomment
         let transaction = toggle_line_comments(&state.doc, &state.selection, None);
         transaction.apply(&mut state.doc);
-        state.selection = state.selection.clone().map(transaction.changes());
+        state.selection = state.selection.map(transaction.changes());
+        assert_eq!(state.doc, "  1\n\n  2\n  3");
+
+        // 0 margin comments
+        state.doc = Rope::from("  //1\n\n  //2\n  //3");
+        // reset the selection.
+        state.selection = Selection::single(0, state.doc.len_chars() - 1);
+
+        let transaction = toggle_line_comments(&state.doc, &state.selection, None);
+        transaction.apply(&mut state.doc);
+        state.selection = state.selection.map(transaction.changes());
         assert_eq!(state.doc, "  1\n\n  2\n  3");
 
-        // TODO: account for no margin after comment
         // TODO: account for uncommenting with uneven comment indentation
     }
 }
diff --git a/helix-core/src/indent.rs b/helix-core/src/indent.rs
index 0ca05fb3..5ae66769 100644
--- a/helix-core/src/indent.rs
+++ b/helix-core/src/indent.rs
@@ -47,7 +47,7 @@ fn calculate_indentation(query: &IndentQuery, node: Option<Node>, newline: bool)
     // NOTE: can't use contains() on query because of comparing Vec<String> and &str
     // https://doc.rust-lang.org/std/vec/struct.Vec.html#method.contains
 
-    let mut increment: i32 = 0;
+    let mut increment: isize = 0;
 
     let mut node = match node {
         Some(node) => node,
@@ -93,9 +93,7 @@ fn calculate_indentation(query: &IndentQuery, node: Option<Node>, newline: bool)
         node = parent;
     }
 
-    assert!(increment >= 0);
-
-    increment as usize
+    increment.max(0) as usize
 }
 
 #[allow(dead_code)]
diff --git a/helix-core/src/register.rs b/helix-core/src/register.rs
index cc881a17..c3e6652e 100644
--- a/helix-core/src/register.rs
+++ b/helix-core/src/register.rs
@@ -22,13 +22,17 @@ impl Register {
         self.name
     }
 
-    pub fn read(&self) -> &Vec<String> {
+    pub fn read(&self) -> &[String] {
         &self.values
     }
 
     pub fn write(&mut self, values: Vec<String>) {
         self.values = values;
     }
+
+    pub fn push(&mut self, value: String) {
+        self.values.push(value);
+    }
 }
 
 /// Currently just wraps a `HashMap` of `Register`s
@@ -42,11 +46,7 @@ impl Registers {
         self.inner.get(&name)
     }
 
-    pub fn get_mut(&mut self, name: char) -> Option<&mut Register> {
-        self.inner.get_mut(&name)
-    }
-
-    pub fn get_or_insert(&mut self, name: char) -> &mut Register {
+    pub fn get_mut(&mut self, name: char) -> &mut Register {
         self.inner
             .entry(name)
             .or_insert_with(|| Register::new(name))
@@ -57,7 +57,7 @@ impl Registers {
             .insert(name, Register::new_with_values(name, values));
     }
 
-    pub fn read(&self, name: char) -> Option<&Vec<String>> {
+    pub fn read(&self, name: char) -> Option<&[String]> {
         self.get(name).map(|reg| reg.read())
     }
 }