Fix function detection for TODO comments preceding method definitions

ai-code-change.el

@@ -31,6 +31,42 @@ ignoring leading whitespace."
                               "+")
                       (string-trim-left line)))))
 
+(defun ai-code--get-function-name-for-comment ()
+  "Get the appropriate function name when cursor is on a comment line.
+If the comment precedes a function definition or is inside a function body,
+returns that function's name. Otherwise returns the result of `which-function`."
+  (let ((current-func (which-function)))
+    (save-excursion
+      ;; Move to next non-comment, non-blank line
+      (forward-line 1)
+      (while (and (not (eobp))
+                  (or (looking-at-p "^[ \t]*$")
+                      (ai-code--is-comment-line
+                       (buffer-substring-no-properties
+                        (line-beginning-position)
+                        (line-end-position)))))
+        (forward-line 1))
+      ;; Get function name at this position
+      (unless (eobp)
+        (let ((next-func (which-function)))
+          ;; If we found a function name at the next code line,
+          ;; prefer it if it's more specific than current
+          (cond
+           ;; No current function, use next if available
+           ((not current-func) next-func)
+           ;; No next function, keep current
+           ((not next-func) current-func)
+           ;; Both exist: check if next is more specific
+           ;; More specific means it's longer and contains the current as a prefix
+           ;; (handling delimiters like . # or ::)
+           ((and (> (length next-func) (length current-func))
+                 (or (string-prefix-p (concat current-func ".") next-func)
+                     (string-prefix-p (concat current-func "#") next-func)
+                     (string-prefix-p (concat current-func "::") next-func)))
+            next-func)
+           ;; Otherwise keep current
+           (t current-func)))))))
+
 ;;;###autoload
 (defun ai-code-code-change (arg)
   "Generate prompt to change code under cursor or in selected region.
@@ -87,7 +123,9 @@ Argument ARG is the prefix argument."
     (let* ((current-line (string-trim (thing-at-point 'line t)))
            (current-line-number (line-number-at-pos (point)))
            (is-comment (ai-code--is-comment-line current-line))
-           (function-name (which-function))
+           (function-name (if is-comment
+                              (ai-code--get-function-name-for-comment)
+                            (which-function)))
            (function-context (if function-name
                                  (format "\nFunction: %s" function-name)
                                ""))

test_ai-code-change.el

@@ -0,0 +1,122 @@
+;;; test_ai-code-change.el --- Tests for ai-code-change.el -*- lexical-binding: t; -*-
+
+;; Author: Kang Tu <tninja@gmail.com>
+;; SPDX-License-Identifier: Apache-2.0
+
+;;; Commentary:
+;; Tests for the ai-code-change module, specifically testing
+;; the function detection logic for TODO comments.
+
+;;; Code:
+
+(require 'ert)
+(require 'ai-code-change)
+
+(ert-deftest test-ai-code--get-function-name-for-comment-basic ()
+  "Test function name detection when on a comment line before function body.
+This simulates the Ruby example from the issue where a TODO comment
+is between the function definition and its body."
+  (with-temp-buffer
+    ;; Simulate Ruby mode comment syntax
+    (setq-local comment-start "# ")
+    (insert "module Foo\n")
+    (insert "  class Bar\n")
+    (insert "    def baz\n")
+    (insert "    end\n")
+    (insert "\n")
+    (insert "    # TODO remove this function\n")  ;; Line 6 - cursor will be here
+    (insert "    def click_first_available(driver, selectors)\n")
+    (insert "      wait = Selenium::WebDriver::Wait.new(timeout: 10)\n")
+    (insert "    end\n")
+    (insert "  end\n")
+    (insert "end\n")
+    ;; Move cursor to the TODO comment line (line 6)
+    (goto-char (point-min))
+    (forward-line 5)  ;; Move 5 lines forward from line 1 to reach line 6
+    ;; Mock which-function to simulate the actual behavior
+    ;; When on line 6, which-function might return "Bar" (class)
+    ;; When on line 7 (def line), it should return "Bar#click_first_available"
+    (cl-letf (((symbol-function 'which-function)
+               (lambda ()
+                 (save-excursion
+                   (let ((line-num (line-number-at-pos (point))))
+                     (cond
+                      ((= line-num 6) "Bar")  ;; On comment, returns class
+                      ((= line-num 7) "Bar#click_first_available")  ;; On def, returns method
+                      ((>= line-num 8) "Bar#click_first_available")  ;; Inside method body
+                      (t nil)))))))
+      ;; Test that on the comment line, we get the correct function name
+      (let ((result (ai-code--get-function-name-for-comment)))
+        (should (string= result "Bar#click_first_available"))))))
+
+(ert-deftest test-ai-code--get-function-name-for-comment-no-function ()
+  "Test function name detection when comment is not followed by a function."
+  (with-temp-buffer
+    (setq-local comment-start "# ")
+    (insert "# TODO some task\n")
+    (insert "x = 1\n")
+    (goto-char (point-min))
+    (cl-letf (((symbol-function 'which-function) (lambda () nil)))
+      (let ((result (ai-code--get-function-name-for-comment)))
+        (should (null result))))))
+
+(ert-deftest test-ai-code--get-function-name-for-comment-multiple-comments ()
+  "Test function name detection with multiple comment lines before function."
+  (with-temp-buffer
+    (setq-local comment-start "# ")
+    (insert "    # TODO task 1\n")  ;; Line 1 - cursor here
+    (insert "    # TODO task 2\n")  ;; Line 2
+    (insert "    def my_function()\n")  ;; Line 3
+    (insert "      x = 1\n")
+    (insert "    end\n")
+    (goto-char (point-min))
+    ;; Mock which-function
+    (cl-letf (((symbol-function 'which-function)
+               (lambda ()
+                 (save-excursion
+                   (let ((line-num (line-number-at-pos (point))))
+                     (cond
+                      ((<= line-num 2) nil)  ;; On comments, no function context
+                      ((>= line-num 3) "my_function")  ;; On/in function
+                      (t nil)))))))
+      (let ((result (ai-code--get-function-name-for-comment)))
+        (should (string= result "my_function"))))))
+
+(ert-deftest test-ai-code--get-function-name-for-comment-same-function ()
+  "Test that when comment and next line are in same function, we get that function."
+  (with-temp-buffer
+    (setq-local comment-start "# ")
+    (insert "    def my_function()\n")  ;; Line 1
+    (insert "      # TODO implement this\n")  ;; Line 2 - cursor here
+    (insert "      x = 1\n")  ;; Line 3
+    (insert "    end\n")
+    (goto-char (point-min))
+    (forward-line 1)  ;; Move 1 line forward from line 1 to reach line 2 (the comment)
+    ;; Mock which-function - both lines return same function
+    (cl-letf (((symbol-function 'which-function) (lambda () "my_function")))
+      (let ((result (ai-code--get-function-name-for-comment)))
+        (should (string= result "my_function"))))))
+
+(ert-deftest test-ai-code--is-comment-line ()
+  "Test comment line detection."
+  ;; Test with hash comment
+  (let ((comment-start "# "))
+    (should (ai-code--is-comment-line "# This is a comment"))
+    (should (ai-code--is-comment-line "  # This is an indented comment"))
+    (should (ai-code--is-comment-line "## Multiple hashes"))
+    (should-not (ai-code--is-comment-line "This is not a comment"))
+    (should-not (ai-code--is-comment-line "  x = 1  # inline comment")))
+  ;; Test with semicolon comment (Lisp)
+  (let ((comment-start "; "))
+    (should (ai-code--is-comment-line "; This is a comment"))
+    (should (ai-code--is-comment-line "  ;; This is a comment"))
+    (should-not (ai-code--is-comment-line "This is not a comment")))
+  ;; Test with double slash comment (C/Java)
+  (let ((comment-start "// "))
+    (should (ai-code--is-comment-line "// This is a comment"))
+    (should (ai-code--is-comment-line "  // This is an indented comment"))
+    (should-not (ai-code--is-comment-line "This is not a comment"))))
+
+(provide 'test_ai-code-change)
+
+;;; test_ai-code-change.el ends here