From b444573aa22072159ed0c728af063504e5765c5d Mon Sep 17 00:00:00 2001
From: aycabta <aycabta@gmail.com>
Date: Sat, 6 Feb 2021 21:07:11 +0900
Subject: [PATCH] Enable to reassign a new block with "measure" command

---
 lib/irb/init.rb      | 8 ++++++++
 test/irb/test_cmd.rb | 9 ++++++++-
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/lib/irb/init.rb b/lib/irb/init.rb
index cd57012c6..78ef2fa3c 100644
--- a/lib/irb/init.rb
+++ b/lib/irb/init.rb
@@ -157,6 +157,14 @@ def IRB.set_measure_callback(type = nil, arg = nil, &block)
       added = [:CUSTOM, IRB.conf[:MEASURE_PROC][:CUSTOM], arg]
     elsif block_given?
       added = [:BLOCK, block, arg]
+      found = IRB.conf[:MEASURE_CALLBACKS].find{ |m| m[0] == added[0] && m[2] == added[2] }
+      if found
+        found[1] = block
+        return added
+      else
+        IRB.conf[:MEASURE_CALLBACKS] << added
+        return added
+      end
     else
       added = [:TIME, IRB.conf[:MEASURE_PROC][:TIME], arg]
     end
diff --git a/test/irb/test_cmd.rb b/test/irb/test_cmd.rb
index 1b319105a..341fdfbbf 100644
--- a/test/irb/test_cmd.rb
+++ b/test/irb/test_cmd.rb
@@ -294,6 +294,13 @@ def test_measure_with_proc
         "measure { |context, code, line_no, &block|\n",
         "  result = block.()\n",
         "  puts 'aaa' if IRB.conf[:MEASURE]\n",
+        "  result\n",
+        "}\n",
+        "3\n",
+        "measure { |context, code, line_no, &block|\n",
+        "  result = block.()\n",
+        "  puts 'bbb' if IRB.conf[:MEASURE]\n",
+        "  result\n",
         "}\n",
         "3\n",
         "measure :off\n",
@@ -306,7 +313,7 @@ def test_measure_with_proc
         irb.eval_input
       end
       assert_empty err
-      assert_match(/\A=> 3\nBLOCK is added\.\n=> nil\naaa\n=> 3\n=> nil\n=> 3\n/, out)
+      assert_match(/\A=> 3\nBLOCK is added\.\n=> nil\naaa\n=> 3\nBLOCK is added.\naaa\n=> nil\nbbb\n=> 3\n=> nil\n=> 3\n/, out)
       assert_empty(c.class_variables)
     end