Stack allocated strings

When coding various “scripty” things in Zig, I often need to format short strings. I hate using gpa for that purpose, as I always mess up freeing stuff, especially when there’s error handling involved.

And, most of the time, I actually know that the string is going to be short, even if I can’t immediately say up-front just how short exactly. Like, if I am converting u64 to string, the result would surely be shorter than 1KiB.

What is the most convenient API for solving this “short strings on the stack” problem?

Here’s what I came up with:

const std = @import("std");
const assert = std.debug.assert;

pub fn example(runtime_value: u32) void {
    const arg = stack_print("--my-flag={}", .{runtime_value}).text();
    std.debug.print("arg={s}\n", .{arg});
}

/// Prints to stack-allocated buffer, guaranteeing static upper bound.
pub fn stack_print(comptime fmt: []const u8, args: anytype) StackBufferType(stack_size_max(fmt, @TypeOf(args))) {
    var result: StackBufferType(stack_size_max(fmt, @TypeOf(args))) = .{};
    const text = std.fmt.bufPrint(&result.buffer, fmt, args) catch |err| switch (err) {
        error.NoSpaceLeft => unreachable,
    };
    result.size = @intCast(text.len);
    return result;
}

fn stack_size_max(comptime fmt: []const u8, Args: type) u16 {
    var args_worst_case: Args = undefined;
    for (@typeInfo(Args).@"struct".fields, 0..) |field, index| {
        const arg_worst_case = switch (field.type) {
            u8, u16, u32, u64, u128 => std.math.maxInt(field.type),
            else => @compileError("array_print: unsupported type: " ++ @typeName(field.type)),
        };
        args_worst_case[index] = arg_worst_case;
    }
    const size_max = std.fmt.count(fmt, args_worst_case);
    assert(size_max <= std.math.maxInt(u16));
    assert(size_max <= 4096); // Safety check, just so that you don't alloc a gigantic buffer. 
    return @intCast(size_max);
}


fn StackBufferType(comptime size_max: u16) type {
    return struct {
        buffer: [size_max]u8 = undefined,
        size: u16 = 0,

        const StackString = @This();

        pub fn text(stack_string: *const StackString) []const u8 {
            assert(stack_string.size <= size_max);
            return stack_string.buffer[0..stack_string.size];
        }
    };
}

pub fn main() void {
    example(92);
}

Can you do better? Are there any problems in the above?

7 Likes

Neat trick, but I have one question:

Is this guaranteed to work or does it depend on some undocumented behavior? By my intuition, I’d have expected the StackBufferType to go out of scope when you return from text(). Is arg guaranteed to stay valid on the line after?

2 Likes

You were quick to spot that! This is indeed the biggest o_O of the API!

I think this legitimate. The current rules are:

  • Named locals live until the end of declaring block.
  • Unnamed temporaries live until the end of the function, or the back edge of the enclosing loop

Source: #compiler > Lifetime of locals @ :speech_balloon:

2 Likes

One option that might be for stackPrint to return [stackSizeMax(fmt, @TypeOf(args)):0]u8 instead. You’d set the zero sentinel after the printBuf length, if it doesn’t use the full pre built in sentinel length. Though, as I type this out, I’m not sure it’ll be obvious if code is correctly only using [*:0]const u8 when passing to functions which want strings.

(Also don’t forget if you add support for negative numbers to use minInt.)

1 Like

I personally made an allocator I call “Dynamic stack allocator” specifically for problems like that:

pub fn main() !void {
    const sf = DSA.captureStackFrame();
    defer DSA.rewindStackFrame(sf);

    const myString = try std.fmt.allocPrint(DSA.tostdAllocator, "Hello world!!!", .{});
    _ = myString;
}

The implementation is just a big buffer, and if it runs out of space then it starts allocating pages with a page allocator.

Why not use std.heap.BufferFirstAllocator? (known as std.heap.stackFallback in older versions of Zig)

1 Like

First time I hear about it, maybe I just reinvented the wheel.
It’s not a huge under taking to write this allocator and I use it a lot, so it’s probably best I wrote it myself.

That’s a nice little snippet you got there.

I’ve experimented with a similar idea but wanted to be able to format both utf8 and windows wtf16 strings, and also specify substring parts with a max length. It’s defined by a sequence of parts like this:

pub fn StringPart(comptime encoding: Encoding) type {
    return union(enum) {
        static: []const encoding.Char(),
        runtime_value: struct { name: [:0]const u8, type: type },
        runtime_utf8: struct { name: [:0]const u8, max_len: usize, max_wtf16: ?usize = null },
        runtime_wtf16: struct { name: [:0]const u8, max_len: usize, max_utf8: ?usize = null },
    };
}
pub const Encoding = enum {
    utf8,
    wtf16,
    pub fn Char(encoding: Encoding) type {
        return switch (encoding) {
            .utf8 => u8,
            .wtf16 => u16,
        };
    }
};

Here’s an example:

// definition
const max_appdata_len = 400;
const CrashesDirPath = MaxString(.utf8, .yes_sentinel, &.{
    // each element is a StringPart
    .{ .runtime_utf8 = .{ .name = "localappdata", .max_len = max } },
    .{ .static = "/" },
    .{ .runtime_value = .{ .name = "instance", .type = u8 } },
    .{ .static = "/crashes" },
});

// usage
const crashes_dir_path = CrashesDirPath.format(.{
    .localappdata = appdata.getPath(),
    .instance = instance.value,
});
// crashes_dir_path is a BoundedString which is basically
// an array and a len

Now that I’m looking at this, I suppose the MaxString function wouldn’t need to specify utf8 nor whether it requires a sentinel, and those could just be separate format methods (or comptime parameters) instead. Anyway, here’s the source:

const Sentinel = enum { no_sentinel, yes_sentinel };

pub const Encoding = enum {
    utf8,
    wtf16,
    pub fn Char(encoding: Encoding) type {
        return switch (encoding) {
            .utf8 => u8,
            .wtf16 => u16,
        };
    }
};

pub fn StringPart(comptime encoding: Encoding) type {
    return union(enum) {
        static: []const encoding.Char(),
        runtime_value: struct { name: [:0]const u8, type: type },
        runtime_utf8: struct { name: [:0]const u8, max_len: usize, max_wtf16: ?usize = null },
        runtime_wtf16: struct { name: [:0]const u8, max_len: usize, max_utf8: ?usize = null },
    };
}

fn maxLen(comptime encoding: Encoding, Type: type) usize {
    switch (@typeInfo(Type)) {
        .int => |info| {
            if (info.bits == 8 and info.signedness == .unsigned) return 3;
            if (info.bits == 16 and info.signedness == .unsigned) return 6;
            if (info.bits == 32 and info.signedness == .unsigned) return 10;
            if (info.bits == 32 and info.signedness == .signed) return 11;
            if (info.bits == 64 and info.signedness == .unsigned) return 20;
        },
        .@"enum" => |info| {
            var max_name: usize = 0;
            for (info.fields) |field| {
                const name_len = switch (encoding) {
                    .utf8 => field.name.len,
                    .wtf16 => std.unicode.wtf8ToWtf16LeStringLiteral(field.name).len,
                };
                max_name = @max(max_name, name_len);
            }
            return max_name;
        },
        else => {},
    }
    @compileError("todo: implement maxLen for type " ++ @typeName(Type));
}

fn FmtArg(comptime Type: type) type {
    switch (@typeInfo(Type)) {
        .int => return Type,
        .@"enum" => return [:0]const u8,
        else => {},
    }
    @compileError("todo: implement FmtArg for type " ++ @typeName(Type));
}
fn fmtArg(comptime Type: type, value: Type) FmtArg(Type) {
    switch (@typeInfo(Type)) {
        .int => return value,
        .@"enum" => return @tagName(value),
        else => {},
    }
    @compileError("todo: implement fmtArg for type " ++ @typeName(Type));
}
fn fmtSpec(comptime Type: type) [:0]const u8 {
    return switch (@typeInfo(Type)) {
        .@"enum" => return "{s}",
        else => "{}",
    };
}

pub fn MaxString(comptime encoding: Encoding, sentinel: Sentinel, comptime parts: []const StringPart(encoding)) type {
    var struct_fields: [parts.len]std.builtin.Type.StructField = undefined;
    var field_count: usize = 0;
    inline for (parts) |part| {
        const maybe_field: ?std.builtin.Type.StructField = switch (part) {
            .static => null,
            .runtime_value => |d| .{
                .name = d.name,
                .type = d.type,
                .default_value_ptr = null,
                .is_comptime = false,
                .alignment = @alignOf(d.type),
            },
            .runtime_utf8 => |d| .{
                .name = d.name,
                .type = []const u8,
                .default_value_ptr = null,
                .is_comptime = false,
                .alignment = @alignOf([]const u8),
            },
            .runtime_wtf16 => |d| .{
                .name = d.name,
                .type = []const u16,
                .default_value_ptr = null,
                .is_comptime = false,
                .alignment = @alignOf([]const u16),
            },
        };
        if (maybe_field) |field| {
            struct_fields[field_count] = field;
            field_count += 1;
        }
    }
    const FormatArgs = @Type(std.builtin.Type{
        .@"struct" = .{
            .layout = .auto,
            .fields = struct_fields[0..field_count],
            .decls = &.{},
            .is_tuple = false,
        },
    });

    return struct {
        pub const max_len = blk: {
            var len: usize = 0;
            for (parts) |part| {
                len += switch (part) {
                    .static => |s| s.len,
                    .runtime_value => |d| maxLen(encoding, d.type),
                    .runtime_utf8 => |d| switch (encoding) {
                        .utf8 => d.max_len,
                        .wtf16 => @panic("todo"),
                    },
                    .runtime_wtf16 => |d| switch (encoding) {
                        .utf8 => @panic("todo"),
                        .wtf16 => d.max_len,
                    },
                };
            }
            break :blk len;
        };

        pub fn format(args: FormatArgs) BoundedString(encoding, sentinel, max_len) {
            var result: BoundedString(encoding, sentinel, max_len) = .{
                .buffer = undefined,
                .len = 0,
            };

            var max_possible_len: usize = 0;
            inline for (parts) |part| {
                switch (part) {
                    .static => |s| {
                        @memcpy(result.buffer[result.len..][0..s.len], s);
                        result.len += s.len;
                        max_possible_len += s.len;
                    },
                    .runtime_value => |d| {
                        const dst = result.buffer[result.len..];
                        const max_value_len = comptime maxLen(encoding, d.type);
                        std.debug.assert(dst.len >= max_value_len);
                        switch (encoding) {
                            .utf8 => {
                                const len = (std.fmt.bufPrint(dst, fmtSpec(d.type), .{fmtArg(d.type, @field(args, d.name))}) catch unreachable).len;
                                std.debug.assert(len <= max_value_len);
                                result.len += len;
                            },
                            .wtf16 => {
                                // TODO: we'll need more len if the formatted value contains
                                //       utf8 chars that need more than 1 byte
                                var buf: [max_value_len]u8 = undefined;
                                const len = (std.fmt.bufPrint(&buf, fmtSpec(d.type), .{fmtArg(d.type, @field(args, d.name))}) catch unreachable).len;
                                std.debug.assert(len <= max_value_len);
                                const wtf16_len = std.unicode.wtf8ToWtf16Le(dst, buf[0..len]) catch unreachable;
                                // for now we'll assume all formatted values are made up of 1-byte utf8 chars
                                std.debug.assert(wtf16_len == len);
                                result.len += len;
                            },
                        }
                        max_possible_len += max_value_len;
                    },
                    .runtime_utf8 => |d| switch (encoding) {
                        .utf8 => {
                            const s = @field(args, d.name);
                            std.debug.assert(s.len <= d.max_len);
                            @memcpy(result.buffer[result.len..][0..s.len], s);
                            result.len += s.len;
                            max_possible_len += d.max_len;
                        },
                        .wtf16 => @panic("todo"),
                    },
                    .runtime_wtf16 => |d| switch (encoding) {
                        .utf8 => @panic("todo"),
                        .wtf16 => {
                            const s = @field(args, d.name);
                            std.debug.assert(s.len <= d.max_len);
                            @memcpy(result.buffer[result.len..][0..s.len], s);
                            result.len += s.len;
                            max_possible_len += d.max_len;
                        },
                    },
                }
            }
            std.debug.assert(max_possible_len == max_len);
            switch (sentinel) {
                .no_sentinel => return result,
                .yes_sentinel => {
                    result.buffer[result.len] = 0;
                    return result;
                },
            }
            return result;
        }
    };
}

pub fn BoundedString(comptime encoding: Encoding, comptime sentinel: Sentinel, comptime capacity: usize) type {
    return struct {
        const total_capacity = capacity + switch (sentinel) {
            .no_sentinel => 0,
            .yes_sentinel => 1,
        };
        buffer: [total_capacity]encoding.Char(),
        len: usize,

        const Self = @This();
        pub fn slice(self: *const Self) switch (sentinel) {
            .no_sentinel => []const encoding.Char(),
            .yes_sentinel => [:0]const encoding.Char(),
        } {
            return switch (sentinel) {
                .no_sentinel => self.buffer[0..self.len],
                .yes_sentinel => self.buffer[0..self.len :0],
            };
        }

        pub fn format(self: Self, writer: *std.Io.Writer) error{WriteFailed}!void {
            switch (encoding) {
                .utf8 => try writer.writeAll(self.slice()),
                .wtf16 => try writer.print("{}", .{std.unicode.fmtUtf16Le(self.slice())}),
            }
        }
    };
}

const builtin = @import("builtin");
const std = @import("std");

In my C header libraries, whenever I need to copy an externally provided string I store them in a ‘value struct’ like this:

typedef struct {
    char buf[_SG_STRING_SIZE];
} _bla_str_t;

…and at least a creation function which takes a C string pointer and returns a str_t value, and which may clamp the string:

str_t _bla_make_string(const char* c_str);

Upside is that it saves you from a lot of hassle and potential memory corruption issues and you can directly copy those value-string-structs by simple assignment, downside is of course that you need to decide on a max string size, and a too big string size may waste memory. But in general I’ve used this idea for years, it works well and the upside outweigh the downsides (IMHO).

std.BoundedArray did nothing wrong.

3 Likes

I use TempAllocator.
Long living arena that I can reset at any point of time.