ripgrep.ts 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232
  1. import path from "path"
  2. import { Global } from "../global"
  3. import fs from "fs/promises"
  4. import { z } from "zod"
  5. import { NamedError } from "../util/error"
  6. import { lazy } from "../util/lazy"
  7. import { $ } from "bun"
  8. import { Fzf } from "./fzf"
  9. export namespace Ripgrep {
  10. const PLATFORM = {
  11. darwin: { platform: "apple-darwin", extension: "tar.gz" },
  12. linux: { platform: "unknown-linux-musl", extension: "tar.gz" },
  13. win32: { platform: "pc-windows-msvc", extension: "zip" },
  14. } as const
  15. export const ExtractionFailedError = NamedError.create(
  16. "RipgrepExtractionFailedError",
  17. z.object({
  18. filepath: z.string(),
  19. stderr: z.string(),
  20. }),
  21. )
  22. export const UnsupportedPlatformError = NamedError.create(
  23. "RipgrepUnsupportedPlatformError",
  24. z.object({
  25. platform: z.string(),
  26. }),
  27. )
  28. export const DownloadFailedError = NamedError.create(
  29. "RipgrepDownloadFailedError",
  30. z.object({
  31. url: z.string(),
  32. status: z.number(),
  33. }),
  34. )
  35. const state = lazy(async () => {
  36. let filepath = Bun.which("rg")
  37. if (filepath) return { filepath }
  38. filepath = path.join(
  39. Global.Path.bin,
  40. "rg" + (process.platform === "win32" ? ".exe" : ""),
  41. )
  42. const file = Bun.file(filepath)
  43. if (!(await file.exists())) {
  44. const archMap = { x64: "x86_64", arm64: "aarch64" } as const
  45. const arch = archMap[process.arch as keyof typeof archMap] ?? process.arch
  46. const config = PLATFORM[process.platform as keyof typeof PLATFORM]
  47. if (!config)
  48. throw new UnsupportedPlatformError({ platform: process.platform })
  49. const version = "14.1.1"
  50. const filename = `ripgrep-${version}-${arch}-${config.platform}.${config.extension}`
  51. const url = `https://github.com/BurntSushi/ripgrep/releases/download/${version}/${filename}`
  52. const response = await fetch(url)
  53. if (!response.ok)
  54. throw new DownloadFailedError({ url, status: response.status })
  55. const buffer = await response.arrayBuffer()
  56. const archivePath = path.join(Global.Path.bin, filename)
  57. await Bun.write(archivePath, buffer)
  58. if (config.extension === "tar.gz") {
  59. const args = ["tar", "-xzf", archivePath, "--strip-components=1"]
  60. if (process.platform === "darwin") args.push("--include=*/rg")
  61. if (process.platform === "linux") args.push("--wildcards", "*/rg")
  62. const proc = Bun.spawn(args, {
  63. cwd: Global.Path.bin,
  64. stderr: "pipe",
  65. stdout: "pipe",
  66. })
  67. await proc.exited
  68. if (proc.exitCode !== 0)
  69. throw new ExtractionFailedError({
  70. filepath,
  71. stderr: await Bun.readableStreamToText(proc.stderr),
  72. })
  73. }
  74. if (config.extension === "zip") {
  75. const proc = Bun.spawn(
  76. ["unzip", "-j", archivePath, "*/rg.exe", "-d", Global.Path.bin],
  77. {
  78. cwd: Global.Path.bin,
  79. stderr: "pipe",
  80. stdout: "ignore",
  81. },
  82. )
  83. await proc.exited
  84. if (proc.exitCode !== 0)
  85. throw new ExtractionFailedError({
  86. filepath: archivePath,
  87. stderr: await Bun.readableStreamToText(proc.stderr),
  88. })
  89. }
  90. await fs.unlink(archivePath)
  91. if (process.platform !== "win32") await fs.chmod(filepath, 0o755)
  92. }
  93. return {
  94. filepath,
  95. }
  96. })
  97. export async function filepath() {
  98. const { filepath } = await state()
  99. return filepath
  100. }
  101. export async function files(input: {
  102. cwd: string
  103. query?: string
  104. glob?: string
  105. limit?: number
  106. }) {
  107. const commands = [
  108. `${await filepath()} --files --hidden --glob='!.git/*' ${input.glob ? `--glob='${input.glob}'` : ``}`,
  109. ]
  110. if (input.query)
  111. commands.push(`${await Fzf.filepath()} --filter=${input.query}`)
  112. if (input.limit) commands.push(`head -n ${input.limit}`)
  113. const joined = commands.join(" | ")
  114. const result = await $`${{ raw: joined }}`.cwd(input.cwd).nothrow().text()
  115. return result.split("\n").filter(Boolean)
  116. }
  117. export async function tree(input: { cwd: string; limit?: number }) {
  118. const files = await Ripgrep.files({ cwd: input.cwd })
  119. interface Node {
  120. path: string[]
  121. children: Node[]
  122. }
  123. function getPath(node: Node, parts: string[], create: boolean) {
  124. if (parts.length === 0) return node
  125. let current = node
  126. for (const part of parts) {
  127. let existing = current.children.find((x) => x.path.at(-1) === part)
  128. if (!existing) {
  129. if (!create) return
  130. existing = {
  131. path: current.path.concat(part),
  132. children: [],
  133. }
  134. current.children.push(existing)
  135. }
  136. current = existing
  137. }
  138. return current
  139. }
  140. const root: Node = {
  141. path: [],
  142. children: [],
  143. }
  144. for (const file of files) {
  145. const parts = file.split(path.sep)
  146. getPath(root, parts, true)
  147. }
  148. function sort(node: Node) {
  149. node.children.sort((a, b) => {
  150. if (!a.children.length && b.children.length) return 1
  151. if (!b.children.length && a.children.length) return -1
  152. return a.path.at(-1)!.localeCompare(b.path.at(-1)!)
  153. })
  154. for (const child of node.children) {
  155. sort(child)
  156. }
  157. }
  158. sort(root)
  159. let current = [root]
  160. const result: Node = {
  161. path: [],
  162. children: [],
  163. }
  164. let processed = 0
  165. const limit = input.limit ?? 50
  166. while (current.length > 0) {
  167. const next = []
  168. for (const node of current) {
  169. if (node.children.length) next.push(...node.children)
  170. }
  171. const max = Math.max(...current.map((x) => x.children.length))
  172. for (let i = 0; i < max && processed < limit; i++) {
  173. for (const node of current) {
  174. const child = node.children[i]
  175. if (!child) continue
  176. getPath(result, child.path, true)
  177. processed++
  178. if (processed >= limit) break
  179. }
  180. }
  181. if (processed >= limit) {
  182. for (const node of [...current, ...next]) {
  183. const compare = getPath(result, node.path, false)
  184. if (!compare) continue
  185. if (compare?.children.length !== node.children.length) {
  186. const diff = node.children.length - compare.children.length
  187. compare.children.push({
  188. path: compare.path.concat(`[${diff} truncated]`),
  189. children: [],
  190. })
  191. }
  192. }
  193. break
  194. }
  195. current = next
  196. }
  197. const lines: string[] = []
  198. function render(node: Node, depth: number) {
  199. const indent = "\t".repeat(depth)
  200. lines.push(indent + node.path.at(-1) + (node.children.length ? "/" : ""))
  201. for (const child of node.children) {
  202. render(child, depth + 1)
  203. }
  204. }
  205. result.children.map((x) => render(x, 0))
  206. return lines.join("\n")
  207. }
  208. }