From cdac0b6615ca8157f7c426e84e421254410753d6 Mon Sep 17 00:00:00 2001
From: Jason Gross <jgross@mit.edu>
Date: Sat, 8 Apr 2023 15:15:11 -0700
Subject: [PATCH] Optimize `tmBind`

Fixes #441
---
 template-coq/theories/TemplateMonad/Core.v | 29 +++++++++++++++++++---
 test-suite/bug441.v                        |  2 --
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/template-coq/theories/TemplateMonad/Core.v b/template-coq/theories/TemplateMonad/Core.v
index bcb69b52a..59397aeae 100644
--- a/template-coq/theories/TemplateMonad/Core.v
+++ b/template-coq/theories/TemplateMonad/Core.v
@@ -83,13 +83,16 @@ Definition tmOptimize@{t u} {A : Type@{t}} (v : TemplateMonad@{t u} A) : Templat
   := tmOptimizedBind v tmReturn.
 
 (** This allow to use notations of MonadNotation *)
-Global Instance TemplateMonad_Monad@{t u} : Monad@{t u} TemplateMonad@{t u} :=
+Definition TemplateMonad_UnoptimizedMonad@{t u} : Monad@{t u} TemplateMonad@{t u} :=
   {| ret := @tmReturn ; bind := @tmBind |}.
 
-(* We don't want to make this an instance, becuase it blows up performance in some cases *)
 Definition TemplateMonad_OptimizedMonad@{t u} : Monad@{t u} TemplateMonad@{t u} :=
   {| ret := @tmReturn ; bind := @tmOptimizedBind |}.
 
+Definition TemplateMonad_Monad@{t u} : Monad@{t u} TemplateMonad@{t u} :=
+  Eval hnf in TemplateMonad_OptimizedMonad.
+Global Existing Instance TemplateMonad_Monad.
+
 Polymorphic Definition tmDefinitionRed
 : ident -> option reductionStrategy -> forall {A:Type}, A -> TemplateMonad A :=
   @tmDefinitionRed_ false.
@@ -155,7 +158,7 @@ Definition tmMkDefinition (id : ident) (tm : term) : TemplateMonad unit
      tmDefinitionRed id (Some (unfold (Common_kn "my_projT1"))) t'' ;;
      tmReturn tt.
 
-Definition TypeInstance : Common.TMInstance :=
+Definition TypeInstanceUnoptimized : Common.TMInstance :=
   {| Common.TemplateMonad := TemplateMonad
    ; Common.tmReturn:=@tmReturn
    ; Common.tmBind:=@tmBind
@@ -172,6 +175,26 @@ Definition TypeInstance : Common.TMInstance :=
    ; 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 TypeInstanceOptimized.
+
 Definition tmQuoteUniverse@{U t u} : TemplateMonad@{t u} Universe.t
   := u <- @tmQuote Prop (Type@{U} -> True);;
      match u with
diff --git a/test-suite/bug441.v b/test-suite/bug441.v
index 623edbf37..27f86d922 100644
--- a/test-suite/bug441.v
+++ b/test-suite/bug441.v
@@ -1,8 +1,6 @@
 From MetaCoq Require Import Template.All.
 Import MCMonadNotation.
 
-#[local] Existing Instance TemplateMonad_OptimizedMonad.
-
 Fixpoint tm_double n : TemplateMonad nat :=
   match n with
   | 0 => ret 0