102 lines
3.3 KiB
Text
102 lines
3.3 KiB
Text
/-
|
|
Copyright (c) 2020 Microsoft Corporation. All rights reserved.
|
|
Released under Apache 2.0 license as described in the file LICENSE.
|
|
Author: Leonardo de Moura
|
|
-/
|
|
import Lean.Expr
|
|
|
|
namespace Lean
|
|
|
|
inductive HeadIndex where
|
|
| fvar (fvarId : FVarId)
|
|
| mvar (mvarId : MVarId)
|
|
| const (constName : Name)
|
|
| proj (structName : Name) (idx : Nat)
|
|
| lit (litVal : Literal)
|
|
| sort
|
|
| lam
|
|
| forallE
|
|
deriving Inhabited, BEq, Repr
|
|
|
|
namespace HeadIndex
|
|
|
|
protected def HeadIndex.hash : HeadIndex → UInt64
|
|
| fvar fvarId => mixHash 11 <| hash fvarId
|
|
| mvar mvarId => mixHash 13 <| hash mvarId
|
|
| const constName => mixHash 17 <| hash constName
|
|
| proj structName idx => mixHash 19 <| mixHash (hash structName) (hash idx)
|
|
| lit litVal => mixHash 23 <| hash litVal
|
|
| sort => 29
|
|
| lam => 31
|
|
| forallE => 37
|
|
|
|
instance : Hashable HeadIndex := ⟨HeadIndex.hash⟩
|
|
|
|
end HeadIndex
|
|
|
|
namespace Expr
|
|
|
|
def head : Expr → Expr
|
|
| app f _ _ => head f
|
|
| letE _ _ _ b _ => head b
|
|
| mdata _ e _ => head e
|
|
| e => e
|
|
|
|
private def headNumArgsAux : Expr → Nat → Nat
|
|
| app f _ _, n => headNumArgsAux f (n + 1)
|
|
| letE _ _ _ b _, n => headNumArgsAux b n
|
|
| mdata _ e _, n => headNumArgsAux e n
|
|
| _, n => n
|
|
|
|
def headNumArgs (e : Expr) : Nat :=
|
|
headNumArgsAux e 0
|
|
|
|
/-
|
|
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
|
|
| proj structName idx _ _ => HeadIndex.proj structName idx
|
|
| sort _ _ => HeadIndex.sort
|
|
| lam _ _ _ _ => HeadIndex.lam
|
|
| forallE _ _ _ _ => HeadIndex.forallE
|
|
| lit v _ => HeadIndex.lit v
|
|
| 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
|