From 5ed54ac2f33875742bb583d65cda795ccc1f28d9 Mon Sep 17 00:00:00 2001 From: Mikhail Tavarez Date: Mon, 25 Mar 2024 10:40:00 -0500 Subject: [PATCH] fix incorrect resizing logic in bytes struct --- gojo/bufio/bufio.mojo | 2 ++ gojo/builtins/_bytes.mojo | 17 ++++++++++++++--- gojo/builtins/attributes.mojo | 6 +++++- gojo/bytes/buffer.mojo | 20 ++++++++++---------- tests/test_bufio.mojo | 22 ++++++++++++++++++++++ 5 files changed, 53 insertions(+), 14 deletions(-) diff --git a/gojo/bufio/bufio.mojo b/gojo/bufio/bufio.mojo index 783be23..d81c1cf 100644 --- a/gojo/bufio/bufio.mojo +++ b/gojo/bufio/bufio.mojo @@ -727,6 +727,8 @@ struct Writer[W: io.Writer]( self.err = error return error + # Reset the buffer + self.buf = Bytes(size=self.buf.size()) self.bytes_written = 0 return None diff --git a/gojo/builtins/_bytes.mojo b/gojo/builtins/_bytes.mojo index 0b5acea..12b02ef 100644 --- a/gojo/builtins/_bytes.mojo +++ b/gojo/builtins/_bytes.mojo @@ -1,4 +1,3 @@ -from time import now from .errors import panic @@ -89,11 +88,17 @@ struct Bytes(Stringable, Sized, CollectionElement): return new_bytes fn __setitem__(inout self, index: Int, value: Int8): + if index >= len(self._vector): + panic("Bytes.__setitem__: Tried setting index out of range. Vector length is " + str(len(self._vector)) + " but tried to set index " + str(index) + ".") + self._vector[index] = value if index >= self.write_position: self.write_position = index + 1 fn __setitem__(inout self, index: Int, value: Self): + if index >= len(self._vector): + panic("Bytes.__setitem__: Tried setting index out of range. Vector length is " + str(len(self._vector)) + " but tried to set index " + str(index) + ".") + self._vector[index] = value[0] if index >= self.write_position: self.write_position = index + 1 @@ -139,6 +144,12 @@ struct Bytes(Stringable, Sized, CollectionElement): fn __repr__(self) -> String: return self.__str__() + + fn __contains__(self, item: Int8) -> Bool: + for i in range(len(self)): + if self[i] == item: + return True + return False fn append(inout self, value: Byte): """Appends the value to the end of the Bytes. @@ -147,7 +158,7 @@ struct Bytes(Stringable, Sized, CollectionElement): value: The value to append. """ # If the vector is full, resize it. - if self.write_position == self.size(): + if self.write_position >= self.size(): self.resize(self.size() * 2) self[self.write_position] = value @@ -239,4 +250,4 @@ struct Bytes(Stringable, Sized, CollectionElement): """ var new_bytes = self.copy()._vector new_bytes.append(0) - return new_bytes \ No newline at end of file + return new_bytes diff --git a/gojo/builtins/attributes.mojo b/gojo/builtins/attributes.mojo index c116970..13b15cc 100644 --- a/gojo/builtins/attributes.mojo +++ b/gojo/builtins/attributes.mojo @@ -38,8 +38,12 @@ fn copy(inout target: Bytes, source: Bytes, start: Int = 0) -> Int: """ var count = 0 + # If trying to set a value outside of the bounds of the buffer, use append instead. It will resize the buffer before setting the value. for i in range(len(source)): - target[i + start] = source[i] + if i + start < target.size(): + target[i + start] = source[i] + else: + target.append(source[i]) count += 1 return count diff --git a/gojo/bytes/buffer.mojo b/gojo/bytes/buffer.mojo index 55624e6..c193f26 100644 --- a/gojo/bytes/buffer.mojo +++ b/gojo/bytes/buffer.mojo @@ -146,7 +146,7 @@ struct Buffer( It returns the index where bytes should be written and whether it succeeded.""" var buffer_already_used = len(self.buf) - if n <= cap(self.buf) - buffer_already_used: + if n <= self.buf.size() - buffer_already_used: # FIXME: It seems like reslicing in go can extend the length of the slice. Doens't work like that for my get slice impl. # Instead, just add bytes of len(n) to the end of the buffer for now. # self.buf = self.buf[: l + n] @@ -159,7 +159,7 @@ struct Buffer( """Grows the buffer to guarantee space for n more bytes. It returns the index where bytes should be written. If the buffer can't grow it will panic with ERR_TOO_LARGE.""" - var write_at: Int = self.buf.size() + var write_at: Int = len(self.buf) # If buffer is empty, reset to recover space. if write_at == 0 and self.off != 0: self.reset() @@ -182,15 +182,14 @@ struct Buffer( # slice. We only need m+n <= c to slide, but # we instead var capacity get twice as large so we # don't spend all our time copying. - _ = copy(self.buf, self.buf[self.off :]) elif c > MAX_INT - c - n: panic(ERR_TOO_LARGE) - else: - # Add self.off to account for self.buf[:self.off] being sliced off the front. - var sl = self.buf[self.off :] - self.buf = self.grow_slice(sl, self.off + n) - # self.buf = sl + # TODO: Commented out this branch because growing the slice here and then at the end is redundant? + # else: + # # Add self.off to account for self.buf[:self.off] being sliced off the front. + # # var sl = self.buf[self.off :] + # # self.buf = self.grow_slice(sl, self.off + n) # Restore self.off and len(self.buf). self.off = 0 @@ -231,7 +230,8 @@ struct Buffer( if not ok: write_at = self.grow(len(src)) - return Result(copy(self.buf, src, write_at), None) + var bytes_written = copy(self.buf, src, write_at) + return Result(bytes_written, None) fn write_string(inout self, src: String) -> Result[Int]: """Appends the contents of s to the buffer, growing the buffer as @@ -302,7 +302,7 @@ struct Buffer( # we could rely purely on append to determine the growth rate. c = 2 * cap(b) - var resized_buffer = Bytes(c) + var resized_buffer = Bytes(size=c) _ = copy(resized_buffer, b) # var b2: Bytes = Bytes() # b2._vector.reserve(c) diff --git a/tests/test_bufio.mojo b/tests/test_bufio.mojo index 0a971c6..c71b2c9 100644 --- a/tests/test_bufio.mojo +++ b/tests/test_bufio.mojo @@ -136,6 +136,27 @@ fn test_write() raises: test.assert_equal(str(writer.writer), "0123456789") +fn test_several_writes() raises: + var test = MojoTest("Testing several bufio.Writer.write") + + # Create a new Bytes Buffer Writer and use it to create the buffered Writer + var buf = buffer.new_buffer() + var writer = Writer(buf) + + # Write the content from src to the buffered writer's internal buffer and flush it to the Bytes Buffer Writer. + var src = Bytes("0123456789") + var result = Result(0) + for i in range(500): + result = writer.write(src) + _ = writer.flush() + + test.assert_equal(result.value, 10) + test.assert_equal(len(writer.writer), 5000) + var text = str(writer.writer) + test.assert_equal(text[0], "0") + test.assert_equal(text[4999], "9") + + fn test_write_byte() raises: var test = MojoTest("Testing bufio.Writer.write_byte") @@ -193,6 +214,7 @@ fn main() raises: test_peek() test_discard() test_write() + test_several_writes() test_write_byte() test_write_string() test_read_from()