1use crate::resolver::{PlainScalarType, resolve_plain_scalar, value_tag_error};
4#[cfg(test)]
5use crate::scanner::Scanner;
6use crate::tag::TagResolver;
7use crate::version::YamlVersion;
8use crate::{
9 BasicParser, Error, Limits, Parser, Position, ResourceTracker, Result, Value, parser::EventType,
10};
11use indexmap::IndexMap;
12use std::collections::HashMap;
13
14pub(crate) fn calculate_value_complexity(value: &Value) -> Result<usize> {
20 let mut total: usize = 0;
21 let mut stack: Vec<&Value> = vec![value];
22 while let Some(node) = stack.pop() {
23 match node {
24 Value::Sequence(seq) => {
25 total = total.saturating_add(1usize.saturating_add(seq.len()));
26 for item in seq {
27 stack.push(item);
28 }
29 }
30 Value::Mapping(map) => {
31 total = total.saturating_add(1usize.saturating_add(map.len().saturating_mul(2)));
32 for (k, v) in map {
33 stack.push(k);
34 stack.push(v);
35 }
36 }
37 _ => total = total.saturating_add(1),
38 }
39 }
40 Ok(total)
41}
42
43fn calculate_structure_depth(value: &Value) -> usize {
48 let mut max_depth: usize = 1;
49 let mut stack: Vec<(&Value, usize)> = vec![(value, 1)];
50 while let Some((node, depth)) = stack.pop() {
51 if depth > max_depth {
52 max_depth = depth;
53 }
54 let next = depth.saturating_add(1);
55 match node {
56 Value::Sequence(seq) => {
57 for item in seq {
58 stack.push((item, next));
59 }
60 }
61 Value::Mapping(map) => {
62 for (_, v) in map {
63 stack.push((v, next));
64 }
65 }
66 _ => {}
67 }
68 }
69 max_depth
70}
71
72pub trait Composer {
74 fn check_document(&self) -> bool;
76
77 fn compose_document(&mut self) -> Result<Option<Value>>;
82
83 fn position(&self) -> Position;
85
86 fn reset(&mut self);
88}
89
90#[derive(Debug)]
92pub struct BasicComposer {
93 parser: BasicParser,
94 position: Position,
95 anchors: HashMap<String, Value>,
96 limits: Limits,
97 resource_tracker: ResourceTracker,
98 alias_expansion_stack: Vec<String>,
99 current_depth: usize,
100 tag_resolver: TagResolver,
101 yaml_version: YamlVersion,
105}
106
107impl BasicComposer {
108 #[must_use]
110 pub fn new(input: String) -> Self {
111 Self::with_limits(input, Limits::default())
112 }
113
114 #[must_use]
116 pub fn with_limits(input: String, limits: Limits) -> Self {
117 Self {
118 parser: BasicParser::with_limits(input, limits.clone()),
119 position: Position::new(),
120 anchors: HashMap::new(),
121 limits,
122 resource_tracker: ResourceTracker::new(),
123 alias_expansion_stack: Vec::new(),
124 current_depth: 0,
125 tag_resolver: TagResolver::new(),
126 yaml_version: YamlVersion::default(),
127 }
128 }
129
130 #[must_use]
132 pub fn new_eager(input: String) -> Self {
133 Self::new_eager_with_limits(input, Limits::default())
134 }
135
136 #[must_use]
138 pub fn new_eager_with_limits(input: String, limits: Limits) -> Self {
139 Self {
140 parser: BasicParser::new_eager_with_limits(input, limits.clone()),
141 position: Position::new(),
142 anchors: HashMap::new(),
143 limits,
144 resource_tracker: ResourceTracker::new(),
145 alias_expansion_stack: Vec::new(),
146 current_depth: 0,
147 tag_resolver: TagResolver::new(),
148 yaml_version: YamlVersion::default(),
149 }
150 }
151
152 fn compose_node(&mut self) -> Result<Option<Value>> {
154 if !self.parser.check_event() {
155 return Ok(None);
156 }
157
158 let Some(event) = self.parser.get_event()? else {
159 return Ok(None);
160 };
161
162 self.position = event.position;
163
164 match event.event_type {
165 EventType::StreamStart | EventType::StreamEnd => {
166 self.compose_node()
168 }
169
170 EventType::DocumentStart { version, .. } => {
171 self.yaml_version = version
177 .map(|(maj, min)| YamlVersion::from_directive(maj, min))
178 .unwrap_or_default();
179 self.compose_node()
180 }
181
182 EventType::DocumentEnd { .. } => {
183 Ok(None)
185 }
186
187 EventType::Scalar {
188 value,
189 anchor,
190 tag,
191 style,
192 ..
193 } => {
194 let scalar_value = if let Some(tag_str) = tag {
195 self.compose_tagged_scalar(value, tag_str)?
197 } else {
198 self.compose_scalar(value, style, event.position)?
200 };
201
202 if let Some(anchor_name) = anchor {
204 self.resource_tracker.add_anchor(&self.limits)?;
205 self.anchors.insert(anchor_name, scalar_value.clone());
206 }
207
208 Ok(Some(scalar_value))
209 }
210
211 EventType::SequenceStart { anchor, .. } => {
212 let sequence = self.compose_sequence()?;
213
214 if let Some(anchor_name) = anchor {
216 if let Some(ref seq) = sequence {
217 self.resource_tracker.add_anchor(&self.limits)?;
218 self.anchors.insert(anchor_name, seq.clone());
219 }
220 }
221
222 Ok(sequence)
223 }
224
225 EventType::MappingStart { anchor, .. } => {
226 let mapping = self.compose_mapping()?;
227
228 if let Some(anchor_name) = anchor {
230 if let Some(ref map) = mapping {
231 self.resource_tracker.add_anchor(&self.limits)?;
232 self.anchors.insert(anchor_name, map.clone());
233 }
234 }
235
236 Ok(mapping)
237 }
238
239 EventType::SequenceEnd | EventType::MappingEnd => {
240 Ok(None)
245 }
246
247 EventType::Alias { anchor } => {
248 if self.alias_expansion_stack.contains(&anchor) {
250 return Err(Error::construction(
251 event.position,
252 format!("Cyclic alias reference detected: '{anchor}'"),
253 ));
254 }
255
256 if self.alias_expansion_stack.len() >= self.limits.max_alias_depth {
258 return Err(Error::construction(
259 event.position,
260 format!(
261 "Maximum alias expansion depth {} exceeded",
262 self.limits.max_alias_depth
263 ),
264 ));
265 }
266
267 self.resource_tracker.enter_alias(&self.limits)?;
269 self.alias_expansion_stack.push(anchor.clone());
270
271 let result = match self.anchors.get(&anchor) {
273 Some(value) => {
274 let structure_depth = calculate_structure_depth(value);
276 if structure_depth > self.limits.max_alias_depth {
277 return Err(Error::construction(
278 event.position,
279 format!(
280 "Alias '{}' creates structure with depth {} exceeding max_alias_depth {}",
281 anchor, structure_depth, self.limits.max_alias_depth
282 ),
283 ));
284 }
285
286 let nodes = calculate_value_complexity(value)?;
287 self.resource_tracker
292 .add_alias_materialization(&self.limits, nodes)?;
293 self.resource_tracker.add_complexity(&self.limits, nodes)?;
294 Ok(Some(value.clone()))
295 }
296 None => Err(Error::construction(
297 event.position,
298 format!("Unknown anchor '{anchor}'"),
299 )),
300 };
301
302 self.alias_expansion_stack.pop();
304 self.resource_tracker.exit_alias();
305
306 result
307 }
308 }
309 }
310
311 fn compose_scalar(
321 &self,
322 value: String,
323 style: crate::parser::ScalarStyle,
324 position: crate::Position,
325 ) -> Result<Value> {
326 match style {
327 crate::parser::ScalarStyle::SingleQuoted | crate::parser::ScalarStyle::DoubleQuoted => {
328 return Ok(Value::String(value));
329 }
330 _ => {}
331 }
332
333 Ok(match resolve_plain_scalar(&value, self.yaml_version) {
334 PlainScalarType::Null => Value::Null,
335 PlainScalarType::Bool(b) => Value::Bool(b),
336 PlainScalarType::Int(i) => Value::Int(i),
337 PlainScalarType::Float(f) => Value::Float(f),
338 PlainScalarType::Str => Value::String(value),
339 PlainScalarType::Value => return Err(value_tag_error(position)),
340 })
341 }
342
343 fn compose_tagged_scalar(&mut self, value: String, tag_str: String) -> Result<Value> {
345 let tag = self.tag_resolver.resolve(&tag_str)?;
347
348 self.tag_resolver.apply_tag(&tag, &value)
350 }
351
352 fn compose_sequence(&mut self) -> Result<Option<Value>> {
354 self.current_depth += 1;
356 self.resource_tracker
357 .check_depth(&self.limits, self.current_depth)?;
358
359 let mut sequence = Vec::new();
360
361 while self.parser.check_event() {
362 if let Ok(Some(event)) = self.parser.peek_event() {
364 if matches!(event.event_type, EventType::SequenceEnd) {
365 self.parser.get_event()?;
367 break;
368 } else if matches!(
369 event.event_type,
370 EventType::DocumentEnd { .. }
371 | EventType::DocumentStart { .. }
372 | EventType::StreamEnd
373 ) {
374 break;
376 }
377 }
378
379 if let Some(node) = self.compose_node()? {
381 self.resource_tracker.add_collection_item(&self.limits)?;
382 self.resource_tracker.add_complexity(&self.limits, 1)?;
383 sequence.push(node);
384 } else {
385 break;
387 }
388 }
389
390 self.current_depth -= 1;
391 Ok(Some(Value::Sequence(sequence)))
392 }
393
394 fn compose_mapping(&mut self) -> Result<Option<Value>> {
396 self.current_depth += 1;
398 self.resource_tracker
399 .check_depth(&self.limits, self.current_depth)?;
400
401 let mut mapping = IndexMap::new();
402
403 while self.parser.check_event() {
404 if let Ok(Some(event)) = self.parser.peek_event() {
406 if matches!(event.event_type, EventType::MappingEnd) {
407 self.parser.get_event()?;
409 break;
410 } else if matches!(
411 event.event_type,
412 EventType::DocumentEnd { .. }
413 | EventType::DocumentStart { .. }
414 | EventType::StreamEnd
415 ) {
416 break;
418 }
419 }
420
421 let Some(key) = self.compose_node()? else {
423 break;
424 };
425
426 let value = self.compose_node()?.unwrap_or(Value::Null);
428
429 if let Value::String(key_str) = &key {
431 if key_str == "<<" {
432 self.process_merge_key(&mut mapping, &value)?;
434 continue;
435 }
436 }
437
438 self.resource_tracker.add_collection_item(&self.limits)?;
439 self.resource_tracker.add_complexity(&self.limits, 2)?; mapping.insert(key, value);
441 }
442
443 self.current_depth -= 1;
444 Ok(Some(Value::Mapping(mapping)))
445 }
446
447 fn process_merge_key(
450 &self,
451 mapping: &mut IndexMap<Value, Value>,
452 merge_value: &Value,
453 ) -> Result<()> {
454 match merge_value {
455 Value::Mapping(source_map) => {
457 for (key, value) in source_map {
458 mapping.entry(key.clone()).or_insert_with(|| value.clone());
460 }
461 }
462
463 Value::Sequence(sources) => {
465 for source in sources {
466 if let Value::Mapping(source_map) = source {
467 for (key, value) in source_map {
468 mapping.entry(key.clone()).or_insert_with(|| value.clone());
470 }
471 } else {
472 return Err(Error::construction(
473 self.position,
474 "Merge key sequence can only contain mappings",
475 ));
476 }
477 }
478 }
479
480 _ => {
481 return Err(Error::construction(
482 self.position,
483 "Merge key value must be a mapping or sequence of mappings",
484 ));
485 }
486 }
487
488 Ok(())
489 }
490}
491
492impl Composer for BasicComposer {
493 fn check_document(&self) -> bool {
494 if let Ok(Some(event)) = self.parser.peek_event() {
496 !matches!(event.event_type, EventType::StreamEnd)
497 } else {
498 false
499 }
500 }
501
502 fn compose_document(&mut self) -> Result<Option<Value>> {
503 if let Some(error) = self.parser.take_scanning_error() {
505 return Err(error);
506 }
507
508 while let Ok(Some(event)) = self.parser.peek_event() {
510 if let EventType::DocumentStart { tags, version, .. } = &event.event_type {
511 self.yaml_version = version
513 .map(|(maj, min)| YamlVersion::from_directive(maj, min))
514 .unwrap_or_default();
515
516 self.tag_resolver.clear_directives();
518
519 for (handle, prefix) in tags {
521 self.tag_resolver
522 .add_directive(handle.clone(), prefix.clone());
523 }
524
525 self.parser.get_event()?; } else if matches!(event.event_type, EventType::DocumentStart { .. }) {
527 self.parser.get_event()?; } else {
529 break;
530 }
531 }
532
533 let document = self.compose_node()?;
535
536 while let Ok(Some(event)) = self.parser.peek_event() {
538 if matches!(event.event_type, EventType::DocumentEnd { .. }) {
539 self.parser.get_event()?; } else {
541 break;
542 }
543 }
544
545 Ok(document)
546 }
547
548 fn position(&self) -> Position {
549 self.position
550 }
551
552 fn reset(&mut self) {
553 self.position = Position::new();
554 self.anchors.clear();
555 self.resource_tracker.reset();
556 self.alias_expansion_stack.clear();
557 self.current_depth = 0;
558 self.tag_resolver = TagResolver::new();
559 }
560}
561
562impl Default for BasicComposer {
563 fn default() -> Self {
564 Self::new(String::new())
565 }
566}
567
568#[cfg(test)]
569mod tests {
570 use super::*;
571 use indexmap::IndexMap;
572
573 fn build_deep_sequence(depth: usize) -> Value {
577 let mut v = Value::Int(1);
578 for _ in 0..depth {
579 v = Value::Sequence(vec![v]);
580 }
581 v
582 }
583
584 fn drop_deep(mut v: Value) {
588 let mut stack: Vec<Value> = Vec::new();
589 stack.push(std::mem::replace(&mut v, Value::Null));
590 while let Some(node) = stack.pop() {
591 if let Value::Sequence(seq) = node {
592 for item in seq {
593 stack.push(item);
594 }
595 }
596 }
597 }
598
599 #[test]
600 fn iterative_complexity_handles_100k_deep_value() {
601 let deep = build_deep_sequence(100_000);
606 let complexity = calculate_value_complexity(&deep).expect("must not error");
607 assert_eq!(complexity, 200_001);
610 drop_deep(deep);
611 }
612
613 #[test]
614 fn iterative_structure_depth_handles_100k_deep_value() {
615 let deep = build_deep_sequence(100_000);
617 let depth = calculate_structure_depth(&deep);
618 assert_eq!(depth, 100_001);
620 drop_deep(deep);
621 }
622
623 #[test]
624 fn test_check_document() {
625 let mut composer = BasicComposer::new_eager("42".to_string());
626 assert!(composer.check_document());
627
628 let _document = composer.compose_document().unwrap();
629 }
631
632 #[test]
633 fn test_scalar_composition() {
634 let mut composer = BasicComposer::new_eager("42".to_string());
635 let document = composer.compose_document().unwrap().unwrap();
636 assert_eq!(document, Value::Int(42));
637 }
638
639 #[test]
640 fn test_boolean_composition() {
641 let mut composer = BasicComposer::new_eager("true".to_string());
642 let document = composer.compose_document().unwrap().unwrap();
643 assert_eq!(document, Value::Bool(true));
644 }
645
646 #[test]
647 fn test_null_composition() {
648 let mut composer = BasicComposer::new_eager("~".to_string());
649 let document = composer.compose_document().unwrap().unwrap();
650 assert_eq!(document, Value::Null);
651 }
652
653 #[test]
654 fn test_string_composition() {
655 let mut composer = BasicComposer::new_eager("hello world".to_string());
656 let document = composer.compose_document().unwrap().unwrap();
657 assert_eq!(document, Value::String("hello world".to_string()));
658 }
659
660 #[test]
661 fn test_flow_sequence_composition() {
662 let mut composer = BasicComposer::new_eager("[1, 2, 3]".to_string());
663 let document = composer.compose_document().unwrap().unwrap();
664
665 let expected = Value::Sequence(vec![Value::Int(1), Value::Int(2), Value::Int(3)]);
666 assert_eq!(document, expected);
667 }
668
669 #[test]
670 fn test_flow_mapping_composition() {
671 let mut composer = BasicComposer::new_eager("{'key': 'value', 'number': 42}".to_string());
672 let document = composer.compose_document().unwrap().unwrap();
673
674 let mut expected_map = IndexMap::new();
675 expected_map.insert(
676 Value::String("key".to_string()),
677 Value::String("value".to_string()),
678 );
679 expected_map.insert(Value::String("number".to_string()), Value::Int(42));
680 let expected = Value::Mapping(expected_map);
681
682 assert_eq!(document, expected);
683 }
684
685 #[test]
686 fn test_nested_composition() {
687 let yaml_content = "{'users': [{'name': 'Alice', 'age': 30}]}";
688 let mut composer = BasicComposer::new_eager(yaml_content.to_string());
689 let document = composer.compose_document().unwrap().unwrap();
690
691 let mut user = IndexMap::new();
693 user.insert(
694 Value::String("name".to_string()),
695 Value::String("Alice".to_string()),
696 );
697 user.insert(Value::String("age".to_string()), Value::Int(30));
698
699 let users = Value::Sequence(vec![Value::Mapping(user)]);
700
701 let mut expected = IndexMap::new();
702 expected.insert(Value::String("users".to_string()), users);
703
704 assert_eq!(document, Value::Mapping(expected));
705 }
706
707 #[test]
708 fn test_multiple_types() {
709 let yaml_content = "[42, 'hello', true, null]";
710 let mut composer = BasicComposer::new_eager(yaml_content.to_string());
711 let document = composer.compose_document().unwrap().unwrap();
712
713 let expected = Value::Sequence(vec![
714 Value::Int(42),
715 Value::String("hello".to_string()),
716 Value::Bool(true),
717 Value::Null,
718 ]);
719
720 assert_eq!(document, expected);
721 }
722
723 #[test]
724 fn test_merge_keys_simple() {
725 let yaml_content = r"
726base: &base
727 key: value
728 timeout: 30
729
730test:
731 <<: *base
732 environment: prod
733";
734 let mut composer = BasicComposer::new_eager(yaml_content.to_string());
735 let document = composer.compose_document().unwrap().unwrap();
736
737 if let Value::Mapping(ref map) = document {
738 assert!(map.contains_key(&Value::String("base".to_string())));
740
741 if let Some(Value::Mapping(test_map)) = map.get(&Value::String("test".to_string())) {
743 assert!(test_map.contains_key(&Value::String("key".to_string())));
744 assert!(test_map.contains_key(&Value::String("timeout".to_string())));
745 assert!(test_map.contains_key(&Value::String("environment".to_string())));
746
747 assert_eq!(
749 test_map.get(&Value::String("key".to_string())),
750 Some(&Value::String("value".to_string()))
751 );
752 assert_eq!(
753 test_map.get(&Value::String("timeout".to_string())),
754 Some(&Value::Int(30))
755 );
756 assert_eq!(
757 test_map.get(&Value::String("environment".to_string())),
758 Some(&Value::String("prod".to_string()))
759 );
760 } else {
761 panic!("test mapping not found or not a mapping");
762 }
763 } else {
764 panic!("Document should be a mapping, got: {:?}", document);
765 }
766 }
767
768 #[test]
769 fn test_debug_alias_tokens() {
770 let yaml_content = r"
771base: &base
772 key: value
773
774ref: *base
775";
776
777 let mut scanner = crate::BasicScanner::new_eager(yaml_content.to_string());
778
779 println!("Scanning tokens for alias test:");
780 let mut token_count = 0;
781
782 while scanner.check_token() {
783 if let Ok(Some(token)) = scanner.get_token() {
784 token_count += 1;
785 println!(
786 "{}: {:?} at {:?}-{:?}",
787 token_count, token.token_type, token.start_position, token.end_position
788 );
789 } else {
790 println!("No more tokens");
791 break;
792 }
793 }
794
795 println!("Total tokens: {}", token_count);
796 }
797
798 #[test]
799 fn test_debug_alias_events() {
800 let yaml_content = r"
801base: &base
802 key: value
803
804ref: *base
805";
806
807 let mut parser = BasicParser::new_eager(yaml_content.to_string());
808
809 println!("Parsing events for alias test:");
810 let mut event_count = 0;
811
812 while parser.check_event() {
813 if let Ok(Some(event)) = parser.get_event() {
814 event_count += 1;
815 println!(
816 "{}: {:?} at {:?}",
817 event_count, event.event_type, event.position
818 );
819 } else {
820 println!("No more events");
821 break;
822 }
823 }
824
825 println!("Total events: {}", event_count);
826 }
827
828 #[test]
829 fn test_simple_scalar_alias_resolution() {
830 let yaml_content = r"
832base: &base 'hello world'
833ref: *base
834";
835 let mut composer = BasicComposer::new_eager(yaml_content.to_string());
836 let document = composer.compose_document().unwrap().unwrap();
837
838 println!("Simple alias document: {:?}", document);
839
840 if let Value::Mapping(ref map) = document {
841 println!("Mapping keys: {:?}", map.keys().collect::<Vec<_>>());
842
843 let base_value = map
844 .get(&Value::String("base".to_string()))
845 .expect("base should exist");
846 let ref_value = map
847 .get(&Value::String("ref".to_string()))
848 .expect("ref should exist");
849
850 println!("base_value: {:?}", base_value);
851 println!("ref_value: {:?}", ref_value);
852
853 assert_eq!(base_value, ref_value);
854 } else {
855 panic!("Document should be a mapping, got: {:?}", document);
856 }
857 }
858
859 #[test]
860 fn test_basic_alias_resolution() {
861 let yaml_content = r"
862base: &base
863 key: value
864
865ref: *base
866";
867 let mut composer = BasicComposer::new_eager(yaml_content.to_string());
868 let document = composer.compose_document().unwrap().unwrap();
869
870 println!("Composed document: {:?}", document);
871
872 if let Value::Mapping(ref map) = document {
873 println!("Mapping keys: {:?}", map.keys().collect::<Vec<_>>());
874
875 let base_value = map
877 .get(&Value::String("base".to_string()))
878 .expect("base should exist");
879 let ref_value = map
880 .get(&Value::String("ref".to_string()))
881 .expect("ref should exist");
882
883 println!("base_value: {:?}", base_value);
884 println!("ref_value: {:?}", ref_value);
885
886 assert_eq!(
888 base_value, ref_value,
889 "Alias should resolve to the same value as the anchor"
890 );
891
892 if let Value::Mapping(nested) = base_value {
894 assert_eq!(
895 nested.get(&Value::String("key".to_string())),
896 Some(&Value::String("value".to_string()))
897 );
898 } else {
899 panic!("base value should be a nested mapping");
900 }
901
902 println!("✅ Alias resolution working correctly!");
903 } else {
904 panic!("Document should be a mapping, got: {:?}", document);
905 }
906 }
907}