fix(to_dual): also store reorder for the reverse translation (#31716)

This PR fixes the bug in `to_dual` that `@[to_dual]` forgets to add the `reorder := ...` flag to the reverse translation. The correct behaviour is that it adds the inverse of the `reorder` permutation for the reverse translation.

To fix this, I refactor `insertTranslation`/`insertTranslationAndInfo`. Since `insertTranslationAndInfo` with an empty `ArgInfo` does the same as `insertTranslation`, I decided to merge these two into one function `insertTranslation`.

I also merged the `reorderAttr` and `relevantArgAttr` environment extensions into a single `argInfoAttr` environment extension, because this is more convenient.

I add a new `ToDual` test file which includes a test for this fix.

I replace the implementation of `insert_to_additive_translation` with something more principled, instead of using `insertTranslation`.

Author: JovanGerb

Mathlib/Tactic/Translate/Core.lean

@@ -135,17 +135,22 @@ initialize changeNumeralAttr : NameMapExtension (List Nat) ←
       pure <| arg.map (·.1.isNatLit?.get!.pred) |>.toList
     | _, _ => throwUnsupportedSyntax }
 
+/-- `ArgInfo` stores information about how a constant should be translated. -/
+structure ArgInfo where
+  /-- The arguments that should be reordered when translating, using cycle notation. -/
+  reorder : List (List Nat) := []
+  /-- The argument used to determine whether this constant should be translated. -/
+  relevantArg : Nat := 0
+
 /-- `TranslateData` is a structure that holds all data required for a translation attribute. -/
 structure TranslateData : Type where
   /-- An attribute that tells that certain arguments of this definition are not
   involved when translating.
   This helps the translation heuristic by also transforming definitions if `ℕ` or another
   fixed type occurs as one of these arguments. -/
   ignoreArgsAttr : NameMapExtension (List Nat)
-  /-- `reorderAttr` stores the declarations that need their arguments reordered when translating.
-  This is specified using the `(reorder := ...)` syntax. -/
-  reorderAttr : NameMapExtension (List <| List Nat)
-  relevantArgAttr : NameMapExtension Nat
+  /-- `argInfoAttr` stores the declarations that need some extra information to be translated. -/
+  argInfoAttr : NameMapExtension ArgInfo
   /-- The global `dont_translate` attribute specifies that operations on the given type
   should not be translated. This can be either for types that are translated,
   such as `MonoidAlgebra` -> `AddMonoidAlgebra`, or for fixed types, such as `Fin n`/`ZMod n`.
@@ -165,7 +170,6 @@ structure TranslateData : Type where
   isDual : Bool
   guessNameData : GuessName.GuessNameData
 
-attribute [inherit_doc relevantArgOption] TranslateData.relevantArgAttr
 attribute [inherit_doc GuessName.GuessNameData] TranslateData.guessNameData
 
 /-- Get the translation for the given name. -/
@@ -178,44 +182,36 @@ This allows translating automatically generated declarations such as `IsRegular.
 def findPrefixTranslation (env : Environment) (nm : Name) (t : TranslateData) : Name :=
   nm.mapPrefix (findTranslation? env t)
 
-/-- Add a name translation to the translations map. -/
-def insertTranslation (t : TranslateData) (src tgt : Name) (failIfExists := true) : CoreM Unit := do
-  if let some tgt' := findTranslation? (← getEnv) t src then
-    if failIfExists then
-      throwError "The translation {src} ↦ {tgt'} already exists"
-    else
-      trace[translate] "The translation {src} ↦ {tgt'} already exists"
-      return
-  modifyEnv (t.translations.addEntry · (src, tgt))
-  trace[translate] "Added translation {src} ↦ {tgt}"
-  -- For an attribute like `to_dual`, we also insert the reverse direction of the translation
+/-- Compute the `ArgInfo` for the reverse translation. The `reorder` permutation is inverted.
+In practice, `relevantArg` does not overlap with `reorder` for dual translations,
+so we don't bother applying the permutation to `relevantArg`. -/
+def ArgInfo.reverse (i : ArgInfo) : ArgInfo where
+  reorder := i.reorder.map (·.reverse)
+  relevantArg := i.relevantArg
+
+/-- Add a name translation to the translations map and add the `argInfo` information to `src`.
+If the translation attribute is dual, also add the reverse translation. -/
+def insertTranslation (t : TranslateData) (src tgt : Name) (argInfo : ArgInfo)
+    (failIfExists := true) : CoreM Unit := do
+  insertTranslationAux t src tgt failIfExists argInfo
   if t.isDual && src != tgt then
-    if let some src' := findTranslation? (← getEnv) t tgt then
+    insertTranslationAux t tgt src failIfExists argInfo.reverse
+where
+  /-- Insert only one direction of a translation. -/
+  insertTranslationAux (t : TranslateData) (src tgt : Name) (failIfExists : Bool)
+      (argInfo : ArgInfo) : CoreM Unit := do
+    if let some tgt' := findTranslation? (← getEnv) t src then
       if failIfExists then
-        throwError "The translation {tgt} ↦ {src'} already exists"
+        throwError "The translation {src} ↦ {tgt'} already exists"
       else
-        trace[translate] "The translation {tgt} ↦ {src'} already exists"
-        return
-    modifyEnv (t.translations.addEntry · (tgt, src))
-    trace[translate] "Also added translation {tgt} ↦ {src}"
-
-/-- `ArgInfo` stores information about how a constant should be translated. -/
-structure ArgInfo where
-  /-- The arguments that should be reordered when translating, using cycle notation. -/
-  reorder : List (List Nat) := []
-  /-- The argument used to determine whether this constant should be translated. -/
-  relevantArg : Nat := 0
-
-/-- Add a name translation to the translations map and add the `argInfo` information to `src`. -/
-def insertTranslationAndInfo (t : TranslateData) (src tgt : Name) (argInfo : ArgInfo)
-    (failIfExists := true) : CoreM Unit := do
-  insertTranslation t src tgt failIfExists
-  if argInfo.reorder != [] then
-    trace[translate] "@[{t.attrName}] will reorder the arguments of {tgt} by {argInfo.reorder}."
-    t.reorderAttr.add src argInfo.reorder
-  if argInfo.relevantArg != 0 then
-    trace[translate_detail] "Setting relevant_arg for {src} to be {argInfo.relevantArg}."
-    t.relevantArgAttr.add src argInfo.relevantArg
+        trace[translate] "The translation {src} ↦ {tgt'} already exists"
+    else
+      modifyEnv (t.translations.addEntry · (src, tgt))
+      trace[translate] "Added translation {src} ↦ {tgt}"
+    unless argInfo matches {} do
+      trace[translate] "@[{t.attrName}] will reorder the arguments of {src} by {argInfo.reorder}."
+      trace[translate_detail] "Setting relevant_arg for {src} to be {argInfo.relevantArg}."
+      modifyEnv (t.argInfoAttr.addEntry · (src, argInfo))
 
 /-- `Config` is the type of the arguments that can be provided to `to_additive`. -/
 structure Config : Type where
@@ -266,7 +262,6 @@ They are expanded until they are applied to one more argument than the maximum i
 It also expands all kernel projections that have as head a constant `n` in `reorder`. -/
 def expand (t : TranslateData) (e : Expr) : MetaM Expr := do
   let env ← getEnv
-  let reorderFn : Name → List (List ℕ) := fun nm ↦ (t.reorderAttr.find? env nm |>.getD [])
   let e₂ ← Lean.Meta.transform (input := e) (skipConstInApp := true)
     (post := fun e => return .done e) fun e ↦
     e.withApp fun f args ↦ do
@@ -281,11 +276,11 @@ def expand (t : TranslateData) (e : Expr) : MetaM Expr := do
       return .visit <| (← whnfD (← inferType s)).withApp fun sf sargs ↦
         mkAppN (mkApp (mkAppN (.const projName sf.constLevels!) sargs) s) args
     | .const c _ =>
-      let reorder := reorderFn c
-      if reorder.isEmpty then
+      let some info := t.argInfoAttr.find? env c | return .continue
+      if info.reorder.isEmpty then
         -- no need to expand if nothing needs reordering
         return .continue
-      let needed_n := reorder.flatten.foldr Nat.max 0 + 1
+      let needed_n := info.reorder.flatten.foldr Nat.max 0 + 1
       if needed_n ≤ args.size then
         return .continue
       else
@@ -387,19 +382,18 @@ def applyReplacementFun (t : TranslateData) (e : Expr) (dontTranslate : Array FV
   return e'
 where /-- Implementation of `applyReplacementFun`. -/
   aux (env : Environment) (trace : Bool) : Expr → Expr :=
-  let reorderFn : Name → List (List ℕ) := fun nm ↦ (t.reorderAttr.find? env nm |>.getD [])
-  let relevantArg : Name → ℕ := fun nm ↦ (t.relevantArgAttr.find? env nm).getD 0
   Lean.Expr.replaceRec fun r e ↦ Id.run do
     if trace then
       dbg_trace s!"replacing at {e}"
     match e with
     | .const n₀ ls₀ => do
       let n₁ := findPrefixTranslation env n₀ t
-      let ls₁ : List Level := if 0 ∈ (reorderFn n₀).flatten then ls₀.swapFirstTwo else ls₀
+      let swapUniv := (t.argInfoAttr.find? env n₀).elim false (·.reorder.any (·.contains 0))
+      let ls₁ : List Level := if swapUniv then ls₀.swapFirstTwo else ls₀
       if trace then
         if n₀ != n₁ then
           dbg_trace s!"changing {n₀} to {n₁}"
-        if 0 ∈ (reorderFn n₀).flatten then
+        if swapUniv then
           dbg_trace s!"reordering the universe variables from {ls₀} to {ls₁}"
       return some <| .const n₁ ls₁
     | .app g x => do
@@ -412,9 +406,9 @@ where /-- Implementation of `applyReplacementFun`. -/
       let some nm := gf.constName? | return mkAppN (← r gf) (← gAllArgs.mapM r)
       -- e = `(nm y₁ .. yₙ x)
       /- Test if the head should not be replaced. -/
-      let relevantArgId := relevantArg nm
-      if h : relevantArgId < gAllArgs.size then
-        if let some fxd := shouldTranslate env t gAllArgs[relevantArgId] dontTranslate then
+      let { reorder, relevantArg } := t.argInfoAttr.find? env nm |>.getD {}
+      if h : relevantArg < gAllArgs.size then
+        if let some fxd := shouldTranslate env t gAllArgs[relevantArg] dontTranslate then
           if trace then
             match fxd with
             | .inl fxd => dbg_trace s!"The application of {nm} contains the fixed type \
@@ -424,7 +418,6 @@ where /-- Implementation of `applyReplacementFun`. -/
         else
           gf ← r gf
           /- Test if arguments should be reordered. -/
-          let reorder := reorderFn nm
           if !reorder.isEmpty then
             gAllArgs := gAllArgs.permute! reorder
             if trace then
@@ -653,7 +646,7 @@ partial def transformDeclAux (t : TranslateData) (cfg : Config) (pre tgt_pre : N
   -- if the auxiliary declaration doesn't have prefix `pre`, then we have to add this declaration
   -- to the translation dictionary, since otherwise we cannot translate the name.
   if !pre.isPrefixOf src then
-    insertTranslation t src tgt
+    insertTranslation t src tgt {}
   -- now transform the source declaration
   let trgDecl : ConstantInfo ← MetaM.run' <|
     if src == pre then
@@ -761,7 +754,7 @@ def translateLemmas {m : Type → Type} [Monad m] [MonadError m] [MonadLiftT Cor
       throwError "{names[0]!} and {nm} do not generate the same number of {desc}."
   for (srcLemmas, tgtLemmas) in auxLemmas.zip <| auxLemmas.eraseIdx! 0 do
     for (srcLemma, tgtLemma) in srcLemmas.zip tgtLemmas do
-      insertTranslationAndInfo t srcLemma tgtLemma argInfo
+      insertTranslation t srcLemma tgtLemma argInfo
 
 /--
 Find the argument of `nm` that appears in the first translatable (type-class) argument.
@@ -778,7 +771,7 @@ def findRelevantArg (t : TranslateData) (nm : Name) : MetaM Nat := do
     let relevantArg? (tgt : Expr) : Option Nat := do
       let c ← tgt.getAppFn.constName?
       guard (findTranslation? env t c).isSome
-      let relevantArg := (t.relevantArgAttr.find? env c).getD 0
+      let relevantArg := (t.argInfoAttr.find? env c).elim 0 (·.relevantArg)
       let arg ← tgt.getArg? relevantArg
       xs.findIdx? (arg.containsFVar ·.fvarId!)
     -- run the above check on all hypotheses and on the conclusion
@@ -825,7 +818,7 @@ def proceedFieldsAux (t : TranslateData) (src tgt : Name) (argInfo : ArgInfo)
     throwError "Failed to map fields of {src}, {tgt} with {srcFields} ↦ {tgtFields}.\n \
       Lengths do not match."
   for srcField in srcFields, tgtField in tgtFields do
-    insertTranslationAndInfo t srcField tgtField argInfo
+    insertTranslation t srcField tgtField argInfo
 
 /-- Add the structure fields of `src` to the translations dictionary
 so that they will be translated correctly. -/
@@ -1068,15 +1061,10 @@ partial def addTranslationAttr (t : TranslateData) (src : Name) (cfg : Config)
     -- If `tgt` is not in the environment, the translation to `tgt` was added only for
     -- translating the namespace, and `src` wasn't actually tagged.
     if (← getEnv).contains tgt then
-      let mut updated := false
-      if cfg.reorder != [] then
-        modifyEnv (t.reorderAttr.addEntry · (src, cfg.reorder))
-        updated := true
-      if let some relevantArg := cfg.relevantArg? then
-        modifyEnv (t.relevantArgAttr.addEntry · (src, relevantArg))
-        updated := true
-      if updated then
+      if cfg.reorder != [] || cfg.relevantArg?.isSome then
         MetaM.run' <| checkExistingType t src tgt cfg.reorder cfg.dontTranslate
+        let argInfo := { reorder := cfg.reorder, relevantArg := cfg.relevantArg?.getD 0 }
+        insertTranslation t src tgt argInfo false
         return #[tgt]
       throwError
       "Cannot apply attribute @[{t.attrName}] to '{src}': it is already translated to '{tgt}'. \n\
@@ -1095,7 +1083,7 @@ partial def addTranslationAttr (t : TranslateData) (src : Name) (cfg : Config)
     MetaM.run' <| checkExistingType t src tgt cfg.reorder cfg.dontTranslate
   let relevantArg ← cfg.relevantArg?.getDM <| MetaM.run' <| findRelevantArg t src
   let argInfo := { reorder := cfg.reorder, relevantArg }
-  insertTranslationAndInfo t src tgt argInfo alreadyExists
+  insertTranslation t src tgt argInfo alreadyExists
   let nestedNames ←
     if alreadyExists then
       -- since `tgt` already exists, we just need to copy metadata and

Mathlib/Tactic/Translate/ToAdditive.lean

@@ -258,28 +258,8 @@ initialize ignoreArgsAttr : NameMapExtension (List Nat) ←
           | _ => throwUnsupportedSyntax
         return ids.toList }
 
-/-- An extension that stores all the declarations that need their arguments reordered when
-applying `@[to_additive]`. It is applied using the `to_additive (reorder := ...)` syntax. -/
-initialize reorderAttr : NameMapExtension (List (List Nat)) ←
-  registerNameMapExtension _
-
-/-- Linter to check that the `relevant_arg` attribute is not given manually -/
-register_option linter.toAdditiveRelevantArg : Bool := {
-  defValue := true
-  descr := "Linter to check that the `relevant_arg` attribute is not given manually." }
-
-@[inherit_doc to_additive_relevant_arg]
-initialize relevantArgAttr : NameMapExtension Nat ←
-  registerNameMapAttribute {
-    name := `to_additive_relevant_arg
-    descr := "Auxiliary attribute for `to_additive` stating \
-      which arguments are the types with a multiplicative structure."
-    add := fun
-    | _, stx@`(attr| to_additive_relevant_arg $id) => do
-      Linter.logLintIf linter.toAdditiveRelevantArg stx
-        m!"This attribute is deprecated. Use `@[to_additive (relevant_arg := ...)]` instead."
-      pure <| id.getNat.pred
-    | _, _ => throwUnsupportedSyntax }
+@[inherit_doc TranslateData.argInfoAttr]
+initialize argInfoAttr : NameMapExtension ArgInfo ← registerNameMapExtension _
 
 @[inherit_doc to_additive_dont_translate]
 initialize dontTranslateAttr : NameMapExtension Unit ←
@@ -389,11 +369,7 @@ def abbreviationDict : Std.HashMap String String := .ofList [
 
 /-- The bundle of environment extensions for `to_additive` -/
 def data : TranslateData where
-  ignoreArgsAttr := ignoreArgsAttr
-  reorderAttr := reorderAttr
-  relevantArgAttr := relevantArgAttr
-  dontTranslateAttr := dontTranslateAttr
-  translations := translations
+  ignoreArgsAttr; argInfoAttr; dontTranslateAttr; translations
   attrName := `to_additive
   changeNumeral := true
   isDual := false
@@ -412,6 +388,6 @@ initialize registerBuiltinAttribute {
 into the `to_additive` dictionary. This is useful for translating namespaces that don't (yet)
 have a corresponding translated declaration. -/
 elab "insert_to_additive_translation" src:ident tgt:ident : command => do
-  Command.liftCoreM <| insertTranslation data src.getId tgt.getId
+  translations.add src.getId tgt.getId
 
 end Mathlib.Tactic.ToAdditive

Mathlib/Tactic/Translate/ToDual.lean

@@ -96,20 +96,8 @@ initialize ignoreArgsAttr : NameMapExtension (List Nat) ←
           | _ => throwUnsupportedSyntax
         return ids.toList }
 
-/-- An extension that stores all the declarations that need their arguments reordered when
-applying `@[to_dual]`. It is applied using the `to_dual (reorder := ...)` syntax. -/
-initialize reorderAttr : NameMapExtension (List (List Nat)) ←
-  registerNameMapExtension _
-
-@[inherit_doc to_dual_relevant_arg]
-initialize relevantArgAttr : NameMapExtension Nat ←
-  registerNameMapAttribute {
-    name := `to_dual_relevant_arg
-    descr := "Auxiliary attribute for `to_dual` stating \
-      which arguments are the types with a dual structure."
-    add := fun
-    | _, `(attr| to_dual_relevant_arg $id) => pure <| id.1.isNatLit?.get!.pred
-    | _, _ => throwUnsupportedSyntax }
+@[inherit_doc TranslateData.argInfoAttr]
+initialize argInfoAttr : NameMapExtension ArgInfo ← registerNameMapExtension _
 
 @[inherit_doc to_dual_dont_translate]
 initialize dontTranslateAttr : NameMapExtension Unit ←
@@ -177,11 +165,7 @@ def abbreviationDict : Std.HashMap String String := .ofList []
 
 /-- The bundle of environment extensions for `to_dual` -/
 def data : TranslateData where
-  ignoreArgsAttr := ignoreArgsAttr
-  reorderAttr := reorderAttr
-  relevantArgAttr := relevantArgAttr
-  dontTranslateAttr := dontTranslateAttr
-  translations := translations
+  ignoreArgsAttr; argInfoAttr; dontTranslateAttr; translations
   attrName := `to_dual
   changeNumeral := false
   isDual := true

MathlibTest/ToDual.lean

@@ -0,0 +1,39 @@
+import Mathlib.Order.Defs.PartialOrder
+import Mathlib.Order.Notation
+
+-- test that we can translate between structures, reordering the arguments of the fields
+class SemilatticeInf (α : Type) extends PartialOrder α, Min α where
+  le_inf : ∀ a b c : α, a ≤ b → a ≤ c → a ≤ b ⊓ c
+
+class SemilatticeSup (α : Type) extends PartialOrder α, Max α where
+  protected sup_le : ∀ a b c : α, a ≤ c → b ≤ c → a ⊔ b ≤ c
+
+attribute [to_dual] SemilatticeInf
+attribute [to_dual (reorder := 3 4 5)] SemilatticeSup.sup_le
+
+@[to_dual]
+lemma SemilatticeInf.le_inf' {α : Type} [SemilatticeInf α] (a b c : α) : a ≤ b → a ≤ c → a ≤ b ⊓ c :=
+  SemilatticeInf.le_inf a b c
+
+@[to_dual]
+lemma SemilatticeSup.sup_le' {α : Type} [SemilatticeSup α] (a b c : α) : a ≤ c → b ≤ c → a ⊔ b ≤ c :=
+  SemilatticeSup.sup_le a b c
+
+-- we still cannot reorder arguments of arguments, so `SemilatticeInf.mk` is not tranlatable
+/--
+error: @[to_dual] failed. The translated value is not type correct. For help, see the docstring of `to_additive`, section `Troubleshooting`. Failed to add declaration
+instSemilatticeSupOfForallLeForallMax:
+Application type mismatch: The argument
+  le_inf
+has type
+  ∀ (a b c : α), b ≤ a → c ≤ a → b ⊔ c ≤ a
+but is expected to have type
+  ∀ (a b c : α), a ≤ c → b ≤ c → a ⊔ b ≤ c
+in the application
+  { toPartialOrder := inst✝¹, toMax := inst✝, sup_le := le_inf }
+-/
+#guard_msgs in
+@[to_dual]
+instance {α : Type} [PartialOrder α] [Min α]
+    (le_inf : ∀ a b c : α, a ≤ b → a ≤ c → a ≤ b ⊓ c) : SemilatticeInf α where
+  le_inf

MathlibTest/toAdditive.lean

@@ -156,14 +156,14 @@ example {x} (h : 1 = x) : baz20 = x := by simp; guard_target = 1 = x; exact h
 @[to_additive bar21]
 def foo21 {N} {A} [Pow A N] (a : A) (n : N) : A := a ^ n
 
-run_cmd liftCoreM <| MetaM.run' <| guard <| relevantArgAttr.find? (← getEnv) `Test.foo21 == some 1
+run_meta guard <| argInfoAttr.find? (← getEnv) `Test.foo21 matches some ⟨[], 1⟩
 
 @[to_additive bar22]
 abbrev foo22 {α} [Monoid α] (a : α) : ℕ → α
   | 0 => 1
   | _ => a
 
-run_cmd liftCoreM <| MetaM.run' <| do
+run_meta do
   -- make `abbrev` definition `reducible` automatically
   guard <| (← getReducibilityStatus `Test.bar22) == .reducible
   -- make `abbrev` definition `inline` automatically