(* SPDX-License-Identifier: MPL-2.0 *)

(* SPDX-FileCopyrightText: 2024-2026 Jonathan D.A. Jewell (hyperpolymath) *)

(** Python-face: source-level transformer for Python-style AffineScript syntax.

Maps Python surface syntax to canonical AffineScript before lexing and

parsing with the standard parser. The compiler is face-agnostic (ADR-010):

only this module and the error-vocabulary layer know about the Python face.

Surface mappings:

{v

def name(...) → fn name(...)

True / False → true / false

None → ()

and / or → && / ||

not EXPR → !EXPR

class Name → type Name

pass → ()

# comment → // comment

import a.b → use a::b

from a import b → use a::b

if cond: → if cond { (block-opening colon → brace)

else: → } else {

elif cond: → } else if cond {

while cond: → while cond {

for x in e: → for x in e {

match e: → match e {

handle e: → handle e {

INDENT → (block opened by preceding {)

DEDENT → }

statement line → line;

v}

Implementation note: this is a line-by-line text preprocessor.

Span information reported by the compiler refers to the canonical

AffineScript text, not the original Python-style source. Face-aware

error vocabulary (ADR-010 §3) is a follow-up task.

*)

(* ─── Character helpers ──────────────────────────────────────────────── *)

let is_id_char c =

(c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')

|| (c >= '0' && c <= '9') || c = '_'

let starts_with s prefix =

let sl = String.length s and pl = String.length prefix in

sl >= pl && String.sub s 0 pl = prefix

(** Count leading spaces/tabs (each tab = 1 unit for simplicity). *)

let indent_of line =

let len = String.length line in

let i = ref 0 in

while !i < len && (line.[!i] = ' ' || line.[!i] = '\t') do incr i done;

!i

(* ─── Comment stripping ──────────────────────────────────────────────── *)

(** Split a line at a Python-style [#] comment, respecting string literals.

Returns [(code_part, comment_text option)]. The returned [code_part]

still has its trailing whitespace. *)

let strip_py_comment line =

let len = String.length line in

let i = ref 0 in

let in_str = ref false in

let str_delim = ref '"' in

while !i < len do

let c = line.[!i] in

if !in_str then begin

if c = !str_delim && (!i = 0 || line.[!i - 1] <> '\\') then

in_str := false

end else begin

if c = '#' then begin

(* Found comment start — exit the loop *)

let code = String.sub line 0 !i in

let comment = String.sub line (!i + 1) (len - !i - 1) in

i := len; (* break *)

(* We need to return here but OCaml loops can't early-return;

use a reference instead. *)

ignore (code, comment) (* handled below via exception-free approach *)

end else if c = '"' || c = '\'' then begin

in_str := true;

str_delim := c

end

end;

incr i

done;

(* Re-scan cleanly without mutable trickery *)

let result = ref (line, None) in

let j = ref 0 in

let in_s = ref false in

let sd = ref '"' in

while !j < len && not (not !in_s && line.[!j] = '#') do

let c = line.[!j] in

if !in_s then begin

if c = !sd && (!j = 0 || line.[!j - 1] <> '\\') then in_s := false

end else begin

if c = '"' || c = '\'' then begin in_s := true; sd := c end

end;

incr j

done;

if !j < len && line.[!j] = '#' then

result := (String.sub line 0 !j,

Some (String.sub line (!j + 1) (len - !j - 1)));

!result

(* ─── Word-level keyword substitution ───────────────────────────────── *)

(** Replace all whole-word occurrences of [from_w] with [to_w] in [s].

"Whole-word" means not immediately preceded or followed by an id char.

Does not track string literals — callers should apply to code-only text. *)

let replace_word ~from_w ~to_w s =

let flen = String.length from_w in

let slen = String.length s in

let buf = Buffer.create slen in

let i = ref 0 in

while !i < slen do

if slen - !i >= flen && String.sub s !i flen = from_w then begin

let before_ok = !i = 0 || not (is_id_char s.[!i - 1]) in

let after_ok = !i + flen >= slen || not (is_id_char s.[!i + flen]) in

if before_ok && after_ok then begin

Buffer.add_string buf to_w;

i := !i + flen

end else begin

Buffer.add_char buf s.[!i];

incr i

end

end else begin

Buffer.add_char buf s.[!i];

incr i

end

done;

Buffer.contents buf

(** Apply all Python→AffineScript keyword substitutions to a code fragment.

Order matters: longer / more-specific patterns first. *)

let apply_keywords s =

(* `not` as a whole word; `not x` → `! x` (lexer treats `!` and `x` correctly) *)

let s = replace_word ~from_w:"not" ~to_w:"!" s in

let s = replace_word ~from_w:"def" ~to_w:"fn" s in

let s = replace_word ~from_w:"True" ~to_w:"true" s in

let s = replace_word ~from_w:"False" ~to_w:"false" s in

let s = replace_word ~from_w:"None" ~to_w:"()" s in

let s = replace_word ~from_w:"and" ~to_w:"&&" s in

let s = replace_word ~from_w:"or" ~to_w:"||" s in

let s = replace_word ~from_w:"class" ~to_w:"type" s in

let s = replace_word ~from_w:"pass" ~to_w:"()" s in

s

(* ─── Module path helpers ────────────────────────────────────────────── *)

(** Replace `.` with `::` in a dotted module path (e.g. `a.b.c` → `a::b::c`). *)

let dots_to_colons s =

let len = String.length s in

let buf = Buffer.create (len + 4) in

String.iter (fun c ->

if c = '.' then Buffer.add_string buf "::"

else Buffer.add_char buf c

) s;

Buffer.contents buf

(** Try to transform an [import] or [from … import] line.

Returns [Some canonical_line] or [None] if the line is not an import. *)

let transform_import_line stripped =

if starts_with stripped "import " then begin

let path = String.trim (String.sub stripped 7 (String.length stripped - 7)) in

Some ("use " ^ dots_to_colons path ^ ";")

end else if starts_with stripped "from " then begin

let rest = String.trim (String.sub stripped 5 (String.length stripped - 5)) in

(* Expect: MODULE_PATH import NAME [, NAME2 ...] *)

(match String.split_on_char ' ' rest with

| mod_part :: "import" :: names when names <> [] ->

let name = String.concat "::" (List.filter (fun s -> s <> "") names) in

Some ("use " ^ dots_to_colons mod_part ^ "::" ^ name ^ ";")

| _ -> None)

end else None

(* ─── Block-opening colon detection ─────────────────────────────────── *)

(** True if [stripped] ends with `:` (the block-opening colon pattern).

We require the line to start with a known block-opening keyword or

be one of the bare clause keywords so we don't eat type annotation colons

that happen to be at line-end. *)

let block_opening_prefixes = [

"if "; "while "; "for "; "with "; "match "; "handle ";

"fn "; "def "; "class "; "type ";

"try"; "except"; "finally"; "loop"; "effect ";

]

let is_block_opener stripped =

let slen = String.length stripped in

slen > 0 && stripped.[slen - 1] = ':'

&& (List.exists (fun p -> starts_with stripped p) block_opening_prefixes

|| stripped = "else:" || stripped = "try:" || stripped = "finally:")

(** Strip the trailing `:` from a block-opener and return the body. *)

let strip_block_colon stripped =

let len = String.length stripped in

assert (len > 0 && stripped.[len - 1] = ':');

String.trim (String.sub stripped 0 (len - 1))

(* ─── Else / elif detection ──────────────────────────────────────────── *)

let is_else_clause stripped =

stripped = "else:" || starts_with stripped "else :"

let is_elif_clause stripped =

starts_with stripped "elif "

(** Extract the condition from an `elif COND:` line. *)

let elif_condition stripped =

(* stripped = "elif COND:" *)

let rest = String.sub stripped 5 (String.length stripped - 5) in

let rest = String.trim rest in

if String.length rest > 0 && rest.[String.length rest - 1] = ':' then

String.trim (String.sub rest 0 (String.length rest - 1))

else rest

(* ─── Main transformer ───────────────────────────────────────────────── *)

(** True if [raw_line] is blank or comment-only (carries no code). *)

let is_blank_line raw =

let (code, _) = strip_py_comment (String.trim raw) in

String.trim code = ""

(** Transform Python-style AffineScript source text to canonical AffineScript.

The result is valid input for the standard lexer + Menhir parser.

Tail-position detection: a regular statement (non-block-opener) in the

last position of a block — i.e. the next meaningful line's indent is

strictly less than the current line's indent — is emitted WITHOUT a

trailing `;`. This preserves the expression-as-return-value semantics

that AffineScript blocks require (a trailing `;` would make the block

yield unit rather than the expression's value). *)

let transform_source source =

let lines = Array.of_list (String.split_on_char '\n' source) in

let n = Array.length lines in

let out = Buffer.create (String.length source + 256) in

(* Indentation stack: innermost level at head, outermost (0) at tail. *)

let stack = ref [0] in

let top () = match !stack with h :: _ -> h | [] -> 0 in

let emit_dedents target =

while top () > target do

Buffer.add_string out "}\n";

stack := List.tl !stack

done

in

(* The indent level of the next non-blank/non-comment line after index [i],

or [-1] when there is no such line (EOF). *)

let next_meaningful_indent i =

let j = ref (i + 1) in

while !j < n && is_blank_line lines.(!j) do incr j done;

if !j >= n then -1 else indent_of lines.(!j)

in

for i = 0 to n - 1 do

let raw_line = lines.(i) in

let ind = indent_of raw_line in

let (code_part, comment_opt) = strip_py_comment (String.trim raw_line) in

let stripped = String.trim code_part in

(* Append optional trailing comment in canonical // style *)

let with_comment line_text =

match comment_opt with

| None -> line_text ^ "\n"

| Some c -> line_text ^ " // " ^ String.trim c ^ "\n"

in

if stripped = "" then begin

(* Blank or comment-only line *)

(match comment_opt with

| Some c -> Buffer.add_string out ("// " ^ String.trim c ^ "\n")

| None -> Buffer.add_char out '\n')

end else if is_else_clause stripped then begin

(* `else:` appears at the same indent as the matching `if`.

emit_dedents closes the body of the then-branch (emitting `}`).

We then continue with `else {` — no leading `}` here because

emit_dedents already supplied it. *)

emit_dedents ind;

Buffer.add_string out (with_comment "else {")

end else if is_elif_clause stripped then begin

(* `elif COND:` — same structure as `else:` *)

emit_dedents ind;

let cond = apply_keywords (elif_condition stripped) in

Buffer.add_string out (with_comment ("else if " ^ cond ^ " {"))

end else begin

(* Normal line (statement or block-opener) *)

emit_dedents ind;

(* Push a new indent level when we step in *)

if ind > top () then stack := ind :: !stack;

let indent_str = String.make ind ' ' in

(* Tail-position check: the next meaningful line is less indented (or

EOF), meaning this is the last expression in its block. Omit `;`

so the block's value is this expression, not unit. *)

let next_ind = next_meaningful_indent i in

let is_tail = next_ind < ind in (* -1 (EOF) satisfies this for ind > 0 *)

let line_text = match transform_import_line stripped with

| Some s -> s (* imports are always top-level statements *)

| None ->

if is_block_opener stripped then

(* Replace trailing `:` with ` {` *)

apply_keywords (strip_block_colon stripped) ^ " {"

else if is_tail then

(* Tail expression: no `;` — this is the block's return value *)

apply_keywords stripped

else

(* Mid-block statement: terminate with `;` *)

apply_keywords stripped ^ ";"

in

Buffer.add_string out (indent_str ^ with_comment line_text)

end

done;

(* Close any blocks still open at EOF *)

emit_dedents 0;

Buffer.contents out

(* ─── Entry points ───────────────────────────────────────────────────── *)

(** Parse a Python-face program from a string, returning a canonical AST.

Raises {!Parse_driver.Parse_error} or {!Lexer.Lexer_error} on failure. *)

let parse_string_python ~file content =

let canonical = transform_source content in

Parse_driver.parse_string ~file canonical

(** Parse a Python-face program from a file, returning a canonical AST. *)

let parse_file_python path =

let chan = open_in_bin path in

Fun.protect

~finally:(fun () -> close_in chan)

(fun () ->

let content = really_input_string chan (in_channel_length chan) in

parse_string_python ~file:path content)

(** Expose the text-to-text transform for debugging and testing. *)

let preview_transform source = transform_source source