runtime(lua): improve foldexpr, add vim9 script version

closes: #17049

Signed-off-by: Konfekt <Konfekt@users.noreply.github.com>
Signed-off-by: Christian Brabandt <cb@256bit.org>
diff --git a/runtime/ftplugin/lua.vim b/runtime/ftplugin/lua.vim
index cc042c5..87f54dd 100644
--- a/runtime/ftplugin/lua.vim
+++ b/runtime/ftplugin/lua.vim
@@ -1,12 +1,13 @@
 " Vim filetype plugin file.
-" Language:		Lua
-" Maintainer:		Doug Kearns <dougkearns@gmail.com>
-" Previous Maintainer:	Max Ischenko <mfi@ukr.net>
-" Contributor:		Dorai Sitaram <ds26@gte.com>
-"			C.D. MacEachern <craig.daniel.maceachern@gmail.com>
-"			Tyler Miller <tmillr@proton.me>
-"			Phạm Bình An <phambinhanctb2004@gmail.com>
-" Last Change:		2025 Feb 27
+" Language:             Lua
+" Maintainer:           Doug Kearns <dougkearns@gmail.com>
+" Previous Maintainer:  Max Ischenko <mfi@ukr.net>
+" Contributor:          Dorai Sitaram <ds26@gte.com>
+"                       C.D. MacEachern <craig.daniel.maceachern@gmail.com>
+"                       Tyler Miller <tmillr@proton.me>
+"                       Phạm Bình An <phambinhanctb2004@gmail.com>
+"                       @konfekt
+" Last Change:          2025 Apr 04
 
 if exists("b:did_ftplugin")
   finish
@@ -41,11 +42,11 @@
 if exists("loaded_matchit") && !exists("b:match_words")
   let b:match_ignorecase = 0
   let b:match_words =
-	\ '\<\%(do\|function\|if\)\>:' ..
-	\ '\<\%(return\|else\|elseif\)\>:' ..
-	\ '\<end\>,' ..
-	\ '\<repeat\>:\<until\>,' ..
-	\ '\%(--\)\=\[\(=*\)\[:]\1]'
+        \ '\<\%(do\|function\|if\)\>:' ..
+        \ '\<\%(return\|else\|elseif\)\>:' ..
+        \ '\<end\>,' ..
+        \ '\<repeat\>:\<until\>,' ..
+        \ '\%(--\)\=\[\(=*\)\[:]\1]'
   let b:undo_ftplugin ..= " | unlet! b:match_words b:match_ignorecase"
 endif
 
@@ -61,7 +62,7 @@
 
 if has("folding") && get(g:, "lua_folding", 0)
   setlocal foldmethod=expr
-  setlocal foldexpr=s:LuaFold(v:lnum)
+  setlocal foldexpr=s:LuaFold()
   let b:lua_lasttick = -1
   let b:undo_ftplugin ..= " | setl foldexpr< foldmethod< | unlet! b:lua_lasttick b:lua_foldlists"
 endif
@@ -87,19 +88,19 @@
 endfunction
 
 let s:patterns = [
-      \ ['do', 'end'],
-      \ ['if\s+.+\s+then', 'end'],
-      \ ['repeat', 'until\s+.+'],
-      \ ['for\s+.+\s+do', 'end'],
-      \ ['while\s+.+\s+do', 'end'],
-      \ ['function.+', 'end'],
-      \ ['return\s+function.+', 'end'],
-      \ ['local\s+function\s+.+', 'end'],
-      \ ]
+    \ ['do', 'end'],
+    \ ['if\s+.+\s+then', 'end'],
+    \ ['repeat', 'until\s+.+'],
+    \ ['for\s+.+\s+do', 'end'],
+    \ ['while\s+.+\s+do', 'end'],
+    \ ['function.+', 'end'],
+    \ ['return\s+function.+', 'end'],
+    \ ['local\s+function\s+.+', 'end'],
+    \ ]
 
-function s:LuaFold(lnum) abort
+function s:LuaFold() abort
   if b:lua_lasttick == b:changedtick
-    return b:lua_foldlists[a:lnum - 1]
+    return b:lua_foldlists[v:lnum - 1]
   endif
   let b:lua_lasttick = b:changedtick
 
@@ -108,27 +109,78 @@
   let buf = getline(1, "$")
   for line in buf
     for t in s:patterns
+      let open = 0
+      let end = 0
       let tagopen  = '\v^\s*' .. t[0] ..'\s*$'
-      let tagclose = '\v^\s*' .. t[1] ..'\s*$'
+      let tagend = '\v^\s*' .. t[1] ..'\s*$'
       if line =~# tagopen
-	call add(foldlist, t)
-	break
-      elseif line =~# tagclose
-	if len(foldlist) > 0 && line =~# foldlist[-1][1]
-	  call remove(foldlist, -1)
-	else
-	  let foldlist = []
-	endif
-	break
+        call add(foldlist, t)
+        let open = 1
+        break
+      elseif line =~# tagend
+        if len(foldlist) > 0 && line =~# foldlist[-1][1]
+          call remove(foldlist, -1)
+          let end = 1
+        else
+          let foldlist = []
+        endif
+        break
       endif
     endfor
-    call add(b:lua_foldlists, len(foldlist))
+    let prefix = ""
+    if open == 1 | let prefix = ">" | endif
+    if end == 1 | let prefix = "<" | endif
+    let b:lua_foldlists += [prefix..(len(foldlist) + end)]
   endfor
 
-  return lua_foldlists[a:lnum - 1]
+  return b:lua_foldlists[v:lnum - 1]
 endfunction
 
+if !has('vim9script')
+  let &cpo = s:cpo_save
+  unlet s:cpo_save
+
+  finish
+endif
+
+delfunction! s:LuaFold
+def s:LuaFold(): string
+  if b:lua_lasttick == b:changedtick
+    return b:lua_foldlists[v:lnum - 1]
+  endif
+  b:lua_lasttick = b:changedtick
+
+  b:lua_foldlists = []
+  var foldlist = []
+  var buf = getline(1, "$")
+  for line in buf
+    var open = 0
+    var end = 0
+    for t in patterns
+      var tagopen  = '\v^\s*' .. t[0] .. '\s*$'
+      var tagend = '\v^\s*' .. t[1] .. '\s*$'
+      if line =~# tagopen
+        add(foldlist, t)
+        open = 1
+        break
+      elseif line =~# tagend
+        if len(foldlist) > 0 && line =~# foldlist[-1][1]
+          end = 1
+          remove(foldlist, -1)
+        else
+          foldlist = []
+        endif
+        break
+      endif
+    endfor
+    var prefix = ""
+    if open == 1 | prefix = ">" | endif
+    if end == 1 | prefix = "<" | endif
+    b:lua_foldlists += [prefix .. (len(foldlist) + end)]
+  endfor
+  return b:lua_foldlists[v:lnum - 1]
+enddef
+
 let &cpo = s:cpo_save
 unlet s:cpo_save
-
 " vim: nowrap sw=2 sts=2 ts=8 noet: