triez 0.2

data/ext/triez.cc

@@ -0,0 +1,261 @@
+#include <hat-trie.h>
+#include <ruby.h>
+#include <ruby/encoding.h>
+
+static VALUE hat_class;
+static rb_encoding* u8_enc;
+static rb_encoding* bin_enc;
+
+static inline VALUE unify_key(VALUE key) {
+    rb_encoding* enc = rb_enc_get(key);
+    if (enc != u8_enc && enc != bin_enc) {
+        return rb_funcall(key, rb_intern("encode"), 1, rb_enc_from_encoding(u8_enc));
+    } else {
+        return key;
+    }
+}
+
+static inline long long V2LL(VALUE v) {
+    union {VALUE v; long long l;} u;
+    u.v = v;
+    return u.l;
+}
+
+static inline VALUE LL2V(long long l) {
+    union {VALUE v; long long l;} u;
+    u.l = l;
+    return u.v;
+}
+
+struct HatTrie {
+    hattrie_t* p;
+    bool obj_value;
+    bool suffix;
+    bool initialized;
+
+    HatTrie() : obj_value(false), suffix(false), initialized(false) {
+        p = hattrie_create();
+    }
+
+    ~HatTrie() {
+        hattrie_free(p);
+    }
+};
+
+static void hat_mark(void* p_ht) {
+    HatTrie* ht = (HatTrie*)p_ht;
+    if (!ht->obj_value) {
+        return;
+    }
+    hattrie_t* p = ht->p;
+    hattrie_iter_t* it = hattrie_iter_begin(p, false);
+    while (!hattrie_iter_finished(it)) {
+        value_t* v = hattrie_iter_val(it);
+        if (!IMMEDIATE_P(*v)) {
+            rb_gc_mark(*v);
+        }
+    }
+    hattrie_iter_free(it);
+}
+
+static void hat_free(void* p) {
+    delete (HatTrie*)p;
+}
+
+static VALUE hat_alloc(VALUE self) {
+    HatTrie* ht = new HatTrie();
+    return Data_Wrap_Struct(hat_class, hat_mark, hat_free, ht);
+}
+
+#define PRE_HAT\
+    hattrie_t* p;\
+    HatTrie* ht;\
+    Data_Get_Struct(self, HatTrie, ht);\
+    p = ht->p;\
+    Check_Type(key, T_STRING);\
+    key = unify_key(key);
+
+static VALUE hat_set_type(VALUE self, VALUE obj_value, VALUE suffix) {
+    HatTrie* ht;
+    Data_Get_Struct(self, HatTrie, ht);
+    if (ht->initialized) {
+        rb_raise(rb_eRuntimeError, "Already initialized");
+        return self;
+    }
+    ht->obj_value = RTEST(obj_value);
+    ht->suffix = RTEST(suffix);
+    ht->initialized = true;
+    return self;
+}
+
+static VALUE hat_obj_value_p(VALUE self) {
+    HatTrie* ht;
+    Data_Get_Struct(self, HatTrie, ht);
+    return ht->obj_value ? Qtrue : Qfalse;
+}
+
+static VALUE hat_suffix_p(VALUE self) {
+    HatTrie* ht;
+    Data_Get_Struct(self, HatTrie, ht);
+    return ht->suffix ? Qtrue : Qfalse;
+}
+
+static VALUE hat_size(VALUE self) {
+    HatTrie* ht;
+    Data_Get_Struct(self, HatTrie, ht);
+    return ULL2NUM(hattrie_size(ht->p));
+}
+
+static VALUE hat_set(VALUE self, VALUE key, VALUE value) {
+    PRE_HAT;
+    long long v = ht->obj_value ? value : NUM2LL(value);
+    char* s = RSTRING_PTR(key);
+    size_t len = RSTRING_LEN(key);
+    if (ht->suffix) {
+        char* s_end = s + len;
+        long n;
+        for (; s < s_end; s += n, len -= n) {
+            n = rb_enc_mbclen(s, s_end, u8_enc);
+            hattrie_get(p, s, len)[0] = v;
+        }
+    } else {
+        hattrie_get(p, s, len)[0] = v;
+    }
+    return self;
+}
+
+static VALUE hat_alt(VALUE self, VALUE key) {
+    PRE_HAT;
+    char* s = RSTRING_PTR(key);
+    size_t len = RSTRING_LEN(key);
+    if (ht->suffix) {
+        char* s_end = s + len;
+        long n;
+        for (; s < s_end; s += n, len -= n) {
+            n = rb_enc_mbclen(s, s_end, u8_enc);
+            value_t* vp = hattrie_tryget(p, s, len);
+            long long v;
+            if (ht->obj_value) {
+                VALUE value = vp ? LL2V(vp[0]) : Qnil;
+                v = V2LL(rb_yield(value));
+            } else {
+                VALUE value = vp ? LL2NUM(vp[0]) : LL2NUM(0);
+                v = NUM2LL(rb_yield(value));
+            }
+            hattrie_get(p, s, len)[0] = v;
+        }
+    } else {
+        value_t* vp = hattrie_tryget(p, s, len);
+        long long v;
+        if (ht->obj_value) {
+            VALUE value = vp ? LL2V(vp[0]) : Qnil;
+            v = V2LL(rb_yield(value));
+        } else {
+            VALUE value = vp ? LL2NUM(vp[0]) : LL2NUM(0);
+            v = NUM2LL(rb_yield(value));
+        }
+        hattrie_get(p, s, len)[0] = v;
+    }
+    return self;
+}
+
+static VALUE hat_append(VALUE self, VALUE key) {
+    HatTrie* ht;
+    Data_Get_Struct(self, HatTrie, ht);
+    return hat_set(self, key, ht->obj_value ? Qnil : LL2NUM(0));
+}
+
+static VALUE hat_get(VALUE self, VALUE key) {
+    PRE_HAT;
+    value_t* vt = hattrie_tryget(p, RSTRING_PTR(key), RSTRING_LEN(key));
+    if (vt) {
+        return ht->obj_value ? (*vt) : LL2NUM(*vt);
+    } else {
+        return Qnil;
+    }
+}
+
+static VALUE hat_del(VALUE self, VALUE key) {
+    PRE_HAT;
+    const char* s = RSTRING_PTR(key);
+    size_t len = RSTRING_LEN(key);
+    value_t* vt = hattrie_tryget(p, s, len);
+    if (vt) {
+        hattrie_del(p, RSTRING_PTR(key), RSTRING_LEN(key));
+        return ht->obj_value ? (*vt) : LL2NUM(*vt);
+    } else {
+        return Qnil;
+    }
+}
+
+static VALUE hat_check(VALUE self, VALUE key) {
+    PRE_HAT;
+    value_t* vt = hattrie_tryget(p, RSTRING_PTR(key), RSTRING_LEN(key));
+    return vt ? Qtrue : Qfalse;
+}
+
+struct SearchCbData {
+    VALUE callback;
+    VALUE suffix;
+    VALUE value;
+};
+
+static VALUE hat_search_callback(VALUE data) {
+    SearchCbData* p = (SearchCbData*)data;
+    return rb_funcall(p->callback, rb_intern("call"), 2, p->suffix, p->value);
+}
+
+// returns: true if an error occured
+static VALUE hat_search(VALUE self, VALUE key, VALUE vlimit, VALUE vsort, VALUE callback) {
+    PRE_HAT;
+    long limit = 0;
+    if (vlimit != Qnil) {
+        limit = NUM2LONG(vlimit);
+    }
+
+    hattrie_iter_t* it = hattrie_iter_with_prefix(p, RTEST(vsort), RSTRING_PTR(key), RSTRING_LEN(key));
+    int error = 0;
+    SearchCbData data = {callback};
+    while (!hattrie_iter_finished(it)) {
+        if (vlimit != Qnil && limit-- <= 0) {
+            break;
+        }
+        size_t suffix_len;
+        const char* suffix_s = hattrie_iter_key(it, &suffix_len);
+        value_t* v = hattrie_iter_val(it);
+        data.suffix = rb_enc_str_new(suffix_s, suffix_len, u8_enc);
+        data.value = ht->obj_value ? (*v) : LL2NUM(*v);
+        rb_protect(hat_search_callback, (VALUE)&data, &error);
+        if (error) {
+            break;
+        }
+        hattrie_iter_next(it);
+    }
+    hattrie_iter_free(it);
+    if (error) {
+        rb_funcall(rb_mKernel, rb_intern("raise"), 0);
+    }
+    return self;
+}
+
+#define DEF(k,n,f,c) rb_define_method(k,n,RUBY_METHOD_FUNC(f),c)
+
+extern "C"
+void Init_triez() {
+    hat_class = rb_define_class("Triez", rb_cObject);
+    u8_enc = rb_utf8_encoding();
+    bin_enc = rb_ascii8bit_encoding();
+
+    rb_define_alloc_func(hat_class, hat_alloc);
+    DEF(hat_class, "_internal_set_type", hat_set_type, 2);
+    DEF(hat_class, "obj_value?", hat_obj_value_p, 0);
+    DEF(hat_class, "suffix?", hat_suffix_p, 0);
+    DEF(hat_class, "size", hat_size, 0);
+    DEF(hat_class, "has_key?", hat_check, 1);
+    DEF(hat_class, "[]=", hat_set, 2);
+    DEF(hat_class, "alt", hat_alt, 1);
+    DEF(hat_class, "<<", hat_append, 1);
+    DEF(hat_class, "[]", hat_get, 1);
+    DEF(hat_class, "delete", hat_del, 1);
+    DEF(hat_class, "_internal_search", hat_search, 4);
+}

data/lib/triez.rb

@@ -0,0 +1,45 @@
+require_relative "../ext/triez"
+
+class Triez
+  VERSION = '0.2'
+
+  private :_internal_set_type
+  private :_internal_search
+
+  def initialize opts={}
+    opts = opts.dup
+    obj_value = opts.delete :obj_value
+    obj_value = false if obj_value.nil?
+    suffix    = opts.delete :suffix
+    suffix    = false if suffix.nil?
+    unless opts.empty?
+      raise ArgumentError, "Unknown options: #{opts.keys.inspect}, only [:suffix, :obj_value] are allowed"
+    end
+    _internal_set_type obj_value, suffix
+  end
+
+  def each &p
+    raise ArgumentError, 'Need a block' unless p
+    _internal_search '', nil, true, p
+  end
+
+  def search_with_prefix prefix, opts={}, &p
+    opts  = opts.dup
+    limit = opts.delete :limit
+    if !limit.nil? and limit < 0
+      raise ArgumentError, "Limit should be > 0"
+    end
+    sort = opts.delete :sort
+    unless opts.empty?
+      raise ArgumentError, "Unknown options: #{opts.keys.inspect}, only [:limit, :sort] are allowed"
+    end
+
+    if p
+      _internal_search prefix, limit, sort, p
+    else
+      a = []
+      _internal_search prefix, limit, sort, -> k, v {a << [k, v]}
+      a
+    end
+  end
+end

data/readme.md

@@ -0,0 +1,174 @@
+## What
+
+Pragmatic [trie](http://en.wikipedia.org/wiki/Trie) for Ruby.
+
+It is fast, memory efficient, unicode aware.
+
+The backend of *triez* is a cache oblivious data structure: the [HAT trie](https://github.com/dcjones/hat-trie). It is generally faster and more memory efficient than double arrays or burst tries.
+
+## Requirement
+
+- Ruby 1.9
+- `g++` or `clang`
+
+## Install
+
+``` bash
+gem ins triez
+```
+
+## Synopsis
+
+``` ruby
+require 'triez'
+
+t = Triez.new
+
+# insertion
+t['key'] = 100
+
+# insert a key with default value (0 for normal triez, nil for obj_valued triez)
+t << 'key'
+
+# search
+t.has_key? 'key'
+t['key']
+
+# iterate over values under a prefix.
+t.search_with_prefix(prefix, limit: 10, sort: true) do |suffix, value|
+  ...
+end
+
+# enumerate (NOTE it is unordered)
+t.each do |key, value|
+  ...
+end
+```
+
+---
+
+By default, a *triez* stores signed integers within 64bits, you can use it as weights, counts or database IDs, and doesn't cost any time in GC marking phase. In case you need to store arbitrary object in a node, use `obj_value: true`:
+
+``` ruby
+t = Triez.new obj_value: true
+t['Tom'] = {name: 'Tom', sex: 'Female'}
+t['Tree'] = [:leaf, :trunk, :root]
+```
+
+---
+
+When a *triez* is initialized with `suffix: true`, it inserts all suffices of a key
+
+``` ruby
+t = Triez.new suffix: true
+t['万塘路一锅鸡'] = 2
+t['万塘路一锅鸡'] #=> 2
+t['塘路一锅鸡'] #=> 2
+t['路一锅鸡']  #=> 2
+t['一锅鸡']   #=> 2
+t['锅鸡']    #=> 2
+t['鸡']     #=> 2
+```
+
+You can batch change values with a block
+
+``` ruby
+# v *= 5 for 'abcd', 'bcd', 'cd', 'd'
+t.alt 'abcd' do |v|
+  v * 5
+end
+t['abcd'] #=> 10
+t['cd']   #=> 10
+```
+
+---
+
+Misc methods
+
+``` ruby
+# if it is a suffix trie
+t.suffix?
+# if the value type is object
+t.obj_value?
+```
+
+## Examples
+
+Prefix-based autocompletion:
+
+``` ruby
+require 'triez'
+words = %w[readme, rot, red, rah, rasterization]
+t = Triez.new
+words.each do |word|
+  t[word] = 1
+end
+t.search_with_prefix 're' do |word|
+  puts "candidate: #{word}"
+end
+```
+
+The output:
+
+```bash
+candidate: readme
+candidate: red
+```
+
+---
+
+Efficiently search for strings containing a substring:
+
+``` ruby
+require 'triez'
+sequences = {
+  'ACTGAAAAAAACTG' => 1,
+  'ATACGGTCCA' => 2,
+  'GCTTGTACGT' => 3
+}
+t = Triez.new suffix: true
+sequences.each do |seq, id|
+  t[seq] = id
+end
+t.search_with_prefix 'CGGT' do |_, id|
+  puts id #=> 2
+end
+```
+
+The search time is linear to the length of the substring.
+
+## Benchmarks
+
+Here's a benchmark on
+
+```ruby
+ruby 1.9.3p374 (2013-01-15 revision 38858) [x86_64-darwin12.2.1]
+2.3 GHz Intel Core i7
+```
+
+The test data is 3 milion titles of wikipedia articles (from http://dumps.wikimedia.org/enwiki/20121101/)
+
+```
+thing/backend      | memory  | insertion time | 3 M query
+-------------------|---------|----------------|----------
+hash/linked hash   | 340.2 M |    4.369 s     | 0.2800 s
+trie/double array* | 155.6 M |    130.7 s     | 0.4359 s
+triez/HAT trie     | 121.7 M |    3.872 s     | 0.3472 s
+```
+
+NOTE: `trie/double array` -> https://github.com/tyler/trie
+
+## Caveats
+
+- `sort` orders keys with binary collation, not unicode codepoint collation in string comparison.
+- For some rare case of many threads modifying the same trie, you may need a mutex.
+- If you still feel memory not enough, you may consider [MARISA-trie](https://code.google.com/p/marisa-trie/) (NOTE that MARISA is immutable) or a database.
+
+## Development
+
+``` bash
+git clone git://github.com/luikore/triez.git
+cd triez
+rake glob_src
+rake
+```