diff options
| -rw-r--r-- | build.zig | 2 | ||||
| -rw-r--r-- | src/anywriter.zig | 77 | ||||
| -rw-r--r-- | src/bencode.zig | 50 |
3 files changed, 128 insertions, 1 deletions
@@ -27,7 +27,7 @@ pub fn build(b: *std.Build) void { .target = target, .optimize = optimize, }); - //b.installArtifact(unit_tests); // Useful if you want to debug the test binary with lldb or something + b.installArtifact(unit_tests); // Useful if you want to debug the test binary with lldb or something const run_unit_tests = b.addRunArtifact(unit_tests); const test_step = b.step("test", "Run unit tests"); diff --git a/src/anywriter.zig b/src/anywriter.zig new file mode 100644 index 0000000..7018ac6 --- /dev/null +++ b/src/anywriter.zig @@ -0,0 +1,77 @@ +//! std.io.Writer drop in replacement that uses function pointer instead of +//! compile time construct. Analogous to std.io.AnyReader. + +const std = @import("std"); +const AnyWriter = @This(); +pub const Error = anyerror; + +context: *const anyopaque, +writeFn: *const fn (context: *const anyopaque, buffer: []const u8) anyerror!usize, + +// Must create a wrapper around the original writer in order to add a typeErasedWriteFn +// and keep a reference to it alive until we're done writing to it +// Usage might look like this: +// var wrap = AnyWriter.wrapper(myWriter); +// var aw = wrap.any(); +pub fn wrapper(writer: anytype) Wrapper(@TypeOf(writer)) { + return .{ .ww = writer }; +} +pub fn Wrapper(comptime writerType: type) type { + return struct { + ww: writerType, + fn typeErasedWriteFn(context: *const anyopaque, buffer: []const u8) anyerror!usize { + const self: *const @This() = @alignCast(@ptrCast(context)); + return self.ww.write(buffer); + } + pub fn any(self: *@This()) AnyWriter { + return .{ + .context = @ptrCast(self), + .writeFn = &typeErasedWriteFn, + }; + } + }; +} + +pub fn write(self: AnyWriter, bytes: []const u8) anyerror!usize { + return self.writeFn(self.context, bytes); +} + +pub fn writeAll(self: AnyWriter, bytes: []const u8) anyerror!void { + var index: usize = 0; + while (index != bytes.len) { + index += try self.write(bytes[index..]); + } +} + +pub fn print(self: AnyWriter, comptime format: []const u8, args: anytype) anyerror!void { + return std.fmt.format(self, format, args); +} + +pub fn writeByte(self: AnyWriter, byte: u8) anyerror!void { + const array = [1]u8{byte}; + return self.writeAll(&array); +} + +pub fn writeByteNTimes(self: AnyWriter, byte: u8, n: usize) anyerror!void { + var bytes: [256]u8 = undefined; + @memset(bytes[0..], byte); + + var remaining: usize = n; + while (remaining > 0) { + const to_write = @min(remaining, bytes.len); + try self.writeAll(bytes[0..to_write]); + remaining -= to_write; + } +} + +pub inline fn writeInt(self: AnyWriter, comptime T: type, value: T, endian: std.builtin.Endian) anyerror!void { + var bytes: [@divExact(@typeInfo(T).Int.bits, 8)]u8 = undefined; + std.mem.writeInt(std.math.ByteAlignedInt(@TypeOf(value)), &bytes, value, endian); + return self.writeAll(&bytes); +} + +pub fn writeStruct(self: AnyWriter, value: anytype) anyerror!void { + // Only extern and packed structs have defined in-memory layout. + comptime std.debug.assert(@typeInfo(@TypeOf(value)).Struct.layout != .Auto); + return self.writeAll(std.mem.asBytes(&value)); +} diff --git a/src/bencode.zig b/src/bencode.zig index b78ee9f..44dec0d 100644 --- a/src/bencode.zig +++ b/src/bencode.zig @@ -2,6 +2,7 @@ //! See specification here https://wiki.theory.org/BitTorrentSpecification#Bencoding const std = @import("std"); +const AnyWriter = @import("anywriter.zig"); pub const Error = error.Malformatted || std.io.AnyReader.Error; @@ -12,6 +13,41 @@ pub const BValue = union(enum) { list: std.ArrayList(BValue), dict: std.StringArrayHashMap(BValue), + pub fn bencode(self: *const BValue, base_writer: anytype) !void { + var wrap = AnyWriter.wrapper(base_writer); + var writer = wrap.any(); + try self.bencodeInner(writer); + } + + // Note: uses defined types only to avoid trying to recursively evaulate this function + // at compile time, otherwise we run into https://github.com/ziglang/zig/issues/13724 + fn bencodeInner(self: *const BValue, writer: AnyWriter) !void { + switch (self.*) { + .string => |s| { + try std.fmt.format(writer, "{}:{s}", .{ s.len, s }); + }, + .list => |l| { + try writer.writeByte('l'); + for (l.items) |i| { + try i.bencodeInner(writer); + } + try writer.writeByte('e'); + }, + .dict => |d| { + try writer.writeByte('d'); + var it = d.iterator(); + while (it.next()) |entry| { + try std.fmt.format(writer, "{}:{s}", .{ entry.key_ptr.*.len, entry.key_ptr.* }); + try entry.value_ptr.*.bencodeInner(writer); + } + try writer.writeByte('e'); + }, + .int => |i| { + try std.fmt.format(writer, "i{}e", .{i}); + }, + } + } + pub fn deinit(self: *BValue, a: std.mem.Allocator) void { switch (self.*) { .string => |s| { @@ -237,3 +273,17 @@ test "nested structure" { try std.testing.expectEqualDeep(v2.*.list.items[1], BValue{ .int = 456 }); try std.testing.expectEqualStrings("nest", v2.*.list.items[2].list.items[0].string); } + +test "round trip" { + var a = std.testing.allocator; + const in = "d5:hello5:world2:hili123ei456el4:nesteee"; + var bval = try bdecodeBuf(a, in); + defer bval.deinit(a); + var bw = std.ArrayList(u8).init(a); + defer bw.deinit(); + var writer = bw.writer(); + try bval.bencode(writer); + var out = try bw.toOwnedSlice(); + defer a.free(out); + try std.testing.expectEqualStrings(in, out); +} |