codecrafters submit [skip ci]

3 files changed, 144 insertions, 112 deletions

diff --git a/src/executor.zig b/src/executor.zig
index 354e724..f776e74 100644
--- a/src/executor.zig
+++ b/src/executor.zig
@@ -134,81 +134,82 @@ fn runBuiltinInChild(
 
     std.posix.exit(0);
 }
-
 pub fn runPipeline(
     allocator: std.mem.Allocator,
-    left_path: ?[]const u8,
-    left_builtin: bool,
-    left_name: []const u8,
-    left_args: ?[]const u8,
-    left_argv: ?[]const []const u8,
-    right_path: ?[]const u8,
-    right_builtin: bool,
-    right_name: []const u8,
-    right_args: ?[]const u8,
-    right_argv: ?[]const []const u8,
+    stages: []const Stage,
 ) !void {
-    const fds = try std.posix.pipe();
-
-    const pid1 = try std.posix.fork();
-    if (pid1 == 0) {
-        std.posix.close(fds[0]);
-        try std.posix.dup2(fds[1], 1);
-        std.posix.close(fds[1]);
-
-        if (left_builtin) {
-            runBuiltinInChild(allocator, left_name, left_args, 0, 1);
+    if (stages.len == 0) return;
+    if (stages.len == 1) {
+        const s = stages[0];
+        if (s.is_builtin) {
+            runBuiltinInChild(allocator, s.name, s.args, 0, 1);
         } else {
-            const argv1_z = try allocator.allocSentinel(?[*:0]const u8, left_argv.?.len, @as(?[*:0]const u8, null));
-            defer allocator.free(argv1_z);
-            for (left_argv.?, 0..) |arg, i| {
-                argv1_z[i] = (try allocator.dupeZ(u8, arg)).ptr;
-            }
+            const argv_z = try allocator.allocSentinel(?[*:0]const u8, s.argv.?.len, @as(?[*:0]const u8, null));
+            defer allocator.free(argv_z);
+            for (s.argv.?, 0..) |arg, i| argv_z[i] = (try allocator.dupeZ(u8, arg)).ptr;
             defer {
-                for (argv1_z[0..left_argv.?.len]) |arg_ptr| {
-                    if (arg_ptr) |ptr| allocator.free(std.mem.span(ptr));
-                }
+                for (argv_z[0..s.argv.?.len]) |arg_ptr| if (arg_ptr) |ptr| allocator.free(std.mem.span(ptr));
             }
-            const program1_path_z = try allocator.dupeZ(u8, left_path.?);
-            defer allocator.free(program1_path_z);
-            _ = std.posix.execveZ(program1_path_z, argv1_z, std.c.environ) catch {
-                std.posix.exit(1);
-            };
-            unreachable;
+            const path_z = try allocator.dupeZ(u8, s.path.?);
+            defer allocator.free(path_z);
+            _ = std.posix.execveZ(path_z, argv_z, std.c.environ) catch std.posix.exit(1);
         }
     }
 
-    const pid2 = try std.posix.fork();
-    if (pid2 == 0) {
-        std.posix.close(fds[1]);
-        try std.posix.dup2(fds[0], 0);
-        std.posix.close(fds[0]);
+    const pipes = try allocator.alloc([2]std.posix.fd_t, stages.len - 1);
+    defer allocator.free(pipes);
+    for (pipes) |*p| p.* = try std.posix.pipe();
 
-        if (right_builtin) {
-            runBuiltinInChild(allocator, right_name, right_args, 0, 1);
-        } else {
-            const argv2_z = try allocator.allocSentinel(?[*:0]const u8, right_argv.?.len, @as(?[*:0]const u8, null));
-            defer allocator.free(argv2_z);
-            for (right_argv.?, 0..) |arg, i| {
-                argv2_z[i] = (try allocator.dupeZ(u8, arg)).ptr;
+    var pids = try allocator.alloc(std.posix.pid_t, stages.len);
+    defer allocator.free(pids);
+
+    for (stages, 0..) |s, idx| {
+        const pid = try std.posix.fork();
+        if (pid == 0) {
+            if (idx > 0) {
+                try std.posix.dup2(pipes[idx - 1][0], 0);
             }
-            defer {
-                for (argv2_z[0..right_argv.?.len]) |arg_ptr| {
-                    if (arg_ptr) |ptr| allocator.free(std.mem.span(ptr));
+            if (idx + 1 < stages.len) {
+                try std.posix.dup2(pipes[idx][1], 1);
+            }
+
+            for (pipes) |p| {
+                std.posix.close(p[0]);
+                std.posix.close(p[1]);
+            }
+
+            if (s.is_builtin) {
+                runBuiltinInChild(allocator, s.name, s.args, 0, 1);
+            } else {
+                const argv_z = try allocator.allocSentinel(?[*:0]const u8, s.argv.?.len, @as(?[*:0]const u8, null));
+                defer allocator.free(argv_z);
+                for (s.argv.?, 0..) |arg, i| argv_z[i] = (try allocator.dupeZ(u8, arg)).ptr;
+                defer {
+                    for (argv_z[0..s.argv.?.len]) |arg_ptr| if (arg_ptr) |ptr| allocator.free(std.mem.span(ptr));
                 }
+                const path_z = try allocator.dupeZ(u8, s.path.?);
+                defer allocator.free(path_z);
+                _ = std.posix.execveZ(path_z, argv_z, std.c.environ) catch std.posix.exit(1);
             }
-            const program2_path_z = try allocator.dupeZ(u8, right_path.?);
-            defer allocator.free(program2_path_z);
-            _ = std.posix.execveZ(program2_path_z, argv2_z, std.c.environ) catch {
-                std.posix.exit(1);
-            };
             unreachable;
         }
+        pids[idx] = pid;
     }
 
-    std.posix.close(fds[0]);
-    std.posix.close(fds[1]);
+    for (pipes) |p| {
+        std.posix.close(p[0]);
+        std.posix.close(p[1]);
+    }
 
-    _ = std.posix.waitpid(pid1, 0);
-    _ = std.posix.waitpid(pid2, 0);
+    for (pids) |pid| {
+        _ = std.posix.waitpid(pid, 0);
+    }
 }
+
+pub const Stage = struct {
+    is_builtin: bool,
+    name: []const u8,
+    args: ?[]const u8,
+    path: ?[]const u8,
+    argv: ?[]const []const u8,
+};
diff --git a/src/main.zig b/src/main.zig
index 20acf0f..74140cd 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -297,18 +297,39 @@ pub fn main() !void {
             var stdout_writer = std.fs.File.stdout().writerStreaming(&.{});
             const stdout_iface = &stdout_writer.interface;
 
-            if (std.mem.indexOfScalar(u8, cmd, '|')) |pipe_pos| {
-                var left = cmd[0..pipe_pos];
-                var right = if (pipe_pos + 1 < cmd.len) cmd[pipe_pos + 1 ..] else cmd[pipe_pos..pipe_pos];
-                while (left.len > 0 and left[0] == ' ') left = left[1..];
-                while (left.len > 0 and left[left.len - 1] == ' ') left = left[0 .. left.len - 1];
-                while (right.len > 0 and right[0] == ' ') right = right[1..];
-                while (right.len > 0 and right[right.len - 1] == ' ') right = right[0 .. right.len - 1];
-
-                if (left.len == 0 or right.len == 0) {
+            if (std.mem.indexOfScalar(u8, cmd, '|')) |_| {
+                var segments = std.ArrayList([]const u8){};
+                defer segments.deinit(allocator);
+
+                var invalid = false;
+                var start: usize = 0;
+                for (cmd, 0..) |ch, i| {
+                    if (ch == '|') {
+                        const seg = std.mem.trim(u8, cmd[start..i], " ");
+                        if (seg.len == 0) {
+                            try stdout.writeAll("pipe: command not found\n");
+                            invalid = true;
+                            break;
+                        }
+                        try segments.append(allocator, seg);
+                        start = i + 1;
+                    }
+                }
+                if (invalid) continue;
+
+                const last_seg = std.mem.trim(u8, cmd[start..], " ");
+                if (last_seg.len == 0) {
                     try stdout.writeAll("pipe: command not found\n");
+                    continue;
+                }
+                try segments.append(allocator, last_seg);
+
+                if (segments.items.len == 1) {
+                    const parsed = parser.parseCommand(cmd);
+                    const result = try shell.executeCommand(allocator, stdout_iface, parsed.name, parsed.args, parsed.output_redirect, parsed.error_redirect, parsed.append_output, parsed.append_error);
+                    if (result == .exit_shell) break;
                 } else {
-                    const result = try shell.executePipeline(allocator, stdout_iface, left, right);
+                    const result = try shell.executePipeline(allocator, stdout_iface, segments.items);
                     if (result == .exit_shell) break;
                 }
             } else {
diff --git a/src/shell.zig b/src/shell.zig
index 15046d5..271e3ea 100644
--- a/src/shell.zig
+++ b/src/shell.zig
@@ -4,6 +4,12 @@ const builtins = @import("builtins.zig");
 const path = @import("path.zig");
 const executor = @import("executor.zig");
 
+fn freeArgv(allocator: std.mem.Allocator, argv: []const []const u8) void {
+    // argv[0] is cmd_name (not allocated), argv[1..] are allocated by parseArgs
+    for (argv[1..]) |arg| allocator.free(arg);
+    allocator.free(argv);
+}
+
 pub fn executeCommand(
     allocator: std.mem.Allocator,
     stdout: anytype,
@@ -121,7 +127,10 @@ pub fn executeCommand(
         defer allocator.free(program_path);
 
         const argv = try parser.parseArgs(allocator, cmd_name, args);
-        defer allocator.free(argv);
+        defer {
+            for (argv[1..]) |arg| allocator.free(arg);
+            allocator.free(argv);
+        }
 
         if (output_redirect != null or error_redirect != null or append_output != null or append_error != null) {
             try executor.runExternalProgramWithRedirect(allocator, program_path, argv, output_redirect, error_redirect, append_output, append_error);
@@ -138,56 +147,57 @@ pub fn executeCommand(
 pub fn executePipeline(
     allocator: std.mem.Allocator,
     stdout: anytype,
-    left_cmd: []const u8,
-    right_cmd: []const u8,
+    commands: []const []const u8,
 ) !builtins.CommandResult {
-    const left_parsed = parser.parseCommand(left_cmd);
-    const right_parsed = parser.parseCommand(right_cmd);
+    var stages = std.ArrayList(executor.Stage){};
+    defer stages.deinit(allocator);
 
-    if (left_parsed.name.len == 0) {
-        try stdout.print("{s}: command not found\n", .{left_cmd});
-        return .continue_loop;
+    var owned_paths = std.ArrayList(?[]const u8){};
+    defer {
+        for (owned_paths.items) |p| if (p) |path_buf| allocator.free(path_buf);
+        owned_paths.deinit(allocator);
     }
-    if (right_parsed.name.len == 0) {
-        try stdout.print("{s}: command not found\n", .{right_cmd});
-        return .continue_loop;
+
+    var owned_argvs = std.ArrayList(?[]const []const u8){};
+    defer {
+        for (owned_argvs.items) |argv_opt| if (argv_opt) |argv| freeArgv(allocator, argv);
+        owned_argvs.deinit(allocator);
     }
 
-    const left_is_builtin = builtins.isBuiltin(left_parsed.name);
-    const right_is_builtin = builtins.isBuiltin(right_parsed.name);
+    for (commands) |cmd_part| {
+        const parsed = parser.parseCommand(cmd_part);
+        if (parsed.name.len == 0) {
+            try stdout.print("{s}: command not found\n", .{cmd_part});
+            return .continue_loop;
+        }
+
+        const is_builtin = builtins.isBuiltin(parsed.name);
+        var cmd_path: ?[]const u8 = null;
+        if (!is_builtin) {
+            cmd_path = try path.findInPath(allocator, parsed.name);
+        }
+        if (!is_builtin and cmd_path == null) {
+            try stdout.print("{s}: command not found\n", .{parsed.name});
+            return .continue_loop;
+        }
 
-    const left_path = if (left_is_builtin) null else try path.findInPath(allocator, left_parsed.name);
-    defer if (left_path) |p| allocator.free(p);
-    if (!left_is_builtin and left_path == null) {
-        try stdout.print("{s}: command not found\n", .{left_parsed.name});
-        return .continue_loop;
-    }
+        try owned_paths.append(allocator, cmd_path);
 
-    const right_path = if (right_is_builtin) null else try path.findInPath(allocator, right_parsed.name);
-    defer if (right_path) |p| allocator.free(p);
-    if (!right_is_builtin and right_path == null) {
-        try stdout.print("{s}: command not found\n", .{right_parsed.name});
-        return .continue_loop;
+        var argv: ?[]const []const u8 = null;
+        if (!is_builtin) {
+            argv = try parser.parseArgs(allocator, parsed.name, parsed.args);
+        }
+        try owned_argvs.append(allocator, argv);
+
+        try stages.append(allocator, .{
+            .is_builtin = is_builtin,
+            .name = parsed.name,
+            .args = parsed.args,
+            .path = cmd_path,
+            .argv = argv,
+        });
     }
 
-    const left_argv = if (left_is_builtin) null else try parser.parseArgs(allocator, left_parsed.name, left_parsed.args);
-    defer if (left_argv) |argv| allocator.free(argv);
-
-    const right_argv = if (right_is_builtin) null else try parser.parseArgs(allocator, right_parsed.name, right_parsed.args);
-    defer if (right_argv) |argv| allocator.free(argv);
-
-    try executor.runPipeline(
-        allocator,
-        left_path,
-        left_is_builtin,
-        left_parsed.name,
-        left_parsed.args,
-        left_argv,
-        right_path,
-        right_is_builtin,
-        right_parsed.name,
-        right_parsed.args,
-        right_argv,
-    );
+    try executor.runPipeline(allocator, stages.items);
     return .continue_loop;
 }