ansi_align/

lib.rs

use string_width::string_width;

/// Alignment options for text
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Alignment {
    Left,
    Center,
    Right,
}

/// A type-safe wrapper for display width values
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct Width(usize);

impl Width {
    /// Create a new Width value
    pub fn new(value: usize) -> Self {
        Self(value)
    }

    /// Get the underlying usize value
    pub fn get(self) -> usize {
        self.0
    }
}

impl From<usize> for Width {
    fn from(value: usize) -> Self {
        Self(value)
    }
}

/// Options for text alignment
#[derive(Debug, Clone)]
pub struct AlignOptions {
    /// The alignment type (left, center, right)
    pub align: Alignment,
    /// The string to split lines on (default: "\n")
    pub split: String,
    /// The padding character to use (default: " ")
    pub pad: char,
}

impl Default for AlignOptions {
    fn default() -> Self {
        Self {
            align: Alignment::Center,
            split: "\n".to_string(),
            pad: ' ',
        }
    }
}

impl AlignOptions {
    /// Create new options with specified alignment
    pub fn new(align: Alignment) -> Self {
        Self {
            align,
            ..Default::default()
        }
    }

    /// Set the split string using builder pattern
    pub fn with_split<S: Into<String>>(mut self, split: S) -> Self {
        self.split = split.into();
        self
    }

    /// Set the padding character using builder pattern
    pub fn with_pad(mut self, pad: char) -> Self {
        self.pad = pad;
        self
    }
}

/// Efficiently create padding string for alignment
fn create_padding(pad_char: char, count: usize) -> String {
    match count {
        0 => String::new(),
        1 => pad_char.to_string(),
        2..=8 => pad_char.to_string().repeat(count),
        _ => {
            let mut padding = String::with_capacity(count);
            for _ in 0..count {
                padding.push(pad_char);
            }
            padding
        }
    }
}

/// Align text with center alignment (default behavior)
#[must_use]
pub fn ansi_align(text: &str) -> String {
    ansi_align_with_options(text, AlignOptions::default())
}

/// Align text with support for ANSI escape sequences
///
/// This function handles text containing ANSI escape sequences (like color codes)
/// by calculating display width correctly, ignoring the escape sequences.
///
/// # Examples
///
/// ```
/// use ansi_align::{ansi_align, ansi_align_with_options, Alignment, AlignOptions};
///
/// // Basic center alignment
/// let result = ansi_align("hello\nworld");
///
/// // Right alignment with custom padding
/// let opts = AlignOptions::new(Alignment::Right).with_pad('.');
/// let result = ansi_align_with_options("hi\nhello", opts);
/// assert_eq!(result, "...hi\nhello");
///
/// // Works with ANSI escape sequences
/// let colored = "\x1b[31mred\x1b[0m\n\x1b[32mgreen\x1b[0m";
/// let result = ansi_align_with_options(colored, AlignOptions::new(Alignment::Center));
/// // ANSI codes are preserved but don't affect alignment calculation
/// ```
///
/// # Performance
///
/// This function makes a single pass through the input text for optimal performance.
/// For left alignment, it returns the input unchanged as an optimization.
#[must_use]
pub fn ansi_align_with_options(text: &str, opts: AlignOptions) -> String {
    if text.is_empty() {
        return text.to_string();
    }

    // Short-circuit left alignment as no-op
    if opts.align == Alignment::Left {
        return text.to_string();
    }

    // Single pass: collect line data and find max width simultaneously
    let line_data: Vec<(&str, Width)> = text
        .split(&opts.split)
        .map(|line| (line, Width::from(string_width(line))))
        .collect();

    let max_width = line_data
        .iter()
        .map(|(_, width)| width.get())
        .max()
        .unwrap_or(0);

    let aligned_lines: Vec<String> = line_data
        .into_iter()
        .map(|(line, width)| {
            let padding_needed = match opts.align {
                Alignment::Left => 0, // Already handled above
                Alignment::Center => (max_width - width.get()) / 2,
                Alignment::Right => max_width - width.get(),
            };

            if padding_needed == 0 {
                line.to_string()
            } else {
                let mut result = create_padding(opts.pad, padding_needed);
                result.push_str(line);
                result
            }
        })
        .collect();

    aligned_lines.join(&opts.split)
}

/// Align text to the left (no-op, returns original text)
#[must_use]
pub fn left(text: &str) -> String {
    ansi_align_with_options(text, AlignOptions::new(Alignment::Left))
}

/// Align text to the center
#[must_use]
pub fn center(text: &str) -> String {
    ansi_align_with_options(text, AlignOptions::new(Alignment::Center))
}

/// Align text to the right
#[must_use]
pub fn right(text: &str) -> String {
    ansi_align_with_options(text, AlignOptions::new(Alignment::Right))
}

#[cfg(test)]
mod tests {
    use super::*;

    // Basic alignment tests
    #[test]
    fn test_left_alignment() {
        let text = "hello\nworld";
        let result = left(text);
        assert_eq!(result, text); // Left alignment should be no-op
    }

    #[test]
    fn test_center_alignment() {
        let text = "hi\nhello";
        let result = center(text);
        let lines: Vec<&str> = result.split('\n').collect();
        assert_eq!(lines[0], " hi"); // 1 space padding for "hi"
        assert_eq!(lines[1], "hello"); // No padding for "hello"
    }

    #[test]
    fn test_right_alignment() {
        let text = "hi\nhello";
        let result = right(text);
        let lines: Vec<&str> = result.split('\n').collect();
        assert_eq!(lines[0], "   hi"); // 3 spaces padding for "hi"
        assert_eq!(lines[1], "hello"); // No padding for "hello"
    }

    // Unicode and ANSI tests
    #[test]
    fn test_unicode_characters() {
        let text = "古\n古古古";
        let result = center(text);
        let lines: Vec<&str> = result.split('\n').collect();
        assert_eq!(lines[0], "  古"); // 2 spaces padding (CJK char is width 2)
        assert_eq!(lines[1], "古古古"); // No padding
    }

    #[test]
    fn test_ansi_escape_sequences() {
        let text = "hello\n\u{001B}[1mworld\u{001B}[0m";
        let result = center(text);
        let lines: Vec<&str> = result.split('\n').collect();
        assert_eq!(lines[0], "hello");
        assert_eq!(lines[1], "\u{001B}[1mworld\u{001B}[0m"); // ANSI codes preserved
    }

    #[test]
    fn test_complex_ansi_sequences() {
        // Test with multiple ANSI codes and colors
        let text = "\x1b[31m\x1b[1mred\x1b[0m\n\x1b[32mgreen text\x1b[0m";
        let result = right(text);
        let lines: Vec<&str> = result.split('\n').collect();
        // "red" has display width 3, "green text" has width 10
        assert_eq!(lines[0], "       \x1b[31m\x1b[1mred\x1b[0m"); // 7 spaces padding
        assert_eq!(lines[1], "\x1b[32mgreen text\x1b[0m"); // No padding
    }

    // Edge cases
    #[test]
    fn test_empty_string() {
        assert_eq!(ansi_align_with_options("", AlignOptions::default()), "");
        assert_eq!(left(""), "");
        assert_eq!(center(""), "");
        assert_eq!(right(""), "");
    }

    #[test]
    fn test_single_line() {
        let text = "hello";
        assert_eq!(left(text), "hello");
        assert_eq!(center(text), "hello");
        assert_eq!(right(text), "hello");
    }

    #[test]
    fn test_single_character() {
        let text = "a\nb";
        let result = center(text);
        assert_eq!(result, "a\nb"); // Both lines same width, no padding needed
    }

    #[test]
    fn test_whitespace_only() {
        let text = "   \n ";
        let result = center(text);
        let lines: Vec<&str> = result.split('\n').collect();
        assert_eq!(lines[0], "   "); // 3 spaces, no padding needed
        assert_eq!(lines[1], "  "); // 1 space + 1 padding space
    }

    // Custom options tests
    #[test]
    fn test_custom_split_and_pad() {
        let text = "a|bb";
        let opts = AlignOptions::new(Alignment::Right)
            .with_split("|")
            .with_pad('.');
        let result = ansi_align_with_options(text, opts);
        assert_eq!(result, ".a|bb");
    }

    #[test]
    fn test_custom_split_multichar() {
        let text = "short<->very long line";
        let opts = AlignOptions::new(Alignment::Center).with_split("<->");
        let result = ansi_align_with_options(text, opts);
        assert_eq!(result, "    short<->very long line");
    }

    #[test]
    fn test_different_padding_chars() {
        let text = "hi\nhello";

        // Test dot padding
        let opts = AlignOptions::new(Alignment::Right).with_pad('.');
        let result = ansi_align_with_options(text, opts);
        assert_eq!(result, "...hi\nhello");

        // Test underscore padding
        let opts = AlignOptions::new(Alignment::Center).with_pad('_');
        let result = ansi_align_with_options(text, opts);
        assert_eq!(result, "_hi\nhello");

        // Test zero padding
        let opts = AlignOptions::new(Alignment::Right).with_pad('0');
        let result = ansi_align_with_options(text, opts);
        assert_eq!(result, "000hi\nhello");
    }

    // Performance and memory optimization tests
    #[test]
    fn test_large_padding() {
        let text = format!("a\n{}", "b".repeat(100));
        let result = right(&text);
        let lines: Vec<&str> = result.split('\n').collect();
        assert_eq!(lines[0].len(), 100); // 99 spaces + "a"
        assert!(lines[0].starts_with(&" ".repeat(99)));
        assert!(lines[0].ends_with("a"));
        assert_eq!(lines[1], "b".repeat(100));
    }

    #[test]
    fn test_no_padding_optimization() {
        // Test that lines requiring no padding are handled efficiently
        let text = "same\nsame\nsame";
        let result = center(text);
        assert_eq!(result, text); // Should be unchanged
    }

    // Width type tests
    #[test]
    fn test_width_type() {
        let width = Width::new(42);
        assert_eq!(width.get(), 42);

        let width_from_usize: Width = 24.into();
        assert_eq!(width_from_usize.get(), 24);

        // Test ordering
        assert!(Width::new(10) < Width::new(20));
        assert_eq!(Width::new(15), Width::new(15));
    }

    // Comprehensive alignment scenarios
    #[test]
    fn test_mixed_width_lines() {
        let text = "a\nbb\nccc\ndddd\neeeee";

        // Center alignment
        let result = center(text);
        let lines: Vec<&str> = result.split('\n').collect();

        // The longest line is "eeeee" with 5 chars
        // So padding for center should be: (5-1)/2=2, (5-2)/2=1, (5-3)/2=1, (5-4)/2=0, (5-5)/2=0
        assert_eq!(lines[0], "  a"); // 2 spaces + "a"
        assert_eq!(lines[1], " bb"); // 1 space + "bb"
        assert_eq!(lines[2], " ccc"); // 1 space + "ccc" (corrected)
        assert_eq!(lines[3], "dddd"); // no padding (corrected)
        assert_eq!(lines[4], "eeeee"); // no padding

        // Right alignment
        let result = right(text);
        let lines: Vec<&str> = result.split('\n').collect();
        assert_eq!(lines[0], "    a"); // 4 spaces + "a"
        assert_eq!(lines[1], "   bb"); // 3 spaces + "bb"
        assert_eq!(lines[2], "  ccc"); // 2 spaces + "ccc"
        assert_eq!(lines[3], " dddd"); // 1 space + "dddd"
        assert_eq!(lines[4], "eeeee"); // no padding
    }

    #[test]
    fn test_center_odd_padding() {
        // Test center alignment with odd padding amounts
        let text = "a\nbbbb";
        let result = center(text);
        let lines: Vec<&str> = result.split('\n').collect();
        assert_eq!(lines[0], " a"); // (4-1)/2 = 1 space
        assert_eq!(lines[1], "bbbb"); // no padding
    }

    #[test]
    fn test_multiline_with_empty_lines() {
        let text = "hello\n\nworld";
        let result = center(text);
        let lines: Vec<&str> = result.split('\n').collect();
        assert_eq!(lines[0], "hello");
        assert_eq!(lines[1], "  "); // 2 spaces for empty line (center of 5-char width)
        assert_eq!(lines[2], "world");
    }

    // Regression tests for performance improvements
    #[test]
    fn test_no_unnecessary_allocations() {
        // This test ensures we don't regress on the performance improvements
        let text = "line1\nline2\nline3";
        let result = left(text);
        // Left alignment should return original string (no allocations for processing)
        assert_eq!(result, text);
    }

    #[test]
    fn test_padding_efficiency() {
        // Test the efficient padding creation for different sizes
        let text = format!("a\n{}", "b".repeat(20));

        // Small padding (should use repeat)
        let opts = AlignOptions::new(Alignment::Right);
        let result = ansi_align_with_options("a\nbb", opts.clone());
        assert_eq!(result, " a\nbb");

        // Large padding (should use with_capacity)
        let result = ansi_align_with_options(&text, opts);
        let lines: Vec<&str> = result.split('\n').collect();
        assert_eq!(lines[0].len(), 20); // 19 spaces + "a"
        assert!(lines[0].ends_with("a"));
    }

    // Integration tests
    #[test]
    fn test_real_world_scenario() {
        // Simulate aligning a simple table or menu
        let menu = "Home\nAbout Us\nContact\nServices";
        let result = center(menu);
        let lines: Vec<&str> = result.split('\n').collect();

        // "About Us" and "Services" are both 8 chars (longest)
        assert_eq!(lines[0], "  Home"); // 2 spaces (8-4)/2 = 2
        assert_eq!(lines[1], "About Us"); // no padding (8 chars)
        assert_eq!(lines[2], "Contact"); // no padding - "Contact" is 7 chars, (8-7)/2 = 0
        assert_eq!(lines[3], "Services"); // no padding (8 chars)
    }

    #[test]
    fn test_code_alignment() {
        // Test aligning code-like content
        let code = "if x:\n    return y\nelse:\n    return z";
        let result = right(code);
        let lines: Vec<&str> = result.split('\n').collect();

        // "    return y" and "    return z" are longest at 12 chars
        assert_eq!(lines[0], "       if x:"); // 7 spaces + "if x:" (12-5=7)
        assert_eq!(lines[1], "    return y"); // no padding (12 chars)
        assert_eq!(lines[2], "       else:"); // 7 spaces + "else:" (12-5=7)
        assert_eq!(lines[3], "    return z"); // no padding (12 chars)
    }
}