Refactor prompt generation logic

1 files changed, 17 insertions, 34 deletions

diff --git a/lib/reline/line_editor.rb b/lib/reline/line_editor.rb
index 65ad21fbe7..3a7ba3e414 100644
--- a/lib/reline/line_editor.rb
+++ b/lib/reline/line_editor.rb
@@ -60,11 +60,20 @@ class Reline::LineEditor
     reset_variables
   end
 
-  private def check_multiline_prompt(buffer, prompt, special_prompt)
+  private def check_multiline_prompt(buffer, prompt)
+    special_prompt = nil
+    if @vi_arg
+      prompt = "(arg: #{@vi_arg}) "
+      @rerender_all = true
+    elsif @searching_prompt
+      prompt = @searching_prompt
+      @rerender_all = true
+    else
+      prompt = @prompt
+    end
     if @prompt_proc
       prompt_list = @prompt_proc.(buffer)
-      prompt_list[@line_index] = special_prompt if special_prompt
-      prompt_list.map!{ special_prompt } if @searching_prompt
+      prompt_list.map!{ prompt } if @vi_arg or @searching_prompt
       prompt = prompt_list[@line_index]
       prompt_width = calculate_width(prompt, true)
       [prompt, prompt_width, prompt_list]
@@ -90,22 +99,9 @@ class Reline::LineEditor
         @rerender_all = true
         rerender
       else
-        special_prompt = nil
-        if @vi_arg
-          prompt = "(arg: #{@vi_arg}) "
-          prompt_width = calculate_width(prompt)
-          special_prompt = prompt
-        elsif @searching_prompt
-          prompt = @searching_prompt
-          prompt_width = calculate_width(prompt)
-          special_prompt = prompt
-        else
-          prompt = @prompt
-          prompt_width = calculate_width(prompt, true)
-        end
         back = 0
         new_buffer = whole_lines
-        prompt, prompt_width, prompt_list = check_multiline_prompt(new_buffer, prompt, special_prompt)
+        prompt, prompt_width, prompt_list = check_multiline_prompt(new_buffer, prompt)
         new_buffer.each_with_index do |line, index|
           prompt_width = calculate_width(prompt_list[index], true) if @prompt_proc
           width = prompt_width + calculate_width(line)
@@ -320,24 +316,11 @@ class Reline::LineEditor
       move_cursor_up(@highest_in_all - 1 - @first_line_started_from)
       @menu_info = nil
     end
-    special_prompt = nil
-    if @vi_arg
-      prompt = "(arg: #{@vi_arg}) "
-      prompt_width = calculate_width(prompt)
-      special_prompt = prompt
-    elsif @searching_prompt
-      prompt = @searching_prompt
-      prompt_width = calculate_width(prompt)
-      special_prompt = prompt
-    else
-      prompt = @prompt
-      prompt_width = calculate_width(prompt, true)
-    end
+    prompt, prompt_width, prompt_list = check_multiline_prompt(whole_lines, prompt)
     if @cleared
       Reline::IOGate.clear_screen
       @cleared = false
       back = 0
-      prompt, prompt_width, prompt_list = check_multiline_prompt(whole_lines, prompt, special_prompt)
       modify_lines(whole_lines).each_with_index do |line, index|
         if @prompt_proc
           pr = prompt_list[index]
@@ -363,7 +346,7 @@ class Reline::LineEditor
       else
         new_lines = whole_lines
       end
-      prompt, prompt_width, prompt_list = check_multiline_prompt(new_lines, prompt, special_prompt)
+      prompt, prompt_width, prompt_list = check_multiline_prompt(new_lines, prompt)
       all_height = new_lines.inject(0) { |result, line|
         result + calculate_height_by_width(prompt_width + calculate_width(line)) # TODO prompt_list
       }
@@ -427,7 +410,7 @@ class Reline::LineEditor
       Reline::IOGate.move_cursor_column(0)
       back = 0
       new_buffer = whole_lines
-      prompt, prompt_width, prompt_list = check_multiline_prompt(new_buffer, prompt, special_prompt)
+      prompt, prompt_width, prompt_list = check_multiline_prompt(new_buffer, prompt)
       new_buffer.each_with_index do |line, index|
         prompt_width = calculate_width(prompt_list[index], true) if @prompt_proc
         width = prompt_width + calculate_width(line)
@@ -479,7 +462,7 @@ class Reline::LineEditor
     end
     line = modify_lines(whole_lines)[@line_index]
     if @is_multiline
-      prompt, prompt_width, prompt_list = check_multiline_prompt(whole_lines, prompt, special_prompt)
+      prompt, prompt_width, prompt_list = check_multiline_prompt(whole_lines, prompt)
       if finished?
         # Always rerender on finish because output_modifier_proc may return a different output.
         render_partial(prompt, prompt_width, line)