split: implement round-robin arg to --number

jfinkels · sylvestre · commit 32155e157706 · 2022-10-05T13:39:26.000+02:00
Implement distributing lines of a file in a round-robin manner to a
specified number of chunks. For example,

    $ (seq 1 10 | split -n r/3) &amp;&amp; head -v xa[abc]
    ==&gt; xaa &lt;==
    1
    4
    7
    10

    ==&gt; xab &lt;==
    2
    5
    8

    ==&gt; xac &lt;==
    3
    6
    9
diff --git a/src/uu/split/src/split.rs b/src/uu/split/src/split.rs
@@ -1195,6 +1195,46 @@ where
     Ok(())
 }
 
+fn split_into_n_chunks_by_line_round_robin<R>(
+    settings: &Settings,
+    reader: &mut R,
+    num_chunks: u64,
+) -> UResult<()>
+where
+    R: BufRead,
+{
+    // This object is responsible for creating the filename for each chunk.
+    let mut filename_iterator = FilenameIterator::new(
+        &settings.prefix,
+        &settings.additional_suffix,
+        settings.suffix_length,
+        settings.suffix_type,
+    );
+
+    // Create one writer for each chunk. This will create each
+    // of the underlying files (if not in `--filter` mode).
+    let mut writers = vec![];
+    for _ in 0..num_chunks {
+        let filename = filename_iterator
+            .next()
+            .ok_or_else(|| USimpleError::new(1, "output file suffixes exhausted"))?;
+        let writer = settings.instantiate_current_writer(filename.as_str())?;
+        writers.push(writer);
+    }
+
+    let num_chunks: usize = num_chunks.try_into().unwrap();
+    for (i, line_result) in reader.lines().enumerate() {
+        let line = line_result.unwrap();
+        let maybe_writer = writers.get_mut(i % num_chunks);
+        let writer = maybe_writer.unwrap();
+        let bytes = line.as_bytes();
+        writer.write_all(bytes)?;
+        writer.write_all(b"\n")?;
+    }
+
+    Ok(())
+}
+
 fn split(settings: &Settings) -> UResult<()> {
     let mut reader = BufReader::new(if settings.input == "-" {
         Box::new(stdin()) as Box<dyn Read>
@@ -1221,6 +1261,9 @@ fn split(settings: &Settings) -> UResult<()> {
             let chunk_number = chunk_number - 1;
             kth_chunk_by_line(settings, &mut reader, chunk_number, num_chunks)
         }
+        Strategy::Number(NumberType::RoundRobin(num_chunks)) => {
+            split_into_n_chunks_by_line_round_robin(settings, &mut reader, num_chunks)
+        }
         Strategy::Number(_) => Err(USimpleError::new(1, "-n mode not yet fully implemented")),
         Strategy::Lines(chunk_size) => {
             let mut writer = LineChunkWriter::new(chunk_size, settings)?;
diff --git a/tests/by-util/test_split.rs b/tests/by-util/test_split.rs
@@ -701,3 +701,19 @@ fn test_multiple_of_input_chunk() {
     }
     assert_eq!(glob.collate(), at.read_bytes(name));
 }
+
+#[test]
+fn test_round_robin() {
+    let (at, mut ucmd) = at_and_ucmd!();
+
+    let file_read = |f| {
+        let mut s = String::new();
+        at.open(f).read_to_string(&mut s).unwrap();
+        s
+    };
+
+    ucmd.args(&["-n", "r/2", "fivelines.txt"]).succeeds();
+
+    assert_eq!(file_read("xaa"), "1\n3\n5\n");
+    assert_eq!(file_read("xab"), "2\n4\n");
+}
