From a274eacbbc62164a6567faeaa8d18ea5993a133f Mon Sep 17 00:00:00 2001 From: Matthew Wild Date: Sat, 23 Mar 2019 08:47:55 +0000 Subject: util.queue: Add 'consume()' convenience iterator --- spec/util_queue_spec.lua | 37 +++++++++++++++++++++++++++++++++++++ util/queue.lua | 3 +++ 2 files changed, 40 insertions(+) diff --git a/spec/util_queue_spec.lua b/spec/util_queue_spec.lua index 7cd3d695..d73f523d 100644 --- a/spec/util_queue_spec.lua +++ b/spec/util_queue_spec.lua @@ -100,4 +100,41 @@ describe("util.queue", function() end); end); + describe("consume()", function () + it("should work", function () + local q = queue.new(10); + for i = 1, 5 do + q:push(i); + end + local c = 0; + for i in q:consume() do + assert(i == c + 1); + assert(q:count() == (5-i)); + c = i; + end + end); + + it("should work even if items are pushed in the loop", function () + local q = queue.new(10); + for i = 1, 5 do + q:push(i); + end + local c = 0; + for i in q:consume() do + assert(i == c + 1); + if c < 3 then + assert(q:count() == (5-i)); + else + assert(q:count() == (6-i)); + end + + c = i; + + if c == 3 then + q:push(6); + end + end + assert.equal(c, 6); + end); + end); end); diff --git a/util/queue.lua b/util/queue.lua index 728e905f..e63b3f1c 100644 --- a/util/queue.lua +++ b/util/queue.lua @@ -64,6 +64,9 @@ local function new(size, allow_wrapping) return pos+1, t._items[read_pos]; end, self, 0; end; + consume = function (self) + return self.pop, self; + end; }; end -- cgit v1.2.3