module_loader.ml

(* SPDX-License-Identifier: MPL-2.0 *)
(* SPDX-FileCopyrightText: 2025 hyperpolymath *)

(** Module loading and resolution.

    This module handles finding, loading, and caching modules from the file system.
*)

open Ast

(** Module loading errors *)
type load_error =
  | ModuleNotFound of string list  (** Module path not found *)
  | ModuleParseError of string * string  (** File path and parse error *)
  | ModuleCycle of string list  (** Circular dependency detected *)
  | InvalidModulePath of string list
[@@deriving show]

(** Helper for Result bind *)
let ( let* ) = Result.bind

(** Module loader configuration *)
type config = {
  stdlib_path : string;  (** Path to standard library *)
  search_paths : string list;  (** Additional module search paths *)
  current_dir : string;  (** Current working directory *)
}

(** Loaded module information *)
type loaded_module = {
  mod_path : string list;  (** Module path (e.g., ["Math", "Geometry"]) *)
  mod_program : program;  (** Parsed AST *)
  mod_file : string;  (** Source file path *)
}

(** Module loader state *)
type t = {
  config : config;
  loaded : (string list, loaded_module) Hashtbl.t;  (** Cache of loaded modules *)
  loading : (string list, unit) Hashtbl.t;  (** Currently loading (for cycle detection) *)
}

(** Create a new module loader *)
let create (config : config) : t =
  {
    config;
    loaded = Hashtbl.create 32;
    loading = Hashtbl.create 16;
  }

(** Test whether a candidate stdlib directory contains the prelude. *)
let stdlib_dir_valid path =
  Sys.file_exists (Filename.concat path "prelude.affine")

(** Walk up from [start_dir] looking for an ancestor whose [stdlib/]
    sub-directory contains the prelude. Returns [Some path] on hit, [None]
    after exhausting the parent chain. Used by [default_config] to find
    the stdlib when [affinescript check] is invoked from a sub-directory
    of an affinescript repo (issue #415). *)
let find_stdlib_walking_up start_dir =
  let rec walk dir =
    let candidate = Filename.concat dir "stdlib" in
    if stdlib_dir_valid candidate then Some candidate
    else
      let parent = Filename.dirname dir in
      if parent = dir then None
      else walk parent
  in
  walk start_dir

(** Discover the stdlib directory by trying, in order:
    1. [$AFFINESCRIPT_STDLIB] — explicit operator override
    2. [./stdlib] — current behaviour, for in-repo dev
    3. Walk up from CWD looking for [stdlib/prelude.affine] — catches
       sub-directory invocations of [affinescript check]
    4. [<binary_dir>/../share/affinescript/stdlib/] — XDG-style installed
       location alongside an installed binary
    5. [$HOME/.local/share/affinescript/stdlib/] — user-local install
    Falls back to ["./stdlib"] (and lets the loader's own
    [ModuleNotFound] surface for a clean error) if nothing matches.
    Issue #415. *)
let discover_stdlib () =
  match Sys.getenv_opt "AFFINESCRIPT_STDLIB" with
  | Some p -> p
  | None ->
    let cwd = Sys.getcwd () in
    let cwd_local = Filename.concat cwd "stdlib" in
    if stdlib_dir_valid cwd_local then cwd_local
    else match find_stdlib_walking_up cwd with
    | Some p -> p
    | None ->
      let binary_share =
        let bin = Sys.executable_name in
        if bin <> "" then
          Filename.concat
            (Filename.dirname (Filename.dirname bin))
            "share/affinescript/stdlib"
        else ""
      in
      if binary_share <> "" && stdlib_dir_valid binary_share then binary_share
      else
        let user_share =
          match Sys.getenv_opt "HOME" with
          | Some h -> Filename.concat h ".local/share/affinescript/stdlib"
          | None -> ""
        in
        if user_share <> "" && stdlib_dir_valid user_share then user_share
        else "./stdlib"  (* preserves the historical default error path *)

(** Create default configuration *)
let default_config () : config =
  {
    stdlib_path = discover_stdlib ();
    search_paths = [];
    current_dir = Sys.getcwd ();
  }

(** Convert module path to file path candidates

    For module path ["Math", "Geometry"], try:
    1. Math/Geometry.affine (nested module)
    2. Math.affine (check if it exports Geometry)
*)
let path_to_candidates (mod_path : string list) : string list =
  match mod_path with
  | [] -> []
  | [name] -> [name ^ ".affine"]
  | _ ->
    (* Try nested path: A/B/C.affine *)
    let nested = String.concat "/" mod_path ^ ".affine" in
    (* Try parent with submodule: A/B.affine (if B declares module C) *)
    let parent_parts = List.rev mod_path |> List.tl |> List.rev in
    let parent = String.concat "/" parent_parts ^ ".affine" in
    [nested; parent]

(** Search for a module file in search paths *)
let find_module_file (loader : t) (mod_path : string list) : string option =
  let candidates = path_to_candidates mod_path in
  let search_in_dir dir =
    List.find_map (fun candidate ->
      let full_path = Filename.concat dir candidate in
      if Sys.file_exists full_path then Some full_path else None
    ) candidates
  in
  (* Search order: 1) current dir, 2) stdlib, 3) additional search paths *)
  let all_paths =
    loader.config.current_dir ::
    loader.config.stdlib_path ::
    loader.config.search_paths
  in
  List.find_map search_in_dir all_paths

(** Load and parse a module file *)
let parse_module_file (file_path : string) : (program, load_error) result =
  try
    let prog = Parse_driver.parse_file file_path in
    Ok prog
  with
  | Lexer.Lexer_error (msg, _pos) ->
    Error (ModuleParseError (file_path, "Lexer error: " ^ msg))
  | Parse_driver.Parse_error (msg, _span) ->
    Error (ModuleParseError (file_path, "Parse error: " ^ msg))
  | Sys_error msg ->
    Error (ModuleParseError (file_path, "System error: " ^ msg))

(** Load a module and all its dependencies

    Returns the loaded module or an error if loading fails.
    Note: Does NOT resolve symbols - that's done by the caller.
*)
let rec load_module (loader : t) (mod_path : string list) : (loaded_module, load_error) result =
  (* Check if already loaded *)
  match Hashtbl.find_opt loader.loaded mod_path with
  | Some m -> Ok m
  | None ->
    (* Check for circular dependency *)
    if Hashtbl.mem loader.loading mod_path then
      Error (ModuleCycle mod_path)
    else begin
      (* Mark as loading *)
      Hashtbl.add loader.loading mod_path ();

      (* Find the module file *)
      match find_module_file loader mod_path with
      | None -> Error (ModuleNotFound mod_path)
      | Some file_path ->
        (* Parse the module *)
        match parse_module_file file_path with
        | Error e -> Error e
        | Ok prog ->
          (* Load dependencies first (imports) *)
          let* () = load_dependencies loader prog.prog_imports in

          let loaded_mod = {
            mod_path;
            mod_program = prog;
            mod_file = file_path;
          } in

          (* Cache the loaded module *)
          Hashtbl.add loader.loaded mod_path loaded_mod;
          Hashtbl.remove loader.loading mod_path;

          Ok loaded_mod
    end

(** Load all dependencies for a module *)
and load_dependencies (loader : t) (imports : import_decl list) : (unit, load_error) result =
  List.fold_left (fun acc import ->
    match acc with
    | Error e -> Error e
    | Ok () ->
      let mod_paths = match import with
        | ImportSimple (path, _alias) -> [List.map (fun id -> id.name) path]
        | ImportList (path, _items) -> [List.map (fun id -> id.name) path]
        | ImportGlob path -> [List.map (fun id -> id.name) path]
      in
      List.fold_left (fun acc' mod_path ->
        match acc' with
        | Error e -> Error e
        | Ok () ->
          match load_module loader mod_path with
          | Ok _ -> Ok ()
          | Error e -> Error e
      ) acc mod_paths
  ) (Ok ()) imports

(** Get a loaded module by path *)
let get_module (loader : t) (mod_path : string list) : loaded_module option =
  Hashtbl.find_opt loader.loaded mod_path

(** Get the loaded module's program *)
let get_program (loaded_mod : loaded_module) : program =
  loaded_mod.mod_program

(** Pretty-print module path *)
let show_module_path (path : string list) : string =
  String.concat "::" path

(** Check if a module is loaded *)
let is_loaded (loader : t) (mod_path : string list) : bool =
  Hashtbl.mem loader.loaded mod_path

(** Clear the module cache (for testing/reloading) *)
let clear_cache (loader : t) : unit =
  Hashtbl.clear loader.loaded;
  Hashtbl.clear loader.loading

(** Flatten cross-module imports by inlining the public top-level decls of
    each imported module into [prog.prog_decls], with deduplication by name.

    The WASM backend handles cross-module imports via a separate Wasm-import
    section ([Codegen.gen_imports]). Every other codegen ([Julia / JS / C /
    Rust / OCaml / ReScript / Lua / Bash / Nickel / LLVM / ...]) iterates
    [prog.prog_decls] only and would otherwise fail to find imported
    functions at codegen time. Inlining here is the simplest correct
    semantics for those targets and saves implementing per-backend module
    systems.

    The loader must already have been used by [Resolve.resolve_program_with_loader]
    so that the cache is populated; this function does NOT re-load on
    demand. If a referenced module isn't cached, its imports are silently
    skipped — the resolver would have reported the error already.

    Imports are processed in declaration order; later imports override
    earlier ones with the same fn name. Local decls in [prog.prog_decls]
    always win over imported ones. *)
let flatten_imports (loader : t) (prog : program) : program =
  (* Local-decl names suppress same-named imports of any kind. *)
  let local_names =
    List.filter_map (function
      | TopFn fd -> Some fd.fd_name.name
      | TopConst { tc_name; _ } -> Some tc_name.name
      | _ -> None
    ) prog.prog_decls
  in
  let already_in = Hashtbl.create 32 in
  List.iter (fun n -> Hashtbl.add already_in n ()) local_names;
  (* A flattened import is either a function or a constant; both share the
     same name-collision rule against [already_in]. *)
  let imported_decls =
    List.concat_map (fun imp ->
      let path_strs path =
        List.map (fun (id : ident) -> id.name) path
      in
      let mod_path = match imp with
        | ImportSimple (p, _) | ImportList (p, _) | ImportGlob p -> path_strs p
      in
      match Hashtbl.find_opt loader.loaded mod_path with
      | None -> []
      | Some lm ->
        let public_decls = List.filter_map (fun decl ->
          match decl with
          | TopFn fd when fd.fd_vis = Public || fd.fd_vis = PubCrate ->
            Some (fd.fd_name.name, `Fn fd)
          | TopConst { tc_vis; tc_name; _ }
            when tc_vis = Public || tc_vis = PubCrate ->
            Some (tc_name.name, `Const decl)
          | _ -> None
        ) lm.mod_program.prog_decls in
        let select = match imp with
          | ImportGlob _ -> public_decls
          | ImportSimple _ ->
            (* `use Foo` brings the namespace into scope but doesn't import
               specific symbols. For codegens that need them inlined we still
               include all public top-level bindings — same as glob. The
               resolver determines what's referenceable; codegen just needs
               the bodies present. *)
            public_decls
          | ImportList (_, items) ->
            List.filter_map (fun item ->
              let target = item.ii_name.name in
              List.find_opt (fun (n, _) -> n = target) public_decls
              |> Option.map (fun (_, found) ->
                let bound_name = match item.ii_alias with
                  | Some a -> a.name
                  | None -> target
                in
                let renamed = match found with
                  | `Fn fd ->
                    `Fn { fd with fd_name = { fd.fd_name with name = bound_name } }
                  | `Const (TopConst { tc_vis; tc_mut; tc_name; tc_ty; tc_value }) ->
                    `Const (TopConst {
                      tc_vis;
                      tc_mut;
                      tc_name = { tc_name with name = bound_name };
                      tc_ty;
                      tc_value;
                    })
                  | `Const _ ->
                    (* Unreachable: public_decls only stores TopConst under `Const`. *)
                    found
                in
                (bound_name, renamed))
            ) items
        in
        List.filter_map (fun (name, decl_kind) ->
          if Hashtbl.mem already_in name then None
          else begin
            Hashtbl.add already_in name ();
            match decl_kind with
            | `Fn fd -> Some (TopFn fd)
            | `Const decl -> Some decl
          end
        ) select
    ) prog.prog_imports
  in
  (* #138 follow-up: imported TYPE decls are intentionally NOT inlined here.
     An earlier #138 revision carried imported public [TopType]s so the
     prog_decls-iterating backends could register their constructors — but the
     non-Wasm backends (Deno-ESM / JS / Julia / C / Rust / ...) already emit
     the Option/Result constructors from a built-in runtime preamble, so
     carrying the prelude types declared `Some`/`None`/`Ok`/`Err` twice
     (`SyntaxError: Identifier 'Some' has already been declared` when the
     emitted Deno-ESM module is run under node). The Wasm backend learns
     imported variant tags natively via [Codegen.gen_imports] (it consumes the
     original, un-flattened [prog]); the other backends rely on their preamble.
     Re-introducing type-carrying for *user-defined* cross-module enums would
     need per-backend constructor dedup first. *)
  { prog with prog_decls = imported_decls @ prog.prog_decls }