ripgrep.ts 9.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353
  1. // Ripgrep utility functions
  2. import path from "path"
  3. import { Global } from "../global"
  4. import fs from "fs/promises"
  5. import { z } from "zod"
  6. import { NamedError } from "../util/error"
  7. import { lazy } from "../util/lazy"
  8. import { $ } from "bun"
  9. import { Fzf } from "./fzf"
  10. export namespace Ripgrep {
  11. const Stats = z.object({
  12. elapsed: z.object({
  13. secs: z.number(),
  14. nanos: z.number(),
  15. human: z.string(),
  16. }),
  17. searches: z.number(),
  18. searches_with_match: z.number(),
  19. bytes_searched: z.number(),
  20. bytes_printed: z.number(),
  21. matched_lines: z.number(),
  22. matches: z.number(),
  23. })
  24. const Begin = z.object({
  25. type: z.literal("begin"),
  26. data: z.object({
  27. path: z.object({
  28. text: z.string(),
  29. }),
  30. }),
  31. })
  32. const Match = z.object({
  33. type: z.literal("match"),
  34. data: z.object({
  35. path: z.object({
  36. text: z.string(),
  37. }),
  38. lines: z.object({
  39. text: z.string(),
  40. }),
  41. line_number: z.number(),
  42. absolute_offset: z.number(),
  43. submatches: z.array(
  44. z.object({
  45. match: z.object({
  46. text: z.string(),
  47. }),
  48. start: z.number(),
  49. end: z.number(),
  50. }),
  51. ),
  52. }),
  53. })
  54. const End = z.object({
  55. type: z.literal("end"),
  56. data: z.object({
  57. path: z.object({
  58. text: z.string(),
  59. }),
  60. binary_offset: z.number().nullable(),
  61. stats: Stats,
  62. }),
  63. })
  64. const Summary = z.object({
  65. type: z.literal("summary"),
  66. data: z.object({
  67. elapsed_total: z.object({
  68. human: z.string(),
  69. nanos: z.number(),
  70. secs: z.number(),
  71. }),
  72. stats: Stats,
  73. }),
  74. })
  75. const Result = z.union([Begin, Match, End, Summary])
  76. export type Result = z.infer<typeof Result>
  77. export type Match = z.infer<typeof Match>
  78. export type Begin = z.infer<typeof Begin>
  79. export type End = z.infer<typeof End>
  80. export type Summary = z.infer<typeof Summary>
  81. const PLATFORM = {
  82. "arm64-darwin": { platform: "aarch64-apple-darwin", extension: "tar.gz" },
  83. "arm64-linux": {
  84. platform: "aarch64-unknown-linux-gnu",
  85. extension: "tar.gz",
  86. },
  87. "x64-darwin": { platform: "x86_64-apple-darwin", extension: "tar.gz" },
  88. "x64-linux": { platform: "x86_64-unknown-linux-musl", extension: "tar.gz" },
  89. "x64-win32": { platform: "x86_64-pc-windows-msvc", extension: "zip" },
  90. } as const
  91. export const ExtractionFailedError = NamedError.create(
  92. "RipgrepExtractionFailedError",
  93. z.object({
  94. filepath: z.string(),
  95. stderr: z.string(),
  96. }),
  97. )
  98. export const UnsupportedPlatformError = NamedError.create(
  99. "RipgrepUnsupportedPlatformError",
  100. z.object({
  101. platform: z.string(),
  102. }),
  103. )
  104. export const DownloadFailedError = NamedError.create(
  105. "RipgrepDownloadFailedError",
  106. z.object({
  107. url: z.string(),
  108. status: z.number(),
  109. }),
  110. )
  111. const state = lazy(async () => {
  112. let filepath = Bun.which("rg")
  113. if (filepath) return { filepath }
  114. filepath = path.join(
  115. Global.Path.bin,
  116. "rg" + (process.platform === "win32" ? ".exe" : ""),
  117. )
  118. const file = Bun.file(filepath)
  119. if (!(await file.exists())) {
  120. const platformKey =
  121. `${process.arch}-${process.platform}` as keyof typeof PLATFORM
  122. const config = PLATFORM[platformKey]
  123. if (!config) throw new UnsupportedPlatformError({ platform: platformKey })
  124. const version = "14.1.1"
  125. const filename = `ripgrep-${version}-${config.platform}.${config.extension}`
  126. const url = `https://github.com/BurntSushi/ripgrep/releases/download/${version}/${filename}`
  127. const response = await fetch(url)
  128. if (!response.ok)
  129. throw new DownloadFailedError({ url, status: response.status })
  130. const buffer = await response.arrayBuffer()
  131. const archivePath = path.join(Global.Path.bin, filename)
  132. await Bun.write(archivePath, buffer)
  133. if (config.extension === "tar.gz") {
  134. const args = ["tar", "-xzf", archivePath, "--strip-components=1"]
  135. if (platformKey.endsWith("-darwin")) args.push("--include=*/rg")
  136. if (platformKey.endsWith("-linux")) args.push("--wildcards", "*/rg")
  137. const proc = Bun.spawn(args, {
  138. cwd: Global.Path.bin,
  139. stderr: "pipe",
  140. stdout: "pipe",
  141. })
  142. await proc.exited
  143. if (proc.exitCode !== 0)
  144. throw new ExtractionFailedError({
  145. filepath,
  146. stderr: await Bun.readableStreamToText(proc.stderr),
  147. })
  148. }
  149. if (config.extension === "zip") {
  150. const proc = Bun.spawn(
  151. ["unzip", "-j", archivePath, "*/rg.exe", "-d", Global.Path.bin],
  152. {
  153. cwd: Global.Path.bin,
  154. stderr: "pipe",
  155. stdout: "ignore",
  156. },
  157. )
  158. await proc.exited
  159. if (proc.exitCode !== 0)
  160. throw new ExtractionFailedError({
  161. filepath: archivePath,
  162. stderr: await Bun.readableStreamToText(proc.stderr),
  163. })
  164. }
  165. await fs.unlink(archivePath)
  166. if (!platformKey.endsWith("-win32")) await fs.chmod(filepath, 0o755)
  167. }
  168. return {
  169. filepath,
  170. }
  171. })
  172. export async function filepath() {
  173. const { filepath } = await state()
  174. return filepath
  175. }
  176. export async function files(input: {
  177. cwd: string
  178. query?: string
  179. glob?: string
  180. limit?: number
  181. }) {
  182. const commands = [
  183. `${await filepath()} --files --hidden --glob='!.git/*' ${input.glob ? `--glob='${input.glob}'` : ``}`,
  184. ]
  185. if (input.query)
  186. commands.push(`${await Fzf.filepath()} --filter=${input.query}`)
  187. if (input.limit) commands.push(`head -n ${input.limit}`)
  188. const joined = commands.join(" | ")
  189. const result = await $`${{ raw: joined }}`.cwd(input.cwd).nothrow().text()
  190. return result.split("\n").filter(Boolean)
  191. }
  192. export async function tree(input: { cwd: string; limit?: number }) {
  193. const files = await Ripgrep.files({ cwd: input.cwd })
  194. interface Node {
  195. path: string[]
  196. children: Node[]
  197. }
  198. function getPath(node: Node, parts: string[], create: boolean) {
  199. if (parts.length === 0) return node
  200. let current = node
  201. for (const part of parts) {
  202. let existing = current.children.find((x) => x.path.at(-1) === part)
  203. if (!existing) {
  204. if (!create) return
  205. existing = {
  206. path: current.path.concat(part),
  207. children: [],
  208. }
  209. current.children.push(existing)
  210. }
  211. current = existing
  212. }
  213. return current
  214. }
  215. const root: Node = {
  216. path: [],
  217. children: [],
  218. }
  219. for (const file of files) {
  220. const parts = file.split(path.sep)
  221. getPath(root, parts, true)
  222. }
  223. function sort(node: Node) {
  224. node.children.sort((a, b) => {
  225. if (!a.children.length && b.children.length) return 1
  226. if (!b.children.length && a.children.length) return -1
  227. return a.path.at(-1)!.localeCompare(b.path.at(-1)!)
  228. })
  229. for (const child of node.children) {
  230. sort(child)
  231. }
  232. }
  233. sort(root)
  234. let current = [root]
  235. const result: Node = {
  236. path: [],
  237. children: [],
  238. }
  239. let processed = 0
  240. const limit = input.limit ?? 50
  241. while (current.length > 0) {
  242. const next = []
  243. for (const node of current) {
  244. if (node.children.length) next.push(...node.children)
  245. }
  246. const max = Math.max(...current.map((x) => x.children.length))
  247. for (let i = 0; i < max && processed < limit; i++) {
  248. for (const node of current) {
  249. const child = node.children[i]
  250. if (!child) continue
  251. getPath(result, child.path, true)
  252. processed++
  253. if (processed >= limit) break
  254. }
  255. }
  256. if (processed >= limit) {
  257. for (const node of [...current, ...next]) {
  258. const compare = getPath(result, node.path, false)
  259. if (!compare) continue
  260. if (compare?.children.length !== node.children.length) {
  261. const diff = node.children.length - compare.children.length
  262. compare.children.push({
  263. path: compare.path.concat(`[${diff} truncated]`),
  264. children: [],
  265. })
  266. }
  267. }
  268. break
  269. }
  270. current = next
  271. }
  272. const lines: string[] = []
  273. function render(node: Node, depth: number) {
  274. const indent = "\t".repeat(depth)
  275. lines.push(indent + node.path.at(-1) + (node.children.length ? "/" : ""))
  276. for (const child of node.children) {
  277. render(child, depth + 1)
  278. }
  279. }
  280. result.children.map((x) => render(x, 0))
  281. return lines.join("\n")
  282. }
  283. export async function search(input: {
  284. cwd: string
  285. pattern: string
  286. glob?: string[]
  287. limit?: number
  288. }) {
  289. const args = [
  290. `${await filepath()}`,
  291. "--json",
  292. "--hidden",
  293. "--glob='!.git/*'",
  294. ]
  295. if (input.glob) {
  296. for (const g of input.glob) {
  297. args.push(`--glob=${g}`)
  298. }
  299. }
  300. if (input.limit) {
  301. args.push(`--max-count=${input.limit}`)
  302. }
  303. args.push(input.pattern)
  304. const command = args.join(" ")
  305. const result = await $`${{ raw: command }}`.cwd(input.cwd).quiet().nothrow()
  306. if (result.exitCode !== 0) {
  307. return []
  308. }
  309. const lines = result.text().trim().split("\n").filter(Boolean)
  310. // Parse JSON lines from ripgrep output
  311. return lines
  312. .map((line) => JSON.parse(line))
  313. .map((parsed) => Result.parse(parsed))
  314. .filter((r) => r.type === "match")
  315. .map((r) => r.data)
  316. }
  317. }