Compiler Explorer

Source code

const std = @import("std");
const testing = std.testing;
const Allocator = std.mem.Allocator;

const AMP_KEY_LIMIT: usize = 0xff;

const ParseState = enum {
    KeyLenRead,
    KeyLenRead2,
    KeyDataRead,
    ValueLenRead,
    ValueLenRead2,
    ValueDataRead,
};

const ParseError = error{
    NoError,
    KeyTooLong,
    EmptyKey,
    ValueTooLong,
    Unsupported,
    PartialData,
};

const AmpBox = std.hash_map.StringHashMap([]const u8);

const AmpDecoder = struct {
    alloc: *Allocator,
    state: ParseState,
    box_buffer: AmpBoxFifo,

key_len: usize,
    key_data: std.ArrayList(u8),
    val_len: usize,
    val_data: std.ArrayList(u8),
    err: ParseError,
    key_data_fetched: usize,
    val_data_fetched: usize,
    partial_box: bool,

pub const AmpBoxFifo = std.fifo.LinearFifo(AmpBox, .Dynamic);

pub fn init(allocator: *Allocator) AmpDecoder {
        return AmpDecoder{
            .alloc = allocator,
            .state = .KeyLenRead,
            .box_buffer = AmpBoxFifo.init(allocator),
            .key_data = std.ArrayList(u8).init(allocator), // TODO: use ArrayListUnmanaged
            .key_len = 0,
            .val_len = 0,
            .val_data = std.ArrayList(u8).init(allocator), // TODO: use ArrayListUnmanaged
            .val_data_fetched = 0,
            .err = ParseError.NoError,
            .key_data_fetched = 0,
            .partial_box = false,
        };
    }

pub fn deinit(self: *AmpDecoder) void {
        while (self.box_buffer.readItem()) |*item| {
            item.deinit();
        }
        self.box_buffer.deinit();
        self.key_data.deinit();
        self.val_data.deinit();
    }

// pipe data in chunks, and receive boxes. pipe empty array in to
    // flush final boxes out, if any.
    pub fn decode_box(self: *AmpDecoder, bytes: [*]const u8, size: usize) (ParseError || std.mem.Allocator.Error)!?AmpBox {
        if (size == 0) {
            if (self.box_buffer.readItem()) |item| {
                return item;
            }
        }
        const slice = bytes[0..size];

var idx: usize = 0;
        var bytes_remaining: usize = 0;

while (idx < size) {
            bytes_remaining = size - idx;

switch (self.state) {
                .KeyLenRead => {
                    std.debug.print("KeyLenRead\n", .{});
                    if (!self.partial_box) {
                        std.debug.print("writing new AmpBox to box_buffer\n", .{});
                        try self.box_buffer.writeItem(AmpBox.init(self.alloc));
                        self.partial_box = true;
                    }

if (slice[idx] != 0) {
                        self.err = ParseError.KeyTooLong;
                        idx += 1;
                        return self.err;
                    }

self.key_len = 0;
                    idx += 1;
                    self.state = .KeyLenRead2;
                },
                .KeyLenRead2 => {
                    std.debug.print("KeyLenRead2\n", .{});
                    self.key_len = slice[idx];
                    idx += 1;

if (self.key_len == 0) {
                        self.state = .KeyLenRead;
                        std.debug.print("RETURNING\n", .{});
                        return self.box_buffer.readItem();
                    } else {
                        self.state = .KeyDataRead;
                    }
                },
                .KeyDataRead => {
                    std.debug.print("KeyDataRead\n", .{});
                    const bytes_needed = self.key_len - self.key_data_fetched;
                    if (bytes_remaining < bytes_needed) {
                        try self.key_data.appendSlice(slice[self.key_data_fetched..bytes_remaining]);
                        idx += bytes_remaining;
                        self.key_data_fetched += bytes_remaining;
                    } else {
                        try self.key_data.appendSlice(slice[self.key_data_fetched..bytes_needed]);
                        idx += bytes_needed;
                        self.key_data_fetched += bytes_needed;

try self.key_data.append(0);
                        self.state = .ValueLenRead;
                    }
                },
                .ValueLenRead => {
                    std.debug.print("ValueLenRead\n", .{});
                    self.val_len = slice[idx];
                    self.val_len = self.val_len << 8;
                    idx += 1;
                    self.state = .ValueLenRead2;
                },
                .ValueLenRead2 => {
                    std.debug.print("ValueLenRead2\n", .{});
                    self.val_len |= slice[idx];
                    idx += 1;
                    self.state = .ValueDataRead;
                },
                .ValueDataRead => {
                    std.debug.print("ValueDataRead\n", .{});
                    const bytes_needed = self.val_len - self.val_data_fetched;
                    if (bytes_remaining < bytes_needed) {
                        try self.val_data.appendSlice(slice[self.val_data_fetched..bytes_remaining]);
                        idx += bytes_remaining;
                        self.val_data_fetched += bytes_remaining;
                    } else {
                        try self.val_data.appendSlice(slice[self.val_data_fetched..bytes_needed]);
                        idx += bytes_needed;
                        self.val_data_fetched += bytes_needed;

const box = blk: { // mutable peek on fifo
                            const offset = self.box_buffer.count - 1;
                            var index = self.box_buffer.head + offset;
                            index %= self.box_buffer.buf.len;
                            break :blk &self.box_buffer.buf[index];
                        };
                        try box.put(self.key_data.items[0..], self.val_data.items[0..]);

std.debug.print("HERE!\n", .{});

if (self.err != ParseError.NoError) {
                            return self.err;
                        }

self.key_len = 0;
                        self.val_len = 0;
                        self.key_data_fetched = 0;
                        self.val_data_fetched = 0;

self.state = .KeyLenRead;
                    }
                },
            }
        }
        // if we haven't returned a box yet, go ahead and return the oldest queued one
        // (the head of the buffer is the "current" box and is still being filled in and
        //  should not be returned yet)
        if (self.box_buffer.count > 1) {
            return self.box_buffer.readItem();
        }
        return null;
    }
};

pub fn main() !void { // if you use try in main it needs to mention errors in the type with !
    var buffer: [256]u8 = undefined; // placeholder
    const alloc = std.heap.page_allocator;
    var a = AmpDecoder.init(alloc);
    defer a.deinit();

if (try a.decode_box(&buffer, 0)) |*box| {
        defer box.deinit();
        // probably want to use the value somehow
    }
}

test "blah" {
    const b1 = &[_]u8{
        0x00, 0x07, // key of length 7 follows
        0x5F, 0x61, 0x6E, 0x73, 0x77, 0x65, 0x72, // "_answer"
        0x00, 0x02, // value of length 2 follows
        0x32, 0x33, // "23"
        0x00, 0x05, // key of length 5 follows
        0x74, 0x6F, 0x74, 0x61, 0x6C, // "total"
        0x00, 0x02, // value of length 2 follows
        0x39, 0x34, // "94"
        0x00, 0x00, // empty key terminates the box
    };

const badKeySizeBox = &[_]u8{
        0x04, 0x07, // key length > 255
        0x5F, 0x61,
        0x6E, 0x73,
        0x77, 0x65,
        0x72, 0x00,
        0x02, 0x32,
        0x33, 0x00,
        0x05, 0x74,
        0x6F, 0x74,
        0x61, 0x6C,
        0x00, 0x02,
        0x39, 0x34,
        0x00, 0x00,
    };

const emptyBox = &[_]u8{ 0x00, 0x00 };

var dec = AmpDecoder.init(testing.allocator);
    defer dec.deinit();

var b = (try dec.decode_box(b1, b1.len)).?;
    defer b.deinit();

std.debug.print("\n\n{}\n\n", .{b.count()});
    testing.expectEqual(@as(usize, 2), b.count());
}