From 346378b3a37b06ccebee6b9bae414e8e2b0d5085 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sun, 23 Feb 2020 10:35:14 -0800
Subject: [PATCH] feat: add `[recursor <major-pos>]` attribute

---
 src/Init/Lean/Attributes.lean        |  7 ++++--
 src/Init/Lean/Meta/Basic.lean        |  5 ++++
 src/Init/Lean/Meta/RecursorInfo.lean | 35 +++++++++++++++++++++++++++-
 3 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/src/Init/Lean/Attributes.lean b/src/Init/Lean/Attributes.lean
index e2f807c009..0df48e33b6 100644
--- a/src/Init/Lean/Attributes.lean
+++ b/src/Init/Lean/Attributes.lean
@@ -309,8 +309,10 @@ structure ParametricAttribute (α : Type) :=
 (ext  : PersistentEnvExtension (Name × α) (Name × α) (NameMap α))
 
 def registerParametricAttribute {α : Type} [Inhabited α] (name : Name) (descr : String)
-       (getParam : Environment → Name → Syntax → Except String α)
-       (afterSet : Environment → Name → α → Except String Environment := fun env _ _ => Except.ok env) : IO (ParametricAttribute α) := do
+    (getParam : Environment → Name → Syntax → Except String α)
+    (afterSet : Environment → Name → α → Except String Environment := fun env _ _ => Except.ok env)
+    (appTime := AttributeApplicationTime.afterTypeChecking)
+    : IO (ParametricAttribute α) := do
 ext : PersistentEnvExtension (Name × α) (Name × α) (NameMap α) ← registerPersistentEnvExtension {
   name            := name,
   mkInitial       := pure {},
@@ -324,6 +326,7 @@ ext : PersistentEnvExtension (Name × α) (Name × α) (NameMap α) ← register
 let attrImpl : AttributeImpl := {
   name  := name,
   descr := descr,
+  applicationTime := appTime,
   add   := fun env decl args persistent => do
     unless persistent $ throw (IO.userError ("invalid attribute '" ++ toString name ++ "', must be persistent"));
     unless (env.getModuleIdxFor? decl).isNone $
diff --git a/src/Init/Lean/Meta/Basic.lean b/src/Init/Lean/Meta/Basic.lean
index 54e7df5be1..8396036296 100644
--- a/src/Init/Lean/Meta/Basic.lean
+++ b/src/Init/Lean/Meta/Basic.lean
@@ -805,6 +805,11 @@ finally x (modify $ fun s => { mctx := mctx', .. s })
 @[init] private def regTraceClasses : IO Unit :=
 registerTraceClass `Meta
 
+def run {α} (env : Environment) (x : MetaM α) (maxRecDepth := 10000) : Except Exception α :=
+match x { maxRecDepth := maxRecDepth, currRecDepth := 0 } { env := env } with
+| EStateM.Result.ok a _     => Except.ok a
+| EStateM.Result.error ex _ => Except.error ex
+
 end Meta
 end Lean
 
diff --git a/src/Init/Lean/Meta/RecursorInfo.lean b/src/Init/Lean/Meta/RecursorInfo.lean
index c073affe1e..10d9a795f1 100644
--- a/src/Init/Lean/Meta/RecursorInfo.lean
+++ b/src/Init/Lean/Meta/RecursorInfo.lean
@@ -8,6 +8,7 @@ import Init.Data.Nat.Control
 import Init.Lean.AuxRecursor
 import Init.Lean.Util.FindExpr
 import Init.Lean.Meta.ExprDefEq
+import Init.Lean.Meta.Message
 
 namespace Lean
 namespace Meta
@@ -264,11 +265,43 @@ forallTelescopeReducing cinfo.type $ fun xs type => type.withApp $ fun motive mo
     ("invalid user defined recursor '" ++ toString declName
      ++ "', type of the major premise must be of the form (I ...), where I is a constant")
 
-def mkRecursorInfo (declName : Name) (majorPos? : Option Nat := none) : MetaM RecursorInfo := do
+private def syntaxToMajorPos (stx : Syntax) : Except String Nat :=
+match stx with
+| Syntax.node _ args =>
+  if args.size == 0 then Except.error "unexpected kind of argument"
+  else match (args.get! 0).isNatLit? with
+    | some idx =>
+      if idx == 0 then Except.error "major premise position must be greater than 0"
+      else pure (idx - 1)
+    | none => Except.error "unexpected kind of argument"
+| _ => Except.error "unexpected kind of argument"
+
+private def mkRecursorInfoCore (declName : Name) (majorPos? : Option Nat := none) : MetaM RecursorInfo := do
 cinfo ← getConstInfo declName;
 match cinfo with
 | ConstantInfo.recInfo val => mkRecursorInfoForKernelRec declName val
 | _                        => mkRecursorInfoAux cinfo majorPos?
 
+def mkRecursorAttr : IO (ParametricAttribute Nat) :=
+registerParametricAttribute `recursor "user defined recursor, numerical parameter specifies position of the major premise"
+  (fun _ _ => syntaxToMajorPos)
+  (fun env declName majorPos =>
+    match Meta.run env (mkRecursorInfoCore declName (some majorPos)) with
+    | Except.ok _     => pure env
+    | Except.error ex => Except.error $ toString ex)
+
+@[init mkRecursorAttr] constant recursorAttribute : ParametricAttribute Nat := arbitrary _
+
+def getMajorPos? (env : Environment) (declName : Name) : Option Nat :=
+recursorAttribute.getParam env declName
+
+def mkRecursorInfo (declName : Name) (majorPos? : Option Nat := none) : MetaM RecursorInfo := do
+cinfo ← getConstInfo declName;
+match cinfo with
+| ConstantInfo.recInfo val => mkRecursorInfoForKernelRec declName val
+| _                        => match majorPos? with
+  | none => do env ← getEnv; mkRecursorInfoAux cinfo (getMajorPos? env declName)
+  | _    => mkRecursorInfoAux cinfo majorPos?
+
 end Meta
 end Lean