From 0913a35b10db0d91baf9af4f2cac30354009d689 Mon Sep 17 00:00:00 2001
From: Yusuke Endoh <mame@ruby-lang.org>
Date: Fri, 3 Apr 2026 20:00:04 +0900
Subject: [PATCH] Add module_function support

When `module_function` is called without arguments in a module body,
subsequent method definitions are registered as both instance methods
and singleton methods, matching Ruby's behavior.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 lib/typeprof/core/ast.rb         |  4 ++++
 lib/typeprof/core/ast/meta.rb    |  6 ++++++
 lib/typeprof/core/ast/method.rb  |  3 +++
 lib/typeprof/core/env.rb         |  1 +
 scenario/misc/module_function.rb | 16 ++++++++++++++++
 5 files changed, 30 insertions(+)
 create mode 100644 scenario/misc/module_function.rb

diff --git a/lib/typeprof/core/ast.rb b/lib/typeprof/core/ast.rb
index d1ec4ef8c..2f03228a3 100644
--- a/lib/typeprof/core/ast.rb
+++ b/lib/typeprof/core/ast.rb
@@ -280,6 +280,10 @@ def self.create_node(raw_node, lenv, use_result = true, allow_meta = false)
             return AttrReaderMetaNode.new(raw_node, lenv)
           when :attr_accessor
             return AttrAccessorMetaNode.new(raw_node, lenv)
+          when :module_function
+            if raw_node.arguments.nil?
+              return ModuleFunctionMetaNode.new(raw_node, lenv)
+            end
           end
         end
         CallNode.new(raw_node, lenv)
diff --git a/lib/typeprof/core/ast/meta.rb b/lib/typeprof/core/ast/meta.rb
index 2da05462f..6864e31a0 100644
--- a/lib/typeprof/core/ast/meta.rb
+++ b/lib/typeprof/core/ast/meta.rb
@@ -146,5 +146,11 @@ def install0(genv)
         Source.new
       end
     end
+    class ModuleFunctionMetaNode < Node
+      def install0(genv)
+        @lenv.module_function = true
+        Source.new
+      end
+    end
   end
 end
diff --git a/lib/typeprof/core/ast/method.rb b/lib/typeprof/core/ast/method.rb
index 843dcca8f..443be522f 100644
--- a/lib/typeprof/core/ast/method.rb
+++ b/lib/typeprof/core/ast/method.rb
@@ -289,6 +289,9 @@ def install0(genv)
         )
 
         @changes.add_method_def_box(genv, @lenv.cref.cpath, @singleton, @mid, f_args, @body.lenv.return_boxes)
+        if @lenv.module_function && !@singleton
+          @changes.add_method_def_box(genv, @lenv.cref.cpath, true, @mid, f_args, @body.lenv.return_boxes)
+        end
 
         Source.new(Type::Symbol.new(genv, @mid))
       end
diff --git a/lib/typeprof/core/env.rb b/lib/typeprof/core/env.rb
index 911174687..6ef3510b1 100644
--- a/lib/typeprof/core/env.rb
+++ b/lib/typeprof/core/env.rb
@@ -331,6 +331,7 @@ def initialize(file_context, cref, locals, return_boxes)
     end
 
     attr_reader :file_context, :cref, :locals, :return_boxes, :break_vtx, :next_boxes, :strict_const_scope
+    attr_accessor :module_function
 
     def path = @file_context&.path
     def code_range_from_node(node)
diff --git a/scenario/misc/module_function.rb b/scenario/misc/module_function.rb
new file mode 100644
index 000000000..bd20c8dfd
--- /dev/null
+++ b/scenario/misc/module_function.rb
@@ -0,0 +1,16 @@
+## update
+module Foo
+  module_function
+
+  def bar
+    42
+  end
+end
+
+Foo.bar
+
+## assert
+module Foo
+  def bar: -> Integer
+  def self.bar: -> Integer
+end