From b43f35d8f36ff84d73144160c7fc7fbceeb071e4 Mon Sep 17 00:00:00 2001
From: aycabta <aycabta@gmail.com>
Date: Tue, 22 Dec 2020 04:56:03 +0900
Subject: [PATCH] Support arg for measure command

---
 lib/irb.rb             |  5 +++--
 lib/irb/cmd/measure.rb | 12 ++++++------
 lib/irb/init.rb        | 14 +++++++-------
 3 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/lib/irb.rb b/lib/irb.rb
index 26c5d2ebe..4eef8be61 100644
--- a/lib/irb.rb
+++ b/lib/irb.rb
@@ -544,9 +544,10 @@ def eval_input
             if IRB.conf[:MEASURE] && !IRB.conf[:MEASURE_CALLBACKS].empty?
               result = nil
               last_proc = proc{ result = @context.evaluate(line, line_no, exception: exc) }
-              IRB.conf[:MEASURE_CALLBACKS].map{ |s| s.last }.inject(last_proc) { |chain, item|
+              IRB.conf[:MEASURE_CALLBACKS].inject(last_proc) { |chain, item|
+                _name, callback, arg = item
                 proc {
-                  item.(@context, line, line_no, exception: exc) do
+                  callback.(@context, line, line_no, arg, exception: exc) do
                     chain.call
                   end
                 }
diff --git a/lib/irb/cmd/measure.rb b/lib/irb/cmd/measure.rb
index 6161d15bc..244879c10 100644
--- a/lib/irb/cmd/measure.rb
+++ b/lib/irb/cmd/measure.rb
@@ -14,17 +14,17 @@ def execute(type = nil, arg = nil)
           IRB.conf[:MEASURE] = nil
           IRB.unset_measure_callback(arg)
         when :list
-          IRB.conf[:MEASURE_CALLBACKS].each do |type_name, _|
-            puts "- #{type_name}"
+          IRB.conf[:MEASURE_CALLBACKS].each do |type_name, _, arg|
+            puts "- #{type_name}" + (arg ? "(#{arg.inspect})" : '')
           end
         when :on
           IRB.conf[:MEASURE] = true
-          added = IRB.set_measure_callback(type)
-          puts "#{added.first} is added."
+          added = IRB.set_measure_callback(type, arg)
+          puts "#{added[0]} is added."
         else
           IRB.conf[:MEASURE] = true
-          added = IRB.set_measure_callback(type)
-          puts "#{added.first} is added."
+          added = IRB.set_measure_callback(type, arg)
+          puts "#{added[0]} is added."
         end
         nil
       end
diff --git a/lib/irb/init.rb b/lib/irb/init.rb
index f17d316c4..45a35cf94 100644
--- a/lib/irb/init.rb
+++ b/lib/irb/init.rb
@@ -120,7 +120,7 @@ def IRB.init_config(ap_path)
       puts 'processing time: %fs' % (now - time) if IRB.conf[:MEASURE]
       result
     }
-    @CONF[:MEASURE_PROC][:STACKPROF] = proc { |context, code, line_no, &block|
+    @CONF[:MEASURE_PROC][:STACKPROF] = proc { |context, code, line_no, arg, &block|
       success = false
       begin
         require 'stackprof'
@@ -130,7 +130,7 @@ def IRB.init_config(ap_path)
       end
       if success
         result = nil
-        stackprof_result = StackProf.run(mode: :cpu) do
+        stackprof_result = StackProf.run(mode: arg ? arg : :cpu) do
           result = block.()
         end
         StackProf::Report.new(stackprof_result).print_text if IRB.conf[:MEASURE]
@@ -146,17 +146,17 @@ def IRB.init_config(ap_path)
     @CONF[:AT_EXIT] = []
   end
 
-  def IRB.set_measure_callback(type = nil)
+  def IRB.set_measure_callback(type = nil, arg = nil)
     added = nil
     if type
       type_sym = type.upcase.to_sym
       if IRB.conf[:MEASURE_PROC][type_sym]
-        added = [type_sym, IRB.conf[:MEASURE_PROC][type_sym]]
+        added = [type_sym, IRB.conf[:MEASURE_PROC][type_sym], arg]
       end
     elsif IRB.conf[:MEASURE_PROC][:CUSTOM]
-      added = [:CUSTOM, IRB.conf[:MEASURE_PROC][:CUSTOM]]
+      added = [:CUSTOM, IRB.conf[:MEASURE_PROC][:CUSTOM], arg]
     else
-      added = [:TIME, IRB.conf[:MEASURE_PROC][:TIME]]
+      added = [:TIME, IRB.conf[:MEASURE_PROC][:TIME], arg]
     end
     IRB.conf[:MEASURE_CALLBACKS] << added if added
     added
@@ -167,7 +167,7 @@ def IRB.unset_measure_callback(type = nil)
       IRB.conf[:MEASURE_CALLBACKS].clear
     else
       type_sym = type.upcase.to_sym
-      IRB.conf[:MEASURE_CALLBACKS].reject!{ |t, c| t == type_sym }
+      IRB.conf[:MEASURE_CALLBACKS].reject!{ |t, | t == type_sym }
     end
   end