From 68e8ae6413f6514290d13323ffddfa9906b275ef Mon Sep 17 00:00:00 2001
From: Jacek Sieka <jacek@status.im>
Date: Tue, 17 Sep 2024 13:54:28 +0200
Subject: [PATCH] arraybuf: fix iterator, setLen (#231)

---
 stew/arraybuf.nim       | 15 +++++++++++++--
 tests/test_arraybuf.nim | 22 ++++++++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/stew/arraybuf.nim b/stew/arraybuf.nim
index df1e409..573eb25 100644
--- a/stew/arraybuf.nim
+++ b/stew/arraybuf.nim
@@ -51,7 +51,7 @@ template setLen*(b: var ArrayBuf, newLenParam: int) =
   newLenParam.evalOnceAs(newLen)
   let nl = typeof(b.n)(newLen)
   for i in newLen ..< b.len():
-    reset(b.buf[b.len() - i - 1]) # reset cleared items when shrinking
+    reset(b.buf[i]) # reset cleared items when shrinking
   b.n = nl
 
 template data*(bParam: ArrayBuf): openArray =
@@ -64,7 +64,7 @@ template data*(bParam: var ArrayBuf): var openArray =
 
 iterator items*[N, T](b: ArrayBuf[N, T]): lent T =
   for i in 0 ..< b.len:
-    yield b.d[i]
+    yield b.buf[i]
 
 iterator mitems*[N, T](b: var ArrayBuf[N, T]): var T =
   for i in 0 ..< b.len:
@@ -83,6 +83,9 @@ template `[]`*[N, T](b: var ArrayBuf[N, T], i: int): var T =
 template `[]=`*[N, T](b: var ArrayBuf[N, T], i: int, v: T) =
   b.buf[i] = v
 
+template `[]`*[N, T](b: ArrayBuf[N, T], i: BackwardsIndex): lent T =
+  b.buf[b.len - int(i)]
+
 template `==`*(a, b: ArrayBuf): bool =
   a.data() == b.data()
 
@@ -102,3 +105,11 @@ template add*[N, T](b: var ArrayBuf[N, T], v: openArray[T]) =
   # TODO `b` is evaluated multiple times but since it's a `var` this should
   #      _hopefully_ be fine..
   b.n += typeof(b.n)(b.buf.toOpenArray(b.len, N - 1).copyFrom(v))
+
+template pop*[N, T](b: var ArrayBuf[N, T]): T =
+  ## Return the last item while removing it from the buffer
+  # TODO `b` is evaluated multiple times but since it's a `var` this should
+  #      _hopefully_ be fine..
+  assert b.n > 0, "pop from empty ArrayBuf"
+  b.n -= 1
+  move(b.buf[b.n])
\ No newline at end of file
diff --git a/tests/test_arraybuf.nim b/tests/test_arraybuf.nim
index ec4cf30..a80ad5c 100644
--- a/tests/test_arraybuf.nim
+++ b/tests/test_arraybuf.nim
@@ -34,3 +34,25 @@ suite "ArrayBuf":
 
     v.add(byte 4)
     doAssert v.data() == [byte 0, 1]
+
+  test "setLen clearing":
+    var v: ArrayBuf[5, byte]
+
+    v.add(1)
+    v.add(2)
+
+    check: v.pop() == 2
+
+    v.setLen(2)
+    check:
+      v[^1] == 0 # not 2!
+
+
+    v[1] = 42
+
+    v.setLen(1)
+    check:
+      v[0] == 1
+      v.len == 1
+    v.setLen(2)
+    doAssert v.data() == [byte 1, 0]