1#!/usr/bin/env node
2
3// ./script/cargo is a transparent wrapper around cargo that:
4// - When running in a clone of `./zed-industries/zed`
5// - outputs build timings to the ZED_DATA_DIR/build_timings
6// When Zed starts for staff-members it uploads the build timings to Snowflake
7// To use it:
8// ./script/cargo --init
9// This will add a wrapper to your shell configuration files.
10// (Otherwise set up an alias `cargo=PATH_TO_THIS_FILE`)
11
12// We need to ignore SIGINT in this process so that we can continue
13// processing timing files after the child cargo process exits.
14// The signal will still be delivered to the child process.
15process.on("SIGINT", () => {});
16
17const { spawn, spawnSync } = require("child_process");
18const fs = require("fs");
19const os = require("os");
20const path = require("path");
21const readline = require("readline");
22
23const SUBCOMMANDS_WITH_TIMINGS = new Set(["build", "check", "run", "test"]);
24
25// Built-in cargo aliases
26const CARGO_ALIASES = {
27 b: "build",
28 c: "check",
29 t: "test",
30 r: "run",
31 d: "doc",
32};
33
34function expandAlias(subcommand) {
35 return CARGO_ALIASES[subcommand] || subcommand;
36}
37
38function detectShell() {
39 // Check for PowerShell first (works when running from pwsh)
40 if (process.env.PSModulePath && !process.env.BASH_VERSION) {
41 return "powershell";
42 }
43
44 const shell = process.env.SHELL || "";
45 if (shell.endsWith("/zsh")) return "zsh";
46 if (shell.endsWith("/bash")) return "bash";
47 if (shell.endsWith("/fish")) return "fish";
48 if (shell.endsWith("/pwsh") || shell.endsWith("/powershell")) return "powershell";
49 return path.basename(shell) || "unknown";
50}
51
52function getShellConfigPath(shell) {
53 const home = os.homedir();
54 switch (shell) {
55 case "zsh":
56 return path.join(process.env.ZDOTDIR || home, ".zshrc");
57 case "bash":
58 // Prefer .bashrc, fall back to .bash_profile
59 const bashrc = path.join(home, ".bashrc");
60 if (fs.existsSync(bashrc)) return bashrc;
61 return path.join(home, ".bash_profile");
62 case "fish":
63 return path.join(home, ".config", "fish", "config.fish");
64 case "powershell":
65 // PowerShell Core (pwsh) profile locations
66 if (process.platform === "win32") {
67 return path.join(home, "Documents", "PowerShell", "Microsoft.PowerShell_profile.ps1");
68 } else {
69 return path.join(home, ".config", "powershell", "Microsoft.PowerShell_profile.ps1");
70 }
71 default:
72 return null;
73 }
74}
75
76function generateAlias(shell, scriptDir) {
77 const cargoWrapper = path.join(scriptDir, "cargo");
78
79 switch (shell) {
80 case "zsh":
81 case "bash":
82 return `\n# Zed cargo timing wrapper\ncargo() { local w="${cargoWrapper}"; [[ -x "$w" ]] && "$w" "$@" || command cargo "$@"; }\n`;
83 case "fish":
84 return `\n# Zed cargo timing wrapper\nfunction cargo\n set -l w "${cargoWrapper}"\n if test -x "$w"\n "$w" $argv\n else\n command cargo $argv\n end\nend\n`;
85 case "powershell":
86 return `\n# Zed cargo timing wrapper\nfunction cargo {\n \$wrapper = "${cargoWrapper}"\n if (Test-Path \$wrapper) {\n & \$wrapper @args\n } else {\n & (Get-Command -Name cargo -CommandType Application | Select-Object -First 1).Source @args\n }\n}\n`;
87 default:
88 return `cargo() { local w="${cargoWrapper}"; [[ -x "$w" ]] && "$w" "$@" || command cargo "$@"; }`;
89 }
90}
91
92function initShellAlias() {
93 const scriptDir = __dirname;
94 const shell = detectShell();
95 const configPath = getShellConfigPath(shell);
96 const alias = generateAlias(shell, scriptDir);
97
98 if (!configPath) {
99 console.log(`Unsupported shell: ${shell}`);
100 console.log("\nAdd the following to your shell configuration:\n");
101 console.log(alias);
102 return;
103 }
104
105 // Check if alias already exists
106 if (fs.existsSync(configPath)) {
107 const content = fs.readFileSync(configPath, "utf-8");
108 if (content.includes("Zed cargo timing wrapper")) {
109 console.log(`Alias already exists in ${configPath}`);
110 console.log("To update, remove the existing alias and run --init again.");
111 return;
112 }
113 }
114
115 // Create parent directory if needed (for PowerShell on Linux/macOS)
116 const configDir = path.dirname(configPath);
117 if (!fs.existsSync(configDir)) {
118 fs.mkdirSync(configDir, { recursive: true });
119 }
120
121 // Append alias to config file
122 fs.appendFileSync(configPath, alias);
123 console.log(`Added cargo timing alias to ${configPath}`);
124
125 if (shell === "powershell") {
126 console.log(`\nRestart PowerShell or run: . "${configPath}"`);
127 } else {
128 console.log(`\nRestart your shell or run: source ${configPath}`);
129 }
130}
131
132function isZedRepo() {
133 try {
134 const result = spawnSync("git", ["remote", "-v"], {
135 encoding: "utf-8",
136 stdio: ["pipe", "pipe", "pipe"],
137 timeout: 5000,
138 });
139 if (result.status !== 0 || !result.stdout) {
140 return false;
141 }
142 return result.stdout.includes("zed-industries/zed");
143 } catch {
144 return false;
145 }
146}
147
148function findSubcommand(args) {
149 for (let i = 0; i < args.length; i++) {
150 const arg = args[i];
151 // Skip flags and their values
152 if (arg.startsWith("-")) {
153 // If this flag takes a value and it's not using = syntax, skip the next arg too
154 if (!arg.includes("=") && i + 1 < args.length && !args[i + 1].startsWith("-")) {
155 i++;
156 }
157 continue;
158 }
159 // First non-flag argument is the subcommand
160 return { subcommand: arg, index: i };
161 }
162 return null;
163}
164
165function findLatestTimingFile(targetDir) {
166 const timingsDir = path.join(targetDir, "cargo-timings");
167 if (!fs.existsSync(timingsDir)) {
168 return null;
169 }
170
171 const files = fs
172 .readdirSync(timingsDir)
173 .filter((f) => f.startsWith("cargo-timing-") && f.endsWith(".html") && f !== "cargo-timing.html")
174 .map((f) => ({
175 name: f,
176 path: path.join(timingsDir, f),
177 mtime: fs.statSync(path.join(timingsDir, f)).mtime.getTime(),
178 }))
179 .sort((a, b) => b.mtime - a.mtime);
180
181 return files.length > 0 ? files[0].path : null;
182}
183
184function getTargetDir(args) {
185 // Check for --target-dir flag
186 for (let i = 0; i < args.length; i++) {
187 if (args[i] === "--target-dir" && i + 1 < args.length) {
188 return args[i + 1];
189 }
190 if (args[i].startsWith("--target-dir=")) {
191 return args[i].substring("--target-dir=".length);
192 }
193 }
194 // Default target directory
195 return "target";
196}
197
198function runCargoPassthrough(args) {
199 const cargoCmd = process.env.CARGO || "cargo";
200 const result = spawnSync(cargoCmd, args, {
201 stdio: "inherit",
202 shell: false,
203 });
204 process.exit(result.status ?? 1);
205}
206
207async function main() {
208 const args = process.argv.slice(2);
209
210 // Handle --init flag
211 if (args[0] === "--init") {
212 if (process.env.NIX_WRAPPER === "1") {
213 console.error("`--init` not supported when going through the nix wrapper");
214 process.exit(1);
215 }
216 initShellAlias();
217 return;
218 }
219
220 // If not in zed repo, just pass through to cargo
221 if (!isZedRepo()) {
222 runCargoPassthrough(args);
223 return;
224 }
225
226 const cargoCmd = process.env.CARGO || "cargo";
227 const subcommandInfo = findSubcommand(args);
228 const expandedSubcommand = subcommandInfo ? expandAlias(subcommandInfo.subcommand) : null;
229 const shouldAddTimings = expandedSubcommand && SUBCOMMANDS_WITH_TIMINGS.has(expandedSubcommand);
230
231 // Build the final args array
232 let finalArgs = [...args];
233 if (shouldAddTimings) {
234 // Check if --timings is already present
235 const hasTimings = args.some((arg) => arg === "--timings" || arg.startsWith("--timings="));
236 if (!hasTimings) {
237 // Insert --timings after the subcommand
238 finalArgs.splice(subcommandInfo.index + 1, 0, "--timings");
239 }
240 }
241
242 // Run cargo asynchronously so we can handle signals properly
243 const child = spawn(cargoCmd, finalArgs, {
244 stdio: "inherit",
245 shell: false,
246 });
247
248 // Wait for the child to exit
249 const result = await new Promise((resolve) => {
250 child.on("exit", (code, signal) => {
251 resolve({ status: code, signal });
252 });
253 });
254
255 // If we added timings, try to process the timing file (regardless of cargo's exit status)
256 if (shouldAddTimings) {
257 const targetDir = getTargetDir(args);
258 const timingFile = findLatestTimingFile(targetDir);
259
260 if (timingFile) {
261 // Run cargo-timing-info.js in the background
262 const scriptDir = __dirname;
263 const timingScript = path.join(scriptDir, "cargo-timing-info.js");
264
265 if (fs.existsSync(timingScript)) {
266 const timingChild = spawn(process.execPath, [timingScript, timingFile, `cargo ${expandedSubcommand}`], {
267 detached: true,
268 stdio: "ignore",
269 });
270 timingChild.unref();
271 }
272 }
273 }
274
275 // Exit with cargo's exit code, or re-raise the signal if it was killed
276 if (result.signal) {
277 // Reset signal handler and re-raise so parent sees the signal
278 process.removeAllListeners(result.signal);
279 process.kill(process.pid, result.signal);
280 } else {
281 process.exit(result.status ?? 1);
282 }
283}
284
285main();