Skip to main content

ytsaurus_yson/
node.rs

1use serde::ser::{Serialize, SerializeMap, SerializeSeq, SerializeStruct, Serializer};
2use std::{borrow::Cow, collections::BTreeMap};
3
4/// Represents a complete YSON value, including its optional attributes and data node.
5#[derive(Debug, Clone, PartialEq)]
6pub struct YsonValue {
7    /// Optional attributes associated with this value.
8    /// In YSON, attributes are stored as a map of byte strings to other YSON values.
9    pub attributes: Option<BTreeMap<Vec<u8>, YsonValue>>,
10    /// The data content of this YSON node.
11    pub node: YsonNode,
12}
13
14impl YsonValue {
15    /// Attempts to interpret the node as a UTF-8 string.
16    /// Returns `None` if the node is not a string or if the bytes are not valid UTF-8.
17    #[must_use]
18    pub fn as_str(&self) -> Option<&str> {
19        if let YsonNode::String(bytes) = &self.node {
20            std::str::from_utf8(bytes).ok()
21        } else {
22            None
23        }
24    }
25
26    /// Attempts to interpret the node as a 64-bit signed integer.
27    /// Returns `None` if the node is not an `Int64`.
28    #[must_use]
29    pub fn as_i64(&self) -> Option<i64> {
30        if let YsonNode::Int64(v) = self.node {
31            Some(v)
32        } else {
33            None
34        }
35    }
36
37    /// Retrieves an attribute by its string key.
38    /// Returns `None` if attributes are missing or the key is not found.
39    #[must_use]
40    pub fn attr(&self, key: &str) -> Option<&YsonValue> {
41        self.attributes.as_ref()?.get(key.as_bytes())
42    }
43}
44
45impl<'a> std::ops::Index<&'a str> for YsonValue {
46    type Output = YsonValue;
47
48    /// Provides convenient access to map elements or attributes using index notation.
49    ///
50    /// # Panics
51    /// Panics if the key is not found or if the value is not a map.
52    ///
53    /// # Examples
54    /// ```
55    /// use ytsaurus_yson::{YsonValue, from_slice, YsonFormat};
56    ///
57    /// let input = b"<status=\"ok\">{id=1u}";
58    /// let value: YsonValue = from_slice(input, YsonFormat::Text).unwrap();
59    ///
60    /// // Access an attribute with '@' prefix
61    /// assert_eq!(value["@status"].as_str(), Some("ok"));
62    ///
63    /// // Access a map field directly
64    /// // Note: value["id"] would work if it were a map
65    /// ```
66    fn index(&self, key: &'a str) -> &Self::Output {
67        if let Some(attr_name) = key.strip_prefix('@') {
68            return self
69                .attributes
70                .as_ref()
71                .and_then(|a| a.get(attr_name.as_bytes()))
72                .expect("Attribute not found");
73        }
74        if let YsonNode::Map(m) = &self.node {
75            return m.get(key.as_bytes()).expect("Key not found in map");
76        }
77        panic!("Value is not a map");
78    }
79}
80
81/// Represents the data variants available in the YSON data model.
82#[derive(Debug, Clone, PartialEq)]
83pub enum YsonNode {
84    /// An empty value, represented by `#` in text format.
85    Entity,
86    /// A boolean value (`%true` or `%false`).
87    Boolean(bool),
88    /// A signed 64-bit integer.
89    Int64(i64),
90    /// An unsigned 64-bit integer, followed by `u` in text format (e.g., `42u`).
91    Uint64(u64),
92    /// A double-precision floating point number.
93    Double(f64),
94    /// A byte string.
95    String(Vec<u8>),
96    /// A list of YSON values, enclosed in `[...]`.
97    List(Vec<YsonValue>),
98    /// A map of byte strings to YSON values, enclosed in `{...}`.
99    Map(BTreeMap<Vec<u8>, YsonValue>),
100}
101
102/// Serializes a YSON string key/value, which is an arbitrary byte string.
103///
104/// Valid UTF-8 goes through `serialize_str` so that text output can use the
105/// unquoted-identifier form where possible; anything else goes through
106/// `serialize_bytes`, which never loses non-UTF-8 bytes. In binary format both
107/// paths emit exactly `0x01 + zigzag(len) + raw bytes`.
108struct ByteString<'a>(&'a [u8]);
109
110impl Serialize for ByteString<'_> {
111    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
112        match std::str::from_utf8(self.0) {
113            Ok(s) => serializer.serialize_str(s),
114            Err(_) => serializer.serialize_bytes(self.0),
115        }
116    }
117}
118
119/// Serializes a `BTreeMap<Vec<u8>, YsonValue>` preserving non-UTF-8 keys.
120struct ByteKeyedMap<'a>(&'a BTreeMap<Vec<u8>, YsonValue>);
121
122impl Serialize for ByteKeyedMap<'_> {
123    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
124        let mut map = serializer.serialize_map(Some(self.0.len()))?;
125        for (key, value) in self.0 {
126            map.serialize_entry(&ByteString(key), value)?;
127        }
128        map.end()
129    }
130}
131
132impl Serialize for YsonNode {
133    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
134        match self {
135            YsonNode::Entity => serializer.serialize_unit(),
136            YsonNode::Boolean(v) => serializer.serialize_bool(*v),
137            YsonNode::Int64(v) => serializer.serialize_i64(*v),
138            YsonNode::Uint64(v) => serializer.serialize_u64(*v),
139            YsonNode::Double(v) => serializer.serialize_f64(*v),
140            YsonNode::String(bytes) => ByteString(bytes).serialize(serializer),
141            YsonNode::List(items) => {
142                let mut seq = serializer.serialize_seq(Some(items.len()))?;
143                for item in items {
144                    seq.serialize_element(item)?;
145                }
146                seq.end()
147            }
148            YsonNode::Map(entries) => ByteKeyedMap(entries).serialize(serializer),
149        }
150    }
151}
152
153impl Serialize for YsonValue {
154    /// Round-trips through [`crate::to_vec`]/[`crate::from_slice`].
155    ///
156    /// Note that maps are stored in a `BTreeMap`, so keys come back out in
157    /// sorted order rather than in the order they appeared in the input. The
158    /// round-trip therefore preserves the *value*, not the exact byte layout.
159    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
160        match &self.attributes {
161            Some(attributes) if !attributes.is_empty() => {
162                // `$__yson_attributes` is the marker the serializer looks for to
163                // emit `<attrs>value` rather than a plain map; see `ser.rs`.
164                let mut state = serializer.serialize_struct("$__yson_attributes", 2)?;
165                state.serialize_field("$attributes", &ByteKeyedMap(attributes))?;
166                state.serialize_field("$value", &self.node)?;
167                state.end()
168            }
169            _ => self.node.serialize(serializer),
170        }
171    }
172}
173
174/// Represents individual lexical units (tokens) produced by the YSON lexer.
175#[derive(Debug, Clone, PartialEq)]
176pub enum Token<'a> {
177    /// Opening bracket for attributes: `<`.
178    BeginAttributes,
179    /// Closing bracket for attributes: `>`.
180    EndAttributes,
181    /// Opening bracket for a list: `[`.
182    BeginList,
183    /// Closing bracket for a list: `]`.
184    EndList,
185    /// Opening bracket for a map: `{`.
186    BeginMap,
187    /// Closing bracket for a map: `}`.
188    EndMap,
189
190    /// A string literal, either quoted or unquoted. Uses `Cow` for zero-copy borrowing.
191    String(Cow<'a, [u8]>),
192    /// A signed 64-bit integer literal.
193    Int64(i64),
194    /// An unsigned 64-bit integer literal.
195    Uint64(u64),
196    /// A floating point literal.
197    Double(f64),
198    /// A boolean literal.
199    Boolean(bool),
200    /// The entity literal: `#`.
201    Entity,
202
203    /// Key-value separator: `=`.
204    KeyValueSeparator,
205    /// Item separator: `;`.
206    ItemSeparator,
207}