mte/unikernel/duniverse/ocaml-syntax-shims/src/pp.real.ml
2025-11-11 02:07:51 +01:00

239 lines
7.8 KiB
OCaml

open StdLabels
open Shims
let prog_name = Filename.basename Sys.executable_name
let dump_ast = ref false
(* Table from positions to custom operators at these positions *)
let custom_operators = Hashtbl.create 128
module Wrap_lexer = struct
let save_loc = Location.curr
let restore_loc (lexbuf : Lexing.lexbuf) (loc : Location.t) =
lexbuf.lex_start_p <- loc.loc_start;
lexbuf.lex_curr_p <- loc.loc_end
let encode_op (tok : Parser.token) op =
(match tok with LET -> "let__" | AND -> "and__" | _ -> assert false) ^ op
let pending = Queue.create ()
let add (x : Parser.token * _) = Queue.push x pending
let register_custom_operator tok op (loc1 : Location.t) (loc2 : Location.t) =
let op = encode_op tok op in
Hashtbl.add custom_operators loc1.loc_start
({ loc1 with loc_end = loc2.loc_end }, op)
let wrap (lexer : Lexing.lexbuf -> Parser.token) lb =
if not (Queue.is_empty pending) then (
let tok, loc = Queue.pop pending in
restore_loc lb loc;
tok )
else
match lexer lb with
| (LET | AND) as tok ->
let loc = save_loc lb in
( match Let_trail.op lb with
| None -> ()
| Some op -> register_custom_operator tok op loc (save_loc lb) );
restore_loc lb loc;
tok
| LPAREN ->
let loc1 = save_loc lb in
let tok2 = lexer lb in
let loc2 = save_loc lb in
let tok, loc =
match tok2 with
| LET | AND -> (
match Let_trail.op lb with
| None ->
add (tok2, loc2);
(Parser.LPAREN, loc1)
| Some op -> (
let loc3 = save_loc lb in
match lexer lb with
| RPAREN ->
( LIDENT (encode_op tok2 op),
{ loc2 with loc_end = loc3.loc_end } )
| tok4 ->
let loc4 = save_loc lb in
add (tok2, loc2);
add (tok4, loc4);
register_custom_operator tok2 op loc2 loc3;
(LPAREN, loc1) ) )
| _ ->
add (tok2, loc2);
(LPAREN, loc1)
in
restore_loc lb loc;
tok
| tok -> tok
let () = Lexer.set_preprocessor (fun () -> Queue.clear pending) wrap
end
module Map_ast = struct
open Ast_mapper
open Asttypes
open Parsetree
open Ast_helper
let get_op vb =
match Hashtbl.find custom_operators vb.pvb_loc.loc_start with
| exception Not_found -> None
| loc, op -> Some (Exp.ident ~loc { txt = Lident op; loc })
let mapper =
let super = default_mapper in
let expr self expr =
let expr =
match expr.pexp_desc with
| Pexp_let (rf, (vb :: _ as vbs), body) -> (
match get_op vb with
| None -> expr
| Some op ->
if rf = Recursive then
Location.raise_errorf ~loc:expr.pexp_loc
"Custom 'let' operators cannot be recursive";
let patts, exprs =
List.map vbs ~f:(fun vb ->
let {
pvb_pat = patt;
pvb_expr = expr;
pvb_attributes = attrs;
pvb_loc = loc;
} =
vb
in
( match attrs with
| [] -> ()
| ({ loc; _ }, _) :: _ ->
Location.raise_errorf ~loc
"This attribute will be discarded" );
let op =
match get_op vb with
| Some op ->
Hashtbl.remove custom_operators vb.pvb_loc.loc_start;
op
| None ->
Location.raise_errorf ~loc
"Custom 'and' operator expected, got stantard \
'and' keyword"
in
(patt, (loc, op, expr)))
|> List.split
in
let patt =
List.fold_left (List.tl patts) ~init:(List.hd patts)
~f:(fun acc patt ->
let loc = patt.ppat_loc in
Pat.tuple ~loc [ acc; patt ])
in
let vars =
List.mapi exprs ~f:(fun i _ ->
Printf.sprintf "__future_syntax__%d__" i)
in
let pvars =
List.map2 vars patts ~f:(fun v p ->
let loc = { p.ppat_loc with loc_ghost = true } in
Pat.var ~loc { txt = v; loc })
in
let evars =
List.map2 vars exprs ~f:(fun v (_, _, e) ->
let loc = { e.pexp_loc with loc_ghost = true } in
Exp.ident ~loc { txt = Lident v; loc })
in
let expr =
List.fold_left2 (List.tl evars) (List.tl exprs)
~init:(List.hd evars) ~f:(fun acc var (loc, op, _) ->
Exp.apply ~loc op [ (nolabel, acc); (nolabel, var) ])
in
let body =
let loc = expr.pexp_loc in
Exp.apply ~loc op
[
(nolabel, expr);
(nolabel, Exp.fun_ ~loc nolabel None patt body);
]
in
List.fold_right2 pvars exprs ~init:body
~f:(fun var (loc, _, expr) acc ->
Exp.let_ Nonrecursive ~loc [ Vb.mk ~loc var expr ] acc) )
| _ -> expr
in
super.expr self expr
in
{ super with expr }
let map f ast =
let ast = f mapper ast in
let fail _ (loc, _) =
Location.raise_errorf ~loc "Invalid use of custom 'let' or 'and' operator"
in
Hashtbl.iter fail custom_operators;
ast
let structure = mapper.structure mapper
let signature = mapper.signature mapper
end
let process_file fn ~magic ~parse ~print ~map ~mk_ext =
let lexbuf = Lexing.from_channel (open_in_bin fn) in
Location.init lexbuf fn;
Location.input_lexbuf := Some lexbuf;
let ast =
try map (parse lexbuf)
with exn -> (
match error_of_exn exn with
| Some error ->
if !dump_ast then
[
mk_ext ?loc:None ?attrs:None (Ast_mapper.extension_of_error error);
]
else (
Location.report_error Format.err_formatter error;
exit 1 )
| None -> raise exn )
in
if !dump_ast then (
set_binary_mode_out stdout true;
output_string stdout magic;
output_value stdout fn;
output_value stdout ast;
flush stdout )
else Format.printf "%a@?" print ast
let process_file fn =
let ext =
match String.rindex fn '.' with
| exception Not_found -> ""
| i -> String.sub fn ~pos:i ~len:(String.length fn - i)
in
match ext with
| ".ml" ->
process_file fn ~magic:Config.ast_impl_magic_number
~parse:Parse.implementation ~print:Pprintast.structure
~map:Map_ast.structure ~mk_ext:Ast_helper.Str.extension
| ".mli" ->
process_file fn ~magic:Config.ast_intf_magic_number ~parse:Parse.interface
~print:Pprintast.signature ~map:Map_ast.signature
~mk_ext:Ast_helper.Sig.extension
| _ ->
Printf.eprintf "%s: Don't know what to do with %s.\n%!" prog_name fn;
exit 2
let () =
let args =
Arg.align
[
( "-dump-ast",
Arg.Set dump_ast,
" Output a binary AST rather than a pretty-printed source file" );
]
in
let usage = Printf.sprintf "Usage: %s [-dump-ast] FILES" prog_name in
Arg.parse args process_file usage