finished wrapper to use (comptime) allocators from capstone

2 files changed, 149 insertions, 13 deletions

diff --git a/src/elfo.zig b/src/elfo.zig
index 20f798a..5e0d780 100644
--- a/src/elfo.zig
+++ b/src/elfo.zig
@@ -1,29 +1,50 @@
 const std = @import("std");
 const cs = @import("capstone");
+const hooks = @import("hooks.zig");
 const meta_opts = @import("meta");
 
-pub fn main(init: std.process.Init) !void {
-    var args = try init.minimal.args.iterateAllocator(init.gpa);
+pub fn main(init: std.process.Init.Minimal) !void {
+    var debug_alloc: std.heap.DebugAllocator(.{}) = .init;
+    const alloc = debug_alloc.allocator();
+
+    var threaded: std.Io.Threaded = .init(alloc, .{
+        .argv0 = .init(init.args),
+        .environ = init.environ,
+    });
+    defer threaded.deinit();
+    const io = threaded.io();
+
+    var args = try init.args.iterateAllocator(alloc);
     defer args.deinit();
     _ = args.next(); // skip argv[0]
 
     var buffer: [64]u8 = undefined;
-    const stderr = try init.io.lockStderr(&buffer, .escape_codes);
+    const stderr = try io.lockStderr(&buffer, .escape_codes);
 
-    // TODO: finish passing custom alloc operations here
-    // const mem_config: cs.cs_opt_mem = undefined;
-    // std.debug.assert(cs.cs_option(0, cs.CS_OPT_MEM, @intFromPtr(&mem_config)) == cs.CS_ERR_OK);
+    const hook = hooks.init(
+        if (meta_opts.gdb) .breakpoint else .none,
+        // TODO: this needs to be comptime, it'd be cool to don't need that
+        std.heap.page_allocator,
+        false,
+    );
 
-    if (meta_opts.gdb) @breakpoint();
+    const mem_config: cs.cs_opt_mem = .{
+        .malloc = @ptrCast(&hook.malloc),
+        .free = @ptrCast(&hook.free),
+        .calloc = @ptrCast(&hook.calloc),
+        .realloc = @ptrCast(&hook.realloc),
+        .vsnprintf = @ptrCast(&hook.vsnprintf),
+    };
+    std.debug.assert(cs.cs_option(0, cs.CS_OPT_MEM, @intFromPtr(&mem_config)) == cs.CS_ERR_OK);
 
     try printElf(
-        init.gpa,
-        init.io,
+        alloc,
+        io,
         args.next() orelse "./study-samples/split",
         stderr.terminal(),
         .{
             .show_unaddressable_sections = true,
-            // .skip_sections_content = true,
+            .skip_sections_content = true,
         },
     );
 }
@@ -156,7 +177,7 @@ pub fn printElf(
 
                 // x64: symtab: .{ .sh_name = 1, .sh_type = 2, .sh_flags = 0, .sh_addr = 0, .sh_offset = 4256, .sh_size = 1680, .sh_link = 27, .sh_info = 45, .sh_addralign = 8, .sh_entsize = 24 }
                 // arm32: .{ .sh_name = 1, .sh_type = 2, .sh_flags = 0, .sh_addr = 0, .sh_offset = 4264, .sh_size = 1856, .sh_link = 27, .sh_info = 87, .sh_addralign = 4, .sh_entsize = 16 }
-                try bw.print("symtab: {any}\n", .{s});
+                // try bw.print("symtab: {any}\n", .{s});
                 break :blk s;
             }
         }
@@ -225,7 +246,7 @@ pub fn printElf(
 
     for (symbols_index.items) |sym| {
         if (sym.kind == std.elf.STT_FUNC and sym.name.len > 0)
-            try bw.print("{s} {x}-{x}\n", .{ sym.name, sym.start, sym.end });
+            try bw.print("{x}-{x} {s}\n", .{ sym.start, sym.end, sym.name });
     }
 
     for (sections.items) |section| {
@@ -240,7 +261,7 @@ pub fn printElf(
             try term.setColor(.bright_green);
             if (elf_shstrtab_slice != null)
                 try bw.print("{s}", .{std.mem.sliceTo(elf_shstrtab_slice.?[section.sh_name..], 0)});
-            try bw.print("\n", .{});
+            try bw.print("", .{});
             try term.setColor(.reset);
 
             // --
diff --git a/src/hooks.zig b/src/hooks.zig
new file mode 100644
index 0000000..549d871
--- /dev/null
+++ b/src/hooks.zig
@@ -0,0 +1,115 @@
+const std = @import("std");
+
+pub const Kind = enum {
+    none,
+    breakpoint,
+};
+const Self = @This();
+pub const Hook = @This();
+
+// pub fn init(comptime kind: Kind, comptime alloc: std.mem.Allocator) type {
+//     return struct {
+//         alloc: std.mem.Allocator = alloc,
+
+//         pub fn malloc(size: usize) callconv(.c) ?*anyopaque {
+//             if (kind == .breakpoint) @breakpoint();
+//             return std.c.malloc(size);
+//         }
+
+//         pub fn calloc(n: usize, size: usize) callconv(.c) ?*anyopaque {
+//             if (kind == .breakpoint) @breakpoint();
+//             return std.c.calloc(n, size);
+//         }
+//         pub fn realloc(ptr: ?*anyopaque, size: usize) callconv(.c) ?*anyopaque {
+//             if (kind == .breakpoint) @breakpoint();
+//             return std.c.realloc(ptr, size);
+//         }
+//         pub fn free(ptr: ?*anyopaque) callconv(.c) void {
+//             if (kind == .breakpoint) @breakpoint();
+//             return std.c.free(ptr);
+//         }
+//         pub fn vsnprintf(_: [*c]u8, _: usize, _: [*c]const u8, _: [*c]u8) callconv(.c) ?*anyopaque {
+//             if (kind == .breakpoint) @breakpoint();
+//             return null;
+//         }
+//     };
+// }
+//
+
+pub fn init(comptime kind: Kind, comptime alloc: std.mem.Allocator, comptime uselibc: bool) type {
+    return struct {
+        alloc: std.mem.Allocator = alloc,
+
+        const Header = struct {
+            len: usize,
+        };
+        const header_size = std.mem.alignForward(usize, @sizeOf(Header), @alignOf(std.c.max_align_t));
+        const alignment: std.mem.Alignment = .fromByteUnits(@alignOf(std.c.max_align_t));
+
+        inline fn ptrFromRaw(raw: []u8) ?*anyopaque {
+            const h: *Header = @ptrCast(@alignCast(raw.ptr));
+            h.* = .{ .len = raw.len };
+            return @ptrFromInt(@intFromPtr(raw.ptr) + header_size);
+        }
+
+        inline fn rawFromPtr(ptr: *anyopaque) []u8 {
+            const base = @intFromPtr(ptr) - header_size;
+            const h: *Header = @ptrFromInt(base);
+            return @as([*]u8, @ptrFromInt(base))[0..h.len];
+        }
+
+        pub fn malloc(size: usize) callconv(.c) ?*anyopaque {
+            if (kind == .breakpoint) @breakpoint();
+            if (uselibc) return std.c.malloc(size);
+            const total = header_size + size;
+            const raw = alloc.alignedAlloc(u8, alignment, total) catch return null;
+            return ptrFromRaw(raw);
+        }
+
+        pub fn calloc(n: usize, size: usize) callconv(.c) ?*anyopaque {
+            if (kind == .breakpoint) @breakpoint();
+            if (uselibc) return std.c.calloc(n, size);
+            const total_data = std.math.mul(usize, n, size) catch return null;
+            const total = header_size + total_data;
+            const raw = alloc.alignedAlloc(u8, alignment, total) catch return null;
+            @memset(raw[header_size..], 0);
+            return ptrFromRaw(raw);
+        }
+
+        pub fn realloc(ptr: ?*anyopaque, size: usize) callconv(.c) ?*anyopaque {
+            if (kind == .breakpoint) @breakpoint();
+            if (uselibc) return std.c.realloc(ptr, size);
+            if (ptr) |p| {
+                const old_raw = rawFromPtr(p);
+                const total = header_size + size;
+                if (alloc.resize(old_raw, total)) {
+                    const h: *Header = @ptrCast(@alignCast(old_raw.ptr));
+                    h.* = .{ .len = total };
+                    return p;
+                }
+                const new_raw = alloc.alignedAlloc(u8, alignment, total) catch return null;
+                const old_data = old_raw[header_size..];
+                const copy_len = @min(old_data.len, size);
+                @memcpy(new_raw[header_size..][0..copy_len], old_data[0..copy_len]);
+                alloc.free(old_raw);
+                return ptrFromRaw(new_raw);
+            } else {
+                return malloc(size);
+            }
+        }
+
+        pub fn free(ptr: ?*anyopaque) callconv(.c) void {
+            if (kind == .breakpoint) @breakpoint();
+            if (uselibc) return std.c.free(ptr);
+            if (ptr) |p| {
+                const raw = rawFromPtr(p);
+                alloc.free(raw);
+            }
+        }
+
+        pub fn vsnprintf(_: [*c]u8, _: usize, _: [*c]const u8, _: [*c]u8) callconv(.c) ?*anyopaque {
+            if (kind == .breakpoint) @breakpoint();
+            return null;
+        }
+    };
+}