Generalize count_lets

1 files changed, 20 insertions, 8 deletions

diff --git a/src/Reflection/CountLets.v b/src/Reflection/CountLets.v
index 43f152da1..00673ae34 100644
--- a/src/Reflection/CountLets.v
+++ b/src/Reflection/CountLets.v
@@ -27,29 +27,41 @@ Section language.
     Local Notation expr := (@expr base_type_code interp_base_type op var).
 
     Section gen.
-      Context (count_type : flat_type -> nat).
+      Context (count_type_let : flat_type -> nat).
+      Context (count_type_abs : base_type_code -> nat).
 
       Fixpoint count_lets_genf {t} (e : exprf t) : nat
         := match e with
            | LetIn tx _ _ eC
-             => count_type tx + @count_lets_genf _ (eC (SmartVal var mkVar tx))
+             => count_type_let tx + @count_lets_genf _ (eC (SmartVal var mkVar tx))
            | _ => 0
            end.
       Fixpoint count_lets_gen {t} (e : expr t) : nat
         := match e with
            | Return _ x
              => count_lets_genf x
-           | Abs tx _ f => @count_lets_gen _ (f (mkVar tx))
+           | Abs tx _ f => count_type_abs tx + @count_lets_gen _ (f (mkVar tx))
            end.
     End gen.
 
-    Definition count_letsf {t} (e : exprf t) : nat
+    Definition count_let_bindersf {t} (e : exprf t) : nat
       := count_lets_genf count_pairs e.
+    Definition count_letsf {t} (e : exprf t) : nat
+      := count_lets_genf (fun _ => 1) e.
+    Definition count_let_binders {t} (e : expr t) : nat
+      := count_lets_gen count_pairs (fun _ => 0) e.
     Definition count_lets {t} (e : expr t) : nat
-      := count_lets_gen count_pairs e.
+      := count_lets_gen (fun _ => 1) (fun _ => 0) e.
+    Definition count_binders {t} (e : expr t) : nat
+      := count_lets_gen count_pairs (fun _ => 1) e.
   End with_var.
 
-  Definition CountLetsGen (count_type : flat_type -> nat) {t} (e : Expr t) : nat
-    := count_lets_gen (fun _ => tt) count_type (e _).
-  Definition CountLets {t} (e : Expr t) : nat := CountLetsGen count_pairs e.
+  Definition CountLetsGen (count_type_let : flat_type -> nat) (count_type_abs : base_type_code -> nat) {t} (e : Expr t) : nat
+    := count_lets_gen (fun _ => tt) count_type_let count_type_abs (e _).
+  Definition CountLetBinders {t} (e : Expr t) : nat
+    := count_let_binders (fun _ => tt) (e _).
+  Definition CountLets {t} (e : Expr t) : nat
+    := count_lets (fun _ => tt) (e _).
+  Definition CountBinders {t} (e : Expr t) : nat
+    := count_binders (fun _ => tt) (e _).
 End language.