1use crate::options::{ForceStyle, FormatOptions};
7use crate::scalar::{can_be_bare, count_escapes, count_newlines, escape_quoted};
8
9#[derive(Debug, Clone)]
11pub enum Context {
12 Struct {
14 first: bool,
15 is_root: bool,
16 force_multiline: bool,
17 inline_start: bool,
19 comma_positions: Vec<usize>,
21 open_brace_pos: Option<usize>,
23 },
24 Seq {
26 first: bool,
27 inline_start: bool,
29 },
30}
31
32pub struct StyxWriter {
39 out: Vec<u8>,
40 stack: Vec<Context>,
41 options: FormatOptions,
42 skip_next_before_value: bool,
44}
45
46impl StyxWriter {
47 pub fn new() -> Self {
49 Self::with_options(FormatOptions::default())
50 }
51
52 pub fn with_options(options: FormatOptions) -> Self {
54 Self {
55 out: Vec::new(),
56 stack: Vec::new(),
57 skip_next_before_value: false,
58 options,
59 }
60 }
61
62 pub fn finish(self) -> Vec<u8> {
64 self.out
65 }
66
67 pub fn finish_document(mut self) -> Vec<u8> {
70 if !self.out.is_empty() && !self.out.ends_with(b"\n") {
71 self.out.push(b'\n');
72 }
73 self.out
74 }
75
76 pub fn finish_string(self) -> String {
81 String::from_utf8(self.out).expect("Styx output should always be valid UTF-8")
82 }
83
84 pub fn depth(&self) -> usize {
86 self.stack.len()
87 }
88
89 fn indent_depth(&self) -> usize {
91 let mut depth = 0;
92 for ctx in &self.stack {
93 match ctx {
94 Context::Struct { is_root: true, .. } => {
95 }
97 Context::Struct {
98 inline_start: true,
99 force_multiline: false,
100 ..
101 } => {
102 }
104 Context::Seq {
105 inline_start: true, ..
106 } => {
107 }
109 _ => {
110 depth += 1;
111 }
112 }
113 }
114 depth
115 }
116
117 pub fn available_width(&self) -> usize {
119 let used = self.depth() * self.options.indent.len();
120 self.options.max_width.saturating_sub(used)
121 }
122
123 pub fn should_inline(&self) -> bool {
125 if self.options.force_style == ForceStyle::Inline {
126 return true;
127 } else if self.options.force_style == ForceStyle::Multiline {
128 return false;
129 }
130 if self.depth() == 0 {
132 return false;
133 }
134 if let Some(Context::Struct { is_root: true, .. }) = self.stack.first()
136 && self.depth() == 1
137 {
138 return false;
139 }
140 if let Some(Context::Struct {
142 force_multiline: true,
143 ..
144 }) = self.stack.last()
145 {
146 return false;
147 }
148 self.available_width() >= self.options.min_inline_width
150 }
151
152 pub fn write_indent(&mut self) {
154 for _ in 0..self.indent_depth() {
155 self.out.extend_from_slice(self.options.indent.as_bytes());
156 }
157 }
158
159 pub fn write_newline_indent(&mut self) {
161 self.out.push(b'\n');
162 self.write_indent();
163 }
164
165 pub fn write_raw(&mut self, bytes: &[u8]) {
167 self.out.extend_from_slice(bytes);
168 }
169
170 pub fn write_str(&mut self, s: &str) {
172 self.out.extend_from_slice(s.as_bytes());
173 }
174
175 pub fn write_byte(&mut self, b: u8) {
177 self.out.push(b);
178 }
179
180 pub fn begin_struct(&mut self, is_root: bool) {
184 self.begin_struct_with_options(is_root, false);
185 }
186
187 pub fn begin_struct_with_options(&mut self, is_root: bool, force_multiline: bool) {
192 self.before_value();
193
194 let inline_start = !is_root;
197
198 if is_root {
199 self.stack.push(Context::Struct {
200 first: true,
201 is_root: true,
202 force_multiline,
203 inline_start: false,
204 comma_positions: Vec::new(),
205 open_brace_pos: None,
206 });
207 } else {
208 self.out.push(b'{');
209 let open_pos = self.out.len(); self.stack.push(Context::Struct {
211 first: true,
212 is_root: false,
213 force_multiline,
214 inline_start,
215 comma_positions: Vec::new(),
216 open_brace_pos: Some(open_pos),
217 });
218 }
219 }
220
221 pub fn begin_struct_after_tag(&mut self, force_multiline: bool) {
223 self.out.push(b'{');
225 let open_pos = self.out.len(); self.stack.push(Context::Struct {
227 first: true,
228 is_root: false,
229 force_multiline,
230 inline_start: true,
231 comma_positions: Vec::new(),
232 open_brace_pos: Some(open_pos),
233 });
234 }
235
236 pub fn field_key_raw(&mut self, key: &str) -> Result<(), &'static str> {
240 let (is_struct, is_first, is_root, inline_start, force_multiline) = match self.stack.last()
242 {
243 Some(Context::Struct {
244 first,
245 is_root,
246 inline_start,
247 force_multiline,
248 ..
249 }) => (true, *first, *is_root, *inline_start, *force_multiline),
250 _ => (false, true, false, false, false),
251 };
252
253 if !is_struct {
254 return Err("field_key_raw called outside of struct");
255 }
256
257 let struct_is_inline = inline_start && !force_multiline;
259 let should_inline = struct_is_inline || self.should_inline();
260
261 if !is_first {
262 if should_inline && !is_root {
263 let comma_pos = self.out.len();
265 self.out.extend_from_slice(b", ");
266 if let Some(Context::Struct {
267 comma_positions, ..
268 }) = self.stack.last_mut()
269 {
270 comma_positions.push(comma_pos);
271 }
272 } else {
273 if is_root {
275 self.out.push(b'\n');
276 }
277 self.write_newline_indent();
278 }
279 } else {
280 if !is_root && !should_inline {
282 self.write_newline_indent();
283 }
284 }
285
286 if let Some(Context::Struct { first, .. }) = self.stack.last_mut() {
288 *first = false;
289 }
290
291 self.out.extend_from_slice(key.as_bytes());
293 self.out.push(b' ');
294 Ok(())
295 }
296
297 pub fn field_key(&mut self, key: &str) -> Result<(), &'static str> {
301 let (is_struct, is_first, is_root, inline_start, force_multiline) = match self.stack.last()
303 {
304 Some(Context::Struct {
305 first,
306 is_root,
307 inline_start,
308 force_multiline,
309 ..
310 }) => (true, *first, *is_root, *inline_start, *force_multiline),
311 _ => (false, true, false, false, false),
312 };
313
314 if !is_struct {
315 return Err("field_key called outside of struct");
316 }
317
318 let struct_is_inline = inline_start && !force_multiline;
320 let should_inline = struct_is_inline || self.should_inline();
321
322 if !is_first {
323 if should_inline && !is_root {
324 let comma_pos = self.out.len();
326 self.out.extend_from_slice(b", ");
327 if let Some(Context::Struct {
328 comma_positions, ..
329 }) = self.stack.last_mut()
330 {
331 comma_positions.push(comma_pos);
332 }
333 } else {
334 if is_root {
336 self.out.push(b'\n');
337 }
338 self.write_newline_indent();
339 }
340 } else {
341 if !is_root && !should_inline {
343 self.write_newline_indent();
344 }
345 }
346
347 if let Some(Context::Struct { first, .. }) = self.stack.last_mut() {
349 *first = false;
350 }
351
352 if can_be_bare(key) {
354 self.out.extend_from_slice(key.as_bytes());
355 } else {
356 self.write_quoted_string(key);
357 }
358 self.out.push(b' ');
359 Ok(())
360 }
361
362 pub fn end_struct(&mut self) -> Result<(), &'static str> {
366 let should_inline = self.should_inline();
368
369 match self.stack.pop() {
370 Some(Context::Struct {
371 first,
372 is_root,
373 force_multiline,
374 inline_start,
375 ..
376 }) => {
377 if is_root {
378 if !first {
380 self.out.push(b'\n');
381 }
382 } else {
383 let needs_newline =
387 !first && (force_multiline || (!inline_start && !should_inline));
388 if needs_newline {
389 self.out.push(b'\n');
391 self.write_indent();
393 }
394 self.out.push(b'}');
395 }
396 Ok(())
397 }
398 _ => Err("end_struct called without matching begin_struct"),
399 }
400 }
401
402 pub fn begin_seq(&mut self) {
404 self.before_value();
405 self.out.push(b'(');
406 self.stack.push(Context::Seq {
408 first: true,
409 inline_start: true,
410 });
411 }
412
413 pub fn end_seq(&mut self) -> Result<(), &'static str> {
417 let should_inline = self.should_inline();
419
420 match self.stack.pop() {
421 Some(Context::Seq {
422 first,
423 inline_start,
424 }) => {
425 if !first && !inline_start && !should_inline {
428 self.write_newline_indent();
429 }
430 self.out.push(b')');
431 Ok(())
432 }
433 _ => Err("end_seq called without matching begin_seq"),
434 }
435 }
436
437 pub fn write_null(&mut self) {
439 self.before_value();
440 self.out.push(b'@');
441 }
442
443 pub fn write_bool(&mut self, v: bool) {
445 self.before_value();
446 if v {
447 self.out.extend_from_slice(b"true");
448 } else {
449 self.out.extend_from_slice(b"false");
450 }
451 }
452
453 pub fn write_i64(&mut self, v: i64) {
455 self.before_value();
456 self.out.extend_from_slice(v.to_string().as_bytes());
457 }
458
459 pub fn write_u64(&mut self, v: u64) {
461 self.before_value();
462 self.out.extend_from_slice(v.to_string().as_bytes());
463 }
464
465 pub fn write_i128(&mut self, v: i128) {
467 self.before_value();
468 self.out.extend_from_slice(v.to_string().as_bytes());
469 }
470
471 pub fn write_u128(&mut self, v: u128) {
473 self.before_value();
474 self.out.extend_from_slice(v.to_string().as_bytes());
475 }
476
477 pub fn write_f64(&mut self, v: f64) {
479 self.before_value();
480 self.out.extend_from_slice(v.to_string().as_bytes());
481 }
482
483 pub fn write_string(&mut self, s: &str) {
485 self.before_value();
486 self.write_scalar_string(s);
487 }
488
489 pub fn write_char(&mut self, c: char) {
491 self.before_value();
492 let mut buf = [0u8; 4];
493 let s = c.encode_utf8(&mut buf);
494 self.write_scalar_string(s);
495 }
496
497 pub fn write_bytes(&mut self, bytes: &[u8]) {
499 self.before_value();
500 self.out.push(b'"');
501 for byte in bytes.iter() {
502 let hex = |d: u8| {
503 if d < 10 { b'0' + d } else { b'a' + (d - 10) }
504 };
505 self.out.push(hex(byte >> 4));
506 self.out.push(hex(byte & 0xf));
507 }
508 self.out.push(b'"');
509 }
510
511 pub fn write_variant_tag(&mut self, name: &str) {
513 self.before_value();
514 self.out.push(b'@');
515 self.out.extend_from_slice(name.as_bytes());
516 self.skip_next_before_value = true;
518 }
519
520 pub fn write_scalar(&mut self, s: &str) {
523 self.write_string(s);
524 }
525
526 pub fn write_tag(&mut self, name: &str) {
528 self.write_variant_tag(name);
529 }
530
531 pub fn clear_skip_before_value(&mut self) {
534 self.skip_next_before_value = false;
535 }
536
537 pub fn begin_seq_after_tag(&mut self) {
539 self.out.push(b'(');
540 self.stack.push(Context::Seq {
542 first: true,
543 inline_start: true,
544 });
545 }
546
547 pub fn write_doc_comment_and_key(&mut self, doc: &str, key: &str) {
550 let (is_first, is_root) = match self.stack.last() {
552 Some(Context::Struct { first, is_root, .. }) => (*first, *is_root),
553 _ => (true, false),
554 };
555
556 if let Some(Context::Struct {
558 first,
559 force_multiline,
560 ..
561 }) = self.stack.last_mut()
562 {
563 *first = false;
564 *force_multiline = true;
565 }
566
567 self.fix_comma_separators();
569
570 self.propagate_multiline_to_parents();
572
573 if !is_first && is_root {
576 self.out.push(b'\n');
578 }
579
580 let need_leading_newline = !is_first || !is_root;
582
583 for (i, line) in doc.lines().enumerate() {
584 if i > 0 || need_leading_newline {
585 self.write_newline_indent();
586 }
587 self.out.extend_from_slice(b"/// ");
588 self.out.extend_from_slice(line.as_bytes());
589 }
590
591 if is_first && is_root {
593 self.out.push(b'\n');
594 } else {
595 self.write_newline_indent();
596 }
597
598 if can_be_bare(key) {
600 self.out.extend_from_slice(key.as_bytes());
601 } else {
602 self.write_quoted_string(key);
603 }
604 self.out.push(b' ');
605 }
606
607 pub fn write_doc_comment_and_key_raw(&mut self, doc: &str, key: &str) {
610 let (is_first, is_root) = match self.stack.last() {
612 Some(Context::Struct { first, is_root, .. }) => (*first, *is_root),
613 _ => (true, false),
614 };
615
616 if let Some(Context::Struct {
618 first,
619 force_multiline,
620 ..
621 }) = self.stack.last_mut()
622 {
623 *first = false;
624 *force_multiline = true;
625 }
626
627 self.fix_comma_separators();
629
630 self.propagate_multiline_to_parents();
632
633 if !is_first && is_root {
636 self.out.push(b'\n');
638 }
639
640 let need_leading_newline = !is_first || !is_root;
642
643 for (i, line) in doc.lines().enumerate() {
644 if i > 0 || need_leading_newline {
645 self.write_newline_indent();
646 }
647 self.out.extend_from_slice(b"/// ");
648 self.out.extend_from_slice(line.as_bytes());
649 }
650
651 if is_first && is_root {
653 self.out.push(b'\n');
654 } else {
655 self.write_newline_indent();
656 }
657
658 self.out.extend_from_slice(key.as_bytes());
660 self.out.push(b' ');
661 }
662
663 fn propagate_multiline_to_parents(&mut self) {
670 let mut all_comma_positions: Vec<usize> = Vec::new();
673 let mut structs_to_fix: Vec<usize> = Vec::new(); for (idx, ctx) in self.stack.iter().enumerate() {
676 if let Context::Struct {
677 inline_start,
678 force_multiline,
679 is_root,
680 comma_positions,
681 ..
682 } = ctx
683 && !*is_root
684 && *inline_start
685 && !*force_multiline
686 {
687 structs_to_fix.push(idx);
688 all_comma_positions.extend(comma_positions.iter().copied());
689 }
690 }
691
692 if !all_comma_positions.is_empty() {
694 all_comma_positions.sort_unstable();
696 for &comma_pos in all_comma_positions.iter().rev() {
697 if comma_pos + 2 <= self.out.len()
699 && self.out[comma_pos] == b','
700 && self.out[comma_pos + 1] == b' '
701 {
702 let indent = self.options.indent.repeat(self.indent_depth());
704 let newline_indent = format!("\n{}", indent);
705 self.out.drain(comma_pos..comma_pos + 2);
706 let bytes = newline_indent.as_bytes();
707 for (i, &b) in bytes.iter().enumerate() {
708 self.out.insert(comma_pos + i, b);
709 }
710 }
711 }
712 }
713
714 for &idx in &structs_to_fix {
716 if let Some(Context::Struct {
717 comma_positions, ..
718 }) = self.stack.get_mut(idx)
719 {
720 comma_positions.clear();
721 }
722 }
723
724 let mut fixes: Vec<(usize, usize)> = Vec::new(); let mut effective_depth = 0;
728 for ctx in self.stack.iter_mut() {
729 if let Context::Struct {
730 inline_start,
731 force_multiline,
732 is_root,
733 open_brace_pos: Some(pos),
734 ..
735 } = ctx
736 {
737 if !*is_root {
739 if *inline_start && !*force_multiline {
741 *force_multiline = true;
742 fixes.push((*pos, effective_depth + 1));
744 }
745 effective_depth += 1;
747 }
748 }
749 }
750
751 for (pos, indent_depth) in fixes.into_iter().rev() {
753 let indent = self.options.indent.repeat(indent_depth);
755 let insert = format!("\n{}", indent);
756 for (i, b) in insert.bytes().enumerate() {
757 self.out.insert(pos + i, b);
758 }
759
760 for ctx in self.stack.iter_mut() {
762 if let Context::Struct {
763 open_brace_pos: Some(p),
764 comma_positions,
765 ..
766 } = ctx
767 {
768 if *p > pos {
769 *p += insert.len();
770 }
771 for cp in comma_positions.iter_mut() {
772 if *cp >= pos {
773 *cp += insert.len();
774 }
775 }
776 }
777 }
778 }
779 }
780
781 fn fix_comma_separators(&mut self) {
784 let comma_positions = match self.stack.last_mut() {
786 Some(Context::Struct {
787 comma_positions, ..
788 }) => std::mem::take(comma_positions),
789 _ => return,
790 };
791
792 if comma_positions.is_empty() {
793 return;
794 }
795
796 let indent = self.options.indent.repeat(self.indent_depth());
798 let newline_indent = format!("\n{}", indent);
799
800 for &comma_pos in comma_positions.iter().rev() {
802 if comma_pos + 2 <= self.out.len()
805 && self.out[comma_pos] == b','
806 && self.out[comma_pos + 1] == b' '
807 {
808 self.out.drain(comma_pos..comma_pos + 2);
810 let bytes = newline_indent.as_bytes();
811 for (i, &b) in bytes.iter().enumerate() {
812 self.out.insert(comma_pos + i, b);
813 }
814 }
815 }
816 }
817
818 pub fn before_value(&mut self) {
820 if self.skip_next_before_value {
822 self.skip_next_before_value = false;
823 if let Some(Context::Seq { first, .. }) = self.stack.last_mut() {
825 *first = false;
826 }
827 return;
828 }
829
830 let (is_seq, is_first, inline_start) = match self.stack.last() {
832 Some(Context::Seq {
833 first,
834 inline_start,
835 }) => (true, *first, *inline_start),
836 _ => (false, true, false),
837 };
838
839 if is_seq && !is_first {
840 if inline_start || self.should_inline() {
842 self.out.push(b' ');
843 } else {
844 self.write_newline_indent();
845 }
846 }
847
848 if let Some(Context::Seq { first, .. }) = self.stack.last_mut() {
850 *first = false;
851 }
852 }
853
854 fn write_scalar_string(&mut self, s: &str) {
856 if can_be_bare(s) {
858 self.out.extend_from_slice(s.as_bytes());
859 return;
860 }
861
862 let newline_count = count_newlines(s);
863 let escape_count = count_escapes(s);
864
865 if newline_count >= self.options.heredoc_line_threshold {
867 self.write_heredoc(s);
868 return;
869 }
870
871 if escape_count > 3 && !s.contains("\"#") {
873 self.write_raw_string(s);
874 return;
875 }
876
877 self.write_quoted_string(s);
879 }
880
881 fn write_quoted_string(&mut self, s: &str) {
883 self.out.push(b'"');
884 let escaped = escape_quoted(s);
885 self.out.extend_from_slice(escaped.as_bytes());
886 self.out.push(b'"');
887 }
888
889 fn write_raw_string(&mut self, s: &str) {
891 let mut hashes = 0;
893 let mut check = String::from("\"");
894 while s.contains(&check) {
895 hashes += 1;
896 check = format!("\"{}#", "#".repeat(hashes - 1));
897 }
898
899 self.out.push(b'r');
900 for _ in 0..hashes {
901 self.out.push(b'#');
902 }
903 self.out.push(b'"');
904 self.out.extend_from_slice(s.as_bytes());
905 self.out.push(b'"');
906 for _ in 0..hashes {
907 self.out.push(b'#');
908 }
909 }
910
911 fn write_heredoc(&mut self, s: &str) {
913 let delimiters = ["TEXT", "END", "HEREDOC", "DOC", "STR", "CONTENT"];
915 let delimiter = delimiters
916 .iter()
917 .find(|d| !s.contains(*d))
918 .unwrap_or(&"TEXT");
919
920 self.out.extend_from_slice(b"<<");
921 self.out.extend_from_slice(delimiter.as_bytes());
922 self.out.push(b'\n');
923 self.out.extend_from_slice(s.as_bytes());
924 if !s.ends_with('\n') {
925 self.out.push(b'\n');
926 }
927 self.out.extend_from_slice(delimiter.as_bytes());
928 }
929}
930
931impl Default for StyxWriter {
932 fn default() -> Self {
933 Self::new()
934 }
935}
936
937#[cfg(test)]
938mod tests {
939 use super::*;
940
941 #[test]
942 fn test_simple_struct() {
943 let mut w = StyxWriter::new();
944 w.begin_struct(true);
945 w.field_key("name").unwrap();
946 w.write_string("hello");
947 w.field_key("value").unwrap();
948 w.write_i64(42);
949 w.end_struct().unwrap();
950
951 let result = w.finish_string();
952 assert!(result.contains("name hello"));
953 assert!(result.contains("value 42"));
954 }
955
956 #[test]
957 fn test_nested_inline() {
958 let mut w = StyxWriter::with_options(FormatOptions::default());
959 w.begin_struct(true);
960 w.field_key("point").unwrap();
961 w.begin_struct(false);
962 w.field_key("x").unwrap();
963 w.write_i64(10);
964 w.field_key("y").unwrap();
965 w.write_i64(20);
966 w.end_struct().unwrap();
967 w.end_struct().unwrap();
968
969 let result = w.finish_string();
970 assert!(result.contains("{x 10, y 20}"));
972 }
973
974 #[test]
975 fn test_sequence() {
976 let mut w = StyxWriter::new();
977 w.begin_struct(true);
978 w.field_key("items").unwrap();
979 w.begin_seq();
980 w.write_i64(1);
981 w.write_i64(2);
982 w.write_i64(3);
983 w.end_seq().unwrap();
984 w.end_struct().unwrap();
985
986 let result = w.finish_string();
987 assert!(result.contains("items (1 2 3)"));
988 }
989
990 #[test]
991 fn test_quoted_string() {
992 let mut w = StyxWriter::new();
993 w.begin_struct(true);
994 w.field_key("message").unwrap();
995 w.write_string("hello world");
996 w.end_struct().unwrap();
997
998 let result = w.finish_string();
999 assert!(result.contains("message \"hello world\""));
1000 }
1001
1002 #[test]
1003 fn test_force_inline() {
1004 let mut w = StyxWriter::with_options(FormatOptions::default().inline());
1005 w.begin_struct(false);
1006 w.field_key("a").unwrap();
1007 w.write_i64(1);
1008 w.field_key("b").unwrap();
1009 w.write_i64(2);
1010 w.end_struct().unwrap();
1011
1012 let result = w.finish_string();
1013 assert_eq!(result, "{a 1, b 2}");
1014 }
1015
1016 #[test]
1017 fn test_doc_comment_fixes_commas() {
1018 let mut w = StyxWriter::with_options(FormatOptions::default().inline());
1021 w.begin_struct(false);
1022 w.field_key("a").unwrap();
1023 w.write_i64(1);
1024 w.field_key("b").unwrap();
1025 w.write_i64(2);
1026 w.write_doc_comment_and_key("A documented field", "c");
1028 w.write_i64(3);
1029 w.end_struct().unwrap();
1030
1031 let result = w.finish_string();
1032 assert!(
1034 !result.contains(", "),
1035 "Result should not contain commas after doc comment: {}",
1036 result
1037 );
1038 assert!(
1040 result.contains("a 1\n"),
1041 "Expected newline after a: {}",
1042 result
1043 );
1044 }
1045}