iads.rs

use super::{impl_ads::AdsTree, tree::SplittingTree};
use crate::{
    ads::traits::AdaptiveDistinguishingSequence,
    automatadefs::{
        mealy::{InputSymbol, Mealy, State},
        traits::{FiniteStateMachine, ObservationTree},
    },
    learner::apartness::states_are_apart,
    oracles::{
        equivalence::{
            automata::shortest_separating_sequence,
            equivalence_trait::{CounterExample, EquivalenceOracle, InternalEquivalenceOracle},
            incomplete::tree::TreeCons,
        },
        membership::Oracle as OQOracle,
    },
    util::{
        sequences::{FixedInfixGenerator, Order, RandomInfixGenerator},
        toolbox,
    },
};
use fnv::{FnvHashMap, FnvHashSet};
use itertools::Itertools;
use rand::{
    prelude::{SliceRandom, StdRng},
    Rng, SeedableRng,
};
use rand_distr::{Distribution, Uniform};
use rayon::iter::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator};
use rustc_hash::{FxHashMap, FxHashSet};
use std::{
    cell::RefCell,
    collections::{BTreeSet, VecDeque},
    rc::Rc,
    time::Instant,
};

pub struct IadsEO<'a, T> {
    oq_oracle: Rc<RefCell<OQOracle<'a, T>>>,
    lookahead: usize,
    seed: u64,
    ads_map: FnvHashMap<BTreeSet<State>, AdsTree>,
    frontier_states: FnvHashMap<State, Vec<State>>, // frontier -> basis
    // transitions from the basis using an input that goes to the frontier
    frontier_trans: FnvHashSet<(State, InputSymbol)>,
}

impl<'a, T> InternalEquivalenceOracle<'a, T> for IadsEO<'a, T>
where
    T: ObservationTree + Sync + Send,
{
    fn new(oq_oracle: Rc<RefCell<OQOracle<'a, T>>>, lookahead: usize, seed: u64) -> Self {
        Self {
            oq_oracle,
            lookahead,
            seed,
            ads_map: Default::default(),
            frontier_states: Default::default(),
            frontier_trans: Default::default(),
        }
    }
}

impl<'a, T> EquivalenceOracle<'a, T> for IadsEO<'a, T>
where
    T: ObservationTree + Send + Sync,
{
    type Options = ();

    fn get_counts(&mut self) -> (usize, usize) {
        RefCell::borrow_mut(&self.oq_oracle).get_counts()
    }

    fn find_counterexample(&mut self, hypothesis: &Mealy) -> CounterExample {
        self.find_counterexample_with_lookahead(hypothesis, self.lookahead)
    }

    fn find_counterexample_with_lookahead(
        &mut self,
        hyp: &Mealy,
        lookahead: usize,
    ) -> CounterExample {
        let hyp_states = hyp.states();
        let num_states = hyp_states.len();
        if num_states == 1 {
            return self.find_counterexample_initial(hyp);
        }
        self.ads_map.clear();
        self.frontier_trans = self.frontier_transitions(hyp);
        let access_map = self.access_map(hyp);
        // let access_map = random_access_map(hyp, self.seed);
        let rig = RandomInfixGenerator::new(self.seed, hyp.input_alphabet().len(), 0, lookahead);
        log::info!("Running AIC sequences.");
        let splitting_tree =
            TreeCons::construct(hyp, &hyp_states.iter().copied().collect_vec(), false);
        log::info!("Finished making ST-IADS.");
        let mut rng = StdRng::seed_from_u64(self.seed);
        let access_vec = access_map.into_values().collect_vec();

        let mut time_now = Instant::now();
        for (cnt, infix) in rig.generate().enumerate() {
            let st_num = rng.gen_range(0..num_states);
            let access = access_vec.get(st_num).expect("Safe");
            // let access = access_vec.choose(&mut rng).expect("Safe");
            let prefix = toolbox::concat_slices(&[access, &infix]);
            let ce = self.run_ads_test(&prefix, hyp, &splitting_tree, rng.gen());
            if ce.is_some() {
                let mut rng: StdRng = SeedableRng::seed_from_u64(self.seed);
                self.seed = rng.gen::<u64>();
                return ce;
            }
            if cnt % 10000 == 0 {
                let time_spent = time_now.elapsed();
                time_now = Instant::now();
                println!("Spent : {} ms", time_spent.as_millis());
            }
        }
        unreachable!("This process never ends until a CE is found.");
    }
}

/// Checks whether the hypothesis and the observation tree agree after a given prefix.
#[must_use]
fn _check_consistency_from<T: ObservationTree + Sync>(
    prefix: &[InputSymbol],
    o_tree: &T,
    hypothesis: &Mealy,
) -> CounterExample {
    let tree_start = o_tree.get_succ(State::new(0), prefix).expect("Safe");
    let hyp_start = hypothesis.trace(prefix).0;
    let input_size = hypothesis.input_alphabet().len();
    // let tree_start = State::new(0);
    // let hyp_start = State::new(0);
    let mut queue = VecDeque::from([(tree_start, hyp_start)]);
    while let Some((q, r)) = queue.pop_front() {
        for i in toolbox::inputs_iterator(input_size) {
            let tree_hyp_next_pair = o_tree
                .get_out_succ(q, i)
                .zip(Some(hypothesis.step_from(r, i)));
            if let Some(((out_t, dest_t), (dest_h, out_h))) = tree_hyp_next_pair {
                if out_h == out_t {
                    queue.push_back((dest_t, dest_h));
                } else {
                    let mut inputs = o_tree.get_access_seq(q);
                    inputs.push(i);
                    let outputs = o_tree.get_observation(None, &inputs).expect("Safe");
                    return Some((inputs, outputs));
                }
            }
        }
    }
    None
}

impl<'a, T> IadsEO<'a, T>
where
    T: ObservationTree + Send + Sync,
{
    pub fn load_frontier_states(&mut self, frontier: FnvHashMap<State, Vec<State>>) {
        self.frontier_states = frontier;
    }

    #[must_use]
    fn run_ads_test(
        &mut self,
        prefix: &[InputSymbol],
        hyp: &Mealy,
        splitting_tree: &SplittingTree,
        seed: u64,
    ) -> CounterExample {
        let hyp_dest = hyp.trace(prefix).0;
        // State hyp_dest must be separated from the following states.
        let states_to_sep = self.must_sep_from(prefix, hyp);
        log::debug!("States to separate: {:?}", &states_to_sep);
        match &states_to_sep[..] {
            [] => {
                unreachable!("A state cannot be distinguished from itself!");
            }
            // Already separated from everything, no need to run any test.
            [_q_p] => {
                // let temp = RefCell::borrow(&self.oq_oracle);
                // let o_tree = temp.borrow_tree();
                // assert!(_check_consistency_from(&[], o_tree, hyp).is_none());
                None
            }
            [q, q_p] => {
                // We need to split "q" from "q'", so we just run an output query.
                // let q = hyp_dest;
                let sep_word = shortest_separating_sequence(hyp, hyp, *q, *q_p)
                    .expect("Safe")
                    .0;
                let test_input = toolbox::concat_slices(&[prefix, &sep_word]);
                self.run_test(hyp, &test_input)
            }
            _ => {
                // There are at least three states (in total) we need to split,
                // so it makes sense to use an IADS now.
                // states_to_sep.push(hyp_dest);
                let states = states_to_sep.iter().copied().collect();
                let st_ads = self.ads_map.entry(states).or_insert_with(|| {
                    AdsTree::construct(hyp, splitting_tree, &states_to_sep, seed)
                });
                st_ads.randomise_root(hyp_dest);
                log::debug!("IADS:{:?}", &st_ads);
                let ce = RefCell::borrow_mut(&self.oq_oracle).ads_equiv_test(st_ads, prefix, hyp);
                st_ads.reset_to_root();
                ce
            }
        }
    }

    /// Returns the states which `tree_state` is not separated from in the observation tree.
    #[must_use]
    fn must_sep_from(&self, prefix: &[InputSymbol], hyp: &Mealy) -> Vec<State> {
        let temp = RefCell::borrow(&self.oq_oracle);
        let tree = temp.borrow_tree();
        let tree_dest = tree.get_succ(State::new(0), prefix);
        // let hyp_dest = hyp.trace_from(State::new(0), prefix).0;
        let mut other_states = hyp.states().into_iter().collect();
        // .into_par_iter()
        // .filter(|&s| s != hyp_dest)
        // .collect();
        let ts = match tree_dest {
            Some(ts) => ts,
            None => return other_states,
        };
        let are_not_sep = |s1: &State| !states_are_apart(tree, *s1, ts);
        other_states.retain(are_not_sep);
        other_states
    }
    /// Specialised case for when the hypothesis is a single-state FSM.
    #[must_use]
    fn find_counterexample_initial(&mut self, hyp: &Mealy) -> CounterExample {
        // When the hypothesis has size one, there's no point in testing A.C,
        // as we will have no characterising set.
        // A similar argument works for A.I.C, as C is empty.
        // A will contain only the empty sequence,
        // since we are already in the initial state;
        // Therefore, the only thing left is I<=(n+1)
        let input_alphabet = hyp.input_alphabet().into_iter().sorted().collect_vec();
        let fig = FixedInfixGenerator::new(input_alphabet, self.lookahead + 1, Order::ASC);
        for seq in fig.generate() {
            let res = self.run_test(hyp, &seq);
            if res.is_some() {
                return res;
            }
        }
        None
    }

    /// Run a single test sequence and return whether it is a CE or not.
    #[inline]
    #[must_use]
    fn run_test(&mut self, hyp: &Mealy, input_seq: &[InputSymbol]) -> CounterExample {
        let hyp_output = hyp.trace(input_seq).1.to_vec();
        let sut_output = RefCell::borrow_mut(&self.oq_oracle).output_query(input_seq);
        if hyp_output == sut_output {
            return None;
        }
        Some((input_seq.to_vec(), sut_output.to_vec()))
    }

    #[allow(unused)]
    fn num_inputs_for_frontier(
        &self,
        access: &[InputSymbol],
        infix: &[InputSymbol],
        hyp: &Mealy,
    ) -> Option<usize> {
        let mut curr = hyp.trace(access).0;
        for (cnt, input) in infix.iter().copied().enumerate() {
            let cnt = cnt + 1;
            if self.frontier_trans.contains(&(curr, input)) {
                return Some(cnt);
            }
            curr = hyp.step_from(curr, input).0;
        }
        None
        // unreachable!("We should hit a frontier state from the basis states!");
    }

    /// Given a hypothesis and the frontier-basis map, we compute the transitions
    /// that would lead to a frontier state.
    fn frontier_transitions(&mut self, hyp: &Mealy) -> FnvHashSet<(State, InputSymbol)> {
        let temp = RefCell::borrow(&self.oq_oracle);
        let o_tree = temp.borrow_tree();
        hyp.states()
            .into_iter()
            .cartesian_product(hyp.input_alphabet().into_iter())
            .filter(|(s, i)| {
                let dest = o_tree.get_succ(*s, &[*i]).expect("Safe");
                self.frontier_states.contains_key(&dest)
            })
            .collect()
    }

    #[allow(unused)]
    pub(crate) fn access_map(&self, hyp: &Mealy) -> FxHashMap<State, Vec<InputSymbol>> {
        let temp = RefCell::borrow(&self.oq_oracle);
        let o_tree = temp.borrow_tree();
        hyp.states()
            .into_par_iter()
            .map(|s| (s, o_tree.get_access_seq(s)))
            .collect()
    }
}

#[allow(unused)]
fn get_successors(fsm: &Mealy, input: InputSymbol, block: &FxHashSet<State>) -> FxHashSet<State> {
    block
        .par_iter()
        .copied()
        .map(|s| fsm.step_from(s, input).0)
        .collect()
}

/// Randomly compute access sequencs for states in the hypothesis.
///
/// Equivalent to the ``buggy`` mode of H-ADS.
#[must_use]
#[allow(dead_code)]
fn random_access_map(fsm: &Mealy, seed: u64) -> FxHashMap<State, Vec<InputSymbol>> {
    let mut ret = FxHashMap::default();
    let mut rng: StdRng = SeedableRng::seed_from_u64(seed);
    let mut work_list = VecDeque::from([State::new(0)]);
    ret.insert(State::new(0), vec![]);
    let uni_dist = Uniform::new_inclusive(0.0, 1.0);
    while !work_list.is_empty() {
        let state = {
            let sample = uni_dist.sample(&mut rng);
            let scaled_sample = {
                let scaled = (work_list.len() as f64 * sample).floor() as usize;
                scaled.clamp(0, work_list.len() - 1)
            };
            work_list.remove(scaled_sample).expect("Safe")
        };
        let mut inputs = fsm.input_alphabet().into_iter().collect_vec();
        inputs.shuffle(&mut rng);
        for i in inputs {
            let dest = fsm.step_from(state, i).0;
            if ret.contains_key(&dest) {
                continue;
            }
            let dest_acc = {
                let mut x = ret.get(&state).expect("Safe").clone();
                x.push(i);
                x
            };
            work_list.push_back(dest);
            ret.insert(dest, dest_acc);
        }
    }
    ret
}

#[cfg(test)]
mod tests {
    use rstest::rstest;

    use super::random_access_map;
    use crate::{
        automatadefs::traits::FiniteStateMachine, util::parsers::machine::read_mealy_from_file,
    };
    use std::collections::HashSet;

    #[rstest]
    #[case("tests/hypothesis_23.dot")]
    /// Check if all the number of access sequences found is equal to
    /// the number of states of an FSM.
    fn all_states_found(#[case] file_name: &str) {
        let (fsm, _input_map, _) = read_mealy_from_file(file_name);
        let access_states: HashSet<_, _> = random_access_map(&fsm, 10).into_keys().collect();
        assert_eq!(fsm.states(), access_states)
    }
}