synthetic: make bytes generation more flexible

jcollie · jcollie · commit 80c250ab16f9 · 2025-10-14T09:36:02.000-05:00
diff --git a/src/synthetic/Bytes.zig b/src/synthetic/Bytes.zig
@@ -1,4 +1,4 @@
-/// Generates bytes.
+//! Generates bytes.
 const Bytes = @This();
 
 const std = @import("std");
@@ -7,54 +7,180 @@ const Generator = @import("Generator.zig");
 /// Random number generator.
 rand: std.Random,
 
-/// The minimum and maximum length of the generated bytes. The maximum
-/// length will be capped to the length of the buffer passed in if the
-/// buffer length is smaller.
-min_len: usize = 1,
-max_len: usize = std.math.maxInt(usize),
+/// The minimum and maximum length of the generated bytes.
+_min_len: usize = 1,
+_max_len: usize = std.math.maxInt(usize),
 
 /// The possible bytes that can be generated. If a byte is duplicated
 /// in the alphabet, it will be more likely to be generated. That's a
 /// side effect of the generator, not an intended use case.
-alphabet: ?[]const u8 = null,
+_alphabet: ?[]const u8 = null,
 
 /// Predefined alphabets.
 pub const Alphabet = struct {
-    pub const ascii = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*()_+-=[]{}|;':\\\",./<>?`~";
+    fn _generate(comptime func: fn (u8) bool) []const u8 {
+        @setEvalBranchQuota(3000);
+        var count = 0;
+        for (0..256) |c| {
+            if (func(c)) count += 1;
+        }
+        var _alphabet: [count]u8 = undefined;
+        var i = 0;
+        for (0..256) |c| {
+            if (func(c)) {
+                _alphabet[i] = c;
+                i += 1;
+            }
+        }
+        const result = _alphabet;
+        return &result;
+    }
+
+    fn _ascii(c: u8) bool {
+        return switch (c) {
+            ' ' => false,
+            else => std.ascii.isPrint(c),
+        };
+    }
+
+    pub const ascii = _generate(_ascii);
+    pub const alphanumeric = _generate(std.ascii.isAlphanumeric);
+    pub const alphabetic = _generate(std.ascii.isAlphabetic);
+
+    fn _kv(c: u8) bool {
+        return switch (c) {
+            std.ascii.control_code.esc, std.ascii.control_code.bel, ';', '=' => false,
+            else => std.ascii.isPrint(c),
+        };
+    }
+
+    /// The alphabet for random bytes in OSC key/value pairs (omitting 0x1B,
+    /// 0x07, ';', '=').
+    pub const kv = _generate(_kv);
+
+    fn _osc(c: u8) bool {
+        return switch (c) {
+            std.ascii.control_code.esc, std.ascii.control_code.bel => false,
+            else => true,
+        };
+    }
+
+    /// The alphabet for random bytes in OSCs (omitting 0x1B and 0x07).
+    pub const osc = _generate(_osc);
+
+    pub const Keys = keys: {
+        const decls = @typeInfo(Alphabet).@"struct".decls;
+        var count = 0;
+        for (decls) |decl| {
+            if (std.mem.eql(u8, decl.name, "Keys")) continue;
+            if (decl.name[0] == '_') continue;
+            count += 1;
+        }
+        var fields: [count]std.builtin.Type.EnumField = undefined;
+        var i = 0;
+        for (decls) |decl| {
+            if (std.mem.eql(u8, decl.name, "Keys")) continue;
+            if (decl.name[0] == '_') continue;
+            fields[i].name = decl.name;
+            fields[i].value = i;
+            i += 1;
+        }
+        break :keys @Type(
+            .{
+                .@"enum" = .{
+                    .fields = &fields,
+                    .decls = &.{},
+                    .is_exhaustive = true,
+                    .tag_type = std.math.IntFittingRange(0, count),
+                },
+            },
+        );
+    };
 };
 
 pub fn generator(self: *Bytes) Generator {
     return .init(self, next);
 }
 
-pub fn next(self: *Bytes, writer: *std.Io.Writer, max_len: usize) Generator.Error!void {
-    std.debug.assert(max_len >= 1);
-    const len = @min(
-        self.rand.intRangeAtMostBiased(usize, self.min_len, self.max_len),
-        max_len,
-    );
+pub fn init(rand: std.Random) Bytes {
+    return .{
+        .rand = rand,
+    };
+}
+
+pub fn next(self: *const Bytes, writer: *std.Io.Writer, max_len: usize) std.Io.Writer.Error!void {
+    try self.atMost(max_len).format(writer);
+}
+
+/// Return a copy of the Bytes, but with a new alphabet.
+pub fn alphabet(self: *const Bytes, comptime key: Alphabet.Keys) Bytes {
+    return .{
+        ._alphabet = @field(Alphabet, @tagName(key)),
+        ._min_len = self._min_len,
+        ._max_len = self._max_len,
+        .rand = self.rand,
+    };
+}
+
+/// Return a copy of the Bytes, but with a new min_len. The new min
+/// len cannot be more than the previous max_len.
+pub fn atLeast(self: *const Bytes, min_len: usize) Bytes {
+    return .{
+        .rand = self.rand,
+        ._alphabet = self._alphabet,
+        ._min_len = @min(self._max_len, min_len),
+        ._max_len = self._max_len,
+    };
+}
+
+/// Return a copy of the Bytes, but with a new max_len. The new max_len cannot
+/// be more the previous max_len.
+pub fn atMost(self: *const Bytes, max_len: usize) Bytes {
+    return .{
+        .rand = self.rand,
+        ._alphabet = self._alphabet,
+        ._min_len = @min(self._min_len, @min(self._max_len, max_len)),
+        ._max_len = @min(self._max_len, max_len),
+    };
+}
+
+pub fn format(self: *const Bytes, writer: *std.Io.Writer) std.Io.Writer.Error!void {
+    _ = try self.write(writer);
+}
+
+/// Write some random data and return the number of bytes written.
+pub fn write(self: *const Bytes, writer: *std.Io.Writer) std.Io.Writer.Error!usize {
+    std.debug.assert(self._min_len >= 1);
+    std.debug.assert(self._max_len >= self._min_len);
+
+    const len = self.rand.intRangeAtMostBiased(usize, self._min_len, self._max_len);
 
     var buf: [8]u8 = undefined;
+
     var remaining = len;
     while (remaining > 0) {
         const data = buf[0..@min(remaining, buf.len)];
         self.rand.bytes(data);
-        if (self.alphabet) |alphabet| {
-            for (data) |*byte| byte.* = alphabet[byte.* % alphabet.len];
+        if (self._alphabet) |_alphabet| {
+            for (data) |*byte| byte.* = _alphabet[byte.* % _alphabet.len];
         }
         try writer.writeAll(data);
         remaining -= data.len;
     }
+
+    return len;
 }
 
 test "bytes" {
     const testing = std.testing;
     var prng = std.Random.DefaultPrng.init(0);
     var buf: [256]u8 = undefined;
     var writer: std.Io.Writer = .fixed(&buf);
-    var v: Bytes = .{ .rand = prng.random() };
-    v.min_len = buf.len;
-    v.max_len = buf.len;
+    var v: Bytes = .{
+        .rand = prng.random(),
+        ._min_len = buf.len,
+        ._max_len = buf.len,
+    };
     const gen = v.generator();
     try gen.next(&writer, buf.len);
     try testing.expectEqual(buf.len, writer.buffered().len);
diff --git a/src/synthetic/Osc.zig b/src/synthetic/Osc.zig
@@ -35,20 +35,6 @@ p_valid: f64 = 1.0,
 p_valid_kind: std.enums.EnumArray(ValidKind, f64) = .initFill(1.0),
 p_invalid_kind: std.enums.EnumArray(InvalidKind, f64) = .initFill(1.0),
 
-/// The alphabet for random bytes (omitting 0x1B and 0x07).
-const bytes_alphabet: []const u8 = alphabet: {
-    var alphabet: [256]u8 = undefined;
-    for (0..alphabet.len) |i| {
-        if (i == 0x1B or i == 0x07) {
-            alphabet[i] = @intCast(i + 1);
-        } else {
-            alphabet[i] = @intCast(i);
-        }
-    }
-    const result = alphabet;
-    break :alphabet &result;
-};
-
 pub fn generator(self: *Osc) Generator {
     return .init(self, next);
 }
@@ -100,30 +86,32 @@ fn nextUnwrapped(self: *Osc, writer: *std.Io.Writer, max_len: usize) Generator.E
 fn nextUnwrappedValidExact(self: *const Osc, writer: *std.Io.Writer, k: ValidKind, max_len: usize) Generator.Error!void {
     switch (k) {
         .change_window_title => {
-            try writer.writeAll("0;"); // Set window title
-            var bytes_gen = self.bytes();
-            try bytes_gen.next(writer, max_len - 2);
+            try writer.print("0;{f}", .{self.bytes().atMost(max_len - 3)}); // Set window title
         },
 
         .prompt_start => {
+            var remaining = max_len;
+
             try writer.writeAll("133;A"); // Start prompt
+            remaining -= 4;
 
             // aid
             if (self.rand.boolean()) {
-                var bytes_gen = self.bytes();
-                bytes_gen.max_len = 16;
                 try writer.writeAll(";aid=");
-                try bytes_gen.next(writer, max_len);
+                remaining -= 5;
+                remaining -= try self.bytes().alphabet(.kv).atMost(@min(16, remaining)).write(writer);
             }
 
             // redraw
-            if (self.rand.boolean()) {
+            if (self.rand.boolean()) redraw: {
+                if (remaining < 9) break :redraw;
                 try writer.writeAll(";redraw=");
                 if (self.rand.boolean()) {
                     try writer.writeAll("1");
                 } else {
                     try writer.writeAll("0");
                 }
+                remaining -= 9;
             }
         },
 
@@ -139,23 +127,17 @@ fn nextUnwrappedInvalidExact(
 ) Generator.Error!void {
     switch (k) {
         .random => {
-            var bytes_gen = self.bytes();
-            try bytes_gen.next(writer, max_len);
+            try self.bytes().atMost(max_len).format(writer);
         },
 
         .good_prefix => {
-            try writer.writeAll("133;");
-            var bytes_gen = self.bytes();
-            try bytes_gen.next(writer, max_len - 4);
+            try writer.print("133;{f}", .{self.bytes().atMost(max_len - 4)});
         },
     }
 }
 
 fn bytes(self: *const Osc) Bytes {
-    return .{
-        .rand = self.rand,
-        .alphabet = bytes_alphabet,
-    };
+    return Bytes.init(self.rand).alphabet(.osc);
 }
 
 /// Choose whether to generate a valid or invalid OSC request based
diff --git a/src/synthetic/cli/Ascii.zig b/src/synthetic/cli/Ascii.zig
@@ -26,10 +26,7 @@ pub fn destroy(self: *Ascii, alloc: Allocator) void {
 pub fn run(self: *Ascii, writer: *std.Io.Writer, rand: std.Random) !void {
     _ = self;
 
-    var gen: synthetic.Bytes = .{
-        .rand = rand,
-        .alphabet = synthetic.Bytes.Alphabet.ascii,
-    };
+    var gen = synthetic.Bytes.init(rand).alphabet(.ascii);
 
     while (true) {
         gen.next(writer, 1024) catch |err| {
