Optimize tmBind

Fixes MetaCoq#441

template-coq/theories/TemplateMonad/Core.v

@@ -4,6 +4,7 @@ From MetaCoq.Template Require Import Ast AstUtils Common.
 
 Local Set Universe Polymorphism.
 Local Unset Universe Minimization ToSet.
+Local Unset Asymmetric Patterns.
 Import MCMonadNotation.
 
 (** * The Template Monad
@@ -68,9 +69,17 @@ Cumulative Inductive TemplateMonad@{t u} : Type@{t} -> Prop :=
 | tmInferInstance : option reductionStrategy -> forall A : Type@{t}, TemplateMonad (option_instance A)
 .
 
+Fixpoint tmOptimizedBind@{t u} {A B : Type@{t}} (v : TemplateMonad A) : (A -> TemplateMonad@{t u} B) -> TemplateMonad@{t u} B
+  := match v with
+     | tmReturn x => fun f => f x
+     | tmBind v k => fun f => tmOptimizedBind v (fun v => tmOptimizedBind (k v) f)
+     | tmFail msg => fun _ => tmFail msg
+     | v => tmBind v
+     end.
+
 (** This allow to use notations of MonadNotation *)
 Global Instance TemplateMonad_Monad@{t u} : Monad@{t u} TemplateMonad@{t u} :=
-  {| ret := @tmReturn ; bind := @tmBind |}.
+  {| ret := @tmReturn ; bind := @tmOptimizedBind |}.
 
 Polymorphic Definition tmDefinitionRed
 : ident -> option reductionStrategy -> forall {A:Type}, A -> TemplateMonad A :=
@@ -140,7 +149,7 @@ Definition tmMkDefinition (id : ident) (tm : term) : TemplateMonad unit
 Definition TypeInstance : Common.TMInstance :=
   {| Common.TemplateMonad := TemplateMonad
    ; Common.tmReturn:=@tmReturn
-   ; Common.tmBind:=@tmBind
+   ; Common.tmBind:=@tmOptimizedBind
    ; Common.tmFail:=@tmFail
    ; Common.tmFreshName:=@tmFreshName
    ; Common.tmLocate:=@tmLocate

test-suite/_CoqProject.in

@@ -8,6 +8,7 @@ TypeAnnotationTests.v
 bug1.v
 bug2.v
 bug369.v
+bug441.v
 bug5.v
 bug6.v
 bug7.v

test-suite/bug441.v

@@ -0,0 +1,17 @@
+From MetaCoq Require Import Template.All.
+Import MCMonadNotation.
+
+Fixpoint tm_double n : TemplateMonad nat :=
+  match n with
+  | 0 => ret 0
+  | S n =>
+      n' <- tm_double n;;
+      ret (S (S n'))
+  end.
+
+(* these should all run in under 0.2s; previously the final one took ~30s *)
+Timeout 5 Time MetaCoq Run (tmPrint =<< tm_double 1000).
+Timeout 5 Time MetaCoq Run (tmPrint =<< tm_double 2000).
+Timeout 5 Time MetaCoq Run (tmPrint =<< tm_double 3000).
+Timeout 5 Time MetaCoq Run (tmPrint =<< tm_double 4000).
+Timeout 5 Time MetaCoq Run (tmPrint =<< tm_double 5000).