fix: First miss-handling of free variables

leanprover · Oct 17, 2023 · bc9cba1 · bc9cba1
1 parent f15d561
commit bc9cba1
Showing 1 changed file with 35 additions and 31 deletions.
diff --git a/DocGen4/Process/NameInfo.lean b/DocGen4/Process/NameInfo.lean
@@ -15,38 +15,42 @@ def NameInfo.ofTypedName (n : Name) (t : Expr) : MetaM NameInfo := do
   let env ← getEnv
   return { name := n, type := ← prettyPrintTerm t, doc := ← findDocString? env n}
 
-partial def typeToArgsType (e : Expr) : (Array ((Option Name) × Expr × BinderInfo) × Expr) :=
-  let helper := fun name type body data =>
-    -- Once we hit a name with a macro scope we stop traversing the expression
-    -- and print what is left after the : instead. The only exception
-    -- to this is instances since these almost never have a name
-    -- but should still be printed as arguments instead of after the :.
-    if data.isInstImplicit && name.hasMacroScopes then
-      let arg := (none, type, data)
-      let (args, final) := typeToArgsType (Expr.instantiate1 body (mkFVar ⟨name⟩))
-      (#[arg] ++ args, final)
-    else if name.hasMacroScopes then
-      (#[], e)
-    else
-      let arg := (some name, type, data)
-      let (args, final) := typeToArgsType (Expr.instantiate1 body (mkFVar ⟨name⟩))
-      (#[arg] ++ args, final)
-  match e.consumeMData with
-  | Expr.forallE name type body binderInfo => helper name type body binderInfo
-  | _ => (#[], e)
+partial def argTypeTelescope {α : Type} (e : Expr) (k : Array ((Option Name) × Expr × BinderInfo) → Expr → MetaM α) : MetaM α :=
+  go e #[]
+where
+  go (e : Expr) (args : Array ((Option Name) × Expr × BinderInfo)) : MetaM α := do
+    let helper := fun name type body bi =>
+      -- Once we hit a name with a macro scope we stop traversing the expression
+      -- and print what is left after the : instead. The only exception
+      -- to this is instances since these almost never have a name
+      -- but should still be printed as arguments instead of after the :.
+      if bi.isInstImplicit && name.hasMacroScopes then
+        let arg := (none, type, bi)
+        Meta.withLocalDecl name bi type fun fvar => do
+          go (Expr.instantiate1 body fvar) (args.push arg)
+      else if name.hasMacroScopes then
+        k args e
+      else
+        let arg := (some name, type, bi)
+        Meta.withLocalDecl name bi type fun fvar => do
+          go (Expr.instantiate1 body fvar) (args.push arg)
+    match e.consumeMData with
+    | Expr.forallE name type body binderInfo => helper name type body binderInfo
+    | _ => k args e
 
 def Info.ofConstantVal (v : ConstantVal) : MetaM Info := do
-  let (args, type) := typeToArgsType v.type
-  let args ← args.mapM (fun (n, e, b) => do return Arg.mk n (← prettyPrintTerm e) b)
-  let nameInfo ← NameInfo.ofTypedName v.name type
-  match ← findDeclarationRanges? v.name with
-  -- TODO: Maybe selection range is more relevant? Figure this out in the future
-  | some range => return {
-      toNameInfo := nameInfo,
-      args,
-      declarationRange := range.range,
-      attrs := ← getAllAttributes v.name
-    }
-  | none => panic! s!"{v.name} is a declaration without position"
+  argTypeTelescope v.type fun args type => do
+    let args ← args.mapM (fun (n, e, b) => do return Arg.mk n (← prettyPrintTerm e) b)
+    let nameInfo ← NameInfo.ofTypedName v.name type
+    match ← findDeclarationRanges? v.name with
+    -- TODO: Maybe selection range is more relevant? Figure this out in the future
+    | some range =>
+      return {
+        toNameInfo := nameInfo,
+        args,
+        declarationRange := range.range,
+        attrs := ← getAllAttributes v.name
+      }
+    | none => panic! s!"{v.name} is a declaration without position"
 
 end DocGen4.Process