From de5e07c4d2d0b48dce5cefdb629bc7171134b490 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Sat, 20 Jul 2024 03:13:54 +0200
Subject: [PATCH] perf: `find?` and `findExt?` (#4795)

use the kernel implementation.
---
 src/Lean/Util/FindExpr.lean | 78 +------------------------------------
 1 file changed, 2 insertions(+), 76 deletions(-)

diff --git a/src/Lean/Util/FindExpr.lean b/src/Lean/Util/FindExpr.lean
index caf92c778db9..32795ec30ca5 100644
--- a/src/Lean/Util/FindExpr.lean
+++ b/src/Lean/Util/FindExpr.lean
@@ -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