use std::fs::{create_dir_all, read_to_string, File};
use std::io::Write;
use std::path::PathBuf;
use std::str::FromStr;
use homoglyph_core::confusable;
use homoglyph_core::domain::{SentenceDomain, WordDomain};
use homoglyph_core::glyph::EncodedGlyph;
use homoglyph_core::sentence::EncodedSentence;
use homoglyph_core::word::EncodedWord;
use tantivy::collector::TopDocs;
use tantivy::directory::{MmapDirectory, RamDirectory};
use tantivy::query::{Query, QueryParser};
use tantivy::schema::{STORED, TEXT};
use tantivy::{doc, IndexSettings, ReloadPolicy};
use tantivy::{schema::Schema, Index};
use crate::SearchEngine;
#[allow(dead_code)]
pub struct Tantivy {
index: Index,
schema: Schema,
queries_by_domain: Vec<Vec<Box<dyn Query>>>,
}
static HOMOGLYPHS_DIR: &'static str = "/tmp/homoglyphs";
static HOMOGLYPHS_STATE_FILE: &'static str = "/tmp/homoglyphs/.state";
fn create_schema() -> Schema {
let mut schema_builder = Schema::builder();
schema_builder.add_text_field("glyph", TEXT | STORED);
schema_builder.build()
}
fn build_mmap_path() -> PathBuf {
let mut path = PathBuf::new();
path.push(HOMOGLYPHS_DIR);
path
}
fn create_managed_index(schema: &Schema) -> Index {
let homoglyphs_dir = build_mmap_path();
if !homoglyphs_dir.exists() {
create_dir_all(&homoglyphs_dir).unwrap();
}
let mmap = MmapDirectory::open(homoglyphs_dir).unwrap();
let index = Index::open_or_create(mmap.clone(), schema.to_owned()).unwrap();
index
}
#[allow(dead_code)]
fn create_in_ram_index(ram_directory: RamDirectory, schema: &Schema) -> Index {
Index::create(
ram_directory.to_owned(),
schema.to_owned(),
IndexSettings::default(),
)
.unwrap()
}
#[allow(dead_code)]
impl SearchEngine for Tantivy {
fn init() -> Self {
let schema = create_schema();
let index = create_managed_index(&schema);
Self {
index,
schema,
queries_by_domain: Vec::<Vec<Box<dyn Query>>>::new(),
}
}
fn new() -> Self {
let mut tantivy = Self::init();
let contents = read_to_string(HOMOGLYPHS_STATE_FILE);
let indexed_state = match contents {
Ok(p) => p, Err(_) => "0".to_string(), };
if i32::from_str(indexed_state.as_str()).unwrap() == 0 {
tantivy.index();
}
tantivy
}
fn index(&mut self) {
let confusable = confusable::confusable::HEX_FILE;
let glyph = self.schema.get_field("glyph").unwrap();
let mut index_writer = self.index.writer(50_000_000).unwrap();
for line in confusable.lines() {
index_writer.add_document(doc!(glyph => line)).unwrap();
}
index_writer.commit().unwrap();
let mut file = File::create("/tmp/homoglyphs/.state").unwrap();
file.write_all(b"1").unwrap();
}
fn query(&mut self, mut sentence_enc: EncodedSentence) {
let glyph = self.schema.get_field("glyph").unwrap();
let query_parser = QueryParser::for_index(&self.index, vec![glyph]);
for c in sentence_enc.iter() {
let mut queries = Vec::new();
for cc in c.iter() {
let query = query_parser.parse_query(cc.0.as_str()).unwrap();
queries.push(query);
}
self.queries_by_domain.push(queries)
}
}
fn search(&mut self) -> SentenceDomain {
let reader = self
.index
.reader_builder()
.reload_policy(ReloadPolicy::OnCommit)
.try_into()
.unwrap();
let searcher = reader.searcher();
let mut sentence_domain: Vec<WordDomain> = Vec::new();
for queries in &self.queries_by_domain {
let mut world_domain: Vec<EncodedWord> = Vec::new();
for query in queries {
let mut confusable_word: Vec<EncodedWord> = Vec::new();
let top_docs = searcher.search(query, &TopDocs::with_limit(1)).unwrap();
let glyph = self.schema.get_field("glyph").unwrap();
for (_score, doc_address) in top_docs {
let retrieved_doc = searcher.doc(doc_address).unwrap();
let value = retrieved_doc.get_all(glyph).into_iter().next().unwrap();
let slice = value.as_text().unwrap().split_terminator(",");
let mut domain_words_enc: Vec<EncodedGlyph> = Vec::new();
for s in slice {
let confusable_glyph_enc = EncodedGlyph::from_str(s.trim()).unwrap();
domain_words_enc.push(confusable_glyph_enc);
}
let encode_word = EncodedWord::new(domain_words_enc);
confusable_word.push(encode_word);
}
world_domain.append(&mut confusable_word);
}
sentence_domain.push(WordDomain::new(world_domain));
}
SentenceDomain::new(sentence_domain)
}
}
#[cfg(test)]
mod tests {
use super::*;
use homoglyph_core::sentence::{EncodedSentence, Sentence};
#[test]
fn when_init_then_create_resource() {
let mut tantivy = Tantivy::init();
let sentence_dec = Sentence::from_str("ru best").unwrap();
let sentence_dec: EncodedSentence = EncodedSentence::from(sentence_dec);
tantivy.index();
tantivy.query(sentence_dec);
let domain: SentenceDomain = tantivy.search();
println!("{:#?}", domain);
}
}