From 4b216e8d2a376f95f486f661e28603a9651b0a4f Mon Sep 17 00:00:00 2001
From: Julian Berman <Julian@GrayVines.com>
Date: Tue, 31 Dec 2024 08:53:34 -0500
Subject: [PATCH] Also jump between do/return.

(And therefore highlight them).

Requested on Zulip: https://leanprover.zulipchat.com/#narrow/channel/113488-general/topic/lean.2Envim/near/491389017
---
 ftplugin/lean/lean.lua |  1 +
 spec/matchit_spec.lua  | 28 ++++++++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/ftplugin/lean/lean.lua b/ftplugin/lean/lean.lua
index e8420f46..c880bc78 100644
--- a/ftplugin/lean/lean.lua
+++ b/ftplugin/lean/lean.lua
@@ -27,6 +27,7 @@ if vim.g.loaded_matchit and not vim.b.match_words then
       [[\<\%(namespace\|section\)\s\+\([^«»]\{-}\>\|«.\{-}»\):\<end\s\+\1]],
       [[^\s*section\s*$:^end\s*$]],
       [[\<if\>:\<then\>:\<else\>]],
+      [[\<do\>:\<return\>]],
     })
     :join ','
 end
diff --git a/spec/matchit_spec.lua b/spec/matchit_spec.lua
index c976c5ed..cf84547d 100644
--- a/spec/matchit_spec.lua
+++ b/spec/matchit_spec.lua
@@ -105,4 +105,32 @@ describe('matchit', function()
       end
     )
   )
+
+  it(
+    'jumps between do/return',
+    helpers.clean_buffer(
+      [[
+        example : Nat := Id.run do
+          let _ ← Id.run do return 1
+          return 2
+      ]],
+      function()
+        vim.cmd.normal '1$'
+        assert.current_word.is 'do'
+
+        vim.cmd.normal '%'
+        assert.current_line.is '  return 2'
+
+        vim.cmd.normal '%'
+        assert.current_line.is 'example : Nat := Id.run do'
+
+        vim.cmd.normal '2gg17|'
+        assert.current_word.is 'do'
+
+        vim.cmd.normal '%'
+        assert.current_word.is 'return'
+        assert.current_line.is '  let _ ← Id.run do return 1'
+      end
+    )
+  )
 end)