1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100

use rand::Rng;

use crate::{Markov, CFGrammar, FormattingRule, WorkingSet, SampleSet, LearnError, WordList};
use crate::formatting::format_ws;
use crate::core::ValidationError;

#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Clone)]
enum PartGenerator {
    #[cfg_attr(feature = "serde", serde(rename="markov"))]
    Markov(Markov),
    #[cfg_attr(feature = "serde", serde(rename="cfgrammar"))]
    CFGrammar(CFGrammar),
    #[cfg_attr(feature = "serde", serde(rename="wordlist"))]
    WordList(WordList),
}

impl PartGenerator {
    fn generate(&self, ws: &mut WorkingSet, rng: &mut impl Rng) {
        match self {
            PartGenerator::Markov(m) => m.generate(ws, rng),
            PartGenerator::CFGrammar(c) => c.generate(ws, rng),
            PartGenerator::WordList(wl) => wl.generate(ws, rng),
        }
    }

    fn learn(&mut self, sample_set: &SampleSet) -> Result<(), LearnError> {
        match self {
            PartGenerator::Markov(m) => m.learn(sample_set),
            PartGenerator::CFGrammar(c) => c.learn(sample_set),
            PartGenerator::WordList(wl) => wl.learn(sample_set),
        }
    }

    fn validate(&self) -> Result<(), ValidationError> {
        match self {
            PartGenerator::Markov(m) => m.validate(),
            PartGenerator::CFGrammar(c) => c.validate(),
            PartGenerator::WordList(wl) => wl.validate(),
        }
    }
}

#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "serde", serde(rename_all = "camelCase"))]
#[derive(Clone)]
pub struct NamePart {
    name: String,
    generator: PartGenerator,
    format_rules: Vec<FormattingRule>,
}

impl NamePart {
    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn generate(&self, ws: &mut WorkingSet, rng: &mut impl Rng) {
        self.generator.generate(ws, rng);
        format_ws(ws, &self.format_rules);
    }

    pub fn learn(&mut self, sample_set: &SampleSet) -> Result<(), LearnError> {
        self.generator.learn(sample_set)
    }

    pub fn validate(&self) -> Result<(), ValidationError> {
        self.generator.validate().map_err(|err| err.with_name(&self.name))
    }

    pub fn new_markov<S: AsRef<str>>(name: &str, format_rules: &[FormattingRule], initial_tokens: &[S], lrs: bool, lrm: bool, lre: bool, rlf: bool) -> NamePart {
        NamePart {
            name: name.to_owned(),
            format_rules: format_rules.to_vec(),
            generator: PartGenerator::Markov(
                Markov::with_constraints(initial_tokens, lrs, lrm, lre, rlf),
            ),
        }
    }

    pub fn new_cfgrammar<S: AsRef<str>>(name: &str, format_rules: &[FormattingRule], initial_subtokens: &[S], rlf: bool, ral: bool) -> NamePart {
        NamePart {
            name: name.to_owned(),
            format_rules: format_rules.to_vec(),
            generator: PartGenerator::CFGrammar(
                CFGrammar::new(initial_subtokens, rlf, ral),
            )
        }
    }

    pub fn new_wordlist(name: &str, format_rules: &[FormattingRule]) -> NamePart {
        NamePart {
            name: name.to_owned(),
            format_rules: format_rules.to_vec(),
            generator: PartGenerator::WordList(
                WordList::new(),
            )
        }
    }
}