open Inshellah.Parser open Inshellah.Manpage open Inshellah.Nushell open Inshellah.Store module SSet = Set.Make(String) let usage () = Printf.eprintf {|inshellah - nushell completions engine Usage: inshellah index PREFIX... [--db PATH] Index completions into a SQLite database. PREFIX is a directory containing bin/ and share/man/. Default db: $XDG_CACHE_HOME/inshellah/completions.db inshellah dump [--db PATH] Show stats and commands in the database. inshellah generate BINDIR MANDIR -o OUTDIR Full generation: native completions, manpages, and --help fallback. One .nu file per command. inshellah manpage FILE Parse a manpage and emit nushell extern inshellah manpage-dir DIR Batch-process manpages under DIR inshellah help [--iterative] CMD [ARGS...] Run CMD ARGS --help, parse and emit extern. Recursively resolves subcommands unless --iterative is given. inshellah parse-help CMD Read --help text from stdin, emit extern inshellah demo Run built-in demo |}; exit 1 let command_sections = [1; 8] let contains_str s sub = try ignore (Str.search_forward (Str.regexp_string sub) s 0); true with Not_found -> false let is_nushell_source text = String.length text > 20 && (contains_str text "export extern" || contains_str text "export def" || (contains_str text "module " && contains_str text "export")) let filename_of_cmd cmd = String.map (function | ('a'..'z' | 'A'..'Z' | '0'..'9' | '-' | '_' | '.') as c -> c | _ -> '-') cmd let write_file path contents = let oc = open_out path in output_string oc contents; close_out oc let cmd_name_of_manpage path = let base = Filename.basename path in let base = if Filename.check_suffix base ".gz" then Filename.chop_suffix base ".gz" else base in try Filename.chop_extension base with Invalid_argument _ -> base let safe_env = lazy ( Array.of_list ( List.filter (fun s -> not (String.starts_with ~prefix:"DISPLAY=" s || String.starts_with ~prefix:"WAYLAND_DISPLAY=" s || String.starts_with ~prefix:"DBUS_SESSION_BUS_ADDRESS=" s || String.starts_with ~prefix:"XAUTHORITY=" s)) (Array.to_list (Unix.environment ())))) let run_cmd args timeout_ms = let (rd, wr) = Unix.pipe () in let devnull = Unix.openfile "/dev/null" [Unix.O_RDONLY] 0 in let argv = Array.of_list args in let pid = try Unix.create_process_env (List.hd args) argv (Lazy.force safe_env) devnull wr wr with Unix.Unix_error _ -> Unix.close rd; Unix.close wr; Unix.close devnull; -1 in Unix.close wr; Unix.close devnull; if pid < 0 then (Unix.close rd; None) else begin let buf = Buffer.create 4096 in let deadline = Unix.gettimeofday () +. (float_of_int timeout_ms /. 1000.0) in let chunk = Bytes.create 8192 in let alive = ref true in (try while !alive do let remaining = deadline -. Unix.gettimeofday () in if remaining <= 0.0 then alive := false else match Unix.select [rd] [] [] (min remaining 0.05) with | (_ :: _, _, _) -> let n = Unix.read rd chunk 0 8192 in if n = 0 then raise Exit else Buffer.add_subbytes buf chunk 0 n | _ -> () done with Exit -> ()); Unix.close rd; if not !alive then begin (try Unix.kill pid Sys.sigkill with Unix.Unix_error _ -> ()); ignore (Unix.waitpid [] pid) end else ignore (Unix.waitpid [] pid); if Buffer.length buf > 0 then Some (Buffer.contents buf) else None end let is_executable path = try let st = Unix.stat path in st.st_kind = Unix.S_REG && st.st_perm land 0o111 <> 0 with Unix.Unix_error _ -> false let is_script path = try let real = Unix.realpath path in let ic = open_in_bin real in let has_shebang = try let b = Bytes.create 2 in really_input ic b 0 2; Bytes.get b 0 = '#' && Bytes.get b 1 = '!' with End_of_file -> false in close_in ic; has_shebang with _ -> false let elf_scan path needles = let found = Hashtbl.create 4 in let remaining () = List.filter (fun n -> not (Hashtbl.mem found n)) needles in (try let real = Unix.realpath path in let ic = open_in_bin real in let magic = Bytes.create 4 in really_input ic magic 0 4; if Bytes.get magic 0 = '\x7f' && Bytes.get magic 1 = 'E' && Bytes.get magic 2 = 'L' && Bytes.get magic 3 = 'F' then begin let max_needle = List.fold_left (fun m n -> max m (String.length n)) 0 needles in let chunk_size = 65536 in let buf = Bytes.create (chunk_size + max_needle) in let carry = ref 0 in let eof = ref false in while not !eof && remaining () <> [] do let n = (try input ic buf !carry chunk_size with End_of_file -> 0) in if n = 0 then eof := true else begin let total = !carry + n in List.iter (fun needle -> if not (Hashtbl.mem found needle) then begin let nlen = String.length needle in let i = ref 0 in while !i <= total - nlen do if Bytes.get buf !i = needle.[0] then begin let ok = ref true in for j = 1 to nlen - 1 do if Bytes.get buf (!i + j) <> needle.[j] then ok := false done; if !ok then (Hashtbl.replace found needle true; i := total) else incr i end else incr i done end ) (remaining ()); let new_carry = min max_needle total in Bytes.blit buf (total - new_carry) buf 0 new_carry; carry := new_carry end done end; close_in ic with _ -> List.iter (fun n -> Hashtbl.replace found n true) needles); found let nix_wrapper_target path = try let real = Unix.realpath path in let ic = open_in_bin real in let n = in_channel_length ic in if n > 65536 then (close_in ic; None) else begin let s = Bytes.create n in really_input ic s 0 n; close_in ic; let s = Bytes.to_string s in if not (contains_str s "makeCWrapper") then None else let re = Str.regexp "/nix/store/[a-z0-9]+-[^' \n\r\x00]+/bin/[a-zA-Z0-9._-]+" in try ignore (Str.search_forward re s 0); let target = Str.matched_string s in if Sys.file_exists target then Some target else None with Not_found -> None end with _ -> None let skip_name name = String.length name = 0 || name = "-" || name.[0] = '.' || String.starts_with ~prefix:"lib" name || String.ends_with ~suffix:"-daemon" name || String.ends_with ~suffix:"-wrapped" name || String.ends_with ~suffix:".so" name || not (String.exists (fun c -> (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9')) name) type bin_class = Skip | Try_help | Try_native_and_help let classify_binary bindir name = if is_nushell_builtin name || skip_name name then Skip else let path = Filename.concat bindir name in if not (is_executable path) then Skip else if is_script path then Try_help else let scan = elf_scan path ["-h"; "completion"] in if Hashtbl.mem scan "completion" then Try_native_and_help else if Hashtbl.mem scan "-h" then Try_help else if nix_wrapper_target path <> None then Try_help else Skip let num_cores () = try let ic = open_in "/proc/cpuinfo" in let n = ref 0 in (try while true do if String.starts_with ~prefix:"processor" (input_line ic) then incr n done with End_of_file -> ()); close_in ic; max 1 !n with _ -> 4 let try_native_completion bin_path = let patterns = [ [bin_path; "completions"; "nushell"]; [bin_path; "completion"; "nushell"]; [bin_path; "--completions"; "nushell"]; [bin_path; "--completion"; "nushell"]; [bin_path; "generate-completion"; "nushell"]; [bin_path; "--generate-completion"; "nushell"]; [bin_path; "shell-completions"; "nushell"]; ] in let rec go = function | [] -> None | args :: rest -> match run_cmd args 500 with | Some text when is_nushell_source text -> Some text | _ -> go rest in go patterns let cmd_manpage file = let contents = read_manpage_file file in let fallback = cmd_name_of_manpage file in let cmd = match extract_synopsis_command contents with | Some name -> name | None -> fallback in if not (is_nushell_builtin cmd) then let result = parse_manpage_string contents in if result.entries <> [] then print_string (generate_extern cmd result) let cmd_manpage_dir dir = List.iter (fun section -> let subdir = Filename.concat dir (Printf.sprintf "man%d" section) in if Sys.file_exists subdir && Sys.is_directory subdir then Array.iter (fun file -> (try cmd_manpage (Filename.concat subdir file) with _ -> ()) ) (Sys.readdir subdir) ) command_sections let max_resolve_results = 500 let help_resolve ?(timeout=10_000) cmd rest name = let max_jobs = num_cores () in let queue = Queue.create () in Queue.push (rest, name, 0, "") queue; let results = ref [] in (* pending: (pid, rd, rest, name, depth, desc) *) let pending = ref [] in let collect rd p_rest p_name p_depth p_desc = let ic = Unix.in_channel_of_descr rd in let[@warning "-8"] result : (string * (string * string) list) option = try Marshal.from_channel ic with _ -> None in close_in ic; match result with | None -> if p_desc <> "" then results := Printf.sprintf "export extern \"%s\" [ # %s\n]\n" (escape_nu p_name) (escape_nu p_desc) :: !results | Some (code, subs) -> results := code :: !results; if p_depth < 5 && List.length !results < max_resolve_results then List.iter (fun (sc_name, sc_desc) -> Queue.push (p_rest @ [sc_name], p_name ^ " " ^ sc_name, p_depth + 1, sc_desc) queue ) subs in let reap () = pending := List.filter (fun (pid, rd, p_rest, p_name, p_depth, p_desc) -> match Unix.waitpid [Unix.WNOHANG] pid with | (0, _) -> true | _ -> collect rd p_rest p_name p_depth p_desc; false | exception Unix.Unix_error (Unix.ECHILD, _, _) -> collect rd p_rest p_name p_depth p_desc; false ) !pending in let wait_slot () = while List.length !pending >= max_jobs do reap (); if List.length !pending >= max_jobs then (try ignore (Unix.wait ()) with Unix.Unix_error _ -> ()); reap () done in while not (Queue.is_empty queue) || !pending <> [] do while not (Queue.is_empty queue) do let (q_rest, q_name, q_depth, q_desc) = Queue.pop queue in wait_slot (); let (rd, wr) = Unix.pipe () in let pid = Unix.fork () in if pid = 0 then begin Unix.close rd; List.iter (fun (_, prd, _, _, _, _) -> try Unix.close prd with _ -> ()) !pending; let result = let text = match run_cmd (cmd :: q_rest @ ["--help"]) timeout with | Some _ as r -> r | None -> run_cmd (cmd :: q_rest @ ["-h"]) timeout in match text with | None -> None | Some text -> (match parse_help text with | Error _ -> None | Ok r when r.entries = [] && r.subcommands = [] && r.positionals = [] -> None | Ok r -> (* If the subcommand we just queried appears in its own subcommand list, the command is echoing the parent help (e.g. nil ignores subcommands when --help is present). Abandon this branch to avoid infinite recursion. *) let self_listed = match q_rest with | [] -> false | _ -> let leaf = List.nth q_rest (List.length q_rest - 1) in List.exists (fun (sc : subcommand) -> sc.name = leaf) r.subcommands in if self_listed then None else let at_limit = q_depth >= 5 in let code = generate_extern q_name (if at_limit then r else { r with subcommands = [] }) in let subs = if at_limit then [] else List.map (fun (sc : subcommand) -> (sc.name, sc.desc)) r.subcommands in Some (code, subs)) in let oc = Unix.out_channel_of_descr wr in Marshal.to_channel oc (result : (string * (string * string) list) option) []; close_out oc; exit 0 end else begin Unix.close wr; pending := (pid, rd, q_rest, q_name, q_depth, q_desc) :: !pending end done; if !pending <> [] then begin reap (); if !pending <> [] && Queue.is_empty queue then begin (try ignore (Unix.wait ()) with Unix.Unix_error _ -> ()); reap () end end done; match !results with | [] -> None | rs -> Some (String.concat "\n" (List.rev rs)) let cmd_help args = let iterative, cmd_args = match args with | "--iterative" :: rest -> (true, rest) | _ -> (false, args) in match cmd_args with | [] -> Printf.eprintf "error: help requires a command name\n"; exit 1 | cmd :: rest -> let name = String.concat " " (Filename.basename cmd :: rest) in if iterative then (match run_cmd (cmd :: rest @ ["--help"]) 10_000 with | None -> Printf.eprintf "no output from %s --help\n" name; exit 1 | Some text -> (match parse_help text with | Ok r -> print_string (generate_extern name r) | Error msg -> Printf.eprintf "parse error for %s: %s\n" name msg; exit 1)) else (match help_resolve cmd rest name with | None -> Printf.eprintf "no output from %s --help\n" name; exit 1 | Some output -> print_string output) let cmd_parse_help cmd = let buf = Buffer.create 4096 in (try while true do Buffer.add_string buf (input_line stdin); Buffer.add_char buf '\n' done with End_of_file -> ()); (match parse_help (Buffer.contents buf) with | Ok r -> print_string (generate_extern cmd r) | Error msg -> Printf.eprintf "parse error for %s: %s\n" cmd msg; exit 1) let process_manpage file = try let contents = read_manpage_file file in let fallback = cmd_name_of_manpage file in let cmd = match extract_synopsis_command contents with | Some name -> name | None -> fallback in if is_nushell_builtin cmd then None else let result = parse_manpage_string contents in if result.entries <> [] then Some (cmd, result) else None with _ -> None let manpaged_commands mandir = List.fold_left (fun acc section -> let subdir = Filename.concat mandir (Printf.sprintf "man%d" section) in if Sys.file_exists subdir && Sys.is_directory subdir then Array.fold_left (fun acc f -> SSet.add (cmd_name_of_manpage f) acc) acc (Sys.readdir subdir) else acc ) SSet.empty command_sections let cmd_generate bindir mandir outdir ignorelist = let done_cmds = ref SSet.empty in let bins = Sys.readdir bindir in Array.sort String.compare bins; let manpaged = manpaged_commands mandir in let max_jobs = num_cores () in let classified = Array.map (fun name -> if SSet.mem name manpaged || SSet.mem name ignorelist then (name, Skip) else (name, classify_binary bindir name) ) bins in let pending = ref [] in let reap () = pending := List.filter (fun pid -> match Unix.waitpid [Unix.WNOHANG] pid with | (0, _) -> true | _ -> false | exception Unix.Unix_error (Unix.ECHILD, _, _) -> false ) !pending in let wait_slot () = while List.length !pending >= max_jobs do reap (); if List.length !pending >= max_jobs then (try ignore (Unix.wait ()) with Unix.Unix_error _ -> ()); reap () done in Array.iter (fun (name, cls) -> match cls with | Skip -> () | Try_help | Try_native_and_help -> wait_slot (); let pid = Unix.fork () in if pid = 0 then begin (try let path = Filename.concat bindir name in let native_ok = match cls with | Try_native_and_help -> (match try_native_completion path with | Some src -> write_file (Filename.concat outdir (filename_of_cmd name ^ ".nu")) src; true | None -> false) | _ -> false in if not native_ok then begin match help_resolve ~timeout:200 path [] name with | Some content when String.length content > 0 -> let m = module_name_of name in let src = Printf.sprintf "module %s {\n%s}\n\nuse %s *\n" m content m in write_file (Filename.concat outdir (filename_of_cmd name ^ ".nu")) src | _ -> () end; exit 0 with _ -> exit 1) end else begin pending := pid :: !pending; done_cmds := SSet.add name !done_cmds end ) classified; while !pending <> [] do (try ignore (Unix.wait ()) with Unix.Unix_error _ -> ()); reap () done; List.iter (fun section -> let subdir = Filename.concat mandir (Printf.sprintf "man%d" section) in if Sys.file_exists subdir && Sys.is_directory subdir then begin let files = Sys.readdir subdir in Array.sort String.compare files; Array.iter (fun file -> match process_manpage (Filename.concat subdir file) with | None -> () | Some (cmd, result) -> let base = List.hd (String.split_on_char ' ' cmd) in if SSet.mem cmd !done_cmds then () else begin done_cmds := SSet.add cmd !done_cmds; let outpath = Filename.concat outdir (filename_of_cmd base ^ ".nu") in if Sys.file_exists outpath then begin let existing = let ic = open_in outpath in let n = in_channel_length ic in let s = Bytes.create n in really_input ic s 0 n; close_in ic; Bytes.to_string s in let mod_name = module_name_of base in let use_line = Printf.sprintf "\nuse %s *\n" mod_name in let base_content = if contains_str existing use_line then String.sub existing 0 (Str.search_forward (Str.regexp_string use_line) existing 0) else existing in write_file outpath (String.concat "" [base_content; generate_extern cmd result; use_line]) end else write_file outpath (generate_module base result) end ) files end ) command_sections (* Sequential help resolver for use inside forked children. No forking — just iterates through subcommands with run_cmd directly. *) let help_resolve_seq ?(timeout=200) cmd rest name = let queue = Queue.create () in Queue.push (rest, name, 0) queue; let results = ref [] in while not (Queue.is_empty queue) do let (q_rest, q_name, q_depth) = Queue.pop queue in let text = match run_cmd (cmd :: q_rest @ ["--help"]) timeout with | Some _ as r -> r | None -> run_cmd (cmd :: q_rest @ ["-h"]) timeout in match text with | None -> () | Some text -> (match parse_help text with | Error _ -> () | Ok r when r.entries = [] && r.subcommands = [] && r.positionals = [] -> () | Ok r -> let self_listed = match q_rest with | [] -> false | _ -> let leaf = List.nth q_rest (List.length q_rest - 1) in List.exists (fun (sc : subcommand) -> sc.name = leaf) r.subcommands in if not self_listed then begin let at_limit = q_depth >= 5 || List.length !results >= max_resolve_results in results := (q_name, { r with subcommands = [] }) :: !results; if not at_limit then List.iter (fun (sc : subcommand) -> Queue.push (q_rest @ [sc.name], q_name ^ " " ^ sc.name, q_depth + 1) queue ) r.subcommands end) done; List.rev !results (* Index: mirrors cmd_generate's fork-per-binary pattern. Each child handles one binary completely (including subcommand resolution) and marshals all results back in one shot. No nested forking — children use help_resolve_seq which is purely sequential. *) let cmd_index bindirs mandirs ignorelist db_path = let db = init db_path in begin_transaction db; let done_cmds = ref SSet.empty in let n_results = ref 0 in let index_bindir bindir mandir = if not (Sys.file_exists bindir && Sys.is_directory bindir) then Printf.eprintf "skipping %s (not found)\n" bindir else begin let bins = Sys.readdir bindir in Array.sort String.compare bins; let manpaged = if Sys.file_exists mandir && Sys.is_directory mandir then manpaged_commands mandir else SSet.empty in let max_jobs = num_cores () in let classified = Array.map (fun name -> if SSet.mem name manpaged || SSet.mem name ignorelist then (name, Skip) else (name, classify_binary bindir name) ) bins in let pending = ref [] in let reap () = pending := List.filter (fun (pid, rd, name) -> match Unix.waitpid [Unix.WNOHANG] pid with | (0, _) -> true | _ -> let ic = Unix.in_channel_of_descr rd in (let[@warning "-8"] result : [`Native of string | `Parsed of (string * help_result) list | `None] = try Marshal.from_channel ic with _ -> `None in match result with | `Native src -> upsert_raw db ~source:"native" name src; incr n_results | `Parsed pairs -> List.iter (fun (cmd_name, r) -> if not (SSet.mem cmd_name !done_cmds) then begin upsert db ~source:"help" cmd_name r; done_cmds := SSet.add cmd_name !done_cmds; incr n_results end ) pairs | `None -> ()); close_in ic; done_cmds := SSet.add name !done_cmds; false | exception Unix.Unix_error (Unix.ECHILD, _, _) -> (try Unix.close rd with _ -> ()); false ) !pending in let wait_slot () = while List.length !pending >= max_jobs do reap (); if List.length !pending >= max_jobs then (try ignore (Unix.wait ()) with Unix.Unix_error _ -> ()); reap () done in Array.iter (fun (name, cls) -> match cls with | Skip -> () | Try_help | Try_native_and_help -> wait_slot (); let (rd, wr) = Unix.pipe () in let pid = Unix.fork () in if pid = 0 then begin Unix.close rd; List.iter (fun (_, prd, _) -> try Unix.close prd with _ -> ()) !pending; let result = try let path = Filename.concat bindir name in let native = match cls with | Try_native_and_help -> (match try_native_completion path with | Some src -> Some src | None -> None) | _ -> None in match native with | Some src -> `Native src | None -> let pairs = help_resolve_seq ~timeout:200 path [] name in if pairs <> [] then `Parsed pairs else `None with _ -> `None in let oc = Unix.out_channel_of_descr wr in Marshal.to_channel oc (result : [`Native of string | `Parsed of (string * help_result) list | `None]) []; close_out oc; exit 0 end else begin Unix.close wr; pending := (pid, rd, name) :: !pending end ) classified; while !pending <> [] do (try ignore (Unix.wait ()) with Unix.Unix_error _ -> ()); reap () done; (* Phase 2: manpages *) if Sys.file_exists mandir && Sys.is_directory mandir then List.iter (fun section -> let subdir = Filename.concat mandir (Printf.sprintf "man%d" section) in if Sys.file_exists subdir && Sys.is_directory subdir then begin let files = Sys.readdir subdir in Array.sort String.compare files; Array.iter (fun file -> match process_manpage (Filename.concat subdir file) with | None -> () | Some (cmd, result) -> if not (SSet.mem cmd !done_cmds) then begin upsert db ~source:"manpage" cmd result; done_cmds := SSet.add cmd !done_cmds; incr n_results end ) files end ) command_sections end in List.iter2 index_bindir bindirs mandirs; commit db; Printf.printf "indexed %d commands into %s\n" !n_results db_path; close db let cmd_dump db_path = let db = init db_path in let (count, sources) = stats db in Printf.printf "database: %s\n" db_path; Printf.printf "commands: %d (from %d sources)\n" count sources; let cmds = all_commands db in List.iter (fun cmd -> match lookup db cmd with | None -> () | Some (_data, source) -> Printf.printf " %-40s [%s]\n" cmd source ) cmds; close db let cmd_demo () = Printf.printf "# Generated by: inshellah demo\n\n"; match parse_help {|Usage: ls [OPTION]... [FILE]... -a, --all do not ignore entries starting with . -A, --almost-all do not list implied . and .. --block-size=SIZE with -l, scale sizes by SIZE when printing --color[=WHEN] color the output WHEN -h, --human-readable with -l and -s, print sizes like 1K 234M 2G etc. --help display this help and exit --version output version information and exit |} with | Ok r -> print_string (generate_extern "ls" r) | Error msg -> Printf.eprintf "parse error: %s\n" msg let load_ignorelist path = try let ic = open_in path in let lines = ref [] in (try while true do let line = String.trim (input_line ic) in if String.length line > 0 && line.[0] <> '#' then lines := line :: !lines done with End_of_file -> ()); close_in ic; SSet.of_list !lines with _ -> SSet.empty let parse_index_args args = let rec go prefixes db ignore = function | [] -> (List.rev prefixes, db, ignore) | "--db" :: path :: rest -> go prefixes path ignore rest | "--ignore" :: path :: rest -> go prefixes db (SSet.union ignore (load_ignorelist path)) rest | dir :: rest -> go (dir :: prefixes) db ignore rest in go [] (default_db_path ()) SSet.empty args let () = match Array.to_list Sys.argv |> List.tl with | ["generate"; bindir; mandir; "-o"; outdir] -> cmd_generate bindir mandir outdir SSet.empty | ["generate"; bindir; mandir; "-o"; outdir; "--ignore"; ignore_file] -> cmd_generate bindir mandir outdir (load_ignorelist ignore_file) | "index" :: rest -> let (prefixes, db_path, ignorelist) = parse_index_args rest in if prefixes = [] then (Printf.eprintf "error: index requires at least one prefix dir\n"; exit 1); let bindirs = List.map (fun p -> Filename.concat p "bin") prefixes in let mandirs = List.map (fun p -> Filename.concat p "share/man") prefixes in cmd_index bindirs mandirs ignorelist db_path | "dump" :: rest -> let db_path = match rest with | ["--db"; path] -> path | [] -> default_db_path () | _ -> Printf.eprintf "error: dump [--db PATH]\n"; exit 1 in cmd_dump db_path | ["manpage"; file] -> cmd_manpage file | ["manpage-dir"; dir] -> cmd_manpage_dir dir | "help" :: rest -> cmd_help rest | ["parse-help"; cmd] -> cmd_parse_help cmd | ["demo"] -> cmd_demo () | _ -> usage ()