Improve the syntax of progress: `as ⟨ x, y .. ⟩`

2 files changed, 94 insertions, 44 deletions

diff --git a/backends/lean/Base/Progress/Progress.lean b/backends/lean/Base/Progress/Progress.lean
index 1c509775..c8f94e9e 100644
--- a/backends/lean/Base/Progress/Progress.lean
+++ b/backends/lean/Base/Progress/Progress.lean
@@ -23,7 +23,8 @@ inductive ProgressError
 | Error (msg : MessageData)
 deriving Inhabited
 
-def progressWith (fExpr : Expr) (th : TheoremOrLocal) (keep : Option Name) (ids : Array Name)
+def progressWith (fExpr : Expr) (th : TheoremOrLocal)
+  (keep : Option Name) (ids : Array Name) (splitPost : Bool)
   (asmTac : TacticM Unit) : TacticM ProgressError := do
   /- Apply the theorem
      We try to match the theorem with the goal
@@ -112,24 +113,31 @@ def progressWith (fExpr : Expr) (th : TheoremOrLocal) (keep : Option Name) (ids
         mgoal.tryClearMany #[hEq.fvarId!]
     setGoals (mgoal :: (← getUnsolvedGoals))
     trace[Progress] "Goal after splitting eq and post and simplifying the target: {mgoal}"
-    -- Continue splitting following the ids provided by the user
-    if ¬ ids.isEmpty then
-      let hPost ←
-        match hPost with
-        | none => do return (.Error m!"Too many ids provided ({ids}): there is no postcondition to split")
-        | some hPost => pure hPost
-      let rec splitPost (prevId : Name) (hPost : Expr) (ids : List Name) : TacticM ProgressError := do
+    -- Continue splitting following the post following the user's instructions
+    match hPost with
+    | none =>
+      -- Sanity check
+      if ¬ ids.isEmpty then
+        return (.Error m!"Too many ids provided ({ids}): there is no postcondition to split")
+      else return .Ok
+    | some hPost => do
+      let rec splitPostWithIds (prevId : Name) (hPost : Expr) (ids : List Name) : TacticM ProgressError := do
         match ids with
-        | [] => pure .Ok -- Stop
+        | [] =>
+          /- We used all the user provided ids.
+             Split the remaining conjunctions by using fresh ids if the user
+             instructed to fully split the post-condition, otherwise stop -/
+          if splitPost then
+            splitFullConjTac hPost (λ _ => pure .Ok)
+          else pure .Ok
         | nid :: ids => do
           trace[Progress] "Splitting post: {hPost}"
           -- Split
           if ← isConj (← inferType hPost) then
-            splitConjTac hPost (some (prevId, nid)) (λ _ nhPost => splitPost nid nhPost ids)
+            splitConjTac hPost (some (prevId, nid)) (λ _ nhPost => splitPostWithIds nid nhPost ids)
           else return (.Error m!"Too many ids provided ({nid :: ids}) not enough conjuncts to split in the postcondition")
       let curPostId := (← hPost.fvarId!.getDecl).userName
-      splitPost curPostId hPost ids
-    else return .Ok
+      splitPostWithIds curPostId hPost ids
   match res with
   | .Error _ => return res -- Can we get there? We're using "return"
   | .Ok =>
@@ -160,7 +168,8 @@ def getFirstArg (args : Array Expr) : Option Expr := do
 
 /- Helper: try to lookup a theorem and apply it, or continue with another tactic
    if it fails -/
-def tryLookupApply (keep : Option Name) (ids : Array Name) (asmTac : TacticM Unit) (fExpr : Expr)
+def tryLookupApply (keep : Option Name) (ids : Array Name) (splitPost : Bool)
+  (asmTac : TacticM Unit) (fExpr : Expr)
   (kind : String) (th : Option TheoremOrLocal) (x : TacticM Unit) : TacticM Unit := do
   let res ← do
     match th with
@@ -172,7 +181,7 @@ def tryLookupApply (keep : Option Name) (ids : Array Name) (asmTac : TacticM Uni
       -- Apply the theorem
       let res ← do
         try
-          let res ← progressWith fExpr th keep ids asmTac
+          let res ← progressWith fExpr th keep ids splitPost asmTac
           pure (some res)
         catch _ => none
   match res with
@@ -181,7 +190,8 @@ def tryLookupApply (keep : Option Name) (ids : Array Name) (asmTac : TacticM Uni
   | none => x
 
 -- The array of ids are identifiers to use when introducing fresh variables
-def progressAsmsOrLookupTheorem (keep : Option Name) (withTh : Option TheoremOrLocal) (ids : Array Name) (asmTac : TacticM Unit) : TacticM Unit := do
+def progressAsmsOrLookupTheorem (keep : Option Name) (withTh : Option TheoremOrLocal)
+  (ids : Array Name) (splitPost : Bool) (asmTac : TacticM Unit) : TacticM Unit := do
   withMainContext do
   -- Retrieve the goal
   let mgoal ← Tactic.getMainGoal
@@ -197,7 +207,7 @@ def progressAsmsOrLookupTheorem (keep : Option Name) (withTh : Option TheoremOrL
   -- Otherwise, lookup one.
   match withTh with
   | some th => do
-    match ← progressWith fExpr th keep ids asmTac with
+    match ← progressWith fExpr th keep ids splitPost asmTac with
     | .Ok => return ()
     | .Error msg => throwError msg
   | none =>
@@ -206,7 +216,7 @@ def progressAsmsOrLookupTheorem (keep : Option Name) (withTh : Option TheoremOrL
     let decls ← ctx.getDecls
     for decl in decls.reverse do
       trace[Progress] "Trying assumption: {decl.userName} : {decl.type}"
-      let res ← do try progressWith fExpr (.Local decl) keep ids asmTac catch _ => continue
+      let res ← do try progressWith fExpr (.Local decl) keep ids splitPost asmTac catch _ => continue
       match res with
       | .Ok => return ()
       | .Error msg => throwError msg
@@ -216,7 +226,7 @@ def progressAsmsOrLookupTheorem (keep : Option Name) (withTh : Option TheoremOrL
     let pspec ← do
       let thName ← pspecAttr.find? fName
       pure (thName.map fun th => .Theorem th)
-    tryLookupApply keep ids asmTac fExpr "pspec theorem" pspec do
+    tryLookupApply keep ids splitPost asmTac fExpr "pspec theorem" pspec do
     -- It failed: try to lookup a *class* expr spec theorem (those are more
     -- specific than class spec theorems)
     let pspecClassExpr ← do
@@ -225,7 +235,7 @@ def progressAsmsOrLookupTheorem (keep : Option Name) (withTh : Option TheoremOrL
       | some arg => do
         let thName ← pspecClassExprAttr.find? fName arg
         pure (thName.map fun th => .Theorem th)
-    tryLookupApply keep ids asmTac fExpr "pspec class expr theorem" pspecClassExpr do
+    tryLookupApply keep ids splitPost asmTac fExpr "pspec class expr theorem" pspecClassExpr do
     -- It failed: try to lookup a *class* spec theorem
     let pspecClass ← do
       match ← getFirstArgAppName args with
@@ -233,7 +243,7 @@ def progressAsmsOrLookupTheorem (keep : Option Name) (withTh : Option TheoremOrL
       | some argName => do
         let thName ← pspecClassAttr.find? fName argName
         pure (thName.map fun th => .Theorem th)
-    tryLookupApply keep ids asmTac fExpr "pspec class theorem" pspecClass do
+    tryLookupApply keep ids splitPost asmTac fExpr "pspec class theorem" pspecClass do
     -- Try a recursive call - we try the assumptions of kind "auxDecl"
     let ctx ← Lean.MonadLCtx.getLCtx
     let decls ← ctx.getAllDecls
@@ -241,14 +251,14 @@ def progressAsmsOrLookupTheorem (keep : Option Name) (withTh : Option TheoremOrL
       | .default | .implDetail => false | .auxDecl => true)
     for decl in decls.reverse do
       trace[Progress] "Trying recursive assumption: {decl.userName} : {decl.type}"
-      let res ← do try progressWith fExpr (.Local decl) keep ids asmTac catch _ => continue
+      let res ← do try progressWith fExpr (.Local decl) keep ids splitPost asmTac catch _ => continue
       match res with
       | .Ok => return ()
       | .Error msg => throwError msg
     -- Nothing worked: failed
     throwError "Progress failed"
 
-syntax progressArgs := ("keep" ("as" (ident))?)? ("with" ident)? ("as" " ⟨ " (ident)+ " ⟩")?
+syntax progressArgs := ("keep" ("as" (ident))?)? ("with" ident)? ("as" " ⟨ " ident,* " .."? " ⟩")?
 
 def evalProgress (args : TSyntax `Progress.progressArgs) : TacticM Unit := do
   let args := args.raw
@@ -263,8 +273,8 @@ def evalProgress (args : TSyntax `Progress.progressArgs) : TacticM Unit := do
       else do pure (some (← mkFreshUserName `h))
     else pure none
   trace[Progress] "Keep: {keep}"
-  let withArg := (args.get! 1).getArgs
   let withArg ← do
+    let withArg := (args.get! 1).getArgs
     if withArg.size > 0 then
       let id := withArg.get! 1
       trace[Progress] "With arg: {id}"
@@ -283,20 +293,25 @@ def evalProgress (args : TSyntax `Progress.progressArgs) : TacticM Unit := do
         | id :: _ =>
           pure (some (.Theorem id))
     else pure none
-  let args := (args.get! 2).getArgs
-  let args := (args.get! 2).getArgs
-  let ids := args.map Syntax.getId
+  let ids :=
+    let args := (args.get! 2).getArgs
+    let args := (args.get! 2).getSepArgs
+    args.map Syntax.getId
   trace[Progress] "User-provided ids: {ids}"
-  progressAsmsOrLookupTheorem keep withArg ids (firstTac [assumptionTac, Arith.scalarTac])
+  let splitPost : Bool :=
+    let args := (args.get! 2).getArgs
+    (args.get! 3).getArgs.size > 0
+  trace[Progress] "Split post: {splitPost}"
+  progressAsmsOrLookupTheorem keep withArg ids splitPost (firstTac [assumptionTac, Arith.scalarTac])
 
 elab "progress" args:progressArgs : tactic =>
   evalProgress args
 
-/-
-namespace Test
+/-namespace Test
   open Primitives Result
 
   set_option trace.Progress true
+  set_option pp.rawOnError true
 
   #eval showStoredPSpec
   #eval showStoredPSpecClass
@@ -306,9 +321,9 @@ namespace Test
     (hmax : x.val + y.val ≤ Scalar.max ty) :
     ∃ z, x + y = ret z ∧ z.val = x.val + y.val := by
 --    progress keep as h with Scalar.add_spec as ⟨ z ⟩
-    progress keep as h
+    progress keep as h as ⟨ z, h1 .. ⟩
     simp [*]
 
-end Test -/
+end Test-/
 
 end Progress
diff --git a/backends/lean/Base/Utils.lean b/backends/lean/Base/Utils.lean
index f014e112..3b3d4729 100644
--- a/backends/lean/Base/Utils.lean
+++ b/backends/lean/Base/Utils.lean
@@ -499,7 +499,10 @@ def splitConjTac (h : Expr) (optIds : Option (Name × Name)) (k : Expr → Expr
     -- Try to use the user-provided names
     let altVarNames ←
       match optIds with
-      | none => pure #[]
+      | none => do
+        let id0 ← mkFreshUserName `h
+        let id1 ← mkFreshUserName `h
+        pure #[{ varNames := [id0, id1] }]
       | some (id0, id1) => do
         pure #[{ varNames := [id0, id1] }]
     let newGoals ← goal.cases h.fvarId! altVarNames
@@ -518,29 +521,61 @@ def splitConjTac (h : Expr) (optIds : Option (Name × Name)) (k : Expr → Expr
   else
     throwError "Not a conjunction"
 
-elab "split_conj " n:ident : tactic => do
+-- Tactic to fully split a conjunction
+partial def splitFullConjTacAux [Inhabited α] [Nonempty α] (l : List Expr) (h : Expr) (k : List Expr → TacticM α)  : TacticM α := do
+  try
+    splitConjTac h none (λ h1 h2 =>
+      splitFullConjTacAux l h1 (λ l1 =>
+        splitFullConjTacAux l1 h2 (λ l2 =>
+          k l2)))
+  catch _ =>
+    k (h :: l)
+
+-- Tactic to fully split a conjunction
+def splitFullConjTac [Inhabited α] [Nonempty α] (h : Expr) (k : List Expr → TacticM α)  : TacticM α := do
+  splitFullConjTacAux [] h (λ l => k l.reverse)
+
+syntax optAtArgs := ("at" ident)?
+def elabOptAtArgs (args : TSyntax `Utils.optAtArgs) : TacticM (Option Expr) := do
   withMainContext do
-  let decl ← Lean.Meta.getLocalDeclFromUserName n.getId
-  let fvar := mkFVar decl.fvarId
-  splitConjTac fvar none (fun _ _ =>  pure ())
+  let args := (args.raw.getArgs.get! 0).getArgs
+  if args.size > 0 then do
+    let n := (args.get! 1).getId
+    let decl ← Lean.Meta.getLocalDeclFromUserName n
+    let fvar := mkFVar decl.fvarId
+    pure (some fvar)
+  else
+    pure none
 
-elab "split_all_exists " n:ident : tactic => do
+elab "split_conj" args:optAtArgs : tactic => do
+  withMainContext do
+  match ← elabOptAtArgs args with
+  | some fvar =>
+    splitConjTac fvar none (fun _ _ =>  pure ())
+  | none =>
+    splitConjTarget
+
+elab "split_conjs" args:optAtArgs : tactic => do
+  withMainContext do
+  match ← elabOptAtArgs args with
+  | some fvar =>
+    splitFullConjTac fvar (fun _ =>  pure ())
+  | none =>
+    repeatTac splitConjTarget
+
+elab "split_existsl" " at " n:ident : tactic => do
   withMainContext do
   let decl ← Lean.Meta.getLocalDeclFromUserName n.getId
   let fvar := mkFVar decl.fvarId
   splitAllExistsTac fvar [] (fun _ _ => pure ())
 
-elab "split_target_conjs" : tactic =>
-  withMainContext do
-  repeatTac splitConjTarget
-
 example (h : a ∧ b) : a := by
-  split_all_exists h
-  split_conj h
+  split_existsl at h
+  split_conj at h
   assumption
 
 example (h : ∃ x y z, x + y + z ≥ 0) : ∃ x, x ≥ 0 := by
-  split_all_exists h
+  split_existsl at h
   rename_i x y z
   exists x + y + z