Library MetaRocq.Template.TemplateMonad.Core

(* Distributed under the terms of the MIT license. *)
From MetaRocq.Utils Require Import utils.
From MetaRocq.Template Require Import Ast AstUtils Common.

Local Set Universe Polymorphism.
Local Unset Universe Minimization ToSet.
Local Unset Asymmetric Patterns.
Import MRMonadNotation.

The Template Monad

A monad for programming with Template Rocq structures. Use Run TemplateProgram on a monad action to produce its side-effects.
Uses a reduction strategy specifier reductionStrategy.

The TemplateMonad type

Cumulative Inductive TemplateMonad@{t u} : Type@{t} Prop :=
(* Monadic operations *)
| tmReturn : {A:Type@{t}}, A TemplateMonad A
| tmBind : {A B : Type@{t}}, TemplateMonad A (A TemplateMonad B)
                                TemplateMonad B

(* General commands *)
| tmPrint : {A:Type@{t}}, A TemplateMonad unit
| tmMsg : string TemplateMonad unit
| tmFail : {A:Type@{t}}, string TemplateMonad A
| tmEval : reductionStrategy {A:Type@{t}}, A TemplateMonad A

(* Return the defined constant *)
| tmLemma : ident A : Type@{t}, TemplateMonad A
| tmDefinitionRed_ : (opaque : bool), ident option reductionStrategy {A:Type@{t}}, A TemplateMonad A
| tmAxiomRed : ident option reductionStrategy A : Type@{t}, TemplateMonad A
| tmVariable : ident Type@{t} TemplateMonad unit

(* Guaranteed to not cause "... already declared" error *)
| tmFreshName : ident TemplateMonad ident

(* Returns the list of globrefs corresponding to a qualid,
   the head is the default one if any. *)

| tmLocate : qualid TemplateMonad (list global_reference)
| tmLocateModule : qualid TemplateMonad (list modpath)
| tmLocateModType : qualid TemplateMonad (list modpath)
| tmCurrentModPath : unit TemplateMonad modpath

(* Quoting and unquoting commands *)
(* Similar to MetaRocq Quote Definition ... := ... *)
| tmQuote : {A:Type@{t}}, A TemplateMonad Ast.term
(* Similar to MetaRocq Quote Recursively Definition but takes a boolean "bypass opacity" flag.
  (true - quote bodies of all dependencies (transparent and opaque);
   false -quote bodies of transparent definitions only) *)

| tmQuoteRecTransp : {A:Type@{t}}, A bool(* bypass opacity? *) TemplateMonad program
(* Quote the body of a definition or inductive. Its name need not be fully qualified *)
| tmQuoteInductive : kername TemplateMonad mutual_inductive_body
| tmQuoteUniverses : TemplateMonad ConstraintSet.t
| tmQuoteConstant : kername bool (* bypass opacity? *) TemplateMonad constant_body
| tmQuoteModule : qualid TemplateMonad (list global_reference)
| tmQuoteModFunctor : qualid TemplateMonad (list global_reference)
| tmQuoteModType : qualid TemplateMonad (list global_reference)
(* unquote before making the definition *)
(* FIXME take an optional universe context as well *)
| tmMkInductive : bool mutual_inductive_entry TemplateMonad unit
| tmUnquote : Ast.term TemplateMonad typed_term@{u}
| tmUnquoteTyped : A : Type@{t}, Ast.term TemplateMonad A

(* Typeclass registration and querying for an instance *)
| tmExistingInstance : hint_locality global_reference TemplateMonad unit
| tmInferInstance : option reductionStrategy A : Type@{t}, TemplateMonad (option_instance A)
.

This version of tmBind flattens nesting structure; using it in deeply recursive template programs can speed things up drastically We use tmBind in the recursive position to avoid quadratic blowup in the number of tmOptimizedBinds
Fixpoint tmOptimizedBind@{t u} {A B : Type@{t}} (v : TemplateMonad@{t u} A) : (A TemplateMonad@{t u} B) TemplateMonad@{t u} B
  := match v with
     | tmReturn xfun ff x
     | tmBind v kfun ftmOptimizedBind v (fun vtmBind (k v) f)
     | tmFail msgfun _tmFail msg
     | vtmBind v
     end.

Flatten nested tmBind
Fixpoint tmOptimize'@{t u} {A B : Type@{t}} (v : TemplateMonad@{t u} A) : (A TemplateMonad@{t u} B) TemplateMonad@{t u} B
  := match v with
     | tmReturn xfun ff x
     | tmBind v kfun ftmOptimize' v (fun vtmOptimize' (k v) f)
     | tmFail msgfun _tmFail msg
     | vtmBind v
     end.
Definition tmOptimize@{t u} {A : Type@{t}} (v : TemplateMonad@{t u} A) : TemplateMonad@{t u} A
  := tmOptimize' v tmReturn.

This allow to use notations of MonadNotation
Definition TemplateMonad_UnoptimizedMonad@{t u} : Monad@{t u} TemplateMonad@{t u} :=
  {| ret := @tmReturn ; bind := @tmBind |}.

Definition TemplateMonad_OptimizedMonad@{t u} : Monad@{t u} TemplateMonad@{t u} :=
  {| ret := @tmReturn ; bind := @tmOptimizedBind |}.

(* We don't want to make the optimized monad an instance, because it blows up performance in some cases *)
Definition TemplateMonad_Monad@{t u} : Monad@{t u} TemplateMonad@{t u} :=
  Eval hnf in TemplateMonad_UnoptimizedMonad.
Global Existing Instance TemplateMonad_Monad.

Polymorphic Definition tmDefinitionRed
: ident option reductionStrategy {A:Type}, A TemplateMonad A :=
  @tmDefinitionRed_ false.
Polymorphic Definition tmOpaqueDefinitionRed
: ident option reductionStrategy {A:Type}, A TemplateMonad A :=
  @tmDefinitionRed_ true.

Definition print_nf {A} (msg : A) : TemplateMonad unit
  := tmEval all msg >>= tmPrint.

Definition fail_nf {A} (msg : string) : TemplateMonad A
  := tmEval all msg >>= tmFail.

Definition tmMkInductive' (mind : mutual_inductive_body) : TemplateMonad unit
  := tmMkInductive false (mind_body_to_entry mind).

Definition tmAxiom id := tmAxiomRed id None.
Definition tmDefinition id {A} t := @tmDefinitionRed_ false id None A t.

We keep the original behaviour of tmQuoteRec: it quotes all the dependencies regardless of the opaqueness settings
Definition tmQuoteRec {A} (a : A) := tmQuoteRecTransp a true.

Definition tmLocate1 (q : qualid) : TemplateMonad global_reference :=
  l <- tmLocate q ;;
  match l with
  | []tmFail ("Constant [" ^ q ^ "] not found")
  | x :: _tmReturn x
  end.

Definition tmLocateModule1 (q : qualid) : TemplateMonad modpath :=
  l <- tmLocateModule q ;;
  match l with
  | []tmFail ("Module [" ^ q ^ "] not found")
  | x :: _tmReturn x
  end.

Definition tmLocateModType1 (q : qualid) : TemplateMonad modpath :=
  l <- tmLocateModType q ;;
  match l with
  | []tmFail ("ModType [" ^ q ^ "] not found")
  | x :: _tmReturn x
  end.

Don't remove. Constants used in the implem of the plugin
Definition tmTestQuote {A} (t : A) := tmQuote t >>= tmPrint.

Definition Common_kn (s : ident) :=
  (MPfile ["Common"; "TemplateMonad"; "Template"; "MetaRocq"], s).
Definition tmTestUnquote (t : term) :=
     t' <- tmUnquote t ;;
     t'' <- tmEval (unfold (Common_kn "my_projT2")) (my_projT2 t') ;;
     tmPrint t''.

Definition tmQuoteDefinition id {A} (t : A) := tmQuote t >>= tmDefinition id.
Definition tmQuoteDefinitionRed id rd {A} (t : A)
  := tmEval rd t >>= tmQuoteDefinition id.
Definition tmQuoteRecDefinition id {A} (t : A)
  := tmQuoteRec t >>= tmDefinition id.

Definition tmMkDefinition (id : ident) (tm : term) : TemplateMonad unit
  := t' <- tmUnquote tm ;;
     t'' <- tmEval (unfold (Common_kn "my_projT2")) (my_projT2 t') ;;
     tmDefinitionRed id (Some (unfold (Common_kn "my_projT1"))) t'' ;;
     tmReturn tt.

Definition TypeInstanceUnoptimized : Common.TMInstance :=
  {| Common.TemplateMonad := TemplateMonad
   ; Common.tmReturn:=@tmReturn
   ; Common.tmBind:=@tmBind
   ; Common.tmFail:=@tmFail
   ; Common.tmFreshName:=@tmFreshName
   ; Common.tmLocate:=@tmLocate
   ; Common.tmLocateModule:=@tmLocateModule
   ; Common.tmLocateModType:=@tmLocateModType
   ; Common.tmCurrentModPath:=@tmCurrentModPath
   ; Common.tmQuoteInductive:=@tmQuoteInductive
   ; Common.tmQuoteUniverses:=@tmQuoteUniverses
   ; Common.tmQuoteConstant:=@tmQuoteConstant
   ; Common.tmMkInductive:=@tmMkInductive
   ; Common.tmExistingInstance:=@tmExistingInstance
   |}.

Definition TypeInstanceOptimized : Common.TMInstance :=
  {| Common.TemplateMonad := TemplateMonad
   ; Common.tmReturn:=@tmReturn
   ; Common.tmBind:=@tmOptimizedBind
   ; Common.tmFail:=@tmFail
   ; Common.tmFreshName:=@tmFreshName
   ; Common.tmLocate:=@tmLocate
   ; Common.tmLocateModule:=@tmLocateModule
   ; Common.tmLocateModType:=@tmLocateModType
   ; Common.tmCurrentModPath:=@tmCurrentModPath
   ; Common.tmQuoteInductive:=@tmQuoteInductive
   ; Common.tmQuoteUniverses:=@tmQuoteUniverses
   ; Common.tmQuoteConstant:=@tmQuoteConstant
   ; Common.tmMkInductive:=@tmMkInductive
   ; Common.tmExistingInstance:=@tmExistingInstance
   |}.

Definition TypeInstance : Common.TMInstance :=
  Eval hnf in TypeInstanceUnoptimized.

Definition tmQuoteSort@{U t u} : TemplateMonad@{t u} sort
  := p <- @tmQuote Prop (Type@{U} True);;
     match p with
     | tProd _ (tSort s) _ret s
     | _tmFail "Anomaly: tmQuote (Type -> True) should be (tProd _ (tSort _) _)!"%bs
     end.
Definition tmQuoteUniverse@{U t u} : TemplateMonad@{t u} Universe.t
  := s <- @tmQuoteSort@{U t u};;
     match s with
     | sType uret u
     | _tmFail "Sort does not carry a universe (is not Type)"%bs
     end.
Definition tmQuoteLevel@{U t u} : TemplateMonad@{t u} Level.t
  := u <- tmQuoteUniverse@{U t u};;
     match Universe.get_is_level u with
     | Some lret l
     | NonetmFail "Universe is not a level"%bs
     end.

Definition tmFix'@{a b t u} {A : Type@{a}} {B : Type@{b}} (qtmFix' : Ast.term) (f : (A TemplateMonad@{t u} B) (A TemplateMonad@{t u} B)) : A TemplateMonad@{t u} B
  := f (fun a
        ⇒ tmFix <- tmUnquoteTyped (Ast.term ((A TemplateMonad@{t u} B) (A TemplateMonad@{t u} B)) A TemplateMonad@{t u} B) qtmFix';;
           tmFix qtmFix' f a).
Definition tmFix@{a b t u} {A : Type@{a}} {B : Type@{b}} (f : (A TemplateMonad@{t u} B) (A TemplateMonad@{t u} B)) : A TemplateMonad@{t u} B
  := f (fun a
        ⇒ (qA <- tmQuote A;;
            qB <- tmQuote B;;
            qa <- tmQuoteLevel@{a _ _};;
            qb <- tmQuoteLevel@{b _ _};;
            qt <- tmQuoteLevel@{t _ _};;
            qu <- tmQuoteLevel@{u _ _};;
            let self := tConst (MPfile ["Core"; "TemplateMonad"; "Template"; "MetaRocq"], "tmFix'")%bs [qa;qb;qt;qu] in
            @tmFix'@{a b t u} A B (mkApps self [qA; qB]) f a)).