rose-ash/hosts/ocaml/bootstrap.py

#!/usr/bin/env python3
"""
Bootstrap compiler: SX spec -> OCaml.

Loads the SX-to-OCaml transpiler (transpiler.sx), feeds it the spec files,
and produces sx_ref.ml — the transpiled evaluator as native OCaml.

Usage:
    python3 hosts/ocaml/bootstrap.py --output hosts/ocaml/lib/sx_ref.ml
"""
from __future__ import annotations

import os
import sys

_HERE = os.path.dirname(os.path.abspath(__file__))
_PROJECT = os.path.abspath(os.path.join(_HERE, "..", ".."))
sys.path.insert(0, _PROJECT)

from shared.sx.parser import parse_all
from shared.sx.types import Symbol


def extract_defines(source: str) -> list[tuple[str, list]]:
    """Parse .sx source, return list of (name, define-expr) for top-level defines."""
    exprs = parse_all(source)
    defines = []
    for expr in exprs:
        if isinstance(expr, list) and expr and isinstance(expr[0], Symbol):
            if expr[0].name == "define":
                name = expr[1].name if isinstance(expr[1], Symbol) else str(expr[1])
                defines.append((name, expr))
    return defines


# OCaml preamble — opens and runtime helpers
PREAMBLE = """\
(* sx_ref.ml — Auto-generated from SX spec by hosts/ocaml/bootstrap.py *)
(* Do not edit — regenerate with: python3 hosts/ocaml/bootstrap.py *)

[@@@warning "-26-27"]

open Sx_types
open Sx_runtime

(* Trampoline — forward ref, resolved after eval_expr is defined. *)
let trampoline_fn : (value -> value) ref = ref (fun v -> v)
let trampoline v = !trampoline_fn v



(* === Mutable globals — backing refs for transpiler's !_ref / _ref := === *)
let _strict_ref = ref (Bool false)
let _prim_param_types_ref = ref Nil
let _last_error_kont_ref = ref Nil
let _protocol_registry_ = Dict (Hashtbl.create 0)

"""


# OCaml fixups — wire up trampoline + iterative CEK run + JIT hook
FIXUPS = """\

(* Wire up trampoline to resolve thunks via the CEK machine *)
let () = trampoline_fn := (fun v ->
  match v with
  | Thunk (expr, env) -> eval_expr expr (Env env)
  | _ -> v)

(* Wire up the primitives trampoline so call_any in HO forms resolves Thunks *)
let () = Sx_primitives._sx_trampoline_fn := !trampoline_fn

(* Override recursive cek_run with iterative loop.
   On error, capture the kont from the last state for comp-trace. *)
let cek_run_iterative state =
  let s = ref state in
  (try
    while not (match cek_terminal_p !s with Bool true -> true | _ -> false)
       && not (match cek_suspended_p !s with Bool true -> true | _ -> false) do
      s := cek_step !s
    done;
    (match cek_suspended_p !s with
     | Bool true -> raise (Eval_error "IO suspension in non-IO context")
     | _ -> cek_value !s)
  with Eval_error msg ->
    _last_error_kont_ref := cek_kont !s;
    raise (Eval_error msg))

(* Collect component trace from a kont value *)
let collect_comp_trace kont =
  let trace = ref [] in
  let k = ref kont in
  while (match !k with List (_::_) -> true | _ -> false) do
    (match !k with
     | List (frame :: rest) ->
       (match frame with
        | CekFrame f when f.cf_type = "comp-trace" ->
          let name = match f.cf_name with String s -> s | _ -> "?" in
          let file = match f.cf_env with String s -> s | Nil -> "" | _ -> "" in
          trace := (name, file) :: !trace
        | Dict d when (match Hashtbl.find_opt d "type" with Some (String "comp-trace") -> true | _ -> false) ->
          let name = match Hashtbl.find_opt d "name" with Some (String s) -> s | _ -> "?" in
          let file = match Hashtbl.find_opt d "file" with Some (String s) -> s | _ -> "" in
          trace := (name, file) :: !trace
        | _ -> ());
       k := List rest
     | _ -> k := List [])
  done;
  List.rev !trace

(* Format a comp-trace into a human-readable string *)
let format_comp_trace trace =
  match trace with
  | [] -> ""
  | entries ->
    let lines = List.mapi (fun i (name, file) ->
      let prefix = if i = 0 then "  in " else "  called from " in
      if file = "" then prefix ^ "~" ^ name
      else prefix ^ "~" ^ name ^ " (" ^ file ^ ")"
    ) entries in
    "\n" ^ String.concat "\n" lines

(* Enhance an error message with component trace *)
let enhance_error_with_trace msg =
  let trace = collect_comp_trace !_last_error_kont_ref in
  _last_error_kont_ref := Nil;
  msg ^ (format_comp_trace trace)



"""


def compile_spec_to_ml(spec_dir: str | None = None) -> str:
    """Compile the SX spec to OCaml source."""
    import tempfile
    from shared.sx.ocaml_sync import OcamlSync
    from shared.sx.parser import serialize

    if spec_dir is None:
        spec_dir = os.path.join(_PROJECT, "spec")

    # Load the transpiler into OCaml kernel
    bridge = OcamlSync()
    transpiler_path = os.path.join(_HERE, "transpiler.sx")
    bridge.load(transpiler_path)

    # Spec files to transpile (in dependency order)
    # stdlib.sx functions are already registered as OCaml primitives —
    # only the evaluator needs transpilation.
    sx_files = [
        ("evaluator.sx", "evaluator (frames + eval + CEK)"),
    ]

    parts = [PREAMBLE]

    for filename, label in sx_files:
        filepath = os.path.join(spec_dir, filename)
        if not os.path.exists(filepath):
            print(f"Warning: {filepath} not found, skipping", file=sys.stderr)
            continue

        with open(filepath) as f:
            src = f.read()
        defines = extract_defines(src)

        # Skip defines provided by preamble, fixups, or already-registered primitives
        # Skip: preamble-provided, math primitives, and stdlib functions
        # that use loop/named-let (transpiler can't handle those yet)
        skip = {"trampoline", "ceil", "floor", "round", "abs", "min", "max",
                "debug-log", "debug_log", "range", "chunk-every", "zip-pairs",
                "string-contains?", "starts-with?", "ends-with?",
                "string-replace", "trim", "split", "index-of",
                "pad-left", "pad-right", "char-at", "substring"}
        defines = [(n, e) for n, e in defines if n not in skip]

        # Deduplicate — keep last definition for each name (CEK overrides tree-walk)
        seen = {}
        for i, (n, e) in enumerate(defines):
            seen[n] = i
        defines = [(n, e) for i, (n, e) in enumerate(defines) if seen[n] == i]

        # Build the defines list and known names for the transpiler
        defines_list = [[name, expr] for name, expr in defines]
        known_names = [name for name, _ in defines]

        # Serialize defines + known names to temp file, load into kernel
        defines_sx = serialize(defines_list)
        known_sx = serialize(known_names)
        with tempfile.NamedTemporaryFile(mode="w", suffix=".sx", delete=False) as tmp:
            tmp.write(f"(define _defines \'{defines_sx})\n")
            tmp.write(f"(define _known_defines \'{known_sx})\n")
            tmp_path = tmp.name
        try:
            bridge.load(tmp_path)
        finally:
            os.unlink(tmp_path)

        # Call ml-translate-file — emits as single let rec block
        result = bridge.eval("(ml-translate-file _defines)")

        parts.append(f"\n(* === Transpiled from {label} === *)\n")
        parts.append(result)

    bridge.stop()
    parts.append(FIXUPS)
    output = "\n".join(parts)

    # Mutable globals (*strict*, *prim-param-types*) are now handled by
    # the transpiler directly — it emits !_ref for reads, _ref := for writes.
    import re

    # Remove `and _protocol_registry_ = (Dict ...)` from the let rec block —
    # it's defined in the preamble as a top-level let, and Hashtbl.create
    # is not allowed as a let rec right-hand side.
    output = re.sub(
        r'\n\(\* \*protocol-registry\*.*?\nand _protocol_registry_ =\n  \(Dict \(Hashtbl\.create 0\)\)\n',
        '\n',
        output
    )

    return output


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Bootstrap SX spec -> OCaml")
    parser.add_argument(
        "--output", "-o",
        default=None,
        help="Output file (default: stdout)",
    )
    args = parser.parse_args()

    result = compile_spec_to_ml()

    if args.output:
        with open(args.output, "w") as f:
            f.write(result)
        size = os.path.getsize(args.output)
        print(f"Wrote {args.output} ({size} bytes)", file=sys.stderr)
    else:
        print(result)


if __name__ == "__main__":
    main()