path: root/src/index.rs

use std::collections::{HashSet, HashMap};
use std::fs::File;
use std::io::{BufWriter, BufReader, BufRead, Write};
use std::sync::mpsc::{channel, Sender};
use std::time::Duration;
use walkdir::*;
use std::thread;
use std::option::Option::None;
use crate::vector::FileVector;
use crate::dictionary::Dictionary;
use crate::filecache::FileCache;
use crate::searchresult::SearchResult;
use crate::filecounter::filecount;
use crate::text;
use crate::splitter;
use crate::vector;

/// Represents a Index which is ether generated
/// or read from a file.
#[derive(Clone, Debug)]
pub struct Index {
    dictionary : Dictionary,
    filecache : Vec<FileCache>,
}

impl Default for Index {
    fn default() -> Self {
        Self::empty()
    }
}

#[derive(Clone, Debug, Default, Copy)]
pub enum GenState {
    #[default]
    Fetching,
    Parsing,
    Merging
}

impl Index {
    pub fn empty() -> Self {
        Self {
            dictionary : Dictionary::new(),
            filecache : Vec::new()
        }
    }

    pub fn generate(input_path : &str, callback : impl Fn(GenState, u8)) -> Self {
        let mut nof = 1;
        let mut counter = 0;
        let mut crawler_handles = Vec::new();
        let num_threads = thread::available_parallelism().unwrap().get().min(4);
        let mut tx_vec : Vec<Sender<String>> = Vec::new();
        let mut indexes = Vec::new();

        thread::scope(|s| {
            let mut nof_handle : Option<_> = Some(s.spawn(|| filecount(input_path)));
            let (status_tx, status_rx) = channel();

            for _ in 0..num_threads {
                let (tx, rx) = channel();
                tx_vec.push(tx);
                let status_tx = status_tx.clone();
                crawler_handles.push(thread::spawn(move || {
                    let mut dict = Dictionary::new();
                    let mut filecache : Vec<FileCache> = Vec::new();

                    loop {
                        let path = rx.recv().unwrap();
                        if path.is_empty() {
                            return Self {
                                dictionary : dict,
                                filecache
                            }
                        }

                        let content : String = text::extract_text(path.as_str());

                        let _ = status_tx.send(());

                        if content.is_empty() {
                            continue;
                        }

                        let words : Vec<String> = splitter::split_to_words(content);
                        let fv = dict.insert_words_and_vectorize_word_list(&words.iter().collect());
                        filecache.push(FileCache {
                            path,
                            vector : fv
                        });
                    }
                }));
            }

            let mut next_crawler = 0;
            let mut last_p = 0;

            for entry in WalkDir::new(input_path)
                .into_iter()
                .filter_map(|e| e.ok()) {
                counter += 1;
                if entry.path().is_file() {
                    tx_vec[next_crawler].send(entry.path().to_str().unwrap().to_string()).ok();
                    next_crawler += 1;
                    if next_crawler == num_threads {
                        next_crawler = 0;
                    }

                    match nof_handle {
                        Some(t) => {
                            if t.is_finished() {
                                nof = t.join().unwrap();
                                nof_handle = None;
                            } else {
                                nof_handle = Some(t);
                            }
                        }
                        None => {
                            // Make sure that we only push a update
                            // if there is a visual change to the number
                            // because updating the screen takes a lot
                            // of time.
                            let p = counter * 100 / nof;
                            if p != last_p {
                                callback(GenState::Fetching, p as u8);
                                last_p = p;
                            }
                        }
                    }
                }
            }

            let join_handle = s.spawn(|| {
                for (i, handle) in crawler_handles.into_iter().enumerate() {
                    tx_vec[i].send(String::new()).ok();
                    indexes.push(handle.join().unwrap());
                }
            });


            let mut i = 0;
            let mut last_p = 0;
            while !join_handle.is_finished() {
                if status_rx.recv_timeout(Duration::from_millis(20)).is_ok() {
                    i += 1;
                    let p = i * 100 / nof;
                    if p != last_p {
                        callback(GenState::Parsing, p as u8);
                        last_p = p;
                    }
                }
            }

            callback(GenState::Parsing, 100);

            join_handle.join().ok();
        });

        Index::merge(indexes.iter().collect(), |p| { callback(GenState::Merging, p) })
    }

    pub fn from_file(path : &String) -> Self {
        let index_file = File::open(path).expect("could not open index file");
        let reader = BufReader::new(index_file);
        let mut filecache : Vec<FileCache> = Vec::new();
        let mut dict = Dictionary::new();

        for line in reader.lines() {
            let l = line.unwrap();
            if l.starts_with('#') {
                dict = Dictionary::from_line(l.strip_prefix('#').unwrap());
            } else {
                filecache.push(FileCache::from_line(l));
            }
        }

        Self {
            dictionary : dict,
            filecache
        }
    }

    fn merge_into(&mut self, other : &Index) {
        let mut dict = self.dictionary.clone();
        thread::scope(|s| {
            let mut a_hash : HashSet<&FileCache> = HashSet::new();
            let mut diff : Vec<&FileCache> = Vec::new();

            let converter_handle = s.spawn(|| {
                let mut b_id_to_word : HashMap<u64, &String> = HashMap::new();

                for (value, id) in other.dictionary.iter() {
                    b_id_to_word.insert(*id, value);
                }
                b_id_to_word
            });

            let dict_handle = s.spawn(|| {
                for (word, _) in other.dictionary.iter() {
                   dict.set(word);
                }
                dict
            });

            for file in self.filecache.iter() {
                a_hash.insert(file);
            }

            for file in other.filecache.iter() {
                if !a_hash.contains(file) {
                    diff.push(file);
                }
            }

            let b_id_to_word = converter_handle.join().unwrap();
            self.dictionary = dict_handle.join().unwrap();

            for file in diff {
                let mut words : Vec<&String> = Vec::new();

                for (word_id, i) in file.vector.iter() {
                    for _ in 0..*i {
                        words.push(b_id_to_word.get(word_id).unwrap());
                    }
                }

                self.filecache.push(FileCache {
                    path : file.path.clone(),
                    vector: self.dictionary.vectorize_word_list(&words)
                });
            }
        });
    }

    pub fn merge(mut indexes : Vec<&Index>, callback : impl Fn(u8)) -> Self {
        let max = indexes.len();

        indexes.sort_by(|a, b| a.filecache.len().cmp(&b.filecache.len()));
        let mut merged_index : Index = indexes.pop().unwrap().clone();

        for (i, index) in indexes.into_iter().enumerate() {
            callback((i * 100 / max) as u8);
            merged_index.merge_into(index);
        }
        callback(100);
        merged_index
    }

    pub fn search(&self, search_args : Vec<String>) -> Vec<SearchResult> {
        let mut v : FileVector = FileVector::new();
        let mut opt : FileVector = FileVector::new();

        for arg in search_args {
            let a = arg.trim_start_matches("+");
            if let Some(value) = self.dictionary.get(&a.to_string()) {
                if arg.chars().nth(0).unwrap() == '+' {
                    opt.insert(*value, 1);
                } else {
                    v.insert(*value, 1);
                }
            }
        }

        let mut results : Vec<SearchResult> = Vec::new();

        for filecache in self.filecache.iter() {
            let mut r = SearchResult { priority : 0, path : filecache.path.clone() };
            r.priority = vector::match_vector(&v, &filecache.vector);
            if r.priority > 0 {
                r.priority += vector::scalar_product(&opt, &filecache.vector);
                results.push(r);
            }
        }
        results.sort_by(|a, b| b.priority.cmp(&a.priority));
        results
    }

    pub fn save(&self, path: String) {
        thread::scope(|s| {
            let dict_list_handle = s.spawn(|| {
                self.dictionary.to_list().join(",")
            });
            let mut output : String = self.filecache.iter().map(|c| format!("{}, {}\n", c.path.replace(',', "\0"), c.vector.stringify())).collect();
            output += "#";
            output += dict_list_handle.join().unwrap().as_str();
            output += "\n";

            let index_file = File::create(path).expect("could not open output file");
            let mut file = BufWriter::new(index_file);
            file.write_all(output.as_bytes()).expect("could not write");
            file.flush().ok();
        });
    }

    pub fn num_files(&self) -> usize {
        self.filecache.len()
    }
}