Hi,
I’m looking for some feedback on my example code and hope for some suggestions/tipps what can be done better ![]()
The goal is to search all textfiles of a nested directory tree and return those which match a string pattern. For now, no complex match syntax like Regex or fuzzy matching is wanted, just plain exact matching of the given pattern. Additionally, the following aspects are important:
- return file path on first match. It doesn’t make any difference if the file content matches the pattern once or multiple times.
- All non plain-text files should be skipped.
- Keep memory usage low if possible, especially heap-allocated memory.
Its the first time for me trying to write some SIMD-boosted code. My main sources for getting into it and understanding SIMD stuff were the following:
- SIMD with Zig
- SIMD-friendly algorithms for substring searching
- SIMD Vector to Mask for Substring Search
Here is my example Code:
Longer example code block incoming
// Compiled with Zig 0.16.0
// Run with `zig run simd_search.zig -- <pattern> <root-search-dir>`
const std = @import("std");
const vec_len = std.simd.suggestVectorLength(u8) orelse 8;
const Mask = @Int(.unsigned, vec_len);
pub fn main(init: std.process.Init) !void {
const arena = init.arena.allocator();
const io = init.io;
var args = init.minimal.args.iterate();
_ = args.next();
const pattern = if (args.next()) |pat| pat else return error.NoPattern;
const path = if (args.next()) |p| p else ".";
const root_dir = try std.Io.Dir.cwd().openDir(io, path, .{ .iterate = true });
defer root_dir.close(io);
std.log.debug("Recursivley search for pattern '{s}' in root dir '{s}'", .{ pattern, path });
var matches: std.ArrayList([]const u8) = .empty;
try recursivleySearchFiles(arena, root_dir, io, path, pattern, &matches);
std.debug.print("Matched files {d}:\n", .{matches.items.len});
for (matches.items) |match| {
std.debug.print("- {s}\n", .{match});
}
}
fn recursivleySearchFiles(
arena: std.mem.Allocator,
dir: std.Io.Dir,
io: std.Io,
path: []const u8,
pattern: []const u8,
matches: *std.ArrayList([]const u8),
) !void {
var dir_it = dir.iterate();
const first_char = pattern[0];
const last_char = pattern[pattern.len - 1];
const first_c: @Vector(vec_len, u8) = @splat(first_char);
const last_c: @Vector(vec_len, u8) = @splat(last_char);
walk_dir: while (try dir_it.next(io)) |entry| {
switch (entry.kind) {
.file => {
try matchFileContent(arena, dir, io, path, entry.name, pattern, matches, first_c, last_c);
continue :walk_dir;
},
.directory => {
const subpath = std.fs.path.join(arena, &.{ path, entry.name }) catch entry.name;
const subdir = try dir.openDir(io, entry.name, .{ .iterate = true });
defer subdir.close(io);
try recursivleySearchFiles(arena, subdir, io, subpath, pattern, matches);
continue :walk_dir;
},
else => continue :walk_dir,
}
}
}
const FILE_SIZE_STACK = 1024;
fn matchFileContent(
arena: std.mem.Allocator,
dir: std.Io.Dir,
io: std.Io,
path: []const u8,
file_name: []const u8,
pattern: []const u8,
matches: *std.ArrayList([]const u8),
first_c: @Vector(vec_len, u8),
last_c: @Vector(vec_len, u8),
) !void {
var file = try dir.openFile(io, file_name, .{});
defer file.close(io);
var check_txt_buf: [1024]u8 = undefined;
const i = try file.readPositionalAll(io, &check_txt_buf, 0);
if (!isText(check_txt_buf[0..i])) return;
const file_path = std.fs.path.join(arena, &.{ path, file_name }) catch file_name;
const file_len = try file.length(io);
if (file_len <= FILE_SIZE_STACK) {
var file_buf: [FILE_SIZE_STACK]u8 = undefined;
_ = try file.readPositionalAll(io, &file_buf, 0);
if (findMatch(pattern, &file_buf, first_c, last_c)) |idx| {
_ = idx;
try matches.append(arena, file_path);
return;
}
} else {
const file_buf_alloc = try arena.alloc(u8, file_len);
_ = try file.readPositionalAll(io, file_buf_alloc, 0);
if (findMatch(pattern, file_buf_alloc, first_c, last_c)) |idx| {
_ = idx;
try matches.append(arena, file_path);
return;
}
}
}
fn findMatch(
pattern: []const u8,
data: []const u8,
first_c: @Vector(vec_len, u8),
last_c: @Vector(vec_len, u8),
) ?usize {
var pos: usize = 0;
var rest = data.len;
while (rest > 0) {
if (rest < vec_len + pattern.len) {
return std.mem.findPos(u8, data, pos, pattern);
}
const block_first: @Vector(vec_len, u8) = data[pos..][0..vec_len].*;
const block_last: @Vector(vec_len, u8) = data[pos + pattern.len - 1 ..][0..vec_len].*;
const eq_first = first_c == block_first;
const eq_last = last_c == block_last;
const matched = eq_first & eq_last;
var mask: Mask = @bitCast(matched);
if (mask != 0) {
const bit = @ctz(mask);
const hit = pos + bit;
const match = data[hit .. hit + pattern.len];
if (std.mem.eql(u8, pattern, match)) return hit;
mask = 0;
}
pos += vec_len;
rest -= vec_len;
}
return null;
}
fn isText(buf: []const u8) bool {
if (buf.len == 0) return false;
for (buf) |b| {
if (b == 0) return false;
}
return true;
}
So far, the code works just fine. On a small dir with subdirs containing about 130 Markdownfiles with a total size of round about 500Kb searching for a pattern which is present in 14 files it takes about 7ms. ripgrep takes about 3.5ms.
Of course, the latter is heavily optimized and uses parallel execution etc[1] while my simple code is singe threaded and has no further optimization beside SIMD instructions.[2] Maybe its not that bad already.
However, I’d like to understand all of that stuff better and might be overlooking something really obvious since I’ve no experience with SIMD and low level programming is just a hobby for me.
Thanks already to everyone who read this and might have some feedback ![]()