tantiny 0.2.2

data/src/helpers.rs

@@ -0,0 +1,200 @@
+use std::collections::HashMap;
+use rutie::{AnyException, Array, Exception, RString, Hash, Integer, Float, Boolean, Module};
+use tantivy::schema::{Field};
+use tantivy::tokenizer::Language;
+
+// Macro dependencies:
+pub(super) use paste::paste;
+pub(super) use rutie::{class, wrappable_struct, AnyObject, VerifiedObject, VM, Object, Class};
+
+pub(crate) fn namespace() -> Module {
+    Module::from_existing("Tantiny")
+}
+
+pub(crate) struct LanguageWrapper(pub(crate) Language);
+
+impl std::str::FromStr for LanguageWrapper {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "en" => Ok(LanguageWrapper(Language::English)),
+            "ar" => Ok(LanguageWrapper(Language::Arabic)),
+            "da" => Ok(LanguageWrapper(Language::Danish)),
+            "nl" => Ok(LanguageWrapper(Language::Dutch)),
+            "fi" => Ok(LanguageWrapper(Language::Finnish)),
+            "fr" => Ok(LanguageWrapper(Language::French)),
+            "de" => Ok(LanguageWrapper(Language::German)),
+            "el" => Ok(LanguageWrapper(Language::Greek)),
+            "hu" => Ok(LanguageWrapper(Language::Hungarian)),
+            "it" => Ok(LanguageWrapper(Language::Italian)),
+            "no" => Ok(LanguageWrapper(Language::Norwegian)),
+            "pt" => Ok(LanguageWrapper(Language::Portuguese)),
+            "ro" => Ok(LanguageWrapper(Language::Romanian)),
+            "ru" => Ok(LanguageWrapper(Language::Russian)),
+            "es" => Ok(LanguageWrapper(Language::Spanish)),
+            "sv" => Ok(LanguageWrapper(Language::Swedish)),
+            "ta" => Ok(LanguageWrapper(Language::Tamil)),
+            "tr" => Ok(LanguageWrapper(Language::Turkish)),
+            _ => Err(format!("Language '{}' is not supported.", s)),
+        }
+    }
+}
+
+pub(crate) trait TryUnwrap<T> {
+    fn try_unwrap(self) -> T;
+}
+
+macro_rules! primitive_try_unwrap_impl {
+    ( $ruby_type:ty, $type:ty ) => {
+        paste! {
+            impl TryUnwrap<$type> for $ruby_type {
+                fn try_unwrap(self) -> $type {
+                    self.[<to_ $type:lower>]()
+                }
+            }
+
+            impl TryUnwrap<$type> for AnyObject {
+                fn try_unwrap(self) -> $type {
+                    self.try_convert_to::<$ruby_type>()
+                        .try_unwrap()
+                        .[<to_ $type:lower>]()
+                }
+            }
+        }
+    };
+}
+
+primitive_try_unwrap_impl!(RString, String);
+primitive_try_unwrap_impl!(Integer, i64);
+primitive_try_unwrap_impl!(Float, f64);
+primitive_try_unwrap_impl!(Boolean, bool);
+
+impl<T> TryUnwrap<Vec<T>> for Array where
+    AnyObject: TryUnwrap<T>
+{
+    fn try_unwrap(self) -> Vec<T> {
+        let mut vec = Vec::new();
+
+        for elem in self {
+            vec.push(elem.try_unwrap());
+        }
+
+        vec
+    }
+}
+
+impl<K, V> TryUnwrap<HashMap<K, V>> for Hash where
+    AnyObject: TryUnwrap<K> + TryUnwrap<V>,
+    K: Eq + std::hash::Hash
+{
+    fn try_unwrap(self) -> HashMap<K, V> {
+        let mut hashmap = HashMap::new();
+
+        self.each(|key, value| {
+            hashmap.insert(key.try_unwrap(), value.try_unwrap());
+        });
+
+        hashmap
+    }
+}
+
+impl<T, E> TryUnwrap<T> for Result<T, E>
+where
+    E: ToString,
+{
+    fn try_unwrap(self) -> T {
+        self.map_err(|e| {
+            VM::raise_ex(AnyException::new(
+                "Tantiny::TantivyError",
+                Some(&e.to_string()),
+            ))
+        })
+        .unwrap()
+    }
+}
+
+impl TryUnwrap<Field> for Option<Field> {
+    fn try_unwrap(self) -> Field {
+        if let Some(value) = self {
+            value
+        } else {
+            VM::raise_ex(AnyException::new("Tantiny::UnknownField", None));
+
+            self.unwrap()
+        }
+    }
+}
+
+macro_rules! try_unwrap_params {
+    (
+        $param:ident: $type:ty,
+        $( $rest:tt )*
+    ) => {
+        let _tmp = $param.map_err(|e| $crate::helpers::VM::raise_ex(e)).unwrap();
+        let $param = <_ as $crate::helpers::TryUnwrap<$type>>::try_unwrap(_tmp);
+
+        try_unwrap_params!($($rest)*)
+    };
+    (
+        $param:ident,
+        $( $rest:tt )*
+    ) => {
+        let $param = $param.map_err(|e| $crate::helpers::VM::raise_ex(e)).unwrap();
+
+        try_unwrap_params!($($rest)*)
+    };
+
+    // Handle optional trailing commas.
+    ( $param:ident: $type:ty ) => {
+        try_unwrap_params!($param: $type,)
+    };
+    ( $param:ident ) => {
+        try_unwrap_params!($param,)
+    };
+
+    () => {}
+}
+
+pub(crate) use try_unwrap_params;
+
+macro_rules! scaffold {
+    ( $ruby_type:ident, $type:ty, $klass:literal ) => {
+        $crate::helpers::class!($ruby_type);
+
+        // There is a bug in Rutie which prevents using this macro
+        // by resolving it by a full path, so the only workaround is:
+        use crate::helpers::wrappable_struct;
+         
+        $crate::helpers::paste! {
+            wrappable_struct!(
+                $type,
+                [<$type Wrapper>],
+                [<$type:snake:upper _WRAPPER>]
+            );
+        }
+
+        pub(crate) fn klass() -> $crate::helpers::Class {
+            $crate::helpers::namespace().get_nested_class($klass)
+        }
+
+        impl $crate::helpers::TryUnwrap<$ruby_type> for $crate::helpers::AnyObject {
+            fn try_unwrap(self) -> $ruby_type {
+                let result = self.try_convert_to::<$ruby_type>();
+                <_ as $crate::helpers::TryUnwrap<$ruby_type>>::try_unwrap(result)
+            }
+        }
+
+        impl $crate::helpers::VerifiedObject for $ruby_type {
+            fn is_correct_type<T: $crate::helpers::Object>(object: &T) -> bool {
+                object.class() == klass()
+            }
+
+            fn error_message() -> &'static str {
+                concat!("Error converting to ", stringify!($ruby_type), ".")
+            }
+        }
+    }
+}
+
+pub(crate) use scaffold;

data/src/index.rs

@@ -0,0 +1,261 @@
+use std::collections::HashMap;
+use std::str::FromStr;
+use rutie::{methods, Object, AnyObject, Integer, NilClass, Array, RString, Hash};
+use tantivy::{doc, Document, Term, ReloadPolicy, Index, IndexWriter, IndexReader, DateTime};
+use tantivy::schema::{Schema, TextOptions, TextFieldIndexing, IndexRecordOption, FacetOptions, STRING, STORED, INDEXED, FAST};
+use tantivy::collector::TopDocs;
+use tantivy::directory::MmapDirectory;
+
+use crate::helpers::{scaffold, try_unwrap_params, TryUnwrap};
+use crate::query::{unwrap_query, RTantinyQuery};
+use crate::tokenizer::{unwrap_tokenizer, RTantinyTokenizer};
+
+pub struct TantinyIndex {
+    pub(crate) index_writer: IndexWriter,
+    pub(crate) index_reader: IndexReader,
+    pub(crate) schema: Schema,
+}
+
+scaffold!(RTantinyIndex, TantinyIndex, "Index");
+
+pub(crate) fn unwrap_index(index: &RTantinyIndex) -> &TantinyIndex {
+    index.get_data(&*TANTINY_INDEX_WRAPPER)
+}
+
+#[rustfmt::skip::macros(methods)]
+methods!(
+    RTantinyIndex,
+    _itself,
+
+    fn new_index(
+        path: RString,
+        index_size: Integer,
+        default_tokenizer: AnyObject,
+        field_tokenizers: Hash,
+        text_fields: Array,
+        string_fields: Array,
+        integer_fields: Array,
+        double_fields: Array,
+        date_fields: Array,
+        facet_fields: Array
+    ) -> RTantinyIndex {
+        try_unwrap_params!(
+            path: String,
+            index_size: i64,
+            default_tokenizer: RTantinyTokenizer,
+            field_tokenizers: HashMap<String, RTantinyTokenizer>,
+            text_fields: Vec<String>,
+            string_fields: Vec<String>,
+            integer_fields: Vec<String>,
+            double_fields: Vec<String>,
+            date_fields: Vec<String>,
+            facet_fields: Vec<String>
+        );
+
+        let index_path = MmapDirectory::open(path).try_unwrap();
+        let mut schema_builder = Schema::builder();
+
+        schema_builder.add_text_field("id", STRING | STORED);
+
+        for field in text_fields {
+            let tokenizer_name =
+                if field_tokenizers.contains_key(&field) {
+                    &*field
+                } else {
+                    "default"
+                };
+            let indexing = TextFieldIndexing::default()
+                .set_tokenizer(tokenizer_name)
+                .set_index_option(IndexRecordOption::WithFreqsAndPositions);
+            let options = TextOptions::default()
+                .set_indexing_options(indexing);
+            schema_builder.add_text_field(&field, options);
+        }
+
+        for field in string_fields {
+            schema_builder.add_text_field(&field, STRING);
+        }
+
+        for field in integer_fields {
+            schema_builder.add_i64_field(&field, FAST | INDEXED);
+        }
+
+        for field in double_fields {
+            schema_builder.add_f64_field(&field, FAST | INDEXED);
+        }
+
+        for field in date_fields {
+            schema_builder.add_date_field(&field, FAST | INDEXED);
+        }
+
+        for field in facet_fields {
+            let options = FacetOptions::default().set_indexed();
+            schema_builder.add_facet_field(&field, options);
+        }
+
+        let schema = schema_builder.build();
+        let index = Index::open_or_create(index_path, schema.clone()).try_unwrap();
+        let tokenizers = index.tokenizers();
+
+        tokenizers.register("default", unwrap_tokenizer(&default_tokenizer).clone());
+
+        for (field, tokenizer) in field_tokenizers {
+            tokenizers.register(&field, unwrap_tokenizer(&tokenizer).clone())
+        }
+
+        let mut index_writer = index
+            .writer(index_size as usize)
+            .try_unwrap();
+
+        let index_reader = index
+            .reader_builder()
+            .reload_policy(ReloadPolicy::Manual)
+            .try_into()
+            .try_unwrap();
+        
+        klass().wrap_data(
+            TantinyIndex { index_writer, index_reader, schema },
+            &*TANTINY_INDEX_WRAPPER
+        )
+    }
+
+    fn add_document(
+        id: RString,
+        text_fields: Hash,
+        string_fields: Hash,
+        integer_fields: Hash,
+        double_fields: Hash,
+        date_fields: Hash,
+        facet_fields: Hash
+    ) -> NilClass {
+        try_unwrap_params!(
+            id: String,
+            text_fields: HashMap<String, String>,
+            string_fields: HashMap<String, String>,
+            integer_fields: HashMap<String, i64>,
+            double_fields: HashMap<String, f64>,
+            date_fields: HashMap<String, String>,
+            facet_fields: HashMap<String, String>
+        );
+
+        
+        let internal = unwrap_index(&_itself);
+        let index_writer = &internal.index_writer;
+        let schema = &internal.schema;
+
+        let mut doc = Document::default();
+
+        let id_field = schema.get_field("id").try_unwrap();
+        doc.add_text(id_field, &id);
+
+        for (key, value) in text_fields.iter() {
+            let field = schema.get_field(key).try_unwrap();
+            doc.add_text(field, value);
+        }
+
+        for (key, value) in string_fields.iter() {
+            let field = schema.get_field(key).try_unwrap();
+            doc.add_text(field, value);
+        }
+
+        for (key, &value) in integer_fields.iter() {
+            let field = schema.get_field(key).try_unwrap();
+            doc.add_i64(field, value);
+        }
+
+        for (key, &value) in double_fields.iter() {
+            let field = schema.get_field(key).try_unwrap();
+            doc.add_f64(field, value);
+        }
+
+        for (key, value) in date_fields.iter() {
+            let field = schema.get_field(key).try_unwrap();
+            let value = DateTime::from_str(value).try_unwrap();
+            doc.add_date(field, &value);
+        }
+
+        for (key, value) in facet_fields.iter() {
+            let field = schema.get_field(key).try_unwrap();
+            doc.add_facet(field, &value);
+        }
+
+        let doc_id = Term::from_field_text(id_field, &id);
+        index_writer.delete_term(doc_id.clone());
+
+        index_writer.add_document(doc);
+
+        NilClass::new()
+    }
+
+    fn delete_document(id: RString) -> NilClass {
+        try_unwrap_params!(id: String);
+
+        let internal = unwrap_index(&_itself);
+        let index_writer = &internal.index_writer;
+
+        let id_field = internal.schema.get_field("id").try_unwrap();
+        let doc_id = Term::from_field_text(id_field, &id);
+
+        index_writer.delete_term(doc_id.clone());
+
+        NilClass::new()
+    }
+
+    fn commit() -> NilClass {
+        let internal = _itself.get_data_mut(&*TANTINY_INDEX_WRAPPER);
+        let index_writer = &mut internal.index_writer;
+
+        index_writer.commit().try_unwrap();
+
+        NilClass::new()
+    }
+
+    fn reload() -> NilClass {
+        unwrap_index(&_itself).index_reader.reload().try_unwrap();
+
+        NilClass::new()
+    }
+
+    fn search(
+        query: AnyObject,
+        limit: Integer
+    ) -> Array {
+        try_unwrap_params!(
+            query: RTantinyQuery,
+            limit: i64
+        );
+
+        let internal = unwrap_index(&_itself);
+        let id_field = internal.schema.get_field("id").try_unwrap();
+        let searcher = internal.index_reader.searcher();
+        let query = unwrap_query(&query);
+
+        let top_docs = searcher
+            .search(query, &TopDocs::with_limit(limit as usize))
+            .try_unwrap();
+
+        let mut array = Array::with_capacity(top_docs.len());
+
+        for (_score, doc_address) in top_docs {
+            let doc = searcher.doc(doc_address).try_unwrap();
+            if let Some(value) = doc.get_first(id_field) {
+                if let Some(id) = (&*value).text() {
+                    array.push(RString::from(String::from(id)));
+                }
+            }
+        }
+
+        array
+    }
+);
+
+pub(super) fn init() {
+    klass().define(|klass| {
+        klass.def_self("__new", new_index);
+        klass.def("__add_document", add_document);
+        klass.def("__delete_document", delete_document);
+        klass.def("__commit", commit);
+        klass.def("__reload", reload);
+        klass.def("__search", search);
+    });
+} 

data/src/lib.rs

@@ -0,0 +1,15 @@
+mod helpers;
+#[allow(improper_ctypes_definitions)]
+mod index;
+#[allow(improper_ctypes_definitions)]
+mod query;
+
+#[allow(improper_ctypes_definitions)]
+mod tokenizer;
+
+#[no_mangle]
+pub extern "C" fn Init_tantiny() {
+    index::init();
+    query::init();
+    tokenizer::init();
+}