l_sharp.rs

use crate::{
    ads::{
        lee_yannakakis::adaptive_distinguishing_seq::AdsTree as LyAds,
        traits::AdaptiveDistinguishingSequence, tree::threaded::Ads as OtAds,
    },
    automatadefs::{
        mealy::{InputSymbol, Mealy, MealyBuilder, OutputSymbol, State},
        traits::{FiniteStateMachine, ObservationTree},
    },
    learner::apartness::{compute_witness, states_are_apart},
    oracles::{equivalence::equivalence_trait::CounterExample, membership::Oracle as OutputOracle},
    util::{parsers::logs::Logs, toolbox},
};
use itertools::Itertools;
use rayon::iter::{IntoParallelRefIterator, IntoParallelRefMutIterator, ParallelIterator};
use std::{
    cell::RefCell,
    collections::{HashMap, HashSet, VecDeque},
    fmt::Debug,
    hash::{BuildHasher, BuildHasherDefault},
    rc::Rc,
};

/// Struct for the learner.
pub struct Lsharp<'a, Tree, S> {
    /// OutputQuery Oracle, needed by the Learner.
    oq_oracle: Rc<RefCell<OutputOracle<'a, Tree>>>,
    /// Size of the input alphabet.
    input_size: usize,
    /// List of basis states.
    basis: Vec<State>,
    /// Frontier state &mapsto; set of basis candidate states.
    frontier_to_basis_map: HashMap<State, Vec<State>, S>,
    /// Use the Lee-Yannakakis optimisation.
    use_ly_ads: bool,
    /// Num of `consistency` CEs generated.
    consistency_ce_count: usize,
}

impl<'a, Tree, S> Lsharp<'a, Tree, S>
where
    Tree: ObservationTree + Sync + Debug + Send,
    S: BuildHasher + Default + Send + Sync,
{
    /// Constructor.
    ///
    /// We require an oracle for asking output queries, (wrapped in an `Rc<RefCell<>>`),
    /// the input alphabet of the FSM to be learned, and a flag for enabling the use of
    /// the Lee-Yannakakis optimisation.
    pub fn new(
        oracle: Rc<RefCell<OutputOracle<'a, Tree>>>,
        input_alphabet: &HashSet<InputSymbol, S>,
        use_ly_ads: bool,
    ) -> Self {
        Self {
            oq_oracle: oracle,
            input_size: input_alphabet.len(),
            basis: vec![State::new(0)],
            frontier_to_basis_map: HashMap::default(),
            use_ly_ads,
            consistency_ce_count: 0,
        }
    }

    /// Process the given counterexample in order to find the new state.
    ///
    /// We require the hypothesis in order to find the prefix of input sequence
    /// to the point that the hypothesis and observation tree agree
    /// (i.e., have the same outputs).
    pub fn process_cex(&mut self, counter_example: CounterExample, hypothesis: &Mealy) {
        let (ce_input, ce_output) = {
            let unwrapped_ce = counter_example.expect("Safe");
            (unwrapped_ce.0, unwrapped_ce.1)
        };
        let _ = RefCell::borrow_mut(&self.oq_oracle).add_observation(&ce_input, &ce_output);
        let (_, hyp_output) = hypothesis.trace(&ce_input);
        let prefix_idx = get_ce_prefix_idx(&ce_output, &hyp_output);
        self.process_bin_search(
            &ce_input[..prefix_idx],
            &ce_output[..prefix_idx],
            hypothesis,
        );
    }

    pub fn frontier_basis_map(&self) -> &HashMap<State, Vec<State>, S> {
        &self.frontier_to_basis_map
    }

    #[allow(clippy::many_single_char_names)]
    fn process_bin_search(
        &mut self,
        ce_input: &[InputSymbol],
        ce_output: &[OutputSymbol],
        hypothesis: &Mealy,
    ) {
        let q: State = hypothesis.trace(ce_input).0;
        let r: State = RefCell::borrow(&self.oq_oracle)
            .borrow_tree()
            .get_succ(State::new(0), ce_input)
            .expect("Safe");
        self.update_frontier_and_basis();
        if self.frontier_to_basis_map.contains_key(&r) || self.basis.contains(&r) {
            return;
        }
        let temp = RefCell::borrow(&self.oq_oracle);
        let o_tree = temp.borrow_tree();
        // Get the prefix of the CE input that is leading to a frontier state.
        let idx = {
            let mut src = State::new(0);
            let mut i = 0;
            for (idx, input) in ce_input.iter().enumerate() {
                let dest = o_tree
                    .get_succ(src, &[*input])
                    .expect("Obs-Tree should have CE included!");
                if self.frontier_to_basis_map.contains_key(&dest) {
                    i = idx;
                    break;
                }
                src = dest;
            }
            i + 1
        };

        let eta = compute_witness(o_tree, q, r)
            .expect("Could not find witness when it should exist, during CE processing!");

        // We used idx inputs to get to the frontier.
        let rho = &ce_input[..idx];
        let x = rho.len();
        let y = ce_input.len();
        #[allow(clippy::cast_possible_truncation)]
        #[allow(clippy::cast_precision_loss)]
        #[allow(clippy::cast_sign_loss)]
        let h = ((x as f32 + y as f32) / 2_f32).floor() as usize;

        let sigma_1 = &ce_input[..h];
        let sigma_2 = &ce_input[h..];
        let q_p = hypothesis.trace(sigma_1).0;
        let r_p = o_tree.get_succ(State::new(0), sigma_1).unwrap();
        let access_q_p = o_tree.get_access_seq(q_p);
        let output_query = toolbox::concat_slices(&[&access_q_p, sigma_2, &eta]);
        drop(temp);
        let sul_response = RefCell::borrow_mut(&self.oq_oracle).output_query(&output_query);
        let temp = RefCell::borrow(&self.oq_oracle);
        let o_tree = temp.borrow_tree();
        let wit = compute_witness(o_tree, q_p, r_p);
        drop(temp);
        if wit.is_some() {
            self.process_bin_search(sigma_1, &ce_output[..sigma_1.len()], hypothesis);
        } else {
            let new_inputs = toolbox::concat_slices(&[&access_q_p, sigma_2]);
            self.process_bin_search(&new_inputs, &sul_response[..new_inputs.len()], hypothesis);
        }
    }

    /**
    We loop through the following three rules until the observation tree is adequate:
    1. Explore the frontier (implemented [here](crate::oracles::membership::Oracle::explore_frontier)).
    2. Identify the frontier (implemented [here](crate::oracles::membership::Oracle::identify_frontier)).
    3. Promote a frontier state if possible (internal method of this struct).
    */
    fn make_obs_tree_adequate(&mut self) {
        log::info!(
            "Size of tree after EQ: {}",
            RefCell::borrow(&self.oq_oracle).borrow_tree().size()
        );
        loop {
            log::debug!("Applying rule 2");
            let new_frontier = RefCell::borrow_mut(&self.oq_oracle).explore_frontier(&self.basis);
            for (fs, cands) in new_frontier {
                log::debug!("Adding frontier {:?} with candidates {:?}", fs, cands);
                self.frontier_to_basis_map.insert(fs, cands);
            }

            // Identify the frontier states
            log::debug!("Applying rule 3");
            for (fs, candidates) in self
                .frontier_to_basis_map
                .iter_mut()
                .filter(|(_fs, c)| c.len() > 1)
            {
                log::debug!("Identifying {:?} amongst {:?}", fs, candidates);
                RefCell::borrow_mut(&self.oq_oracle).identify_frontier(*fs, candidates);
            }

            // Promote the frontier states if they're isolated.
            self.promote_frontier_state();

            if self.tree_is_adequate() {
                break;
            }
        }
    }

    fn promote_frontier_state(&mut self) {
        let new_bs = self
            .frontier_to_basis_map
            .iter()
            .find(|(_fs, bs_not_sep)| bs_not_sep.is_empty())
            .map(|(&fs, _)| fs);
        let bs = match new_bs {
            Some(bs) => bs,
            _ => return,
        };
        log::debug!("Applying rule 4");
        log::debug!("Promoting {:?}", bs);
        self.basis.push(bs);
        self.frontier_to_basis_map.remove(&bs);
        let temp = RefCell::borrow(&self.oq_oracle);
        let o_tree = temp.borrow_tree();
        // The new basis state is a candidate for the frontier
        // states it is not apart from.
        self.frontier_to_basis_map
            .par_iter_mut()
            .filter(|(&fs, _cands)| !states_are_apart(o_tree, fs, bs))
            .for_each(|(_fs, cands)| {
                cands.push(bs);
            });
    }

    /// If each frontier state is identified, a tree is adequate.
    ///
    /// *We assume that all basis states have been explored*.
    fn tree_is_adequate(&mut self) -> bool {
        self.check_frontier_consistency();
        log::debug!("Basis set: {:?}", self.basis);
        log::debug!("Frontier->Basis mapping: {:?}", self.frontier_to_basis_map);
        let is_singleton = |c: &Vec<State>| c.len() == 1;
        let all_frontier_identified = self.frontier_to_basis_map.values().all(is_singleton);
        let input_alphabet = toolbox::inputs_iterator(self.input_size).collect_vec();
        let temp = RefCell::borrow(&self.oq_oracle);
        let o_tree = temp.borrow_tree();
        let all_basis_extended = self
            .basis
            .iter()
            .copied()
            .cartesian_product(input_alphabet)
            .all(|(q, i)| o_tree.get_out(q, i).is_some());
        all_frontier_identified && all_basis_extended
    }

    fn update_frontier_and_basis(&mut self) {
        {
            let temp = RefCell::borrow(&self.oq_oracle);
            let o_tree = temp.borrow_tree();
            self.frontier_to_basis_map
                .par_iter_mut()
                .for_each(|(&fs, cands)| cands.retain(|&bs| !states_are_apart(o_tree, fs, bs)));
        }
        self.promote_frontier_state();
        self.check_frontier_consistency();
        {
            let temp = RefCell::borrow(&self.oq_oracle);
            let o_tree = temp.borrow_tree();
            self.frontier_to_basis_map
                .par_iter_mut()
                .for_each(|(&fs, cands)| cands.retain(|&bs| !states_are_apart(o_tree, fs, bs)));
        }
    }

    /// Make the observation tree adequate and return the hypothesis.
    pub fn build_hypothesis(&mut self) -> Mealy {
        loop {
            self.make_obs_tree_adequate();
            let hyp = self.construct_hypothesis();
            if let Some(ce) = self.check_consistency(&hyp) {
                self.consistency_ce_count += 1;
                self.process_cex(Some(ce), &hyp);
            } else {
                if self.use_ly_ads {
                    let mut ly_ads = LyAds::new(&hyp);
                    let temp = RefCell::borrow(&self.oq_oracle);
                    let tree = temp.borrow_tree();
                    let hyp_states: Vec<_> = hyp.states().into_iter().collect();
                    let ot_ads = OtAds::new(tree, &hyp_states);
                    drop(temp);
                    if ly_ads.identification_power() > ot_ads.identification_power() * 10f32 {
                        self.run_ly_ads(&mut ly_ads);
                        continue;
                    }
                }
                println!(
                    "Size of the tree: {}",
                    RefCell::borrow(&self.oq_oracle).borrow_tree().size()
                );
                return hyp;
            }
        }
    }

    fn construct_hypothesis(&mut self) -> Mealy {
        let temp = RefCell::borrow(&self.oq_oracle);
        let o_tree = temp.borrow_tree();
        let basis = self.basis.iter().copied().collect::<HashSet<_, _>>();
        let hyp_input_alphabet = toolbox::inputs_iterator(self.input_size).collect_vec();
        let func_size = basis.len() * hyp_input_alphabet.len();
        let mut hyp_output_alphabet = HashSet::<_, _>::default();
        let mut hyp_trans_fn =
            HashMap::with_capacity_and_hasher(func_size, BuildHasherDefault::default());
        let mut hyp_output_fn =
            HashMap::with_capacity_and_hasher(func_size, BuildHasherDefault::default());
        log::debug!("Basis set: {:?}", self.basis);
        log::debug!("Frontier-Basis sep: {:?}", self.frontier_to_basis_map);

        for (q, i) in
            Itertools::cartesian_product(basis.iter().copied(), hyp_input_alphabet.iter().copied())
        {
            if o_tree.get_out(q, i).is_none() {
                println!("Something is wrong!");
            }
            let output = o_tree.get_out(q, i).expect("Safe");
            hyp_output_alphabet.insert(output);
            let succ = o_tree
                .get_succ(q, &[i])
                .expect("Frontier-Basis sep and observation tree are inconsistent.");
            let dest = {
                if basis.contains(&succ) {
                    succ
                } else {
                    // Destination is another basis state.
                    *self
                        .frontier_to_basis_map
                        .get(&succ)
                        .expect("Missing a frontier state")
                        .iter()
                        .exactly_one()
                        .expect("Multiple basis candidates for a single frontier state.")
                }
            };
            hyp_trans_fn.insert((q, i), dest);
            hyp_output_fn.insert((q, i), output);
        }
        log::info!("Size of the tree: {}", o_tree.size());

        MealyBuilder::default()
            .initial_state(State::new(0))
            .states(basis)
            .input_alphabet(hyp_input_alphabet)
            .output_alphabet(hyp_output_alphabet)
            .trans_function(hyp_trans_fn)
            .output_function(hyp_output_fn)
            .build()
            .expect("Safe")
    }

    #[allow(clippy::unused_self)]
    fn run_ly_ads(&mut self, _ly_ads: &mut LyAds) {
        todo!()
    }

    /**
    Initialise the observation tree.
    If logs are provided, we simply add them all to the tree at this point.
    */
    pub fn init_obs_tree(&mut self, logs: Option<Logs>) {
        if let Some(logs) = logs {
            let mut oracle_ref = self.oq_oracle.borrow_mut();
            for (input_seq, output_seq) in logs {
                let _ = oracle_ref.add_observation(&input_seq, &output_seq);
            }
        }
        self.make_obs_tree_adequate();
    }

    /// Check if the frontier is consistent with our data structures.
    /// CE processing may add a new frontier state that we miss.
    fn check_frontier_consistency(&mut self) {
        let basis_set = &self.basis;
        let input_alph = self.input_size;
        let temp = RefCell::borrow(&self.oq_oracle);
        let o_tree = temp.borrow_tree();
        let ia = toolbox::inputs_iterator(input_alph).collect_vec();
        let state_input_iterator = Itertools::cartesian_product(basis_set.iter().copied(), ia);
        let new_front = state_input_iterator
            .filter_map(|(bs, i)| o_tree.get_succ(bs, &[i]))
            .filter(|x| !self.basis.contains(x))
            .filter(|x| !self.frontier_to_basis_map.contains_key(x))
            .map(|fs| {
                let maps_to: Vec<_> = basis_set
                    .par_iter()
                    .copied()
                    .filter(|&s| !states_are_apart(o_tree, fs, s))
                    .collect();
                (fs, maps_to)
            })
            .collect::<Vec<_>>();
        for (fs, mut bs_set) in new_front {
            self.frontier_to_basis_map
                .entry(fs)
                .and_modify(|bs_vec| bs_set.drain(..).for_each(|x| bs_vec.push(x)))
                .or_insert_with(|| bs_set.into_iter().collect());
        }
    }

    /// Checks whether the hypothesis and the observation tree agree.
    #[must_use]
    pub fn check_consistency(&mut self, hypothesis: &Mealy) -> CounterExample {
        let temp = RefCell::borrow(&self.oq_oracle);
        let o_tree = temp.borrow_tree();
        let mut queue = VecDeque::<_>::new();
        queue.push_back((State::new(0), State::new(0)));
        while !queue.is_empty() {
            let (q, r) = queue.pop_front().expect("Always safe!");
            for i in toolbox::inputs_iterator(self.input_size) {
                if let Some(((out_tree, dest_tree), (dest_hyp, out_hyp))) = o_tree
                    .get_out_succ(q, i)
                    .zip(Some(hypothesis.step_from(r, i)))
                {
                    if out_hyp == out_tree {
                        queue.push_back((dest_tree, dest_hyp));
                    } 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
    }

    #[allow(clippy::unused_self)]
    pub fn get_ads_score(&self) -> f32 {
        0.0
    }

    /// Number of inputs and number of resets
    pub fn get_counts(&self) -> (usize, usize) {
        RefCell::borrow_mut(&self.oq_oracle).get_counts()
    }
}

fn get_ce_prefix_idx(ce_output: &[OutputSymbol], hyp_output: &[OutputSymbol]) -> usize {
    for (idx, (&ce_out, &hyp_out)) in ce_output.iter().zip(hyp_output.iter()).enumerate() {
        if ce_out != hyp_out {
            return idx;
        }
    }
    unreachable!();
}