saphyr-serde/saphyr/examples/gen_large_yaml/main.rs

#![allow(dead_code)]

mod gen;

use std::collections::HashMap;

use rand::{rngs::ThreadRng, Rng};

fn main() -> std::fmt::Result {
    let mut s = String::new();
    let mut g = Generator::new();
    g.gen_record_array(&mut s, 100_000, 100_001)?;
    println!("{s}");
    Ok(())
}

/// YAML Generator.
struct Generator {
    /// The RNG state.
    rng: ThreadRng,
    /// The stack of indentations.
    indents: Vec<usize>,
}

type GenFn<W> = dyn FnOnce(&mut Generator, &mut W) -> std::fmt::Result;

impl Generator {
    /// Create a new generator.
    fn new() -> Self {
        Generator {
            rng: rand::thread_rng(),
            indents: vec![0],
        }
    }

    /// Generate an array of records as per [`Self::gen_record_object`].
    fn gen_record_array<W: std::fmt::Write>(
        &mut self,
        writer: &mut W,
        items_lo: usize,
        items_hi: usize,
    ) -> std::fmt::Result {
        self.gen_array(writer, items_lo, items_hi, Generator::gen_record_object)
    }

    /// Generate a YAML object/mapping containing a record.
    ///
    /// Fields are description, hash, version, home, repository and pdf.
    fn gen_record_object<W: std::fmt::Write>(&mut self, writer: &mut W) -> std::fmt::Result {
        let mut fields = HashMap::<String, Box<GenFn<W>>>::new();
        fields.insert(
            "description".to_string(),
            Box::new(|gen, w| {
                write!(w, "|")?;
                gen.push_indent(2);
                gen.nl(w)?;
                let indent = gen.indent();
                let text = gen::text(&mut gen.rng, 1, 9, 3, 8, 10, 20, 80 - indent);
                gen.write_lines(w, &text)?;
                gen.pop_indent();
                Ok(())
            }),
        );

        fields.insert(
            "authors".to_string(),
            Box::new(|gen, w| {
                gen.push_indent(2);
                gen.nl(w)?;
                gen.gen_authors_array(w, 1, 10)?;
                gen.pop_indent();
                Ok(())
            }),
        );

        fields.insert(
            "hash".to_string(),
            Box::new(|gen, w| write!(w, "{}", gen::hex_string(&mut gen.rng, 64))),
        );
        fields.insert(
            "version".to_string(),
            Box::new(|gen, w| write!(w, "{}", gen::integer(&mut gen.rng, 1, 9))),
        );
        fields.insert(
            "home".to_string(),
            Box::new(|gen, w| write!(w, "{}", gen::url(&mut gen.rng, "https", 0, 1, 0, 0, None))),
        );
        fields.insert(
            "repository".to_string(),
            Box::new(|gen, w| write!(w, "{}", gen::url(&mut gen.rng, "git", 1, 4, 10, 20, None))),
        );
        fields.insert(
            "pdf".to_string(),
            Box::new(|gen, w| {
                write!(
                    w,
                    "{}",
                    gen::url(&mut gen.rng, "https", 1, 4, 10, 30, Some("pdf"))
                )
            }),
        );
        self.gen_object(writer, fields)
    }

    /// Generate an array of authors as per [`Self::gen_author_object`].
    fn gen_authors_array<W: std::fmt::Write>(
        &mut self,
        writer: &mut W,
        items_lo: usize,
        items_hi: usize,
    ) -> std::fmt::Result {
        self.gen_array(writer, items_lo, items_hi, Generator::gen_author_object)
    }

    fn gen_author_object<W: std::fmt::Write>(&mut self, writer: &mut W) -> std::fmt::Result {
        let mut fields = HashMap::<String, Box<GenFn<W>>>::new();
        fields.insert(
            "name".to_string(),
            Box::new(|gen, w| write!(w, "{}", gen::full_name(&mut gen.rng, 10, 15))),
        );
        fields.insert(
            "email".to_string(),
            Box::new(|gen, w| write!(w, "{}", gen::email(&mut gen.rng, 1, 9))),
        );
        self.gen_object(writer, fields)
    }

    /// Generate a YAML array/sequence containing nodes generated by the given function.
    fn gen_array<W: std::fmt::Write, F: FnMut(&mut Generator, &mut W) -> std::fmt::Result>(
        &mut self,
        writer: &mut W,
        len_lo: usize,
        len_hi: usize,
        mut obj_creator: F,
    ) -> std::fmt::Result {
        let mut first = true;
        for _ in 0..self.rng.gen_range(len_lo..len_hi) {
            if first {
                first = false;
            } else {
                self.nl(writer)?;
            }
            write!(writer, "- ")?;
            self.push_indent(2);
            (obj_creator)(self, writer)?;
            self.pop_indent();
        }
        Ok(())
    }

    /// Create a Yaml object with some fields in it.
    fn gen_object<W: std::fmt::Write>(
        &mut self,
        writer: &mut W,
        fields: HashMap<String, Box<GenFn<W>>>,
    ) -> std::fmt::Result {
        let mut first = true;
        for (key, f) in fields {
            if first {
                first = false;
            } else {
                self.nl(writer)?;
            }
            write!(writer, "{key}: ")?;
            f(self, writer)?;
        }
        Ok(())
    }

    /// Write the given lines at the right indentation.
    fn write_lines<W: std::fmt::Write>(
        &mut self,
        writer: &mut W,
        lines: &[String],
    ) -> std::fmt::Result {
        let mut first = true;

        for line in lines {
            if first {
                first = false;
            } else {
                self.nl(writer)?;
            }
            write!(writer, "{line}")?;
        }

        Ok(())
    }

    /// Write a new line to the writer and indent.
    fn nl<W: std::fmt::Write>(&mut self, writer: &mut W) -> std::fmt::Result {
        writeln!(writer)?;
        for _ in 0..self.indent() {
            write!(writer, " ")?;
        }
        Ok(())
    }

    /// Return the given indent.
    fn indent(&self) -> usize {
        *self.indents.last().unwrap()
    }

    /// Push a new indent with the given relative offset.
    fn push_indent(&mut self, offset: usize) {
        self.indents.push(self.indent() + offset);
    }

    /// Pops the last indent.
    fn pop_indent(&mut self) {
        self.indents.pop();
        assert!(!self.indents.is_empty());
    }
}