lib/hypothesis/src/swarm.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const conjecture = @import("conjecture.zig");
  3 
  4 const ConjectureData = conjecture.ConjectureData;
  5 const DrawError = conjecture.DrawError;
  6 
  7 pub const Config = struct {
  8     min_active: u16 = 1,
  9     max_active: ?u16 = null,
 10     max_weight: u16 = 100,
 11 };
 12 
 13 pub fn Weights(comptime Action: type) type {
 14     const actions = actionTags(Action);
 15     return struct {
 16         values: [actions.len]u16,
 17         total: u64,
 18 
 19         const Self = @This();
 20 
 21         pub fn get(self: Self, action: Action) u16 {
 22             inline for (actions, 0..) |tag, index| {
 23                 if (action == tag) return self.values[index];
 24             }
 25             unreachable;
 26         }
 27 
 28         pub fn activeCount(self: Self) usize {
 29             var count: usize = 0;
 30             for (self.values) |weight| {
 31                 if (weight > 0) count += 1;
 32             }
 33             return count;
 34         }
 35 
 36         pub fn choose(self: Self, data: *ConjectureData) DrawError!Action {
 37             std.debug.assert(self.total > 0);
 38             var selection = try data.drawInteger(0, self.total - 1, 0);
 39             inline for (actions, 0..) |tag, index| {
 40                 const weight = self.values[index];
 41                 if (selection < weight) return tag;
 42                 selection -= weight;
 43             }
 44             unreachable;
 45         }
 46     };
 47 }
 48 
 49 pub fn draw(
 50     comptime Action: type,
 51     data: *ConjectureData,
 52     config: Config,
 53 ) DrawError!Weights(Action) {
 54     const actions = actionTags(Action);
 55     const action_count: u16 = @intCast(actions.len);
 56     const max_active = @min(config.max_active orelse action_count, action_count);
 57     std.debug.assert(config.min_active > 0);
 58     std.debug.assert(config.min_active <= max_active);
 59     std.debug.assert(config.max_weight > 0);
 60 
 61     const active_target: u16 = @intCast(try data.drawInteger(
 62         config.min_active,
 63         max_active,
 64         config.min_active,
 65     ));
 66     var result = Weights(Action){
 67         .values = @splat(0),
 68         .total = 0,
 69     };
 70     var selected: u16 = 0;
 71     for (&result.values, 0..) |*value, index| {
 72         const remaining: u16 = action_count - @as(u16, @intCast(index));
 73         const needed = active_target - selected;
 74         const include = if (needed == 0)
 75             false
 76         else if (needed == remaining)
 77             true
 78         else
 79             try data.drawInteger(0, remaining - 1, 0) < needed;
 80         if (include) {
 81             const weight: u16 = @intCast(try data.drawInteger(1, config.max_weight, 1));
 82             value.* = weight;
 83             result.total += weight;
 84             selected += 1;
 85         }
 86     }
 87     std.debug.assert(selected == active_target);
 88     std.debug.assert(result.total > 0);
 89     return result;
 90 }
 91 
 92 fn actionTags(comptime Action: type) []const Action {
 93     const info = @typeInfo(Action);
 94     if (info != .@"enum") @compileError("swarm actions must be an enum");
 95     if (info.@"enum".mode != .exhaustive) {
 96         @compileError("swarm actions must be an exhaustive enum");
 97     }
 98     const tags = std.meta.tags(Action);
 99     if (tags.len == 0) @compileError("swarm actions cannot be empty");
100     if (tags.len > std.math.maxInt(u16)) {
101         @compileError("swarm actions exceed the supported count");
102     }
103     return tags;
104 }
105 
106 const TestAction = enum {
107     put,
108     commit,
109     checkpoint,
110 };
111 
112 test "swarm: one run reuses one deterministic action distribution" {
113     var left = ConjectureData.init(std.testing.allocator, 0x51a2_77);
114     defer left.deinit();
115     var right = ConjectureData.init(std.testing.allocator, 0x51a2_77);
116     defer right.deinit();
117 
118     const left_weights = try draw(TestAction, &left, .{});
119     const right_weights = try draw(TestAction, &right, .{});
120     try std.testing.expectEqualDeep(left_weights, right_weights);
121     for (0..32) |_| {
122         try std.testing.expectEqual(
123             try left_weights.choose(&left),
124             try right_weights.choose(&right),
125         );
126     }
127 }
128 
129 test "swarm: an active subset focuses every action in a run" {
130     var nodes = [_]conjecture.ChoiceNode{
131         .{ .kind = .integer, .value = 1 },
132         .{ .kind = .integer, .value = 2 },
133         .{ .kind = .integer, .value = 1 },
134         .{ .kind = .integer, .value = 7 },
135         .{ .kind = .integer, .value = 0 },
136     };
137     var data = ConjectureData.initReplay(std.testing.allocator, &nodes, null);
138     defer data.deinit();
139 
140     const weights = try draw(TestAction, &data, .{ .min_active = 1, .max_active = 1 });
141     try std.testing.expectEqual(@as(usize, 1), weights.activeCount());
142     try std.testing.expectEqual(@as(u16, 0), weights.get(.put));
143     try std.testing.expectEqual(@as(u16, 0), weights.get(.commit));
144     try std.testing.expectEqual(@as(u16, 7), weights.get(.checkpoint));
145     try std.testing.expectEqual(TestAction.checkpoint, try weights.choose(&data));
146 }