Library TypedExtraction.Common.PrettyPrinterMonad

From Stdlib Require Import List.
From MetaRocq.Utils Require Import monad_utils.
From MetaRocq.Utils Require Import bytestring.
From MetaRocq.SafeChecker Require Import PCUICErrors.
From MetaRocq.Erasure.Typed Require Import Utils.
From MetaRocq.Utils Require Import ResultMonad.
From TypedExtraction.Common Require Import Common.

Import String.
Import MonadNotation.
Import ListNotations.
Local Open Scope bs_scope.

Import Kernames.

Record PrettyPrinterState :=
  {
    indent_stack : list nat;
    used_names : list ident;
    output_lines : list (nat * string);
    cur_output_line : nat * string;
  }.

Definition PrettyPrinter A :=
  PrettyPrinterState -> result (A * PrettyPrinterState) string.

#[export]
Instance Monad_PrettyPrinter : Monad PrettyPrinter :=
  {| ret _ a pps := Ok (a, pps);
     bind _ _ m f pps :=
       match m pps with
       | Ok (a, pps) => f a pps
       | Err err => Err err
       end |}.

Definition map_pps
           (f : list nat -> list nat)
           (g : list ident -> list ident)
           (h : list (nat * string) -> list (nat * string))
           (i : nat * string -> nat * string) : PrettyPrinter unit :=
  fun pps =>
    Ok (tt,
        let (a, b, c, d) := pps in
        Build_PrettyPrinterState (f a) (g b) (h c) (i d)).

Fixpoint prefix_spaces n s :=
  match n with
  | 0 => s
  | S n => prefix_spaces n (String " " s)
  end.

Definition collect_output_lines pps :=
  MRList.rev_map (fun '(n, l) => prefix_spaces n l) (cur_output_line pps :: output_lines pps).

Definition result_string_err A := result A string.

Definition finish_print_lines {A} (pp : PrettyPrinter A) : result_string_err (A * list string) :=
  '(a, pps) <- pp {| indent_stack := [];
                     used_names := [];
                     output_lines := [];
                     cur_output_line := (0, "") |};;
  ret (a, collect_output_lines pps).

Definition collect_output pps :=
  concat nl (collect_output_lines pps).

Definition printer_fail {A} (str : string) : PrettyPrinter A :=
  fun pps => Err (str ++ nl ++ "failed after printing" ++ nl ++ collect_output pps).

Definition finish_print {A} (pp : PrettyPrinter A) : result (A * string) string :=
  '(a, lines) <- finish_print_lines pp;;
  ret (a, concat nl lines).

Definition get_indent : PrettyPrinter nat :=
  fun pps => Ok (hd 0 (indent_stack pps), pps).

Definition map_indent_stack (f : list nat -> list nat) : PrettyPrinter unit :=
  map_pps f id id id.

Definition push_indent (n : nat) : PrettyPrinter unit :=
  map_indent_stack (cons n).

Definition pop_indent : PrettyPrinter unit :=
  map_indent_stack (@tl nat).

Definition get_used_names : PrettyPrinter (list ident) :=
  fun pps => Ok (used_names pps, pps).

Definition map_used_names (f : list ident -> list ident) : PrettyPrinter unit :=
  map_pps id f id id.

Definition push_use (n : ident) : PrettyPrinter unit :=
  map_used_names (cons n).

Definition pop_use : PrettyPrinter unit :=
  map_used_names (@tl _).

Definition get_current_line_length : PrettyPrinter nat :=
  fun pps => Ok ((fst (cur_output_line pps)) + String.length (snd (cur_output_line pps)), pps).

Definition map_cur_output_line (f : nat * string -> nat * string) : PrettyPrinter unit :=
  map_pps id id id f.

Definition append (s : string) : PrettyPrinter unit :=
  map_cur_output_line (fun '(n, prev) => (n, prev ++ s)).

Definition append_nl : PrettyPrinter unit :=
  fun pps =>
    Ok (tt,
        {| indent_stack := indent_stack pps;
           used_names := used_names pps;
           output_lines := cur_output_line pps :: output_lines pps;
           cur_output_line := (hd 0 (indent_stack pps), "") |}).

Definition monad_append_join
           (sep : PrettyPrinter unit)
           (xs : list (PrettyPrinter unit)) : PrettyPrinter unit :=
  monad_fold_left (fun sep' x => sep';; x;; ret sep) xs (ret tt);;
  ret tt.

Definition append_join (sep : string) (s : list string) : PrettyPrinter unit :=
  monad_append_join (append sep) (map append s).

Definition monad_append_concat (xs : list (PrettyPrinter unit)) : PrettyPrinter unit :=
  monad_map id xs;;
  ret tt.

Definition append_concat (xs : list string) : PrettyPrinter unit :=
  monad_append_concat (map append xs).

Definition fresh_name (name : ident) (extra_used : list ident) : PrettyPrinter ident :=
  used <- get_used_names;;
  let used := (extra_used ++ used)%list in
  if existsb (bytestring.String.eqb name) used then
    (fix f n i :=
       match n with
       | 0 => ret "unreachable"
       | S n =>
         let numbered_name := bytestring.String.append name (MRString.string_of_nat i) in
         if existsb (bytestring.String.eqb numbered_name) used then
           f n (S i)
         else
           ret numbered_name
       end) (S (List.length used)) 2
  else
    ret name.

Definition string_of_env_error Σ e :=
  match e with
  | IllFormedDecl s e =>
    "IllFormedDecl " ++ s ++ "\nType error: " ++ string_of_type_error Σ e
  | AlreadyDeclared s => "Alreadydeclared " ++ s
  end.

Definition wrap_EnvCheck {astr A} f (ec : EnvCheck astr A) : PrettyPrinter A :=
  match ec with
  | CorrectDecl a => ret a
  | EnvError Σ err =>
    printer_fail ("EnvError: " ++ string_of_env_error (f Σ) err)
  end.

Module P := PCUICAst.
Definition wrap_typing_result {A} (Σ : P.PCUICEnvironment.global_env) (tr : typing_result A) : PrettyPrinter A :=
  match tr with
  | Checked et => ret et
  | TypeError te =>
    printer_fail ("TypeError: " ++ string_of_type_error (P.PCUICEnvironment.empty_ext Σ) te)
  end.

Definition wrap_option {A} (o : option A) (err : string) : PrettyPrinter A :=
  match o with
  | Some a => ret a
  | None => printer_fail err
  end.

Definition wrap_result {T E} (r : result T E) (err_string : E -> string) : PrettyPrinter T :=
  match r with
  | Ok t => ret t
  | Err e => printer_fail (err_string e)
  end.