#![allow(clippy::module_name_repetitions)]
use std::borrow::Cow;
use std::ops::ControlFlow;
use std::{collections::BTreeMap, convert::TryFrom, mem, ops::Index, ops::IndexMut};
#[cfg(feature = "encoding")]
use encoding_rs::{Decoder, DecoderResult, Encoding};
use hashlink::LinkedHashMap;
use crate::parser::{Event, MarkedEventReceiver, Parser, Tag};
use crate::scanner::{Marker, ScanError, TScalarStyle};
#[derive(Clone, PartialEq, PartialOrd, Debug, Eq, Ord, Hash)]
pub enum Yaml {
Real(String),
Integer(i64),
String(String),
Boolean(bool),
Array(Array),
Hash(Hash),
Alias(usize),
Null,
BadValue,
}
pub type Array = Vec<Yaml>;
pub type Hash = LinkedHashMap<Yaml, Yaml>;
fn parse_f64(v: &str) -> Option<f64> {
match v {
".inf" | ".Inf" | ".INF" | "+.inf" | "+.Inf" | "+.INF" => Some(f64::INFINITY),
"-.inf" | "-.Inf" | "-.INF" => Some(f64::NEG_INFINITY),
".nan" | ".NaN" | ".NAN" => Some(f64::NAN),
_ if v.as_bytes().iter().any(u8::is_ascii_digit) => v.parse::<f64>().ok(),
_ => None,
}
}
#[derive(Default)]
pub struct YamlLoader {
docs: Vec<Yaml>,
doc_stack: Vec<(Yaml, usize)>,
key_stack: Vec<Yaml>,
anchor_map: BTreeMap<usize, Yaml>,
error: Option<ScanError>,
}
impl MarkedEventReceiver for YamlLoader {
fn on_event(&mut self, ev: Event, mark: Marker) {
if self.error.is_some() {
return;
}
if let Err(e) = self.on_event_impl(ev, mark) {
self.error = Some(e);
}
}
}
#[derive(Debug)]
pub enum LoadError {
IO(std::io::Error),
Scan(ScanError),
Decode(std::borrow::Cow<'static, str>),
}
impl From<std::io::Error> for LoadError {
fn from(error: std::io::Error) -> Self {
LoadError::IO(error)
}
}
impl YamlLoader {
fn on_event_impl(&mut self, ev: Event, mark: Marker) -> Result<(), ScanError> {
match ev {
Event::DocumentStart | Event::Nothing | Event::StreamStart | Event::StreamEnd => {
}
Event::DocumentEnd => {
match self.doc_stack.len() {
0 => self.docs.push(Yaml::BadValue),
1 => self.docs.push(self.doc_stack.pop().unwrap().0),
_ => unreachable!(),
}
}
Event::SequenceStart(aid, _) => {
self.doc_stack.push((Yaml::Array(Vec::new()), aid));
}
Event::SequenceEnd => {
let node = self.doc_stack.pop().unwrap();
self.insert_new_node(node, mark)?;
}
Event::MappingStart(aid, _) => {
self.doc_stack.push((Yaml::Hash(Hash::new()), aid));
self.key_stack.push(Yaml::BadValue);
}
Event::MappingEnd => {
self.key_stack.pop().unwrap();
let node = self.doc_stack.pop().unwrap();
self.insert_new_node(node, mark)?;
}
Event::Scalar(v, style, aid, tag) => {
let node = if style != TScalarStyle::Plain {
Yaml::String(v)
} else if let Some(Tag {
ref handle,
ref suffix,
}) = tag
{
if handle == "tag:yaml.org,2002:" {
match suffix.as_ref() {
"bool" => {
match v.parse::<bool>() {
Err(_) => Yaml::BadValue,
Ok(v) => Yaml::Boolean(v),
}
}
"int" => match v.parse::<i64>() {
Err(_) => Yaml::BadValue,
Ok(v) => Yaml::Integer(v),
},
"float" => match parse_f64(&v) {
Some(_) => Yaml::Real(v),
None => Yaml::BadValue,
},
"null" => match v.as_ref() {
"~" | "null" => Yaml::Null,
_ => Yaml::BadValue,
},
_ => Yaml::String(v),
}
} else {
Yaml::String(v)
}
} else {
Yaml::from_str(&v)
};
self.insert_new_node((node, aid), mark)?;
}
Event::Alias(id) => {
let n = match self.anchor_map.get(&id) {
Some(v) => v.clone(),
None => Yaml::BadValue,
};
self.insert_new_node((n, 0), mark)?;
}
}
Ok(())
}
fn insert_new_node(&mut self, node: (Yaml, usize), mark: Marker) -> Result<(), ScanError> {
if node.1 > 0 {
self.anchor_map.insert(node.1, node.0.clone());
}
if self.doc_stack.is_empty() {
self.doc_stack.push(node);
} else {
let parent = self.doc_stack.last_mut().unwrap();
match *parent {
(Yaml::Array(ref mut v), _) => v.push(node.0),
(Yaml::Hash(ref mut h), _) => {
let cur_key = self.key_stack.last_mut().unwrap();
if cur_key.is_badvalue() {
*cur_key = node.0;
} else {
let mut newkey = Yaml::BadValue;
mem::swap(&mut newkey, cur_key);
if h.insert(newkey, node.0).is_some() {
let inserted_key = h.back().unwrap().0;
return Err(ScanError::new_string(
mark,
format!("{inserted_key:?}: duplicated key in mapping"),
));
}
}
}
_ => unreachable!(),
}
}
Ok(())
}
pub fn load_from_str(source: &str) -> Result<Vec<Yaml>, ScanError> {
Self::load_from_iter(source.chars())
}
pub fn load_from_iter<I: Iterator<Item = char>>(source: I) -> Result<Vec<Yaml>, ScanError> {
let mut parser = Parser::new(source);
Self::load_from_parser(&mut parser)
}
pub fn load_from_parser<I: Iterator<Item = char>>(
parser: &mut Parser<I>,
) -> Result<Vec<Yaml>, ScanError> {
let mut loader = YamlLoader::default();
parser.load(&mut loader, true)?;
if let Some(e) = loader.error {
Err(e)
} else {
Ok(loader.docs)
}
}
#[must_use]
pub fn documents(&self) -> &[Yaml] {
&self.docs
}
}
#[cfg(feature = "encoding")]
pub type YAMLDecodingTrapFn = fn(
malformation_length: u8,
bytes_read_after_malformation: u8,
input_at_malformation: &[u8],
output: &mut String,
) -> ControlFlow<Cow<'static, str>>;
#[cfg(feature = "encoding")]
#[derive(Copy, Clone)]
pub enum YAMLDecodingTrap {
Ignore,
Strict,
Replace,
Call(YAMLDecodingTrapFn),
}
impl PartialEq for YAMLDecodingTrap {
fn eq(&self, other: &YAMLDecodingTrap) -> bool {
match (self, other) {
(YAMLDecodingTrap::Call(self_fn), YAMLDecodingTrap::Call(other_fn)) => {
*self_fn as usize == *other_fn as usize
}
(x, y) => x == y,
}
}
}
impl Eq for YAMLDecodingTrap {}
#[cfg(feature = "encoding")]
pub struct YamlDecoder<T: std::io::Read> {
source: T,
trap: YAMLDecodingTrap,
}
#[cfg(feature = "encoding")]
impl<T: std::io::Read> YamlDecoder<T> {
pub fn read(source: T) -> YamlDecoder<T> {
YamlDecoder {
source,
trap: YAMLDecodingTrap::Strict,
}
}
pub fn encoding_trap(&mut self, trap: YAMLDecodingTrap) -> &mut Self {
self.trap = trap;
self
}
pub fn decode(&mut self) -> Result<Vec<Yaml>, LoadError> {
let mut buffer = Vec::new();
self.source.read_to_end(&mut buffer)?;
let (encoding, _) =
Encoding::for_bom(&buffer).unwrap_or_else(|| (detect_utf16_endianness(&buffer), 2));
let mut decoder = encoding.new_decoder();
let mut output = String::new();
decode_loop(&buffer, &mut output, &mut decoder, self.trap)?;
YamlLoader::load_from_str(&output).map_err(LoadError::Scan)
}
}
#[cfg(feature = "encoding")]
fn decode_loop(
input: &[u8],
output: &mut String,
decoder: &mut Decoder,
trap: YAMLDecodingTrap,
) -> Result<(), LoadError> {
output.reserve(input.len());
let mut total_bytes_read = 0;
loop {
match decoder.decode_to_string_without_replacement(&input[total_bytes_read..], output, true)
{
(DecoderResult::InputEmpty, _) => break Ok(()),
(DecoderResult::OutputFull, bytes_read) => {
total_bytes_read += bytes_read;
output.reserve(input.len() / 10);
}
(DecoderResult::Malformed(malformed_len, bytes_after_malformed), bytes_read) => {
total_bytes_read += bytes_read;
match trap {
YAMLDecodingTrap::Ignore => {}
YAMLDecodingTrap::Replace => {
output.push('\u{FFFD}');
}
YAMLDecodingTrap::Strict => {
let malformed_len = malformed_len as usize;
let bytes_after_malformed = bytes_after_malformed as usize;
let byte_idx = total_bytes_read - (malformed_len + bytes_after_malformed);
let malformed_sequence = &input[byte_idx..byte_idx + malformed_len];
break Err(LoadError::Decode(Cow::Owned(format!(
"Invalid character sequence at {byte_idx}: {malformed_sequence:?}",
))));
}
YAMLDecodingTrap::Call(callback) => {
let byte_idx =
total_bytes_read - ((malformed_len + bytes_after_malformed) as usize);
let malformed_sequence =
&input[byte_idx..byte_idx + malformed_len as usize];
if let ControlFlow::Break(error) = callback(
malformed_len,
bytes_after_malformed,
&input[byte_idx..],
output,
) {
if error.is_empty() {
break Err(LoadError::Decode(Cow::Owned(format!(
"Invalid character sequence at {byte_idx}: {malformed_sequence:?}",
))));
}
break Err(LoadError::Decode(error));
}
}
}
}
}
}
}
#[cfg(feature = "encoding")]
fn detect_utf16_endianness(b: &[u8]) -> &'static Encoding {
if b.len() > 1 && (b[0] != b[1]) {
if b[0] == 0 {
return encoding_rs::UTF_16BE;
} else if b[1] == 0 {
return encoding_rs::UTF_16LE;
}
}
encoding_rs::UTF_8
}
macro_rules! define_as (
($name:ident, $t:ident, $yt:ident) => (
#[must_use]
pub fn $name(&self) -> Option<$t> {
match *self {
Yaml::$yt(v) => Some(v),
_ => None
}
}
);
);
macro_rules! define_as_ref (
($name:ident, $t:ty, $yt:ident) => (
#[must_use]
pub fn $name(&self) -> Option<$t> {
match *self {
Yaml::$yt(ref v) => Some(v),
_ => None
}
}
);
);
macro_rules! define_as_mut_ref (
($name:ident, $t:ty, $yt:ident) => (
#[must_use]
pub fn $name(&mut self) -> Option<$t> {
match *self {
Yaml::$yt(ref mut v) => Some(v),
_ => None
}
}
);
);
macro_rules! define_into (
($name:ident, $t:ty, $yt:ident) => (
#[must_use]
pub fn $name(self) -> Option<$t> {
match self {
Yaml::$yt(v) => Some(v),
_ => None
}
}
);
);
impl Yaml {
define_as!(as_bool, bool, Boolean);
define_as!(as_i64, i64, Integer);
define_as_ref!(as_str, &str, String);
define_as_ref!(as_hash, &Hash, Hash);
define_as_ref!(as_vec, &Array, Array);
define_as_mut_ref!(as_mut_hash, &mut Hash, Hash);
define_as_mut_ref!(as_mut_vec, &mut Array, Array);
define_into!(into_bool, bool, Boolean);
define_into!(into_i64, i64, Integer);
define_into!(into_string, String, String);
define_into!(into_hash, Hash, Hash);
define_into!(into_vec, Array, Array);
#[must_use]
pub fn is_null(&self) -> bool {
matches!(*self, Yaml::Null)
}
#[must_use]
pub fn is_badvalue(&self) -> bool {
matches!(*self, Yaml::BadValue)
}
#[must_use]
pub fn is_array(&self) -> bool {
matches!(*self, Yaml::Array(_))
}
#[must_use]
pub fn is_hash(&self) -> bool {
matches!(*self, Yaml::Hash(_))
}
#[must_use]
pub fn as_f64(&self) -> Option<f64> {
if let Yaml::Real(ref v) = self {
parse_f64(v)
} else {
None
}
}
#[must_use]
pub fn into_f64(self) -> Option<f64> {
self.as_f64()
}
#[must_use]
pub fn or(self, other: Self) -> Self {
match self {
Yaml::BadValue | Yaml::Null => other,
this => this,
}
}
#[must_use]
pub fn borrowed_or<'a>(&'a self, other: &'a Self) -> &'a Self {
match self {
Yaml::BadValue | Yaml::Null => other,
this => this,
}
}
}
#[allow(clippy::should_implement_trait)]
impl Yaml {
#[must_use]
pub fn from_str(v: &str) -> Yaml {
if let Some(number) = v.strip_prefix("0x") {
if let Ok(i) = i64::from_str_radix(number, 16) {
return Yaml::Integer(i);
}
} else if let Some(number) = v.strip_prefix("0o") {
if let Ok(i) = i64::from_str_radix(number, 8) {
return Yaml::Integer(i);
}
} else if let Some(number) = v.strip_prefix('+') {
if let Ok(i) = number.parse::<i64>() {
return Yaml::Integer(i);
}
}
match v {
"" | "~" | "null" => Yaml::Null,
"true" => Yaml::Boolean(true),
"false" => Yaml::Boolean(false),
_ => {
if let Ok(integer) = v.parse::<i64>() {
Yaml::Integer(integer)
} else if parse_f64(v).is_some() {
Yaml::Real(v.to_owned())
} else {
Yaml::String(v.to_owned())
}
}
}
}
}
static BAD_VALUE: Yaml = Yaml::BadValue;
impl<'a> Index<&'a str> for Yaml {
type Output = Yaml;
fn index(&self, idx: &'a str) -> &Yaml {
let key = Yaml::String(idx.to_owned());
match self.as_hash() {
Some(h) => h.get(&key).unwrap_or(&BAD_VALUE),
None => &BAD_VALUE,
}
}
}
impl<'a> IndexMut<&'a str> for Yaml {
fn index_mut(&mut self, idx: &'a str) -> &mut Yaml {
let key = Yaml::String(idx.to_owned());
match self.as_mut_hash() {
Some(h) => h.get_mut(&key).unwrap(),
None => panic!("Not a hash type"),
}
}
}
impl Index<usize> for Yaml {
type Output = Yaml;
fn index(&self, idx: usize) -> &Yaml {
if let Some(v) = self.as_vec() {
v.get(idx).unwrap_or(&BAD_VALUE)
} else if let Some(v) = self.as_hash() {
let key = Yaml::Integer(i64::try_from(idx).unwrap());
v.get(&key).unwrap_or(&BAD_VALUE)
} else {
&BAD_VALUE
}
}
}
impl IndexMut<usize> for Yaml {
fn index_mut(&mut self, idx: usize) -> &mut Yaml {
match self {
Yaml::Array(sequence) => sequence.index_mut(idx),
Yaml::Hash(mapping) => {
let key = Yaml::Integer(i64::try_from(idx).unwrap());
mapping.get_mut(&key).unwrap()
}
_ => panic!("Attempting to index but `self` is not a sequence nor a mapping"),
}
}
}
impl IntoIterator for Yaml {
type Item = Yaml;
type IntoIter = YamlIter;
fn into_iter(self) -> Self::IntoIter {
YamlIter {
yaml: self.into_vec().unwrap_or_default().into_iter(),
}
}
}
pub struct YamlIter {
yaml: std::vec::IntoIter<Yaml>,
}
impl Iterator for YamlIter {
type Item = Yaml;
fn next(&mut self) -> Option<Yaml> {
self.yaml.next()
}
}
#[cfg(test)]
mod test {
use super::{YAMLDecodingTrap, Yaml, YamlDecoder};
#[test]
fn test_read_bom() {
let s = b"\xef\xbb\xbf---
a: 1
b: 2.2
c: [1, 2]
";
let out = YamlDecoder::read(s as &[u8]).decode().unwrap();
let doc = &out[0];
assert_eq!(doc["a"].as_i64().unwrap(), 1i64);
assert!((doc["b"].as_f64().unwrap() - 2.2f64).abs() <= f64::EPSILON);
assert_eq!(doc["c"][1].as_i64().unwrap(), 2i64);
assert!(doc["d"][0].is_badvalue());
}
#[test]
fn test_read_utf16le() {
let s = b"\xff\xfe-\x00-\x00-\x00
\x00a\x00:\x00 \x001\x00
\x00b\x00:\x00 \x002\x00.\x002\x00
\x00c\x00:\x00 \x00[\x001\x00,\x00 \x002\x00]\x00
\x00";
let out = YamlDecoder::read(s as &[u8]).decode().unwrap();
let doc = &out[0];
println!("GOT: {doc:?}");
assert_eq!(doc["a"].as_i64().unwrap(), 1i64);
assert!((doc["b"].as_f64().unwrap() - 2.2f64) <= f64::EPSILON);
assert_eq!(doc["c"][1].as_i64().unwrap(), 2i64);
assert!(doc["d"][0].is_badvalue());
}
#[test]
fn test_read_utf16be() {
let s = b"\xfe\xff\x00-\x00-\x00-\x00
\x00a\x00:\x00 \x001\x00
\x00b\x00:\x00 \x002\x00.\x002\x00
\x00c\x00:\x00 \x00[\x001\x00,\x00 \x002\x00]\x00
";
let out = YamlDecoder::read(s as &[u8]).decode().unwrap();
let doc = &out[0];
println!("GOT: {doc:?}");
assert_eq!(doc["a"].as_i64().unwrap(), 1i64);
assert!((doc["b"].as_f64().unwrap() - 2.2f64).abs() <= f64::EPSILON);
assert_eq!(doc["c"][1].as_i64().unwrap(), 2i64);
assert!(doc["d"][0].is_badvalue());
}
#[test]
fn test_read_utf16le_nobom() {
let s = b"-\x00-\x00-\x00
\x00a\x00:\x00 \x001\x00
\x00b\x00:\x00 \x002\x00.\x002\x00
\x00c\x00:\x00 \x00[\x001\x00,\x00 \x002\x00]\x00
\x00";
let out = YamlDecoder::read(s as &[u8]).decode().unwrap();
let doc = &out[0];
println!("GOT: {doc:?}");
assert_eq!(doc["a"].as_i64().unwrap(), 1i64);
assert!((doc["b"].as_f64().unwrap() - 2.2f64).abs() <= f64::EPSILON);
assert_eq!(doc["c"][1].as_i64().unwrap(), 2i64);
assert!(doc["d"][0].is_badvalue());
}
#[test]
fn test_read_trap() {
let s = b"---
a\xa9: 1
b: 2.2
c: [1, 2]
";
let out = YamlDecoder::read(s as &[u8])
.encoding_trap(YAMLDecodingTrap::Ignore)
.decode()
.unwrap();
let doc = &out[0];
println!("GOT: {doc:?}");
assert_eq!(doc["a"].as_i64().unwrap(), 1i64);
assert!((doc["b"].as_f64().unwrap() - 2.2f64).abs() <= f64::EPSILON);
assert_eq!(doc["c"][1].as_i64().unwrap(), 2i64);
assert!(doc["d"][0].is_badvalue());
}
#[test]
fn test_or() {
assert_eq!(Yaml::Null.or(Yaml::Integer(3)), Yaml::Integer(3));
assert_eq!(Yaml::Integer(3).or(Yaml::Integer(7)), Yaml::Integer(3));
}
}