From 57ca5960ad207beb0c4f2788df0e74f8cc6b7cac Mon Sep 17 00:00:00 2001
From: Stan Lo <stan001212@gmail.com>
Date: Fri, 1 Mar 2024 23:51:14 +0800
Subject: [PATCH] [ruby/irb] Restructure workspace management
 (https://github.com/ruby/irb/pull/888)

* Remove dead irb_level method

* Restructure workspace management

Currently, workspace is an attribute of IRB::Context in most use cases.
But when some workspace commands are used, like `pushws` or `popws`, a
workspace will be created and used along side with the original workspace
attribute.

This complexity is not necessary and will prevent us from expanding
multi-workspace support in the future.

So this commit introduces a @workspace_stack ivar to IRB::Context so IRB
can have a more natural way to manage workspaces.

* Fix pushws without args

* Always display workspace stack after related commands are used

https://github.com/ruby/irb/commit/61560b99b3
---
 lib/irb.rb                  | 15 ++++++------
 lib/irb/command/pushws.rb   | 18 +++++++++++++-
 lib/irb/context.rb          | 26 +++++++++++++-------
 lib/irb/ext/change-ws.rb    |  6 ++---
 lib/irb/ext/eval_history.rb |  4 ++--
 lib/irb/ext/use-loader.rb   |  4 ++--
 lib/irb/ext/workspaces.rb   | 43 ++++++++-------------------------
 test/irb/test_command.rb    | 48 +++++++++++++++++++++----------------
 8 files changed, 86 insertions(+), 78 deletions(-)

diff --git a/lib/irb.rb b/lib/irb.rb
index 55988858664ced..0c481ff1dc76ca 100644
--- a/lib/irb.rb
+++ b/lib/irb.rb
@@ -933,7 +933,7 @@ def debug_break
 
     def debug_readline(binding)
       workspace = IRB::WorkSpace.new(binding)
-      context.workspace = workspace
+      context.replace_workspace(workspace)
       context.workspace.load_commands_to_main
       @line_no += 1
 
@@ -1269,12 +1269,11 @@ def suspend_name(path = nil, name = nil)
     # Used by the irb command +irb_load+, see IRB@IRB+Sessions for more
     # information.
     def suspend_workspace(workspace)
-      @context.workspace, back_workspace = workspace, @context.workspace
-      begin
-        yield back_workspace
-      ensure
-        @context.workspace = back_workspace
-      end
+      current_workspace = @context.workspace
+      @context.replace_workspace(workspace)
+      yield
+    ensure
+      @context.replace_workspace current_workspace
     end
 
     # Evaluates the given block using the given +input_method+ as the
@@ -1534,7 +1533,7 @@ def irb(show_code: true)
 
     if debugger_irb
       # If we're already in a debugger session, set the workspace and irb_path for the original IRB instance
-      debugger_irb.context.workspace = workspace
+      debugger_irb.context.replace_workspace(workspace)
       debugger_irb.context.irb_path = irb_path
       # If we've started a debugger session and hit another binding.irb, we don't want to start an IRB session
       # instead, we want to resume the irb:rdbg session.
diff --git a/lib/irb/command/pushws.rb b/lib/irb/command/pushws.rb
index fadd10d6aa4df3..888fe466f1a176 100644
--- a/lib/irb/command/pushws.rb
+++ b/lib/irb/command/pushws.rb
@@ -15,7 +15,23 @@ class Workspaces < Base
       description "Show workspaces."
 
       def execute(*obj)
-        irb_context.workspaces.collect{|ws| ws.main}
+        inspection_resuls = irb_context.instance_variable_get(:@workspace_stack).map do |ws|
+          truncated_inspect(ws.main)
+        end
+
+        puts "[" + inspection_resuls.join(", ") + "]"
+      end
+
+      private
+
+      def truncated_inspect(obj)
+        obj_inspection = obj.inspect
+
+        if obj_inspection.size > 20
+          obj_inspection = obj_inspection[0, 19] + "...>"
+        end
+
+        obj_inspection
       end
     end
 
diff --git a/lib/irb/context.rb b/lib/irb/context.rb
index 9647327037cb19..60dfb9668dac3c 100644
--- a/lib/irb/context.rb
+++ b/lib/irb/context.rb
@@ -22,10 +22,11 @@ class Context
     # +other+::   uses this as InputMethod
     def initialize(irb, workspace = nil, input_method = nil)
       @irb = irb
+      @workspace_stack = []
       if workspace
-        @workspace = workspace
+        @workspace_stack << workspace
       else
-        @workspace = WorkSpace.new
+        @workspace_stack << WorkSpace.new
       end
       @thread = Thread.current
 
@@ -229,15 +230,24 @@ def history_file=(hist)
       IRB.conf[:HISTORY_FILE] = hist
     end
 
+    # Workspace in the current context.
+    def workspace
+      @workspace_stack.last
+    end
+
+    # Replace the current workspace with the given +workspace+.
+    def replace_workspace(workspace)
+      @workspace_stack.pop
+      @workspace_stack.push(workspace)
+    end
+
     # The top-level workspace, see WorkSpace#main
     def main
-      @workspace.main
+      workspace.main
     end
 
     # The toplevel workspace, see #home_workspace
     attr_reader :workspace_home
-    # WorkSpace in the current context.
-    attr_accessor :workspace
     # The current thread in this context.
     attr_reader :thread
     # The current input method.
@@ -489,7 +499,7 @@ def prompting?
     # to #last_value.
     def set_last_value(value)
       @last_value = value
-      @workspace.local_variable_set :_, value
+      workspace.local_variable_set :_, value
     end
 
     # Sets the +mode+ of the prompt in this context.
@@ -585,7 +595,7 @@ def evaluate(line, line_no) # :nodoc:
 
       if IRB.conf[:MEASURE] && !IRB.conf[:MEASURE_CALLBACKS].empty?
         last_proc = proc do
-          result = @workspace.evaluate(line, @eval_path, line_no)
+          result = workspace.evaluate(line, @eval_path, line_no)
         end
         IRB.conf[:MEASURE_CALLBACKS].inject(last_proc) do |chain, item|
           _name, callback, arg = item
@@ -596,7 +606,7 @@ def evaluate(line, line_no) # :nodoc:
           end
         end.call
       else
-        result = @workspace.evaluate(line, @eval_path, line_no)
+        result = workspace.evaluate(line, @eval_path, line_no)
       end
 
       set_last_value(result)
diff --git a/lib/irb/ext/change-ws.rb b/lib/irb/ext/change-ws.rb
index ec29f7a2bc78cb..87fe03e23d953d 100644
--- a/lib/irb/ext/change-ws.rb
+++ b/lib/irb/ext/change-ws.rb
@@ -12,7 +12,7 @@ def home_workspace
       if defined? @home_workspace
         @home_workspace
       else
-        @home_workspace = @workspace
+        @home_workspace = workspace
       end
     end
 
@@ -25,11 +25,11 @@ def home_workspace
     # See IRB::WorkSpace.new for more information.
     def change_workspace(*_main)
       if _main.empty?
-        @workspace = home_workspace
+        replace_workspace(home_workspace)
         return main
       end
 
-      @workspace = WorkSpace.new(_main[0])
+      replace_workspace(WorkSpace.new(_main[0]))
 
       if !(class<<main;ancestors;end).include?(ExtendCommandBundle)
         main.extend ExtendCommandBundle
diff --git a/lib/irb/ext/eval_history.rb b/lib/irb/ext/eval_history.rb
index 316e0b312e8ef8..6c21ff00eea420 100644
--- a/lib/irb/ext/eval_history.rb
+++ b/lib/irb/ext/eval_history.rb
@@ -18,7 +18,7 @@ def set_last_value(value)
 
       if defined?(@eval_history) && @eval_history
         @eval_history_values.push @line_no, @last_value
-        @workspace.evaluate "__ = IRB.CurrentContext.instance_eval{@eval_history_values}"
+        workspace.evaluate "__ = IRB.CurrentContext.instance_eval{@eval_history_values}"
       end
 
       @last_value
@@ -49,7 +49,7 @@ def eval_history=(no)
         else
           @eval_history_values = EvalHistory.new(no)
           IRB.conf[:__TMP__EHV__] = @eval_history_values
-          @workspace.evaluate("__ = IRB.conf[:__TMP__EHV__]")
+          workspace.evaluate("__ = IRB.conf[:__TMP__EHV__]")
           IRB.conf.delete(:__TMP_EHV__)
         end
       else
diff --git a/lib/irb/ext/use-loader.rb b/lib/irb/ext/use-loader.rb
index 77c53b7d57d719..c8a3ea1fe89cf7 100644
--- a/lib/irb/ext/use-loader.rb
+++ b/lib/irb/ext/use-loader.rb
@@ -49,12 +49,12 @@ def use_loader=(opt)
       if IRB.conf[:USE_LOADER] != opt
         IRB.conf[:USE_LOADER] = opt
         if opt
-          (class<<@workspace.main;self;end).instance_eval {
+          (class<<workspace.main;self;end).instance_eval {
             alias_method :load, :irb_load
             alias_method :require, :irb_require
           }
         else
-          (class<<@workspace.main;self;end).instance_eval {
+          (class<<workspace.main;self;end).instance_eval {
             alias_method :load, :__original__load__IRB_use_loader__
             alias_method :require, :__original__require__IRB_use_loader__
           }
diff --git a/lib/irb/ext/workspaces.rb b/lib/irb/ext/workspaces.rb
index f447948fbd44b4..5a1edd89b736e9 100644
--- a/lib/irb/ext/workspaces.rb
+++ b/lib/irb/ext/workspaces.rb
@@ -6,21 +6,6 @@
 
 module IRB # :nodoc:
   class Context
-
-    # Size of the current WorkSpace stack
-    def irb_level
-      workspace_stack.size
-    end
-
-    # WorkSpaces in the current stack
-    def workspaces
-      if defined? @workspaces
-        @workspaces
-      else
-        @workspaces = []
-      end
-    end
-
     # Creates a new workspace with the given object or binding, and appends it
     # onto the current #workspaces stack.
     #
@@ -28,20 +13,16 @@ def workspaces
     # information.
     def push_workspace(*_main)
       if _main.empty?
-        if workspaces.empty?
-          print "No other workspace\n"
-          return nil
+        if @workspace_stack.size > 1
+          # swap the top two workspaces
+          previous_workspace, current_workspace = @workspace_stack.pop(2)
+          @workspace_stack.push current_workspace, previous_workspace
+        end
+      else
+        @workspace_stack.push WorkSpace.new(workspace.binding, _main[0])
+        if !(class<<main;ancestors;end).include?(ExtendCommandBundle)
+          main.extend ExtendCommandBundle
         end
-        ws = workspaces.pop
-        workspaces.push @workspace
-        @workspace = ws
-        return workspaces
-      end
-
-      workspaces.push @workspace
-      @workspace = WorkSpace.new(@workspace.binding, _main[0])
-      if !(class<<main;ancestors;end).include?(ExtendCommandBundle)
-        main.extend ExtendCommandBundle
       end
     end
 
@@ -50,11 +31,7 @@ def push_workspace(*_main)
     #
     # Also, see #push_workspace.
     def pop_workspace
-      if workspaces.empty?
-        print "workspace stack empty\n"
-        return
-      end
-      @workspace = workspaces.pop
+      @workspace_stack.pop if @workspace_stack.size > 1
     end
   end
 end
diff --git a/test/irb/test_command.rb b/test/irb/test_command.rb
index e27fac67b79686..d6c8c534f94529 100644
--- a/test/irb/test_command.rb
+++ b/test/irb/test_command.rb
@@ -482,7 +482,8 @@ class Foo; end
   class CwwsTest < WorkspaceCommandTestCase
     def test_cwws_returns_the_current_workspace_object
       out, err = execute_lines(
-        "cwws.class",
+        "cwws",
+        "self.class"
       )
 
       assert_empty err
@@ -493,51 +494,56 @@ def test_cwws_returns_the_current_workspace_object
   class PushwsTest < WorkspaceCommandTestCase
     def test_pushws_switches_to_new_workspace_and_pushes_the_current_one_to_the_stack
       out, err = execute_lines(
-        "pushws #{self.class}::Foo.new\n",
-        "cwws.class",
+        "pushws #{self.class}::Foo.new",
+        "self.class",
+        "popws",
+        "self.class"
       )
       assert_empty err
-      assert_include(out, "#{self.class}::Foo")
+
+      assert_match(/=> #{self.class}::Foo\n/, out)
+      assert_match(/=> #{self.class}\n$/, out)
     end
 
     def test_pushws_extends_the_new_workspace_with_command_bundle
       out, err = execute_lines(
-        "pushws Object.new\n",
+        "pushws Object.new",
         "self.singleton_class.ancestors"
       )
       assert_empty err
       assert_include(out, "IRB::ExtendCommandBundle")
     end
 
-    def test_pushws_prints_help_message_when_no_arg_is_given
+    def test_pushws_prints_workspace_stack_when_no_arg_is_given
       out, err = execute_lines(
-        "pushws\n",
+        "pushws",
       )
       assert_empty err
-      assert_match(/No other workspace/, out)
+      assert_include(out, "[#<TestIRB::PushwsTe...>]")
     end
-  end
 
-  class WorkspacesTest < WorkspaceCommandTestCase
-    def test_workspaces_returns_the_array_of_non_main_workspaces
+    def test_pushws_without_argument_swaps_the_top_two_workspaces
       out, err = execute_lines(
-        "pushws #{self.class}::Foo.new\n",
-        "workspaces.map { |w| w.class.name }",
+        "pushws #{self.class}::Foo.new",
+        "self.class",
+        "pushws",
+        "self.class"
       )
-
       assert_empty err
-      # self.class::Foo would be the current workspace
-      # self.class would be the old workspace that's pushed to the stack
-      assert_include(out, "=> [\"#{self.class}\"]")
+      assert_match(/=> #{self.class}::Foo\n/, out)
+      assert_match(/=> #{self.class}\n$/, out)
     end
+  end
 
-    def test_workspaces_returns_empty_array_when_no_workspaces_were_added
+  class WorkspacesTest < WorkspaceCommandTestCase
+    def test_workspaces_returns_the_stack_of_workspaces
       out, err = execute_lines(
-        "workspaces.map(&:to_s)",
+        "pushws #{self.class}::Foo.new\n",
+        "workspaces",
       )
 
       assert_empty err
-      assert_include(out, "=> []")
+      assert_match(/\[#<TestIRB::Workspac...>, #<TestIRB::Workspac...>]\n/, out)
     end
   end
 
@@ -557,7 +563,7 @@ def test_popws_prints_help_message_if_the_workspace_is_empty
         "popws\n",
       )
       assert_empty err
-      assert_match(/workspace stack empty/, out)
+      assert_match(/\[#<TestIRB::PopwsTes...>\]\n/, out)
     end
   end