From dd618f10ec750e363ac2f7c89433dd2f8a647c31 Mon Sep 17 00:00:00 2001
From: Evgeni Kolev <evgenysw@gmail.com>
Date: Thu, 29 Dec 2022 17:49:40 +0200
Subject: [PATCH] Improve go-ts-mode Imenu, navigation and electric pair
 configuration

The Imenu items are extended to support "Method", "Struct",
"Interface", "Alias" and "Type".

go-ts-mode is updated to use the Imenu facility added in commit
b39dc7ab27a696a8607ab859aeff3c71509231f5.

Variable electric-indent-chars is set in order to improve integration
with Electric Pair mode.

* lisp/progmodes/go-ts-mode.el (go-ts-mode--imenu-1) (go-ts-mode--imenu):
Remove functions.
(go-ts-mode--defun-name, go-ts-mode--interface-node-p)
(go-ts-mode--struct-node-p, go-ts-mode--other-type-node-p)
(go-ts-mode--alias-node-p): New functions.
(go-ts-mode): Improve Imenu settings, navigation, add Electric Pair
mode settings.
---
 lisp/progmodes/go-ts-mode.el | 107 ++++++++++++++++++++++-------------
 1 file changed, 67 insertions(+), 40 deletions(-)

diff --git a/lisp/progmodes/go-ts-mode.el b/lisp/progmodes/go-ts-mode.el
index 1d6a8a30db5..64e761d2f72 100644
--- a/lisp/progmodes/go-ts-mode.el
+++ b/lisp/progmodes/go-ts-mode.el
@@ -36,6 +36,7 @@
 (declare-function treesit-node-child-by-field-name "treesit.c")
 (declare-function treesit-node-start "treesit.c")
 (declare-function treesit-node-type "treesit.c")
+(declare-function treesit-search-subtree "treesit.c")
 
 (defcustom go-ts-mode-indent-offset 4
   "Number of spaces for each indentation step in `go-ts-mode'."
@@ -173,44 +174,6 @@ go-ts-mode--font-lock-settings
    '((ERROR) @font-lock-warning-face))
   "Tree-sitter font-lock settings for `go-ts-mode'.")
 
-(defun go-ts-mode--imenu ()
-  "Return Imenu alist for the current buffer."
-  (let* ((node (treesit-buffer-root-node))
-         (func-tree (treesit-induce-sparse-tree
-                     node "function_declaration" nil 1000))
-         (type-tree (treesit-induce-sparse-tree
-                     node "type_spec" nil 1000))
-         (func-index (go-ts-mode--imenu-1 func-tree))
-         (type-index (go-ts-mode--imenu-1 type-tree)))
-    (append
-     (when func-index `(("Function" . ,func-index)))
-     (when type-index `(("Type" . ,type-index))))))
-
-(defun go-ts-mode--imenu-1 (node)
-  "Helper for `go-ts-mode--imenu'.
-Find string representation for NODE and set marker, then recurse
-the subtrees."
-  (let* ((ts-node (car node))
-         (children (cdr node))
-         (subtrees (mapcan #'go-ts-mode--imenu-1
-                           children))
-         (name (when ts-node
-                 (treesit-node-text
-                  (pcase (treesit-node-type ts-node)
-                    ("function_declaration"
-                     (treesit-node-child-by-field-name ts-node "name"))
-                    ("type_spec"
-                     (treesit-node-child-by-field-name ts-node "name"))))))
-         (marker (when ts-node
-                   (set-marker (make-marker)
-                               (treesit-node-start ts-node)))))
-    (cond
-     ((or (null ts-node) (null name)) subtrees)
-     (subtrees
-      `((,name ,(cons name marker) ,@subtrees)))
-     (t
-      `((,name . ,marker))))))
-
 ;;;###autoload
 (add-to-list 'auto-mode-alist '("\\.go\\'" . go-ts-mode))
 
@@ -228,14 +191,30 @@ go-ts-mode
     (setq-local comment-end "")
     (setq-local comment-start-skip (rx "//" (* (syntax whitespace))))
 
+    ;; Navigation.
+    (setq-local treesit-defun-type-regexp
+                (regexp-opt '("method_declaration"
+                              "function_declaration"
+                              "type_declaration")))
+    (setq-local treesit-defun-name-function #'go-ts-mode--defun-name)
+
     ;; Imenu.
-    (setq-local imenu-create-index-function #'go-ts-mode--imenu)
-    (setq-local which-func-functions nil)
+    (setq-local treesit-simple-imenu-settings
+                `(("Function" "\\`function_declaration\\'" nil nil)
+                  ("Method" "\\`method_declaration\\'" nil nil)
+                  ("Struct" "\\`type_declaration\\'" go-ts-mode--struct-node-p nil)
+                  ("Interface" "\\`type_declaration\\'" go-ts-mode--interface-node-p nil)
+                  ("Type" "\\`type_declaration\\'" go-ts-mode--other-type-node-p nil)
+                  ("Alias" "\\`type_declaration\\'" go-ts-mode--alias-node-p nil)))
 
     ;; Indent.
     (setq-local indent-tabs-mode t
                 treesit-simple-indent-rules go-ts-mode--indent-rules)
 
+    ;; Electric
+    (setq-local electric-indent-chars
+                (append "{}()" electric-indent-chars))
+
     ;; Font-lock.
     (setq-local treesit-font-lock-settings go-ts-mode--font-lock-settings)
     (setq-local treesit-font-lock-feature-list
@@ -247,6 +226,54 @@ go-ts-mode
 
     (treesit-major-mode-setup)))
 
+(defun go-ts-mode--defun-name (node)
+  "Return the defun name of NODE.
+Return nil if there is no name or if NODE is not a defun node."
+  (pcase (treesit-node-type node)
+    ("function_declaration"
+     (treesit-node-text
+      (treesit-node-child-by-field-name
+       node "name")
+      t))
+    ("method_declaration"
+     (let* ((receiver-node (treesit-node-child-by-field-name node "receiver"))
+            (type-node (treesit-search-subtree receiver-node "type_identifier"))
+            (name-node (treesit-node-child-by-field-name node "name")))
+       (concat
+        "(" (treesit-node-text type-node) ")."
+        (treesit-node-text name-node))))
+    ("type_declaration"
+     (treesit-node-text
+      (treesit-node-child-by-field-name
+       (treesit-node-child node 0 t) "name")
+      t))))
+
+(defun go-ts-mode--interface-node-p (node)
+  "Return t if NODE is an interface."
+  (and
+   (string-equal "type_declaration" (treesit-node-type node))
+   (treesit-search-subtree node "interface_type" nil nil 2)))
+
+(defun go-ts-mode--struct-node-p (node)
+  "Return t if NODE is a struct."
+  (and
+   (string-equal "type_declaration" (treesit-node-type node))
+   (treesit-search-subtree node "struct_type" nil nil 2)))
+
+(defun go-ts-mode--alias-node-p (node)
+  "Return t if NODE is a type alias."
+  (and
+   (string-equal "type_declaration" (treesit-node-type node))
+   (treesit-search-subtree node "type_alias" nil nil 1)))
+
+(defun go-ts-mode--other-type-node-p (node)
+  "Return t if NODE is a type, other than interface, struct or alias."
+  (and
+   (string-equal "type_declaration" (treesit-node-type node))
+   (not (go-ts-mode--interface-node-p node))
+   (not (go-ts-mode--struct-node-p node))
+   (not (go-ts-mode--alias-node-p node))))
+
 ;; go.mod support.
 
 (defvar go-mod-ts-mode--syntax-table
-- 
2.30.2

