From d0391d07c2fb2e3962e16d4e3950d8da87ff675a Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 27 Sep 2021 18:15:57 -0700
Subject: [PATCH] feat: use `PSigma.casesOn` instead of projections at
 `packDomain`

Reason: we want to "refine" the `WellFounded.fix` functional `F` over it.
---
 src/Lean/Elab/PreDefinition/WF/Fix.lean       |  2 ++
 .../Elab/PreDefinition/WF/PackDomain.lean     | 19 ++++++++++++++++---
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/src/Lean/Elab/PreDefinition/WF/Fix.lean b/src/Lean/Elab/PreDefinition/WF/Fix.lean
index 37382a87d7..98d92d5505 100644
--- a/src/Lean/Elab/PreDefinition/WF/Fix.lean
+++ b/src/Lean/Elab/PreDefinition/WF/Fix.lean
@@ -91,6 +91,8 @@ private partial def processSumCasesOn (x F val : Expr) (k : (F : Expr) → (val
     k F val
 
 def mkFix (preDef : PreDefinition) (wfRel : Expr) : TermElabM PreDefinition := do
+  trace[Elab.definition.wf] ">> {preDef.value}"
+  check preDef.value -- TODO remove
   let wfFix ← forallBoundedTelescope preDef.type (some 1) fun x type => do
     let x := x[0]
     let α ← inferType x
diff --git a/src/Lean/Elab/PreDefinition/WF/PackDomain.lean b/src/Lean/Elab/PreDefinition/WF/PackDomain.lean
index c3c205c0f0..a87504b438 100644
--- a/src/Lean/Elab/PreDefinition/WF/PackDomain.lean
+++ b/src/Lean/Elab/PreDefinition/WF/PackDomain.lean
@@ -40,6 +40,18 @@ where
     else
       return args[i]
 
+private partial def mkPSigmaCasesOn (y : Expr) (codomain : Expr) (xs : Array Expr) (value : Expr) : MetaM Expr := do
+  let mvar ← mkFreshExprSyntheticOpaqueMVar codomain
+  let rec go (mvarId : MVarId) (y : FVarId) (ys : Array Expr) : MetaM Unit := do
+    if ys.size < xs.size - 1 then
+      let #[s] ← cases mvarId y | unreachable!
+      go s.mvarId s.fields[1].fvarId! (ys.push s.fields[0])
+    else
+      let ys := ys.push (mkFVar y)
+      assignExprMVar mvarId (value.replaceFVars xs ys)
+  go mvar.mvarId! y.fvarId! #[]
+  instantiateMVars mvar
+
 /--
   Convert the given pre-definitions into unary functions.
   We "pack" the arguments using `PSigma`.
@@ -88,8 +100,9 @@ def packDomain (preDefs : Array PreDefinition) : MetaM (Array PreDefinition) :=
     let arity := arities[i]
     let valueNew ← lambdaTelescope preDef.value fun xs body => do
       if arity > 1 then
-        forallBoundedTelescope preDefNew.type (some 1) fun y _ => do
-          let newBody := body.replaceFVars xs (mkTupleElems y[0] xs.size)
+        forallBoundedTelescope preDefNew.type (some 1) fun y codomain => do
+          let y := y[0]
+          let newBody ← mkPSigmaCasesOn y codomain xs body
           let visit (e : Expr) : MetaM TransformStep := do
             if let some idx := isTargetApp? e |>.run then
               let f := e.getAppFn
@@ -98,7 +111,7 @@ def packDomain (preDefs : Array PreDefinition) : MetaM (Array PreDefinition) :=
               return TransformStep.done <| mkApp fNew argNew
             else
               return TransformStep.done e
-          mkLambdaFVars y (← transform newBody (post := visit))
+          mkLambdaFVars #[y] (← transform newBody (post := visit))
       else
         preDef.value
     if let some bad := valueNew.find? fun e => isAppOfPreDef? e |>.isSome then