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

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

(** Lean 4 emitter (MVP, dependent-type proof target).

Lowers Int/Bool/Float functions to Lean 4 [def] declarations. *)

open Ast

let lean_reserved = [

"abbrev"; "axiom"; "by"; "calc"; "class"; "def"; "deriving"; "do"; "elab";

"else"; "end"; "example"; "fun"; "have"; "if"; "import"; "in"; "inductive";

"instance"; "let"; "macro"; "match"; "mut"; "namespace"; "open"; "section";

"set_option"; "show"; "structure"; "syntax"; "then"; "theorem"; "where";

"with"; "true"; "false";

]

let mangle s = if List.mem s lean_reserved then s ^ "_" else s

let rec lean_type = function

| TyCon id when id.name = "Int" -> "Int"

| TyCon id when id.name = "Bool" -> "Bool"

| TyCon id when id.name = "Float" -> "Float"

| TyCon id when id.name = "String" -> "String"

| TyCon id when id.name = "Unit" -> "Unit"

| TyCon id -> mangle id.name

| TyOwn t | TyRef (_, t) | TyMut (_, t) -> lean_type t

| _ -> "Int"

let ret_type = function None -> "Unit" | Some t -> lean_type t

let rec gen_expr (e : expr) : string =

match e with

| ExprLit lit -> gen_lit lit

| ExprVar id -> mangle id.name

| ExprApp (callee, args) ->

let f = gen_expr callee in

let xs = List.map (fun a -> "(" ^ gen_expr a ^ ")") args in

f ^ " " ^ String.concat " " xs

| ExprBinary (a, op, b) ->

let s = match op with

| OpAdd -> "+" | OpSub -> "-" | OpMul -> "*" | OpDiv -> "/" | OpMod -> "%"

| OpEq -> "==" | OpNe -> "!="

| OpLt -> "<" | OpLe -> "<=" | OpGt -> ">" | OpGe -> ">="

| OpAnd -> "&&" | OpOr -> "||"

| OpConcat -> "++"

| _ -> "+"

in

"(" ^ gen_expr a ^ " " ^ s ^ " " ^ gen_expr b ^ ")"

| ExprUnary (OpNeg, x) -> "(- " ^ gen_expr x ^ ")"

| ExprUnary (OpNot, x) -> "(! " ^ gen_expr x ^ ")"

| ExprUnary _ -> "0"

| ExprIf { ei_cond; ei_then; ei_else } ->

let f = match ei_else with Some e -> gen_expr e | None -> "()" in

Printf.sprintf "(if %s then %s else %s)" (gen_expr ei_cond) (gen_expr ei_then) f

| ExprLet { el_pat; el_value; el_body; _ } ->

let var = match el_pat with PatVar id -> mangle id.name | _ -> "_" in

let body = match el_body with Some e -> gen_expr e | None -> "()" in

Printf.sprintf "(let %s := %s; %s)" var (gen_expr el_value) body

| ExprBlock blk -> gen_block blk

| ExprSpan (inner, _) -> gen_expr inner

| ExprReturn (Some e) -> gen_expr e

| _ -> "0"

and gen_lit = function

| LitInt (n, _) -> string_of_int n

| LitFloat (f, _) ->

let s = string_of_float f in

if String.length s > 0 && s.[String.length s - 1] = '.' then s ^ "0" else s

| LitBool (true, _) -> "true"

| LitBool (false, _) -> "false"

| LitString (s, _) -> "\"" ^ String.escaped s ^ "\""

| LitChar (c, _) -> "'" ^ Char.escaped c ^ "'"

| LitUnit _ -> "()"

and gen_block (blk : block) : string =

let rec fold = function

| [] -> (match blk.blk_expr with Some e -> gen_expr e | None -> "()")

| StmtLet { sl_pat = PatVar id; sl_value; _ } :: rest ->

Printf.sprintf "(let %s := %s; %s)" (mangle id.name) (gen_expr sl_value) (fold rest)

| _ :: rest -> fold rest

in

fold blk.blk_stmts

let gen_function (fd : fn_decl) : string =

let name = mangle fd.fd_name.name in

(* #624: fail loud on `return` rather than silently drop it. The block emitter

skips non-`let` statements, so an early `return` would vanish (the body

would lower to `:= ()`). Refuse instead — these are experimental backends. *)

if fn_body_contains_return fd.fd_body then

failwith

(Printf.sprintf

"Lean backend does not support `return` (in `%s`, Refs #624): early \

returns would silently drop control flow; rewrite as a tail expression"

fd.fd_name.name);

let params = match fd.fd_params with

| [] -> ""

| _ -> " " ^ String.concat " " (List.map (fun (p : param) ->

Printf.sprintf "(%s : %s)" (mangle p.p_name.name) (lean_type p.p_ty))

fd.fd_params)

in

let ret = ret_type fd.fd_ret_ty in

let body = match fd.fd_body with

| FnExpr e -> gen_expr e

| FnBlock b -> gen_block b

in

Printf.sprintf "def %s%s : %s :=\n %s\n\n" name params ret body

let generate (program : program) (_symbols : Symbol.t) : string =

let buf = Buffer.create 1024 in

Buffer.add_string buf "-- Generated by AffineScript compiler (Lean 4)\n";

Buffer.add_string buf "-- SPDX-License-Identifier: MPL-2.0\n\n";

List.iter (function

| TopFn fd -> Buffer.add_string buf (gen_function fd)

| _ -> ()

) program.prog_decls;

Buffer.contents buf

let codegen_lean (program : program) (symbols : Symbol.t) : (string, string) result =

try Ok (generate program symbols)

with

| Failure m -> Error ("Lean codegen error: " ^ m)

| e -> Error ("Lean codegen error: " ^ Printexc.to_string e)