perf: find? and findExt? (#4795)

use the kernel implementation.
This commit is contained in:
Leonardo de Moura 2024-07-20 03:13:54 +02:00 committed by GitHub
parent 327986e6fb
commit de5e07c4d2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -9,53 +9,11 @@ import Lean.Util.PtrSet
namespace Lean
namespace Expr
namespace FindImpl
unsafe abbrev FindM := StateT (PtrSet Expr) Id
@[inline] unsafe def checkVisited (e : Expr) : OptionT FindM Unit := do
if (← get).contains e then
failure
modify fun s => s.insert e
unsafe def findM? (p : Expr → Bool) (e : Expr) : OptionT FindM Expr :=
let rec visit (e : Expr) := do
checkVisited e
if p e then
pure e
else match e with
| .forallE _ d b _ => visit d <|> visit b
| .lam _ d b _ => visit d <|> visit b
| .mdata _ b => visit b
| .letE _ t v b _ => visit t <|> visit v <|> visit b
| .app f a => visit f <|> visit a
| .proj _ _ b => visit b
| _ => failure
visit e
unsafe def findUnsafe? (p : Expr → Bool) (e : Expr) : Option Expr :=
Id.run <| findM? p e |>.run' mkPtrSet
end FindImpl
-- TODO: replace `find?` with this one after update-stage0
@[extern "lean_find_expr"]
opaque findImpl? (p : @& (Expr → Bool)) (e : @& Expr) : Option Expr
@[implemented_by FindImpl.findUnsafe?]
def find? (p : Expr → Bool) (e : Expr) : Option Expr :=
/- This is a reference implementation for the unsafe one above -/
if p e then
some e
else match e with
| .forallE _ d b _ => find? p d <|> find? p b
| .lam _ d b _ => find? p d <|> find? p b
| .mdata _ b => find? p b
| .letE _ t v b _ => find? p t <|> find? p v <|> find? p b
| .app f a => find? p f <|> find? p a
| .proj _ _ b => find? p b
| _ => none
@[inline] def find? (p : Expr → Bool) (e : Expr) : Option Expr := findImpl? p e
/-- Return true if `e` occurs in `t` -/
def occurs (e : Expr) (t : Expr) : Bool :=
@ -69,45 +27,13 @@ inductive FindStep where
/-- Search subterms -/ | visit
/-- Do not search subterms -/ | done
-- TODO: replace `findExt?` with this one after update-stage0
@[extern "lean_find_ext_expr"]
opaque findExtImpl? (p : @& (Expr → FindStep)) (e : @& Expr) : Option Expr
namespace FindExtImpl
unsafe def findM? (p : Expr → FindStep) (e : Expr) : OptionT FindImpl.FindM Expr :=
visit e
where
visitApp (e : Expr) :=
match e with
| .app f a .. => visitApp f <|> visit a
| e => visit e
visit (e : Expr) := do
FindImpl.checkVisited e
match p e with
| .done => failure
| .found => pure e
| .visit =>
match e with
| .forallE _ d b _ => visit d <|> visit b
| .lam _ d b _ => visit d <|> visit b
| .mdata _ b => visit b
| .letE _ t v b _ => visit t <|> visit v <|> visit b
| .app .. => visitApp e
| .proj _ _ b => visit b
| _ => failure
unsafe def findUnsafe? (p : Expr → FindStep) (e : Expr) : Option Expr :=
Id.run <| findM? p e |>.run' mkPtrSet
end FindExtImpl
/--
Similar to `find?`, but `p` can return `FindStep.done` to interrupt the search on subterms.
Remark: Differently from `find?`, we do not invoke `p` for partial applications of an application. -/
@[implemented_by FindExtImpl.findUnsafe?]
opaque findExt? (p : Expr → FindStep) (e : Expr) : Option Expr
@[inline] def findExt? (p : Expr → FindStep) (e : Expr) : Option Expr := findExtImpl? p e
end Expr
end Lean