1 files changed, 81 insertions, 4 deletions

diff --git a/src/lua_tool.zig b/src/lua_tool.zig
index 3064d2b..21b3e51 100644
--- a/src/lua_tool.zig
+++ b/src/lua_tool.zig
@@ -25,6 +25,10 @@ pub const LuaTool = struct {
     allocator: Allocator,
     // Owned, NUL-terminated for `luaL_loadfilex`.
     script_path_z: [:0]u8,
+    /// Optional directory to prepend to `package.path` so the script can
+    /// `require` sibling files. Owned, NUL-terminated. Null for single-file
+    /// extensions.
+    package_root_z: ?[:0]u8,
     name_owned: []u8,
     description_owned: []u8,
     schema_owned: []u8,
@@ -39,16 +43,31 @@ pub const LuaTool = struct {
         name: []const u8,
         description: []const u8,
         schema_json: []const u8,
+        package_root: ?[]const u8,
     ) !panto.Tool {
         const self = try allocator.create(LuaTool);
         errdefer allocator.destroy(self);
 
+        const script_path_z = try allocator.dupeZ(u8, script_path);
+        errdefer allocator.free(script_path_z);
+        const name_owned = try allocator.dupe(u8, name);
+        errdefer allocator.free(name_owned);
+        const description_owned = try allocator.dupe(u8, description);
+        errdefer allocator.free(description_owned);
+        const schema_owned = try allocator.dupe(u8, schema_json);
+        errdefer allocator.free(schema_owned);
+        const package_root_z: ?[:0]u8 = if (package_root) |p|
+            try allocator.dupeZ(u8, p)
+        else
+            null;
+
         self.* = .{
             .allocator = allocator,
-            .script_path_z = try allocator.dupeZ(u8, script_path),
-            .name_owned = try allocator.dupe(u8, name),
-            .description_owned = try allocator.dupe(u8, description),
-            .schema_owned = try allocator.dupe(u8, schema_json),
+            .script_path_z = script_path_z,
+            .package_root_z = package_root_z,
+            .name_owned = name_owned,
+            .description_owned = description_owned,
+            .schema_owned = schema_owned,
         };
 
         return panto.Tool{
@@ -63,6 +82,7 @@ pub const LuaTool = struct {
     fn freeAll(self: *LuaTool) void {
         const a = self.allocator;
         a.free(self.script_path_z);
+        if (self.package_root_z) |p| a.free(p);
         a.free(self.name_owned);
         a.free(self.description_owned);
         a.free(self.schema_owned);
@@ -102,6 +122,10 @@ fn runLuaHandler(
     c.luaL_openlibs(L);
     lua_bridge.install(L);
 
+    if (self.package_root_z) |root| {
+        try prependPackagePath(L, root);
+    }
+
     // Run the script so register_tool fires.
     lua_bridge.loadFile(L, self.script_path_z) catch |err| {
         logTopAsError(L, "lua: failed to load extension");
@@ -176,6 +200,27 @@ pub fn loadExtension(
     registry: *panto.ToolRegistry,
     script_path: []const u8,
 ) !usize {
+    return loadExtensionImpl(allocator, registry, script_path, null);
+}
+
+/// Like `loadExtension`, but extends `package.path` with `<package_root>/?.lua`
+/// and `<package_root>/?/init.lua` so the script can `require` sibling files.
+/// `package_root` should be the directory containing the script.
+pub fn loadExtensionWithPackageRoot(
+    allocator: Allocator,
+    registry: *panto.ToolRegistry,
+    script_path: []const u8,
+    package_root: []const u8,
+) !usize {
+    return loadExtensionImpl(allocator, registry, script_path, package_root);
+}
+
+fn loadExtensionImpl(
+    allocator: Allocator,
+    registry: *panto.ToolRegistry,
+    script_path: []const u8,
+    package_root: ?[]const u8,
+) !usize {
     const path_z = try allocator.dupeZ(u8, script_path);
     defer allocator.free(path_z);
 
@@ -184,6 +229,13 @@ pub fn loadExtension(
     c.luaL_openlibs(L);
     lua_bridge.install(L);
 
+    // Configure require() for directory-style extensions before running.
+    if (package_root) |root| {
+        const root_z = try allocator.dupeZ(u8, root);
+        defer allocator.free(root_z);
+        try prependPackagePath(L, root_z);
+    }
+
     lua_bridge.loadFile(L, path_z) catch |err| {
         logTopAsError(L, "lua: failed to load extension");
         return err;
@@ -200,6 +252,7 @@ pub fn loadExtension(
             r.name,
             r.description,
             r.schema_json,
+            package_root,
         );
         // If registration fails (e.g. duplicate name), free the tool we just
         // built and propagate the error.
@@ -211,6 +264,30 @@ pub fn loadExtension(
     return regs.len;
 }
 
+/// Prepend `<root>/?.lua` and `<root>/?/init.lua` to `package.path` so that
+/// `require("foo")` resolves against files next to the extension's
+/// `init.lua`. We prepend (not replace) so the standard library remains
+/// available.
+///
+/// Stack effect: net zero.
+fn prependPackagePath(L: *c.lua_State, root: [:0]const u8) !void {
+    // Run a tiny Lua snippet rather than fiddle with the stack manually:
+    // string concatenation is so much nicer in Lua.
+    const snippet =
+        \\local root = ...
+        \\package.path = root .. "/?.lua;" .. root .. "/?/init.lua;" .. package.path
+    ;
+    if (c.luaL_loadstring(L, snippet) != 0) {
+        logTopAsError(L, "lua: package.path loader failed to compile");
+        return error.LuaPackagePathLoadFailed;
+    }
+    _ = c.lua_pushlstring(L, root.ptr, root.len);
+    if (c.lua_pcallk(L, 1, 0, 0, 0, null) != 0) {
+        logTopAsError(L, "lua: package.path setup failed");
+        return error.LuaPackagePathSetupFailed;
+    }
+}
+
 // ---------------------------------------------------------------------------
 // Tests
 // ---------------------------------------------------------------------------