Skip to content

Commit

Permalink
arraybuf: fix iterator, setLen
Browse files Browse the repository at this point in the history
  • Loading branch information
arnetheduck committed Sep 16, 2024
1 parent fc09b2e commit 861c173
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 2 deletions.
15 changes: 13 additions & 2 deletions stew/arraybuf.nim
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand All @@ -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:
Expand All @@ -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()

Expand All @@ -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])
22 changes: 22 additions & 0 deletions tests/test_arraybuf.nim
Original file line number Diff line number Diff line change
Expand Up @@ -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 0, 1]

0 comments on commit 861c173

Please sign in to comment.