feat: add Expr.fvarsSubset

src/Lean/Util.lean

@@ -34,3 +34,4 @@ import Lean.Util.SearchPath
 import Lean.Util.SafeExponentiation
 import Lean.Util.NumObjs
 import Lean.Util.NumApps
+import Lean.Util.FVarSubset

src/Lean/Util/CollectFVars.lean

@@ -26,14 +26,14 @@ mutual
     else main e { s with visitedExpr := s.visitedExpr.insert e }
 
   partial def main : Expr → Visitor
-    | Expr.proj _ _ e      => visit e
-    | Expr.forallE _ d b _ => visit b ∘ visit d
-    | Expr.lam _ d b _     => visit b ∘ visit d
-    | Expr.letE _ t v b _  => visit b ∘ visit v ∘ visit t
-    | Expr.app f a         => visit a ∘ visit f
-    | Expr.mdata _ b       => visit b
-    | Expr.fvar fvarId     => fun s => s.add fvarId
-    | _                    => id
+    | .proj _ _ e      => visit e
+    | .forallE _ d b _ => visit b ∘ visit d
+    | .lam _ d b _     => visit b ∘ visit d
+    | .letE _ t v b _  => visit b ∘ visit v ∘ visit t
+    | .app f a         => visit a ∘ visit f
+    | .mdata _ b       => visit b
+    | .fvar fvarId     => fun s => s.add fvarId
+    | _                => id
 end
 
 end CollectFVars

src/Lean/Util/FVarSubset.lean

@@ -0,0 +1,25 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+prelude
+import Lean.Util.CollectFVars
+import Lean.Util.FindExpr
+
+namespace Lean.Expr
+
+/-- Returns true if the free variables in `a` are subset of the free variables in `b`. -/
+def fvarsSubset (a b : Expr) : Bool :=
+  if !a.hasFVar then
+    true -- Empty set is subset of anything
+  else if !b.hasFVar then
+    false -- Empty set is not a superset of a nonempty set.
+  else
+    let s := collectFVars {} b
+    Option.isNone <| a.findExt? fun e =>
+      if !e.hasFVar then .done
+      else if e.isFVar && !s.fvarSet.contains e.fvarId! then .found
+      else .visit
+
+end Lean.Expr

src/Lean/Util/FindExpr.lean

@@ -14,25 +14,24 @@ opaque findImpl? (p : @& (Expr → Bool)) (e : @& Expr) : Option Expr
 
 @[inline] def find? (p : Expr → Bool) (e : Expr) : Option Expr := findImpl? p e
 
-/-- Return true if `e` occurs in `t` -/
+/-- Returns true if `e` occurs in `t` -/
 def occurs (e : Expr) (t : Expr) : Bool :=
   (t.find? fun s => s == e).isSome
 
 /--
-  Return type for `findExt?` function argument.
+Returns type for `findExt?` function argument.
 -/
 inductive FindStep where
-  /-- Found desired subterm -/ | found
-  /-- Search subterms -/ | visit
-  /-- Do not search subterms -/ | done
+  | /-- Found desired subterm -/ found
+  | /-- Search subterms -/ visit
+  | /-- Do not search subterms -/ done
 
 @[extern "lean_find_ext_expr"]
 opaque findExtImpl? (p : @& (Expr → FindStep)) (e : @& Expr) : Option Expr
 
 /--
-  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. -/
+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. -/
 @[inline] def findExt? (p : Expr → FindStep) (e : Expr) : Option Expr := findExtImpl? p e
 
-end Expr
-end Lean
+end Lean.Expr

src/Lean/Util/ForEachExpr.lean

@@ -20,13 +20,13 @@ def visit (g : Expr → m Bool) (e : Expr) : MonadCacheT Expr Unit m Unit :=
   checkCache e fun _ => do
     if (← g e) then
       match e with
-      | Expr.forallE _ d b _   => do visit g d; visit g b
-      | Expr.lam _ d b _       => do visit g d; visit g b
-      | Expr.letE _ t v b _    => do visit g t; visit g v; visit g b
-      | Expr.app f a           => do visit g f; visit g a
-      | Expr.mdata _ b         => visit g b
-      | Expr.proj _ _ b        => visit g b
-      | _                      => pure ()
+      | .forallE _ d b _   => do visit g d; visit g b
+      | .lam _ d b _       => do visit g d; visit g b
+      | .letE _ t v b _    => do visit g t; visit g v; visit g b
+      | .app f a           => do visit g f; visit g a
+      | .mdata _ b         => visit g b
+      | .proj _ _ b        => visit g b
+      | _                  => pure ()
 
 end ForEachExpr
 

tests/lean/run/fvarSubset.lean

@@ -0,0 +1,29 @@
+import Lean
+
+open Lean Meta
+run_meta do
+  let nat := mkConst ``Nat
+  withLocalDeclD `x nat fun x =>
+  withLocalDeclD `y nat fun y =>
+  withLocalDeclD `z nat fun z =>
+  withLocalDeclD `w nat fun w => do
+    let o := mkNatLit 0
+    let e1 ← mkAdd x x     -- x + x
+    let e2 ← mkAdd x y     -- x + y
+    let e3 ← mkAdd e2 z    -- (x + y) + z
+    let e4 ← mkAdd e2 e2   -- (x + y) + (x + y)
+    let e5 ← mkAdd e4 w    -- ((x + y) + (x + y)) + w
+    let e6 ← mkAdd e5 z    -- (((x + y) + (x + y)) + w) + z
+    assert! o.fvarsSubset e6
+    assert! !e6.fvarsSubset o
+    assert! o.fvarsSubset e1
+    assert! !e1.fvarsSubset o
+    assert! e1.fvarsSubset e1
+    assert! e1.fvarsSubset e2
+    assert! e6.fvarsSubset e6
+    assert! !e2.fvarsSubset e1
+    assert! e1.fvarsSubset e6
+    assert! e4.fvarsSubset e6
+    assert! !e3.fvarsSubset e5
+    assert! !e5.fvarsSubset e3
+    assert! e4.fvarsSubset e6