From 651ebdb3a0d5a4b6b3b890c80caf3db5f356dd2b Mon Sep 17 00:00:00 2001
From: Igor <ig_@tutamail.com>
Date: Sat, 26 Oct 2024 23:37:08 +0400
Subject: [PATCH] feat: implement `|&` extension (#240)

* Implements support for bash `|&` operator
* Adds basic compatibility test for `|&` usage
---
 brush-parser/src/parser.rs            | 80 ++++++++++++++++++++++++++-
 brush-parser/src/tokenizer.rs         |  2 +-
 brush-shell/tests/cases/pipeline.yaml |  5 ++
 3 files changed, 84 insertions(+), 3 deletions(-)

diff --git a/brush-parser/src/parser.rs b/brush-parser/src/parser.rs
index 734d363e..5979ce55 100644
--- a/brush-parser/src/parser.rs
+++ b/brush-parser/src/parser.rs
@@ -254,8 +254,22 @@ peg::parser! {
             bang:bang()? seq:pipe_sequence() { ast::Pipeline { bang: bang.is_some(), seq } }
         rule bang() -> bool = specific_word("!") { true }
 
-        rule pipe_sequence() -> Vec<ast::Command> =
-            c:command() ++ (specific_operator("|") linebreak()) { c }
+        pub(crate) rule pipe_sequence() -> Vec<ast::Command> =
+            c:(c:command() r:&pipe_extension_redirection()? {? // check for `|&` without consuming the stream.
+                let mut c = c;
+                if r.is_some() {
+                    add_pipe_extension_redirection(&mut c)?;
+                }
+                Ok(c)
+            }) ++ (pipe_operator() linebreak()) {
+            c
+        }
+        rule pipe_operator() =
+            specific_operator("|") /
+            pipe_extension_redirection()
+
+        rule pipe_extension_redirection() -> &'input Token  =
+            non_posix_extensions_enabled() p:specific_operator("|&") { p }
 
         // N.B. We needed to move the function definition branch up to avoid conflicts with array assignment syntax.
         rule command() -> ast::Command =
@@ -826,6 +840,37 @@ fn parse_assignment_word(word: &str) -> Result<ast::Assignment, &'static str> {
     parse_result.map_err(|_| "not assignment word")
 }
 
+// add `2>&1` to the command if the pipeline is `|&`
+fn add_pipe_extension_redirection(c: &mut ast::Command) -> Result<(), &'static str> {
+    let r = ast::IoRedirect::File(
+        Some(2),
+        ast::IoFileRedirectKind::DuplicateOutput,
+        ast::IoFileRedirectTarget::Fd(1),
+    );
+    match c {
+        ast::Command::Simple(c) => {
+            let r = ast::CommandPrefixOrSuffixItem::IoRedirect(r);
+            if let Some(suffix) = &mut c.suffix {
+                suffix.0.push(r);
+            } else {
+                let v = vec![r];
+                c.suffix = Some(ast::CommandSuffix(v));
+            }
+        }
+        ast::Command::Compound(_, l) => {
+            if let Some(r_list) = l {
+                r_list.0.push(r);
+            } else {
+                let v = vec![r];
+                *l = Some(ast::RedirectList(v));
+            }
+        }
+        ast::Command::ExtendedTest(_) => return Err("|& unimplemented for extended tests"),
+        ast::Command::Function(_) => return Err("|& unimplemented for functions"),
+    };
+    Ok(())
+}
+
 fn parse_array_assignment(
     word: &str,
     elements: &[&String],
@@ -858,9 +903,11 @@ fn parse_array_assignment(
 
 #[cfg(test)]
 mod tests {
+
     use super::*;
     use crate::tokenizer::tokenize_str;
     use anyhow::Result;
+    use assert_matches::assert_matches;
 
     #[test]
     fn parse_case() -> Result<()> {
@@ -911,4 +958,33 @@ esac\
 
         Ok(())
     }
+
+    #[test]
+    fn parse_redirection() -> Result<()> {
+        let input = r"echo |& wc";
+
+        let tokens = tokenize_str(input)?;
+        let seq = super::token_parser::pipe_sequence(
+            &Tokens {
+                tokens: tokens.as_slice(),
+            },
+            &ParserOptions::default(),
+            &SourceInfo::default(),
+        )?;
+
+        assert_eq!(seq.len(), 2);
+        assert_matches!(seq[0], ast::Command::Simple(..));
+        if let ast::Command::Simple(c) = &seq[0] {
+            let c = c.suffix.as_ref().unwrap();
+            assert_matches!(
+                c.0[0],
+                ast::CommandPrefixOrSuffixItem::IoRedirect(ast::IoRedirect::File(
+                    Some(2),
+                    ast::IoFileRedirectKind::DuplicateOutput,
+                    ast::IoFileRedirectTarget::Fd(1)
+                ))
+            )
+        }
+        Ok(())
+    }
 }
diff --git a/brush-parser/src/tokenizer.rs b/brush-parser/src/tokenizer.rs
index 50170477..79ee1d19 100644
--- a/brush-parser/src/tokenizer.rs
+++ b/brush-parser/src/tokenizer.rs
@@ -936,7 +936,7 @@ impl<'a, R: ?Sized + std::io::BufRead> Tokenizer<'a, R> {
 
     fn is_operator(&self, s: &str) -> bool {
         // Handle non-POSIX operators.
-        if !self.options.posix_mode && matches!(s, "<<<" | "&>" | "&>>" | ";;&" | ";&") {
+        if !self.options.posix_mode && matches!(s, "<<<" | "&>" | "&>>" | ";;&" | ";&" | "|&") {
             return true;
         }
 
diff --git a/brush-shell/tests/cases/pipeline.yaml b/brush-shell/tests/cases/pipeline.yaml
index 4a60272d..6b7a39d9 100644
--- a/brush-shell/tests/cases/pipeline.yaml
+++ b/brush-shell/tests/cases/pipeline.yaml
@@ -36,3 +36,8 @@ cases:
       echo "var: ${var}"
       echo hi | { var=3; cat; }
       echo "var: ${var}"
+
+  - name: "pipeline extension"
+    stdin: |
+      echo -e "hello" |& wc -l
+      cat dfdfgdfgdf |& wc -l