1use crate::array::{Array, Data};
5use crate::dtype::DType;
6
7#[derive(Clone, Copy, Debug, PartialEq, Eq)]
9pub enum BoxStyle {
10 Fenced,
12 Spaced,
16}
17
18#[derive(Clone, Copy, Debug)]
20pub struct FmtOpts {
21 pub neg: char,
23 pub imag: char,
26 pub boxes: BoxStyle,
27}
28
29impl FmtOpts {
30 pub const J: FmtOpts = FmtOpts { neg: '_', imag: 'j', boxes: BoxStyle::Fenced };
31 pub const APL: FmtOpts = FmtOpts { neg: '¯', imag: 'J', boxes: BoxStyle::Spaced };
32}
33
34const SIG_DIGITS: usize = 6;
36
37pub fn format_array(a: &Array, opts: &FmtOpts) -> String {
39 if a.shape.contains(&0) {
41 return String::new();
42 }
43 if !a.is_row_major() {
47 return format_array(&a.to_row_major(), opts);
48 }
49 if a.dtype() == DType::Box {
50 match mixed_simple_texts(a, opts) {
54 Some(texts) if opts.boxes == BoxStyle::Spaced => {
55 return laid_out(&a.shape, texts, Cells::Right)
56 }
57 _ => return format_boxed(a, opts),
58 }
59 }
60 let texts: Vec<String> = (0..a.count()).map(|i| format_atom(&a.data, i, opts)).collect();
61 laid_out(&a.shape, texts, Cells::of(a.dtype()))
62}
63
64#[derive(Clone, Copy, PartialEq)]
66enum Cells {
67 Right,
69 Text,
71 Left,
74}
75
76impl Cells {
77 fn of(dtype: DType) -> Cells {
78 match dtype {
79 DType::Char => Cells::Text,
80 DType::Symbol => Cells::Left,
81 _ => Cells::Right,
82 }
83 }
84}
85
86fn mixed_simple_texts(a: &Array, opts: &FmtOpts) -> Option<Vec<String>> {
89 let boxes = a.as_boxes()?;
90 let mut texts = Vec::with_capacity(boxes.len());
91 for b in boxes {
92 if b.rank() != 0 || b.dtype() == DType::Box {
93 return None;
94 }
95 texts.push(format_atom(&b.data, 0, opts));
96 }
97 Some(texts)
98}
99
100fn laid_out(shape: &[usize], texts: Vec<String>, cells: Cells) -> String {
103 let rank = shape.len();
104 let a_shape = shape;
105 match rank {
106 0 => texts.into_iter().next().unwrap_or_default(),
107 1 if cells == Cells::Text => texts.concat(),
108 1 => texts.join(" "),
109 _ => {
110 let ncols = a_shape[rank - 1];
111 let nrows = a_shape[rank - 2];
112 let widths = if cells == Cells::Text {
115 vec![0; ncols]
116 } else {
117 column_widths(&texts, ncols)
118 };
119 let frame = &a_shape[..rank - 2];
120 let plane_size = nrows * ncols;
121 let planes: usize = frame.iter().product();
122 let mut out = String::new();
123 for p in 0..planes {
124 if p > 0 {
125 out.push_str(&"\n".repeat(plane_gap(frame, p) + 1));
127 }
128 for r in 0..nrows {
129 if r > 0 {
130 out.push('\n');
131 }
132 let start = p * plane_size + r * ncols;
133 push_row(&mut out, &texts[start..start + ncols], &widths, cells);
134 }
135 }
136 out
137 }
138 }
139}
140
141fn format_boxed(a: &Array, opts: &FmtOpts) -> String {
145 let boxes = a.as_boxes().expect("boxed data");
146 let blocks: Vec<(Vec<String>, usize)> = boxes.iter().map(|b| block(b, opts)).collect();
147 let rank = a.rank();
148 let (nrows, ncols) = match rank {
149 0 => (1, 1),
150 1 => (1, a.shape[0]),
151 _ => (a.shape[rank - 2], a.shape[rank - 1]),
152 };
153 let mut widths = vec![0usize; ncols];
155 for (i, (_, w)) in blocks.iter().enumerate() {
156 widths[i % ncols] = widths[i % ncols].max(*w);
157 }
158 let frame: &[usize] = if rank > 2 { &a.shape[..rank - 2] } else { &[] };
159 let planes: usize = frame.iter().product();
160 let plane_size = nrows * ncols;
161 let mut out = String::new();
162 for p in 0..planes.max(1) {
163 if p > 0 {
164 out.push_str(&"\n".repeat(plane_gap(frame, p) + 1));
165 }
166 push_boxed_plane(
167 &mut out,
168 &blocks[p * plane_size..(p + 1) * plane_size],
169 nrows,
170 ncols,
171 &widths,
172 opts,
173 );
174 }
175 out
176}
177
178fn block(a: &Array, opts: &FmtOpts) -> (Vec<String>, usize) {
187 if a.count() == 0 && a.rank() > 0 {
188 let rank = a.rank();
189 let rows: usize = a.shape[..rank - 1].iter().product();
190 let w = a.shape[rank - 1];
191 return (vec![" ".repeat(w); rows], w);
192 }
193 let text = format_array(a, opts);
194 if text.is_empty() {
195 return (vec![String::new()], 0);
196 }
197 let lines: Vec<String> = text.lines().map(str::to_string).collect();
198 let w = lines.iter().map(|l| width(l)).max().unwrap_or(0);
199 (lines, w)
200}
201
202fn push_boxed_plane(
203 out: &mut String,
204 blocks: &[(Vec<String>, usize)],
205 nrows: usize,
206 ncols: usize,
207 widths: &[usize],
208 opts: &FmtOpts,
209) {
210 let fence = opts.boxes == BoxStyle::Fenced;
211 let border: String = if fence {
212 let mut s = String::from("+");
213 for &w in widths {
214 s.push_str(&"-".repeat(w));
215 s.push('+');
216 }
217 s
218 } else {
219 String::new()
220 };
221 let mut lines: Vec<String> = Vec::new();
222 for r in 0..nrows {
223 if fence {
224 lines.push(border.clone());
225 }
226 let row = &blocks[r * ncols..(r + 1) * ncols];
227 let height = row.iter().map(|(lines, _)| lines.len()).max().unwrap_or(1);
230 for k in 0..height {
231 let mut line = String::new();
232 line.push(if fence { '|' } else { ' ' });
233 for (c, (cell, _)) in row.iter().enumerate() {
234 if !fence && c > 0 {
235 line.push(' ');
236 }
237 let text = cell.get(k).map(String::as_str).unwrap_or("");
238 line.push_str(text);
239 for _ in 0..widths[c].saturating_sub(width(text)) {
240 line.push(' ');
241 }
242 if fence {
243 line.push('|');
244 }
245 }
246 if !fence {
247 line.push(' ');
248 }
249 lines.push(line);
250 }
251 }
252 if fence {
253 lines.push(border);
254 }
255 out.push_str(&lines.join("\n"));
256}
257
258fn plane_gap(frame: &[usize], p: usize) -> usize {
261 let mut gap = 1;
264 let mut rest = p;
265 for &n in frame.iter().rev() {
266 if rest % n != 0 {
267 break;
268 }
269 rest /= n;
270 gap += 1;
271 }
272 gap
273}
274
275fn column_widths(texts: &[String], ncols: usize) -> Vec<usize> {
277 let mut widths = vec![0usize; ncols];
278 for (i, t) in texts.iter().enumerate() {
279 let j = i % ncols;
280 widths[j] = widths[j].max(width(t));
281 }
282 widths
283}
284
285fn push_row(out: &mut String, row: &[String], widths: &[usize], cells: Cells) {
286 for (j, cell) in row.iter().enumerate() {
287 if cells == Cells::Text {
288 out.push_str(cell);
289 continue;
290 }
291 if j > 0 {
292 out.push(' ');
293 }
294 let pad = widths[j].saturating_sub(width(cell));
295 if cells == Cells::Right {
296 for _ in 0..pad {
297 out.push(' ');
298 }
299 }
300 out.push_str(cell);
301 if cells == Cells::Left {
302 for _ in 0..pad {
303 out.push(' ');
304 }
305 }
306 }
307}
308
309fn width(s: &str) -> usize {
311 s.chars().count()
312}
313
314fn format_atom(data: &Data, i: usize, opts: &FmtOpts) -> String {
315 match data {
316 Data::Bool(v) => (if v[i] != 0 { "1" } else { "0" }).to_string(),
317 Data::I64(v) => format_i64(v[i], opts),
318 Data::Ext(v) => with_neg_sign(&v[i].to_string(), opts),
319 Data::Rat(v) => with_neg_sign(&v[i].to_string(), opts),
320 Data::F64(v) => format_f64(v[i], opts),
321 Data::Complex(v) => format_complex(v[i], opts),
322 Data::Char(v) => v[i].to_string(),
323 Data::Symbol(v) => format!("`{}", crate::symbol::name(v[i])),
325 Data::Box(_) => String::new(),
327 }
328}
329
330fn format_complex(z: crate::complex::Cx, opts: &FmtOpts) -> String {
335 if z[1] == 0.0 {
336 return format_f64(z[0], opts);
337 }
338 format!("{}{}{}", format_f64(z[0], opts), opts.imag, format_f64(z[1], opts))
339}
340
341fn format_i64(v: i64, opts: &FmtOpts) -> String {
342 with_neg_sign(&v.to_string(), opts)
343}
344
345fn with_neg_sign(s: &str, opts: &FmtOpts) -> String {
350 match s.strip_prefix('-') {
351 Some(rest) => with_sign(rest, opts),
352 None => s.to_string(),
353 }
354}
355
356fn format_f64(x: f64, opts: &FmtOpts) -> String {
357 if x.is_nan() {
358 return format!("{}.", opts.neg);
359 }
360 if x.is_infinite() {
361 return match (opts.neg, x > 0.0) {
363 ('_', true) => "_".to_string(),
364 ('_', false) => "__".to_string(),
365 (_, true) => "∞".to_string(),
366 (neg, false) => format!("{neg}∞"),
367 };
368 }
369 let magnitude = x.abs();
370 let sci = format!("{:.*e}", SIG_DIGITS - 1, magnitude);
373 let (mantissa, exponent) = sci.split_once('e').expect("scientific form has an exponent");
374 let digits: String = mantissa.chars().filter(char::is_ascii_digit).collect();
375 let exponent: i32 = exponent.parse().expect("exponent is an integer");
376 let body = if exponent >= 12 || exponent <= -6 {
377 let mut s = trim_fraction(&place_point(&digits, 1));
378 s.push('e');
379 if exponent < 0 {
380 s.push_str(&with_sign(&(-(exponent as i64)).to_string(), opts));
381 } else {
382 s.push_str(&exponent.to_string());
383 }
384 s
385 } else {
386 positional(&digits, exponent)
387 };
388 if x < 0.0 { with_sign(&body, opts) } else { body }
389}
390
391fn positional(digits: &str, exponent: i32) -> String {
393 if exponent < 0 {
394 let zeros = (-exponent - 1) as usize;
395 return trim_fraction(&format!("0.{}{}", "0".repeat(zeros), digits));
396 }
397 let int_len = exponent as usize + 1;
398 if int_len >= digits.len() {
399 return format!("{}{}", digits, "0".repeat(int_len - digits.len()));
402 }
403 trim_fraction(&place_point(digits, int_len))
404}
405
406fn place_point(digits: &str, int_len: usize) -> String {
408 format!("{}.{}", &digits[..int_len], &digits[int_len..])
409}
410
411fn trim_fraction(s: &str) -> String {
413 if !s.contains('.') {
414 return s.to_string();
415 }
416 s.trim_end_matches('0').trim_end_matches('.').to_string()
417}
418
419fn with_sign(body: &str, opts: &FmtOpts) -> String {
420 let mut s = String::with_capacity(body.len() + opts.neg.len_utf8());
421 s.push(opts.neg);
422 s.push_str(body);
423 s
424}
425
426#[cfg(test)]
427mod tests {
428 use super::*;
429 use crate::array::Buf;
430 use rstest::rstest;
431
432 fn j(a: &Array) -> String {
433 format_array(a, &FmtOpts::J)
434 }
435
436 fn apl(a: &Array) -> String {
437 format_array(a, &FmtOpts::APL)
438 }
439
440 fn fj(x: f64) -> String {
441 format_f64(x, &FmtOpts::J)
442 }
443
444 #[rstest]
447 #[case(0, "0")]
448 #[case(7, "7")]
449 #[case(-3, "_3")]
450 #[case(-1234, "_1234")]
451 #[case(i64::MIN, "_9223372036854775808")]
452 fn integers_j(#[case] v: i64, #[case] want: &str) {
453 assert_eq!(format_i64(v, &FmtOpts::J), want);
454 }
455
456 #[rstest]
457 #[case(-3, "¯3")]
458 #[case(3, "3")]
459 fn integers_apl(#[case] v: i64, #[case] want: &str) {
460 assert_eq!(format_i64(v, &FmtOpts::APL), want);
461 }
462
463 #[rstest]
464 #[case(0.0, "0")]
465 #[case(0.5, "0.5")]
466 #[case(2.0, "2")]
467 #[case(-2.0, "_2")]
468 #[case(1.0 / 3.0, "0.333333")]
469 #[case(-1.0 / 3.0, "_0.333333")]
470 #[case(2.0 / 3.0, "0.666667")]
471 #[case(1.25, "1.25")]
472 #[case(100.0, "100")]
473 #[case(1e-5, "0.00001")]
474 #[case(0.000012345678, "0.0000123457")]
475 #[case(1e11, "100000000000")]
476 #[case(123456789.0, "123457000")]
477 fn floats_positional(#[case] x: f64, #[case] want: &str) {
478 assert_eq!(fj(x), want);
479 }
480
481 #[rstest]
482 #[case(1e-7, "1e_7")]
483 #[case(-1e-7, "_1e_7")]
484 #[case(1.5e13, "1.5e13")]
485 #[case(1e12, "1e12")]
486 #[case(-2.5e20, "_2.5e20")]
487 #[case(1.234567e-9, "1.23457e_9")]
488 fn floats_exponent(#[case] x: f64, #[case] want: &str) {
489 assert_eq!(fj(x), want);
490 }
491
492 #[test]
493 fn floats_apl_signs() {
494 assert_eq!(format_f64(-0.5, &FmtOpts::APL), "¯0.5");
495 assert_eq!(format_f64(1e-7, &FmtOpts::APL), "1e¯7");
496 assert_eq!(format_f64(-1e-7, &FmtOpts::APL), "¯1e¯7");
497 }
498
499 #[test]
500 fn negative_zero_prints_unsigned() {
501 assert_eq!(fj(-0.0), "0");
502 }
503
504 #[test]
505 fn nan_and_infinities() {
506 assert_eq!(fj(f64::NAN), "_.");
507 assert_eq!(fj(f64::INFINITY), "_");
508 assert_eq!(fj(f64::NEG_INFINITY), "__");
509 assert_eq!(format_f64(f64::NAN, &FmtOpts::APL), "¯.");
510 assert_eq!(format_f64(f64::INFINITY, &FmtOpts::APL), "∞");
511 assert_eq!(format_f64(f64::NEG_INFINITY, &FmtOpts::APL), "¯∞");
512 }
513
514 #[test]
515 fn scalars() {
516 assert_eq!(j(&Array::scalar_i64(-3)), "_3");
517 assert_eq!(apl(&Array::scalar_i64(-3)), "¯3");
518 assert_eq!(j(&Array::scalar_f64(0.5)), "0.5");
519 assert_eq!(j(&Array::scalar_bool(true)), "1");
520 assert_eq!(j(&Array::scalar_bool(false)), "0");
521 assert_eq!(j(&Array::new(vec![], Data::Char(vec!['q'].into()))), "q");
522 }
523
524 #[test]
527 fn integer_vector() {
528 let a = Array::from_i64(vec![1, -22, 333]);
529 assert_eq!(j(&a), "1 _22 333");
530 assert_eq!(apl(&a), "1 ¯22 333");
531 }
532
533 #[test]
534 fn float_vector_trims_independently() {
535 let a = Array::from_f64(vec![0.5, 2.0, 1.0 / 3.0, -1e-7]);
536 assert_eq!(j(&a), "0.5 2 0.333333 _1e_7");
537 }
538
539 #[test]
540 fn bool_vector() {
541 let a = Array::new(vec![4], Data::Bool(vec![1, 0, 0, 1].into()));
542 assert_eq!(j(&a), "1 0 0 1");
543 }
544
545 #[test]
546 fn char_vector_is_a_plain_string() {
547 let a = Array::from_chars("hello".chars().collect());
548 assert_eq!(j(&a), "hello");
549 }
550
551 #[test]
554 fn matrix_columns_align_right() {
555 let a = Array::new(vec![2, 3], Data::I64(vec![1, 22, 333, 4444, 5, 66].into()));
556 assert_eq!(j(&a), " 1 22 333\n4444 5 66");
557 }
558
559 #[test]
560 fn matrix_negatives_widen_their_column() {
561 let a = Array::new(vec![2, 2], Data::I64(vec![-1, 10, 100, -2].into()));
562 assert_eq!(j(&a), " _1 10\n100 _2");
563 assert_eq!(apl(&a), " ¯1 10\n100 ¯2");
565 }
566
567 #[test]
568 fn matrix_of_floats() {
569 let a = Array::new(vec![2, 2], Data::F64(vec![0.5, 2.0, -1.0 / 3.0, 10.0].into()));
570 assert_eq!(j(&a), " 0.5 2\n_0.333333 10");
571 }
572
573 #[test]
574 fn matrix_of_bools() {
575 let a = Array::new(vec![2, 3], Data::Bool(vec![1, 0, 1, 0, 1, 0].into()));
576 assert_eq!(j(&a), "1 0 1\n0 1 0");
577 }
578
579 #[test]
580 fn single_column_matrix() {
581 let a = Array::new(vec![3, 1], Data::I64(vec![1, -20, 300].into()));
582 assert_eq!(j(&a), " 1\n_20\n300");
583 }
584
585 #[test]
588 fn rank_3_separates_planes_with_one_blank_line() {
589 let a = Array::new(vec![2, 2, 2], Data::I64(vec![1, 2, 3, 4, 5, 6, 7, 8].into()));
590 assert_eq!(j(&a), "1 2\n3 4\n\n5 6\n7 8");
591 }
592
593 #[test]
594 fn rank_3_column_widths_are_global() {
595 let a = Array::new(vec![2, 1, 2], Data::I64(vec![1, 2, 300, 4].into()));
596 assert_eq!(j(&a), " 1 2\n\n300 4");
597 }
598
599 #[test]
600 fn rank_4_separates_groups_with_two_blank_lines() {
601 let a = Array::new(vec![2, 2, 1, 2], Data::I64(vec![1, 2, 3, 4, 5, 6, 7, 8].into()));
602 assert_eq!(j(&a), "1 2\n\n3 4\n\n\n5 6\n\n7 8");
603 }
604
605 #[test]
606 fn rank_5_gap_grows_with_the_axis() {
607 let a = Array::new(vec![2, 1, 1, 1, 1], Data::I64(vec![1, 2].into()));
608 assert_eq!(j(&a), "1\n\n\n\n2");
610 }
611
612 #[rstest]
613 #[case(&[2], 1, 1)]
615 #[case(&[2, 3], 1, 1)]
617 #[case(&[2, 3], 2, 1)]
618 #[case(&[2, 3], 3, 2)]
619 #[case(&[2, 3], 4, 1)]
620 fn plane_gaps(#[case] frame: &[usize], #[case] p: usize, #[case] want: usize) {
621 assert_eq!(plane_gap(frame, p), want);
622 }
623
624 #[test]
627 fn char_matrix_is_lines() {
628 let a = Array::new(vec![2, 3], Data::Char("abcdef".chars().collect()));
629 assert_eq!(j(&a), "abc\ndef");
630 }
631
632 #[test]
633 fn char_matrix_keeps_spaces_unpadded() {
634 let a = Array::new(vec![2, 3], Data::Char("a bcd".chars().collect()));
635 assert_eq!(j(&a), "a \nbcd");
636 }
637
638 #[test]
639 fn char_rank_3_separates_planes() {
640 let a = Array::new(vec![2, 2, 2], Data::Char("abcdefgh".chars().collect()));
641 assert_eq!(j(&a), "ab\ncd\n\nef\ngh");
642 }
643
644 fn boxed(shape: &[usize], items: Vec<Array>) -> Array {
647 Array::new(shape.to_vec(), Data::Box(items.into()))
648 }
649
650 #[test]
651 fn a_box_is_drawn_as_a_fenced_cell() {
652 let a = boxed(&[], vec![Array::from_i64(vec![1, 2])]);
653 assert_eq!(j(&a), "+---+\n|1 2|\n+---+");
654 assert_eq!(apl(&a), " 1 2 ");
656 }
657
658 #[test]
659 fn a_boxed_vector_is_a_row_of_cells() {
660 let a = boxed(
661 &[3],
662 vec![
663 Array::scalar_i64(1),
664 Array::from_i64(vec![2, 3]),
665 Array::from_chars("abc".chars().collect()),
666 ],
667 );
668 assert_eq!(j(&a), "+-+---+---+\n|1|2 3|abc|\n+-+---+---+");
669 assert_eq!(apl(&a), " 1 2 3 abc ");
670 }
671
672 #[test]
673 fn a_tall_cell_pads_the_others_below_it() {
674 let a = boxed(
675 &[2],
676 vec![
677 Array::scalar_i64(1),
678 Array::new(vec![2, 2], Data::I64(vec![1, 2, 3, 4].into())),
679 ],
680 );
681 assert_eq!(j(&a), "+-+---+\n|1|1 2|\n| |3 4|\n+-+---+");
682 }
683
684 #[test]
685 fn a_nested_box_draws_inside_its_cell() {
686 let inner = boxed(&[], vec![Array::scalar_i64(5)]);
687 assert_eq!(j(&boxed(&[], vec![inner])), "+---+\n|+-+|\n||5||\n|+-+|\n+---+");
688 }
689
690 #[test]
691 fn a_box_matrix_fences_every_row() {
692 let a = boxed(&[2, 2], (1..=4).map(Array::scalar_i64).collect());
693 assert_eq!(j(&a), "+-+-+\n|1|2|\n+-+-+\n|3|4|\n+-+-+");
694 assert_eq!(apl(&a), "1 2\n3 4");
697 }
698
699 #[test]
700 fn a_boxed_empty_is_a_cell_of_width_zero() {
701 let a = boxed(&[], vec![Array::empty(DType::I64)]);
702 assert_eq!(j(&a), "++\n||\n++");
703 assert_eq!(j(&Array::new(vec![0], Data::Box(Buf::new()))), "");
705 }
706
707 #[rstest]
710 #[case(DType::Bool)]
711 #[case(DType::I64)]
712 #[case(DType::F64)]
713 #[case(DType::Char)]
714 fn empty_vectors_print_nothing(#[case] dtype: DType) {
715 assert_eq!(j(&Array::empty(dtype)), "");
716 }
717
718 #[rstest]
719 #[case(&[0, 3])]
720 #[case(&[3, 0])]
721 #[case(&[2, 0, 4])]
722 fn any_empty_axis_prints_nothing(#[case] shape: &[usize]) {
723 let a = Array::new(shape.to_vec(), Data::I64(vec![].into()));
724 assert_eq!(j(&a), "");
725 }
726
727 #[test]
728 fn no_trailing_newline_or_spaces() {
729 let a = Array::new(vec![2, 2, 2], Data::I64(vec![1, 22, 3, 4, 5, 6, 7, 8].into()));
730 let s = j(&a);
731 assert!(!s.ends_with('\n'));
732 for line in s.lines() {
733 assert_eq!(line.trim_end(), line, "line has trailing space: {line:?}");
734 }
735 }
736}