From c95d5f25a35d9e4a7d48d587cc0ecd6e4df59eec Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sat, 19 Mar 2022 10:11:50 -0700
Subject: [PATCH] feat: convert `ite`s into `dite`s in the `WF` module

Motivation: the condition is often used in termination proofs.
---
 src/Lean/Elab/PreDefinition/WF/Ite.lean  | 29 ++++++++++++++++++++++++
 src/Lean/Elab/PreDefinition/WF/Main.lean |  4 +++-
 tests/lean/run/lex.lean                  |  4 ++--
 tests/lean/wf2.lean.expected.out         |  1 +
 4 files changed, 35 insertions(+), 3 deletions(-)
 create mode 100644 src/Lean/Elab/PreDefinition/WF/Ite.lean

diff --git a/src/Lean/Elab/PreDefinition/WF/Ite.lean b/src/Lean/Elab/PreDefinition/WF/Ite.lean
new file mode 100644
index 0000000000..d83357d7a4
--- /dev/null
+++ b/src/Lean/Elab/PreDefinition/WF/Ite.lean
@@ -0,0 +1,29 @@
+/-
+Copyright (c) 2022 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+import Lean.Meta.Transform
+
+namespace Lean.Meta
+
+/--
+  Convert `ite` expressions in `e` to `dite`s.
+  It is useful to make this conversion in the `WF` module because the condition is often used in
+  termination proofs. -/
+def iteToDIte (e : Expr) : MetaM Expr := do
+  -- TODO: move this file to `Meta` if we decide to use it in other places.
+  let post (e : Expr) : MetaM TransformStep := do
+    if e.isAppOfArity ``ite 5 then
+      let f    := e.getAppFn
+      let args := e.getAppArgs
+      let c    := args[1]
+      let h    ← mkFreshUserName `h
+      let args := args.set! 3 (Lean.mkLambda h BinderInfo.default c args[3])
+      let args := args.set! 4 (Lean.mkLambda h BinderInfo.default (mkNot c) args[4])
+      return .done <| mkAppN (mkConst ``dite f.constLevels!) args
+    else
+      return .done e
+  transform e (post := post)
+
+end Lean.Meta
diff --git a/src/Lean/Elab/PreDefinition/WF/Main.lean b/src/Lean/Elab/PreDefinition/WF/Main.lean
index e069e45a1f..105697abca 100644
--- a/src/Lean/Elab/PreDefinition/WF/Main.lean
+++ b/src/Lean/Elab/PreDefinition/WF/Main.lean
@@ -10,6 +10,7 @@ import Lean.Elab.PreDefinition.WF.PackMutual
 import Lean.Elab.PreDefinition.WF.Rel
 import Lean.Elab.PreDefinition.WF.Fix
 import Lean.Elab.PreDefinition.WF.Eqns
+import Lean.Elab.PreDefinition.WF.Ite
 
 namespace Lean.Elab
 open WF
@@ -84,7 +85,8 @@ def wfRecursion (preDefs : Array PreDefinition) (wf? : Option TerminationWF) (de
       addAsAxiom preDef
     let fixedPrefixSize ← getFixedPrefix preDefs
     trace[Elab.definition.wf] "fixed prefix: {fixedPrefixSize}"
-    let unaryPreDefs ← packDomain fixedPrefixSize preDefs
+    let preDefsDIte ← preDefs.mapM fun preDef => return { preDef with value := (← iteToDIte preDef.value) }
+    let unaryPreDefs ← packDomain fixedPrefixSize preDefsDIte
     return (← packMutual fixedPrefixSize preDefs unaryPreDefs, fixedPrefixSize)
   let preDefNonRec ← forallBoundedTelescope unaryPreDef.type fixedPrefixSize fun prefixArgs type => do
     let packedArgType := type.bindingDomain!
diff --git a/tests/lean/run/lex.lean b/tests/lean/run/lex.lean
index 561c6f1262..6ca9920d14 100644
--- a/tests/lean/run/lex.lean
+++ b/tests/lean/run/lex.lean
@@ -182,7 +182,7 @@ namespace Ex5
 
 mutual
   def lex [Monad m] [MonadExceptOf LexErr m] (it : String.Iterator) : m (List Token) := do
-    if h : it.hasNext then
+    if it.hasNext then
       match it.curr with
       | '(' => return { text := "(", tok := Tok.lpar } :: (← lex it.next)
       | ')' => return { text := ")", tok := Tok.rpar } :: (← lex it.next)
@@ -195,7 +195,7 @@ mutual
       return []
 
   def lexnumber [Monad m] [MonadExceptOf LexErr m] (soFar : Nat) (text : List Char) (it : String.Iterator) : m (List Token) :=
-    if h : it.hasNext then
+    if it.hasNext then
       let c := it.curr
       match charDigit c with
       | .error _ => return { text := text.reverse.asString, tok := Tok.num soFar } :: (← lex it)
diff --git a/tests/lean/wf2.lean.expected.out b/tests/lean/wf2.lean.expected.out
index 65bda1a461..fe427e05f2 100644
--- a/tests/lean/wf2.lean.expected.out
+++ b/tests/lean/wf2.lean.expected.out
@@ -3,4 +3,5 @@ wf2.lean:3:8-3:17: error: failed to prove termination, possible solutions:
   - Use `termination_by` to specify a different well-founded relation
   - Use `decreasing_by` to specify your own tactic for discharging this kind of goal
 x y : Nat
+ : x < y
 ⊢ x - 1 < x