Add compaction support to `rb_ast_t`

This commit adds compaction support to `rb_ast_t`.

4 files changed, 70 insertions, 4 deletions

diff --git a/gc.c b/gc.c
index c083fbb297..adb86b10a5 100644
--- a/gc.c
+++ b/gc.c
@@ -7918,6 +7918,8 @@ gc_ref_update_imemo(rb_objspace_t *objspace, VALUE obj)
         rb_iseq_update_references((rb_iseq_t *)obj);
         break;
       case imemo_ast:
+        rb_ast_update_references((rb_ast_t *)obj);
+        break;
       case imemo_parser_strterm:
       case imemo_tmpbuf:
         break;
diff --git a/node.c b/node.c
index 4a265c3b9c..ee6a356b06 100644
--- a/node.c
+++ b/node.c
@@ -1262,20 +1262,20 @@ mark_ast_value(void *ctx, NODE * node)
             ID *buf = node->nd_tbl;
             if (buf) {
                 unsigned int size = (unsigned int)*buf;
-                rb_gc_mark((VALUE)buf[size + 1]);
+                rb_gc_mark_movable((VALUE)buf[size + 1]);
             }
             break;
         }
         case NODE_ARYPTN:
         {
             struct rb_ary_pattern_info *apinfo = node->nd_apinfo;
-            rb_gc_mark(apinfo->imemo);
+            rb_gc_mark_movable(apinfo->imemo);
             break;
         }
         case NODE_ARGS:
         {
             struct rb_args_info *args = node->nd_ainfo;
-            rb_gc_mark(args->imemo);
+            rb_gc_mark_movable(args->imemo);
             break;
         }
         case NODE_MATCH:
@@ -1286,13 +1286,62 @@ mark_ast_value(void *ctx, NODE * node)
         case NODE_DXSTR:
         case NODE_DREGX:
         case NODE_DSYM:
-            rb_gc_mark(node->nd_lit);
+            rb_gc_mark_movable(node->nd_lit);
             break;
         default:
             rb_bug("unreachable node %s", ruby_node_name(nd_type(node)));
     }
 }
 
+static void
+update_ast_value(void *ctx, NODE * node)
+{
+    switch (nd_type(node)) {
+        case NODE_SCOPE:
+        {
+            ID *buf = node->nd_tbl;
+            if (buf) {
+                unsigned int size = (unsigned int)*buf;
+                buf[size + 1] = rb_gc_location((VALUE)buf[size + 1]);
+            }
+            break;
+        }
+        case NODE_ARYPTN:
+        {
+            struct rb_ary_pattern_info *apinfo = node->nd_apinfo;
+            apinfo->imemo = rb_gc_location(apinfo->imemo);
+            break;
+        }
+        case NODE_ARGS:
+        {
+            struct rb_args_info *args = node->nd_ainfo;
+            args->imemo = rb_gc_location(args->imemo);
+            break;
+        }
+        case NODE_LIT:
+        case NODE_STR:
+        case NODE_XSTR:
+        case NODE_DSTR:
+        case NODE_DXSTR:
+        case NODE_DREGX:
+        case NODE_DSYM:
+            node->nd_lit = rb_gc_location(node->nd_lit);
+            break;
+        default:
+            rb_bug("unreachable");
+    }
+}
+
+void
+rb_ast_update_references(rb_ast_t *ast)
+{
+    if (ast->node_buffer) {
+        node_buffer_t *nb = ast->node_buffer;
+
+        iterate_node_values(&nb->markable, update_ast_value, NULL);
+    }
+}
+
 void
 rb_ast_mark(rb_ast_t *ast)
 {
diff --git a/node.h b/node.h
index 276b4d419c..55c2984ff3 100644
--- a/node.h
+++ b/node.h
@@ -405,6 +405,7 @@ typedef struct rb_ast_struct {
 } rb_ast_t;
 rb_ast_t *rb_ast_new(void);
 void rb_ast_mark(rb_ast_t*);
+void rb_ast_update_references(rb_ast_t*);
 void rb_ast_dispose(rb_ast_t*);
 void rb_ast_free(rb_ast_t*);
 size_t rb_ast_memsize(const rb_ast_t*);
diff --git a/test/ruby/test_gc_compact.rb b/test/ruby/test_gc_compact.rb
index ca4a285b70..eaffccda08 100644
--- a/test/ruby/test_gc_compact.rb
+++ b/test/ruby/test_gc_compact.rb
@@ -130,4 +130,18 @@ class TestGCCompact < Test::Unit::TestCase
     GC.verify_compaction_references(toward: :empty)
     assert_equal hash, list_of_objects.hash
   end
+
+  def walk_ast ast
+    children = ast.children.grep(RubyVM::AbstractSyntaxTree::Node)
+    children.each do |child|
+      assert child.type
+      walk_ast child
+    end
+  end
+
+  def test_ast_compacts
+    ast = RubyVM::AbstractSyntaxTree.parse_file __FILE__
+    assert GC.compact
+    walk_ast ast
+  end
 end