gemini_memory_rs/okf/
yaml.rs

1//! A strict YAML subset for OKF front matter.
2//!
3//! The engine both writes and reads this front matter, so it needs exactly the
4//! shape it emits — block mappings, block sequences of scalars, and scalars —
5//! and nothing else. Implementing that subset directly (rather than pulling a
6//! general YAML parser) keeps the canonical format dependency-free and lets
7//! parse failures name the offending line.
8//!
9//! Supported:
10//!
11//! ```yaml
12//! key: scalar          # comments, after a value or on their own line
13//! nested:
14//!   key: value
15//! list:
16//!   - item
17//!   - item
18//! empty_list: []
19//! quoted: "a: value with punctuation"
20//! nothing: null
21//! ```
22//!
23//! Not supported (and rejected with a diagnostic): flow mappings, anchors,
24//! aliases, multi-line scalars, tabs for indentation, and documents whose root
25//! is not a mapping.
26
27use std::fmt::Write as _;
28
29/// A parsed YAML-subset value.
30#[derive(Debug, Clone, PartialEq)]
31pub enum Yaml {
32    /// `null`, `~`, or an empty value.
33    Null,
34    /// `true` / `false`.
35    Bool(bool),
36    /// An integer.
37    Int(i64),
38    /// A floating-point number.
39    Float(f64),
40    /// A string, quoted or bare.
41    Str(String),
42    /// A block sequence.
43    Seq(Vec<Yaml>),
44    /// A block mapping. Order is preserved so emitted files are stable.
45    Map(Vec<(String, Yaml)>),
46}
47
48impl Yaml {
49    /// Look up a key in a mapping.
50    pub fn get(&self, key: &str) -> Option<&Yaml> {
51        match self {
52            Self::Map(entries) => entries.iter().find(|(k, _)| k == key).map(|(_, v)| v),
53            _ => None,
54        }
55    }
56
57    /// Borrow as a string, accepting any scalar and rendering it.
58    pub fn as_str(&self) -> Option<&str> {
59        match self {
60            Self::Str(s) => Some(s),
61            _ => None,
62        }
63    }
64
65    /// Read a string, or `None` for null/missing.
66    pub fn as_string(&self) -> Option<String> {
67        match self {
68            Self::Str(s) => Some(s.clone()),
69            Self::Bool(b) => Some(b.to_string()),
70            Self::Int(i) => Some(i.to_string()),
71            Self::Float(f) => Some(f.to_string()),
72            _ => None,
73        }
74    }
75
76    /// Read a boolean.
77    pub fn as_bool(&self) -> Option<bool> {
78        match self {
79            Self::Bool(b) => Some(*b),
80            _ => None,
81        }
82    }
83
84    /// Read a number as `f64`, accepting integers.
85    pub fn as_f64(&self) -> Option<f64> {
86        match self {
87            Self::Float(f) => Some(*f),
88            Self::Int(i) => Some(*i as f64),
89            _ => None,
90        }
91    }
92
93    /// Read a number as `u64`.
94    pub fn as_u64(&self) -> Option<u64> {
95        match self {
96            Self::Int(i) if *i >= 0 => Some(*i as u64),
97            Self::Float(f) if *f >= 0.0 => Some(*f as u64),
98            _ => None,
99        }
100    }
101
102    /// Borrow as a sequence.
103    pub fn as_seq(&self) -> Option<&[Yaml]> {
104        match self {
105            Self::Seq(items) => Some(items),
106            _ => None,
107        }
108    }
109
110    /// Read a sequence of strings, tolerating a missing or null value.
111    pub fn as_string_list(&self) -> Vec<String> {
112        match self {
113            Self::Seq(items) => items.iter().filter_map(Yaml::as_string).collect(),
114            Self::Null => Vec::new(),
115            other => other.as_string().into_iter().collect(),
116        }
117    }
118
119    /// Whether the value is null.
120    pub fn is_null(&self) -> bool {
121        matches!(self, Self::Null)
122    }
123}
124
125impl From<&str> for Yaml {
126    fn from(value: &str) -> Self {
127        Yaml::Str(value.to_string())
128    }
129}
130
131impl From<String> for Yaml {
132    fn from(value: String) -> Self {
133        Yaml::Str(value)
134    }
135}
136
137impl From<bool> for Yaml {
138    fn from(value: bool) -> Self {
139        Yaml::Bool(value)
140    }
141}
142
143impl From<u32> for Yaml {
144    fn from(value: u32) -> Self {
145        Yaml::Int(i64::from(value))
146    }
147}
148
149impl From<u64> for Yaml {
150    fn from(value: u64) -> Self {
151        Yaml::Int(value as i64)
152    }
153}
154
155impl From<f32> for Yaml {
156    fn from(value: f32) -> Self {
157        Yaml::Float(f64::from(value))
158    }
159}
160
161/// A YAML-subset parse failure, with the line it occurred on.
162#[derive(Debug, Clone, PartialEq, Eq)]
163pub struct YamlError {
164    /// 1-based line number.
165    pub line: usize,
166    /// What went wrong.
167    pub message: String,
168}
169
170impl std::fmt::Display for YamlError {
171    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
172        write!(f, "line {}: {}", self.line, self.message)
173    }
174}
175
176impl std::error::Error for YamlError {}
177
178/// One significant (non-blank, non-comment) input line.
179struct Line {
180    number: usize,
181    indent: usize,
182    content: String,
183}
184
185/// Parse a YAML-subset document. The root must be a mapping.
186pub fn parse(input: &str) -> Result<Yaml, YamlError> {
187    let lines = significant_lines(input)?;
188    if lines.is_empty() {
189        return Ok(Yaml::Map(Vec::new()));
190    }
191    let mut cursor = 0usize;
192    let base = lines[0].indent;
193    let value = parse_block(&lines, &mut cursor, base)?;
194    if cursor < lines.len() {
195        return Err(YamlError {
196            line: lines[cursor].number,
197            message: "unexpected dedent — inconsistent indentation".to_string(),
198        });
199    }
200    match value {
201        Yaml::Map(_) => Ok(value),
202        _ => Err(YamlError {
203            line: lines[0].number,
204            message: "OKF front matter must be a mapping".to_string(),
205        }),
206    }
207}
208
209fn significant_lines(input: &str) -> Result<Vec<Line>, YamlError> {
210    let mut out = Vec::new();
211    for (idx, raw) in input.lines().enumerate() {
212        let number = idx + 1;
213        if raw.contains('\t') && raw.trim_start().len() != raw.len() {
214            return Err(YamlError {
215                line: number,
216                message: "tabs may not be used for indentation".to_string(),
217            });
218        }
219        let trimmed = raw.trim_end();
220        let indent = trimmed.len() - trimmed.trim_start().len();
221        let content = trimmed.trim_start().to_string();
222        if content.is_empty() || content.starts_with('#') || content == "---" {
223            continue;
224        }
225        out.push(Line {
226            number,
227            indent,
228            content,
229        });
230    }
231    Ok(out)
232}
233
234fn parse_block(lines: &[Line], cursor: &mut usize, indent: usize) -> Result<Yaml, YamlError> {
235    if *cursor >= lines.len() {
236        return Ok(Yaml::Null);
237    }
238    if lines[*cursor].content.starts_with("- ") || lines[*cursor].content == "-" {
239        parse_sequence(lines, cursor, indent)
240    } else {
241        parse_mapping(lines, cursor, indent)
242    }
243}
244
245fn parse_sequence(lines: &[Line], cursor: &mut usize, indent: usize) -> Result<Yaml, YamlError> {
246    let mut items = Vec::new();
247    while *cursor < lines.len() && lines[*cursor].indent == indent {
248        let line = &lines[*cursor];
249        let Some(rest) = line
250            .content
251            .strip_prefix("- ")
252            .or_else(|| (line.content == "-").then_some(""))
253        else {
254            break;
255        };
256        *cursor += 1;
257        let scalar = strip_comment(rest.trim());
258        if scalar.is_empty() {
259            // A nested block under a bare `-`.
260            if *cursor < lines.len() && lines[*cursor].indent > indent {
261                let child_indent = lines[*cursor].indent;
262                items.push(parse_block(lines, cursor, child_indent)?);
263            } else {
264                items.push(Yaml::Null);
265            }
266        } else {
267            items.push(parse_scalar(scalar));
268        }
269    }
270    Ok(Yaml::Seq(items))
271}
272
273fn parse_mapping(lines: &[Line], cursor: &mut usize, indent: usize) -> Result<Yaml, YamlError> {
274    let mut entries: Vec<(String, Yaml)> = Vec::new();
275    while *cursor < lines.len() && lines[*cursor].indent == indent {
276        let line = &lines[*cursor];
277        if line.content.starts_with("- ") {
278            break;
279        }
280        let (key, rest) = split_key(&line.content).ok_or_else(|| YamlError {
281            line: line.number,
282            message: format!("expected `key: value`, found `{}`", line.content),
283        })?;
284        let line_number = line.number;
285        *cursor += 1;
286
287        let inline = strip_comment(rest.trim());
288        let value = if inline.is_empty() {
289            if *cursor < lines.len() && lines[*cursor].indent > indent {
290                let child_indent = lines[*cursor].indent;
291                parse_block(lines, cursor, child_indent)?
292            } else {
293                Yaml::Null
294            }
295        } else {
296            parse_scalar(inline)
297        };
298
299        if entries.iter().any(|(k, _)| k == &key) {
300            return Err(YamlError {
301                line: line_number,
302                message: format!("duplicate key `{key}`"),
303            });
304        }
305        entries.push((key, value));
306    }
307    Ok(Yaml::Map(entries))
308}
309
310/// Split `key: value`, respecting quoted keys and not splitting inside quotes.
311fn split_key(content: &str) -> Option<(String, &str)> {
312    let bytes = content.as_bytes();
313    let mut quote: Option<u8> = None;
314    for (i, b) in bytes.iter().enumerate() {
315        match quote {
316            Some(q) if *b == q => quote = None,
317            Some(_) => {}
318            None if *b == b'"' || *b == b'\'' => quote = Some(*b),
319            None if *b == b':' => {
320                let is_last = i + 1 == bytes.len();
321                if is_last || bytes[i + 1] == b' ' {
322                    let key = unquote(content[..i].trim());
323                    return Some((key, &content[i + 1..]));
324                }
325            }
326            None => {}
327        }
328    }
329    None
330}
331
332/// Remove a trailing `# comment` when it is not inside quotes.
333fn strip_comment(value: &str) -> &str {
334    let bytes = value.as_bytes();
335    let mut quote: Option<u8> = None;
336    for (i, b) in bytes.iter().enumerate() {
337        match quote {
338            Some(q) if *b == q => quote = None,
339            Some(_) => {}
340            None if *b == b'"' || *b == b'\'' => quote = Some(*b),
341            None if *b == b'#' && (i == 0 || bytes[i - 1] == b' ') => {
342                return value[..i].trim_end();
343            }
344            None => {}
345        }
346    }
347    value
348}
349
350fn unquote(raw: &str) -> String {
351    let bytes = raw.as_bytes();
352    if bytes.len() >= 2
353        && ((bytes[0] == b'"' && bytes[bytes.len() - 1] == b'"')
354            || (bytes[0] == b'\'' && bytes[bytes.len() - 1] == b'\''))
355    {
356        raw[1..raw.len() - 1]
357            .replace("\\\"", "\"")
358            .replace("\\n", "\n")
359    } else {
360        raw.to_string()
361    }
362}
363
364fn parse_scalar(raw: &str) -> Yaml {
365    if raw == "[]" {
366        return Yaml::Seq(Vec::new());
367    }
368    if raw == "{}" {
369        return Yaml::Map(Vec::new());
370    }
371    let bytes = raw.as_bytes();
372    let quoted = bytes.len() >= 2
373        && ((bytes[0] == b'"' && bytes[bytes.len() - 1] == b'"')
374            || (bytes[0] == b'\'' && bytes[bytes.len() - 1] == b'\''));
375    if quoted {
376        return Yaml::Str(unquote(raw));
377    }
378    match raw {
379        "null" | "~" | "" => Yaml::Null,
380        "true" => Yaml::Bool(true),
381        "false" => Yaml::Bool(false),
382        _ => {
383            if let Ok(i) = raw.parse::<i64>() {
384                Yaml::Int(i)
385            } else if let Ok(f) = raw.parse::<f64>() {
386                Yaml::Float(f)
387            } else {
388                Yaml::Str(raw.to_string())
389            }
390        }
391    }
392}
393
394/// Render a value as YAML-subset text. The root must be a mapping.
395pub fn emit(value: &Yaml) -> String {
396    let mut out = String::new();
397    emit_into(value, 0, &mut out);
398    out
399}
400
401fn emit_into(value: &Yaml, indent: usize, out: &mut String) {
402    match value {
403        Yaml::Map(entries) => {
404            for (key, child) in entries {
405                let pad = " ".repeat(indent);
406                match child {
407                    Yaml::Map(inner) if inner.is_empty() => {
408                        let _ = writeln!(out, "{pad}{key}: {{}}");
409                    }
410                    Yaml::Seq(items) if items.is_empty() => {
411                        let _ = writeln!(out, "{pad}{key}: []");
412                    }
413                    Yaml::Map(_) | Yaml::Seq(_) => {
414                        let _ = writeln!(out, "{pad}{key}:");
415                        emit_into(child, indent + 2, out);
416                    }
417                    scalar => {
418                        let _ = writeln!(out, "{pad}{key}: {}", emit_scalar(scalar));
419                    }
420                }
421            }
422        }
423        Yaml::Seq(items) => {
424            for item in items {
425                let pad = " ".repeat(indent);
426                match item {
427                    Yaml::Map(_) | Yaml::Seq(_) => {
428                        let _ = writeln!(out, "{pad}-");
429                        emit_into(item, indent + 2, out);
430                    }
431                    scalar => {
432                        let _ = writeln!(out, "{pad}- {}", emit_scalar(scalar));
433                    }
434                }
435            }
436        }
437        scalar => {
438            let pad = " ".repeat(indent);
439            let _ = writeln!(out, "{pad}{}", emit_scalar(scalar));
440        }
441    }
442}
443
444fn emit_scalar(value: &Yaml) -> String {
445    match value {
446        Yaml::Null => "null".to_string(),
447        Yaml::Bool(b) => b.to_string(),
448        Yaml::Int(i) => i.to_string(),
449        Yaml::Float(f) => {
450            if f.fract() == 0.0 && f.is_finite() {
451                format!("{f:.1}")
452            } else {
453                format!("{f}")
454            }
455        }
456        Yaml::Str(s) => quote_if_needed(s),
457        // Nested collections are handled by `emit_into`; reaching here means an
458        // empty collection was inlined.
459        Yaml::Seq(_) => "[]".to_string(),
460        Yaml::Map(_) => "{}".to_string(),
461    }
462}
463
464/// Quote a string when leaving it bare would change how it re-parses.
465fn quote_if_needed(raw: &str) -> String {
466    let needs_quotes = raw.is_empty()
467        || raw.trim() != raw
468        || raw.contains(": ")
469        || raw.ends_with(':')
470        || raw.contains('\n')
471        || raw.contains(" #")
472        || raw.starts_with([
473            '-', '?', '&', '*', '!', '|', '>', '%', '@', '`', '[', '{', '"', '\'',
474        ])
475        || matches!(raw, "null" | "true" | "false" | "~")
476        || raw.parse::<f64>().is_ok();
477    if needs_quotes {
478        format!("\"{}\"", raw.replace('"', "\\\"").replace('\n', "\\n"))
479    } else {
480        raw.to_string()
481    }
482}
483
484/// Build a mapping from an ordered list of entries, dropping `None` values.
485pub fn map(entries: Vec<(&str, Option<Yaml>)>) -> Yaml {
486    Yaml::Map(
487        entries
488            .into_iter()
489            .filter_map(|(k, v)| v.map(|v| (k.to_string(), v)))
490            .collect(),
491    )
492}
493
494/// Build a sequence of strings.
495pub fn seq_of_strings<I, S>(items: I) -> Yaml
496where
497    I: IntoIterator<Item = S>,
498    S: Into<String>,
499{
500    Yaml::Seq(items.into_iter().map(|s| Yaml::Str(s.into())).collect())
501}
502
503#[cfg(test)]
504mod tests {
505    use super::*;
506
507    #[test]
508    fn parses_the_okf_front_matter_shape() {
509        let src = r#"
510okf: memory/v1
511id: mem_01K4D8P3
512confidence: 1.0
513source:
514  type: explicit_user_statement
515  turn_id: turn_17
516retrieval:
517  tags:
518    - food
519    - diet
520  aliases: []
521privacy:
522  deletable: true
523  sensitivity: normal
524valid_to: null
525"#;
526        let parsed = parse(src).unwrap();
527        assert_eq!(parsed.get("okf").unwrap().as_str(), Some("memory/v1"));
528        assert_eq!(parsed.get("confidence").unwrap().as_f64(), Some(1.0));
529        assert_eq!(
530            parsed
531                .get("source")
532                .unwrap()
533                .get("turn_id")
534                .unwrap()
535                .as_str(),
536            Some("turn_17")
537        );
538        assert_eq!(
539            parsed
540                .get("retrieval")
541                .unwrap()
542                .get("tags")
543                .unwrap()
544                .as_string_list(),
545            vec!["food", "diet"]
546        );
547        assert_eq!(
548            parsed
549                .get("retrieval")
550                .unwrap()
551                .get("aliases")
552                .unwrap()
553                .as_seq(),
554            Some(&[][..])
555        );
556        assert_eq!(
557            parsed
558                .get("privacy")
559                .unwrap()
560                .get("deletable")
561                .unwrap()
562                .as_bool(),
563            Some(true)
564        );
565        assert!(parsed.get("valid_to").unwrap().is_null());
566    }
567
568    #[test]
569    fn round_trips_through_emit_and_parse() {
570        let original = map(vec![
571            ("okf", Some("memory/v1".into())),
572            ("confidence", Some(Yaml::Float(0.85))),
573            (
574                "retrieval",
575                Some(map(vec![
576                    ("tags", Some(seq_of_strings(["food", "diet"]))),
577                    ("aliases", Some(seq_of_strings(Vec::<String>::new()))),
578                ])),
579            ),
580            ("valid_to", Some(Yaml::Null)),
581        ]);
582        let text = emit(&original);
583        assert_eq!(parse(&text).unwrap(), original);
584    }
585
586    #[test]
587    fn strings_that_would_re_parse_as_other_types_are_quoted() {
588        let value = map(vec![
589            ("a", Some("true".into())),
590            ("b", Some("2026".into())),
591            ("c", Some("a: b".into())),
592            ("d", Some("- leading dash".into())),
593        ]);
594        let text = emit(&value);
595        let reparsed = parse(&text).unwrap();
596        assert_eq!(reparsed.get("a").unwrap().as_str(), Some("true"));
597        assert_eq!(reparsed.get("b").unwrap().as_str(), Some("2026"));
598        assert_eq!(reparsed.get("c").unwrap().as_str(), Some("a: b"));
599        assert_eq!(reparsed.get("d").unwrap().as_str(), Some("- leading dash"));
600    }
601
602    #[test]
603    fn comments_are_ignored_but_hashes_inside_quotes_survive() {
604        let src = "# leading comment\nkey: value # trailing\nquoted: \"has # hash\"\n";
605        let parsed = parse(src).unwrap();
606        assert_eq!(parsed.get("key").unwrap().as_str(), Some("value"));
607        assert_eq!(parsed.get("quoted").unwrap().as_str(), Some("has # hash"));
608    }
609
610    #[test]
611    fn timestamps_stay_strings() {
612        let parsed = parse("created_at: 2026-07-26T09:12:14Z\n").unwrap();
613        assert_eq!(
614            parsed.get("created_at").unwrap().as_str(),
615            Some("2026-07-26T09:12:14Z")
616        );
617    }
618
619    #[test]
620    fn tabs_and_duplicate_keys_are_rejected_with_a_line_number() {
621        let tabbed = parse("a:\n\tb: 1\n").unwrap_err();
622        assert_eq!(tabbed.line, 2);
623
624        let duped = parse("a: 1\na: 2\n").unwrap_err();
625        assert_eq!(duped.line, 2);
626        assert!(duped.message.contains("duplicate"));
627    }
628
629    #[test]
630    fn a_non_mapping_root_is_rejected() {
631        assert!(parse("- just\n- a list\n").is_err());
632    }
633
634    #[test]
635    fn nested_maps_inside_sequences_round_trip() {
636        let value = Yaml::Map(vec![(
637            "items".to_string(),
638            Yaml::Seq(vec![
639                Yaml::Map(vec![("name".into(), "a".into())]),
640                Yaml::Map(vec![("name".into(), "b".into())]),
641            ]),
642        )]);
643        let text = emit(&value);
644        assert_eq!(parse(&text).unwrap(), value);
645    }
646}