bug fix, code clean-up

Fixed bug in stream parser

Author: mtolman

src/column.zig

@@ -277,6 +277,9 @@ pub fn Parser(comptime Reader: type) type {
         /// Returns the next row in the CSV file
         /// Heap memory is owend by the Row, so call Row.deinit()
         pub fn next(self: *@This()) ?Row {
+            if (self._done) {
+                return null;
+            }
             return self.nextImpl() catch |err| {
                 switch (err) {
                     Error.EndOfInput => {
@@ -299,16 +302,22 @@ pub fn Parser(comptime Reader: type) type {
             }
 
             if (self._buffer.done()) {
-                return null;
+                return Error.EndOfInput;
             }
 
+            std.debug.assert(!self._done);
+
             var row = Row{
                 ._fields = std.ArrayList(RowField).init(self._allocator),
                 ._bytes = std.ArrayList(u8).init(self._allocator),
             };
 
-            // Doing this defer since we don't want to access a deinitialized
-            // row when there's an error
+            // Doing this defer since so we can track our max field for
+            // fewer allocation in the future
+            // Not worried about excess memory as much since rows should be
+            // short-lived
+            // Rows that are long-lived should have a `clone()` called which
+            // will trim excess memory on copy
             defer {
                 self._row_field_count = @max(self._row_field_count, row._fields.items.len);
                 self._row_byte_count = @max(self._row_byte_count, row._bytes.items.len);

src/common.zig

@@ -28,9 +28,13 @@ const std = @import("std");
 /// Unquotes a quoted CSV field
 /// Assumes the field is a valid CSV field
 pub fn unquoteQuoted(data: []const u8) []const u8 {
-    if (data.len > 1 and data[0] == '"') {
+    if (data.len > 2 and data[0] == '"') {
+        // If we get here, we should be working on quoted CSV data
+        std.debug.assert(data[data.len - 1] == '"');
         return data[1..(data.len - 1)];
     } else {
+        // If we get here, we should be working on unquoted CSV data
+        std.debug.assert(std.mem.indexOf(u8, data, "\"") == null);
         return data;
     }
 }

src/stream.zig

@@ -79,6 +79,9 @@ pub fn Parser(
 
         /// Consume a character and move forward our reader
         fn consume(self: *@This()) ReadError!void {
+            // While we should behave properly, we shouldn't be consuming
+            // when we're done
+            std.debug.assert(!self.done());
             self._cur = self._next;
 
             // This checks to see if we already hit the end of the stream
@@ -128,31 +131,40 @@ pub fn Parser(
                 return;
             }
 
+            // When we start next we always are at the start or middle of a row
             self._flags._row_end = false;
 
             // We won't technically hit an infinite loop,
             // but we practically will since this is a lot
             const MAX_ITER = self._opts.max_iter;
             var index: usize = 0;
-            for (0..(MAX_ITER + 1)) |i| {
-                index = i;
-
+            while (index < MAX_ITER) : (index += 1) {
                 if (self._flags._in_quote) {
                     // Handle quoted strings
                     if (self.current()) |cur| {
                         switch (cur) {
                             '"' => {
-                                if (self.peek() == ',' or self.peek() == '\n' or self.peek() == '\r') {
+                                if (self.peek()) |p| {
+                                    switch (p) {
+                                        ',', '\r', '\n' => {
+                                            self._flags._in_quote = false;
+                                            try self.consume();
+                                            continue;
+                                        },
+                                        '"' => {
+                                            try self.consume();
+                                            try self.consume();
+                                            try writer.writeByte('"');
+                                        },
+                                        else => {
+                                            return CsvReadError.QuotePrematurelyTerminated;
+                                        },
+                                    }
+                                } else {
                                     self._flags._in_quote = false;
                                     try self.consume();
                                     continue;
                                 }
-                                if (self.peek() != '"') {
-                                    return CsvReadError.QuotePrematurelyTerminated;
-                                }
-                                try self.consume();
-                                try self.consume();
-                                try writer.writeByte('"');
                             },
                             else => |c| {
                                 try writer.writeByte(c);
@@ -380,3 +392,33 @@ test "crlf,\" at 63" {
 
     try testing.expectEqual(fieldCount, cnt);
 }
+
+test "End with quote" {
+    const testing = @import("std").testing;
+    var buff = std.ArrayList(u8).init(std.testing.allocator);
+    defer buff.deinit();
+
+    var input = std.io.fixedBufferStream("\"hello, world\"");
+
+    const fieldCount = 1;
+
+    const reader = input.reader();
+    var stream = init(
+        @TypeOf(reader),
+        @TypeOf(buff.writer()),
+        reader,
+        .{},
+        .{},
+    );
+
+    var cnt: usize = 0;
+    while (!stream.done()) {
+        try stream.next(buff.writer());
+        defer {
+            buff.clearRetainingCapacity();
+        }
+        cnt += 1;
+    }
+
+    try testing.expectEqual(fieldCount, cnt);
+}

src/stream_fast.zig

@@ -45,8 +45,7 @@ pub fn Parser(comptime Reader: type, comptime Writer: type) type {
         }
 
         pub fn atRowEnd(self: *const @This()) bool {
-            const res = self.done() or self._row_end;
-            return res;
+            return self.done() or self._row_end;
         }
 
         /// Returns if a parser is done
@@ -90,12 +89,18 @@ pub fn Parser(comptime Reader: type, comptime Writer: type) type {
                 return;
             }
 
+            // Don't read from a reader when we're done
+            std.debug.assert(!self.done());
             const amt = try self._reader.readAll(&self._state.next_bytes);
             self._state.next_bytes_len = amt;
         }
 
         /// Gets the next CSV field
         pub fn next(self: *@This(), writer: Writer) Error!void {
+            self._state.field_start = false;
+            if (self.done()) {
+                return;
+            }
             self.nextImpl(writer) catch |err| {
                 self._state._erred = true;
                 return err;
@@ -104,24 +109,20 @@ pub fn Parser(comptime Reader: type, comptime Writer: type) type {
 
         /// Gets the next CSV field
         fn nextImpl(self: *@This(), writer: Writer) !void {
-            self._state.field_start = false;
+            std.debug.assert(!self.done());
             // lazy init our parser
             if (self._state._need_init) {
                 // Consume twice to populate our cur and next registers
                 try self.nextChunk();
                 self._state.cur_bytes_pos = 64;
                 std.debug.assert(!self._state._need_init);
-            } else if (self.done()) {
-                return;
             }
 
             self._row_end = false;
 
             const MAX_ITER = self._opts.max_iter;
             var index: usize = 0;
-            for (0..(MAX_ITER + 1)) |i| {
-                index = i;
-
+            while (index < MAX_ITER) : (index += 1) {
                 // If we have a field to print, print it
                 if (self._state.field_separators != 0) {
                     const chunk_end = @ctz(self._state.field_separators);
@@ -267,10 +268,9 @@ pub fn Parser(comptime Reader: type, comptime Writer: type) type {
         fn quotedRegions(m: u64) u64 {
             var x: u64 = m;
             var res: u64 = x;
-            while (x != 0) {
+            while (x != 0) : (x = x & x - 1) {
                 const x1: u64 = @bitCast(-%@as(i64, @bitCast(x)));
                 res = res ^ (x1 ^ x);
-                x = x & x - 1;
             }
             return res;
         }
@@ -613,3 +613,31 @@ test "crlfa, at 63" {
 
     try testing.expectEqual(fieldCount, cnt);
 }
+
+test "End with quote" {
+    const testing = @import("std").testing;
+    var buff = std.ArrayList(u8).init(std.testing.allocator);
+    defer buff.deinit();
+
+    var input = std.io.fixedBufferStream("\"hello, world\"");
+
+    const fieldCount = 1;
+
+    const reader = input.reader();
+    var stream = init(
+        reader,
+        @TypeOf(buff.writer()),
+        .{},
+    );
+
+    var cnt: usize = 0;
+    while (!stream.done()) {
+        try stream.next(buff.writer());
+        defer {
+            buff.clearRetainingCapacity();
+        }
+        cnt += 1;
+    }
+
+    try testing.expectEqual(fieldCount, cnt);
+}