Rework the replacements flag (#267)

* Add replacements tests

* Honor `-n` when using `--preview`

* Rework CLI for replacements flag

* Remove dead code

* Remove lingering TODO

10 files changed, 157 insertions, 41 deletions

diff --git a/gen/completions/_sd b/gen/completions/_sd
index e8f2804..1a8c150 100644
--- a/gen/completions/_sd
+++ b/gen/completions/_sd
@@ -15,7 +15,8 @@ _sd() {
 
     local context curcontext="$curcontext" state line
     _arguments "${_arguments_options[@]}" \
-'-n+[Limit the number of replacements]:REPLACEMENTS: ' \
+'-n+[Limit the number of replacements that can occur per file. 0 indicates unlimited replacements]:LIMIT: ' \
+'--max-replacements=[Limit the number of replacements that can occur per file. 0 indicates unlimited replacements]:LIMIT: ' \
 '-f+[Regex flags. May be combined (like \`-f mc\`).]:FLAGS: ' \
 '--flags=[Regex flags. May be combined (like \`-f mc\`).]:FLAGS: ' \
 '-p[Display changes in a human reviewable format (the specifics of the format are likely to change in the future)]' \
diff --git a/gen/completions/_sd.ps1 b/gen/completions/_sd.ps1
index 4704172..4ac2dfb 100644
--- a/gen/completions/_sd.ps1
+++ b/gen/completions/_sd.ps1
@@ -21,7 +21,8 @@ Register-ArgumentCompleter -Native -CommandName 'sd' -ScriptBlock {
 
     $completions = @(switch ($command) {
         'sd' {
-            [CompletionResult]::new('-n', 'n', [CompletionResultType]::ParameterName, 'Limit the number of replacements')
+            [CompletionResult]::new('-n', 'n', [CompletionResultType]::ParameterName, 'Limit the number of replacements that can occur per file. 0 indicates unlimited replacements')
+            [CompletionResult]::new('--max-replacements', 'max-replacements', [CompletionResultType]::ParameterName, 'Limit the number of replacements that can occur per file. 0 indicates unlimited replacements')
             [CompletionResult]::new('-f', 'f', [CompletionResultType]::ParameterName, 'Regex flags. May be combined (like `-f mc`).')
             [CompletionResult]::new('--flags', 'flags', [CompletionResultType]::ParameterName, 'Regex flags. May be combined (like `-f mc`).')
             [CompletionResult]::new('-p', 'p', [CompletionResultType]::ParameterName, 'Display changes in a human reviewable format (the specifics of the format are likely to change in the future)')
diff --git a/gen/completions/sd.bash b/gen/completions/sd.bash
index 78eae05..b3e8b71 100644
--- a/gen/completions/sd.bash
+++ b/gen/completions/sd.bash
@@ -19,12 +19,16 @@ _sd() {
 
     case "${cmd}" in
         sd)
-            opts="-p -F -n -f -h -V --preview --fixed-strings --flags --help --version <FIND> <REPLACE_WITH> [FILES]..."
+            opts="-p -F -n -f -h -V --preview --fixed-strings --max-replacements --flags --help --version <FIND> <REPLACE_WITH> [FILES]..."
             if [[ ${cur} == -* || ${COMP_CWORD} -eq 1 ]] ; then
                 COMPREPLY=( $(compgen -W "${opts}" -- "${cur}") )
                 return 0
             fi
             case "${prev}" in
+                --max-replacements)
+                    COMPREPLY=($(compgen -f "${cur}"))
+                    return 0
+                    ;;
                 -n)
                     COMPREPLY=($(compgen -f "${cur}"))
                     return 0
diff --git a/gen/completions/sd.elv b/gen/completions/sd.elv
index e455cc0..d737837 100644
--- a/gen/completions/sd.elv
+++ b/gen/completions/sd.elv
@@ -18,7 +18,8 @@ set edit:completion:arg-completer[sd] = {|@words|
     }
     var completions = [
         &'sd'= {
-            cand -n 'Limit the number of replacements'
+            cand -n 'Limit the number of replacements that can occur per file. 0 indicates unlimited replacements'
+            cand --max-replacements 'Limit the number of replacements that can occur per file. 0 indicates unlimited replacements'
             cand -f 'Regex flags. May be combined (like `-f mc`).'
             cand --flags 'Regex flags. May be combined (like `-f mc`).'
             cand -p 'Display changes in a human reviewable format (the specifics of the format are likely to change in the future)'
diff --git a/gen/completions/sd.fish b/gen/completions/sd.fish
index a0e03ab..7d324ff 100644
--- a/gen/completions/sd.fish
+++ b/gen/completions/sd.fish
@@ -1,4 +1,4 @@
-complete -c sd -s n -d 'Limit the number of replacements' -r
+complete -c sd -s n -l max-replacements -d 'Limit the number of replacements that can occur per file. 0 indicates unlimited replacements' -r
 complete -c sd -s f -l flags -d 'Regex flags. May be combined (like `-f mc`).' -r
 complete -c sd -s p -l preview -d 'Display changes in a human reviewable format (the specifics of the format are likely to change in the future)'
 complete -c sd -s F -l fixed-strings -d 'Treat FIND and REPLACE_WITH args as literal strings'
diff --git a/gen/sd.1 b/gen/sd.1
index 3c3916b..a60e1c8 100644
--- a/gen/sd.1
+++ b/gen/sd.1
@@ -8,7 +8,7 @@ sd
 .ie \n(.g .ds Aq \(aq
 .el .ds Aq '
 .SH SYNOPSIS
-\fBsd\fR [\fB\-p\fR|\fB\-\-preview\fR] [\fB\-F\fR|\fB\-\-fixed\-strings\fR] [\fB\-n \fR] [\fB\-f\fR|\fB\-\-flags\fR] [\fB\-h\fR|\fB\-\-help\fR] [\fB\-V\fR|\fB\-\-version\fR] <\fIFIND\fR> <\fIREPLACE_WITH\fR> [\fIFILES\fR] 
+\fBsd\fR [\fB\-p\fR|\fB\-\-preview\fR] [\fB\-F\fR|\fB\-\-fixed\-strings\fR] [\fB\-n\fR|\fB\-\-max\-replacements\fR] [\fB\-f\fR|\fB\-\-flags\fR] [\fB\-h\fR|\fB\-\-help\fR] [\fB\-V\fR|\fB\-\-version\fR] <\fIFIND\fR> <\fIREPLACE_WITH\fR> [\fIFILES\fR] 
 .ie \n(.g .ds Aq \(aq
 .el .ds Aq '
 .SH DESCRIPTION
@@ -22,8 +22,8 @@ Display changes in a human reviewable format (the specifics of the format are li
 \fB\-F\fR, \fB\-\-fixed\-strings\fR
 Treat FIND and REPLACE_WITH args as literal strings
 .TP
-\fB\-n\fR=\fIREPLACEMENTS\fR
-Limit the number of replacements
+\fB\-n\fR, \fB\-\-max\-replacements\fR=\fILIMIT\fR [default: 0]
+Limit the number of replacements that can occur per file. 0 indicates unlimited replacements
 .TP
 \fB\-f\fR, \fB\-\-flags\fR=\fIFLAGS\fR
 Regex flags. May be combined (like `\-f mc`).
diff --git a/src/cli.rs b/src/cli.rs
index dff7678..e96e609 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -29,9 +29,15 @@ pub struct Options {
     /// Treat FIND and REPLACE_WITH args as literal strings
     pub literal_mode: bool,
 
-    #[arg(short = 'n')]
-    /// Limit the number of replacements
-    pub replacements: Option<usize>,
+    #[arg(
+        short = 'n',
+        long = "max-replacements",
+        value_name = "LIMIT",
+        default_value_t
+    )]
+    /// Limit the number of replacements that can occur per file. 0 indicates
+    /// unlimited replacements.
+    pub replacements: usize,
 
     #[arg(short, long, verbatim_doc_comment)]
     #[rustfmt::skip]
diff --git a/src/replacer/mod.rs b/src/replacer/mod.rs
index b6e0c0b..d0a90a2 100644
--- a/src/replacer/mod.rs
+++ b/src/replacer/mod.rs
@@ -1,4 +1,4 @@
-use std::{fs, fs::File, io::prelude::*, path::Path};
+use std::{borrow::Cow, fs, fs::File, io::prelude::*, path::Path};
 
 use crate::{utils, Error, Result};
 
@@ -23,7 +23,7 @@ impl Replacer {
         replace_with: String,
         is_literal: bool,
         flags: Option<String>,
-        replacements: Option<usize>,
+        replacements: usize,
     ) -> Result<Self> {
         let (look_for, replace_with) = if is_literal {
             (regex::escape(&look_for), replace_with.into_bytes())
@@ -70,7 +70,7 @@ impl Replacer {
             regex: regex.build()?,
             replace_with,
             is_literal,
-            replacements: replacements.unwrap_or(0),
+            replacements,
         })
     }
 
@@ -88,46 +88,92 @@ impl Replacer {
         &'a self,
         content: &'a [u8],
     ) -> std::borrow::Cow<'a, [u8]> {
+        let regex = &self.regex;
+        let limit = self.replacements;
+        let use_color = false;
         if self.is_literal {
-            self.regex.replacen(
+            Self::replacen(
+                regex,
+                limit,
                 content,
-                self.replacements,
+                use_color,
                 regex::bytes::NoExpand(&self.replace_with),
             )
         } else {
-            self.regex
-                .replacen(content, self.replacements, &*self.replace_with)
+            Self::replacen(
+                regex,
+                limit,
+                content,
+                use_color,
+                &*self.replace_with,
+            )
         }
     }
 
-    pub(crate) fn replace_preview<'a>(
-        &'a self,
-        content: &[u8],
-    ) -> std::borrow::Cow<'a, [u8]> {
-        let mut v = Vec::<u8>::new();
-        let mut captures = self.regex.captures_iter(content);
-
-        self.regex.split(content).for_each(|sur_text| {
-            use regex::bytes::Replacer;
-
-            v.extend(sur_text);
-            if let Some(capture) = captures.next() {
-                v.extend_from_slice(
+    /// A modified form of [`regex::bytes::Regex::replacen`] that supports
+    /// coloring replacements
+    pub(crate) fn replacen<'haystack, R: regex::bytes::Replacer>(
+        regex: &regex::bytes::Regex,
+        limit: usize,
+        haystack: &'haystack [u8],
+        use_color: bool,
+        mut rep: R,
+    ) -> Cow<'haystack, [u8]> {
+        let mut it = regex.captures_iter(haystack).enumerate().peekable();
+        if it.peek().is_none() {
+            return Cow::Borrowed(haystack);
+        }
+        let mut new = Vec::with_capacity(haystack.len());
+        let mut last_match = 0;
+        for (i, cap) in it {
+            // unwrap on 0 is OK because captures only reports matches
+            let m = cap.get(0).unwrap();
+            new.extend_from_slice(&haystack[last_match..m.start()]);
+            if use_color {
+                new.extend_from_slice(
                     ansi_term::Color::Green.prefix().to_string().as_bytes(),
                 );
-                if self.is_literal {
-                    regex::bytes::NoExpand(&self.replace_with)
-                        .replace_append(&capture, &mut v);
-                } else {
-                    (&*self.replace_with).replace_append(&capture, &mut v);
-                }
-                v.extend_from_slice(
+            }
+            rep.replace_append(&cap, &mut new);
+            if use_color {
+                new.extend_from_slice(
                     ansi_term::Color::Green.suffix().to_string().as_bytes(),
                 );
             }
-        });
+            last_match = m.end();
+            if limit > 0 && i >= limit - 1 {
+                break;
+            }
+        }
+        new.extend_from_slice(&haystack[last_match..]);
+        Cow::Owned(new)
+    }
 
-        return std::borrow::Cow::Owned(v);
+    pub(crate) fn replace_preview<'a>(
+        &self,
+        content: &'a [u8],
+    ) -> std::borrow::Cow<'a, [u8]> {
+        let regex = &self.regex;
+        let limit = self.replacements;
+        // TODO: refine this condition more
+        let use_color = true;
+        if self.is_literal {
+            Self::replacen(
+                regex,
+                limit,
+                content,
+                use_color,
+                regex::bytes::NoExpand(&self.replace_with),
+            )
+        } else {
+            Self::replacen(
+                regex,
+                limit,
+                content,
+                use_color,
+                &*self.replace_with,
+            )
+        }
     }
 
     pub(crate) fn replace_file(&self, path: &Path) -> Result<()> {
diff --git a/src/replacer/tests.rs b/src/replacer/tests.rs
index 9304a17..1db23a1 100644
--- a/src/replacer/tests.rs
+++ b/src/replacer/tests.rs
@@ -29,12 +29,13 @@ fn replace(
     src: &'static str,
     target: &'static str,
 ) {
+    const UNLIMITED_REPLACEMENTS: usize = 0;
     let replacer = Replacer::new(
         look_for.into(),
         replace_with.into(),
         literal,
         flags.map(ToOwned::to_owned),
-        None,
+        UNLIMITED_REPLACEMENTS,
     )
     .unwrap();
     assert_eq!(
diff --git a/tests/cli.rs b/tests/cli.rs
index d0af23f..a727ed5 100644
--- a/tests/cli.rs
+++ b/tests/cli.rs
@@ -194,4 +194,60 @@ mod cli {
           <b>^^^^</b>
         "###);
     }
+
+    #[test]
+    fn limit_replacements_file() -> Result<()> {
+        let mut file = tempfile::NamedTempFile::new()?;
+        file.write_all(b"foo\nfoo\nfoo")?;
+        let path = file.into_temp_path();
+
+        sd().args(["-n", "1", "foo", "bar", path.to_str().unwrap()])
+            .assert()
+            .success();
+        assert_file(&path, "bar\nfoo\nfoo");
+
+        Ok(())
+    }
+
+    #[test]
+    fn limit_replacements_file_preview() -> Result<()> {
+        let mut file = tempfile::NamedTempFile::new()?;
+        file.write_all(b"foo\nfoo\nfoo")?;
+        let path = file.into_temp_path();
+
+        sd().args([
+            "--preview",
+            "-n",
+            "1",
+            "foo",
+            "bar",
+            path.to_str().unwrap(),
+        ])
+        .assert()
+        .success()
+        .stdout(format!(
+            "{}\nfoo\nfoo\n",
+            ansi_term::Color::Green.paint("bar")
+        ));
+
+        Ok(())
+    }
+
+    #[test]
+    fn limit_replacements_stdin() {
+        sd().args(["-n", "1", "foo", "bar"])
+            .write_stdin("foo\nfoo\nfoo")
+            .assert()
+            .success()
+            .stdout("bar\nfoo\nfoo");
+    }
+
+    #[test]
+    fn limit_replacements_stdin_preview() {
+        sd().args(["--preview", "-n", "1", "foo", "bar"])
+            .write_stdin("foo\nfoo\nfoo")
+            .assert()
+            .success()
+            .stdout("bar\nfoo\nfoo");
+    }
 }