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(),
)
}
}
}