pr: fix column behavior for short files

src/uu/pr/src/pr.rs

@@ -1017,6 +1017,74 @@ fn print_page(
     Ok(lines_written)
 }
 
+/// Group the lines of the input file in columns read left-to-right.
+fn to_table_across(
+    content_lines_per_page: usize,
+    columns: usize,
+    lines: &[FileLine],
+) -> Vec<Vec<Option<&FileLine>>> {
+    (0..content_lines_per_page)
+        .map(|i| (0..columns).map(|j| lines.get(i * columns + j)).collect())
+        .collect()
+}
+
+/// Group the lines of the input files in columns, one column per file.
+fn to_table_merged(
+    content_lines_per_page: usize,
+    columns: usize,
+    filled_lines: Vec<Option<&FileLine>>,
+) -> Vec<Vec<Option<&FileLine>>> {
+    (0..content_lines_per_page)
+        .map(|i| {
+            (0..columns)
+                .map(|j| {
+                    *filled_lines
+                        .get(content_lines_per_page * j + i)
+                        .unwrap_or(&None)
+                })
+                .collect()
+        })
+        .collect()
+}
+
+/// Group lines of the file in columns, going top-to-bottom then left-to-right.
+///
+/// This function should be applied when there are more lines than the
+/// total number of cells in the table.
+fn to_table(
+    content_lines_per_page: usize,
+    columns: usize,
+    lines: &[FileLine],
+) -> Vec<Vec<Option<&FileLine>>> {
+    (0..content_lines_per_page)
+        .map(|i| {
+            (0..columns)
+                .map(|j| lines.get(content_lines_per_page * j + i))
+                .collect()
+        })
+        .collect()
+}
+
+/// Group lines of the file in columns, going top-to-bottom then left-to-right.
+///
+/// This function should be applied when there are fewer lines than the
+/// total number of cells in the table.
+fn to_table_short_file(
+    content_lines_per_page: usize,
+    columns: usize,
+    lines: &[FileLine],
+) -> Vec<Vec<Option<&FileLine>>> {
+    let num_rows = lines.len() / columns;
+    let mut table: Vec<Vec<_>> = (0..num_rows)
+        .map(|i| (0..columns).map(|j| lines.get(num_rows * j + i)).collect())
+        .collect();
+    // Fill the rest with Nones.
+    for _ in num_rows..content_lines_per_page {
+        table.push(vec![None; columns]);
+    }
+    table
+}
+
 #[allow(clippy::cognitive_complexity)]
 fn write_columns(
     lines: &[FileLine],
@@ -1044,7 +1112,7 @@ fn write_columns(
         .as_ref()
         .is_some_and(|i| i.across_mode);
 
-    let mut filled_lines = Vec::new();
+    let mut filled_lines: Vec<Option<&FileLine>> = Vec::new();
     if options.merge_files_print.is_some() {
         let mut offset = 0;
         for col in 0..columns {
@@ -1064,23 +1132,19 @@ fn write_columns(
         }
     }
 
-    let table: Vec<Vec<_>> = (0..content_lines_per_page)
-        .map(move |a| {
-            (0..columns)
-                .map(|i| {
-                    if across_mode {
-                        lines.get(a * columns + i)
-                    } else if options.merge_files_print.is_some() {
-                        *filled_lines
-                            .get(content_lines_per_page * i + a)
-                            .unwrap_or(&None)
-                    } else {
-                        lines.get(content_lines_per_page * i + a)
-                    }
-                })
-                .collect()
-        })
-        .collect();
+    // Group the flat list of lines into a 2-dimensional table of
+    // cells, where each row will be printed as a single line in the
+    // output.
+    let merge = options.merge_files_print.is_some();
+    let table = if !merge && (lines.len() < (content_lines_per_page * columns)) {
+        to_table_short_file(content_lines_per_page, columns, lines)
+    } else if across_mode {
+        to_table_across(content_lines_per_page, columns, lines)
+    } else if merge {
+        to_table_merged(content_lines_per_page, columns, filled_lines)
+    } else {
+        to_table(content_lines_per_page, columns, lines)
+    };
 
     let blank_line = FileLine::default();
     for row in table {

tests/by-util/test_pr.rs

@@ -7,8 +7,8 @@
 use jiff::{Timestamp, ToSpan};
 use regex::Regex;
 use std::fs::metadata;
-use uutests::new_ucmd;
 use uutests::util::UCommand;
+use uutests::{at_and_ucmd, new_ucmd};
 
 const DATE_TIME_FORMAT_DEFAULT: &str = "%Y-%m-%d %H:%M";
 
@@ -666,3 +666,93 @@ fn test_form_feed_followed_by_new_line() {
         .succeeds()
         .stdout_matches(&regex);
 }
+
+#[test]
+fn test_columns() {
+    let whitespace = " ".repeat(50);
+    let datetime_pattern = r"\d\d\d\d-\d\d-\d\d \d\d:\d\d";
+    let header = format!("\n\n{datetime_pattern}{whitespace}Page 1\n\n\n");
+    // TODO Our output still does not match the behavior of GNU
+    // pr. The correct output should be:
+    //
+    //     "a\t\t\t\t    b\n";
+    //
+    let data = "a                                  \tb                                  \n";
+    let blank_lines_60 = "\n".repeat(60);
+    let pattern = format!("{header}{data}{blank_lines_60}");
+    let regex = Regex::new(&pattern).unwrap();
+
+    // Command line: `printf "a\nb\n" | pr -2`.
+    new_ucmd!()
+        .arg("-2")
+        .pipe_in("a\nb\n")
+        .succeeds()
+        .stdout_matches(&regex);
+}
+
+#[test]
+fn test_merge() {
+    // Create the two files to merge.
+    let (at, mut ucmd) = at_and_ucmd!();
+    at.write("f", "a\n");
+    at.write("g", "b\n");
+
+    let whitespace = " ".repeat(50);
+    let datetime_pattern = r"\d\d\d\d-\d\d-\d\d \d\d:\d\d";
+    let header = format!("\n\n{datetime_pattern}{whitespace}Page 1\n\n\n");
+    // TODO Our output still does not match the behavior of GNU
+    // pr. The correct output should be:
+    //
+    //     "a\t\t\t\t    b\n";
+    //
+    // and the blank lines should actually be empty lines.
+    let data = "a                                  \tb                                  \n";
+    let blank_lines_55 =
+        "                                   \t                                   \n".repeat(55);
+    let footer = "\n".repeat(5);
+    let pattern = format!("{header}{data}{blank_lines_55}{footer}");
+    let regex = Regex::new(&pattern).unwrap();
+
+    // Command line: `(echo "a" > f; echo "b" > g; pr -m f g)`.
+    ucmd.args(&["-m", "f", "g"])
+        .succeeds()
+        .stdout_matches(&regex);
+}
+
+#[test]
+fn test_merge_one_long_one_short() {
+    // Create the two files to merge.
+    let (at, mut ucmd) = at_and_ucmd!();
+    at.write("f", "a\na\n");
+    at.write("g", "b\n");
+
+    // Page 1 should have the first line of `f` and the first line of
+    // `b` side-by-side.
+    let whitespace = " ".repeat(50);
+    let datetime_pattern = r"\d\d\d\d-\d\d-\d\d \d\d:\d\d";
+    let header = format!("\n\n{datetime_pattern}{whitespace}Page 1\n\n\n");
+    let data = "a                                  \tb                                  \n";
+    let footer = "\n".repeat(5);
+    let page1 = format!("{header}{data}{footer}");
+
+    // Page 2 should have just the second line of `f`.
+    let header = format!("\n\n{datetime_pattern}{whitespace}Page 2\n\n\n");
+    let data = "a                                  \t                                   \n";
+    let page2 = format!("{header}{data}{footer}");
+
+    let pattern = format!("{page1}{page2}");
+    let regex = Regex::new(&pattern).unwrap();
+
+    // Command line:
+    //
+    //     printf "a\na\n" > f
+    //     printf "b\n" > g
+    //     pr -l11 -m f g
+    //
+    // The line length of 11 leaves room for a 5-line header, a 5-line
+    // footer, and one line of data from the input files. The extra
+    // line from the file `f` will be on the second page.
+    ucmd.args(&["-l", "11", "-m", "f", "g"])
+        .succeeds()
+        .stdout_matches(&regex);
+}