Merge pull request #31 from AeneasVerif/son_lean_backend

Improve the Lean backend

2 files changed, 693 insertions, 0 deletions

diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
new file mode 100644
index 00000000..6f820a84
--- /dev/null
+++ b/backends/lean/Base/Progress/Base.lean
@@ -0,0 +1,316 @@
+import Lean
+import Std.Lean.HashSet
+import Base.Utils
+import Base.Primitives.Base
+
+namespace Progress
+
+open Lean Elab Term Meta
+open Utils
+
+-- We can't define and use trace classes in the same file
+initialize registerTraceClass `Progress
+
+/- # Progress tactic -/
+
+structure PSpecDesc where
+  -- The universally quantified variables
+  fvars : Array Expr
+  -- The existentially quantified variables
+  evars : Array Expr
+  -- The function
+  fExpr : Expr
+  fName : Name
+  -- The function arguments
+  fLevels : List Level
+  args : Array Expr
+  -- The universally quantified variables which appear in the function arguments
+  argsFVars : Array FVarId
+  -- The returned value
+  ret : Expr
+  -- The postcondition (if there is)
+  post : Option Expr
+
+section Methods
+  variable [MonadLiftT MetaM m] [MonadControlT MetaM m] [Monad m] [MonadOptions m]
+  variable [MonadTrace m] [MonadLiftT IO m] [MonadRef m] [AddMessageContext m]
+  variable [MonadError m]
+  variable {a : Type}
+
+  /- Analyze a pspec theorem to decompose its arguments.
+
+     PSpec theorems should be of the following shape:
+     ```
+     ∀ x1 ... xn, H1 → ... Hn → ∃ y1 ... ym. f x1 ... xn = .ret ... ∧ Post1 ∧ ... ∧ Postk
+     ```
+
+     The continuation `k` receives the following inputs:
+     - universally quantified variables
+     - assumptions
+     - existentially quantified variables
+     - function name
+     - function arguments
+     - return
+     - postconditions
+
+     TODO: generalize for when we do inductive proofs
+  -/
+  partial
+  def withPSpec [Inhabited (m a)] [Nonempty (m a)] (th : Expr) (k : PSpecDesc → m a)
+    (sanityChecks : Bool := false) :
+    m a := do
+    trace[Progress] "Proposition: {th}"
+    -- Dive into the quantified variables and the assumptions
+    forallTelescope th.consumeMData fun fvars th => do
+    trace[Progress] "Universally quantified arguments and assumptions: {fvars}"
+    -- Dive into the existentials
+    existsTelescope th.consumeMData fun evars th => do
+    trace[Progress] "Existentials: {evars}"
+    trace[Progress] "Proposition after stripping the quantifiers: {th}"
+    -- Take the first conjunct
+    let (th, post) ← optSplitConj th.consumeMData
+    trace[Progress] "After splitting the conjunction:\n- eq: {th}\n- post: {post}"
+    -- Destruct the equality
+    let (mExpr, ret) ← destEq th.consumeMData
+    trace[Progress] "After splitting the equality:\n- lhs: {th}\n- rhs: {ret}"
+    -- Destruct the monadic application to dive into the bind, if necessary (this
+    -- is for when we use `withPSpec` inside of the `progress` tactic), and
+    -- destruct the application to get the function name
+    mExpr.consumeMData.withApp fun mf margs => do
+    trace[Progress] "After stripping the arguments of the monad expression:\n- mf: {mf}\n- margs: {margs}"
+    let (fExpr, f, args) ← do
+      if mf.isConst ∧ mf.constName = ``Bind.bind then do
+        -- Dive into the bind
+        let fExpr := (margs.get! 4).consumeMData
+        fExpr.withApp fun f args => pure (fExpr, f, args)
+      else pure (mExpr, mf, margs)
+    trace[Progress] "After stripping the arguments of the function call:\n- f: {f}\n- args: {args}"
+    if ¬ f.isConst then throwError "Not a constant: {f}"
+    -- Compute the set of universally quantified variables which appear in the function arguments
+    let allArgsFVars ← args.foldlM (fun hs arg => getFVarIds arg hs) HashSet.empty
+    -- Sanity check
+    if sanityChecks then
+      -- All the variables which appear in the inputs given to the function are
+      -- universally quantified (in particular, they are not *existentially* quantified)
+      let fvarsSet : HashSet FVarId := HashSet.ofArray (fvars.map (fun x => x.fvarId!))
+      let filtArgsFVars := allArgsFVars.toArray.filter (fun fvar => ¬ fvarsSet.contains fvar)
+      if ¬ filtArgsFVars.isEmpty then
+        let filtArgsFVars := filtArgsFVars.map (fun fvarId => Expr.fvar fvarId)
+        throwError "Some of the function inputs are not universally quantified: {filtArgsFVars}"
+    let argsFVars := fvars.map (fun x => x.fvarId!)
+    let argsFVars := argsFVars.filter (fun fvar => allArgsFVars.contains fvar)
+    -- Return
+    trace[Progress] "Function: {f.constName!}";
+    let thDesc := {
+      fvars := fvars
+      evars := evars
+      fExpr
+      fName := f.constName!
+      fLevels := f.constLevels!
+      args := args
+      argsFVars
+      ret := ret
+      post := post
+    }
+    k thDesc
+
+end Methods
+
+def getPSpecFunName (th : Expr) : MetaM Name :=
+  withPSpec th (fun d => do pure d.fName) true
+
+def getPSpecClassFunNames (th : Expr) : MetaM (Name × Name) :=
+  withPSpec th (fun d => do
+    let arg0 := d.args.get! 0
+    arg0.withApp fun f _ => do
+    if ¬ f.isConst then throwError "Not a constant: {f}"
+    pure (d.fName, f.constName)
+    ) true
+
+def getPSpecClassFunNameArg (th : Expr) : MetaM (Name × Expr) :=
+  withPSpec th (fun d => do
+    let arg0 := d.args.get! 0
+    pure (d.fName, arg0)
+    ) true
+
+-- "Regular" pspec attribute
+structure PSpecAttr where
+  attr : AttributeImpl
+  ext  : MapDeclarationExtension Name
+  deriving Inhabited
+
+/- pspec attribute for type classes: we use the name of the type class to
+   lookup another map. We use the *first* argument of the type class to lookup
+   into this second map.
+
+   Example:
+   ========
+   We use type classes for addition. For instance, the addition between two
+   U32 is written (without syntactic sugar) as `HAdd.add (Scalar ty) x y`. As a consequence,
+   we store the theorem through the bindings: HAdd.add → Scalar → ...
+
+   SH: TODO: this (and `PSpecClassExprAttr`) is a bit ad-hoc. For now it works for the
+   specs of the scalar operations, which is what I really need, but I'm not sure it
+   applies well to other situations. A better way would probably to use type classes, but
+   I couldn't get them to work on those cases. It is worth retrying.
+-/
+structure PSpecClassAttr where
+  attr : AttributeImpl
+  ext  : MapDeclarationExtension (NameMap Name)
+  deriving Inhabited
+
+/- Same as `PSpecClassAttr` but we use the full first argument (it works when it
+   is a constant). -/
+structure PSpecClassExprAttr where
+  attr : AttributeImpl
+  ext  : MapDeclarationExtension (HashMap Expr Name)
+  deriving Inhabited
+
+-- TODO: the original function doesn't define correctly the `addImportedFn`. Do a PR?
+def mkMapDeclarationExtension [Inhabited α] (name : Name := by exact decl_name%) : IO (MapDeclarationExtension α) :=
+  registerSimplePersistentEnvExtension {
+    name          := name,
+    addImportedFn := fun a => a.foldl (fun s a => a.foldl (fun s (k, v) => s.insert k v) s) RBMap.empty,
+    addEntryFn    := fun s n => s.insert n.1 n.2 ,
+    toArrayFn     := fun es => es.toArray.qsort (fun a b => Name.quickLt a.1 b.1)
+  }
+
+/- The persistent map from function to pspec theorems. -/
+initialize pspecAttr : PSpecAttr ← do
+  let ext ← mkMapDeclarationExtension `pspecMap
+  let attrImpl : AttributeImpl := {
+    name := `pspec
+    descr := "Marks theorems to use with the `progress` tactic"
+    add := fun thName stx attrKind => do
+      Attribute.Builtin.ensureNoArgs stx
+      -- TODO: use the attribute kind
+      unless attrKind == AttributeKind.global do
+        throwError "invalid attribute 'pspec', must be global"
+      -- Lookup the theorem
+      let env ← getEnv
+      let thDecl := env.constants.find! thName
+      let fName ← MetaM.run' (getPSpecFunName thDecl.type)
+      trace[Progress] "Registering spec theorem for {fName}"
+      let env := ext.addEntry env (fName, thName)
+      setEnv env
+      pure ()
+  }
+  registerBuiltinAttribute attrImpl
+  pure { attr := attrImpl, ext := ext }
+
+/- The persistent map from type classes to pspec theorems -/
+initialize pspecClassAttr : PSpecClassAttr ← do
+  let ext : MapDeclarationExtension (NameMap Name) ← mkMapDeclarationExtension `pspecClassMap
+  let attrImpl : AttributeImpl  := {
+    name := `cpspec
+    descr := "Marks theorems to use for type classes with the `progress` tactic"
+    add := fun thName stx attrKind => do
+      Attribute.Builtin.ensureNoArgs stx
+      -- TODO: use the attribute kind
+      unless attrKind == AttributeKind.global do
+        throwError "invalid attribute 'cpspec', must be global"
+      -- Lookup the theorem
+      let env ← getEnv
+      let thDecl := env.constants.find! thName
+      let (fName, argName) ← MetaM.run' (getPSpecClassFunNames thDecl.type)
+      trace[Progress] "Registering class spec theorem for ({fName}, {argName})"
+      -- Update the entry if there is one, add an entry if there is none
+      let env :=
+        match (ext.getState (← getEnv)).find? fName with
+        | none =>
+          let m := RBMap.ofList [(argName, thName)]
+          ext.addEntry env (fName, m)
+        | some m =>
+          let m := m.insert argName thName
+          ext.addEntry env (fName, m)
+      setEnv env
+      pure ()
+  }
+  registerBuiltinAttribute attrImpl
+  pure { attr := attrImpl, ext := ext }
+
+/- The 2nd persistent map from type classes to pspec theorems -/
+initialize pspecClassExprAttr : PSpecClassExprAttr ← do
+  let ext : MapDeclarationExtension (HashMap Expr Name) ← mkMapDeclarationExtension `pspecClassExprMap
+  let attrImpl : AttributeImpl  := {
+    name := `cepspec
+    descr := "Marks theorems to use for type classes with the `progress` tactic"
+    add := fun thName stx attrKind => do
+      Attribute.Builtin.ensureNoArgs stx
+      -- TODO: use the attribute kind
+      unless attrKind == AttributeKind.global do
+        throwError "invalid attribute 'cpspec', must be global"
+      -- Lookup the theorem
+      let env ← getEnv
+      let thDecl := env.constants.find! thName
+      let (fName, arg) ← MetaM.run' (getPSpecClassFunNameArg thDecl.type)
+      -- Sanity check: no variables appear in the argument
+      MetaM.run' do
+        let fvars ← getFVarIds arg
+        if ¬ fvars.isEmpty then throwError "The first argument ({arg}) contains variables"
+      -- We store two bindings:
+      -- - arg to theorem name
+      -- - reduced arg to theorem name
+      let rarg ← MetaM.run' (reduceAll arg)
+      trace[Progress] "Registering class spec theorem for ({fName}, {arg}) and ({fName}, {rarg})"
+      -- Update the entry if there is one, add an entry if there is none
+      let env :=
+        match (ext.getState (← getEnv)).find? fName with
+        | none =>
+          let m := HashMap.ofList [(arg, thName), (rarg, thName)]
+          ext.addEntry env (fName, m)
+        | some m =>
+          let m := m.insert arg thName
+          let m := m.insert rarg thName
+          ext.addEntry env (fName, m)
+      setEnv env
+      pure ()
+  }
+  registerBuiltinAttribute attrImpl
+  pure { attr := attrImpl, ext := ext }
+
+
+def PSpecAttr.find? (s : PSpecAttr) (name : Name) : MetaM (Option Name) := do
+  return (s.ext.getState (← getEnv)).find? name
+
+def PSpecClassAttr.find? (s : PSpecClassAttr) (className argName : Name) : MetaM (Option Name) := do
+  match (s.ext.getState (← getEnv)).find? className with
+  | none => return none
+  | some map => return map.find? argName
+
+def PSpecClassExprAttr.find? (s : PSpecClassExprAttr) (className : Name) (arg : Expr) : MetaM (Option Name) := do
+  match (s.ext.getState (← getEnv)).find? className with
+  | none => return none
+  | some map => return map.find? arg
+
+def PSpecAttr.getState (s : PSpecAttr) : MetaM (NameMap Name) := do
+  pure (s.ext.getState (← getEnv))
+
+def PSpecClassAttr.getState (s : PSpecClassAttr) : MetaM (NameMap (NameMap Name)) := do
+  pure (s.ext.getState (← getEnv))
+
+def PSpecClassExprAttr.getState (s : PSpecClassExprAttr) : MetaM (NameMap (HashMap Expr Name)) := do
+  pure (s.ext.getState (← getEnv))
+
+def showStoredPSpec : MetaM Unit := do
+  let st ← pspecAttr.getState
+  let s := st.toList.foldl (fun s (f, th) => f!"{s}\n{f} → {th}") f!""
+  IO.println s
+
+def showStoredPSpecClass : MetaM Unit := do
+  let st ← pspecClassAttr.getState
+  let s := st.toList.foldl (fun s (f, m) =>
+    let ms := m.toList.foldl (fun s (f, th) =>
+      f!"{s}\n  {f} → {th}") f!""
+    f!"{s}\n{f} → [{ms}]") f!""
+  IO.println s
+
+def showStoredPSpecExprClass : MetaM Unit := do
+  let st ← pspecClassExprAttr.getState
+  let s := st.toList.foldl (fun s (f, m) =>
+    let ms := m.toList.foldl (fun s (f, th) =>
+      f!"{s}\n  {f} → {th}") f!""
+    f!"{s}\n{f} → [{ms}]") f!""
+  IO.println s
+
+end Progress
diff --git a/backends/lean/Base/Progress/Progress.lean b/backends/lean/Base/Progress/Progress.lean
new file mode 100644
index 00000000..6a4729dc
--- /dev/null
+++ b/backends/lean/Base/Progress/Progress.lean
@@ -0,0 +1,377 @@
+import Lean
+import Base.Arith
+import Base.Progress.Base
+import Base.Primitives -- TODO: remove?
+
+namespace Progress
+
+open Lean Elab Term Meta Tactic
+open Utils
+
+inductive TheoremOrLocal where
+| Theorem (thName : Name)
+| Local (asm : LocalDecl)
+
+instance : ToMessageData TheoremOrLocal where
+  toMessageData := λ x => match x with | .Theorem thName => m!"{thName}" | .Local asm => m!"{asm.userName}"
+
+/- Type to propagate the errors of `progressWith`.
+   We need this because we use the exceptions to backtrack, when trying to
+   use the assumptions for instance. When there is actually an error we want
+   to propagate to the user, we return it. -/
+inductive ProgressError
+| Ok
+| Error (msg : MessageData)
+deriving Inhabited
+
+def progressWith (fExpr : Expr) (th : TheoremOrLocal)
+  (keep : Option Name) (ids : Array (Option Name)) (splitPost : Bool)
+  (asmTac : TacticM Unit) : TacticM ProgressError := do
+  /- Apply the theorem
+     We try to match the theorem with the goal
+     In order to do so, we introduce meta-variables for all the parameters
+     (i.e., quantified variables and assumpions), and unify those with the goal.
+     Remark: we do not introduce meta-variables for the quantified variables
+     which don't appear in the function arguments (we want to let them
+     quantified).
+     We also make sure that all the meta variables which appear in the
+     function arguments have been instantiated
+   -/
+  let env ← getEnv
+  let thTy ← do
+    match th with
+    | .Theorem thName =>
+      let thDecl := env.constants.find! thName
+      -- We have to introduce fresh meta-variables for the universes already
+      let ul : List (Name × Level) ←
+        thDecl.levelParams.mapM (λ x => do pure (x, ← mkFreshLevelMVar))
+      let ulMap : HashMap Name Level := HashMap.ofList ul
+      let thTy := thDecl.type.instantiateLevelParamsCore (λ x => ulMap.find! x)
+      pure thTy
+    | .Local asmDecl => pure asmDecl.type
+  trace[Progress] "Looked up theorem/assumption type: {thTy}"
+  -- TODO: the tactic fails if we uncomment withNewMCtxDepth
+  -- withNewMCtxDepth do
+  let (mvars, binders, thExBody) ← forallMetaTelescope thTy
+  trace[Progress] "After stripping foralls: {thExBody}"
+  -- Introduce the existentially quantified variables and the post-condition
+  -- in the context
+  let thBody ←
+    existsTelescope thExBody.consumeMData fun _evars thBody => do
+    trace[Progress] "After stripping existentials: {thBody}"
+    let (thBody, _) ← optSplitConj thBody
+    trace[Progress] "After splitting the conjunction: {thBody}"
+    let (thBody, _) ← destEq thBody
+    trace[Progress] "After splitting equality: {thBody}"
+    -- There shouldn't be any existential variables in thBody
+    pure thBody.consumeMData
+  -- Match the body with the target
+  trace[Progress] "Matching:\n- body:\n{thBody}\n- target:\n{fExpr}"
+  let ok ← isDefEq thBody fExpr
+  if ¬ ok then throwError "Could not unify the theorem with the target:\n- theorem: {thBody}\n- target: {fExpr}"
+  let mgoal ← Tactic.getMainGoal
+  postprocessAppMVars `progress mgoal mvars binders true true
+  Term.synthesizeSyntheticMVarsNoPostponing
+  let thBody ← instantiateMVars thBody
+  trace[Progress] "thBody (after instantiation): {thBody}"
+  -- Add the instantiated theorem to the assumptions (we apply it on the metavariables).
+  let th ← do
+    match th with
+    | .Theorem thName => mkAppOptM thName (mvars.map some)
+    | .Local decl => mkAppOptM' (mkFVar decl.fvarId) (mvars.map some)
+  let asmName ← do match keep with | none => mkFreshAnonPropUserName | some n => do pure n
+  let thTy ← inferType th
+  let thAsm ← Utils.addDeclTac asmName th thTy (asLet := false)
+  withMainContext do -- The context changed - TODO: remove once addDeclTac is updated
+  let ngoal ← getMainGoal
+  trace[Progress] "current goal: {ngoal}"
+  trace[Progress] "current goal: {← ngoal.isAssigned}"
+  -- The assumption should be of the shape:
+  -- `∃ x1 ... xn, f args = ... ∧ ...`
+  -- We introduce the existentially quantified variables and split the top-most
+  -- conjunction if there is one. We use the provided `ids` list to name the
+  -- introduced variables.
+  let res ← splitAllExistsTac thAsm ids.toList fun h ids => do
+    -- Split the conjunctions.
+    -- For the conjunctions, we split according once to separate the equality `f ... = .ret ...`
+    -- from the postcondition, if there is, then continue to split the postcondition if there
+    -- are remaining ids.
+    let splitEqAndPost (k : Expr → Option Expr → List (Option Name) → TacticM ProgressError) : TacticM ProgressError := do
+      if ← isConj (← inferType h) then do
+        let hName := (← h.fvarId!.getDecl).userName
+        let (optIds, ids) ← do
+          match ids with
+          | [] => do pure (some (hName, ← mkFreshAnonPropUserName), [])
+          | none :: ids => do pure (some (hName, ← mkFreshAnonPropUserName), ids)
+          | some id :: ids => do pure (some (hName, id), ids)
+        splitConjTac h optIds (fun hEq hPost => k hEq (some hPost) ids)
+      else k h none ids
+    -- Simplify the target by using the equality and some monad simplifications,
+    -- then continue splitting the post-condition
+    splitEqAndPost fun hEq hPost ids => do
+    trace[Progress] "eq and post:\n{hEq} : {← inferType hEq}\n{hPost}"
+    simpAt [] [``Primitives.bind_tc_ret, ``Primitives.bind_tc_fail, ``Primitives.bind_tc_div]
+           [hEq.fvarId!] (.targets #[] true)
+    -- Clear the equality, unless the user requests not to do so
+    let mgoal ← do
+      if keep.isSome then getMainGoal
+      else do
+        let mgoal ← getMainGoal
+        mgoal.tryClearMany #[hEq.fvarId!]
+    setGoals (mgoal :: (← getUnsolvedGoals))
+    trace[Progress] "Goal after splitting eq and post and simplifying the target: {mgoal}"
+    -- 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) (ids0 : List (Option Name)) : TacticM ProgressError := do
+        match ids0 with
+        | [] =>
+          /- 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 true hPost (λ _ => pure .Ok)
+          else pure .Ok
+        | nid :: ids => do
+          trace[Progress] "Splitting post: {← inferType hPost}"
+          -- Split
+          let nid ← do
+            match nid with
+            | none => mkFreshAnonPropUserName
+            | some nid => pure nid
+          trace[Progress] "\n- prevId: {prevId}\n- nid: {nid}\n- remaining ids: {ids}"
+          if ← isConj (← inferType hPost) then
+            splitConjTac hPost (some (prevId, nid)) (λ _ nhPost => splitPostWithIds nid nhPost ids)
+          else return (.Error m!"Too many ids provided ({ids0}) not enough conjuncts to split in the postcondition")
+      let curPostId := (← hPost.fvarId!.getDecl).userName
+      splitPostWithIds curPostId hPost ids
+  match res with
+  | .Error _ => return res -- Can we get there? We're using "return"
+  | .Ok =>
+    -- Update the set of goals
+    let curGoals ← getUnsolvedGoals
+    let newGoals := mvars.map Expr.mvarId!
+    let newGoals ← newGoals.filterM fun mvar => not <$> mvar.isAssigned
+    trace[Progress] "new goals: {newGoals}"
+    setGoals newGoals.toList
+    allGoals asmTac
+    let newGoals ← getUnsolvedGoals
+    setGoals (newGoals ++ curGoals)
+    trace[Progress] "progress: replaced the goals"
+    --
+    pure .Ok
+
+-- Small utility: if `args` is not empty, return the name of the app in the first
+-- arg, if it is a const.
+def getFirstArgAppName (args : Array Expr) : MetaM (Option Name) := do
+  if args.size = 0 then pure none
+  else
+    (args.get! 0).withApp fun f _ => do
+    if f.isConst then pure (some f.constName)
+    else pure none
+
+def getFirstArg (args : Array Expr) : Option Expr := do
+  if args.size = 0 then none
+  else some (args.get! 0)
+
+/- Helper: try to lookup a theorem and apply it, or continue with another tactic
+   if it fails -/
+def tryLookupApply (keep : Option Name) (ids : Array (Option 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
+    | none =>
+      trace[Progress] "Could not find a {kind}"
+      pure none
+    | some th => do
+      trace[Progress] "Lookuped up {kind}: {th}"
+      -- Apply the theorem
+      let res ← do
+        try
+          let res ← progressWith fExpr th keep ids splitPost asmTac
+          pure (some res)
+        catch _ => none
+  match res with
+  | some .Ok => return ()
+  | some (.Error msg) => throwError msg
+  | none => x
+
+-- The array of ids are identifiers to use when introducing fresh variables
+def progressAsmsOrLookupTheorem (keep : Option Name) (withTh : Option TheoremOrLocal)
+  (ids : Array (Option Name)) (splitPost : Bool) (asmTac : TacticM Unit) : TacticM Unit := do
+  withMainContext do
+  -- Retrieve the goal
+  let mgoal ← Tactic.getMainGoal
+  let goalTy ← mgoal.getType
+  trace[Progress] "goal: {goalTy}"
+  -- Dive into the goal to lookup the theorem
+  let (fExpr, fName, args) ← do
+    withPSpec goalTy fun desc =>
+    -- TODO: check that no quantified variables in the arguments
+    pure (desc.fExpr, desc.fName, desc.args)
+  trace[Progress] "Function: {fName}"
+  -- If the user provided a theorem/assumption: use it.
+  -- Otherwise, lookup one.
+  match withTh with
+  | some th => do
+    match ← progressWith fExpr th keep ids splitPost asmTac with
+    | .Ok => return ()
+    | .Error msg => throwError msg
+  | none =>
+    -- Try all the assumptions one by one and if it fails try to lookup a theorem.
+    let ctx ← Lean.MonadLCtx.getLCtx
+    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 splitPost asmTac catch _ => continue
+      match res with
+      | .Ok => return ()
+      | .Error msg => throwError msg
+    -- It failed: try to lookup a theorem
+    -- TODO: use a list of theorems, and try them one by one?
+    trace[Progress] "No assumption succeeded: trying to lookup a theorem"
+    let pspec ← do
+      let thName ← pspecAttr.find? fName
+      pure (thName.map fun th => .Theorem th)
+    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
+      match getFirstArg args with
+      | none => pure none
+      | some arg => do
+        let thName ← pspecClassExprAttr.find? fName arg
+        pure (thName.map fun th => .Theorem th)
+    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
+      | none => pure none
+      | some argName => do
+        let thName ← pspecClassAttr.find? fName argName
+        pure (thName.map fun th => .Theorem th)
+    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
+    let decls := decls.filter (λ decl => match decl.kind with
+      | .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 splitPost asmTac catch _ => continue
+      match res with
+      | .Ok => return ()
+      | .Error msg => throwError msg
+    -- Nothing worked: failed
+    throwError "Progress failed"
+
+syntax progressArgs := ("keep" (ident <|> "_"))? ("with" ident)? ("as" " ⟨ " (ident <|> "_"),* " .."? " ⟩")?
+
+def evalProgress (args : TSyntax `Progress.progressArgs) : TacticM Unit := do
+  let args := args.raw
+  -- Process the arguments to retrieve the identifiers to use
+  trace[Progress] "Progress arguments: {args}"
+  let (keepArg, withArg, asArgs) ←
+    match args.getArgs.toList with
+    | [keepArg, withArg, asArgs] => do pure (keepArg, withArg, asArgs)
+    | _ => throwError "Unexpected: invalid arguments"
+  let keep : Option Name ← do
+    trace[Progress] "Keep arg: {keepArg}"
+    let args := keepArg.getArgs
+    if args.size > 0 then do
+      trace[Progress] "Keep args: {args}"
+      let arg := args.get! 1
+      trace[Progress] "Keep arg: {arg}"
+      if arg.isIdent then pure (some arg.getId)
+      else do pure (some (← mkFreshAnonPropUserName))
+    else do pure none
+  trace[Progress] "Keep: {keep}"
+  let withArg ← do
+    let withArg := withArg.getArgs
+    if withArg.size > 0 then
+      let id := withArg.get! 1
+      trace[Progress] "With arg: {id}"
+      -- Attempt to lookup a local declaration
+      match (← getLCtx).findFromUserName? id.getId with
+      | some decl => do
+        trace[Progress] "With arg: local decl"
+        pure (some (.Local decl))
+      | none => do
+        -- Not a local declaration: should be a theorem
+        trace[Progress] "With arg: theorem"
+        addCompletionInfo <| CompletionInfo.id id id.getId (danglingDot := false) {} none
+        let cs ← resolveGlobalConstWithInfos id
+        match cs with
+        | [] => throwError "Could not find theorem {id}"
+        | id :: _ =>
+          pure (some (.Theorem id))
+    else pure none
+  let ids :=
+    let args := asArgs.getArgs
+    let args := (args.get! 2).getSepArgs
+    args.map (λ s => if s.isIdent then some s.getId else none)
+  trace[Progress] "User-provided ids: {ids}"
+  let splitPost : Bool :=
+    let args := asArgs.getArgs
+    (args.get! 3).getArgs.size > 0
+  trace[Progress] "Split post: {splitPost}"
+  /- For scalarTac we have a fast track: if the goal is not a linear
+     arithmetic goal, we skip (note that otherwise, scalarTac would try
+     to prove a contradiction) -/
+  let scalarTac : TacticM Unit := do
+    if ← Arith.goalIsLinearInt then
+      -- Also: we don't try to split the goal if it is a conjunction
+      -- (it shouldn't be)
+      Arith.scalarTac false
+    else
+      throwError "Not a linear arithmetic goal"
+  progressAsmsOrLookupTheorem keep withArg ids splitPost (
+    withMainContext do
+    trace[Progress] "trying to solve assumption: {← getMainGoal}"
+    firstTac [assumptionTac, scalarTac])
+  trace[Diverge] "Progress done"
+
+elab "progress" args:progressArgs : tactic =>
+  evalProgress args
+
+namespace Test
+  open Primitives Result
+
+  set_option trace.Progress true
+  set_option pp.rawOnError true
+
+  #eval showStoredPSpec
+  #eval showStoredPSpecClass
+
+  example {ty} {x y : Scalar ty}
+    (hmin : Scalar.min ty ≤ x.val + y.val)
+    (hmax : x.val + y.val ≤ Scalar.max ty) :
+    ∃ z, x + y = ret z ∧ z.val = x.val + y.val := by
+    progress keep _ as ⟨ z, h1 .. ⟩
+    simp [*, h1]
+
+  example {ty} {x y : Scalar ty}
+    (hmin : Scalar.min ty ≤ x.val + y.val)
+    (hmax : x.val + y.val ≤ Scalar.max ty) :
+    ∃ z, x + y = ret z ∧ z.val = x.val + y.val := by
+    progress keep h with Scalar.add_spec as ⟨ z ⟩
+    simp [*, h]
+
+  /- Checking that universe instantiation works: the original spec uses
+     `α : Type u` where u is quantified, while here we use `α : Type 0` -/
+  example {α : Type} (v: Vec α) (i: Usize) (x : α)
+    (hbounds : i.val < v.length) :
+    ∃ nv, v.index_mut_back α i x = ret nv ∧
+    nv.val = v.val.update i.val x := by
+    progress
+    simp [*]
+
+end Test
+
+end Progress