[nongnu] elpa/treesit-fold ad1d9b2412 258/417: feat(fold): Add support for Lua (#52)

[ELPA Syncer]

branch: elpa/treesit-fold
commit ad1d9b24127fe0105642790a1cacc779d70ec7a0
Author: Jen-Chieh Shen <jcs090218@gmail.com>
Commit: GitHub <noreply@github.com>

    feat(fold): Add support for Lua (#52)
    
    * feat(fold): Add support for Lua
    
    * wip
    
    * Display nicely on next line
    
    * Fix comiple error
    
    * complete
---
 CHANGELOG.md       |   1 +
 README.md          |   1 +
 ts-fold-parsers.el |  23 +++++++++-
 ts-fold-util.el    |   5 ++-
 ts-fold.el         | 124 ++++++++++++++++++++++++++++++++++++++++++++++++++---
 5 files changed, 146 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 19deaf481e..310749db7c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,7 @@ Check [Keep a Changelog](http://keepachangelog.com/) for 
recommendations on how
 * Change global mode to turn on with tree-sitter (#41)
 * Add minor-mode `ts-fold-line-comment-mode` for line comment folding (#45)
 * Improve folding for C preprocessor operators (#46)
+* Add support for Lua (#52)
 
 ## 0.1.0
 > Released Oct 18, 2021
diff --git a/README.md b/README.md
index 9e0ab75151..eab80abbb9 100644
--- a/README.md
+++ b/README.md
@@ -119,6 +119,7 @@ These languages are fairly complete:
 - Go
 - HTML
 - Java / JavaScript / JSX / JSON / Julia
+- Lua
 - Nix
 - PHP / Python
 - R / Ruby / Rust
diff --git a/ts-fold-parsers.el b/ts-fold-parsers.el
index aa26a183a4..3815ae8d99 100644
--- a/ts-fold-parsers.el
+++ b/ts-fold-parsers.el
@@ -46,6 +46,13 @@
 (declare-function ts-fold-range-c-preproc-else "ts-fold.el")
 (declare-function ts-fold-range-html "ts-fold.el")
 (declare-function ts-fold-range-julia "ts-fold.el")
+(declare-function ts-fold-range-lua-comment "ts-fold.el")
+(declare-function ts-fold-range-lua-function "ts-fold.el")
+(declare-function ts-fold-range-lua-if "ts-fold.el")
+(declare-function ts-fold-range-lua-elseif "ts-fold.el")
+(declare-function ts-fold-range-lua-else "ts-fold.el")
+(declare-function ts-fold-range-lua-do-loop "ts-fold.el")
+(declare-function ts-fold-range-lua-repeat "ts-fold.el")
 (declare-function ts-fold-range-ocaml "ts-fold.el")
 (declare-function ts-fold-range-python "ts-fold.el")
 (declare-function ts-fold-range-ruby-class-def "ts-fold.el")
@@ -178,6 +185,18 @@
     (try_statement       . (ts-fold-range-seq 2 -2))
     (while_statement     . ts-fold-range-julia)))
 
+(defun ts-fold-parsers-lua ()
+  "Rule set for Lua."
+  '((expression_list      . ts-fold-range-seq)
+    (function_declaration . ts-fold-range-lua-function)
+    (if_statement         . ts-fold-range-lua-if)
+    (elseif_statement     . ts-fold-range-lua-elseif)
+    (else_statement       . ts-fold-range-lua-else)
+    (while_statement      . ts-fold-range-lua-do-loop)
+    (for_statement        . ts-fold-range-lua-do-loop)
+    (repeat_statement     . ts-fold-range-lua-repeat)
+    (comment              . ts-fold-range-lua-comment)))
+
 (defun ts-fold-parsers-nix ()
   "Rule set for Nix."
   '((attrset       . ts-fold-range-seq)
@@ -223,8 +242,8 @@
   '((class    . ts-fold-range-ruby-class-def)
     (method   . ts-fold-range-ruby-class-def)
     (array    . ts-fold-range-seq)
-    (do       . (ts-fold-range-seq 1 -2))  ; match with `end`
-    (do_block . (ts-fold-range-seq 1 -2))  ; match with `end`, in spec file
+    (do       . (ts-fold-range-seq 1 -2))     ; match with `end`
+    (do_block . (ts-fold-range-seq 1 -2))     ; match with `end`, in spec file
     (then     . ts-fold-range-ruby-if)        ; `if` and `elsif` block
     (else     . (ts-fold-range-ruby-if 4 0))  ; `else` block
     (comment
diff --git a/ts-fold-util.el b/ts-fold-util.el
index 929030f936..b3500c612a 100644
--- a/ts-fold-util.el
+++ b/ts-fold-util.el
@@ -116,7 +116,10 @@ then return that list."
                     (ts-fold-get-children node)))
 
 (defun ts-fold-find-children-traverse (node type)
-  "Like function `ts-fold-find-children' but traverse it."
+  "Like function `ts-fold-find-children' but traverse it.
+
+For arguments NODE and TYPE, see function `ts-fold-find-children' for more
+information."
   (cl-remove-if-not (lambda (child) (ts-fold--compare-type child type))
                     (ts-fold-get-children-traverse node)))
 
diff --git a/ts-fold.el b/ts-fold.el
index 49f2574973..ed060af687 100644
--- a/ts-fold.el
+++ b/ts-fold.el
@@ -77,6 +77,7 @@
     (json-mode       . ,(ts-fold-parsers-json))
     (jsonc-mode      . ,(ts-fold-parsers-json))
     (julia-mode      . ,(ts-fold-parsers-julia))
+    (lua-mode        . ,(ts-fold-parsers-lua))
     (nix-mode        . ,(ts-fold-parsers-nix))
     (ocaml-mode      . ,(ts-fold-parsers-ocaml))
     (php-mode        . ,(ts-fold-parsers-php))
@@ -106,6 +107,14 @@ the fold in a cons cell.  See `ts-fold-range-python' for 
an example."
   :type 'hook
   :group 'ts-fold)
 
+(defcustom ts-fold-on-next-line t
+  "If non-nil, we leave ending keywords on the next line.
+
+This is only used in languages that uses keyword to end the scope.
+For example, Lua, Ruby, etc."
+  :type 'boolean
+  :group 'ts-fold)
+
 (defcustom ts-fold-replacement "..."
   "Show this string instead of the folded text."
   :type 'string
@@ -403,7 +412,7 @@ Argument OFFSET can be used to tweak the final beginning 
and end position."
     (ts-fold--cons-add (cons beg end) offset)))
 
 (defun ts-fold-range-markers (node offset start-seq &optional end-seq)
-  "Returns the fold range for NODE with an OFFSET where the range starts at
+  "Return the fold range for NODE with an OFFSET where the range starts at
 the end of the first occurence of START-SEQ and ends at the end of the node
 or the start of the last occurence of the optional parameter LAST-SEQ.
 
@@ -456,6 +465,12 @@ more information."
           (ts-fold-range-line-comment node offset "///")
         (ts-fold-range-line-comment node offset "//")))))
 
+(defun ts-fold-point-before-line-break (pos)
+  "Go to POS then find previous line break, and return its position."
+  (save-excursion
+    (goto-char pos)
+    (max 1 (1- (line-beginning-position)))))
+
 ;;
 ;; (@* "Languages" )
 ;;
@@ -607,16 +622,24 @@ more information."
                               (tsc-get-child-by-field node :parameters)
                               (tsc-get-child-by-field node :name)))
               (beg (tsc-node-end-position named-node))
-              (end (tsc-node-end-position node)))
-    (ts-fold--cons-add (cons beg (- end 3)) offset)))
+              (end (tsc-node-end-position node))
+              (end (- end 3)))
+    (when ts-fold-on-next-line  ; display nicely
+      (setq end (ts-fold-point-before-line-break end)))
+    (ts-fold--cons-add (cons beg end) offset)))
 
 (defun ts-fold-range-ruby-if (node offset)
   "Define fold range for `if' (then), `elsif', and `else' in Ruby.
 
 For arguments NODE and OFFSET, see function `ts-fold-range-seq' for
 more information."
-  (when-let ((beg (tsc-node-start-position node))
-             (end (tsc-node-end-position node)))
+  (when-let* ((beg (tsc-node-start-position node))
+              (end (cond ((when-let ((next (tsc-get-next-sibling node)))
+                            (tsc-node-start-position next)))
+                         ((when-let ((parent (ts-fold-find-parent node "if")))
+                            (- (tsc-node-end-position parent) 3))))))
+    (when ts-fold-on-next-line  ; display nicely
+      (setq end (ts-fold-point-before-line-break end)))
     (ts-fold--cons-add (cons beg end) offset)))
 
 (defun ts-fold-range-rust-macro (node offset)
@@ -655,5 +678,96 @@ information."
          (fold-begin (1- (- end-position start-position))))
     (ts-fold-range-seq node (ts-fold--cons-add (cons fold-begin -2) offset))))
 
+(defun ts-fold-range-lua-comment (node offset)
+  "Define fold range for Lua comemnt.
+
+For arguments NODE and OFFSET, see function `ts-fold-range-seq' for
+more information."
+  (let ((text (tsc-node-text node)))
+    (if (and (string-match-p "\n" text) (string-prefix-p "--[[" text))
+        (ts-fold-range-block-comment node
+                                     ;; XXX: Add 2 to for ]] at the end
+                                     (ts-fold--cons-add (cons 2 0) offset))
+      (ts-fold-range-line-comment node offset "--"))))
+
+(defun ts-fold-range-lua-function (node offset)
+  "Define fold range for Lua `function' declaration.
+
+For arguments NODE and OFFSET, see function `ts-fold-range-seq' for
+more information."
+  (let* ((params (tsc-get-child-by-field node :parameters))
+         (beg (tsc-node-end-position params))
+         (end (- (tsc-node-end-position node) 3)))  ; fit identifier `end'
+    (when ts-fold-on-next-line  ; display nicely
+      (setq end (ts-fold-point-before-line-break end)))
+    (ts-fold--cons-add (cons beg end) offset)))
+
+(defun ts-fold-range-lua-if (node offset)
+  "Define fold range for Lua `if' statement.
+
+For arguments NODE and OFFSET, see function `ts-fold-range-seq' for
+more information."
+  (let* ((then (car (ts-fold-find-children node "then")))
+         (beg (tsc-node-end-position then))
+         (next (or (ts-fold-find-children-traverse node "elseif_statement")
+                   (ts-fold-find-children-traverse node "else_statement")))
+         (end (if next
+                  (tsc-node-start-position (car next))
+                (- (tsc-node-end-position node) 3))))
+    (when ts-fold-on-next-line  ; display nicely
+      (setq end (ts-fold-point-before-line-break end)))
+    (ts-fold--cons-add (cons beg end) offset)))
+
+(defun ts-fold-range-lua-elseif (node offset)
+  "Define fold range for Lua `elseif' statement.
+
+For arguments NODE and OFFSET, see function `ts-fold-range-seq' for
+more information."
+  (let* ((then (car (ts-fold-find-children node "then")))
+         (beg (tsc-node-end-position then))
+         (next (tsc-get-next-sibling node))
+         (end (if next
+                  (tsc-node-start-position next)
+                (tsc-node-end-position node))))
+    (when ts-fold-on-next-line  ; display nicely
+      (setq end (ts-fold-point-before-line-break end)))
+    (ts-fold--cons-add (cons beg end) offset)))
+
+(defun ts-fold-range-lua-else (node offset)
+  "Define fold range for Lua `else' statement.
+
+For arguments NODE and OFFSET, see function `ts-fold-range-seq' for
+more information."
+  (let* ((beg (+ (tsc-node-start-position node) 4))  ; fit `else', 4 letters
+         (next (tsc-get-next-sibling node))          ; the `end' node
+         (end (tsc-node-start-position next)))
+    (when ts-fold-on-next-line  ; display nicely
+      (setq end (ts-fold-point-before-line-break end)))
+    (ts-fold--cons-add (cons beg end) offset)))
+
+(defun ts-fold-range-lua-do-loop (node offset)
+  "Define fold range for Lua `while' and `for' statement.
+
+For arguments NODE and OFFSET, see function `ts-fold-range-seq' for
+more information."
+  (let* ((do (car (ts-fold-find-children node "do")))
+         (beg (tsc-node-end-position do))
+         (end (- (tsc-node-end-position node) 3)))
+    (when ts-fold-on-next-line  ; display nicely
+      (setq end (ts-fold-point-before-line-break end)))
+    (ts-fold--cons-add (cons beg end) offset)))
+
+(defun ts-fold-range-lua-repeat (node offset)
+  "Define fold range for Lua `repeat' statement.
+
+For arguments NODE and OFFSET, see function `ts-fold-range-seq' for
+more information."
+  (let* ((beg (+ (tsc-node-start-position node) 6))  ; fit `repeat', 6 letters
+         (until (car (ts-fold-find-children node "until")))
+         (end (tsc-node-start-position until)))
+    (when ts-fold-on-next-line  ; display nicely
+      (setq end (ts-fold-point-before-line-break end)))
+    (ts-fold--cons-add (cons beg end) offset)))
+
 (provide 'ts-fold)
 ;;; ts-fold.el ends here