Skip to main content

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<_, (usize, Mark)> = 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                        if let Some((_, first_mark)) = anchors.get(&anchor) {
106                            document.error = Some(
107                                error::new(ErrorImpl::DuplicateAnchor {
108                                    first: (*first_mark).into(),
109                                    second: mark.into(),
110                                })
111                                .shared(),
112                            );
113                            return Some(document);
114                        }
115                        let id = anchors.len();
116                        anchors.insert(anchor, (id, mark));
117                        document.aliases.insert(id, document.events.len());
118                    }
119                    Event::Scalar(scalar)
120                }
121                YamlEvent::SequenceStart(mut sequence_start) => {
122                    if let Some(anchor) = sequence_start.anchor.take() {
123                        if let Some((_, first_mark)) = anchors.get(&anchor) {
124                            document.error = Some(
125                                error::new(ErrorImpl::DuplicateAnchor {
126                                    first: (*first_mark).into(),
127                                    second: mark.into(),
128                                })
129                                .shared(),
130                            );
131                            return Some(document);
132                        }
133                        let id = anchors.len();
134                        anchors.insert(anchor, (id, mark));
135                        document.aliases.insert(id, document.events.len());
136                    }
137                    Event::SequenceStart(sequence_start)
138                }
139                YamlEvent::SequenceEnd => Event::SequenceEnd,
140                YamlEvent::MappingStart(mut mapping_start) => {
141                    if let Some(anchor) = mapping_start.anchor.take() {
142                        if let Some((_, first_mark)) = anchors.get(&anchor) {
143                            document.error = Some(
144                                error::new(ErrorImpl::DuplicateAnchor {
145                                    first: (*first_mark).into(),
146                                    second: mark.into(),
147                                })
148                                .shared(),
149                            );
150                            return Some(document);
151                        }
152                        let id = anchors.len();
153                        anchors.insert(anchor, (id, mark));
154                        document.aliases.insert(id, document.events.len());
155                    }
156                    Event::MappingStart(mapping_start)
157                }
158                YamlEvent::MappingEnd => Event::MappingEnd,
159            };
160            document.events.push((event, mark));
161        }
162    }
163}