From 55bbaa55d8a1b32c092534e144a31d37dcaece3e Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Fri, 22 Oct 2021 14:54:01 -0700
Subject: [PATCH] fix: `toHeadIndex`

---
 src/Lean/HeadIndex.lean | 41 +++++++++++++++++++++++++++++++++++++----
 1 file changed, 37 insertions(+), 4 deletions(-)

diff --git a/src/Lean/HeadIndex.lean b/src/Lean/HeadIndex.lean
index 6fa3e28d83..8b22dac104 100644
--- a/src/Lean/HeadIndex.lean
+++ b/src/Lean/HeadIndex.lean
@@ -51,7 +51,15 @@ private def headNumArgsAux : Expr → Nat → Nat
 def headNumArgs (e : Expr) : Nat :=
   headNumArgsAux e 0
 
-def toHeadIndex : Expr → HeadIndex
+/-
+  Quick version that may fail if it "hits" a loose bound variable.
+  This can happen, for example, if the input expression is of the form.
+  ```
+  let f := fun x => x + 1;
+  f 0
+  ```
+-/
+private def toHeadIndexQuick? : Expr → Option HeadIndex
   | mvar mvarId _           => HeadIndex.mvar mvarId
   | fvar fvarId _           => HeadIndex.fvar fvarId
   | const constName _ _     => HeadIndex.const constName
@@ -60,10 +68,35 @@ def toHeadIndex : Expr → HeadIndex
   | lam _ _ _ _             => HeadIndex.lam
   | forallE _ _ _ _         => HeadIndex.forallE
   | lit v _                 => HeadIndex.lit v
-  | app f _ _               => toHeadIndex f
-  | letE _ _ _ b _          => toHeadIndex b
-  | mdata _ e _             => toHeadIndex e
+  | app f _ _               => toHeadIndexQuick? f
+  | letE _ _ _ b _          => toHeadIndexQuick? b
+  | mdata _ e _             => toHeadIndexQuick? e
+  | _                       => none
+
+/-
+  Slower version of `toHeadIndexQuick?` that "expands" let-declarations to make
+  sure we never hit a loose bound variable.
+  The performance of the `letE` alternative can be improved, but this function should not be in the hotpath
+  since `toHeadIndexQuick?` succeeds most of the time.
+-/
+private partial def toHeadIndexSlow : Expr → HeadIndex
+  | mvar mvarId _           => HeadIndex.mvar mvarId
+  | fvar fvarId _           => HeadIndex.fvar fvarId
+  | const constName _ _     => HeadIndex.const constName
+  | proj structName idx _ _ => HeadIndex.proj structName idx
+  | sort _ _                => HeadIndex.sort
+  | lam _ _ _ _             => HeadIndex.lam
+  | forallE _ _ _ _         => HeadIndex.forallE
+  | lit v _                 => HeadIndex.lit v
+  | app f _ _               => toHeadIndexSlow f
+  | letE _ _ v b _          => toHeadIndexSlow (b.instantiate1 v)
+  | mdata _ e _             => toHeadIndexSlow e
   | _                       => panic! "unexpected expression kind"
 
+def toHeadIndex (e : Expr) : HeadIndex :=
+  match toHeadIndexQuick? e with
+  | some i => i
+  | none => toHeadIndexSlow e
+
 end Expr
 end Lean