TMP

Gemfile.lock

@@ -40,7 +40,7 @@ GEM
     racc (1.8.0)
     rainbow (3.1.1)
     rake (13.2.1)
-    rbi (0.1.13)
+    rbi (0.1.14)
       prism (>= 0.18.0, < 1.0.0)
       sorbet-runtime (>= 0.5.9204)
     rdoc (6.6.3.1)

lib/spoom/cfg.rb

@@ -122,6 +122,10 @@ def to_dot(cluster_id)
             "class << self"
           when Prism::DefNode
             "def #{i.name}"
+          when Prism::CallNode
+            "#{i.receiver&.slice || "<self>"}.#{i.name}(#{i.arguments&.slice})"
+          when Prism::BlockNode
+            "<block-call>"
           else
             i.slice
           end
@@ -168,6 +172,10 @@ def inspect
             "class << self"
           when Prism::DefNode
             "def #{i.name}"
+          when Prism::CallNode
+            "#{i.receiver&.slice || "<self>"}.#{i.name}(#{i.arguments&.slice})"
+          when Prism::BlockNode
+            "<block-call>"
           else
             i.slice
           end
@@ -213,6 +221,7 @@ def initialize
       def visit_program_node(node)
         builder = Builder.new("<main>", node)
         builder.visit(node.statements)
+        builder.finalize!
         @cluster.cfgs << builder.cfg
 
         super
@@ -222,6 +231,7 @@ def visit_program_node(node)
       def visit_class_node(node)
         builder = Builder.new("#{node.name}::<static-init>", node)
         builder.visit(node.body)
+        builder.finalize!
         @cluster.cfgs << builder.cfg
 
         super
@@ -231,6 +241,7 @@ def visit_class_node(node)
       def visit_module_node(node)
         builder = Builder.new("#{node.name}::<static-init>", node)
         builder.visit(node.body)
+        builder.finalize!
         @cluster.cfgs << builder.cfg
 
         super
@@ -240,6 +251,7 @@ def visit_module_node(node)
       def visit_singleton_class_node(node)
         builder = Builder.new("class << self::<static-init>", node)
         builder.visit(node.body)
+        builder.finalize!
         @cluster.cfgs << builder.cfg
 
         super
@@ -249,6 +261,7 @@ def visit_singleton_class_node(node)
       def visit_def_node(node)
         builder = Builder.new(node.name.to_s, node)
         builder.visit(node.body)
+        builder.finalize!
         @cluster.cfgs << builder.cfg
       end
     end
@@ -264,9 +277,15 @@ def initialize(name, node)
         super()
 
         @block_count = T.let(0, Integer)
-        @current_block = T.let(new_block, BasicBlock)
-        @cfg = T.let(CFG.new(name, node, @current_block), CFG)
         @loop_stack = T.let([], T::Array[Loop])
+        @current_block = T.let(new_block, BasicBlock)
+        @exit_block = T.let(new_block, BasicBlock)
+        @cfg = T.let(CFG.new(name, node, @current_block, @exit_block), CFG)
+      end
+
+      sig { void }
+      def finalize!
+        @current_block.add_edge(@exit_block) unless @current_block.exits
       end
 
       # sig { override.params(node: Prism::AndNode).void }
@@ -298,27 +317,38 @@ def visit_break_node(node)
         raise Error.new("Unexpected break outside of loop", node) unless current_loop
 
         @current_block.instructions << node
-        @current_block.exits = true
         @current_block.add_edge(current_loop.merge)
+
+        after_block = new_block
+        @current_block.add_edge(after_block)
+        @current_block = after_block
+        @current_block.exits = true
       end
 
       sig { override.params(node: Prism::CallNode).void }
       def visit_call_node(node)
         @current_block.instructions << node
+        before_block = @current_block
 
-        # block_node = node.block
+        block_node = node.block
+        if block_node
+          call_block = new_block
+          before_block.add_edge(call_block)
+          call_block.instructions << block_node
+
+          block_block = new_block
+          call_block.add_edge(block_block)
+          @current_block = block_block
+          @loop_stack << Loop.new(call_block, @exit_block)
+          visit(block_node)
+          @loop_stack.pop
+          block_block = @current_block
+          block_block.add_edge(call_block) unless block_block.exits
 
-        # if block_node
-        #   block_block = new_block
-        #   @loop_stack << block_block
-        #   @current_block.add_edge(block_block)
-        #   @current_block = block_block
-        #   visit(block_node)
-        #   merge_block = new_block
-        #   @current_block.add_edge(merge_block)
-        #   @current_block = merge_block
-        #   @loop_stack.pop
-        # end
+          merge_block = new_block
+          call_block.add_edge(merge_block)
+          @current_block = merge_block
+        end
       end
 
       sig { override.params(node: Prism::CaseNode).void }
@@ -397,9 +427,13 @@ def visit_next_node(node)
         current_loop = @loop_stack.last
         raise Error.new("Unexpected next outside of loop", node) unless current_loop
 
-        @current_block.exits = true
         @current_block.instructions << node
         @current_block.add_edge(current_loop.header)
+
+        after_block = new_block
+        @current_block.add_edge(after_block)
+        @current_block = after_block
+        @current_block.exits = true
       end
 
       sig { override.params(node: Prism::IfNode).void }
@@ -471,6 +505,7 @@ def visit_module_node(node)
       sig { override.params(node: Prism::ReturnNode).void }
       def visit_return_node(node)
         @current_block.instructions << node
+        @current_block.add_edge(@exit_block)
         after_block = new_block
         @current_block.add_edge(after_block)
         @current_block = after_block
@@ -586,19 +621,20 @@ def new_block
     attr_reader :node
 
     sig { returns(BasicBlock) }
-    attr_reader :root
+    attr_reader :root_block, :exit_block
 
-    sig { params(name: String, node: Prism::Node, root: BasicBlock).void }
-    def initialize(name, node, root)
+    sig { params(name: String, node: Prism::Node, root_block: BasicBlock, exit_block: BasicBlock).void }
+    def initialize(name, node, root_block, exit_block)
       @name = name
       @node = node
-      @root = root
+      @root_block = root_block
+      @exit_block = exit_block
     end
 
     sig { returns(T::Array[BasicBlock]) }
     def blocks
       blocks = T.let([], T::Array[BasicBlock])
-      queue = T.let([root], T::Array[BasicBlock])
+      queue = T.let([root_block], T::Array[BasicBlock])
       seen = T.let(Set.new, T::Set[BasicBlock])
 
       until queue.empty?
@@ -622,6 +658,8 @@ def compact!
       return self if blocks.size == 1
 
       blocks.each do |block|
+        next if block == @root_block
+        next if block == @exit_block
         next unless block.empty?
 
         block.ins.each do |block_in|

lib/spoom/cli/cfg.rb

@@ -12,6 +12,7 @@ class CFG < Thor
       default_task :cfg
 
       desc "cfg PATH...", "Show the control flow graph of a Ruby file"
+      option :compact, type: :boolean, default: true, desc: "Compact empty blocks away"
       def cfg(*paths)
         files = files_to_typecheck(paths)
 
@@ -30,6 +31,7 @@ def cfg(*paths)
 
           node = result.value
           cfgs = Spoom::CFG.from_node(node)
+          cfgs.compact! if options[:compact]
           cfgs.show_dot
         end
       end

test.rb

@@ -1,9 +1,13 @@
 # typed: true
 # frozen_string_literal: true
 
-def foo
-  puts "before"
-  return if bar?
-
-  puts "after"
+puts "before"
+foo do |x|
+  if foo?
+    puts "will return"
+    return x
+    puts "dead"
+  end
+  puts "after return"
 end
+puts "after"