fzf_files.ml

#!/usr/bin/env -S mach run --
#require "unix"
#require "cmdliner"

open Printf

type command = {
  mode: string;
  command: string;
  nth: int option;
  accept_nth: int option;
}

let command ?nth ?accept_nth ~command mode =
  {mode; command; nth; accept_nth}

let git_status =
  "git -c color.ui=always status --short --untracked-files=no | grep -v '#'"

(* The git mode lists the current status. When a base ref is given, it also
   lists the files introduced by the current branch relative to the base
   (the three-dot diff against the merge base, so changes made only in the
   base are excluded), deduplicating against the status by path. *)
let git_command ?git_base () =
  match git_base with
  | None -> git_status
  | Some base ->
    sprintf
      "{ %s; git -c color.ui=always diff --name-status %s...HEAD; } | \
       awk '{ k=$NF; gsub(/\\033\\[[0-9;]*m/,\"\",k); if (!seen[k]++) print }'"
      git_status base

let commands ?git_base () =
  [
    command "all"
      ~command:"rg --files --hidden --glob='!.git' --color=never";
    command "git"
      ~nth:2 ~accept_nth:2
      ~command:(git_command ?git_base ());
  ]

let bold s = sprintf "\027[1m%s\027[0m" s
let inverse s = sprintf "\027[7m%s\027[0m" s
let dim s = sprintf "\027[2m%s\027[0m" s

let header commands current =
  let modes =
    List.map
      (fun {mode;_} ->
        if mode = current then bold mode else mode)
      commands
    |> String.concat " "
  in
  sprintf "%s %s" modes (dim "ctrl-o: toggle mode")

let write_file path content =
  Out_channel.with_open_bin path (fun oc -> output_string oc content)

let temp_file ?chmod ~content suffix =
  let path = Filename.temp_file "fzf" suffix in
  at_exit (fun () -> Sys.remove path);
  write_file path content;
  Option.iter (Unix.chmod path) chmod;
  path

let actionf name fmt = ksprintf (sprintf "%s(%s)" name) fmt
let many actions = String.concat "+" actions
let reload_sync fmt = actionf "reload-sync" fmt
let execute_silent fmt = actionf "execute-silent" fmt
let change_nth fmt = actionf "change-nth" fmt
let change_header fmt = actionf "change-header" fmt
let become fmt = actionf "become" fmt

let read_file path =
  In_channel.with_open_bin path In_channel.input_all

let find_next_command commands current_mode =
  let n = List.length commands in
  let rec find i = function
    | [] -> List.hd commands
    | {mode; _} :: rest ->
      if mode = current_mode then List.nth commands ((i + 1) mod n)
      else find (i + 1) rest
  in
  find 0 commands

let find_command commands mode =
  List.find_opt (fun c -> c.mode = mode) commands

let find_command_or_exit commands mode =
  match find_command commands mode with
  | Some cmd -> cmd
  | None -> eprintf "unknown mode: %s\n%!" mode; exit 1

let run_toggle git_base (mode_filename, mode) =
  let commands = commands ?git_base () in
  let next = find_next_command commands mode in
  let actions = [
    reload_sync "%s" next.command;
    execute_silent "echo '%s' > %s" next.mode mode_filename;
    change_header "%s" (header commands next.mode);
  ] in
  let actions =
    match next.nth with
    | Some nth -> actions @ [change_nth "%d" nth]
    | None -> actions @ [change_nth ""]
  in
  print_endline (many actions);
  0

let run_choose git_base (_mode_filename, mode) =
  let commands = commands ?git_base () in
  let cmd = find_command_or_exit commands mode in
  print_endline (
    match cmd.accept_nth with
    | Some nth -> many [become "echo '{%d}'" nth; "abort"]
    | None -> many [become "echo '{}'"; "abort"]
  );
  0

let run_fzf git_base mode =
  let commands = commands ?git_base () in
  let default = match mode with
    | Some m -> Option.value (find_command commands m) ~default:(List.hd commands)
    | None -> List.hd commands
  in
  let mode_state = temp_file "mode" ~content:default.mode in
  let nth_args =
    match default.nth with
    | Some nth -> [| "--nth"; string_of_int nth |]
    | None -> [||]
  in
  (* git base is baked into each mode's command, but the internal _toggle and
     _choose invocations rebuild the command list, so they need it too. *)
  let base_arg =
    match git_base with
    | Some base -> sprintf " --git-base=%s" base
    | None -> ""
  in
  let pid =
    Unix.create_process "fzf"
      (Array.concat
         [
           [|
             "fzf";
             "--ansi";
             "--header-first";
             "--prompt"; "> ";
             "--header"; header commands default.mode;
             "--bind"; sprintf "start:reload-sync(%s)" default.command;
             "--bind"; sprintf "enter:transform:%s _choose %s%s" Sys.executable_name mode_state base_arg;
             "--bind"; sprintf "ctrl-o:transform:%s _toggle %s%s" Sys.executable_name mode_state base_arg;
           |];
           nth_args;
         ])
      Unix.stdin Unix.stdout Unix.stderr
  in
  match Unix.waitpid [] pid with
  | _, Unix.WEXITED code -> code
  | _, Unix.WSIGNALED _ -> 1
  | _, Unix.WSTOPPED _ -> 1

open Cmdliner

let mode_state =
  Term.(
    const (fun filename -> filename, String.trim (read_file filename))
    $ Arg.(required & pos 0 (some file) None & info [] ~docv:"MODE_STATE"))

let git_base =
  let doc =
    "Git base ref (commit, branch, ...). The git mode then also lists files \
     introduced by the current branch relative to $(docv) (three-dot diff), \
     in addition to the current status."
  in
  Arg.(value & opt (some string) None & info ["git-base"] ~doc ~docv:"REF")

let choose_cmd =
  let doc = "choose an item (internal command)" in
  let info = Cmd.info "_choose" ~doc in
  Cmd.v info Term.(const run_choose $ git_base $ mode_state)

let toggle_cmd =
  let doc = "toggle mode (internal command)" in
  let info = Cmd.info "_toggle" ~doc in
  Cmd.v info Term.(const run_toggle $ git_base $ mode_state)

let default_cmd =
  let doc = "fzf file picker, using preconfigured sources" in
  let info = Cmd.info "pick" ~doc in
  let mode =
    let doc = "Starting mode" in
    Arg.(value & opt (some string) None & info ["mode"] ~doc ~docv:"MODE")
  in
  Cmd.group ~default:Term.(const run_fzf $ git_base $ mode) info [toggle_cmd; choose_cmd]

let () = exit (Cmd.eval' default_cmd)