Add MDB_MULTIPLE support (fixes #24 and fixes #28)

bp-alex · bp-alex · commit c31d125cdef1 · 2017-07-18T10:22:23.000+10:00
diff --git a/src/main/java/org/lmdbjava/Cursor.java b/src/main/java/org/lmdbjava/Cursor.java
@@ -29,6 +29,7 @@
 import static org.lmdbjava.Library.LIB;
 import static org.lmdbjava.MaskedFlag.isSet;
 import static org.lmdbjava.MaskedFlag.mask;
+import static org.lmdbjava.PutFlags.MDB_MULTIPLE;
 import static org.lmdbjava.PutFlags.MDB_NODUPDATA;
 import static org.lmdbjava.PutFlags.MDB_NOOVERWRITE;
 import static org.lmdbjava.PutFlags.MDB_RESERVE;
@@ -224,6 +225,42 @@ public boolean put(final T key, final T val, final PutFlags... op) {
     return true;
   }
 
+  /**
+   * Put multiple values into the database in one <code>MDB_MULTIPLE</code>
+   * operation.
+   *
+   * <p>
+   * The database must have been opened with {@link DbiFlags#MDB_DUPFIXED}. The
+   * buffer must contain fixed-sized values to be inserted. The size of each
+   * element is calculated from the buffer's size divided by the given element
+   * count. For example, to populate 10 X 4 byte integers at once, present a
+   * buffer of 40 bytes and specify the element as 10.
+   *
+   * @param key      key to store in the database (not null)
+   * @param val      value to store in the database (not null)
+   * @param elements number of elements contained in the passed value buffer
+   * @param op       options for operation (must set <code>MDB_MULTIPLE</code>)
+   */
+  public void putMultiple(final T key, final T val, final int elements,
+                          final PutFlags... op) {
+    if (SHOULD_CHECK) {
+      requireNonNull(txn);
+      requireNonNull(key);
+      requireNonNull(val);
+      txn.checkReady();
+      txn.checkWritesAllowed();
+    }
+    final int mask = mask(op);
+    if (SHOULD_CHECK && !isSet(mask, MDB_MULTIPLE)) {
+      throw new IllegalArgumentException("Must set " + MDB_MULTIPLE + " flag");
+    }
+    txn.kv().keyIn(key);
+    final Pointer dataPtr = txn.kv().valInMulti(val, elements);
+    final int rc = LIB.mdb_cursor_put(ptrCursor, txn.kv().pointerKey(),
+                                      dataPtr, mask);
+    checkRc(rc);
+  }
+
   /**
    * Renew a cursor handle.
    *
diff --git a/src/main/java/org/lmdbjava/KeyVal.java b/src/main/java/org/lmdbjava/KeyVal.java
@@ -24,6 +24,7 @@
 import jnr.ffi.Pointer;
 import jnr.ffi.provider.MemoryManager;
 import static org.lmdbjava.BufferProxy.MDB_VAL_STRUCT_SIZE;
+import static org.lmdbjava.BufferProxy.STRUCT_FIELD_OFFSET_SIZE;
 import static org.lmdbjava.Library.RUNTIME;
 
 /**
@@ -37,6 +38,7 @@ final class KeyVal<T> implements AutoCloseable {
   private boolean closed;
   private T k;
   private final BufferProxy<T> proxy;
+  private final Pointer ptrArray;
   private final Pointer ptrKey;
   private final long ptrKeyAddr;
   private final Pointer ptrVal;
@@ -50,7 +52,8 @@ final class KeyVal<T> implements AutoCloseable {
     this.v = proxy.allocate();
     ptrKey = MEM_MGR.allocateTemporary(MDB_VAL_STRUCT_SIZE, false);
     ptrKeyAddr = ptrKey.address();
-    ptrVal = MEM_MGR.allocateTemporary(MDB_VAL_STRUCT_SIZE, false);
+    ptrArray = MEM_MGR.allocateTemporary(MDB_VAL_STRUCT_SIZE * 2, false);
+    ptrVal = ptrArray.slice(0, MDB_VAL_STRUCT_SIZE);
     ptrValAddr = ptrVal.address();
   }
 
@@ -97,6 +100,35 @@ void valIn(final int size) {
     proxy.in(v, size, ptrVal, ptrValAddr);
   }
 
+  /**
+   * Prepares an array suitable for presentation as the data argument to a
+   * <code>MDB_MULTIPLE</code> put.
+   *
+   * <p>
+   * The returned array is equivalent of two <code>MDB_val</code>s as follows:
+   *
+   * <ul>
+   * <li>ptrVal1.data = pointer to the data address of passed buffer</li>
+   * <li>ptrVal1.size = size of each individual data element</li>
+   * <li>ptrVal2.data = unused</li>
+   * <li>ptrVal2.size = number of data elements (as passed to this method)</li>
+   * </ul>
+   *
+   * @param val      a user-provided buffer with data elements (required)
+   * @param elements number of data elements the user has provided
+   * @return a properly-prepared pointer to an array for the operation
+   */
+  Pointer valInMulti(final T val, final int elements) {
+    final long ptrVal2SizeOff = MDB_VAL_STRUCT_SIZE + STRUCT_FIELD_OFFSET_SIZE;
+    ptrArray.putLong(ptrVal2SizeOff, elements); // ptrVal2.size
+    proxy.in(val, ptrVal, ptrValAddr); // ptrVal1.data
+    final long totalBufferSize = ptrVal.getLong(STRUCT_FIELD_OFFSET_SIZE);
+    final long elemSize = totalBufferSize / elements;
+    ptrVal.putLong(STRUCT_FIELD_OFFSET_SIZE, elemSize); // ptrVal1.size
+
+    return ptrArray;
+  }
+
   T valOut() {
     v = proxy.out(v, ptrVal, ptrValAddr);
     return v;
diff --git a/src/test/java/org/lmdbjava/CursorTest.java b/src/test/java/org/lmdbjava/CursorTest.java
@@ -26,6 +26,7 @@
 import static java.lang.Long.BYTES;
 import static java.lang.Long.MIN_VALUE;
 import java.nio.ByteBuffer;
+import static java.nio.ByteBuffer.allocateDirect;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.MatcherAssert.assertThat;
 import org.junit.After;
@@ -38,10 +39,12 @@
 import static org.lmdbjava.ByteBufferProxy.PROXY_OPTIMAL;
 import org.lmdbjava.Cursor.ClosedException;
 import static org.lmdbjava.DbiFlags.MDB_CREATE;
+import static org.lmdbjava.DbiFlags.MDB_DUPFIXED;
 import static org.lmdbjava.DbiFlags.MDB_DUPSORT;
 import static org.lmdbjava.Env.create;
 import static org.lmdbjava.EnvFlags.MDB_NOSUBDIR;
 import static org.lmdbjava.PutFlags.MDB_APPENDDUP;
+import static org.lmdbjava.PutFlags.MDB_MULTIPLE;
 import static org.lmdbjava.PutFlags.MDB_NODUPDATA;
 import static org.lmdbjava.PutFlags.MDB_NOOVERWRITE;
 import static org.lmdbjava.SeekOp.MDB_FIRST;
@@ -169,6 +172,36 @@ public void delete() {
     }
   }
 
+  @Test
+  public void putMultiple() {
+    final Dbi<ByteBuffer> db = env.openDbi(DB_1, MDB_CREATE, MDB_DUPSORT,
+                                           MDB_DUPFIXED);
+    final int elemCount = 20;
+
+    final ByteBuffer values = allocateDirect(Integer.BYTES * elemCount);
+    for (int i = 1; i <= elemCount; i++) {
+      values.putInt(i);
+    }
+    values.flip();
+
+    final int key = 100;
+    final ByteBuffer k = bb(key);
+    try (Txn<ByteBuffer> txn = env.txnWrite()) {
+      final Cursor<ByteBuffer> c = db.openCursor(txn);
+      c.putMultiple(k, values, elemCount, MDB_MULTIPLE);
+      assertThat(c.count(), is((long) elemCount));
+    }
+  }
+
+  @Test(expected = IllegalArgumentException.class)
+  public void putMultipleWithoutMdbMultipleFlag() {
+    final Dbi<ByteBuffer> db = env.openDbi(DB_1, MDB_CREATE, MDB_DUPSORT);
+    try (Txn<ByteBuffer> txn = env.txnWrite()) {
+      final Cursor<ByteBuffer> c = db.openCursor(txn);
+      c.putMultiple(bb(100), bb(1), 1);
+    }
+  }
+
   @Test
   public void renewTxRo() {
     final Dbi<ByteBuffer> db = env.openDbi(DB_1, MDB_CREATE);
