1use std::collections::{BTreeMap, BTreeSet};
2use std::fmt::{self, Write as _};
3
4use crate::{NodeId, SemanticKind, YamlDoc, YamlError, YamlScalarStyle};
5
6#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct YamlFragment {
9 doc: YamlDoc,
10 root: NodeId,
11}
12
13impl YamlFragment {
14 pub fn parse_owned(input: String) -> Result<Self, FragmentError> {
21 let doc = YamlDoc::parse_owned(input).map_err(FragmentError::from)?;
22 if doc.document_count() != 1 {
23 return Err(FragmentError::new(format!(
24 "a YAML value must contain exactly one document, found {}",
25 doc.document_count()
26 )));
27 }
28 let root = doc
29 .document_root(0)
30 .map_err(FragmentError::from)?
31 .ok_or_else(|| FragmentError::new("a YAML value must contain one root node"))?;
32 let fragment = Self { doc, root };
33 fragment.validate_alias_scope()?;
34 Ok(fragment)
35 }
36
37 pub fn parse(input: &str) -> Result<Self, FragmentError> {
43 Self::parse_owned(input.to_owned())
44 }
45
46 #[must_use]
48 pub fn document(&self) -> &YamlDoc {
49 &self.doc
50 }
51
52 #[must_use]
54 pub const fn root(&self) -> NodeId {
55 self.root
56 }
57
58 pub fn to_yaml(&self) -> Result<String, FragmentError> {
64 self.doc
65 .extract_node(self.root)
66 .map_err(FragmentError::from)
67 }
68
69 pub(crate) fn contains_anchor(&self) -> bool {
70 self.subtree_nodes()
71 .any(|node| self.doc.anchor(node).is_some())
72 }
73
74 pub(crate) fn from_document_node(doc: &YamlDoc, root: NodeId) -> Result<Self, FragmentError> {
75 doc.node(root)
76 .ok_or_else(|| FragmentError::new("fragment source node is missing"))?;
77 Ok(Self {
78 doc: doc.clone(),
79 root,
80 })
81 }
82
83 pub(crate) fn prepared(&self, target: &YamlDoc) -> Result<Self, FragmentError> {
84 self.prepared_with_anchor_exemption(target, None)
85 }
86
87 pub(crate) fn prepared_for_replacement(
88 &self,
89 target: &YamlDoc,
90 replaced: NodeId,
91 ) -> Result<Self, FragmentError> {
92 self.prepared_with_anchor_exemption(target, target.anchor(replaced))
93 }
94
95 fn prepared_with_anchor_exemption(
96 &self,
97 target: &YamlDoc,
98 exempt_anchor: Option<&str>,
99 ) -> Result<Self, FragmentError> {
100 let mut used = target.anchor_names();
101 if let Some(anchor) = exempt_anchor {
102 used.remove(anchor);
103 }
104 let mut renamed = BTreeMap::new();
105 for node in self.subtree_nodes() {
106 let Some(name) = self.doc.anchor(node) else {
107 continue;
108 };
109 if used.insert(name.to_owned()) {
110 continue;
111 }
112 let mut suffix = 1_u64;
113 let replacement = loop {
114 let candidate = format!("{name}_{suffix}");
115 if used.insert(candidate.clone()) {
116 break candidate;
117 }
118 suffix = suffix.saturating_add(1);
119 };
120 renamed.insert(name.to_owned(), replacement);
121 }
122 if renamed.is_empty() {
123 return Ok(self.clone());
124 }
125
126 let mut doc = self.doc.clone();
127 let nodes = self.subtree_nodes().collect::<Vec<_>>();
128 for node in nodes {
129 let Some(properties) = doc.semantic_properties(node) else {
130 continue;
131 };
132 if let Some(span) = properties.anchor {
133 let old = doc.source.slice(span);
134 if let Some(new) = renamed.get(old) {
135 doc.queue_edit(span, new.clone())
136 .map_err(FragmentError::from)?;
137 }
138 }
139 if let Some(span) = properties.alias {
140 let old = doc.source.slice(span);
141 if let Some(new) = renamed.get(old) {
142 doc.queue_edit(span, new.clone())
143 .map_err(FragmentError::from)?;
144 }
145 }
146 }
147 doc.commit_edits().map_err(FragmentError::from)?;
148 let root = doc
149 .document_root(0)
150 .map_err(FragmentError::from)?
151 .ok_or_else(|| FragmentError::new("prepared fragment lost its root node"))?;
152 Ok(Self { doc, root })
153 }
154
155 pub(crate) fn render_flow(&self, target: &YamlDoc) -> Result<String, FragmentError> {
156 let prepared = self.prepared(target)?;
157 prepared.render_node_flow(prepared.root, 0)
158 }
159
160 pub(crate) fn render_flow_for_replacement(
161 &self,
162 target: &YamlDoc,
163 replaced: NodeId,
164 ) -> Result<String, FragmentError> {
165 let prepared = self.prepared_for_replacement(target, replaced)?;
166 prepared.render_node_flow(prepared.root, 0)
167 }
168
169 fn render_node_flow(&self, node: NodeId, depth: usize) -> Result<String, FragmentError> {
170 enum RenderAction {
171 Node(NodeId, usize),
172 Text(&'static str),
173 }
174
175 let mut output = String::new();
176 let mut pending = vec![RenderAction::Node(node, depth)];
177 while let Some(action) = pending.pop() {
178 let RenderAction::Node(node, depth) = action else {
179 let RenderAction::Text(text) = action else {
180 unreachable!();
181 };
182 output.push_str(text);
183 continue;
184 };
185 if depth > 1024 {
186 return Err(FragmentError::new(
187 "fragment rendering recursion limit exceeded",
188 ));
189 }
190 match self.doc.semantic_kind(node) {
191 Some(SemanticKind::Alias) => {
192 output.push('*');
193 output.push_str(self.doc.alias_name(node).unwrap_or_default());
194 }
195 Some(SemanticKind::Scalar { style }) => {
196 let prefix = self.property_prefix(node);
197 if matches!(style, YamlScalarStyle::Literal | YamlScalarStyle::Folded)
198 || self
199 .doc
200 .node(node)
201 .is_some_and(|node| node.span().is_empty())
202 {
203 let value = self.doc.scalar_value(node).map_err(FragmentError::from)?;
204 output.push_str(&prefix);
205 output.push_str("e_string(&value));
206 } else {
207 let source = self.doc.extract_node(node).map_err(FragmentError::from)?;
208 if source.contains(['\n', '\r'])
209 || style == YamlScalarStyle::Plain
210 && source.contains(['[', ']', '{', '}', ','])
211 {
212 let value = self.doc.scalar_value(node).map_err(FragmentError::from)?;
213 output.push_str(&prefix);
214 output.push_str("e_string(&value));
215 } else {
216 output.push_str(&source);
217 }
218 }
219 }
220 Some(SemanticKind::Sequence { .. }) => {
221 output.push_str(&self.property_prefix(node));
222 output.push('[');
223 pending.push(RenderAction::Text("]"));
224 let items = self.doc.sequence_items(node).collect::<Vec<_>>();
225 for (index, item) in items.into_iter().enumerate().rev() {
226 pending.push(RenderAction::Node(item, depth + 1));
227 if index > 0 {
228 pending.push(RenderAction::Text(", "));
229 }
230 }
231 }
232 Some(SemanticKind::Mapping { .. }) => {
233 output.push_str(&self.property_prefix(node));
234 output.push('{');
235 pending.push(RenderAction::Text("}"));
236 let entries = self.doc.mapping_entries(node).collect::<Vec<_>>();
237 for (index, (key, value)) in entries.into_iter().enumerate().rev() {
238 pending.push(RenderAction::Node(value, depth + 1));
239 pending.push(RenderAction::Text(": "));
240 pending.push(RenderAction::Node(key, depth + 1));
241 if index > 0 {
242 pending.push(RenderAction::Text(", "));
243 }
244 }
245 }
246 Some(SemanticKind::Document) | None => {
247 return Err(FragmentError::new("cannot render unknown YAML node"));
248 }
249 }
250 }
251 Ok(output)
252 }
253
254 fn property_prefix(&self, node: NodeId) -> String {
255 let mut prefix = String::new();
256 if let Some(tag) = self.doc.raw_tag(node) {
257 prefix.push_str(tag);
258 prefix.push(' ');
259 }
260 if let Some(anchor) = self.doc.anchor(node) {
261 prefix.push('&');
262 prefix.push_str(anchor);
263 prefix.push(' ');
264 }
265 prefix
266 }
267
268 fn subtree_nodes(&self) -> impl Iterator<Item = NodeId> + '_ {
269 let span = self.doc.node(self.root).map(super::syntax::Node::span);
270 self.doc
271 .nodes
272 .iter()
273 .enumerate()
274 .map(|(index, _)| NodeId::from_usize(index))
275 .filter(move |node| {
276 let Some(root_span) = span else {
277 return false;
278 };
279 self.doc.node(*node).is_some_and(|node| {
280 node.span().start >= root_span.start && node.span().end <= root_span.end
281 })
282 })
283 }
284
285 fn validate_alias_scope(&self) -> Result<(), FragmentError> {
286 let root_span = self
287 .doc
288 .node(self.root)
289 .map(super::syntax::Node::span)
290 .ok_or_else(|| FragmentError::new("fragment root node is missing"))?;
291 for node in self.subtree_nodes() {
292 if !matches!(self.doc.semantic_kind(node), Some(SemanticKind::Alias)) {
293 continue;
294 }
295 let target = self.doc.resolve_alias(node).ok_or_else(|| {
296 FragmentError::new(format!(
297 "unresolved value alias `*{}`",
298 self.doc.alias_name(node).unwrap_or_default()
299 ))
300 })?;
301 let target_span = self
302 .doc
303 .node(target)
304 .map(super::syntax::Node::span)
305 .ok_or_else(|| FragmentError::new("alias target is missing"))?;
306 if target_span.start < root_span.start || target_span.end > root_span.end {
307 return Err(FragmentError::new(
308 "a value alias cannot reference a node outside the value root",
309 ));
310 }
311 }
312 Ok(())
313 }
314}
315
316impl YamlDoc {
317 pub fn extract_node(&self, node: NodeId) -> Result<String, YamlError> {
323 let node = self.expect_node(node)?;
324 let source = self.source.slice(node.span);
325 let line_start = self.source.as_str()[..node.span.start as usize]
326 .rfind(['\n', '\r'])
327 .map_or(0, |index| index + 1);
328 let base_indent = self.source.as_str()[line_start..node.span.start as usize]
329 .bytes()
330 .take_while(|byte| *byte == b' ')
331 .count();
332 Ok(deindent_continuation_lines(source, base_indent))
333 }
334
335 pub(crate) fn anchor_names(&self) -> BTreeSet<String> {
336 self.nodes
337 .iter()
338 .enumerate()
339 .filter_map(|(index, _)| self.anchor(NodeId::from_usize(index)).map(str::to_owned))
340 .collect()
341 }
342}
343
344fn deindent_continuation_lines(source: &str, indent: usize) -> String {
345 if indent == 0 || !source.contains(['\n', '\r']) {
346 return source.to_owned();
347 }
348 let bytes = source.as_bytes();
349 let mut output = String::with_capacity(source.len());
350 let mut position = 0;
351 let mut first = true;
352 while position < bytes.len() {
353 if !first {
354 let mut removed = 0;
355 while removed < indent && bytes.get(position) == Some(&b' ') {
356 position += 1;
357 removed += 1;
358 }
359 }
360 first = false;
361 let line_end = source[position..]
362 .find(['\n', '\r'])
363 .map_or(source.len(), |offset| position + offset);
364 output.push_str(&source[position..line_end]);
365 position = line_end;
366 if bytes.get(position) == Some(&b'\r') {
367 output.push('\r');
368 position += 1;
369 if bytes.get(position) == Some(&b'\n') {
370 output.push('\n');
371 position += 1;
372 }
373 } else if bytes.get(position) == Some(&b'\n') {
374 output.push('\n');
375 position += 1;
376 }
377 }
378 output
379}
380
381pub(crate) fn indent_text(source: &str, indent: usize) -> String {
382 if indent == 0 || source.is_empty() {
383 return source.to_owned();
384 }
385 let prefix = " ".repeat(indent);
386 let mut output = String::with_capacity(source.len() + prefix.len());
387 let mut at_line_start = true;
388 for character in source.chars() {
389 if at_line_start && !matches!(character, '\r' | '\n') {
390 output.push_str(&prefix);
391 at_line_start = false;
392 }
393 output.push(character);
394 if character == '\n' || character == '\r' {
395 at_line_start = true;
396 }
397 }
398 output
399}
400
401pub(crate) fn quote_string(value: &str) -> String {
402 let mut output = String::from("\"");
403 for character in value.chars() {
404 match character {
405 '"' => output.push_str("\\\""),
406 '\\' => output.push_str("\\\\"),
407 '\n' => output.push_str("\\n"),
408 '\r' => output.push_str("\\r"),
409 '\t' => output.push_str("\\t"),
410 character if character.is_control() => {
411 write!(output, "\\u{:04X}", u32::from(character))
412 .expect("writing to a String cannot fail");
413 }
414 character => output.push(character),
415 }
416 }
417 output.push('"');
418 output
419}
420
421#[derive(Debug, Clone, PartialEq, Eq)]
423pub struct FragmentError {
424 message: String,
425}
426
427impl FragmentError {
428 pub(crate) fn new(message: impl Into<String>) -> Self {
429 Self {
430 message: message.into(),
431 }
432 }
433}
434
435impl fmt::Display for FragmentError {
436 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
437 formatter.write_str(&self.message)
438 }
439}
440
441impl std::error::Error for FragmentError {}
442
443impl From<YamlError> for FragmentError {
444 fn from(error: YamlError) -> Self {
445 Self::new(error.to_string())
446 }
447}
448
449#[cfg(test)]
450mod tests {
451 use super::*;
452
453 fn nested_block_sequence(depth: usize) -> String {
454 let mut yaml = String::new();
455 for level in 0..depth {
456 yaml.push_str(&" ".repeat(level));
457 yaml.push_str("-\n");
458 }
459 yaml.push_str(&" ".repeat(depth));
460 yaml.push_str("value\n");
461 yaml
462 }
463
464 #[test]
465 fn fragment_requires_one_nonempty_document() {
466 assert!(YamlFragment::parse("").is_err());
467 assert!(YamlFragment::parse("--- a\n--- b\n").is_err());
468 assert!(YamlFragment::parse("[a, b]").is_ok());
469 }
470
471 #[test]
472 fn extraction_deindents_nested_block_nodes() {
473 let doc = YamlDoc::parse("outer:\n one: 1\n two:\n - a\n - b\n").unwrap();
474 let node = doc
475 .get_mapping_value(doc.document_root_mapping(0).unwrap(), "outer")
476 .unwrap()
477 .unwrap();
478 assert_eq!(
479 doc.extract_node(node).unwrap(),
480 "one: 1\ntwo:\n - a\n - b"
481 );
482 }
483
484 #[test]
485 fn rejects_unresolved_value_aliases() {
486 assert!(YamlFragment::parse("*outside").is_err());
487 }
488
489 #[test]
490 fn flow_rendering_normalizes_block_collections_only() {
491 let fragment = YamlFragment::parse("one:\n - a\n - b\n").unwrap();
492 let target = YamlDoc::parse("target: []\n").unwrap();
493 assert_eq!(fragment.render_flow(&target).unwrap(), "{one: [a, b]}");
494 }
495
496 #[test]
497 fn flow_rendering_preserves_the_existing_depth_limit() {
498 std::thread::Builder::new()
499 .stack_size(32 * 1024 * 1024)
500 .spawn(|| {
501 let target = YamlDoc::parse("target: []\n").unwrap();
502 let accepted = YamlFragment::parse(&nested_block_sequence(1024)).unwrap();
503 assert!(accepted.render_flow(&target).is_ok());
504
505 let rejected = YamlFragment::parse(&nested_block_sequence(1025)).unwrap();
506 assert!(
507 rejected
508 .render_flow(&target)
509 .unwrap_err()
510 .to_string()
511 .contains("recursion limit")
512 );
513 })
514 .unwrap()
515 .join()
516 .unwrap();
517 }
518}