1use std::collections::{BTreeMap, BTreeSet};
2use std::fmt::{self, Write as _};
3
4use crate::{NodeId, SemanticKind, YamlDoc, YamlError, YamlScalarStyle, strip_inline_comment};
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 let source = strip_inline_comment(&source).trim_end();
209 if source.contains(['\n', '\r'])
210 || style == YamlScalarStyle::Plain
211 && source.contains(['[', ']', '{', '}', ','])
212 {
213 let value = self.doc.scalar_value(node).map_err(FragmentError::from)?;
214 output.push_str(&prefix);
215 output.push_str("e_string(&value));
216 } else {
217 output.push_str(source);
218 }
219 }
220 }
221 Some(SemanticKind::Sequence { .. }) => {
222 output.push_str(&self.property_prefix(node));
223 output.push('[');
224 pending.push(RenderAction::Text("]"));
225 let items = self.doc.sequence_items(node).collect::<Vec<_>>();
226 for (index, item) in items.into_iter().enumerate().rev() {
227 pending.push(RenderAction::Node(item, depth + 1));
228 if index > 0 {
229 pending.push(RenderAction::Text(", "));
230 }
231 }
232 }
233 Some(SemanticKind::Mapping { .. }) => {
234 output.push_str(&self.property_prefix(node));
235 output.push('{');
236 pending.push(RenderAction::Text("}"));
237 let entries = self.doc.mapping_entries(node).collect::<Vec<_>>();
238 for (index, (key, value)) in entries.into_iter().enumerate().rev() {
239 pending.push(RenderAction::Node(value, depth + 1));
240 pending.push(RenderAction::Text(": "));
241 pending.push(RenderAction::Node(key, depth + 1));
242 if index > 0 {
243 pending.push(RenderAction::Text(", "));
244 }
245 }
246 }
247 Some(SemanticKind::Document) | None => {
248 return Err(FragmentError::new("cannot render unknown YAML node"));
249 }
250 }
251 }
252 Ok(output)
253 }
254
255 fn property_prefix(&self, node: NodeId) -> String {
256 let mut prefix = String::new();
257 if let Some(tag) = self.doc.raw_tag(node) {
258 prefix.push_str(tag);
259 prefix.push(' ');
260 }
261 if let Some(anchor) = self.doc.anchor(node) {
262 prefix.push('&');
263 prefix.push_str(anchor);
264 prefix.push(' ');
265 }
266 prefix
267 }
268
269 fn subtree_nodes(&self) -> impl Iterator<Item = NodeId> + '_ {
270 let span = self.doc.node(self.root).map(super::syntax::Node::span);
271 self.doc
272 .nodes
273 .iter()
274 .enumerate()
275 .map(|(index, _)| NodeId::from_usize(index))
276 .filter(move |node| {
277 let Some(root_span) = span else {
278 return false;
279 };
280 self.doc.node(*node).is_some_and(|node| {
281 node.span().start >= root_span.start && node.span().end <= root_span.end
282 })
283 })
284 }
285
286 fn validate_alias_scope(&self) -> Result<(), FragmentError> {
287 let root_span = self
288 .doc
289 .node(self.root)
290 .map(super::syntax::Node::span)
291 .ok_or_else(|| FragmentError::new("fragment root node is missing"))?;
292 for node in self.subtree_nodes() {
293 if !matches!(self.doc.semantic_kind(node), Some(SemanticKind::Alias)) {
294 continue;
295 }
296 let target = self.doc.resolve_alias(node).ok_or_else(|| {
297 FragmentError::new(format!(
298 "unresolved value alias `*{}`",
299 self.doc.alias_name(node).unwrap_or_default()
300 ))
301 })?;
302 let target_span = self
303 .doc
304 .node(target)
305 .map(super::syntax::Node::span)
306 .ok_or_else(|| FragmentError::new("alias target is missing"))?;
307 if target_span.start < root_span.start || target_span.end > root_span.end {
308 return Err(FragmentError::new(
309 "a value alias cannot reference a node outside the value root",
310 ));
311 }
312 }
313 Ok(())
314 }
315}
316
317impl YamlDoc {
318 pub fn extract_node(&self, node: NodeId) -> Result<String, YamlError> {
324 let node = self.expect_node(node)?;
325 let source = self.source.slice(node.span);
326 let line_start = self.source.as_str()[..node.span.start as usize]
327 .rfind(['\n', '\r'])
328 .map_or(0, |index| index + 1);
329 let base_indent = self.source.as_str()[line_start..node.span.start as usize]
330 .bytes()
331 .take_while(|byte| *byte == b' ')
332 .count();
333 Ok(deindent_continuation_lines(source, base_indent))
334 }
335
336 pub(crate) fn anchor_names(&self) -> BTreeSet<String> {
337 self.nodes
338 .iter()
339 .enumerate()
340 .filter_map(|(index, _)| self.anchor(NodeId::from_usize(index)).map(str::to_owned))
341 .collect()
342 }
343}
344
345fn deindent_continuation_lines(source: &str, indent: usize) -> String {
346 if indent == 0 || !source.contains(['\n', '\r']) {
347 return source.to_owned();
348 }
349 let bytes = source.as_bytes();
350 let mut output = String::with_capacity(source.len());
351 let mut position = 0;
352 let mut first = true;
353 while position < bytes.len() {
354 if !first {
355 let mut removed = 0;
356 while removed < indent && bytes.get(position) == Some(&b' ') {
357 position += 1;
358 removed += 1;
359 }
360 }
361 first = false;
362 let line_end = source[position..]
363 .find(['\n', '\r'])
364 .map_or(source.len(), |offset| position + offset);
365 output.push_str(&source[position..line_end]);
366 position = line_end;
367 if bytes.get(position) == Some(&b'\r') {
368 output.push('\r');
369 position += 1;
370 if bytes.get(position) == Some(&b'\n') {
371 output.push('\n');
372 position += 1;
373 }
374 } else if bytes.get(position) == Some(&b'\n') {
375 output.push('\n');
376 position += 1;
377 }
378 }
379 output
380}
381
382pub(crate) fn indent_text(source: &str, indent: usize) -> String {
383 if indent == 0 || source.is_empty() {
384 return source.to_owned();
385 }
386 let prefix = " ".repeat(indent);
387 let mut output = String::with_capacity(source.len() + prefix.len());
388 let mut at_line_start = true;
389 for character in source.chars() {
390 if at_line_start && !matches!(character, '\r' | '\n') {
391 output.push_str(&prefix);
392 at_line_start = false;
393 }
394 output.push(character);
395 if character == '\n' || character == '\r' {
396 at_line_start = true;
397 }
398 }
399 output
400}
401
402pub(crate) fn quote_string(value: &str) -> String {
403 let mut output = String::from("\"");
404 for character in value.chars() {
405 match character {
406 '"' => output.push_str("\\\""),
407 '\\' => output.push_str("\\\\"),
408 '\n' => output.push_str("\\n"),
409 '\r' => output.push_str("\\r"),
410 '\t' => output.push_str("\\t"),
411 character if character.is_control() => {
412 write!(output, "\\u{:04X}", u32::from(character))
413 .expect("writing to a String cannot fail");
414 }
415 character => output.push(character),
416 }
417 }
418 output.push('"');
419 output
420}
421
422#[derive(Debug, Clone, PartialEq, Eq)]
424pub struct FragmentError {
425 message: String,
426}
427
428impl FragmentError {
429 pub(crate) fn new(message: impl Into<String>) -> Self {
430 Self {
431 message: message.into(),
432 }
433 }
434}
435
436impl fmt::Display for FragmentError {
437 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
438 formatter.write_str(&self.message)
439 }
440}
441
442impl std::error::Error for FragmentError {}
443
444impl From<YamlError> for FragmentError {
445 fn from(error: YamlError) -> Self {
446 Self::new(error.to_string())
447 }
448}
449
450#[cfg(test)]
451mod tests {
452 use super::*;
453
454 fn nested_block_sequence(depth: usize) -> String {
455 let mut yaml = String::new();
456 for level in 0..depth {
457 yaml.push_str(&" ".repeat(level));
458 yaml.push_str("-\n");
459 }
460 yaml.push_str(&" ".repeat(depth));
461 yaml.push_str("value\n");
462 yaml
463 }
464
465 #[test]
466 fn fragment_requires_one_nonempty_document() {
467 assert!(YamlFragment::parse("").is_err());
468 assert!(YamlFragment::parse("--- a\n--- b\n").is_err());
469 assert!(YamlFragment::parse("[a, b]").is_ok());
470 }
471
472 #[test]
473 fn extraction_deindents_nested_block_nodes() {
474 let doc = YamlDoc::parse("outer:\n one: 1\n two:\n - a\n - b\n").unwrap();
475 let node = doc
476 .get_mapping_value(doc.document_root_mapping(0).unwrap(), "outer")
477 .unwrap()
478 .unwrap();
479 assert_eq!(
480 doc.extract_node(node).unwrap(),
481 "one: 1\ntwo:\n - a\n - b"
482 );
483 }
484
485 #[test]
486 fn rejects_unresolved_value_aliases() {
487 assert!(YamlFragment::parse("*outside").is_err());
488 }
489
490 #[test]
491 fn flow_rendering_normalizes_block_collections_only() {
492 let fragment = YamlFragment::parse("one:\n - a\n - b\n").unwrap();
493 let target = YamlDoc::parse("target: []\n").unwrap();
494 assert_eq!(fragment.render_flow(&target).unwrap(), "{one: [a, b]}");
495 }
496
497 #[test]
498 fn flow_rendering_preserves_the_existing_depth_limit() {
499 std::thread::Builder::new()
500 .stack_size(32 * 1024 * 1024)
501 .spawn(|| {
502 let target = YamlDoc::parse("target: []\n").unwrap();
503 let accepted = YamlFragment::parse(&nested_block_sequence(1024)).unwrap();
504 assert!(accepted.render_flow(&target).is_ok());
505
506 let rejected = YamlFragment::parse(&nested_block_sequence(1025)).unwrap();
507 assert!(
508 rejected
509 .render_flow(&target)
510 .unwrap_err()
511 .to_string()
512 .contains("recursion limit")
513 );
514 })
515 .unwrap()
516 .join()
517 .unwrap();
518 }
519}