perf: Replacement.apply (#4776)

Avoid potentially expensive `e.replace` if it is not applicable.
This commit is contained in:
Leonardo de Moura 2024-07-17 18:17:47 +02:00 committed by GitHub
parent 933445608c
commit 41b4914836
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 26 additions and 5 deletions

View file

@ -80,6 +80,10 @@ protected def max : RBNode α β → Option (Sigma (fun k => β k))
def singleton (k : α) (v : β k) : RBNode α β :=
node red leaf k v leaf
def isSingleton : RBNode α β → Bool
| node _ leaf _ _ leaf => true
| _ => false
-- the first half of Okasaki's `balance`, concerning red-red sequences in the left child
@[inline] def balance1 : RBNode α β → (a : α) → β a → RBNode α β → RBNode α β
| node red (node red a kx vx b) ky vy c, kz, vz, d
@ -269,6 +273,9 @@ variable {α : Type u} {β : Type v} {σ : Type w} {cmp : αα → Ordering
def depth (f : Nat → Nat → Nat) (t : RBMap α β cmp) : Nat :=
t.val.depth f
def isSingleton (t : RBMap α β cmp) : Bool :=
t.val.isSingleton
@[inline] def fold (f : σα → β → σ) : (init : σ) → RBMap α β cmp → σ
| b, ⟨t, _⟩ => t.fold f b

View file

@ -728,12 +728,26 @@ def insertReplacementForLetRecs (r : Replacement) (letRecClosures : List LetRecC
letRecClosures.foldl (init := r) fun r c =>
r.insert c.toLift.fvarId c.closed
def isApplicable (r : Replacement) (e : Expr) : Bool :=
Option.isSome <| e.findExt? fun e =>
if e.hasFVar then
match e with
| .fvar fvarId => if r.contains fvarId then .found else .done
| _ => .visit
else
.done
def Replacement.apply (r : Replacement) (e : Expr) : Expr :=
e.replace fun e => match e with
| .fvar fvarId => match r.find? fvarId with
| some c => some c
| _ => none
| _ => none
-- Remark: if `r` is not a singlenton, then declaration is using `mutual` or `let rec`,
-- and there is a big chance `isApplicable r e` is true.
if r.isSingleton && !isApplicable r e then
e
else
e.replace fun e => match e with
| .fvar fvarId => match r.find? fvarId with
| some c => some c
| _ => none
| _ => none
def pushMain (preDefs : Array PreDefinition) (sectionVars : Array Expr) (mainHeaders : Array DefViewElabHeader) (mainVals : Array Expr)
: TermElabM (Array PreDefinition) :=