diff --git a/src/Lean/Elab/Match.lean b/src/Lean/Elab/Match.lean index c802fac250..1f7947d945 100644 --- a/src/Lean/Elab/Match.lean +++ b/src/Lean/Elab/Match.lean @@ -422,6 +422,54 @@ private def withMVar (mvarId : MVarId) (x : M α) : M α := do else x +/-- + Creating a mapping containing `b ↦ e'` where `patternWithRef e' = some (stx, b)`, + and `e'` is a subterm of `e`. + + This is a helper function for `whnfPreservingPatternRef`. -/ +private def mkPatternRefMap (e : Expr) : ExprMap Expr := + runST go +where + go (σ) : ST σ (ExprMap Expr) := do + let map : ST.Ref σ (ExprMap Expr) ← ST.mkRef {} + e.forEach fun e => do + match patternWithRef? e with + | some (ref, b) => map.modify (·.insert b e) + | none => return () + map.get + +/-- + Try to restore `Syntax` ref information stored in `map` after + applying `whnf` at `whnfPreservingPatternRef`. + It assumes `map` has been constructed using `mkPatternRefMap`. +-/ +private def applyRefMap (e : Expr) (map : ExprMap Expr) : Expr := + e.replace fun e => + match patternWithRef? e with + | some _ => some e -- stop `e` already has annotation + | none => match map.find? e with + | some eWithRef => some eWithRef -- stop `e` found annotation + | none => none -- continue + +/-- + Applies `whnf` but tries to preserve `PatternWithRef` information. + This is a bit hackish, but it is necessary for providing proper + jump-to-definition information in examples such as + ``` + def f (x : Nat) : Nat := + match x with + | 0 => 1 + | y + 1 => y + ``` + Without this trick, the `PatternWithRef` is lost for the `y` at the pattern `y+1`. +-/ +private def whnfPreservingPatternRef (e : Expr) : MetaM Expr := do + let eNew ← whnf e + if eNew.isConstructorApp (← getEnv) then + return eNew + else + return applyRefMap eNew (mkPatternRefMap e) + /-- Normalize the pattern and collect all patterns variables (explicit and implicit). This method is the one that decides where the inaccessible annotations must be inserted. @@ -464,7 +512,7 @@ partial def normalize (e : Expr) : M Expr := do else throwInvalidPattern e else - let eNew ← whnf e + let eNew ← whnfPreservingPatternRef e if eNew != e then normalize eNew else diff --git a/tests/lean/interactive/definition.lean b/tests/lean/interactive/definition.lean index f302e1fe66..00aeb837d5 100644 --- a/tests/lean/interactive/definition.lean +++ b/tests/lean/interactive/definition.lean @@ -4,3 +4,9 @@ example : Foo := let c := Foo.foo c --^ textDocument/typeDefinition + +def f (x : Nat) : Nat := + match x with + | 0 => 1 + | y + 1 => y + --^ textDocument/declaration diff --git a/tests/lean/interactive/definition.lean.expected.out b/tests/lean/interactive/definition.lean.expected.out index de7824b6d1..b7d96fb332 100644 --- a/tests/lean/interactive/definition.lean.expected.out +++ b/tests/lean/interactive/definition.lean.expected.out @@ -7,3 +7,13 @@ {"start": {"line": 0, "character": 0}, "end": {"line": 1, "character": 7}}, "originSelectionRange": {"start": {"line": 4, "character": 2}, "end": {"line": 4, "character": 3}}}] +{"textDocument": {"uri": "file://definition.lean"}, + "position": {"line": 10, "character": 13}} +[{"targetUri": "file://definition.lean", + "targetSelectionRange": + {"start": {"line": 10, "character": 4}, "end": {"line": 10, "character": 5}}, + "targetRange": + {"start": {"line": 10, "character": 4}, "end": {"line": 10, "character": 5}}, + "originSelectionRange": + {"start": {"line": 10, "character": 13}, + "end": {"line": 10, "character": 14}}}]