2 files changed, 62 insertions, 19 deletions

diff --git a/ext/dbm/dbm.c b/ext/dbm/dbm.c
index 7196c18149..062ae692ab 100644
--- a/ext/dbm/dbm.c
+++ b/ext/dbm/dbm.c
@@ -84,7 +84,7 @@ fdbm_initialize(argc, argv, obj)
     else {
 	mode = NUM2INT(vmode);
     }
-    Check_SafeStr(file);
+    SafeStringValue(file);
 
     dbm = 0;
     if (mode >= 0) {
@@ -150,7 +150,7 @@ fdbm_fetch(obj, keystr, ifnone)
     struct dbmdata *dbmp;
     DBM *dbm;
 
-    Check_Type(keystr, T_STRING);
+    StringValue(keystr);
     key.dptr = RSTRING(keystr)->ptr;
     key.dsize = RSTRING(keystr)->len;
 
@@ -196,7 +196,7 @@ fdbm_index(obj, valstr)
     struct dbmdata *dbmp;
     DBM *dbm;
 
-    Check_Type(valstr, T_STRING);
+    StringValue(valstr);
     val.dptr = RSTRING(valstr)->ptr;
     val.dsize = RSTRING(valstr)->len;
 
@@ -238,7 +238,7 @@ fdbm_delete(obj, keystr)
     DBM *dbm;
 
     rb_secure(4);
-    Check_Type(keystr, T_STRING);
+    StringValue(keystr);
     key.dptr = RSTRING(keystr)->ptr;
     key.dsize = RSTRING(keystr)->len;
 
@@ -273,14 +273,15 @@ fdbm_shift(obj)
     rb_secure(4);
     GetDBM(obj, dbmp);
     dbm = dbmp->di_dbm;
+    dbmp->di_size = -1;
 
     key = dbm_firstkey(dbm); 
     if (!key.dptr) return Qnil;
     val = dbm_fetch(dbm, key);
-    dbm_delete(dbm, key);
-
     keystr = rb_tainted_str_new(key.dptr, key.dsize);
     valstr = rb_tainted_str_new(val.dptr, val.dsize);
+    dbm_delete(dbm, key);
+
     return rb_assoc_new(keystr, valstr);
 }
 
@@ -292,20 +293,35 @@ fdbm_delete_if(obj)
     struct dbmdata *dbmp;
     DBM *dbm;
     VALUE keystr, valstr;
+    VALUE ret, ary = rb_ary_new();
+    int i, status = 0, n;
 
     rb_secure(4);
     GetDBM(obj, dbmp);
     dbm = dbmp->di_dbm;
+    n = dbmp->di_size;
+    dbmp->di_size = -1;
+
     for (key = dbm_firstkey(dbm); key.dptr; key = dbm_nextkey(dbm)) {
 	val = dbm_fetch(dbm, key);
 	keystr = rb_tainted_str_new(key.dptr, key.dsize);
 	valstr = rb_tainted_str_new(val.dptr, val.dsize);
-	if (RTEST(rb_yield(rb_assoc_new(keystr, valstr)))) {
-	    if (dbm_delete(dbm, key)) {
-		rb_raise(rb_eDBMError, "dbm_delete failed");
-	    }
+        ret = rb_protect(rb_yield, rb_assoc_new(rb_str_dup(keystr), valstr), &status);
+        if (status != 0) break;
+	if (RTEST(ret)) rb_ary_push(ary, keystr);
+    }
+
+    for (i = 0; i < RARRAY(ary)->len; i++) {
+	keystr = RARRAY(ary)->ptr[i];
+	key.dptr = RSTRING(keystr)->ptr;
+	key.dsize = RSTRING(keystr)->len;
+	if (dbm_delete(dbm, key)) {
+	    rb_raise(rb_eDBMError, "dbm_delete failed");
 	}
     }
+    if (status) rb_jump_tag(status);
+    if (n > 0) dbmp->di_size = n - RARRAY(ary)->len;
+
     return obj;
 }
 
@@ -321,11 +337,13 @@ fdbm_clear(obj)
     GetDBM(obj, dbmp);
     dbm = dbmp->di_dbm;
     dbmp->di_size = -1;
-    for (key = dbm_firstkey(dbm); key.dptr; key = dbm_nextkey(dbm)) {
+    while (key = dbm_firstkey(dbm), key.dptr) {
 	if (dbm_delete(dbm, key)) {
 	    rb_raise(rb_eDBMError, "dbm_delete failed");
 	}
     }
+    dbmp->di_size = 0;
+
     return obj;
 }
 
@@ -569,7 +587,7 @@ fdbm_has_key(obj, keystr)
     struct dbmdata *dbmp;
     DBM *dbm;
 
-    Check_Type(keystr, T_STRING);
+    StringValue(keystr);
     key.dptr = RSTRING(keystr)->ptr;
     key.dsize = RSTRING(keystr)->len;
 
@@ -588,7 +606,7 @@ fdbm_has_value(obj, valstr)
     struct dbmdata *dbmp;
     DBM *dbm;
 
-    Check_Type(valstr, T_STRING);
+    StringValue(valstr);
     val.dptr = RSTRING(valstr)->ptr;
     val.dsize = RSTRING(valstr)->len;
 
@@ -682,7 +700,7 @@ Init_dbm()
     rb_define_method(cDBM, "each_pair", fdbm_each_pair, 0);
     rb_define_method(cDBM, "keys", fdbm_keys, 0);
     rb_define_method(cDBM, "values", fdbm_values, 0);
-    rb_define_method(cDBM, "shift", fdbm_shift, 1);
+    rb_define_method(cDBM, "shift", fdbm_shift, 0);
     rb_define_method(cDBM, "delete", fdbm_delete, 1);
     rb_define_method(cDBM, "delete_if", fdbm_delete_if, 0);
     rb_define_method(cDBM, "reject!", fdbm_delete_if, 0);
diff --git a/ext/dbm/extconf.rb b/ext/dbm/extconf.rb
index 9ce235a2b5..7afa7d92d4 100644
--- a/ext/dbm/extconf.rb
+++ b/ext/dbm/extconf.rb
@@ -1,13 +1,38 @@
 require 'mkmf'
 
 dir_config("dbm")
-if have_library("gdbm", "dbm_open")
-  gdbm = true
+
+$db_hdr = "ndbm.h"
+$db_prefix = ""
+dblib = with_config("dbm-type", nil)
+
+def db_check(db)
+  if /^db2?$/ =~ db
+    $db_prefix = "__db_n"
+    $db_hdr = db+".h"
+  end
+  r = have_library(db, db_prefix("dbm_open"))
+  if db == "gdbm"
+    $have_gdbm = true
+  end
+  return r
+end
+
+def db_prefix(func)
+  $db_prefix+func
 end
-gdbm or have_library("db", "dbm_open") or have_library("dbm", "dbm_open")
+
+if dblib
+  db_check(dblib)
+else
+  for dblib in %w(db db2 db1 dbm gdbm)
+    db_check(dblib) and break
+  end
+end
+
 have_header("cdefs.h") 
 have_header("sys/cdefs.h") 
-if have_header("ndbm.h") and have_func("dbm_open")
-  have_func("dbm_clearerr") unless gdbm
+if have_header($db_hdr) and have_func(db_prefix("dbm_open"))
+  have_func(db_prefix("dbm_clearerr")) unless $have_gdbm
   create_makefile("dbm")
 end