From af143d8a749ae0ba0f394521dd46bea824a354fa Mon Sep 17 00:00:00 2001
From: ydah <t.yudai92@gmail.com>
Date: Tue, 3 Sep 2024 17:53:55 +0900
Subject: [PATCH] Implement UNDEF NODE keyword locations

---
 ast.c                 | 4 ++++
 node_dump.c           | 1 +
 parse.y               | 2 ++
 rubyparser.h          | 1 +
 test/ruby/test_ast.rb | 8 ++++++++
 5 files changed, 16 insertions(+)

diff --git a/ast.c b/ast.c
index 338ab00819..fd6a6dea89 100644
--- a/ast.c
+++ b/ast.c
@@ -775,6 +775,10 @@ node_locations(VALUE ast_value, const NODE *node)
 {
     enum node_type type = nd_type(node);
     switch (type) {
+      case NODE_UNDEF:
+        return rb_ary_new_from_args(2,
+                                    location_new(nd_code_loc(node)),
+                                    location_new(&RNODE_UNDEF(node)->keyword_loc));
       case NODE_UNLESS:
         return rb_ary_new_from_args(4,
                                     location_new(nd_code_loc(node)),
diff --git a/node_dump.c b/node_dump.c
index f23aeca6f1..2b9322bdb7 100644
--- a/node_dump.c
+++ b/node_dump.c
@@ -940,6 +940,7 @@ dump_node(VALUE buf, VALUE indent, int comment, const NODE * node)
         ANN("example: undef foo");
         LAST_NODE;
         F_ARRAY(nd_undefs, RNODE_UNDEF, "nd_undefs");
+        F_LOC(keyword_loc, RNODE_UNDEF);
         return;
 
       case NODE_CLASS:
diff --git a/parse.y b/parse.y
index 5139ee99e6..f8bca1ee52 100644
--- a/parse.y
+++ b/parse.y
@@ -3156,6 +3156,7 @@ stmt		: keyword_alias fitem {SET_LEX_STATE(EXPR_FNAME|EXPR_FITEM);} fitem
                 | keyword_undef undef_list
                     {
                         nd_set_first_loc($2, @1.beg_pos);
+                        RNODE_UNDEF($2)->keyword_loc = @1;
                         $$ = $2;
                     /*% ripper: undef!($:2) %*/
                     }
@@ -12320,6 +12321,7 @@ rb_node_undef_new(struct parser_params *p, NODE *nd_undef, const YYLTYPE *loc)
 {
     rb_node_undef_t *n = NODE_NEWNODE(NODE_UNDEF, rb_node_undef_t, loc);
     n->nd_undefs = rb_parser_ary_new_capa_for_node(p, 1);
+    n->keyword_loc = NULL_LOC;
     rb_parser_ary_push_node(p, n->nd_undefs, nd_undef);
 
     return n;
diff --git a/rubyparser.h b/rubyparser.h
index 0af0198500..e46e99616d 100644
--- a/rubyparser.h
+++ b/rubyparser.h
@@ -833,6 +833,7 @@ typedef struct RNode_UNDEF {
     NODE node;
 
     rb_parser_ary_t *nd_undefs;
+    rb_code_location_t keyword_loc;
 } rb_node_undef_t;
 
 typedef struct RNode_CLASS {
diff --git a/test/ruby/test_ast.rb b/test/ruby/test_ast.rb
index 936dd6d3e6..d21c57d88f 100644
--- a/test/ruby/test_ast.rb
+++ b/test/ruby/test_ast.rb
@@ -1338,6 +1338,14 @@ dummy
       assert_locations(node.children[-1].locations, [[1, 0, 1, 10], [1, 2, 1, 8], nil, nil])
     end
 
+    def test_undef_locations
+      node = RubyVM::AbstractSyntaxTree.parse("undef foo")
+      assert_locations(node.children[-1].locations, [[1, 0, 1, 9], [1, 0, 1, 5]])
+
+      node = RubyVM::AbstractSyntaxTree.parse("undef foo, bar")
+      assert_locations(node.children[-1].locations, [[1, 0, 1, 14], [1, 0, 1, 5]])
+    end
+
     private
     def assert_locations(locations, expected)
       ary = locations.map {|loc| loc && [loc.first_lineno, loc.first_column, loc.last_lineno, loc.last_column] }