markov.rs

extern crate rand;

use std::collections::HashMap;
use std::io::{self, Read};

struct Markov {
    rng: rand::ThreadRng,
    seed: String,
    freqs: HashMap<String, Vec<String>>,
}

impl Markov {
    fn new(freqs: HashMap<String, Vec<String>>) -> Self {
        let mut rng = rand::thread_rng();
        let seed = rand::sample(&mut rng, freqs.keys(), 1)[0];
        Markov {
            rng: rng,
            seed: seed.clone(),
            freqs: freqs.clone(),
        }
    }
}

impl Iterator for Markov {
    type Item = String;

    fn next(&mut self) -> Option<String> {
        let empty = Vec::new();
        let next = rand::sample(&mut self.rng,
                                self.freqs.get(&self.seed).unwrap_or(&empty),
                                1);
        self.seed = match next.get(0) {
            Some(x) => x.to_string(),
            None => {
                rand::sample(&mut self.rng, self.freqs.keys(), 1)[0].clone()
            }
        };
        Some(self.seed.clone())
    }
}

fn main() {
    let mut buffer = String::new();
    io::stdin().read_to_string(&mut buffer).unwrap();

    let words: Vec<_> = buffer.split_whitespace().collect();
    let mut freqs: HashMap<String, Vec<String>> = HashMap::new();
    for slice in words.as_slice().windows(2) {
        (*freqs.entry(slice[0].into()).or_insert(Vec::new()))
            .push(slice[1].into())
    }

    let markov = Markov::new(freqs);
    let sentence: Vec<String> = markov.take(12).collect();
    println!("{}", sentence.join(" "));
}