Source file quill_top.ml

(*---------------------------------------------------------------------------
  Copyright (c) 2026 The Raven authors. All rights reserved.
  SPDX-License-Identifier: ISC
  ---------------------------------------------------------------------------*)

(* ───── Toplevel primitives ───── *)

let findlib_predicates = ref [ "byte"; "toploop" ]

let ensure_findlib () =
  match Findlib.init () with () -> true | exception _ -> false

(* Mark packages that are already linked into the executable. Their .cmi files
   need to be on the search path, but we must not try to load their .cma
   again. *)
let add_packages pkgs =
  if ensure_findlib () then
    List.iter
      (fun pkg ->
        (match Findlib.package_directory pkg with
        | dir -> Topdirs.dir_directory dir
        | exception _ -> ());
        if not (Findlib.is_recorded_package pkg) then
          Findlib.record_package Findlib.Record_core pkg)
      pkgs
  else
    (* Findlib unavailable — fall back to OCAMLPATH. *)
    let sep = if Sys.win32 then ';' else ':' in
    match Sys.getenv_opt "OCAMLPATH" with
    | None -> ()
    | Some ocamlpath ->
        let roots = String.split_on_char sep ocamlpath in
        List.iter
          (fun pkg ->
            let subdir =
              String.concat Filename.dir_sep (String.split_on_char '.' pkg)
            in
            List.iter
              (fun root ->
                let dir = Filename.concat root subdir in
                if Sys.file_exists dir && Sys.is_directory dir then
                  Topdirs.dir_directory dir)
              roots)
          pkgs

(* Try loading a single ancestor package. Returns [true] on success, [false] if
   the archive references an undefined global (dependency not yet loaded). *)
let try_load_ancestor p =
  let loaded =
    Findlib.is_recorded_package p
    && Findlib.type_of_recorded_package p = Findlib.Record_load
  in
  if loaded then true
  else
    let incore =
      Findlib.is_recorded_package p
      && Findlib.type_of_recorded_package p = Findlib.Record_core
    in
    let d = Findlib.package_directory p in
    Topdirs.dir_directory d;
    if incore then begin
      Findlib.record_package Findlib.Record_load p;
      true
    end
    else
      let archive =
        try Findlib.package_property !findlib_predicates p "archive"
        with Not_found -> ""
      in
      let archives =
        String.split_on_char ' ' archive |> List.filter (fun s -> s <> "")
      in
      try
        List.iter
          (fun arch ->
            let path = Findlib.resolve_path ~base:d arch in
            Topdirs.dir_load Format.err_formatter path)
          archives;
        Findlib.record_package Findlib.Record_load p;
        true
      with Symtable.Error (Undefined_global _) -> false

(* Load a package: resolve its dependency chain and load archives.

   Findlib's topological sort does not account for virtual library
   implementations (a virtual package has no archive; its implementation archive
   may appear later in the ancestor list). This causes [Undefined_global] when a
   dependent is loaded before the implementation.

   We handle this with a fixpoint loop: load what we can, collect failures, and
   retry until either everything succeeds or no progress is made. *)
let load_package pkg =
  if not (ensure_findlib ()) then
    Printf.eprintf "[quill] #require: findlib unavailable\n%!"
  else
    let ancestors =
      Findlib.package_deep_ancestors !findlib_predicates [ pkg ]
    in
    let rec loop remaining =
      let deferred =
        List.filter (fun p -> not (try_load_ancestor p)) remaining
      in
      match deferred with
      | [] -> ()
      | _ when List.length deferred < List.length remaining -> loop deferred
      | _ ->
          (* No progress — report the packages we cannot load *)
          List.iter
            (fun p -> Printf.eprintf "[quill] failed to load package %s\n%!" p)
            deferred
    in
    loop ancestors

(* ───── Initialization ───── *)

let initialized = ref false
let init_mutex = Mutex.create ()

let initialize_if_needed () =
  Mutex.lock init_mutex;
  Fun.protect
    ~finally:(fun () -> Mutex.unlock init_mutex)
    (fun () ->
      if not !initialized then (
        Sys.interactive := false;
        Topeval.init ();
        Toploop.initialize_toplevel_env ();
        Toploop.input_name := "//toplevel//";
        (* Register #require directive for loading packages at runtime. *)
        Toploop.add_directive "require"
          (Directive_string (fun pkg -> load_package pkg))
          { section = "Loading code"; doc = "Load a findlib package" };
        Sys.interactive := true;
        initialized := true))

let install_printer name =
  try
    let phrase =
      Printf.sprintf "#install_printer %s;;" name
      |> Lexing.from_string
      |> !Toploop.parse_toplevel_phrase
    in
    ignore (Toploop.execute_phrase false Format.err_formatter phrase)
  with _ -> ()

let install_printer_fn ~ty f =
  try
    let parts = String.split_on_char '.' ty in
    match Longident.unflatten parts with
    | None -> ()
    | Some lid ->
        let path, _decl = Env.find_type_by_name lid !Toploop.toplevel_env in
        let ty_expr = Ctype.newconstr path [] in
        let printer_path = Path.Pident (Ident.create_local ty) in
        Toploop.install_printer printer_path ty_expr f
  with _ -> ()

(* ───── Output capture ───── *)

(** Pre-allocated read buffer for the poll thread. Avoids major heap allocations
    (4096 > minor heap max) that could trigger GC while the execute thread is
    inside Nx C code. *)
let poll_buf = Bytes.create 4096

(** [read_available fd buf] reads whatever bytes are currently available on [fd]
    into [buf] without blocking indefinitely (the caller uses [Unix.select]
    first). Returns [None] on EOF. *)
let read_available fd buf =
  match Unix.read fd buf 0 (Bytes.length buf) with
  | 0 -> None
  | n -> Some (Bytes.sub_string buf 0 n)
  | exception Unix.Unix_error (Unix.EAGAIN, _, _) -> Some ""

(** [drain_remaining fd] reads all remaining bytes after the write end is
    closed. *)
let drain_remaining fd =
  let buf = Buffer.create 256 in
  let tmp = Bytes.create 4096 in
  let rec loop () =
    match Unix.read fd tmp 0 4096 with
    | 0 -> ()
    | n ->
        Buffer.add_subbytes buf tmp 0 n;
        loop ()
  in
  loop ();
  Unix.close fd;
  Buffer.contents buf

let capture ~on_stdout ~on_stderr ~on_display f =
  let buf_out = Buffer.create 256 in
  let buf_err = Buffer.create 256 in
  let ppf_out = Format.formatter_of_buffer buf_out in
  let ppf_err = Format.formatter_of_buffer buf_err in
  (* Intercept Display_tag semantic tags on the toplevel formatter *)
  Format.pp_set_print_tags ppf_out true;
  Format.pp_set_formatter_stag_functions ppf_out
    {
      mark_open_stag = (fun _ -> "");
      mark_close_stag = (fun _ -> "");
      print_open_stag =
        (fun stag ->
          match stag with
          | Quill.Cell.Display_tag { mime; data } ->
              on_display (Quill.Cell.Display { mime; data })
          | _ -> ());
      print_close_stag = (fun _ -> ());
    };
  (* Pipes for raw stdout/stderr from user code (e.g. print_string) *)
  let rd_out, wr_out = Unix.pipe ~cloexec:true () in
  let rd_err, wr_err = Unix.pipe ~cloexec:true () in
  let stdout_backup = Unix.dup ~cloexec:true Unix.stdout in
  let stderr_backup = Unix.dup ~cloexec:true Unix.stderr in
  (* Poll pipes in a background thread, streaming output as it arrives. Uses
     Unix.select with a 50ms timeout so training progress prints (Printf.printf
     "\rstep %d loss: %.4f%!" ...) appear in real time. *)
  let stop = Atomic.make false in
  let poll_thread =
    Thread.create
      (fun () ->
        while not (Atomic.get stop) do
          let ready, _, _ =
            try Unix.select [ rd_out; rd_err ] [] [] 0.05
            with Unix.Unix_error (Unix.EINTR, _, _) -> ([], [], [])
          in
          List.iter
            (fun fd ->
              match read_available fd poll_buf with
              | Some s when s <> "" ->
                  if fd == rd_out then on_stdout s else on_stderr s
              | _ -> ())
            ready
        done)
      ()
  in
  let result = ref None in
  Fun.protect
    (fun () ->
      flush stdout;
      flush stderr;
      Unix.dup2 ~cloexec:false wr_out Unix.stdout;
      Unix.dup2 ~cloexec:false wr_err Unix.stderr;
      result := Some (f ppf_out ppf_err))
    ~finally:(fun () ->
      Format.pp_print_flush ppf_out ();
      Format.pp_print_flush ppf_err ();
      flush stdout;
      flush stderr;
      Unix.dup2 ~cloexec:false stdout_backup Unix.stdout;
      Unix.dup2 ~cloexec:false stderr_backup Unix.stderr;
      Unix.close stdout_backup;
      Unix.close stderr_backup;
      (* Close write ends so poll thread and drain see EOF *)
      Unix.close wr_out;
      Unix.close wr_err);
  (* Stop the poll thread and drain any remaining bytes *)
  Atomic.set stop true;
  Thread.join poll_thread;
  let rest_out = drain_remaining rd_out in
  let rest_err = drain_remaining rd_err in
  if rest_out <> "" then on_stdout rest_out;
  if rest_err <> "" then on_stderr rest_err;
  (* Format buffer output (toplevel results like "val x = ...") *)
  let toplevel_out = Buffer.contents buf_out in
  let toplevel_err = Buffer.contents buf_err in
  match !result with
  | None -> failwith "capture: unreachable"
  | Some ok -> (ok, toplevel_out, toplevel_err)

(* ───── Execution ───── *)

let ensure_terminator code =
  let trimmed = String.trim code in
  if trimmed = "" || String.ends_with ~suffix:";;" trimmed then code
  else code ^ ";;"

let execute_code ppf_out ppf_err code =
  let code = ensure_terminator code in
  let lb = Lexing.from_string code in
  lb.lex_curr_p <-
    { pos_fname = "//toplevel//"; pos_lnum = 1; pos_bol = 0; pos_cnum = 0 };
  let old_warnings_fmt = !Location.formatter_for_warnings in
  Location.formatter_for_warnings := ppf_err;
  let orig_input_lexbuf = !Location.input_lexbuf in
  Location.input_lexbuf := Some lb;
  let phrases = ref [] in
  let parse_ok =
    try
      while true do
        let phr = !Toploop.parse_toplevel_phrase lb in
        phrases := phr :: !phrases
      done;
      assert false
    with
    | End_of_file -> true
    | e ->
        Location.report_exception ppf_err e;
        false
  in
  let phrases = List.rev !phrases in
  let num_phrases = List.length phrases in
  let success = ref parse_ok in
  Fun.protect
    (fun () ->
      List.iteri
        (fun i phr ->
          try
            let is_last = i = num_phrases - 1 in
            let ok = Toploop.execute_phrase is_last ppf_out phr in
            success := !success && ok
          with
          | Sys.Break ->
              success := false;
              Format.fprintf ppf_err "Interrupted.@."
          | x ->
              success := false;
              Errors.report_error ppf_err x)
        phrases)
    ~finally:(fun () ->
      Location.formatter_for_warnings := old_warnings_fmt;
      Location.input_lexbuf := orig_input_lexbuf;
      Format.pp_print_flush ppf_out ();
      Format.pp_print_flush ppf_err ());
  !success

(* ───── Completion ───── *)

let clamp lo hi x = if x < lo then lo else if x > hi then hi else x

let starts_with ~prefix s =
  let lp = String.length prefix and ls = String.length s in
  lp <= ls && String.sub s 0 lp = prefix

let is_ident_char = function
  | 'a' .. 'z' | 'A' .. 'Z' | '0' .. '9' | '_' | '\'' -> true
  | _ -> false

let is_path_char c = is_ident_char c || Char.equal c '.'

let parse_completion_context code pos =
  let len = String.length code in
  let pos = clamp 0 len pos in
  let i = ref (pos - 1) in
  while !i >= 0 && is_path_char code.[!i] do
    decr i
  done;
  let start = !i + 1 in
  let token = if pos > start then String.sub code start (pos - start) else "" in
  let token =
    if String.starts_with ~prefix:"." token then
      String.sub token 1 (String.length token - 1)
    else token
  in
  if token = "" then (None, "")
  else
    let trailing_dot = String.ends_with ~suffix:"." token in
    let parts = String.split_on_char '.' token |> List.filter (( <> ) "") in
    if trailing_dot then (Longident.unflatten parts, "")
    else
      match List.rev parts with
      | [] -> (None, "")
      | prefix :: rev_qual ->
          let qualifier = Longident.unflatten (List.rev rev_qual) in
          (qualifier, prefix)

let format_type env ty =
  Printtyp.wrap_printing_env ~error:false env (fun () ->
      Format.asprintf "%a" Printtyp.type_scheme ty)

let collect_env_items env qualifier =
  let open Quill.Kernel in
  let add label kind detail acc =
    if String.length label = 0 then acc else { label; kind; detail } :: acc
  in
  let items =
    Env.fold_values
      (fun name _path (vd : Types.value_description) acc ->
        add name Value (format_type env vd.val_type) acc)
      qualifier env []
  in
  let items =
    Env.fold_types
      (fun name _path (td : Types.type_declaration) acc ->
        let detail =
          match td.type_manifest with
          | Some ty -> "= " ^ format_type env ty
          | None -> (
              match td.type_kind with
              | Type_abstract _ -> "abstract"
              | Type_record _ -> "record"
              | Type_variant _ -> "variant"
              | Type_open -> "open")
        in
        add name Type detail acc)
      qualifier env items
  in
  let items =
    Env.fold_modules
      (fun name _path (_md : Types.module_declaration) acc ->
        add name Module "module" acc)
      qualifier env items
  in
  let items =
    Env.fold_modtypes
      (fun name _path (_mtd : Types.modtype_declaration) acc ->
        add name Module_type "module type" acc)
      qualifier env items
  in
  let items =
    Env.fold_constructors
      (fun (c : Data_types.constructor_description) acc ->
        let detail = format_type env c.cstr_res in
        add c.cstr_name Constructor detail acc)
      qualifier env items
  in
  Env.fold_labels
    (fun (l : Data_types.label_description) acc ->
      let detail = format_type env l.lbl_arg in
      add l.lbl_name Label detail acc)
    qualifier env items

let complete_names ~code ~pos =
  let qualifier, prefix = parse_completion_context code pos in
  let env = !Toploop.toplevel_env in
  collect_env_items env qualifier
  |> List.filter (fun (item : Quill.Kernel.completion_item) ->
      String.length prefix = 0 || starts_with ~prefix item.label)
  |> List.sort_uniq (fun (a : Quill.Kernel.completion_item) b ->
      String.compare a.label b.label)

(* ───── Parse and typecheck ───── *)

let parse_phrases code =
  let code = ensure_terminator code in
  let lb = Lexing.from_string code in
  lb.lex_curr_p <-
    { pos_fname = "//toplevel//"; pos_lnum = 1; pos_bol = 0; pos_cnum = 0 };
  let phrases = ref [] in
  (try
     while true do
       let phr = !Toploop.parse_toplevel_phrase lb in
       phrases := phr :: !phrases
     done
   with End_of_file -> ());
  List.rev !phrases

let typecheck_structure env structure =
  let tstr, _sig, _names, _shape, _env =
    Typemod.type_toplevel_phrase env structure
  in
  tstr

(* ───── Type at position ───── *)

let loc_contains (loc : Location.t) pos =
  (not loc.loc_ghost)
  && loc.loc_start.pos_cnum <= pos
  && pos <= loc.loc_end.pos_cnum

let loc_span (loc : Location.t) = loc.loc_end.pos_cnum - loc.loc_start.pos_cnum

let find_type_at_pos env (tstr : Typedtree.structure) pos =
  let best = ref None in
  let update loc ty =
    if loc_contains loc pos then
      match !best with
      | Some (_, prev_loc, _) when loc_span loc >= loc_span prev_loc -> ()
      | _ ->
          let typ = format_type env ty in
          best := Some (typ, loc, None)
  in
  let iter =
    {
      Tast_iterator.default_iterator with
      expr =
        (fun self (e : Typedtree.expression) ->
          update e.exp_loc e.exp_type;
          Tast_iterator.default_iterator.expr self e);
      pat =
        (fun (type k) self (p : k Typedtree.general_pattern) ->
          update p.pat_loc p.pat_type;
          Tast_iterator.default_iterator.pat self p);
    }
  in
  iter.structure iter tstr;
  match !best with
  | None -> None
  | Some (typ, loc, doc) ->
      Some
        Quill.Kernel.
          {
            typ;
            doc;
            from_pos = loc.loc_start.pos_cnum;
            to_pos = loc.loc_end.pos_cnum;
          }

let type_at_pos ~code ~pos =
  let env = !Toploop.toplevel_env in
  let phrases = parse_phrases code in
  let rec try_phrases = function
    | [] -> None
    | Parsetree.Ptop_def structure :: rest -> (
        match typecheck_structure env structure with
        | tstr -> (
            match find_type_at_pos env tstr pos with
            | Some _ as result -> result
            | None -> try_phrases rest)
        | exception _ -> try_phrases rest)
    | _ :: rest -> try_phrases rest
  in
  try_phrases phrases

(* ───── Diagnostics ───── *)

let loc_to_positions (loc : Location.t) =
  (loc.loc_start.pos_cnum, loc.loc_end.pos_cnum)

let error_loc_of_exn exn =
  match exn with
  | Location.Error report -> report.main.loc
  | _ -> Location.in_file "//toplevel//"

let format_exn exn =
  match Location.error_of_exn exn with
  | Some (`Ok report) -> Format.asprintf "%a" Location.print_report report
  | _ -> Printexc.to_string exn

let compute_diagnostics ~code =
  let diags = ref [] in
  let len = String.length code in
  let add_diag severity loc message =
    let from_pos, to_pos = loc_to_positions loc in
    (* Clamp to valid range; skip diagnostics with no usable location *)
    let from_pos = clamp 0 len from_pos in
    let to_pos = clamp 0 len to_pos in
    let to_pos =
      if to_pos <= from_pos then min (from_pos + 1) len else to_pos
    in
    if from_pos < len then
      diags := Quill.Kernel.{ from_pos; to_pos; severity; message } :: !diags
  in
  (match parse_phrases code with
  | _ -> ()
  | exception exn -> add_diag Error (error_loc_of_exn exn) (format_exn exn));
  List.rev !diags

(* ───── Phrase completeness ───── *)

let is_complete_phrase code =
  let trimmed = String.trim code in
  if trimmed = "" then false
  else if String.ends_with ~suffix:";;" trimmed then true
  else
    (* Try parsing with ";;" appended. If it parses, the phrase is complete. If
       End_of_file, the parser consumed the phrase and wants more. If syntax
       error, the code is broken -- submit to show the error. *)
    let code_term = trimmed ^ ";;" in
    let lb = Lexing.from_string code_term in
    lb.lex_curr_p <-
      { pos_fname = "//toplevel//"; pos_lnum = 1; pos_bol = 0; pos_cnum = 0 };
    match !Toploop.parse_toplevel_phrase lb with
    | _ -> true
    | exception End_of_file -> false
    | exception _ -> true

(* ───── Markdown image detection ───── *)

(** Scan [s] for markdown data-URI images [![...](data:MIME;base64,DATA)] and
    emit each as a Display output. Surrounding text is emitted as Stdout. This
    allows pretty-printers (e.g. hugin.top) to render rich images in quill
    without depending on quill. *)
let emit_with_images ~emit s =
  let len = String.length s in
  let text_start = ref 0 in
  let i = ref 0 in
  while !i < len - 1 do
    if
      Char.equal (String.unsafe_get s !i) '!'
      && Char.equal (String.unsafe_get s (!i + 1)) '['
    then begin
      let start = !i in
      (* Skip past alt text to find ]( *)
      let j = ref (!i + 2) in
      while !j < len && not (Char.equal (String.unsafe_get s !j) ']') do
        incr j
      done;
      if
        !j < len - 1
        && Char.equal (String.unsafe_get s !j) ']'
        && Char.equal (String.unsafe_get s (!j + 1)) '('
      then begin
        let paren_start = !j + 2 in
        (* Check for data: URI *)
        let prefix = "data:" in
        let prefix_len = String.length prefix in
        if
          paren_start + prefix_len < len
          && String.sub s paren_start prefix_len = prefix
        then begin
          (* Find ;base64, *)
          let k = ref (paren_start + prefix_len) in
          let base64_marker = ";base64," in
          let marker_len = String.length base64_marker in
          let found_marker = ref false in
          let mime_end = ref 0 in
          while !k < len - marker_len && not !found_marker do
            if String.sub s !k marker_len = base64_marker then begin
              found_marker := true;
              mime_end := !k
            end
            else incr k
          done;
          if !found_marker then begin
            let data_start = !mime_end + marker_len in
            (* Find closing ) *)
            let m = ref data_start in
            while !m < len && not (Char.equal (String.unsafe_get s !m) ')') do
              incr m
            done;
            if !m < len then begin
              let mime =
                String.sub s (paren_start + prefix_len)
                  (!mime_end - paren_start - prefix_len)
              in
              let data = String.sub s data_start (!m - data_start) in
              (* Emit text before this image *)
              if start > !text_start then
                emit
                  (Quill.Cell.Stdout
                     (String.sub s !text_start (start - !text_start)));
              emit (Quill.Cell.Display { mime; data });
              i := !m + 1;
              text_start := !i
            end
            else incr i
          end
          else incr i
        end
        else incr i
      end
      else incr i
    end
    else incr i
  done;
  (* Emit remaining text *)
  if !text_start < len then begin
    let rest = String.sub s !text_start (len - !text_start) in
    if String.trim rest <> "" then emit (Quill.Cell.Stdout rest)
  end

(* ───── Kernel interface ───── *)

let status_ref = ref Quill.Kernel.Idle

let create ?setup ~on_event () =
  let setup_done = ref false in
  let ensure_setup () =
    if not !setup_done then (
      setup_done := true;
      initialize_if_needed ();
      match setup with Some f -> f () | None -> ())
  in
  let execute ~cell_id ~code =
    ensure_setup ();
    status_ref := Quill.Kernel.Busy;
    on_event (Quill.Kernel.Status_changed Busy);
    let emit output = on_event (Quill.Kernel.Output { cell_id; output }) in
    let ok, toplevel_out, toplevel_err =
      capture
        ~on_stdout:(fun s -> emit (Quill.Cell.Stdout s))
        ~on_stderr:(fun s -> emit (Quill.Cell.Stderr s))
        ~on_display:emit
        (fun ppf_out ppf_err -> execute_code ppf_out ppf_err code)
    in
    (* Emit toplevel formatter output (val bindings, type info). Scan for
       markdown data-URI images and convert to Display outputs. *)
    if toplevel_out <> "" then emit_with_images ~emit toplevel_out;
    if toplevel_err <> "" then emit (Quill.Cell.Stderr toplevel_err);
    (* Signal completion *)
    on_event (Quill.Kernel.Finished { cell_id; success = ok });
    status_ref := Quill.Kernel.Idle;
    on_event (Quill.Kernel.Status_changed Idle)
  in
  let interrupt () =
    (* Send SIGINT to the current thread - this will cause Sys.Break *)
    try Unix.kill (Unix.getpid ()) Sys.sigint with _ -> ()
  in
  let complete ~code ~pos =
    try
      ensure_setup ();
      complete_names ~code ~pos
    with exn ->
      Printf.eprintf "[quill-top] complete error: %s\n%!"
        (Printexc.to_string exn);
      []
  in
  let status () = !status_ref in
  let shutdown () =
    status_ref := Quill.Kernel.Shutting_down;
    on_event (Quill.Kernel.Status_changed Shutting_down)
  in
  {
    Quill.Kernel.execute;
    interrupt;
    complete;
    type_at =
      Some
        (fun ~code ~pos ->
          try
            ensure_setup ();
            type_at_pos ~code ~pos
          with exn ->
            Printf.eprintf "[quill-top] type_at error: %s\n%!"
              (Printexc.to_string exn);
            None);
    diagnostics =
      Some
        (fun ~code ->
          try
            ensure_setup ();
            compute_diagnostics ~code
          with exn ->
            Printf.eprintf "[quill-top] diagnostics error: %s\n%!"
              (Printexc.to_string exn);
            []);
    is_complete =
      Some
        (fun code ->
          try
            ensure_setup ();
            is_complete_phrase code
          with _ -> false);
    status;
    shutdown;
  }