Implement auto indent for multiline

4 files changed, 79 insertions, 2 deletions

diff --git a/lib/irb/input-method.rb b/lib/irb/input-method.rb
index 68d2ad280c..dbed7da14b 100644
--- a/lib/irb/input-method.rb
+++ b/lib/irb/input-method.rb
@@ -244,6 +244,10 @@ module IRB
       @prompt_proc = block
     end
 
+    def auto_indent(&block)
+      @auto_indent_proc = block
+    end
+
     # Reads the next line from this input method.
     #
     # See IO#gets for more information.
@@ -251,6 +255,7 @@ module IRB
       Reline.input = @stdin
       Reline.output = @stdout
       Reline.prompt_proc = @prompt_proc
+      Reline.auto_indent_proc = @auto_indent_proc
       if l = readmultiline(@prompt, false, &@check_termination_proc)
         HISTORY.push(l) if !l.empty?
         @line[@line_no += 1] = l + "\n"
diff --git a/lib/irb/ruby-lex.rb b/lib/irb/ruby-lex.rb
index 5f138701ef..036b831591 100644
--- a/lib/irb/ruby-lex.rb
+++ b/lib/irb/ruby-lex.rb
@@ -53,6 +53,34 @@ class RubyLex
         result
       end
     end
+    if @io.respond_to?(:auto_indent)
+      @io.auto_indent do |lines, line_index, byte_pointer, is_newline|
+        if is_newline
+          md = lines[line_index - 1].match(/(\A +)/)
+          prev_spaces = md.nil? ? 0 : md[1].count(' ')
+          indent_list = []
+          code = ''
+          lines.each_with_index { |l, i|
+            code << l + "\n"
+            @tokens = Ripper.lex(code)
+            indent_list << process_nesting_level
+          }
+          prev_indent = (line_index - 1).zero? ? 0 : indent_list[line_index - 2]
+          indent = indent_list[line_index - 1]
+          prev_spaces + (indent - prev_indent) * 2
+        else
+          code = line_index.zero? ? '' : lines[0..(line_index - 1)].map{ |l| l + "\n" }.join
+          code += lines[line_index].byteslice(0, byte_pointer)
+          @tokens = Ripper.lex(code)
+          indent, close_token = process_nesting_level(check_closing: true)
+          if close_token
+            indent * 2
+          else
+            nil
+          end
+        end
+      end
+    end
     if p.respond_to?(:call)
       @input = p
     elsif block_given?
@@ -257,13 +285,16 @@ class RubyLex
     false
   end
 
-  def process_nesting_level
-    @tokens.inject(0) { |indent, t|
+  def process_nesting_level(check_closing: false)
+    close_token = false
+    indent = @tokens.inject(0) { |indent, t|
+      close_token = false
       case t[1]
       when :on_lbracket, :on_lbrace, :on_lparen
         indent += 1
       when :on_rbracket, :on_rbrace, :on_rparen
         indent -= 1
+        close_token = true
       when :on_kw
         case t[2]
         when 'def', 'do', 'case', 'for', 'begin', 'class', 'module'
@@ -273,11 +304,17 @@ class RubyLex
           indent += 1 unless t[3].allbits?(Ripper::EXPR_LABEL)
         when 'end'
           indent -= 1
+          close_token = true
         end
       end
       # percent literals are not indented
       indent
     }
+    if check_closing
+      [indent, close_token]
+    else
+      indent
+    end
   end
 
   def check_string_literal
diff --git a/lib/reline.rb b/lib/reline.rb
index ff29a2539b..191d8bb158 100644
--- a/lib/reline.rb
+++ b/lib/reline.rb
@@ -186,6 +186,15 @@ module Reline
     @@prompt_proc = p
   end
 
+  @@auto_indent_proc = nil
+  def self.auto_indent_proc
+    @@auto_indent_proc
+  end
+  def self.auto_indent_proc=(p)
+    raise ArgumentError unless p.is_a?(Proc)
+    @@auto_indent_proc = p
+  end
+
   @@pre_input_hook = nil
   def self.pre_input_hook
     @@pre_input_hook
@@ -330,6 +339,7 @@ module Reline
     @@line_editor.completion_proc = @@completion_proc
     @@line_editor.output_modifier_proc = @@output_modifier_proc
     @@line_editor.prompt_proc = @@prompt_proc
+    @@line_editor.auto_indent_proc = @@auto_indent_proc
     @@line_editor.dig_perfect_match_proc = @@dig_perfect_match_proc
     @@line_editor.pre_input_hook = @@pre_input_hook
     @@line_editor.rerender
diff --git a/lib/reline/line_editor.rb b/lib/reline/line_editor.rb
index be774a3e06..629d561686 100644
--- a/lib/reline/line_editor.rb
+++ b/lib/reline/line_editor.rb
@@ -12,6 +12,7 @@ class Reline::LineEditor
   attr_accessor :completion_proc
   attr_accessor :output_modifier_proc
   attr_accessor :prompt_proc
+  attr_accessor :auto_indent_proc
   attr_accessor :pre_input_hook
   attr_accessor :dig_perfect_match_proc
   attr_writer :output
@@ -116,6 +117,7 @@ class Reline::LineEditor
     @highest_in_all = 1
     @line_backup_in_history = nil
     @multibyte_buffer = String.new(encoding: 'ASCII-8BIT')
+    @check_new_auto_indent = false
   end
 
   def multiline_on
@@ -781,6 +783,28 @@ class Reline::LineEditor
     unless completion_occurs
       @completion_state = CompletionState::NORMAL
     end
+    if @is_multiline and @auto_indent_proc
+      if @previous_line_index
+        new_lines = whole_lines(index: @previous_line_index, line: @line)
+      else
+        new_lines = whole_lines
+      end
+      new_indent = @auto_indent_proc.(new_lines, @line_index, @byte_pointer, @check_new_auto_indent)
+      if new_indent
+        md = @buffer_of_lines[@line_index].match(/\A */)
+        prev_indent = md[0].count(' ')
+        if @check_new_auto_indent
+          @buffer_of_lines[@line_index] = ' ' * new_indent + @buffer_of_lines[@line_index].gsub(/\A */, '')
+          @cursor = new_indent
+          @byte_pointer = new_indent
+        else
+          @line = ' ' * new_indent + @line.gsub(/\A */, '')
+          @cursor -= prev_indent - new_indent
+          @byte_pointer -= prev_indent - new_indent
+        end
+      end
+      @check_new_auto_indent = false
+    end
   end
 
   def retrieve_completion_block
@@ -977,6 +1001,7 @@ class Reline::LineEditor
       cursor_line = @line.byteslice(0, @byte_pointer)
       insert_new_line(cursor_line, next_line)
       @cursor = 0
+      @check_new_auto_indent = true
     end
   end