diff --git a/src/uu/tr/src/operation.rs b/src/uu/tr/src/operation.rs index 504f8ac3a..04850aabf 100644 --- a/src/uu/tr/src/operation.rs +++ b/src/uu/tr/src/operation.rs @@ -99,17 +99,11 @@ impl Sequence { } } - pub fn last(&self) -> Option { - match self { - Sequence::CharStar(c) => Some(*c), - rest => rest.flatten().last(), - } - } - // Hide all the nasty sh*t in here pub fn solve_set_characters( - set1: &Vec, - set2: &Vec, + set1: Vec, + set2: Vec, + truncate_set1_flag: bool, ) -> Result<(Vec, Vec), String> { let is_char_star = |s: &&Sequence| -> bool { match s { @@ -139,7 +133,8 @@ impl Sequence { .flat_map(Sequence::flatten) .count(); let star_compensate_len = set1_len.saturating_sub(set2_len); - let set2_solved = match (partition.next(), partition.next()) { + let (left, right) = (partition.next(), partition.next()); + let set2_solved: Vec = match (left, right) { (None, None) => match char_star { Some(c) => std::iter::repeat(*c).take(star_compensate_len).collect(), None => std::iter::empty().collect(), @@ -176,7 +171,10 @@ impl Sequence { .collect(), }, }; - let set1_solved = set1.iter().flat_map(Sequence::flatten).collect(); + let mut set1_solved: Vec = set1.iter().flat_map(Sequence::flatten).collect(); + if truncate_set1_flag { + set1_solved.truncate(set2_solved.len()); + } return Ok((set1_solved, set2_solved)); } else { Err(format!( @@ -407,9 +405,9 @@ pub struct DeleteOperation { } impl DeleteOperation { - pub fn new(set: Vec, complement_flag: bool) -> DeleteOperation { + pub fn new(set: Vec, complement_flag: bool) -> DeleteOperation { DeleteOperation { - set: set.iter().flat_map(Sequence::flatten).collect::>(), + set, complement_flag, } } @@ -427,18 +425,16 @@ pub struct TranslateOperationComplement { set2_iter: usize, set1: Vec, set2: Vec, - fallback: char, translation_map: HashMap, } impl TranslateOperationComplement { - fn new(set1: Vec, set2: Vec, fallback: char) -> TranslateOperationComplement { + fn new(set1: Vec, set2: Vec) -> TranslateOperationComplement { TranslateOperationComplement { iter: 0, set2_iter: 0, set1, set2, - fallback, translation_map: HashMap::new(), } } @@ -450,12 +446,22 @@ pub struct TranslateOperationStandard { } impl TranslateOperationStandard { - fn new(set1: Vec, set2: Vec, fallback: char) -> TranslateOperationStandard { - TranslateOperationStandard { - translation_map: set1 - .into_iter() - .zip(set2.into_iter().chain(std::iter::repeat(fallback))) - .collect::>(), + fn new(set1: Vec, set2: Vec) -> Result { + if let Some(fallback) = set2.last().map(|s| *s) { + Ok(TranslateOperationStandard { + translation_map: set1 + .into_iter() + .zip(set2.into_iter().chain(std::iter::repeat(fallback))) + .collect::>(), + }) + } else { + if set1.is_empty() && set2.is_empty() { + Ok(TranslateOperationStandard { + translation_map: HashMap::new(), + }) + } else { + Err("when not truncating set1, string2 must be non-empty".to_string()) + } } } } @@ -478,29 +484,17 @@ impl TranslateOperation { impl TranslateOperation { pub fn new( - set1: Vec, - set2: Vec, - truncate_set1_flag: bool, + set1: Vec, + set2: Vec, complement: bool, ) -> Result { - let (mut set1_solved, set2_solved) = Sequence::solve_set_characters(&set1, &set2)?; - if truncate_set1_flag { - set1_solved.truncate(set2_solved.len()); - } - let fallback = set2.last().map(Sequence::last).flatten().expect( - format!( - "{}: when not truncating set1, string2 must be non-empty", - executable!() - ) - .as_str(), - ); if complement { Ok(TranslateOperation::Complement( - TranslateOperationComplement::new(set1_solved, set2_solved, fallback), + TranslateOperationComplement::new(set1, set2), )) } else { Ok(TranslateOperation::Standard( - TranslateOperationStandard::new(set1_solved, set2_solved, fallback), + TranslateOperationStandard::new(set1, set2)?, )) } } @@ -520,7 +514,6 @@ impl SymbolTranslator for TranslateOperation { set2_iter, set1, set2, - fallback, translation_map, }) => { // First, try to see if current char is already mapped @@ -539,7 +532,7 @@ impl SymbolTranslator for TranslateOperation { *set2_iter = set2_iter.saturating_add(1); translation_map.insert(next_key, *value); } else { - translation_map.insert(current, *fallback); + translation_map.insert(current, *set2.last().unwrap()); } } Some(*translation_map.get(¤t).unwrap()) @@ -557,9 +550,9 @@ pub struct SqueezeOperation { } impl SqueezeOperation { - pub fn new(set1: Vec, complement: bool) -> SqueezeOperation { + pub fn new(set1: Vec, complement: bool) -> SqueezeOperation { SqueezeOperation { - set1: set1.iter().flat_map(Sequence::flatten).collect(), + set1, complement, previous: None, } diff --git a/src/uu/tr/src/tr.rs b/src/uu/tr/src/tr.rs index 59e4852b2..99d7b5132 100644 --- a/src/uu/tr/src/tr.rs +++ b/src/uu/tr/src/tr.rs @@ -66,6 +66,7 @@ pub fn uumain(args: impl uucore::Args) -> i32 { .values_of(options::SETS) .map(|v| v.map(ToString::to_string).collect::>()) .unwrap_or_default(); + let sets_len = sets.len(); if sets.is_empty() { show_error!( @@ -75,7 +76,7 @@ pub fn uumain(args: impl uucore::Args) -> i32 { return 1; } - if !(delete_flag || squeeze_flag) && sets.len() < 2 { + if !(delete_flag || squeeze_flag) && sets_len < 2 { show_error!( "missing operand after '{}'\nTry '{} --help' for more information.", sets[0], @@ -84,7 +85,7 @@ pub fn uumain(args: impl uucore::Args) -> i32 { return 1; } - if sets.len() > 2 { + if sets_len > 2 { show_error!( "extra operand '{}'\nTry '{} --help' for more information.", sets[0], @@ -99,37 +100,44 @@ pub fn uumain(args: impl uucore::Args) -> i32 { let locked_stdout = stdout.lock(); let mut buffered_stdout = BufWriter::new(locked_stdout); + let mut sets_iter = sets.into_iter(); + let (set1, set2) = match Sequence::solve_set_characters( + Sequence::from_str(sets_iter.next().unwrap_or_default().as_str()), + Sequence::from_str(sets_iter.next().unwrap_or_default().as_str()), + truncate_set1_flag, + ) { + Ok(r) => r, + Err(s) => { + show_error!("{}", s); + return 1; + } + }; if delete_flag { if squeeze_flag { let mut delete_buffer = vec![]; { let mut delete_writer = BufWriter::new(&mut delete_buffer); - let delete_op = DeleteOperation::new(Sequence::from_str(&sets[0]), complement_flag); + let delete_op = DeleteOperation::new(set1.clone(), complement_flag); translate_input(&mut locked_stdin, &mut delete_writer, delete_op); } { let mut squeeze_reader = BufReader::new(delete_buffer.as_bytes()); - let op = SqueezeOperation::new(Sequence::from_str(&sets[1]), complement_flag); + let op = SqueezeOperation::new(set2, complement_flag); translate_input(&mut squeeze_reader, &mut buffered_stdout, op); } } else { - let op = DeleteOperation::new(Sequence::from_str(&sets[0]), complement_flag); + let op = DeleteOperation::new(set1, complement_flag); translate_input(&mut locked_stdin, &mut buffered_stdout, op); } } else if squeeze_flag { - if sets.len() < 2 { - let op = SqueezeOperation::new(Sequence::from_str(&sets[0]), complement_flag); + if sets_len < 2 { + let op = SqueezeOperation::new(set1, complement_flag); translate_input(&mut locked_stdin, &mut buffered_stdout, op); } else { let mut translate_buffer = vec![]; { let mut writer = BufWriter::new(&mut translate_buffer); - match TranslateOperation::new( - Sequence::from_str(&sets[0]), - Sequence::from_str(&sets[1]), - truncate_set1_flag, - complement_flag, - ) { + match TranslateOperation::new(set1.clone(), set2.clone(), complement_flag) { Ok(op) => translate_input(&mut locked_stdin, &mut writer, op), Err(s) => { show_error!("{}", s); @@ -139,17 +147,12 @@ pub fn uumain(args: impl uucore::Args) -> i32 { } { let mut reader = BufReader::new(translate_buffer.as_bytes()); - let squeeze_op = SqueezeOperation::new(Sequence::from_str(&sets[1]), false); + let squeeze_op = SqueezeOperation::new(set2, false); translate_input(&mut reader, &mut buffered_stdout, squeeze_op); } } } else { - match TranslateOperation::new( - Sequence::from_str(&sets[0]), - Sequence::from_str(&sets[1]), - truncate_set1_flag, - complement_flag, - ) { + match TranslateOperation::new(set1, set2, complement_flag) { Ok(op) => translate_input(&mut locked_stdin, &mut buffered_stdout, op), Err(s) => { show_error!("{}", s);