msgpack#105 Add MessageUnpacker.readPayloadAsReference

xerial · xerial · commit cf28209472fc · 2014-06-10T16:43:58.000+09:00
diff --git a/msgpack-core/src/main/java/org/msgpack/core/MessageUnpacker.java b/msgpack-core/src/main/java/org/msgpack/core/MessageUnpacker.java
@@ -202,6 +202,9 @@ private MessageBuffer takeNextBuffer() throws IOException {
      * @throws IOException
      */
     private boolean ensure(int byteSizeToRead) throws IOException {
+        if(byteSizeToRead == 0)
+            return true;
+
         if(!ensureBuffer())
             return false;
 
@@ -1016,6 +1019,16 @@ public void readPayload(byte[] dst, int off, int len) throws IOException {
         }
     }
 
+    public MessageBuffer readPayloadAsReference(int length) throws IOException {
+        checkArgument(length >= 0);
+        if(!ensure(length))
+            throw new EOFException();
+
+        MessageBuffer ref = buffer.slice(position, length);
+        position += length;
+        return ref;
+    }
+
 
     private int readNextLength8() throws IOException {
         byte u8 = readByte();
diff --git a/msgpack-core/src/test/scala/org/msgpack/core/MessageUnpackerTest.scala b/msgpack-core/src/test/scala/org/msgpack/core/MessageUnpackerTest.scala
@@ -445,11 +445,47 @@ class MessageUnpackerTest extends MessagePackSpec {
           }
           unpacker.close()
         }
-      }
 
+        block("v7-ref") {
+          val unpacker = new MessageUnpacker(new ByteArrayInputStream(b))
+          var i = 0
+          while(i < R) {
+            val len = unpacker.unpackBinaryHeader()
+            val out = unpacker.readPayloadAsReference(len)
+            i += 1
+          }
+          unpacker.close()
+        }
+      }
     }
 
+    "read payload as a reference" taggedAs("ref") in {
+
+      val dataSizes = Seq(0, 1, 5, 8, 16, 32, 128, 256, 1024, 2000, 10000, 100000)
+
+      for(s <- dataSizes) {
+        When(f"data size is $s%,d")
+        val data = new Array[Byte](s)
+        Random.nextBytes(data)
+        val b = new ByteArrayOutputStream()
+        val packer = new MessagePacker(b)
+        packer.packBinaryHeader(s)
+        packer.writePayload(data)
+        packer.close()
+
+        val unpacker = new MessageUnpacker(b.toByteArray)
+        val len = unpacker.unpackBinaryHeader()
+        len shouldBe s
+        val ref = unpacker.readPayloadAsReference(len)
+        unpacker.close()
+        ref.size() shouldBe s
+        val stored = new Array[Byte](len)
+        ref.getBytes(0, stored, 0, len)
+
+        stored shouldBe data
+      }
 
+    }
 
   }
 
