Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

perf: find? and findExt? #4795

Merged
merged 1 commit into from
Jul 20, 2024
Merged

perf: find? and findExt? #4795

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 2 additions & 76 deletions src/Lean/Util/FindExpr.lean
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 :=
Expand All @@ -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
Loading