fix: fix to_additive handling of Nat and Int (#689)

* Remove `[to_additive]` attributes from `Nat` and `Int`. That attribute disables all useful heuristics for `to_additive` on these types.
* deal with literals in specific types, like `Nat`. This uses a new attribute `@[to_additive_fixed_numeral]`.
* Add `to_additive_relevant_arg` attribute for heterogenous operations
* Fix `to_additive_reorder` arguments on `HPow.hPow`
* Fix parsing of `to_additive_relevant_arg` to use 1-indexed numbers instead of 0-indexed numbers
* Do not add `relevantArgAttr` on many declarations if they would have the default value anyway.
* Revert test `foo2` in `tests/ToAdditive` to the Lean 3 behavior. Previosuly the test would only succeed if the heuristics for `Nat` and `Int` were disabled.
* Some cleanup in code, comments and docs

Author: fpvandoorn

Mathlib/Algebra/Group/Defs.lean

@@ -60,20 +60,16 @@ infixl:65 " +ᵥ " => VAdd.vadd
 infixl:65 " -ᵥ " => HasVsub.vsub
 infixr:73 " • " => SMul.smul
 
--- [todo] is this correct? I think it's needed to ensure that additiveTest
--- succeeds if the relevant arg involves Nat.
-attribute [to_additive] Nat
-attribute [to_additive] Int
-
 attribute [to_additive] Mul
 attribute [to_additive] Div
 attribute [to_additive] HMul
 attribute [to_additive] instHMul
 attribute [to_additive] HDiv
 attribute [to_additive] instHDiv
 
+attribute [to_additive_relevant_arg 3] HMul HAdd HAdd.hAdd HMul.hMul
 attribute [to_additive_reorder 1] HPow
-attribute [to_additive_reorder 1 4] HPow.hPow
+attribute [to_additive_reorder 1 5] HPow.hPow
 attribute [to_additive] HPow
 
 universe u
@@ -482,7 +478,7 @@ need right away.
 
 In the definition, we use `n.succ` instead of `n + 1` in the `nsmul_succ'` and `npow_succ'` fields
 to make sure that `to_additive` is not confused (otherwise, it would try to convert `1 : ℕ`
-to `0 : ℕ`).
+to `0 : ℕ`). Todo: fix this in `to_additive`
 -/
 
 

Mathlib/Data/Fin/Basic.lean

@@ -74,6 +74,7 @@ section
 
 variable {n : Nat} [Nonempty (Fin n)]
 
+@[to_additive_fixed_numeral]
 instance : OfNat (Fin n) a where
   ofNat := Fin.ofNat' a Fin.size_positive'
 

Mathlib/Init/ZeroOne.lean

@@ -24,14 +24,17 @@ class One (α : Type u) where
   one : α
 #align has_one One
 
-@[to_additive Zero.toOfNat0]
+@[to_additive Zero.toOfNat0, to_additive_fixed_numeral ?]
 instance One.toOfNat1 {α} [One α] : OfNat α (nat_lit 1) where
   ofNat := ‹One α›.1
-@[to_additive Zero.ofOfNat0]
+@[to_additive Zero.ofOfNat0, to_additive_fixed_numeral ?]
 instance One.ofOfNat1 {α} [OfNat α (nat_lit 1)] : One α where
   one := 1
 
 @[deprecated, match_pattern] def bit0 {α : Type u} [Add α] (a : α) : α := a + a
 
 set_option linter.deprecated false in
 @[deprecated, match_pattern] def bit1 {α : Type u} [One α] [Add α] (a : α) : α := bit0 a + 1
+
+attribute [to_additive_fixed_numeral ?] OfNat OfNat.ofNat
+attribute [to_additive_fixed_numeral] instOfNatNat instOfNatInt

Mathlib/Tactic/ToAdditive.lean

@@ -33,6 +33,8 @@ syntax (name := to_additive_ignore_args) "to_additive_ignore_args" num* : attr
 syntax (name := to_additive_relevant_arg) "to_additive_relevant_arg" num : attr
 /-- The  `to_additive_reorder` attribute. -/
 syntax (name := to_additive_reorder) "to_additive_reorder" num* : attr
+/-- The  `to_additive_fixed_numeral` attribute. -/
+syntax (name := to_additive_fixed_numeral) "to_additive_fixed_numeral" "?"? : attr
 /-- Remaining arguments of `to_additive`. -/
 syntax to_additiveRest := (ppSpace ident)? (ppSpace str)?
 /-- The `to_additive` attribute. -/
@@ -104,7 +106,7 @@ initialize ignoreArgsAttr : NameMapExtension (List Nat) ←
     add   := fun _ stx ↦ do
         let ids ← match stx with
           | `(attr| to_additive_ignore_args $[$ids:num]*) => pure <| ids.map (·.1.isNatLit?.get!)
-          | _ => throwError "unexpected to_additive_ignore_args syntax {stx}"
+          | _ => throwUnsupportedSyntax
         return ids.toList
   }
 
@@ -131,8 +133,7 @@ initialize reorderAttr : NameMapExtension (List Nat) ←
     add := fun
     | _, `(attr| to_additive_reorder $[$ids:num]*) =>
       pure <| Array.toList <| ids.map (·.1.isNatLit?.get!)
-    | _, stx => throwError "unexpected to_additive_reorder syntax {stx}"
-  }
+    | _, _ => throwUnsupportedSyntax }
 
 /-- Get the reorder list (defined using `@[to_additive_reorder ...]`) for the given declaration. -/
 def getReorder [Functor M] [MonadEnv M]: Name →  M (List Nat)
@@ -170,9 +171,8 @@ initialize relevantArgAttr : NameMapExtension Nat ←
     descr := "Auxiliary attribute for `to_additive` stating" ++
       " which arguments are the types with a multiplicative structure."
     add := fun
-    | _, `(attr| to_additive_relevant_arg $id) => pure <| id.1.isNatLit?.get!
-    | _, stx => throwError "unexpected to_additive_relevant_arg syntax {stx}"
-  }
+    | _, `(attr| to_additive_relevant_arg $id) => pure <| id.1.isNatLit?.get!.pred
+    | _, _ => throwUnsupportedSyntax }
 
 /-- Given a declaration name and an argument index, determines whether it
 is relevant. This is used in `applyReplacementFun` where more detail on what it
@@ -182,6 +182,25 @@ def isRelevant [Monad M] [MonadEnv M] (n : Name) (i : Nat) : M Bool := do
   | some j => return i == j
   | none => return i == 0
 
+/--
+An attribute that stores all the declarations that deal with numeric literals on fixed types.
+*  `@[to_additive_fixed_numeral]` should be added to all functions that take a numeral as argument
+  that should never be changed by `@[to_additive]` (because it represents a numeral in a fixed
+  type).
+* `@[to_additive_fixed_numeral?]` should be added to all functions that take a numeral as argument
+  that should only be changed if `additiveTest` succeeds on the first argument, i.e. when the
+  numeral is only translated if the first argument is a variable (or consists of variables).
+-/
+initialize fixedNumeralAttr : NameMapExtension Bool ←
+  registerNameMapAttribute {
+    name := `to_additive_fixed_numeral
+    descr :=
+      "Auxiliary attribute for `to_additive` that stores functions that have numerals as argument."
+    add := fun
+    | _, `(attr| to_additive_fixed_numeral $[?%$conditional]?) =>
+      pure <| conditional.isSome
+    | _, _ => throwUnsupportedSyntax }
+
 /-- Maps multiplicative names to their additive counterparts. -/
 initialize translations : NameMapExtension Name ← registerNameMapExtension _
 
@@ -238,18 +257,23 @@ def additiveTest (e : Expr) : M Bool := do
   else
     additiveTestAux false e
 
+/-- Checks whether a numeral should be translated. -/
+def shouldTranslateNumeral [Monad M] [MonadEnv M] (n : Name) (firstArg : Expr) : M Bool := do
+  match fixedNumeralAttr.find? (← getEnv) n with
+  | some true => additiveTest firstArg
+  | some false => return false
+  | none => return true
+
 /--
-`e.applyReplacementFun f test` applies `f` to each identifier
-(inductive type, defined function etc) in an expression, unless
+`applyReplacementFun e` replaces the expression `e` with its additive counterpart.
+It translates each identifier (inductive type, defined function etc) in an expression, unless
 * The identifier occurs in an application with first argument `arg`; and
 * `test arg` is false.
 However, if `f` is in the dictionary `relevant`, then the argument `relevant.find f`
 is tested, instead of the first argument.
 
-Reorder contains the information about what arguments to reorder:
-e.g. `g x₁ x₂ x₃ ... xₙ` becomes `g x₂ x₁ x₃ ... xₙ` if `reorder.find g = some [1]`.
-We assume that all functions where we want to reorder arguments are fully applied.
-This can be done by applying `etaExpand` first.
+It will also reorder arguments of certain functions, using `shouldReorder`:
+e.g. `g x₁ x₂ x₃ ... xₙ` becomes `g x₂ x₁ x₃ ... xₙ` if `reorderAttr.find? env g = some [1]`.
 -/
 def applyReplacementFun : Expr → MetaM Expr :=
   Lean.Expr.replaceRecMeta fun r e ↦ do
@@ -258,7 +282,8 @@ def applyReplacementFun : Expr → MetaM Expr :=
     | .lit (.natVal 1) => pure <| mkRawNatLit 0
     | .const n₀ ls => do
       let n₁ := Name.mapPrefix (findTranslation? <|← getEnv) n₀
-      trace[to_additive_detail] "applyReplacementFun: {n₀} → {n₁}"
+      if n₀ != n₁ then
+        trace[to_additive_detail] "applyReplacementFun: {n₀} → {n₁}"
       let ls : List Level ← (do -- [todo] just get Lean to figure out the levels?
         if ← shouldReorder n₀ 1 then
             return ls.get! 1::ls.head!::ls.drop 2
@@ -270,6 +295,7 @@ def applyReplacementFun : Expr → MetaM Expr :=
         let gArgs := g.getAppArgs
         -- e = `(nm y₁ .. yₙ x)
         trace[to_additive_detail] "applyReplacementFun: app {nm} {gArgs} {x}"
+        /- Test if arguments should be reordered. -/
         if h : gArgs.size > 0 then
           let c1 ← shouldReorder nm gArgs.size
           let c2 ← additiveTest gArgs[0]
@@ -282,16 +308,25 @@ def applyReplacementFun : Expr → MetaM Expr :=
             trace[to_additive_detail]
               "applyReplacementFun: reordering {nm}: {x} ↔ {ga}\nBefore: {e}\nAfter:  {e₂}"
             return some e₂
+        /- Test if the head should not be replaced. -/
         let c1 ← isRelevant nm gArgs.size
         let c2 := gf.isConst
         let c3 ← additiveTest x
+        if c1 && c2 && c3 then
+          trace[to_additive_detail]
+            "applyReplacementFun: {x} doesn't contain a fixed type, so we will change {nm}"
         if c1 && c2 && not c3 then
           -- the test failed, so don't update the function body.
           trace[to_additive_detail]
-            "applyReplacementFun: isRelevant and test failed: {nm} {gArgs} {x}"
+            "applyReplacementFun: {x} contains a fixed type, so {nm} is not changed"
           let x ← r x
           let args ← gArgs.mapM r
           return some $ mkApp (mkAppN gf args) x
+        /- Do not replace numerals in specific types. -/
+        let firstArg := if h : gArgs.size > 0 then gArgs[0] else x
+        if not (← shouldTranslateNumeral nm firstArg) then
+          trace[to_additive_detail] "applyReplacementFun: Do not change numeral {g.app x}"
+          return some <| g.app x
       return e.updateApp! (← r g) (← r x)
     | _ => return none
 
@@ -366,7 +401,7 @@ partial def transformDeclAux
   if env.contains tgt then
     return
   let srcDecl ← getConstInfo src
-  -- we first transform all the declarations of the form `pre._proof_i`
+  -- we first transform all auxilliary declarations generated when elaborating `pre`
   for n in srcDecl.type.listNamesWithPrefix pre do
     transformDeclAux none pre tgt_pre n
   if let some value := srcDecl.value? then
@@ -450,7 +485,7 @@ Find the first argument of `nm` that has a multiplicative type-class on it.
 Returns 1 if there are no types with a multiplicative class as arguments.
 E.g. `prod.group` returns 1, and `pi.has_one` returns 2.
 -/
-def firstMultiplicativeArg (nm : Name) : MetaM (Option Nat) := do
+def firstMultiplicativeArg (nm : Name) : MetaM Nat := do
   forallTelescopeReducing (← getConstInfo nm).type fun xs _ ↦ do
     -- xs are the arguments to the constant
     let xs := xs.toList
@@ -466,8 +501,8 @@ def firstMultiplicativeArg (nm : Name) : MetaM (Option Nat) := do
           xs.findIdx? fun x ↦ Expr.containsFVar tgt_arg x.fvarId!
     trace[to_additive_detail] "firstMultiplicativeArg: {l}"
     match l.join with
-    | [] => return none
-    | (head :: tail) => return some <| tail.foldr Nat.min head
+    | [] => return 0
+    | (head :: tail) => return tail.foldr Nat.min head
 
 /-- `ValueType` is the type of the arguments that can be provided to `to_additive`. -/
 structure ValueType : Type where
@@ -633,12 +668,12 @@ def targetName (src tgt : Name) (allowAutoName : Bool) : CoreM Name := do
     throwError "to_additive: can't transport {src} to itself."
   return res
 
-private def proceedFieldsAux (src tgt : Name) (f : Name → CoreM (List String)) : CoreM Unit := do
+private def proceedFieldsAux (src tgt : Name) (f : Name → CoreM (Array Name)) : CoreM Unit := do
   let srcFields ← f src
   let tgtFields ← f tgt
-  if srcFields.length != tgtFields.length then
+  if srcFields.size != tgtFields.size then
     throwError "Failed to map fields of {src}, {tgt} with {srcFields} ↦ {tgtFields}"
-  for (srcField, tgtField) in List.zip srcFields tgtFields do
+  for (srcField, tgtField) in srcFields.zip tgtFields do
     if srcField != tgtField then
       insertTranslation (src ++ srcField) (tgt ++ tgtField)
 
@@ -647,12 +682,9 @@ so that future uses of `to_additive` will map them to the corresponding `tgt` fi
 def proceedFields (src tgt : Name) : CoreM Unit := do
   let env : Environment ← getEnv
   let aux := proceedFieldsAux src tgt
-  aux (fun n ↦ do
-    let fields := if isStructure env n then getStructureFieldsFlattened env n else #[]
-    return fields |> .map Name.toString |> Array.toList
-  )
-  -- [todo] run to_additive on the constructors of n:
-  -- aux (fun n ↦ (env.constructorsOf n).mmap $ ...
+  aux fun n ↦ pure <| if isStructure env n then getStructureFields env n else #[]
+  -- We don't have to run toAdditive on the constructor of a structure, since the use of
+  -- `Name.mapPrefix` will do that automatically.
 
 private def elabToAdditiveAux (ref : Syntax) (replaceAll trace : Bool) (tgt : Option Syntax)
     (doc : Option Syntax) : ValueType :=
@@ -661,8 +693,7 @@ private def elabToAdditiveAux (ref : Syntax) (replaceAll trace : Bool) (tgt : Op
     tgt := match tgt with | some tgt => tgt.getId | none => Name.anonymous
     doc := doc.bind (·.isStrLit?)
     allowAutoName := false
-    ref
-  }
+    ref }
 
 private def elabToAdditive : Syntax → CoreM ValueType
   | `(attr| to_additive%$tk $[!%$replaceAll]? $[?%$trace]? $[$tgt]? $[$doc]?) =>
@@ -676,7 +707,8 @@ def addToAdditiveAttr (src : Name) (val : ValueType) : AttrM Unit := do
   if let some tgt' := findTranslation? (← getEnv) src then
     throwError "{src} already has a to_additive translation {tgt'}."
   insertTranslation src tgt
-  if let some firstMultArg ← (MetaM.run' <| firstMultiplicativeArg src) then
+  let firstMultArg ← MetaM.run' <| firstMultiplicativeArg src
+  if firstMultArg != 0 then
     trace[to_additive_detail] "Setting relevant_arg for {src} to be {firstMultArg}."
     relevantArgAttr.add src firstMultArg
   if (← getEnv).contains tgt then

test/toAdditive.lean

@@ -41,17 +41,14 @@ theorem bar1_works : bar1 3 4 = 3 * 4 := by decide
 
 infix:80 " ^ " => my_has_pow.pow
 
-instance dummy_pow : my_has_pow ℕ $ PLift ℤ := ⟨fun _ _ => 0⟩
-instance dummy_smul : my_has_scalar (PLift ℤ) ℕ := ⟨fun _ _ => 0⟩
-attribute [to_additive dummy_smul] dummy_pow
+instance dummy_pow : my_has_pow ℕ $ PLift ℤ := ⟨fun _ _ => 5⟩
 
 set_option pp.universes true
 @[to_additive bar2]
 def foo2 {α} [my_has_pow α ℕ] (x : α) (n : ℕ) (m : PLift ℤ) : α := x ^ (n ^ m)
 
-theorem foo2_works : foo2 2 3 (PLift.up 2) = Nat.pow 2 0 := by decide
--- [todo] should it still be using dummy?
-theorem bar2_works : bar2 2 3 (PLift.up 2) =  2 * (dummy_smul.1 (PLift.up 2) 3) := by decide
+theorem foo2_works : foo2 2 3 (PLift.up 2) = Nat.pow 2 5 := by decide
+theorem bar2_works : bar2 2 3 (PLift.up 2) =  2 * 5 := by decide
 
 @[to_additive bar3]
 def foo3 {α} [my_has_pow α ℕ] (x : α) : ℕ → α := @my_has_pow.pow α ℕ _ x
@@ -77,6 +74,28 @@ def foo7 := @my_has_pow.pow
 theorem foo7_works : foo7 2 3 = Nat.pow 2 3 := by decide
 theorem bar7_works : bar7 2 3 =  2 * 3 := by decide
 
+/-- Check that we don't additivize `Nat` expressions. -/
+@[to_additive bar8]
+def foo8 (a b : ℕ) := a * b
+
+theorem bar8_works : bar8 2 3 = 6 := by decide
+
+/-- Check that we don't additivize `Nat` numerals. -/
+@[to_additive bar9]
+def foo9 := 1
+
+theorem bar9_works : bar9 = 1 := by decide
+
+@[to_additive bar10]
+def foo10 (n m : ℕ) := HPow.hPow n m + n * m * 2 + 1 * 0 + 37 * 1 + 2
+
+theorem bar10_works : bar10 = foo10 := by rfl
+
+@[to_additive bar11]
+def foo11 (n : ℕ) (m : ℤ) := n * m * 2 + 1 * 0 + 37 * 1 + 2
+
+theorem bar11_works : bar11 = foo11 := by rfl
+
 /- test the eta-expansion applied on `foo6`. -/
 run_cmd do
   let c ← getConstInfo `Test.foo6
@@ -106,9 +125,9 @@ run_cmd do
   Elab.Command.liftCoreM <| successIfFail (getConstInfo `Test.add_some_def.in_namespace)
 
 -- [todo] currently this test breaks.
--- example : (add_units.mk_of_add_eq_zero 0 0 (by simp) : ℕ)
---         = (add_units.mk_of_add_eq_zero 0 0 (by simp) : ℕ) :=
--- by normCast
+-- example : (AddUnits.mk_of_add_eq_zero 0 0 (by simp) : ℕ)
+--         = (AddUnits.mk_of_add_eq_zero 0 0 (by simp) : ℕ) :=
+-- by norm_cast
 
 section
 
@@ -126,10 +145,10 @@ instance pi.has_one {I : Type} {f : I → Type} [(i : I) → One $ f i] : One ((
 run_cmd do
   let n ← (Elab.Command.liftCoreM <| Lean.Meta.MetaM.run' <| ToAdditive.firstMultiplicativeArg
     `Test.pi.has_one)
-  if n != some 1 then throwError "{n} != 1"
+  if n != 1 then throwError "{n} != 1"
   let n ← (Elab.Command.liftCoreM <| Lean.Meta.MetaM.run' <| ToAdditive.firstMultiplicativeArg
     `Test.foo_mul)
-  if n != some 4 then throwError "{n} != 4"
+  if n != 4 then throwError "{n} != 4"
 
 end
 
@@ -139,6 +158,8 @@ def nat_pi_has_one {α : Type} [One α] : One ((x : Nat) → α) := by infer_ins
 @[to_additive]
 def pi_nat_has_one {I : Type} : One ((x : I) → Nat)  := pi.has_one
 
+example : @pi_nat_has_one = @pi_nat_has_zero := rfl
+
 section noncomputablee
 
 @[to_additive Bar.bar]
@@ -147,7 +168,7 @@ noncomputable def Foo.foo (h : ∃ _ : α, True) : α := Classical.choose h
 @[to_additive Bar.bar']
 def Foo.foo' : ℕ := 2
 
-#eval Bar.bar'
+theorem Bar.bar'_works : Bar.bar' = 2 := by decide
 
 run_cmd (do
   if !isNoncomputable (← getEnv) `Bar.bar then throwError "bar shouldn't be computable"