best_projection.rs

use std::collections::hash_map::Entry::*;
use std::collections::HashMap;
use std::convert::TryFrom;

use crate::HashedSlice;
use crate::Tokenization;

#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Copy, Hash)]
struct Coord {
    next_lcs: usize,
    next_seq: usize,
}

#[derive(Debug)]
pub struct NormalizationInput<'a> {
    pub seq: &'a Tokenization<'a>,
    pub lcs: &'a Tokenization<'a>,
}

#[derive(Debug)]
struct Context<'a> {
    seq_index: HashMap<HashedSlice<'a>, Vec<usize>>,
}

impl<'a> Context<'a> {
    fn new(input: &'a NormalizationInput) -> Self {
        let mut seq_index = HashMap::new();
        for i in 0..to_isize(input.lcs.nb_tokens()) {
            match seq_index.entry(input.lcs.nth_token(i)) {
                Occupied(_) => (),
                Vacant(e) => {
                    e.insert(vec![]);
                }
            }
        }
        for i in 0..input.seq.nb_tokens() {
            match seq_index.entry(input.seq.nth_token(to_isize(i))) {
                Occupied(e) => {
                    e.into_mut().push(i);
                }
                Vacant(_) => (),
            }
        }
        Context { seq_index }
    }

    fn get_indexes(&self, tok: &HashedSlice<'a>, min_value: usize) -> &[usize] {
        match self.seq_index.get(tok) {
            Some(values) => {
                let min_idx = match values.binary_search(&min_value) {
                    Ok(i) | Err(i) => i,
                };
                &values[min_idx..]
            }
            None => &[],
        }
    }
}

#[derive(Debug)]
pub struct NormalizationResult {
    pub path: Vec<isize>,
    pub starts_with_shared: bool,
}

pub fn normalize_partition(input: &NormalizationInput) -> NormalizationResult {
    let context = Context::new(input);
    let root = Coord {
        next_lcs: 0,
        next_seq: 0,
    };
    let target = Coord {
        next_lcs: input.lcs.nb_tokens(),
        next_seq: input.seq.nb_tokens(),
    };
    let mut frontier = vec![root];
    let mut prev = HashMap::new();
    let mut found_seq = None;
    while !frontier.is_empty() && found_seq == None {
        let mut new_frontier = vec![];
        for coord in frontier {
            if coord.next_lcs == target.next_lcs {
                found_seq = Some(coord.next_seq);
                if coord.next_seq == target.next_seq {
                    break;
                } else {
                    // TODO do something more clever here
                    continue;
                }
            }
            let start_lcs = coord.next_lcs;
            let lcs_len = input.lcs.nb_tokens() - start_lcs;
            let mut last_enqueued_snake_len = 0;
            for start_seq in context.get_indexes(
                &input.lcs.nth_token(coord.next_lcs as isize),
                coord.next_seq,
            ) {
                if start_seq + lcs_len > input.seq.nb_tokens() {
                    break;
                }
                let mut snake_len = 1;
                let max_smake_len = lcs_len.min(input.seq.nb_tokens() - start_seq);
                while snake_len < max_smake_len {
                    let lcs_tok = input.lcs.nth_token((start_lcs + snake_len) as isize);
                    let seq_tok = input.seq.nth_token((start_seq + snake_len) as isize);
                    if lcs_tok != seq_tok {
                        break;
                    }
                    snake_len += 1
                }
                let next_coord = Coord {
                    next_lcs: start_lcs + snake_len,
                    next_seq: start_seq + snake_len,
                };
                if last_enqueued_snake_len < snake_len || next_coord == target {
                    if next_coord.next_lcs == target.next_lcs {
                        if next_coord.next_seq == target.next_seq {
                            found_seq = Some(next_coord.next_seq);
                        } else if found_seq == None {
                            found_seq = Some(next_coord.next_seq);
                        }
                    }
                    match prev.entry(next_coord) {
                        Occupied(_) => continue,
                        Vacant(e) => e.insert(coord),
                    };
                    new_frontier.push(next_coord);
                    last_enqueued_snake_len = snake_len;
                }
            }
        }
        frontier = new_frontier;
    }

    let target = found_seq.map(|next_seq| Coord {
        next_lcs: input.lcs.nb_tokens(),
        next_seq,
    });
    let mut path = vec![];
    let mut starts_with_shared = false;
    let mut coord = target.as_ref();
    let mut seq = input.seq.nb_tokens();
    let mut lcs = input.lcs.nb_tokens();
    while let Some(&coord_content) = coord {
        let next_seq = coord_content.next_seq;
        let next_lcs = coord_content.next_lcs;
        let snake_len = lcs - next_lcs;
        push_if_not_last(&mut path, to_isize(seq - snake_len));
        starts_with_shared = !push_if_not_last(&mut path, to_isize(next_seq));

        coord = prev.get(&coord_content);

        seq = next_seq;
        lcs = next_lcs;
    }
    path.reverse();
    NormalizationResult {
        path,
        starts_with_shared,
    }
}

fn push_if_not_last(v: &mut Vec<isize>, val: isize) -> bool {
    let should_push = v.last() != Some(&val);
    if should_push {
        v.push(val);
    }
    should_push
}

fn to_isize(input: usize) -> isize {
    isize::try_from(input).unwrap()
}

pub struct Parts<'a> {
    index: usize,
    normalization: &'a Vec<isize>,
    tokenization: &'a Tokenization<'a>,
}

impl<'a> Parts<'a> {
    pub fn new(normalization: &'a NormalizationResult, tokenization: &'a Tokenization) -> Self {
        Parts {
            index: if normalization.starts_with_shared {
                0
            } else {
                1
            },
            normalization: &normalization.path,
            tokenization,
        }
    }
}

impl<'a> Iterator for Parts<'a> {
    type Item = (usize, usize);
    fn next(&mut self) -> Option<Self::Item> {
        if self.index + 1 < self.normalization.len() {
            let prev = self.normalization[self.index];
            let curr = self.normalization[self.index + 1];
            let from = self.tokenization.nth_span(prev).lo;
            let to = self.tokenization.nth_span(curr - 1).hi;
            self.index += 2;
            Some((from, to))
        } else {
            None
        }
    }
}