learning.rs

use crate::automatadefs::mealy::shortest_separating_sequence;
use crate::learner::l_sharp::{Rule2, Rule3};
use crate::oracles::equivalence::{
    EquivalenceOracle, ExternalEquivalenceOracle, InternalEquivalenceOracle,
};
use crate::{
    automatadefs::{
        mealy::{InputSymbol, Mealy, OutputSymbol},
        FiniteStateMachine,
    },
    learner::{l_sharp::Lsharp, obs_tree::map_tree::ObsTree as ObsMapTree},
    oracles::{
        equivalence::{
            hads_tree::HadsTree,
            incomplete::iads::IadsEO,
            sep_seq::SequenceOracle,
            soucha::{ConfigBuilder as SouchaConfigBuilder, Oracle as SouchaOracle},
        },
        membership::Oracle as OQOracle,
    },
    sul::simulator::Simulator,
    util::{
        learning_config::{EqOracle, LearnResult},
        writers::overall as MealyWriter,
    },
};
use fnv::{FnvHashMap, FnvHashSet};
use rand::{prelude::StdRng, Rng, SeedableRng};
use std::collections::HashMap;
use std::{cell::RefCell, hash::BuildHasherDefault, rc::Rc, sync::Arc};

#[derive(derive_builder::Builder)]
pub struct Options {
    pub rule2_mode: Rule2,
    pub rule3_mode: Rule3,
    pub oracle_choice: EqOracle,
    pub seed: u64,
    pub extra_states: usize,
    // pub expected_rnd_length: usize,
    pub use_ly_ads: bool,
}

/// Main function that learns the FSM.
#[allow(clippy::type_complexity)]
#[allow(clippy::too_many_lines)]
#[must_use]
pub fn learn_fsm<S: ::std::hash::BuildHasher>(
    sul: &Mealy,
    input_map: &HashMap<String, InputSymbol, S>,
    output_map: &HashMap<String, OutputSymbol, S>,
    options: &Options,
    logs: Option<Vec<(Box<[InputSymbol]>, Box<[OutputSymbol]>)>>,
) -> LearnResult {
    let mealy_machine = Arc::new(sul.clone());
    log::info!("Qsize : {}", mealy_machine.as_ref().states().len());
    log::info!("Isize : {}", mealy_machine.as_ref().input_alphabet().len());
    let mut trial = Simulator::new(Arc::clone(&mealy_machine));
    let oq_oracle: Rc<RefCell<OQOracle<'_, ObsMapTree<BuildHasherDefault<rustc_hash::FxHasher>>>>> =
        Rc::new(RefCell::new(OQOracle::new(
            &mut trial,
            options.rule2_mode.clone(),
            options.rule3_mode.clone(),
        )));
    let rev_input_map: FnvHashMap<_, _> = input_map.iter().map(|(x, y)| (*y, x.clone())).collect();

    let rev_output_map: FnvHashMap<_, _> =
        output_map.iter().map(|(x, y)| (*y, x.clone())).collect();
    let mut rng = StdRng::seed_from_u64(options.seed);
    let mut hads_oracle_tree = HadsTree::new(
        Rc::clone(&oq_oracle),
        (),
        options.extra_states,
        rng.gen(),
        rev_input_map.clone(),
        rev_output_map.clone(),
    );

    let soucha_config = SouchaConfigBuilder::default()
        .exec_loc("/Users/bharat/research/FSMlib/fsm_lib")
        .lookahead(options.extra_states)
        .method(options.oracle_choice.clone().into())
        .build()
        .expect("Incorrect/Missing params from Soucha oracles.");
    let mut soucha_oracle = SouchaOracle::with_location(
        Rc::clone(&oq_oracle),
        soucha_config,
        "".to_string(),
        0,
        0,
        rev_input_map.clone(),
        rev_output_map.clone(),
    );
    let mut iads_oracle = IadsEO::new(Rc::clone(&oq_oracle), options.extra_states, rng.gen());
    let mut seq_oracle = SequenceOracle::new(Rc::clone(&oq_oracle), 10, rng.gen());
    let mut learner = Lsharp::new(
        oq_oracle,
        &mealy_machine
            .input_alphabet()
            .into_iter()
            .collect::<FnvHashSet<_>>(),
        options.use_ly_ads,
    );
    let mut idx = 1;
    let success: bool;
    let mut learn_inputs: usize = 0;
    let mut learn_resets: usize = 0;
    let mut test_inputs: usize = 0;
    let mut test_resets: usize = 0;
    let mut hyp_states;
    learner.init_obs_tree(logs);
    loop {
        println!("LOOP {}", idx);
        let hypothesis = learner.build_hypothesis();
        let (l_inputs, l_resets) = learner.get_counts();
        learn_inputs += l_inputs;
        learn_resets += l_resets;
        println!("Learning queries/symbols: {}/{}", l_resets, l_inputs,);
        hyp_states = hypothesis.states().len();
        println!(
            "Hypothesis size: {} and model size: {}",
            hyp_states,
            sul.states().len()
        );
        let writer_config = MealyWriter::WriteConfigBuilder::default()
            .file_name(Some(
                format!("./hypothesis/hypothesis_{}.dot", idx).to_string(),
            ))
            .format(MealyWriter::MealyEncoding::Dot)
            .build()
            .expect("Could not write Hypothesis file.");
        MealyWriter::write_machine(&writer_config, &hypothesis, &rev_input_map, &rev_output_map);
        let ideal_ce = shortest_separating_sequence(&sul.clone(), &hypothesis, None, None);
        if ideal_ce.is_none() {
            success = true;
            println!("Learning finished!");
            break;
        }
        println!("Searching for CEX...");
        let ce = match options.oracle_choice {
            EqOracle::SouchaH
            | EqOracle::SouchaHSI
            | EqOracle::SouchaSPY
            | EqOracle::SouchaSPYH => {
                let word = soucha_oracle.find_counterexample(&hypothesis);
                let (t_inputs, t_resets) = soucha_oracle.get_counts();
                test_inputs += t_inputs;
                test_resets += t_resets;
                println!("Testing queries/symbols: {}/{}", t_resets, t_inputs);
                word
            }
            EqOracle::SepSeq => {
                let word = seq_oracle.find_counterexample(&hypothesis);
                let (t_inputs, t_resets) = seq_oracle.get_counts();
                test_inputs += t_inputs;
                test_resets += t_resets;
                println!("Testing queries/symbols: {}/{}", t_resets, t_inputs);
                word
            }
            EqOracle::Internal => ideal_ce.clone(),
            EqOracle::HadsTree => {
                let word = hads_oracle_tree.find_counterexample(&hypothesis);
                let (t_inputs, t_resets) = hads_oracle_tree.get_counts();
                test_inputs += t_inputs;
                test_resets += t_resets;
                println!("Testing queries/symbols: {}/{}", t_resets, t_inputs);
                word
            }
            EqOracle::Iads => {
                let fs = learner.frontier_basis_map().clone();
                iads_oracle.load_frontier_states(fs);
                let word = iads_oracle.find_counterexample(&hypothesis);
                let (t_inputs, t_resets) = iads_oracle.get_counts();
                test_inputs += t_inputs;
                test_resets += t_resets;
                println!("Testing queries/symbols: {}/{}", t_resets, t_inputs);
                word
            }
        };
        let ce_len = ce.as_ref().map_or(0, |f| f.0.len());
        log::info!("EQStats:{{'Iter':{idx},'HypSize':{hyp_states},'CESize':{ce_len}}}");
        if ce.is_some() {
            println!("CE found, refining observation tree!");
            learner.process_cex(ce, &hypothesis);
        } else {
            println!("No CE found.");
            success = ideal_ce.is_none();
            println!("Model learned: {}", success);
            break;
        }
        idx += 1;
    }

    log::info!("MQ [Resets] : {learn_resets}");
    log::info!("MQ [Symbols] : {learn_inputs}");
    log::info!("EQ [Resets] : {test_resets}");
    log::info!("EQ [Symbols] : {test_inputs}");
    log::info!("Rounds : {idx}");
    log::info!("Searching for counterexample [ms] : {}", 0);
    log::info!("Learning [ms] : {}", 0);
    let equivalent = |learned| {
        if learned {
            "OK"
        } else {
            "NOK"
        }
    };
    log::info!("Equivalent : {}", equivalent(success));
    let ads_score = learner.get_ads_score();
    LearnResult {
        learn_inputs,
        learn_resets,
        eq_oracle: options.oracle_choice.clone(),
        rounds: idx,
        success,
        test_inputs,
        test_resets,
        num_inputs: rev_input_map.len(),
        num_states: hyp_states,
        ads_score,
    }
}