ligo/src/ast_simplified/PP.ml

open Types
open PP_helpers
open Format

let list_sep_d x ppf lst = match lst with
  | [] -> ()
  | _ -> fprintf ppf "@;  @[<v>%a@]@;" (list_sep x (tag "@;")) lst

let smap_sep_d x ppf m =
  if Map.String.is_empty m
  then ()
  else fprintf ppf "@;  @[<v>%a@]@;" (smap_sep x (tag "@;")) m

let rec type_expression ppf (te:type_expression) = match te with
  | T_tuple lst -> fprintf ppf "tuple[%a]" (list_sep_d type_expression) lst
  | T_sum m -> fprintf ppf "sum[%a]" (smap_sep_d type_expression) m
  | T_record m -> fprintf ppf "record[%a]" (smap_sep_d type_expression) m
  | T_function (p, r) -> fprintf ppf "%a -> %a" type_expression p type_expression r
  | T_variable name -> fprintf ppf "%s" name
  | T_constant (name, lst) -> fprintf ppf "%s(%a)" name (list_sep_d type_expression) lst

let literal ppf (l:literal) = match l with
  | Literal_unit -> fprintf ppf "Unit"
  | Literal_bool b -> fprintf ppf "%b" b
  | Literal_int n -> fprintf ppf "%d" n
  | Literal_nat n -> fprintf ppf "+%d" n
  | Literal_tez n -> fprintf ppf "%dtz" n
  | Literal_string s -> fprintf ppf "%S" s
  | Literal_bytes b -> fprintf ppf "0x%s" @@ Bytes.to_string @@ Bytes.escaped b
  | Literal_address s -> fprintf ppf "@%S" s
  | Literal_operation _ -> fprintf ppf "Operation(...bytes)"

let rec expression ppf (e:expression) = match e with
  | E_literal l -> literal ppf l
  | E_variable name -> fprintf ppf "%s" name
  | E_application (f, arg) -> fprintf ppf "(%a)@(%a)" annotated_expression f annotated_expression arg
  | E_constructor (name, ae) -> fprintf ppf "%s(%a)" name annotated_expression ae
  | E_constant (name, lst) -> fprintf ppf "%s(%a)" name (list_sep_d annotated_expression) lst
  | E_tuple lst -> fprintf ppf "tuple[%a]" (list_sep_d annotated_expression) lst
  | E_accessor (ae, p) -> fprintf ppf "%a.%a" annotated_expression ae access_path p
  | E_record m -> fprintf ppf "record[%a]" (smap_sep_d annotated_expression) m
  | E_map m -> fprintf ppf "map[%a]" (list_sep_d assoc_annotated_expression) m
  | E_list lst -> fprintf ppf "list[%a]" (list_sep_d annotated_expression) lst
  | E_look_up (ds, ind) -> fprintf ppf "(%a)[%a]" annotated_expression ds annotated_expression ind
  | E_lambda {binder;input_type;output_type;result} ->
      fprintf ppf "lambda (%s:%a) : %a return %a"
        binder (PP_helpers.option type_expression) input_type (PP_helpers.option type_expression) output_type
        annotated_expression result
  | E_matching (ae, m) ->
      fprintf ppf "match %a with %a" annotated_expression ae (matching annotated_expression) m
  | E_failwith ae ->
      fprintf ppf "failwith %a" annotated_expression ae
  | E_sequence (a , b) ->
    fprintf ppf "%a ; %a"
      annotated_expression a
      annotated_expression b
  | E_loop (expr , body) ->
    fprintf ppf "%a ; %a"
      annotated_expression expr
      annotated_expression body
  | E_assign (name , path , expr) ->
    fprintf ppf "%s.%a := %a"
      name
      PP_helpers.(list_sep access (const ".")) path
      annotated_expression expr
  | E_let_in { binder; rhs; result } ->
      fprintf ppf "let %s = %a in %a" binder annotated_expression rhs annotated_expression result
  | E_skip -> fprintf ppf "skip"

and assoc_annotated_expression ppf : (ae * ae) -> unit = fun (a, b) ->
  fprintf ppf "%a -> %a" annotated_expression a annotated_expression b

and access ppf (a:access) =
  match a with
  | Access_tuple n -> fprintf ppf "%d" n
  | Access_record s -> fprintf ppf "%s" s
  | Access_map s -> fprintf ppf "(%a)" annotated_expression s

and access_path ppf (p:access_path) =
  fprintf ppf "%a" (list_sep access (const ".")) p

and type_annotation ppf (ta:type_expression option) = match ta with
  | None -> fprintf ppf ""
  | Some t -> type_expression ppf t

and annotated_expression ppf (ae:annotated_expression) = match ae.type_annotation with
  | None -> fprintf ppf "%a" expression ae.expression
  | Some t -> fprintf ppf "(%a) : %a" expression ae.expression type_expression t

and value : _ -> value -> unit = fun x -> annotated_expression x

and single_record_patch ppf ((p, ae) : string * ae) =
  fprintf ppf "%s <- %a" p annotated_expression ae

and single_tuple_patch ppf ((p, ae) : int * ae) =
  fprintf ppf "%d <- %a" p annotated_expression ae

and matching_variant_case : type a . (_ -> a -> unit) -> _ -> (constructor_name * name) * a -> unit =
  fun f ppf ((c,n),a) ->
  fprintf ppf "| %s %s -> %a" c n f a

and matching : type a . (formatter -> a -> unit) -> formatter -> a matching -> unit =
  fun f ppf m -> match m with
    | Match_tuple (lst, b) ->
        fprintf ppf "let (%a) = %a" (list_sep_d string) lst f b
    | Match_variant lst ->
        fprintf ppf "%a" (list_sep (matching_variant_case f) (tag "@.")) lst
    | Match_bool {match_true ; match_false} ->
        fprintf ppf "| True -> %a @.| False -> %a" f match_true f match_false
    | Match_list {match_nil ; match_cons = (hd, tl, match_cons)} ->
        fprintf ppf "| Nil -> %a @.| %s :: %s -> %a" f match_nil hd tl f match_cons
    | Match_option {match_none ; match_some = (some, match_some)} ->
        fprintf ppf "| None -> %a @.| Some %s -> %a" f match_none some f match_some

let declaration ppf (d:declaration) = match d with
  | Declaration_type {type_name ; type_expression = te} ->
      fprintf ppf "type %s = %a" type_name type_expression te
  | Declaration_constant {name ; annotated_expression = ae} ->
      fprintf ppf "const %s = %a" name annotated_expression ae

let program ppf (p:program) =
  fprintf ppf "@[<v>%a@]" (list_sep declaration (tag "@;")) (List.map Location.unwrap p)