From af7cb910acbe7dbdd1fa722025e399ad95503a78 Mon Sep 17 00:00:00 2001
From: Justin Balthrop <justin@geni.com>
Date: Fri, 7 Aug 2009 13:18:56 -0700
Subject: [PATCH] add simple_test and enumerable range/dup queries. improve
 absolute sort

---
 lib/bdb/simple.rb | 105 +++++++++++++++++++++++++++++-----------------
 1 file changed, 67 insertions(+), 38 deletions(-)

diff --git a/lib/bdb/simple.rb b/lib/bdb/simple.rb
index 5da5b92..4bfb4df 100644
--- a/lib/bdb/simple.rb
+++ b/lib/bdb/simple.rb
@@ -1,4 +1,4 @@
-require 'bdb'
+require 'bdb' if not defined?(Bdb)
 
 class Bdb::Simple
   include Enumerable
@@ -19,7 +19,7 @@ class Bdb::Simple
       @db = Bdb::Db.new
       @db.flags = Bdb::DB_DUPSORT if dup?
       @db.btree_compare = lambda do |db, key1, key2|
-        compare_absolute(Marshal.load(key1), Marshal.load(key2))
+        self.class.compare_absolute(Marshal.load(key1), Marshal.load(key2))
       end
       @db.open(nil, file, nil, Bdb::Db::BTREE, Bdb::DB_CREATE, 0)    
     end
@@ -30,33 +30,16 @@ class Bdb::Simple
     db[Marshal.dump(key)] = Marshal.dump(value)
   end
 
+  def delete(key)
+    db.del(nil, Marshal.dump(key), 0)
+  end
+
   def [](key)    
-    if key.kind_of?(Range)
-      values = []
-      cursor = db.cursor(nil, 0)
-      k,v = cursor.get(Marshal.dump(key.first), nil, Bdb::DB_SET_RANGE)
-      while k and key.include?(Marshal.load(k))
-        values << Marshal.load(v)
-        k, v = cursor.get(nil, nil, Bdb::DB_NEXT)
-      end
-      cursor.close
-      values
+    if key.kind_of?(Range) or dup?
+      Bdb::SimpleSet.new(db, key)
     else
-      key = Marshal.dump(key)
-      if dup?
-        values = []
-        cursor = db.cursor(nil, 0)
-        k,v = cursor.get(key, nil, Bdb::DB_SET)
-        while data
-          values << Marshal.load(v)
-          k,v = cursor.get(nil, nil, Bdb::DB_NEXT_DUP)
-        end
-        cursor.close
-        values
-      else
-        v = db.get(nil, key, nil, 0)
-        Marshal.load(v) if v
-      end
+      v = db.get(nil, Marshal.dump(key), nil, 0)
+      Marshal.load(v) if v
     end
   end
 
@@ -86,20 +69,66 @@ class Bdb::Simple
   end
 
   def self.compare_absolute(left, right)
-    if left.is_a?(Array) and right.is_a?(Array)
-      left.zip(right) do |l,r|
-        comp = compare_absolute(l, r)
-        return comp unless comp == 0
+    if left.class == right.class
+      if left.is_a?(Array) and right.is_a?(Array)
+        # Arrays: compare one element at a time.
+        left.zip(right) do |l,r|
+          comp = compare_absolute(l, r)
+          return comp unless comp == 0
+        end
+        left.size == right.size ? 0 : -1
+      elsif left.is_a?(Hash) and right.is_a?(Hash)
+        # Hashes: sort the keys and compare as an array of arrays.
+        left  = left.to_a.sort  {|a,b| compare_absolute(a[0],b[0])}
+        right = right.to_a.sort {|a,b| compare_absolute(a[0],b[0])}
+        compare_absolute(left, right)
+      else
+        begin
+          # Try to use the spaceship operator.
+          left <=> right
+        rescue NoMethodError => e
+          left.hash <=> right.hash
+        end
       end
-      left.size == right.size ? 0 : -1
-    elsif left.kind_of?(right.class) or right.kind_of?(left.class)
-      left <=> right rescue 0
-    elsif left.is_a?(NilClass)
-      -1
-    elsif right.is_a?(NilClass)
-      1
     else
+      # Nil is the smallest. Hash is the largest. All other objects are sorted by class name.
+      return -1 if left.is_a?(NilClass)
+      return  1 if right.is_a?(NilClass)
+      return  1 if left.is_a?(Hash)
+      return -1 if right.is_a?(Hash)
+      
+      # Compare class names and hashes as a last resort if that fails.
       right.class.name <=> left.class.name
     end
   end
 end
+
+class Bdb::SimpleSet
+  include Enumerable
+  
+  def initialize(db, key)
+    @db  = db
+    @key = key
+  end
+  attr_reader :db, :key
+  
+  def each
+    if key.kind_of?(Range)
+      cursor = db.cursor(nil, 0)
+      k,v = cursor.get(Marshal.dump(key.first), nil, Bdb::DB_SET_RANGE)
+      while k and key.include?(Marshal.load(k))
+        yield Marshal.load(v)
+        k, v = cursor.get(nil, nil, Bdb::DB_NEXT)
+      end
+      cursor.close
+    else
+      cursor = db.cursor(nil, 0)
+      k,v = cursor.get(Marshal.dump(key), nil, Bdb::DB_SET)
+      while k
+        yield Marshal.load(v)
+        k,v = cursor.get(nil, nil, Bdb::DB_NEXT_DUP)
+      end
+      cursor.close
+    end
+  end
+end