fix(Translate): don't look in arguments of free variables in shouldTranslate (leanprover-community#35549)

This PR addresses a `to_additive` limitation reported at [#mathlib4 > leanprover-community#35179 breaks to_additive call in FLT](https://leanprover.zulipchat.com/#narrow/channel/287929-mathlib4/topic/.2335179.20breaks.20to_additive.20call.20in.20FLT/with/574823087) by making sure `shouldTranslate` doesn't looks at the arguments of free variable applications

This PR refactors `shouldTranslate` in a few ways
- The monad is changed from `OptionT` to `ExceptT`, with the underlying type changing from `Option Expr` to `ExceptT Expr Unit`. Although this is technically a bit less efficient, it makes the program shorter and more intuitive, because a failure now actually corresponds to a monadic failure instead of a monadic success.
- When visiting a `Expr.app` node, we now loop through the whole application at once, instead of doing this step by step and checking the cache at each step. This allows us to detect when the head is a free variable, and skip its arguments. It also lets us deal with `ignoreArgsAttr` more naturally.
- This change removes the need for the `inApp : Bool` argument that used to be passed around in order to tell whether a constant was in an application.

Author: JovanGerb

Mathlib/Tactic/Translate/Core.lean

@@ -350,56 +350,53 @@ def ReplacementM.run {α} (dontTranslate allFVars : Array FVarId) (x : Replaceme
   return (a, allFVars.findIdx? relevantFVars.contains)
 
 /-- Implementation function for `shouldTranslate`.
-Failure means that in that subexpression there is no constant that blocks `e` from being translated.
+Returning `none` means that `e` contains no constant that blocks translation.
 We cache previous applications of the function, using an expression cache using ptr equality
-to avoid visiting the same subexpression many times. Note that we only need to cache the
-expressions without taking the value of `inApp` into account, since `inApp` only matters when
-the expression is a constant. However, for this reason we have to make sure that we never
-cache constant expressions, so that's why the `if`s in the implementation are in this order.
+to avoid visiting the same subexpression many times.
 
 Note that this function is still called many times by `applyReplacementFun`
 and we're not remembering the cache between these calls. -/
 private unsafe def shouldTranslateUnsafe (env : Environment) (t : TranslateData) (e : Expr) :
     ReplacementM (Option Expr) := do
   let visitedFVars : IO.Ref (Array FVarId) ← IO.mkRef #[]
   let dontTranslate ← read
-  let rec visit (e : Expr) (inApp := false) : OptionT (StateT (PtrSet Expr) BaseIO) Expr := do
-    if e.isConst then
-      let doTranslate :=
-        (t.doTranslateAttr.find? env e.constName!).getD <|
-          inApp || (findTranslation? env t e.constName).isSome
-      if doTranslate then failure else return e
+  let rec visit (e : Expr) : ExceptT Expr (StateT (PtrSet Expr) BaseIO) Unit := do
     if (← get).contains e then
-      failure
+      return
     modify fun s => s.insert e
     match e with
-    | .app e a =>
-        visit e true <|> do
-          if let some n := e.getAppFn.constName? then
-            if let some l := t.ignoreArgsAttr.find? env n then
-              if e.getAppNumArgs + 1 ∈ l then
-                failure
-          visit a
+    | .app .. => e.withApp fun f args ↦ do
+      match f with
+      | .const n _ =>
+        -- A constant in an application, e.g. `Prod` in `α × β`, is translated by default.
+        let doTranslate := (t.doTranslateAttr.find? env n).getD true
+        unless doTranslate do throw e
+        let l := (t.ignoreArgsAttr.find? env n).getD []
+        args.size.forM fun i _ ↦ do
+          if !l.contains i then visit args[i]
+      | .fvar .. => visit f -- We don't look in the arguments of free variables.
+      | _ => visit f; args.forM visit
+    | .const n _ =>
+      -- A constant not in an application, e.g. `ℕ`, is not translated by default.
+      let doTranslate := (t.doTranslateAttr.find? env n).getD (findTranslation? env t n).isSome
+      unless doTranslate do throw e
     | .lam _ _ t _       => visit t
     | .forallE _ _ t _   => visit t
-    | .letE _ _ e body _ => visit e <|> visit body
+    | .letE _ _ e body _ => visit e; visit body
     | .mdata _ b         => visit b
     | .proj _ _ b        => visit b
     | .fvar fvarId       =>
       if dontTranslate.contains fvarId then
-        return e
-      else
-        visitedFVars.modify (·.push fvarId)
-        failure
+        throw e
+      visitedFVars.modify (·.push fvarId)
     /- We do not translate the order on `Prop`.
     TODO: We also don't want to translate the category on `Type u`. Unfortunately, replacing
     `.sort 0` with `.sort _` here breaks some uses of `to_additive` on `MonCat`. -/
-    | .sort 0            => return e
-    | _                  => failure
-  let x ← (visit e).run' mkPtrSet
-  match x with
-  | some e => return some e
-  | none =>
+    | .sort 0            => throw e
+    | _                  => pure ()
+  match ← (visit e).run' mkPtrSet with
+  | .error e => return some e
+  | .ok () =>
     /- In the case that we do translate, we mark the visited free variables as relevant for
     the translation by inserting them into the state. -/
     modify (·.insertMany (← visitedFVars.get))

Mathlib/Tactic/Translate/ToAdditive.lean

@@ -257,10 +257,10 @@ initialize ignoreArgsAttr : NameMapExtension (List Nat) ←
     descr :=
       "Auxiliary attribute for `to_additive` stating that certain arguments are not additivized."
     add := fun _ stx ↦ do
-        let ids ← match stx with
-          | `(attr| to_additive_ignore_args $[$ids:num]*) => pure <| ids.map (·.1.isNatLit?.get!)
-          | _ => throwUnsupportedSyntax
-        return ids.toList }
+      let ids ← match stx with
+        | `(attr| to_additive_ignore_args $[$ids:num]*) => pure <| ids.map (·.getNat - 1)
+        | _ => throwUnsupportedSyntax
+      return ids.toList }
 
 @[inherit_doc TranslateData.doTranslateAttr]
 initialize doTranslateAttr : NameMapExtension Bool ← registerNameMapExtension _

Mathlib/Tactic/Translate/ToDual.lean

@@ -112,11 +112,11 @@ initialize ignoreArgsAttr : NameMapExtension (List Nat) ←
     name  := `to_dual_ignore_args
     descr :=
       "Auxiliary attribute for `to_dual` stating that certain arguments are not dualized."
-    add   := fun _ stx ↦ do
-        let ids ← match stx with
-          | `(attr| to_dual_ignore_args $[$ids:num]*) => pure <| ids.map (·.1.isNatLit?.get!)
-          | _ => throwUnsupportedSyntax
-        return ids.toList }
+    add := fun _ stx ↦ do
+      let ids ← match stx with
+        | `(attr| to_dual_ignore_args $[$ids:num]*) => pure <| ids.map (·.getNat - 1)
+        | _ => throwUnsupportedSyntax
+      return ids.toList }
 
 @[inherit_doc TranslateData.unfoldBoundaries?]
 initialize unfoldBoundaries : UnfoldBoundaryExt ← registerUnfoldBoundaryExt

MathlibTest/toAdditive.lean

@@ -863,3 +863,11 @@ attribute [to_additive someOtherTranslation] abstractMul
 -- Test that we don't blindly translate the prefix of a name.
 def Mul.test : Nat := 5
 @[to_additive] def Mul.test' := Mul.test
+
+-- Test that arguments of free variables aren't considered by `shouldTranslate`
+@[to_additive_dont_translate]
+def dontTranslateId {α} : α → α := id
+
+@[to_additive]
+theorem functionTypeMonoid {ι : Type*} {R : ι → Type*} [(i : ι) → Monoid (R i)] (i : ι)
+  (a : R (dontTranslateId i)) : a * a = a * a := rfl