1
use std::cell::RefCell;
2
use std::io::Error;
3
use std::io::ErrorKind;
4
use std::io::Write;
5
use std::rc::Rc;
6
use std::str::from_utf8;
7

            
8
/// An indentation manager that maintains the current indentation level and provides
9
/// methods for formatting text with proper indentation.
10
///
11
/// The indentation level can be increased with `indent()`, which returns an `Indent`
12
/// guard that automatically decreases the indentation when dropped.
13
#[derive(Debug)]
14
pub struct IndentFormatter<'a, W: Write> {
15
    /// The current indentation level (number of tabs), wrapped in `Rc<RefCell>` for interior mutability
16
    level: Rc<RefCell<usize>>,
17
    /// The underlying writer to which indented content will be written
18
    writer: &'a mut W,
19
    /// The string used for a single level of indentation
20
    indent_str: String,
21
    /// Tracks whether we're at the start of a line (where indentation should be applied)
22
    at_line_start: bool,
23
}
24

            
25
impl<'a, W: Write> IndentFormatter<'a, W> {
26
    /// Creates a new IndentFormatter with zero indentation.
27
3
    pub fn new(writer: &'a mut W) -> Self {
28
3
        Self::with_indent_str(writer, "  ".to_string())
29
3
    }
30

            
31
    /// Creates a new IndentFormatter with zero indentation and specified indentation string.
32
3
    pub fn with_indent_str(writer: &'a mut W, indent_str: String) -> Self {
33
3
        Self {
34
3
            writer,
35
3
            level: Rc::new(RefCell::new(0)),
36
3
            indent_str,
37
3
            at_line_start: true, // Start at the beginning of a line
38
3
        }
39
3
    }
40

            
41
    /// Increases the indentation level and returns a guard that will
42
    /// decrease it when dropped.
43
62
    pub fn indent(&self) -> Indent {
44
62
        let mut level_ref = self.level.borrow_mut();
45
62
        *level_ref += 1;
46
62
        drop(level_ref); // Release the borrow before creating the guard
47

            
48
62
        Indent {
49
62
            level: Rc::clone(&self.level),
50
62
        }
51
62
    }
52

            
53
    /// Returns the current indentation level.
54
396
    pub fn level(&self) -> usize {
55
396
        *self.level.borrow()
56
396
    }
57
}
58

            
59
impl<W: Write> Write for IndentFormatter<'_, W> {
60
1219
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
61
        // Convert the byte slice to a string slice
62
1219
        let s = from_utf8(buf).map_err(|_| Error::new(ErrorKind::InvalidData, "Invalid UTF-8"))?;
63

            
64
1219
        let parts = s.split('\n');
65
1219
        let mut first = true;
66

            
67
        // Handle each split line
68
1652
        for part in parts {
69
            // Write the newline at the end of the previous line.
70
1652
            if !first {
71
433
                self.writer.write_all(b"\n")?;
72
433
                self.at_line_start = true;
73
1219
            }
74

            
75
1652
            if !part.is_empty() {
76
                // Add indentation if we're at the start of a line
77
1280
                if self.at_line_start {
78
396
                    for _ in 0..self.level() {
79
369
                        self.writer.write_all(self.indent_str.as_bytes())?;
80
                    }
81
396
                    self.at_line_start = false;
82
884
                }
83
1280
                self.writer.write_all(part.as_bytes())?;
84
372
            }
85

            
86
1652
            first = false;
87
        }
88

            
89
        // Update line start status based on whether the string ends with a newline
90
1219
        if s.ends_with('\n') {
91
280
            self.at_line_start = true;
92
939
        }
93

            
94
1219
        Ok(buf.len())
95
1219
    }
96

            
97
1
    fn flush(&mut self) -> std::io::Result<()> {
98
1
        self.writer.flush()
99
1
    }
100
}
101

            
102
/// A guard object that decreases indentation when dropped.
103
/// Created by calling `IndentFormatter::indent()`.
104
///
105
/// Uses interior mutability to avoid requiring a mutable reference to the IndentFormatter.
106
#[derive(Debug)]
107
pub struct Indent {
108
    /// Reference-counted cell containing the indentation level
109
    level: Rc<RefCell<usize>>,
110
}
111

            
112
impl Drop for Indent {
113
    /// Decreases the indentation level when this guard is dropped.
114
62
    fn drop(&mut self) {
115
62
        let mut level_ref = self.level.borrow_mut();
116
62
        debug_assert!(*level_ref > 0, "Indentation level cannot go below zero");
117
62
        *level_ref -= 1;
118
62
    }
119
}
120

            
121
#[cfg(test)]
122
mod tests {
123
    use super::IndentFormatter;
124
    use super::Write;
125

            
126
    /// Tests that the indenter correctly handles multi-line strings
127
    #[test]
128
1
    fn test_multiline_string() {
129
1
        let mut buffer = Vec::new();
130
1
        {
131
1
            let mut formatter = IndentFormatter::new(&mut buffer);
132
1

            
133
1
            // First level indent
134
1
            let _indent1 = formatter.indent();
135
1

            
136
1
            // Write a multi-line string at once
137
1
            write!(formatter, "First line\nSecond line\nThird line").unwrap();
138
1
        }
139

            
140
1
        let result = String::from_utf8(buffer).unwrap();
141
1
        let expected = "  First line\n  Second line\n  Third line";
142
1
        assert_eq!(result, expected, "Multiline indentation incorrect");
143
1
    }
144

            
145
    /// Tests that the indenter correctly handles line continuation across multiple write calls
146
    #[test]
147
1
    fn test_line_continuation() {
148
1
        let mut buffer = Vec::new();
149
1
        {
150
1
            let mut formatter = IndentFormatter::new(&mut buffer);
151
1

            
152
1
            // First level indent
153
1
            let _indent1 = formatter.indent();
154
1

            
155
1
            // First part of a line - should be indented
156
1
            write!(formatter, "Start of line ").unwrap();
157
1

            
158
1
            // Continuation of the same line - should not be indented
159
1
            write!(formatter, "continued here").unwrap();
160
1

            
161
1
            // New line followed by text - only the new line's content should be indented
162
1
            write!(formatter, "\nSecond line").unwrap();
163
1

            
164
1
            // Another continuation
165
1
            write!(formatter, " continued").unwrap();
166
1

            
167
1
            // A line ending with newline
168
1
            write!(formatter, "\nThird line\n").unwrap();
169
1

            
170
1
            // A new line after previous newline - should be indented
171
1
            write!(formatter, "Fourth line").unwrap();
172
1
        }
173

            
174
1
        let result = String::from_utf8(buffer).unwrap();
175
1
        let expected = "  Start of line continued here\n  Second line continued\n  Third line\n  Fourth line";
176

            
177
1
        assert_eq!(result, expected, "Line continuation handling incorrect");
178
1
    }
179
}