From 56f24941ab7368b278935b75f25465163bb98f52 Mon Sep 17 00:00:00 2001
From: Adam Retter <adam.retter@googlemail.com>
Date: Sat, 3 Jan 2015 20:19:03 +0000
Subject: [PATCH] Simplify the Java API by permitting WriteBatchWithIndex to be
 provided straight to RocksDB#write

---
 java/org/rocksdb/RocksDB.java             | 22 ++++++++++++---
 java/org/rocksdb/WriteBatchWithIndex.java |  5 ----
 java/rocksjni/rocksjni.cc                 | 33 ++++++++++++++++++-----
 3 files changed, 46 insertions(+), 14 deletions(-)

diff --git a/java/org/rocksdb/RocksDB.java b/java/org/rocksdb/RocksDB.java
index 6c6e72260..089882532 100644
--- a/java/org/rocksdb/RocksDB.java
+++ b/java/org/rocksdb/RocksDB.java
@@ -539,7 +539,21 @@ public class RocksDB extends RocksObject {
    */
   public void write(WriteOptions writeOpts, WriteBatch updates)
       throws RocksDBException {
-    write(writeOpts.nativeHandle_, updates.nativeHandle_);
+    write0(writeOpts.nativeHandle_, updates.nativeHandle_);
+  }
+
+  /**
+   * Apply the specified updates to the database.
+   *
+   * @param writeOpts WriteOptions instance
+   * @param updates WriteBatchWithIndex instance
+   *
+   * @throws RocksDBException thrown if error happens in underlying
+   *    native library.
+   */
+  public void write(WriteOptions writeOpts, WriteBatchWithIndex updates)
+      throws RocksDBException {
+    write1(writeOpts.nativeHandle_, updates.nativeHandle_);
   }
 
   /**
@@ -1547,8 +1561,10 @@ public class RocksDB extends RocksObject {
       long handle, long writeOptHandle,
       byte[] key, int keyLen,
       byte[] value, int valueLen, long cfHandle) throws RocksDBException;
-  protected native void write(
-      long writeOptHandle, long batchHandle) throws RocksDBException;
+  protected native void write0(
+      long writeOptHandle, long wbHandle) throws RocksDBException;
+  protected native void write1(
+      long writeOptHandle, long wbwiHandle) throws RocksDBException;
   protected native boolean keyMayExist(byte[] key, int keyLen,
       StringBuffer stringBuffer);
   protected native boolean keyMayExist(byte[] key, int keyLen,
diff --git a/java/org/rocksdb/WriteBatchWithIndex.java b/java/org/rocksdb/WriteBatchWithIndex.java
index bb42dc3d7..f71ba338c 100644
--- a/java/org/rocksdb/WriteBatchWithIndex.java
+++ b/java/org/rocksdb/WriteBatchWithIndex.java
@@ -18,11 +18,6 @@ package org.rocksdb;
  * get an iterator for the database with Read-Your-Own-Writes like capability
  */
 public class WriteBatchWithIndex extends AbstractWriteBatch {
-
-  //TODO(AR) need to cover directly passing WriteBatchWithIndex to {@see org.rocksdb.RocksDB#write(WriteBatch)
-  //this simplifies the Java API beyond the C++ API as you don't need to call
-  //GetWriteBatch on the WriteBatchWithIndex
-
   /**
    * Creates a WriteBatchWithIndex where no bytes
    * are reserved up-front, bytewise comparison is
diff --git a/java/rocksjni/rocksjni.cc b/java/rocksjni/rocksjni.cc
index fdd1b009f..54eef7f53 100644
--- a/java/rocksjni/rocksjni.cc
+++ b/java/rocksjni/rocksjni.cc
@@ -390,18 +390,39 @@ void Java_org_rocksdb_RocksDB_put__JJ_3BI_3BIJ(
 // rocksdb::DB::Write
 /*
  * Class:     org_rocksdb_RocksDB
- * Method:    write
+ * Method:    write0
  * Signature: (JJ)V
  */
-void Java_org_rocksdb_RocksDB_write(
+void Java_org_rocksdb_RocksDB_write0(
     JNIEnv* env, jobject jdb,
-    jlong jwrite_options_handle, jlong jbatch_handle) {
+    jlong jwrite_options_handle, jlong jwb_handle) {
   rocksdb::DB* db = rocksdb::RocksDBJni::getHandle(env, jdb);
-  auto write_options = reinterpret_cast<rocksdb::WriteOptions*>(
+  auto* write_options = reinterpret_cast<rocksdb::WriteOptions*>(
+      jwrite_options_handle);
+  auto* wb = reinterpret_cast<rocksdb::WriteBatch*>(jwb_handle);
+
+  rocksdb::Status s = db->Write(*write_options, wb);
+
+  if (!s.ok()) {
+    rocksdb::RocksDBExceptionJni::ThrowNew(env, s);
+  }
+}
+
+/*
+ * Class:     org_rocksdb_RocksDB
+ * Method:    write1
+ * Signature: (JJ)V
+ */
+void Java_org_rocksdb_RocksDB_write1(
+    JNIEnv* env, jobject jdb,
+    jlong jwrite_options_handle, jlong jwbwi_handle) {
+  rocksdb::DB* db = rocksdb::RocksDBJni::getHandle(env, jdb);
+  auto* write_options = reinterpret_cast<rocksdb::WriteOptions*>(
       jwrite_options_handle);
-  auto batch = reinterpret_cast<rocksdb::WriteBatch*>(jbatch_handle);
+  auto* wbwi = reinterpret_cast<rocksdb::WriteBatchWithIndex*>(jwbwi_handle);
+  auto* wb = wbwi->GetWriteBatch();
 
-  rocksdb::Status s = db->Write(*write_options, batch);
+  rocksdb::Status s = db->Write(*write_options, wb);
 
   if (!s.ok()) {
     rocksdb::RocksDBExceptionJni::ThrowNew(env, s);