Add bencoding as well as bdecoding

2 files changed, 127 insertions, 0 deletions

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);
+}