dbt_yaml/
loader.rs

1use crate::de::{Event, Progress};
2use crate::error::{self, Error, ErrorImpl, Result};
3use crate::libyaml::error::Mark;
4use crate::libyaml::parser::{Event as YamlEvent, Parser};
5use crate::spanned;
6use std::borrow::Cow;
7use std::collections::BTreeMap;
8use std::sync::Arc;
9
10pub(crate) struct Loader<'input> {
11    parser: Option<Parser<'input>>,
12    document_count: usize,
13}
14
15pub(crate) struct Document<'input> {
16    pub events: Vec<(Event<'input>, Mark)>,
17    pub error: Option<Arc<ErrorImpl>>,
18    /// Map from alias id to index in events.
19    pub aliases: BTreeMap<usize, usize>,
20}
21
22impl<'input> Loader<'input> {
23    pub fn new(progress: Progress<'input>) -> Result<Self> {
24        let input = match progress {
25            Progress::Str(s) => Cow::Borrowed(s.as_bytes()),
26            Progress::Slice(bytes) => Cow::Borrowed(bytes),
27            Progress::Read(mut rdr) => {
28                let mut buffer = Vec::new();
29                if let Err(io_error) = rdr.read_to_end(&mut buffer) {
30                    return Err(error::new(ErrorImpl::Io(io_error)));
31                }
32                Cow::Owned(buffer)
33            }
34            Progress::Iterable(_) | Progress::Document(_) => unreachable!(),
35            Progress::Fail(err) => return Err(error::shared(err)),
36        };
37
38        Ok(Loader {
39            parser: Some(Parser::new(input)),
40            document_count: 0,
41        })
42    }
43
44    pub fn next_document(&mut self) -> Option<Document<'input>> {
45        let document = self.next_document_inner()?;
46        if let Some((_event, mark)) = document.events.first() {
47            spanned::set_marker(*mark);
48        }
49
50        Some(document)
51    }
52
53    fn next_document_inner(&mut self) -> Option<Document<'input>> {
54        let parser = match &mut self.parser {
55            Some(parser) => parser,
56            None => return None,
57        };
58
59        let first = self.document_count == 0;
60        self.document_count += 1;
61
62        let mut anchors = BTreeMap::new();
63        let mut document = Document {
64            events: Vec::new(),
65            error: None,
66            aliases: BTreeMap::new(),
67        };
68
69        loop {
70            let (event, mark) = match parser.next() {
71                Ok((event, mark)) => (event, mark),
72                Err(err) => {
73                    document.error = Some(Error::from(err).shared());
74                    return Some(document);
75                }
76            };
77            let event = match event {
78                YamlEvent::StreamStart => continue,
79                YamlEvent::StreamEnd => {
80                    self.parser = None;
81                    return if first {
82                        if document.events.is_empty() {
83                            document.events.push((Event::Void, mark));
84                        }
85                        Some(document)
86                    } else {
87                        None
88                    };
89                }
90                YamlEvent::DocumentStart => continue,
91                YamlEvent::DocumentEnd => {
92                    document.events.push((Event::Void, mark));
93                    return Some(document);
94                }
95                YamlEvent::Alias(alias) => match anchors.get(&alias) {
96                    Some(id) => Event::Alias(*id),
97                    None => {
98                        document.error =
99                            Some(error::new(ErrorImpl::UnknownAnchor(mark.into())).shared());
100                        return Some(document);
101                    }
102                },
103                YamlEvent::Scalar(mut scalar) => {
104                    if let Some(anchor) = scalar.anchor.take() {
105                        let id = anchors.len();
106                        anchors.insert(anchor, id);
107                        document.aliases.insert(id, document.events.len());
108                    }
109                    Event::Scalar(scalar)
110                }
111                YamlEvent::SequenceStart(mut sequence_start) => {
112                    if let Some(anchor) = sequence_start.anchor.take() {
113                        let id = anchors.len();
114                        anchors.insert(anchor, id);
115                        document.aliases.insert(id, document.events.len());
116                    }
117                    Event::SequenceStart(sequence_start)
118                }
119                YamlEvent::SequenceEnd => Event::SequenceEnd,
120                YamlEvent::MappingStart(mut mapping_start) => {
121                    if let Some(anchor) = mapping_start.anchor.take() {
122                        let id = anchors.len();
123                        anchors.insert(anchor, id);
124                        document.aliases.insert(id, document.events.len());
125                    }
126                    Event::MappingStart(mapping_start)
127                }
128                YamlEvent::MappingEnd => Event::MappingEnd,
129            };
130            document.events.push((event, mark));
131        }
132    }
133}
dbt_yaml/loader.rs

dbt_yaml/
loader.rs
