refactor: simplify Kernel patch restoration (#72)

Author: zah

codetracer/kernel_patches.rb

@@ -3,43 +3,38 @@
 module Codetracer
   module KernelPatches
     @@tracers = []
-    @@original_methods = {}
 
     def self.install(tracer)
       @@tracers << tracer
 
-      if @@original_methods.empty?
-        @@original_methods[:p] = Kernel.instance_method(:p)
-        @@original_methods[:puts] = Kernel.instance_method(:puts)
-        @@original_methods[:print] = Kernel.instance_method(:print)
-
+      if @@tracers.length == 1
         Kernel.module_eval do
-          alias_method :old_p, :p unless method_defined?(:old_p)
-          alias_method :old_puts, :puts unless method_defined?(:old_puts)
-          alias_method :old_print, :print unless method_defined?(:old_print)
+          alias_method :codetracer_original_p, :p unless method_defined?(:codetracer_original_p)
+          alias_method :codetracer_original_puts, :puts unless method_defined?(:codetracer_original_puts)
+          alias_method :codetracer_original_print, :print unless method_defined?(:codetracer_original_print)
 
           define_method(:p) do |*args|
-            loc = caller_locations(1,1).first
+            loc = caller_locations(1, 1).first
             @@tracers.each do |t|
               t.record_event(loc.path, loc.lineno, args.map(&:inspect).join("\n"))
             end
-            @@original_methods[:p].bind(self).call(*args)
+            codetracer_original_p(*args)
           end
 
           define_method(:puts) do |*args|
-            loc = caller_locations(1,1).first
+            loc = caller_locations(1, 1).first
             @@tracers.each do |t|
               t.record_event(loc.path, loc.lineno, args.join("\n"))
             end
-            @@original_methods[:puts].bind(self).call(*args)
+            codetracer_original_puts(*args)
           end
 
           define_method(:print) do |*args|
-            loc = caller_locations(1,1).first
+            loc = caller_locations(1, 1).first
             @@tracers.each do |t|
               t.record_event(loc.path, loc.lineno, args.join)
             end
-            @@original_methods[:print].bind(self).call(*args)
+            codetracer_original_print(*args)
           end
         end
       end
@@ -48,13 +43,16 @@ def self.install(tracer)
     def self.uninstall(tracer)
       @@tracers.delete(tracer)
 
-      if @@tracers.empty? && !@@original_methods.empty?
+      if @@tracers.empty? && Kernel.private_method_defined?(:codetracer_original_p)
         Kernel.module_eval do
-          define_method(:p, @@original_methods[:p])
-          define_method(:puts, @@original_methods[:puts])
-          define_method(:print, @@original_methods[:print])
+          alias_method :p, :codetracer_original_p
+          alias_method :puts, :codetracer_original_puts
+          alias_method :print, :codetracer_original_print
+
+          remove_method :codetracer_original_p
+          remove_method :codetracer_original_puts
+          remove_method :codetracer_original_print
         end
-        @@original_methods.clear
       end
     end
   end

test/test_kernel_patches.rb

@@ -40,7 +40,9 @@ def teardown
     end
     # Verify that original methods are restored if no tracers are left
     assert_empty Codetracer::KernelPatches.class_variable_get(:@@tracers), "Tracers should be empty after teardown"
-    assert_empty Codetracer::KernelPatches.class_variable_get(:@@original_methods), "Original methods should be cleared if no tracers are left"
+    refute Kernel.private_method_defined?(:codetracer_original_p), "Original method aliases should be removed"
+    refute Kernel.private_method_defined?(:codetracer_original_puts), "Original method aliases should be removed"
+    refute Kernel.private_method_defined?(:codetracer_original_print), "Original method aliases should be removed"
   end
 
   def test_patching_and_basic_event_recording
@@ -115,6 +117,9 @@ def test_restoration_of_original_methods
     p 'original restored' # This line's output will go to actual stdout
 
     assert_empty @tracer1.events, "Tracer should not record events after being uninstalled and patches removed"
+    refute Kernel.private_method_defined?(:codetracer_original_p), "Original method aliases should be removed"
+    refute Kernel.private_method_defined?(:codetracer_original_puts), "Original method aliases should be removed"
+    refute Kernel.private_method_defined?(:codetracer_original_print), "Original method aliases should be removed"
   end
 
   def test_correct_event_arguments