util.dbuffer: If no bytes parameter passed to read, return remainder of frontmost chunk

2 files changed, 22 insertions, 2 deletions

diff --git a/spec/util_dbuffer_spec.lua b/spec/util_dbuffer_spec.lua
index 1a0c2992..c4fa22b2 100644
--- a/spec/util_dbuffer_spec.lua
+++ b/spec/util_dbuffer_spec.lua
@@ -21,6 +21,21 @@ describe("util.dbuffer", function ()
 		end);
 	end);
 
+	describe(":read", function ()
+		it("supports optional bytes parameter", function ()
+			-- should return the frontmost chunk
+			local b = dbuffer.new();
+			assert.truthy(b:write("hello"));
+			assert.truthy(b:write(" "));
+			assert.truthy(b:write("world"));
+			assert.equal("h", b:read(1));
+
+			assert.equal("ello", b:read());
+			assert.equal(" ", b:read());
+			assert.equal("world", b:read());
+		end);
+	end);
+
 	describe(":discard", function ()
 		local b = dbuffer.new();
 		it("works", function ()
diff --git a/util/dbuffer.lua b/util/dbuffer.lua
index c38a16f5..63dca750 100644
--- a/util/dbuffer.lua
+++ b/util/dbuffer.lua
@@ -24,6 +24,9 @@ function dbuffer_methods:read_chunk(requested_bytes)
 	if not chunk then return; end
 	local chunk_length = #chunk;
 	local remaining_chunk_length = chunk_length - consumed;
+	if not requested_bytes then
+		requested_bytes = remaining_chunk_length;
+	end
 	if remaining_chunk_length <= requested_bytes then
 		self.front_consumed = 0;
 		self._length = self._length - remaining_chunk_length;
@@ -41,12 +44,14 @@ end
 function dbuffer_methods:read(requested_bytes)
 	local chunks;
 
-	if requested_bytes > self._length then
+	if requested_bytes and requested_bytes > self._length then
 		return nil;
 	end
 
 	local chunk, read_bytes = self:read_chunk(requested_bytes);
-	if chunk then
+	if not requested_bytes then
+		return chunk;
+	elseif chunk then
 		requested_bytes = requested_bytes - read_bytes;
 		if requested_bytes == 0 then -- Already read everything we need
 			return chunk;