From 30b80d0c82634451d855865125c99a54f98aad19 Mon Sep 17 00:00:00 2001
From: Jan Scheer <jhscheer@users.noreply.github.com>
Date: Fri, 9 Jul 2021 22:55:00 +0200
Subject: [PATCH] tr: fix squeeze and complement with two sets (#2485)

---
 src/uu/tr/src/tr.rs      | 26 ++++++++------------------
 tests/by-util/test_tr.rs |  9 +++++++++
 2 files changed, 17 insertions(+), 18 deletions(-)

diff --git a/src/uu/tr/src/tr.rs b/src/uu/tr/src/tr.rs
index 28ce70c22..2abc652ff 100644
--- a/src/uu/tr/src/tr.rs
+++ b/src/uu/tr/src/tr.rs
@@ -173,16 +173,14 @@ struct TranslateAndSqueezeOperation {
 }
 
 impl TranslateAndSqueezeOperation {
-    fn new(
-        set1: ExpandSet,
-        set2: &mut ExpandSet,
-        set2_: ExpandSet,
-        truncate: bool,
-        complement: bool,
-    ) -> TranslateAndSqueezeOperation {
+    fn new(sets: Vec<String>, truncate: bool, complement: bool) -> TranslateAndSqueezeOperation {
+        let set1 = ExpandSet::new(sets[0].as_ref());
+        let set1_ = ExpandSet::new(sets[0].as_ref());
+        let mut set2 = ExpandSet::new(sets[1].as_ref());
+        let set2_ = ExpandSet::new(sets[1].as_ref());
         TranslateAndSqueezeOperation {
-            translate: TranslateOperation::new(set1, set2, truncate, complement),
-            squeeze: SqueezeOperation::new(set2_, complement),
+            translate: TranslateOperation::new(set1, &mut set2, truncate, complement),
+            squeeze: SqueezeOperation::new(if complement { set1_ } else { set2_ }, complement),
         }
     }
 }
@@ -302,15 +300,7 @@ pub fn uumain(args: impl uucore::Args) -> i32 {
             let op = SqueezeOperation::new(set1, complement_flag);
             translate_input(&mut locked_stdin, &mut buffered_stdout, op);
         } else {
-            let mut set2 = ExpandSet::new(sets[1].as_ref());
-            let set2_ = ExpandSet::new(sets[1].as_ref());
-            let op = TranslateAndSqueezeOperation::new(
-                set1,
-                &mut set2,
-                set2_,
-                complement_flag,
-                truncate_flag,
-            );
+            let op = TranslateAndSqueezeOperation::new(sets, truncate_flag, complement_flag);
             translate_input(&mut locked_stdin, &mut buffered_stdout, op);
         }
     } else {
diff --git a/tests/by-util/test_tr.rs b/tests/by-util/test_tr.rs
index 936af2ca8..07e3a843e 100644
--- a/tests/by-util/test_tr.rs
+++ b/tests/by-util/test_tr.rs
@@ -127,6 +127,15 @@ fn test_squeeze_complement() {
         .stdout_is("aaBcDcc");
 }
 
+#[test]
+fn test_squeeze_complement_two_sets() {
+    new_ucmd!()
+        .args(&["-sc", "a", "_"])
+        .pipe_in("test a aa with 3 ___ spaaaces +++") // spell-checker:disable-line
+        .run()
+        .stdout_is("_a_aa_aaa_");
+}
+
 #[test]
 fn test_translate_and_squeeze() {
     new_ucmd!()