From c9d101f450def74fa3305b8ca61a1266512be2cb Mon Sep 17 00:00:00 2001
From: Ulysse Buonomo <buonomo.ulysse@gmail.com>
Date: Thu, 27 May 2021 07:35:09 +0200
Subject: [PATCH] Improve stackprof measure

Allow usage of more detailed args when setting stackprof callback.

Signed-off-by: Ulysse Buonomo <buonomo.ulysse@gmail.com>
---
 Gemfile         |  1 +
 lib/irb/init.rb | 16 ++++++++++++++--
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/Gemfile b/Gemfile
index 481469076..6c896316e 100644
--- a/Gemfile
+++ b/Gemfile
@@ -10,4 +10,5 @@ group :development do
   gem 'vterm', '>= 0.0.5' if is_unix && ENV['WITH_VTERM']
   gem 'yamatanooroti', '>= 0.0.6'
   gem "rake"
+  gem "stackprof" if is_unix
 end
diff --git a/lib/irb/init.rb b/lib/irb/init.rb
index 3a5f3de26..5923c54a2 100644
--- a/lib/irb/init.rb
+++ b/lib/irb/init.rb
@@ -120,7 +120,11 @@ def IRB.init_config(ap_path)
       puts 'processing time: %fs' % (now - time) if IRB.conf[:MEASURE]
       result
     }
+    # arg can be either a symbol for the mode (:cpu, :wall, ..) or a hash for
+    # a more complete configuration.
+    # See https://github.com/tmm1/stackprof#all-options.
     @CONF[:MEASURE_PROC][:STACKPROF] = proc { |context, code, line_no, arg, &block|
+      return block.() unless IRB.conf[:MEASURE]
       success = false
       begin
         require 'stackprof'
@@ -130,10 +134,18 @@ def IRB.init_config(ap_path)
       end
       if success
         result = nil
-        stackprof_result = StackProf.run(mode: arg ? arg : :cpu) do
+        arg = { mode: arg || :cpu } unless arg.is_a?(Hash)
+        stackprof_result = StackProf.run(**arg) do
           result = block.()
         end
-        StackProf::Report.new(stackprof_result).print_text if IRB.conf[:MEASURE]
+        case stackprof_result
+        when File
+          puts "StackProf report saved to #{stackprof_result.path}"
+        when Hash
+          StackProf::Report.new(stackprof_result).print_text
+        else
+          puts "Stackprof ran with #{arg.inspect}"
+        end
         result
       else
         block.()