os: add prependEnv(), like appendEnv() (#2983)

We can use this function in setupXdgDataDirs() to simplify the
XDG_DATA_DIRS environment variable code in a more standardized way.

src/os/env.zig

@@ -34,6 +34,23 @@ pub fn appendEnvAlways(
     });
 }
 
+/// Prepend a value to an environment variable such as PATH.
+/// The returned value is always allocated so it must be freed.
+pub fn prependEnv(
+    alloc: Allocator,
+    current: []const u8,
+    value: []const u8,
+) ![]u8 {
+    // If there is no prior value, we return it as-is
+    if (current.len == 0) return try alloc.dupe(u8, value);
+
+    return try std.fmt.allocPrint(alloc, "{s}{c}{s}", .{
+        value,
+        std.fs.path.delimiter,
+        current,
+    });
+}
+
 /// The result of getenv, with a shared deinit to properly handle allocation
 /// on Windows.
 pub const GetEnvResult = struct {
@@ -110,3 +127,25 @@ test "appendEnv existing" {
         try testing.expectEqualStrings(result, "a:b:foo");
     }
 }
+
+test "prependEnv empty" {
+    const testing = std.testing;
+    const alloc = testing.allocator;
+
+    const result = try prependEnv(alloc, "", "foo");
+    defer alloc.free(result);
+    try testing.expectEqualStrings(result, "foo");
+}
+
+test "prependEnv existing" {
+    const testing = std.testing;
+    const alloc = testing.allocator;
+
+    const result = try prependEnv(alloc, "a:b", "foo");
+    defer alloc.free(result);
+    if (builtin.os.tag == .windows) {
+        try testing.expectEqualStrings(result, "foo;a:b");
+    } else {
+        try testing.expectEqualStrings(result, "foo:a:b");
+    }
+}

src/os/main.zig

@@ -27,6 +27,7 @@ pub const CFReleaseThread = @import("cf_release_thread.zig");
 pub const TempDir = @import("TempDir.zig");
 pub const appendEnv = env.appendEnv;
 pub const appendEnvAlways = env.appendEnvAlways;
+pub const prependEnv = env.prependEnv;
 pub const getenv = env.getenv;
 pub const setenv = env.setenv;
 pub const unsetenv = env.unsetenv;

src/termio/shell_integration.zig

@@ -5,6 +5,7 @@ const ArenaAllocator = std.heap.ArenaAllocator;
 const EnvMap = std.process.EnvMap;
 const config = @import("../config.zig");
 const homedir = @import("../os/homedir.zig");
+const internal_os = @import("../os/main.zig");
 
 const log = std.log.scoped(.shell_integration);
 
@@ -435,8 +436,8 @@ test "bash: preserve ENV" {
 /// Setup automatic shell integration for shells that include
 /// their modules from paths in `XDG_DATA_DIRS` env variable.
 ///
-/// Path of shell-integration dir is prepended to `XDG_DATA_DIRS`.
-/// It is also saved in `GHOSTTY_SHELL_INTEGRATION_XDG_DIR` variable
+/// The shell-integration path is prepended to `XDG_DATA_DIRS`.
+/// It is also saved in the `GHOSTTY_SHELL_INTEGRATION_XDG_DIR` variable
 /// so that the shell can refer to it and safely remove this directory
 /// from `XDG_DATA_DIRS` when integration is complete.
 fn setupXdgDataDirs(
@@ -458,11 +459,8 @@ fn setupXdgDataDirs(
     // so that our modifications don't interfere with other commands.
     try env.put("GHOSTTY_SHELL_INTEGRATION_XDG_DIR", integ_dir);
 
-    {
-        const xdg_data_dir_key = "XDG_DATA_DIRS";
-
     // We attempt to avoid allocating by using the stack up to 4K.
-        // Max stack size is considerably larger on macOS and Linux but
+    // Max stack size is considerably larger on mac
     // 4K is a reasonable size for this for most cases. However, env
     // vars can be significantly larger so if we have to we fall
     // back to a heap allocated value.
@@ -473,17 +471,48 @@ fn setupXdgDataDirs(
     // This ensures that the default directories aren't lost by setting
     // our desired integration dir directly. See #2711.
     // <https://specifications.freedesktop.org/basedir-spec/0.6/#variables>
-        const old = env.get(xdg_data_dir_key) orelse "/usr/local/share:/usr/share";
-
-        const prepended = try std.fmt.allocPrint(stack_alloc, "{s}{c}{s}", .{
+    const xdg_data_dirs_key = "XDG_DATA_DIRS";
+    try env.put(
+        xdg_data_dirs_key,
+        try internal_os.prependEnv(
+            stack_alloc,
+            env.get(xdg_data_dirs_key) orelse "/usr/local/share:/usr/share",
             integ_dir,
-            std.fs.path.delimiter,
-            old,
-        });
-        defer stack_alloc.free(prepended);
-
-        try env.put(xdg_data_dir_key, prepended);
+        ),
+    );
 }
+
+test "xdg: empty XDG_DATA_DIRS" {
+    const testing = std.testing;
+
+    var arena = ArenaAllocator.init(testing.allocator);
+    defer arena.deinit();
+    const alloc = arena.allocator();
+
+    var env = EnvMap.init(alloc);
+    defer env.deinit();
+
+    try setupXdgDataDirs(alloc, ".", &env);
+
+    try testing.expectEqualStrings("./shell-integration", env.get("GHOSTTY_SHELL_INTEGRATION_XDG_DIR").?);
+    try testing.expectEqualStrings("./shell-integration:/usr/local/share:/usr/share", env.get("XDG_DATA_DIRS").?);
+}
+
+test "xdg: existing XDG_DATA_DIRS" {
+    const testing = std.testing;
+
+    var arena = ArenaAllocator.init(testing.allocator);
+    defer arena.deinit();
+    const alloc = arena.allocator();
+
+    var env = EnvMap.init(alloc);
+    defer env.deinit();
+
+    try env.put("XDG_DATA_DIRS", "/opt/share");
+    try setupXdgDataDirs(alloc, ".", &env);
+
+    try testing.expectEqualStrings("./shell-integration", env.get("GHOSTTY_SHELL_INTEGRATION_XDG_DIR").?);
+    try testing.expectEqualStrings("./shell-integration:/opt/share", env.get("XDG_DATA_DIRS").?);
 }
 
 /// Setup the zsh automatic shell integration. This works by setting