From 3efaa713ac9fdcb4507e945c2db771f7bb3a1be6 Mon Sep 17 00:00:00 2001
From: Evgeni Kolev <evgenysw@gmail.com>
Date: Mon, 6 Feb 2023 09:49:52 +0200
Subject: [PATCH] Extend go-ts-mode with support for pre-filling return
 statements

Extend go-ts-mode to insert context-aware return statement for the
current function. The return statement is pre-filled with the zero
values of the function's output parameters.

For example, this return statement is added by Emacs:
```
func f() (bool, int, myStruct, *int, chan bool, error) {
    return false, 0, myStruct{}, nil, nil, err
}
```

* lisp/progmodes/go-ts-mode.el (go-ts-mode-error-zero-value): New
defcustom. (go-ts-mode-insert-return, go-ts-mode-return,
go-ts-mode--function-return-types, go-ts-mode--function-at-point,
go-ts-mode--types-from-param-declaration, go-ts-mode--zero-value) New
functions.
---
 lisp/progmodes/go-ts-mode.el | 103 +++++++++++++++++++++++++++++++++++
 1 file changed, 103 insertions(+)

diff --git a/lisp/progmodes/go-ts-mode.el b/lisp/progmodes/go-ts-mode.el
index 2a2783f45f6..0a163753f29 100644
--- a/lisp/progmodes/go-ts-mode.el
+++ b/lisp/progmodes/go-ts-mode.el
@@ -46,6 +46,15 @@ go-ts-mode-indent-offset
   :safe 'integerp
   :group 'go)
 
+(defcustom go-ts-mode-error-zero-value  "err"
+  "The zero value for type \"error\".
+The default is \"err\", assuming there's a varialbe named \"err\" in
+the current scope."
+  :version "30"
+  :type 'string
+  :options '("err", "fmt.Errorf(\"format: %w\", err)", "errors.Wrap(err, \"format\")")
+  :group 'go)
+
 (defvar go-ts-mode--syntax-table
   (let ((table (make-syntax-table)))
     (modify-syntax-entry ?+   "."      table)
@@ -328,6 +337,100 @@ go-ts-mode--comment-on-previous-line-p
      (<= (treesit-node-start node) point (treesit-node-end node))
      (string-equal "comment" (treesit-node-type node)))))
 
+(define-minor-mode go-ts-electric-return-mode
+  "Make `return RET' insert a return statement with zero values."
+  :lighter nil
+  (if go-ts-electric-return-mode
+      (add-hook 'post-self-insert-hook #'go-ts-mode--maybe-insert-return nil t)
+    (remove-hook 'post-self-insert-hook #'go-ts-mode--maybe-insert-return t)))
+
+(defun go-ts-mode--maybe-insert-return ()
+  "If RET is entered after `return', add zero values."
+  (when (and (equal (this-command-keys) (kbd "RET"))
+             (looking-back (rx word-boundary "return" "\n" (* space)) (pos-bol 0)))
+    (delete-region (point) (pos-eol 0)) ;; delete the line added by RET
+    (delete-char (- (length "return")))
+    (insert (go-ts-mode-return))))
+
+(defun go-ts-mode-return()
+  "Return a return statement with zero values for function at point.
+It's an error to call this function outside of a function body."
+  (interactive)
+  (let* ((func-node (go-ts-mode--function-at-point))
+         (body-node (treesit-node-child-by-field-name func-node "body")))
+
+    (unless (and body-node (< (treesit-node-start body-node) (point) (treesit-node-end body-node)))
+      (user-error "Point must be inside a Go function body"))
+
+    (let* ((types (go-ts-mode--function-return-types func-node))
+           (zero-values (mapcar 'go-ts-mode--zero-value types))
+           (return-statement (format "return %s" (string-join zero-values ", "))))
+      (string-trim-right return-statement))))
+
+(defun go-ts-mode--function-at-point ()
+  "Return the treesit function node at point.
+Function, methods and closures (anonymous functions) are recognized.
+If point is not in a recognized function, nil is returned."
+  (interactive)
+  (treesit-parent-until
+   (treesit-node-at (point))
+   (lambda (x)
+     (member (treesit-node-type x)
+             '("func_literal" "function_declaration" "method_declaration"))) t))
+
+(defun go-ts-mode--function-return-types (func-node)
+  "Return the types of the output variables of FUNC-NODE.
+The result is a list which could contain a single element (the Go
+function has one return argument), multiple elements (multiple return
+arguments), or no elements (no return arguments)."
+  (let ((result-node (treesit-node-child-by-field-name func-node "result")))
+    (pcase (treesit-node-type result-node)
+      ("parameter_list" ;; multiple return values
+       (let ((types-nested (treesit-induce-sparse-tree result-node "parameter_declaration" #'go-ts-mode--types-from-param-declaration)))
+         (apply #'append (mapcar 'car (cdr types-nested)))))
+      ((pred null) ;; no return value
+       nil)
+      (_ ;; single return value
+       (list (treesit-node-text result-node t))))))
+
+(defun go-ts-mode--types-from-param-declaration (node)
+ "Extract the Go type(s) from NODE.  It must have type parameter_declaration.
+
+A parameter_declaration node can contain single value `(int)', named
+single value `(sum int)', multiples values `(int, int)', or multiple
+named values `(x, y int)'.
+
+All the above cases are recognized.  Only the types are returned, for
+example, with parameter_declaration `(x, y int)', the result is
+\(\"int\", \"int\")."
+  (let ((type (treesit-node-child-by-field-name node "type"))
+        (count (length (treesit-query-capture node "(parameter_declaration name:(_) @name)" nil nil t))))
+    (make-list (if (zerop count) 1 count) (treesit-node-text type t))))
+
+(defun go-ts-mode--zero-value (go-type)
+  "Return the Go zero value for GO-TYPE.
+For type \"error\", the zero value is determined by variable
+`go-ts-mode-error-zero-value'.
+
+This function assumes arbitrary types are Go structs, for example for
+GO-TYPE \"person\", the retun value is \"person{}\"."
+  (pcase go-type
+    ((or "int"  "int8"  "int16"  "int32"  "int64") "0")
+    ((or "uint" "uint8" "uint16" "uint32" "uint64") "0")
+    ("bool" "false")
+    ("string" "\"\"")
+    ((or "uintptr" "byte" "rune") "0")
+    ((or "float32" "float64") "0.0")
+    ((or "complex64" "complex128") "complex(0, 0)")
+    ((pred (string-prefix-p "map")) "nil")
+    ((pred (string-prefix-p "*")) "nil")
+    ((pred (string-prefix-p "[]")) "nil")
+    ((pred (string-prefix-p "chan")) "nil")
+    ((pred (string-prefix-p "func")) "nil")
+    ("error" go-ts-mode-error-zero-value)
+    ((rx bol (1+ (or alphanumeric "_" ".")) eol) (concat go-type "{}"))
+    (_ (user-error "Unknown Go type \"%s\"" go-type))))
+
 ;; go.mod support.
 
 (defvar go-mod-ts-mode--syntax-table
-- 
2.39.1

