sidekick/src/parser/ast_term.ml
2022-10-10 13:41:08 -04:00

111 lines
2.7 KiB
OCaml

module Pos = Parser_comb.Position
type position = Pos.t
type loc = { start: position; end_: position }
let mk_loc ~start ~end_ : loc = { start; end_ }
let pp_loc out (loc : loc) =
Fmt.fprintf out "%a - %a" Pos.pp loc.start Pos.pp loc.end_
let loc_merge a b : loc =
{ start = Pos.min a.start b.start; end_ = Pos.max a.end_ b.end_ }
type term = { view: term_view; loc: loc }
(** Expressions *)
and term_view =
| Var of string
| Int of string
| App of term * term list
| Let of (var * term) list * term
| Lambda of var list * term
| Pi of var list * term
| Arrow of term list * term
| Error_node of string
and var = { name: string; ty: term option }
(** Variable *)
open struct
let mk_ ~loc view : term = { loc; view }
end
let view (t : term) = t.view
let loc (t : term) = t.loc
let var ?ty name : var = { name; ty }
let mk_var ~loc v : term = mk_ ~loc (Var v)
let mk_app f args : term =
if args = [] then
f
else (
let loc = List.fold_left (fun loc a -> loc_merge loc a.loc) f.loc args in
mk_ ~loc (App (f, args))
)
let mk_arrow ~loc args ret : term =
if args = [] then
ret
else
mk_ ~loc (Arrow (args, ret))
let mk_lam ~loc args bod : term =
if args = [] then
bod
else
mk_ ~loc (Lambda (args, bod))
let mk_int ~loc x : term = mk_ ~loc (Int x)
let mk_pi ~loc args bod : term =
if args = [] then
bod
else
mk_ ~loc (Pi (args, bod))
let mk_let ~loc bs bod : term =
if bs = [] then
bod
else (
let loc = loc_merge loc bod.loc in
mk_ ~loc (Let (bs, bod))
)
let mk_error ~loc msg : term = mk_ ~loc (Error_node msg)
(** Pretty print terms *)
let rec pp_term out (e : term) : unit =
let pp = pp_term in
let pp_sub out e =
match e.view with
| App _ | Arrow _ | Pi _ | Let _ | Lambda _ ->
Fmt.fprintf out "(@[%a@])" pp e
| Var _ | Error_node _ | Int _ -> pp out e
in
let pp_tyvar out x =
match x.ty with
| None -> Fmt.string out x.name
| Some ty -> Fmt.fprintf out "(@[%s : %a@])" x.name pp ty
in
match e.view with
| Var v -> Fmt.string out v
| Error_node msg -> Fmt.fprintf out "<error %s>" msg
| Int i -> Fmt.string out i
| App (f, args) ->
Fmt.fprintf out "@[%a@ %a@]" pp_sub f (Util.pp_list pp_sub) args
| Arrow (args, ret) ->
Fmt.fprintf out "@[%a -> %a@]"
(Util.pp_list ~sep:" -> " pp_sub)
args pp_sub ret
| Let (bs, bod) ->
let ppb out (x, t) = Fmt.fprintf out "@[<2>%s :=@ %a@]" x.name pp t in
Fmt.fprintf out "@[@[<2>let@ %a@] in@ %a@]"
(Util.pp_list ~sep:"and" ppb)
bs pp_sub bod
| Lambda (args, bod) ->
Fmt.fprintf out "@[lam %a.@ %a@]" (Util.pp_list pp_tyvar) args pp_sub bod
| Pi (args, bod) ->
Fmt.fprintf out "@[pi %a.@ %a@]" (Util.pp_list pp_tyvar) args pp_sub bod