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| {
209 match token {
210 Token::Char(chr) => {
211 width = width.saturating_sub(1);
212 colored_segment.push(chr);
213 }
214 Token::Color(col) => {
215 add_styled_segment(&mut whole_string, &colored_segment, *color, bold);
216 colored_segment = String::new();
217 color = colors
218 .get(col as usize)
219 .unwrap_or(&DynColors::Ansi(AnsiColors::Default));
220 }
221 Token::Space => {
222 width = width.saturating_sub(1);
223 colored_segment.push(' ')
224 }
225 };
226 });
227
228 add_styled_segment(&mut whole_string, &colored_segment, *color, bold);
229 (0..width).for_each(|_| whole_string.push(' '));
230 whole_string
231 }
232}
233
234fn succeed_when<I>(predicate: impl FnOnce(I) -> bool) -> impl FnOnce(I) -> Option<()> {
237 |input| {
238 if predicate(input) {
239 Some(())
240 } else {
241 None
242 }
243 }
244}
245
246fn add_styled_segment(base: &mut String, segment: &str, color: DynColors, bold: bool) {
247 let mut style = Style::new().color(color);
248 if bold {
249 style = style.bold();
250 }
251 let formatted_segment = segment.style(style);
252 let _ = write!(base, "{formatted_segment}");
253}
254
255type ParseResult<'a, R> = Option<(&'a str, R)>;
258
259fn token<R>(s: &str, predicate: impl FnOnce(char) -> Option<R>) -> ParseResult<R> {
260 let mut chars = s.chars();
261 let token = chars.next()?;
262 let result = predicate(token)?;
263 Some((chars.as_str(), result))
264}
265
266fn color_token(s: &str) -> ParseResult<Token> {
270 let (s, _) = token(s, succeed_when(|c| c == '{'))?;
271 let (s, color_index) = token(s, |c| c.to_digit(10))?;
272 let (s, _) = token(s, succeed_when(|c| c == '}'))?;
273 Some((s, Token::Color(color_index)))
274}
275
276fn space_token(s: &str) -> ParseResult<Token> {
278 token(s, succeed_when(|c| c == ' ')).map(|(s, _)| (s, Token::Space))
279}
280
281fn char_token(s: &str) -> ParseResult<Token> {
283 token(s, |c| Some(Token::Char(c)))
284}
285
286#[cfg(test)]
287mod test {
288 use super::*;
289
290 #[test]
291 fn test_get_min_start_max_end() {
292 let lines = [
293 " xxx",
294 " xxx",
295 " oo",
296 " o",
297 " xx",
298 ];
299 assert_eq!(get_min_start_max_end(&lines), (3, 29));
300 }
301
302 #[test]
303 fn space_parses() {
304 assert_eq!(space_token(" "), Some(("", Token::Space)));
305 assert_eq!(space_token(" hello"), Some(("hello", Token::Space)));
306 assert_eq!(space_token(" "), Some((" ", Token::Space)));
307 assert_eq!(space_token(" {1}{2}"), Some(("{1}{2}", Token::Space)));
308 }
309
310 #[test]
311 fn color_indicator_parses() {
312 assert_eq!(color_token("{1}"), Some(("", Token::Color(1))));
313 assert_eq!(color_token("{9} "), Some((" ", Token::Color(9))));
314 }
315
316 #[test]
317 fn leading_spaces_counts_correctly() {
318 assert_eq!(Tokens("").leading_spaces(), 0);
319 assert_eq!(Tokens(" ").leading_spaces(), 5);
320 assert_eq!(Tokens(" a;lksjf;a").leading_spaces(), 5);
321 assert_eq!(Tokens(" {1} {5} {9} a").leading_spaces(), 6);
322 }
323
324 #[test]
325 fn render() {
326 let colors_shim = Vec::new();
327
328 assert_eq!(
329 Tokens("").render(&colors_shim, 0, 0, true),
330 "\u{1b}[39;1m\u{1b}[0m"
331 );
332
333 assert_eq!(
334 Tokens(" ").render(&colors_shim, 0, 0, true),
335 "\u{1b}[39;1m\u{1b}[0m"
336 );
337
338 assert_eq!(
339 Tokens(" ").render(&colors_shim, 0, 5, true),
340 "\u{1b}[39;1m \u{1b}[0m"
341 );
342
343 assert_eq!(
344 Tokens(" ").render(&colors_shim, 1, 5, true),
345 "\u{1b}[39;1m \u{1b}[0m"
346 );
347
348 assert_eq!(
349 Tokens(" ").render(&colors_shim, 3, 5, true),
350 "\u{1b}[39;1m \u{1b}[0m"
351 );
352
353 assert_eq!(
354 Tokens(" ").render(&colors_shim, 0, 4, true),
355 "\u{1b}[39;1m \u{1b}[0m"
356 );
357
358 assert_eq!(
359 Tokens(" ").render(&colors_shim, 0, 3, true),
360 "\u{1b}[39;1m \u{1b}[0m"
361 );
362
363 assert_eq!(
365 Tokens("███").render(Vec::new().as_slice(), 0, 3, true),
366 "\u{1b}[39;1m███\u{1b}[0m"
367 );
368
369 assert_eq!(
370 Tokens(" {1} {5} {9} a").render(&colors_shim, 4, 10, true),
371 "\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 "
372 );
373
374 assert_eq!(
376 Tokens(" ").render(&colors_shim, 0, 0, false),
377 "\u{1b}[39m\u{1b}[0m"
378 );
379 assert_eq!(
380 Tokens(" ").render(&colors_shim, 0, 5, false),
381 "\u{1b}[39m \u{1b}[0m"
382 );
383 }
384
385 #[test]
386 fn truncate() {
387 assert_eq!(
388 Tokens("").truncate(0, 0).collect::<Vec<_>>(),
389 Tokens("").collect::<Vec<_>>()
390 );
391
392 assert_eq!(
393 Tokens(" ").truncate(0, 0).collect::<Vec<_>>(),
394 Tokens("").collect::<Vec<_>>()
395 );
396 assert_eq!(
397 Tokens(" ").truncate(0, 5).collect::<Vec<_>>(),
398 Tokens(" ").collect::<Vec<_>>()
399 );
400 assert_eq!(
401 Tokens(" ").truncate(1, 5).collect::<Vec<_>>(),
402 Tokens(" ").collect::<Vec<_>>()
403 );
404 assert_eq!(
405 Tokens(" ").truncate(3, 5).collect::<Vec<_>>(),
406 Tokens(" ").collect::<Vec<_>>()
407 );
408 assert_eq!(
409 Tokens(" ").truncate(0, 4).collect::<Vec<_>>(),
410 Tokens(" ").collect::<Vec<_>>()
411 );
412 assert_eq!(
413 Tokens(" ").truncate(0, 3).collect::<Vec<_>>(),
414 Tokens(" ").collect::<Vec<_>>()
415 );
416
417 assert_eq!(
418 Tokens(" {1} {5} {9} a")
419 .truncate(4, 10)
420 .collect::<Vec<_>>(),
421 Tokens("{1}{5} {9} a").collect::<Vec<_>>()
422 );
423 }
424}