From 8acbf68ed67377ff69a735462d30300c18b32f62 Mon Sep 17 00:00:00 2001
From: Takumi Shotoku <sinsoku.listy@gmail.com>
Date: Sun, 30 Aug 2026 18:09:06 +0900
Subject: [PATCH] Support BEGIN blocks

`BEGIN { ... }` raised "not supported yet: pre_execution_node" and
aborted the whole run, unlike `END { ... }`, which was already
supported. BEGIN now shares its implementation with END but is installed
before the surrounding statements instead of after, so a method or a
local variable it defines reaches the statements that follow.
---
 lib/typeprof/core/ast.rb          |  1 +
 lib/typeprof/core/ast/misc.rb     |  7 ++++++-
 lib/typeprof/core/ast/variable.rb |  8 +++++++-
 scenario/misc/pre-exec.rb         | 24 ++++++++++++++++++++++++
 4 files changed, 38 insertions(+), 2 deletions(-)
 create mode 100644 scenario/misc/pre-exec.rb

diff --git a/lib/typeprof/core/ast.rb b/lib/typeprof/core/ast.rb
index 8420c3881..7f999a1bc 100644
--- a/lib/typeprof/core/ast.rb
+++ b/lib/typeprof/core/ast.rb
@@ -278,6 +278,7 @@ def self.create_node(raw_node, lenv, use_result = true, allow_meta = false)
       when :splat_node then SplatNode.new(raw_node, lenv)
       when :for_node then ForNode.new(raw_node, lenv)
       when :alias_global_variable_node then AliasGlobalVariableNode.new(raw_node, lenv)
+      when :pre_execution_node then PreExecutionNode.new(raw_node, lenv)
       when :post_execution_node then PostExecutionNode.new(raw_node, lenv)
       when :flip_flop_node then FlipFlopNode.new(raw_node, lenv)
       when :shareable_constant_node then create_node(raw_node.write, lenv)
diff --git a/lib/typeprof/core/ast/misc.rb b/lib/typeprof/core/ast/misc.rb
index 56f0321e0..369c79119 100644
--- a/lib/typeprof/core/ast/misc.rb
+++ b/lib/typeprof/core/ast/misc.rb
@@ -23,8 +23,13 @@ def install0(genv)
 
         post_stmts = []
 
+        # BEGIN runs before the surrounding statements, END after them
         @stmts.each do |stmt|
-          next if stmt.nil?
+          stmt.install(genv) if stmt.is_a?(PreExecutionNode)
+        end
+
+        @stmts.each do |stmt|
+          next if stmt.is_a?(PreExecutionNode)
 
           if stmt.is_a?(PostExecutionNode)
             post_stmts << stmt
diff --git a/lib/typeprof/core/ast/variable.rb b/lib/typeprof/core/ast/variable.rb
index f42898b7b..49da27296 100644
--- a/lib/typeprof/core/ast/variable.rb
+++ b/lib/typeprof/core/ast/variable.rb
@@ -257,7 +257,7 @@ def install0(genv)
       end
     end
 
-    class PostExecutionNode < Node
+    class ExecutionBaseNode < Node
       def initialize(raw_node, lenv)
         super(raw_node, lenv)
         @body = raw_node.statements ? AST.create_node(raw_node.statements, lenv) : DummyNilNode.new(TypeProf::CodeRange.new(code_range.last, code_range.last), lenv)
@@ -273,6 +273,12 @@ def install0(genv)
       end
     end
 
+    class PreExecutionNode < ExecutionBaseNode
+    end
+
+    class PostExecutionNode < ExecutionBaseNode
+    end
+
     class ClassVariableWriteNode < Node
       def initialize(raw_node, rhs, lenv)
         super(raw_node, lenv)
diff --git a/scenario/misc/pre-exec.rb b/scenario/misc/pre-exec.rb
new file mode 100644
index 000000000..afe2de491
--- /dev/null
+++ b/scenario/misc/pre-exec.rb
@@ -0,0 +1,24 @@
+## update
+def check(x)
+end
+
+def check2(x)
+end
+
+y = "str"
+
+BEGIN {
+  def hello = "hello"
+  x = 1
+  y = 1
+}
+
+check(x)
+check2(y)
+
+## assert
+class Object
+  def check: (Integer) -> nil
+  def check2: (String) -> nil
+  def hello: -> String
+end