From 487c2a937a892e177ed0248fc6d8cc01d6f37a89 Mon Sep 17 00:00:00 2001
From: Kim Morrison <scott.morrison@gmail.com>
Date: Wed, 16 Oct 2024 14:07:34 +1100
Subject: [PATCH] feat: Expr helper functions (#5729)

`getNumHeadForalls` and `getNumHeadLambdas` were both duplicated
downstream with different names; I'll clean up those next.

Also adds `getAppNumArgs'`.
---
 src/Lean/Expr.lean | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/src/Lean/Expr.lean b/src/Lean/Expr.lean
index b2122c5ba3..a119f49b6d 100644
--- a/src/Lean/Expr.lean
+++ b/src/Lean/Expr.lean
@@ -1038,6 +1038,14 @@ def getForallBinderNames : Expr → List Name
   | forallE n _ b _ => n :: getForallBinderNames b
   | _ => []
 
+/--
+Returns the number of leading `∀` binders of an expression. Ignores metadata.
+-/
+def getNumHeadForalls : Expr → Nat
+  | mdata _ b => getNumHeadForalls b
+  | forallE _ _ body _ => getNumHeadForalls body + 1
+  | _ => 0
+
 /--
 If the given expression is a sequence of
 function applications `f a₁ .. aₙ`, return `f`.
@@ -1085,6 +1093,16 @@ private def getAppNumArgsAux : Expr → Nat → Nat
 def getAppNumArgs (e : Expr) : Nat :=
   getAppNumArgsAux e 0
 
+/-- Like `getAppNumArgs` but ignores metadata. -/
+def getAppNumArgs' (e : Expr) : Nat :=
+  go e 0
+where
+  /-- Auxiliary definition for `getAppNumArgs'`. -/
+  go : Expr → Nat → Nat
+    | mdata _ b, n => go b n
+    | app f _  , n => go f (n + 1)
+    | _        , n => n
+
 /--
 Like `Lean.Expr.getAppFn` but assumes the application has up to `maxArgs` arguments.
 If there are any more arguments than this, then they are returned by `getAppFn` as part of the function.