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 if (process.platform === "win32") {
66 // Spawn PowerShell to get the real $PROFILE path, since os.homedir() doesn't account
67 // for OneDrive folder redirection, and the subdirectory differs between Windows PowerShell
68 // 5.x ("WindowsPowerShell") and PowerShell Core ("PowerShell").
69 const psModulePath = process.env.PSModulePath || "";
70 const psExe = psModulePath.toLowerCase().includes("\\windowspowershell\\") ? "powershell" : "pwsh";
71 const result = spawnSync(psExe, ["-NoProfile", "-Command", "$PROFILE"], {
72 encoding: "utf-8",
73 stdio: ["pipe", "pipe", "pipe"],
74 timeout: 5000,
75 });
76 if (result.status === 0 && result.stdout.trim()) {
77 return result.stdout.trim();
78 }
79 // Fallback if spawning fails
80 return path.join(home, "Documents", "PowerShell", "Microsoft.PowerShell_profile.ps1");
81 } else {
82 return path.join(home, ".config", "powershell", "Microsoft.PowerShell_profile.ps1");
83 }
84 default:
85 return null;
86 }
87}
88
89function generateAlias(shell, scriptDir) {
90 const cargoWrapper = path.join(scriptDir, "cargo");
91
92 switch (shell) {
93 case "zsh":
94 case "bash":
95 return `\n# Zed cargo timing wrapper\ncargo() { local w="${cargoWrapper}"; if [[ -x "$w" ]]; then "$w" "$@"; else command cargo "$@"; fi; }\n`;
96 case "fish":
97 return `\n# Zed cargo timing wrapper\nfunction cargo\n set -l w "${cargoWrapper}"\n if test -x "$w"\n "$w" $argv\n return $status\n else\n command cargo $argv\n end\nend\n`;
98 case "powershell":
99 return `\n# Zed cargo timing wrapper\nfunction cargo {\n \$wrapper = "${cargoWrapper}"\n if (Test-Path \$wrapper) {\n node \$wrapper @args\n } else {\n & (Get-Command -Name cargo -CommandType Application | Select-Object -First 1).Source @args\n }\n}\n`;
100 default:
101 return `cargo() { local w="${cargoWrapper}"; if [[ -x "$w" ]]; then "$w" "$@"; else command cargo "$@"; fi; }`;
102 }
103}
104
105function aliasBlockRegex(shell) {
106 switch (shell) {
107 case "zsh":
108 case "bash":
109 // Comment line + single-line cargo() { ... } function
110 return /\n?# Zed cargo timing wrapper\ncargo\(\) \{[^\n]*\}\n/;
111 case "fish":
112 // Comment line + multi-line function cargo...end block
113 return /\n?# Zed cargo timing wrapper\nfunction cargo\n[\s\S]*?\nend\n/;
114 case "powershell":
115 // Comment line + multi-line function cargo {...} block
116 return /\n?# Zed cargo timing wrapper\nfunction cargo \{[\s\S]*?\n\}\n/;
117 default:
118 return null;
119 }
120}
121
122function initShellAlias() {
123 const scriptDir = __dirname;
124 const shell = detectShell();
125 const configPath = getShellConfigPath(shell);
126 const alias = generateAlias(shell, scriptDir);
127
128 if (!configPath) {
129 console.log(`Unsupported shell: ${shell}`);
130 console.log("\nAdd the following to your shell configuration:\n");
131 console.log(alias);
132 return;
133 }
134
135 // Check if alias already exists; if so, replace it in-place
136 if (fs.existsSync(configPath)) {
137 const content = fs.readFileSync(configPath, "utf-8");
138 if (content.includes("Zed cargo timing wrapper")) {
139 const blockRegex = aliasBlockRegex(shell);
140 const updated = blockRegex ? content.replace(blockRegex, "") : content;
141 fs.writeFileSync(configPath, updated + alias);
142 console.log(`Updated cargo timing alias in ${configPath}`);
143 if (shell === "powershell") {
144 console.log(`\nRestart PowerShell or run: . "${configPath}"`);
145 } else {
146 console.log(`\nRestart your shell or run: source ${configPath}`);
147 }
148 return;
149 }
150 }
151
152 // Create parent directory if needed (for PowerShell on Linux/macOS)
153 const configDir = path.dirname(configPath);
154 if (!fs.existsSync(configDir)) {
155 fs.mkdirSync(configDir, { recursive: true });
156 }
157
158 // Append alias to config file
159 fs.appendFileSync(configPath, alias);
160 console.log(`Added cargo timing alias to ${configPath}`);
161
162 if (shell === "powershell") {
163 console.log(`\nRestart PowerShell or run: . "${configPath}"`);
164 } else {
165 console.log(`\nRestart your shell or run: source ${configPath}`);
166 }
167}
168
169function isZedRepo() {
170 try {
171 const result = spawnSync("git", ["remote", "-v"], {
172 encoding: "utf-8",
173 stdio: ["pipe", "pipe", "pipe"],
174 timeout: 5000,
175 });
176 if (result.status !== 0 || !result.stdout) {
177 return false;
178 }
179 return result.stdout.includes("zed-industries/zed");
180 } catch {
181 return false;
182 }
183}
184
185function findSubcommand(args) {
186 for (let i = 0; i < args.length; i++) {
187 const arg = args[i];
188 // Skip flags and their values
189 if (arg.startsWith("-")) {
190 // If this flag takes a value and it's not using = syntax, skip the next arg too
191 if (!arg.includes("=") && i + 1 < args.length && !args[i + 1].startsWith("-")) {
192 i++;
193 }
194 continue;
195 }
196 // First non-flag argument is the subcommand
197 return { subcommand: arg, index: i };
198 }
199 return null;
200}
201
202function findLatestTimingFile(targetDir) {
203 const timingsDir = path.join(targetDir, "cargo-timings");
204 if (!fs.existsSync(timingsDir)) {
205 return null;
206 }
207
208 const files = fs
209 .readdirSync(timingsDir)
210 .filter((f) => f.startsWith("cargo-timing-") && f.endsWith(".html") && f !== "cargo-timing.html")
211 .map((f) => ({
212 name: f,
213 path: path.join(timingsDir, f),
214 mtime: fs.statSync(path.join(timingsDir, f)).mtime.getTime(),
215 }))
216 .sort((a, b) => b.mtime - a.mtime);
217
218 return files.length > 0 ? files[0].path : null;
219}
220
221function getTargetDir(args) {
222 // Check for --target-dir flag
223 for (let i = 0; i < args.length; i++) {
224 if (args[i] === "--target-dir" && i + 1 < args.length) {
225 return args[i + 1];
226 }
227 if (args[i].startsWith("--target-dir=")) {
228 return args[i].substring("--target-dir=".length);
229 }
230 }
231 // Default target directory
232 return "target";
233}
234
235function runCargoPassthrough(args) {
236 const cargoCmd = process.env.CARGO || "cargo";
237 const result = spawnSync(cargoCmd, args, {
238 stdio: "inherit",
239 shell: false,
240 });
241 process.exit(result.status ?? 1);
242}
243
244async function main() {
245 const args = process.argv.slice(2);
246
247 // Handle --init flag
248 if (args[0] === "--init") {
249 if (process.env.NIX_WRAPPER === "1") {
250 console.error("`--init` not supported when going through the nix wrapper");
251 process.exit(1);
252 }
253 initShellAlias();
254 return;
255 }
256
257 // If not in zed repo, just pass through to cargo
258 if (!isZedRepo()) {
259 runCargoPassthrough(args);
260 return;
261 }
262
263 const cargoCmd = process.env.CARGO || "cargo";
264 const subcommandInfo = findSubcommand(args);
265 const expandedSubcommand = subcommandInfo ? expandAlias(subcommandInfo.subcommand) : null;
266 const shouldAddTimings = expandedSubcommand && SUBCOMMANDS_WITH_TIMINGS.has(expandedSubcommand);
267
268 // Build the final args array
269 let finalArgs = [...args];
270 if (shouldAddTimings) {
271 // Check if --timings is already present
272 const hasTimings = args.some((arg) => arg === "--timings" || arg.startsWith("--timings="));
273 if (!hasTimings) {
274 // Insert --timings after the subcommand
275 finalArgs.splice(subcommandInfo.index + 1, 0, "--timings");
276 }
277 }
278
279 // Run cargo asynchronously so we can handle signals properly
280 const child = spawn(cargoCmd, finalArgs, {
281 stdio: "inherit",
282 shell: false,
283 });
284
285 // Wait for the child to exit
286 const result = await new Promise((resolve) => {
287 child.on("exit", (code, signal) => {
288 resolve({ status: code, signal });
289 });
290 });
291
292 // If we added timings, try to process the timing file (regardless of cargo's exit status)
293 if (shouldAddTimings) {
294 const targetDir = getTargetDir(args);
295 const timingFile = findLatestTimingFile(targetDir);
296
297 if (timingFile) {
298 // Run cargo-timing-info.js in the background
299 const scriptDir = __dirname;
300 const timingScript = path.join(scriptDir, "cargo-timing-info.js");
301
302 if (fs.existsSync(timingScript)) {
303 const timingChild = spawn(process.execPath, [timingScript, timingFile, `cargo ${expandedSubcommand}`], {
304 detached: true,
305 stdio: "ignore",
306 });
307 timingChild.unref();
308 }
309 }
310 }
311
312 // Exit with cargo's exit code, or re-raise the signal if it was killed
313 if (result.signal) {
314 // Reset signal handler and re-raise so parent sees the signal
315 process.removeAllListeners(result.signal);
316 process.kill(process.pid, result.signal);
317 } else {
318 process.exit(result.status ?? 1);
319 }
320}
321
322main();