1use owo_colors::{AnsiColors, DynColors, OwoColorize, Style};
47use std::fmt::Write;
48
49pub struct AsciiArt<'a> {
51 content: Box<dyn 'a + Iterator<Item = &'a str>>,
52 colors: &'a [DynColors],
53 bold: bool,
54 start: usize,
55 end: usize,
56}
57impl<'a> AsciiArt<'a> {
58 pub fn new(input: &'a str, colors: &'a [DynColors], bold: bool) -> AsciiArt<'a> {
59 let mut lines: Vec<_> = input.lines().skip_while(|line| line.is_empty()).collect();
60 while let Some(line) = lines.last() {
61 if Tokens(line).is_empty() {
62 lines.pop();
63 } else {
64 break;
65 }
66 }
67
68 let (start, end) = get_min_start_max_end(&lines);
69
70 AsciiArt {
71 content: Box::new(lines.into_iter()),
72 colors,
73 bold,
74 start,
75 end,
76 }
77 }
78
79 pub fn width(&self) -> usize {
80 assert!(self.end >= self.start);
81 self.end - self.start
82 }
83}
84
85fn get_min_start_max_end(lines: &[&str]) -> (usize, usize) {
86 lines
87 .iter()
88 .map(|line| {
89 let line_start = Tokens(line).leading_spaces();
90 let line_end = Tokens(line).true_length();
91 (line_start, line_end)
92 })
93 .fold((usize::MAX, 0), |(acc_s, acc_e), (line_s, line_e)| {
94 (acc_s.min(line_s), acc_e.max(line_e))
95 })
96}
97
98impl Iterator for AsciiArt<'_> {
101 type Item = String;
102 fn next(&mut self) -> Option<String> {
103 self.content
104 .next()
105 .map(|line| Tokens(line).render(self.colors, self.start, self.end, self.bold))
106 }
107}
108
109#[derive(Clone, Debug, PartialEq, Eq)]
110enum Token {
111 Color(u32),
112 Char(char),
113 Space,
114}
115impl std::fmt::Display for Token {
116 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
117 match *self {
118 Token::Color(c) => write!(f, "{{{c}}}"),
119 Token::Char(c) => write!(f, "{c}"),
120 Token::Space => write!(f, " "),
121 }
122 }
123}
124impl Token {
125 fn is_solid(&self) -> bool {
126 matches!(*self, Token::Char(_))
127 }
128 fn is_space(&self) -> bool {
129 matches!(*self, Token::Space)
130 }
131 fn has_zero_width(&self) -> bool {
132 matches!(*self, Token::Color(_))
133 }
134}
135
136#[derive(Clone, Debug)]
138struct Tokens<'a>(&'a str);
139impl Iterator for Tokens<'_> {
140 type Item = Token;
141 fn next(&mut self) -> Option<Token> {
142 let (s, tok) = color_token(self.0)
143 .or_else(|| space_token(self.0))
144 .or_else(|| char_token(self.0))?;
145
146 self.0 = s;
147 Some(tok)
148 }
149}
150
151impl<'a> Tokens<'a> {
152 fn is_empty(&mut self) -> bool {
153 for token in self {
154 if token.is_solid() {
155 return false;
156 }
157 }
158 true
159 }
160 fn true_length(&mut self) -> usize {
161 let mut last_non_space = 0;
162 let mut last = 0;
163 for token in self {
164 if token.has_zero_width() {
165 continue;
166 }
167 last += 1;
168 if !token.is_space() {
169 last_non_space = last;
170 }
171 }
172 last_non_space
173 }
174 fn leading_spaces(&mut self) -> usize {
175 self.take_while(|token| !token.is_solid())
176 .filter(Token::is_space)
177 .count()
178 }
179 fn truncate(self, mut start: usize, end: usize) -> impl 'a + Iterator<Item = Token> {
180 assert!(start <= end);
181 let mut width = end - start;
182
183 self.filter(move |token| {
184 if start > 0 && !token.has_zero_width() {
185 start -= 1;
186 return false;
187 }
188 true
189 })
190 .take_while(move |token| {
191 if width == 0 {
192 return false;
193 }
194 if !token.has_zero_width() {
195 width -= 1;
196 }
197 true
198 })
199 }
200 fn render(self, colors: &[DynColors], start: usize, end: usize, bold: bool) -> String {
202 assert!(start <= end);
203 let mut width = end - start;
204 let mut colored_segment = String::new();
205 let mut whole_string = String::new();
206 let mut color = &DynColors::Ansi(AnsiColors::Default);
207
208 self.truncate(start, end).for_each(|token| match token {
209 Token::Char(chr) => {
210 width = width.saturating_sub(1);
211 colored_segment.push(chr);
212 }
213 Token::Color(col) => {
214 add_styled_segment(&mut whole_string, &colored_segment, *color, bold);
215 colored_segment = String::new();
216 color = colors
217 .get(col as usize)
218 .unwrap_or(&DynColors::Ansi(AnsiColors::Default));
219 }
220 Token::Space => {
221 width = width.saturating_sub(1);
222 colored_segment.push(' ');
223 }
224 });
225
226 add_styled_segment(&mut whole_string, &colored_segment, *color, bold);
227 (0..width).for_each(|_| whole_string.push(' '));
228 whole_string
229 }
230}
231
232fn succeed_when<I>(predicate: impl FnOnce(I) -> bool) -> impl FnOnce(I) -> Option<()> {
235 |input| {
236 if predicate(input) { Some(()) } else { None }
237 }
238}
239
240fn add_styled_segment(base: &mut String, segment: &str, color: DynColors, bold: bool) {
241 let mut style = Style::new().color(color);
242 if bold {
243 style = style.bold();
244 }
245 let formatted_segment = segment.style(style);
246 let _ = write!(base, "{formatted_segment}");
247}
248
249type ParseResult<'a, R> = Option<(&'a str, R)>;
252
253fn token<'a, R>(s: &'a str, predicate: impl FnOnce(char) -> Option<R>) -> ParseResult<'a, R> {
254 let mut chars = s.chars();
255 let token = chars.next()?;
256 let result = predicate(token)?;
257 Some((chars.as_str(), result))
258}
259
260fn color_token<'a>(s: &'a str) -> ParseResult<'a, Token> {
264 let (s, ()) = token(s, succeed_when(|c| c == '{'))?;
265 let (s, color_index) = token(s, |c| c.to_digit(10))?;
266 let (s, ()) = token(s, succeed_when(|c| c == '}'))?;
267 Some((s, Token::Color(color_index)))
268}
269
270fn space_token<'a>(s: &'a str) -> ParseResult<'a, Token> {
272 token(s, succeed_when(|c| c == ' ')).map(|(s, ())| (s, Token::Space))
273}
274
275fn char_token<'a>(s: &'a str) -> ParseResult<'a, Token> {
277 token(s, |c| Some(Token::Char(c)))
278}
279
280#[cfg(test)]
281mod test {
282 use super::*;
283
284 #[test]
285 fn test_get_min_start_max_end() {
286 let lines = [
287 " xxx",
288 " xxx",
289 " oo",
290 " o",
291 " xx",
292 ];
293 assert_eq!(get_min_start_max_end(&lines), (3, 29));
294 }
295
296 #[test]
297 fn space_parses() {
298 assert_eq!(space_token(" "), Some(("", Token::Space)));
299 assert_eq!(space_token(" hello"), Some(("hello", Token::Space)));
300 assert_eq!(space_token(" "), Some((" ", Token::Space)));
301 assert_eq!(space_token(" {1}{2}"), Some(("{1}{2}", Token::Space)));
302 }
303
304 #[test]
305 fn color_indicator_parses() {
306 assert_eq!(color_token("{1}"), Some(("", Token::Color(1))));
307 assert_eq!(color_token("{9} "), Some((" ", Token::Color(9))));
308 }
309
310 #[test]
311 fn leading_spaces_counts_correctly() {
312 assert_eq!(Tokens("").leading_spaces(), 0);
313 assert_eq!(Tokens(" ").leading_spaces(), 5);
314 assert_eq!(Tokens(" a;lksjf;a").leading_spaces(), 5);
315 assert_eq!(Tokens(" {1} {5} {9} a").leading_spaces(), 6);
316 }
317
318 #[test]
319 fn render() {
320 let colors_shim = Vec::new();
321
322 assert_eq!(
323 Tokens("").render(&colors_shim, 0, 0, true),
324 "\u{1b}[39;1m\u{1b}[0m"
325 );
326
327 assert_eq!(
328 Tokens(" ").render(&colors_shim, 0, 0, true),
329 "\u{1b}[39;1m\u{1b}[0m"
330 );
331
332 assert_eq!(
333 Tokens(" ").render(&colors_shim, 0, 5, true),
334 "\u{1b}[39;1m \u{1b}[0m"
335 );
336
337 assert_eq!(
338 Tokens(" ").render(&colors_shim, 1, 5, true),
339 "\u{1b}[39;1m \u{1b}[0m"
340 );
341
342 assert_eq!(
343 Tokens(" ").render(&colors_shim, 3, 5, true),
344 "\u{1b}[39;1m \u{1b}[0m"
345 );
346
347 assert_eq!(
348 Tokens(" ").render(&colors_shim, 0, 4, true),
349 "\u{1b}[39;1m \u{1b}[0m"
350 );
351
352 assert_eq!(
353 Tokens(" ").render(&colors_shim, 0, 3, true),
354 "\u{1b}[39;1m \u{1b}[0m"
355 );
356
357 assert_eq!(
359 Tokens("███").render(Vec::new().as_slice(), 0, 3, true),
360 "\u{1b}[39;1m███\u{1b}[0m"
361 );
362
363 assert_eq!(
364 Tokens(" {1} {5} {9} a").render(&colors_shim, 4, 10, true),
365 "\u{1b}[39;1m\u{1b}[0m\u{1b}[39;1m\u{1b}[0m\u{1b}[39;1m \u{1b}[0m\u{1b}[39;1m a\u{1b}[0m "
366 );
367
368 assert_eq!(
370 Tokens(" ").render(&colors_shim, 0, 0, false),
371 "\u{1b}[39m\u{1b}[0m"
372 );
373 assert_eq!(
374 Tokens(" ").render(&colors_shim, 0, 5, false),
375 "\u{1b}[39m \u{1b}[0m"
376 );
377 }
378
379 #[test]
380 fn truncate() {
381 assert_eq!(
382 Tokens("").truncate(0, 0).collect::<Vec<_>>(),
383 Tokens("").collect::<Vec<_>>()
384 );
385
386 assert_eq!(
387 Tokens(" ").truncate(0, 0).collect::<Vec<_>>(),
388 Tokens("").collect::<Vec<_>>()
389 );
390 assert_eq!(
391 Tokens(" ").truncate(0, 5).collect::<Vec<_>>(),
392 Tokens(" ").collect::<Vec<_>>()
393 );
394 assert_eq!(
395 Tokens(" ").truncate(1, 5).collect::<Vec<_>>(),
396 Tokens(" ").collect::<Vec<_>>()
397 );
398 assert_eq!(
399 Tokens(" ").truncate(3, 5).collect::<Vec<_>>(),
400 Tokens(" ").collect::<Vec<_>>()
401 );
402 assert_eq!(
403 Tokens(" ").truncate(0, 4).collect::<Vec<_>>(),
404 Tokens(" ").collect::<Vec<_>>()
405 );
406 assert_eq!(
407 Tokens(" ").truncate(0, 3).collect::<Vec<_>>(),
408 Tokens(" ").collect::<Vec<_>>()
409 );
410
411 assert_eq!(
412 Tokens(" {1} {5} {9} a")
413 .truncate(4, 10)
414 .collect::<Vec<_>>(),
415 Tokens("{1}{5} {9} a").collect::<Vec<_>>()
416 );
417 }
418}