118 lines
3.5 KiB
Zig
118 lines
3.5 KiB
Zig
const std = @import("std");
|
|
const rand = std.crypto.random;
|
|
|
|
const markov_data = @embedFile("markov.bin");
|
|
|
|
pub const DataPoint = struct {
|
|
char: u8,
|
|
prob: f32,
|
|
|
|
pub fn desc(context: void, a: DataPoint, b: DataPoint) bool {
|
|
_ = context;
|
|
return a.prob > b.prob;
|
|
}
|
|
};
|
|
|
|
pub const MarkovChain = struct {
|
|
allocator: std.mem.Allocator,
|
|
map: std.AutoHashMap(u8, []DataPoint),
|
|
|
|
pub fn fromFile(path: []const u8, allocator: std.mem.Allocator) !MarkovChain {
|
|
var markovBinFile = try std.fs.cwd().openFile(path, .{ .mode = .read_only });
|
|
const reader = markovBinFile.reader();
|
|
|
|
return try fromReader(&reader, allocator);
|
|
}
|
|
|
|
pub fn fromMemory(data: []const u8, allocator: std.mem.Allocator) !MarkovChain {
|
|
var stream = std.io.fixedBufferStream(data);
|
|
const reader = stream.reader();
|
|
|
|
return try fromReader(reader, allocator);
|
|
}
|
|
|
|
fn fromReader(reader: anytype, allocator: std.mem.Allocator) !MarkovChain {
|
|
var self = MarkovChain{
|
|
.allocator = allocator,
|
|
.map = std.AutoHashMap(u8, []DataPoint).init(allocator),
|
|
};
|
|
|
|
for (0..256) |prevChar| {
|
|
const cnt = try reader.readInt(u8, .little);
|
|
var nextChars: []DataPoint = try self.allocator.alloc(DataPoint, cnt);
|
|
for (0..cnt) |i| {
|
|
const nextByte = try reader.readByte();
|
|
const prob: f32 = @bitCast(try reader.readInt(u32, .little));
|
|
nextChars[i] = DataPoint{
|
|
.char = nextByte,
|
|
.prob = prob,
|
|
};
|
|
}
|
|
try self.map.put(@as(u8, @intCast(prevChar)), nextChars);
|
|
}
|
|
|
|
return self;
|
|
}
|
|
|
|
pub fn deinit(self: *MarkovChain) void {
|
|
var iter = self.map.iterator();
|
|
while (iter.next()) |entry| {
|
|
self.allocator.free(entry.value_ptr.*);
|
|
}
|
|
self.map.deinit();
|
|
}
|
|
|
|
pub fn generate(self: *MarkovChain, size: u8, allocator: std.mem.Allocator) ![]u8 {
|
|
var result = try allocator.alloc(u8, size);
|
|
var previous: u8 = 0;
|
|
for (0..size) |i| {
|
|
const choices = self.map.get(previous).?;
|
|
const randFloat = rand.float(f32);
|
|
var cumul: f32 = 0;
|
|
for (0..choices.len) |j| {
|
|
cumul += choices[j].prob;
|
|
if (randFloat < cumul) {
|
|
result[i] = choices[j].char;
|
|
previous = choices[j].char;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
};
|
|
|
|
test "basic test from file" {
|
|
var gpa = std.heap.GeneralPurposeAllocator(.{}){};
|
|
const allocator = gpa.allocator();
|
|
defer {
|
|
_ = gpa.deinit();
|
|
}
|
|
|
|
var markov = try MarkovChain.fromFile("markov.bin", allocator);
|
|
defer markov.deinit();
|
|
|
|
for (0..6) |_| {
|
|
const randName = try markov.generate(8, allocator);
|
|
defer allocator.free(randName);
|
|
std.debug.print("generated from file: {s}\n", .{randName});
|
|
}
|
|
}
|
|
|
|
test "basic test from embedded file" {
|
|
var gpa = std.heap.GeneralPurposeAllocator(.{}){};
|
|
const allocator = gpa.allocator();
|
|
defer {
|
|
_ = gpa.deinit();
|
|
}
|
|
|
|
var markov = try MarkovChain.fromMemory(markov_data, allocator);
|
|
defer markov.deinit();
|
|
|
|
for (0..6) |_| {
|
|
const randName = try markov.generate(8, allocator);
|
|
defer allocator.free(randName);
|
|
std.debug.print("generated from memory: {s}\n", .{randName});
|
|
}
|
|
}
|