From c1b4add09427f6b7fa715ce29e9598a41ea292ba Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?fn=20=E2=8C=83=20=E2=8C=A5?=
 <70830482+FnControlOption@users.noreply.github.com>
Date: Sun, 24 Oct 2021 11:56:01 -0700
Subject: [PATCH] Add `Source::Corrector` and `Source::Rewriter`

---
 spec/ameba/source/rewriter_spec.cr  | 110 +++++++++++++++++
 src/ameba.cr                        |   1 +
 src/ameba/issue.cr                  |  10 +-
 src/ameba/reportable.cr             |  55 +++++++--
 src/ameba/rule/base.cr              |   4 +-
 src/ameba/source/corrector.cr       |  57 +++++++++
 src/ameba/source/rewriter.cr        | 132 ++++++++++++++++++++
 src/ameba/source/rewriter/action.cr | 181 ++++++++++++++++++++++++++++
 8 files changed, 537 insertions(+), 13 deletions(-)
 create mode 100644 spec/ameba/source/rewriter_spec.cr
 create mode 100644 src/ameba/source/corrector.cr
 create mode 100644 src/ameba/source/rewriter.cr
 create mode 100644 src/ameba/source/rewriter/action.cr

diff --git a/spec/ameba/source/rewriter_spec.cr b/spec/ameba/source/rewriter_spec.cr
new file mode 100644
index 00000000..7d97bb9b
--- /dev/null
+++ b/spec/ameba/source/rewriter_spec.cr
@@ -0,0 +1,110 @@
+require "../../spec_helper"
+
+class Ameba::Source
+  describe Rewriter do
+    code = "puts(:hello, :world)"
+    hello = {5, 11}
+    comma_space = {11, 13}
+    world = {13, 19}
+
+    it "can remove" do
+      rewriter = Rewriter.new(code)
+      rewriter.remove(*hello)
+      rewriter.process.should eq "puts(, :world)"
+    end
+
+    it "can insert before" do
+      rewriter = Rewriter.new(code)
+      rewriter.insert_before(*world, "42, ")
+      rewriter.process.should eq "puts(:hello, 42, :world)"
+    end
+
+    it "can insert after" do
+      rewriter = Rewriter.new(code)
+      rewriter.insert_after(*hello, ", 42")
+      rewriter.process.should eq "puts(:hello, 42, :world)"
+    end
+
+    it "can wrap" do
+      rewriter = Rewriter.new(code)
+      rewriter.wrap(*hello, '[', ']')
+      rewriter.process.should eq "puts([:hello], :world)"
+    end
+
+    it "can replace" do
+      rewriter = Rewriter.new(code)
+      rewriter.replace(*hello, ":hi")
+      rewriter.process.should eq "puts(:hi, :world)"
+    end
+
+    it "accepts crossing deletions" do
+      rewriter = Rewriter.new(code)
+      rewriter.remove(hello[0], comma_space[1])
+      rewriter.remove(comma_space[0], world[1])
+      rewriter.process.should eq "puts()"
+    end
+
+    it "accepts multiple actions" do
+      rewriter = Rewriter.new(code)
+      rewriter.replace(*comma_space, " => ")
+      rewriter.wrap(hello[0], world[1], '{', '}')
+      rewriter.replace(*world, ":everybody")
+      rewriter.wrap(*world, '[', ']')
+      rewriter.process.should eq "puts({:hello => [:everybody]})"
+    end
+
+    it "can wrap the same range" do
+      rewriter = Rewriter.new(code)
+      rewriter.wrap(*hello, '(', ')')
+      rewriter.wrap(*hello, '[', ']')
+      rewriter.process.should eq "puts([(:hello)], :world)"
+    end
+
+    it "can insert on empty ranges" do
+      rewriter = Rewriter.new(code)
+      rewriter.insert_before(hello[0], '{')
+      rewriter.replace(hello[0], hello[0], 'x')
+      rewriter.insert_after(hello[0], '}')
+      rewriter.insert_before(hello[1], '[')
+      rewriter.replace(hello[1], hello[1], 'y')
+      rewriter.insert_after(hello[1], ']')
+      rewriter.process.should eq "puts({x}:hello[y], :world)"
+    end
+
+    it "can replace the same range" do
+      rewriter = Rewriter.new(code)
+      rewriter.replace(*hello, ":hi")
+      rewriter.replace(*hello, ":hey")
+      rewriter.process.should eq "puts(:hey, :world)"
+    end
+
+    it "can swallow insertions" do
+      rewriter = Rewriter.new(code)
+      rewriter.wrap(hello[0] + 1, hello[1], "__", "__")
+      rewriter.replace(world[0], world[1] - 2, "xx")
+      rewriter.replace(hello[0], world[1], ":hi")
+      rewriter.process.should eq "puts(:hi)"
+    end
+
+    it "rejects out-of-range ranges" do
+      rewriter = Rewriter.new(code)
+      expect_raises(IndexError) { rewriter.insert_before(0, 100, "hola") }
+    end
+
+    it "ignores trivial actions" do
+      rewriter = Rewriter.new(code)
+      rewriter.empty?.should be_true
+
+      # This is a trivial wrap
+      rewriter.wrap(2, 5, "", "")
+      rewriter.empty?.should be_true
+
+      # This is a trivial deletion
+      rewriter.remove(2, 2)
+      rewriter.empty?.should be_true
+
+      rewriter.remove(2, 5)
+      rewriter.empty?.should be_false
+    end
+  end
+end
diff --git a/src/ameba.cr b/src/ameba.cr
index f16c62e6..2a34d0ef 100644
--- a/src/ameba.cr
+++ b/src/ameba.cr
@@ -2,6 +2,7 @@ require "./ameba/*"
 require "./ameba/ast/**"
 require "./ameba/rule/**"
 require "./ameba/formatter/*"
+require "./ameba/source/**"
 
 # Ameba's entry module.
 #
diff --git a/src/ameba/issue.cr b/src/ameba/issue.cr
index b913eac3..06619686 100644
--- a/src/ameba/issue.cr
+++ b/src/ameba/issue.cr
@@ -24,12 +24,20 @@ module Ameba
     delegate :enabled?, :disabled?,
       to: status
 
-    def initialize(@rule, @location, @end_location, @message, status : Status? = nil)
+    def initialize(@rule, @location, @end_location, @message, status : Status? = nil, @block : (Source::Corrector ->)? = nil)
       @status = status || Status::Enabled
     end
 
     def syntax?
       rule.is_a?(Rule::Lint::Syntax)
     end
+
+    def correctable?
+      !@block.nil?
+    end
+
+    def correct(corrector)
+      @block.try &.call(corrector)
+    end
   end
 end
diff --git a/src/ameba/reportable.cr b/src/ameba/reportable.cr
index 4632fa91..a85bf28b 100644
--- a/src/ameba/reportable.cr
+++ b/src/ameba/reportable.cr
@@ -5,41 +5,76 @@ module Ameba
     getter issues = [] of Issue
 
     # Adds a new issue to the list of issues.
-    def add_issue(rule, location, end_location, message, status : Issue::Status? = nil) : Issue
+    def add_issue(rule, location, end_location, message, status : Issue::Status? = nil, block : (Source::Corrector ->)? = nil) : Issue
       status ||=
         Issue::Status::Disabled if location_disabled?(location, rule)
 
-      Issue.new(rule, location, end_location, message, status).tap do |issue|
+      Issue.new(rule, location, end_location, message, status, block).tap do |issue|
         issues << issue
       end
     end
 
+    # :ditto:
+    def add_issue(rule, location, end_location, message, status : Issue::Status? = nil, &block : Source::Corrector ->) : Issue
+      add_issue rule, location, end_location, message, status, block
+    end
+
     # Adds a new issue for Crystal AST *node*.
-    def add_issue(rule, node : Crystal::ASTNode, message, status : Issue::Status? = nil) : Issue
-      add_issue rule, node.location, node.end_location, message, status
+    def add_issue(rule, node : Crystal::ASTNode, message, status : Issue::Status? = nil, block : (Source::Corrector ->)? = nil) : Issue
+      add_issue rule, node.location, node.end_location, message, status, block
+    end
+
+    # :ditto:
+    def add_issue(rule, node : Crystal::ASTNode, message, status : Issue::Status? = nil, &block : Source::Corrector ->) : Issue
+      add_issue rule, node, message, status, block
     end
 
     # Adds a new issue for Crystal *token*.
-    def add_issue(rule, token : Crystal::Token, message, status : Issue::Status? = nil) : Issue
-      add_issue rule, token.location, nil, message, status
+    def add_issue(rule, token : Crystal::Token, message, status : Issue::Status? = nil, block : (Source::Corrector ->)? = nil) : Issue
+      add_issue rule, token.location, nil, message, status, block
+    end
+
+    # :ditto:
+    def add_issue(rule, token : Crystal::Token, message, status : Issue::Status? = nil, &block : Source::Corrector ->) : Issue
+      add_issue rule, token, message, status, block
     end
 
     # Adds a new issue for *location* defined by line and column numbers.
-    def add_issue(rule, location : {Int32, Int32}, message, status : Issue::Status? = nil) : Issue
+    def add_issue(rule, location : {Int32, Int32}, message, status : Issue::Status? = nil, block : (Source::Corrector ->)? = nil) : Issue
       location =
         Crystal::Location.new(path, *location)
 
-      add_issue rule, location, nil, message, status
+      add_issue rule, location, nil, message, status, block
+    end
+
+    # :ditto:
+    def add_issue(rule, location : {Int32, Int32}, message, status : Issue::Status? = nil, &block : Source::Corrector ->) : Issue
+      add_issue rule, location, message, status, block
     end
 
     # Adds a new issue for *location* and *end_location* defined by line and column numbers.
-    def add_issue(rule, location : {Int32, Int32}, end_location : {Int32, Int32}, message, status : Issue::Status? = nil) : Issue
+    def add_issue(rule,
+                  location : {Int32, Int32},
+                  end_location : {Int32, Int32},
+                  message,
+                  status : Issue::Status? = nil,
+                  block : (Source::Corrector ->)? = nil) : Issue
       location =
         Crystal::Location.new(path, *location)
       end_location =
         Crystal::Location.new(path, *end_location)
 
-      add_issue rule, location, end_location, message, status
+      add_issue rule, location, end_location, message, status, block
+    end
+
+    # :ditto:
+    def add_issue(rule,
+                  location : {Int32, Int32},
+                  end_location : {Int32, Int32},
+                  message,
+                  status : Issue::Status? = nil,
+                  &block : Source::Corrector ->) : Issue
+      add_issue rule, location, end_location, message, status, block
     end
 
     # Returns `true` if the list of not disabled issues is empty, `false` otherwise.
diff --git a/src/ameba/rule/base.cr b/src/ameba/rule/base.cr
index 7777167c..6a3090d8 100644
--- a/src/ameba/rule/base.cr
+++ b/src/ameba/rule/base.cr
@@ -112,8 +112,8 @@ module Ameba::Rule
       name.hash
     end
 
-    macro issue_for(*args)
-      source.add_issue self, {{*args}}
+    macro issue_for(*args, &block)
+      source.add_issue self, {{*args}} {{block}}
     end
 
     protected def self.rule_name
diff --git a/src/ameba/source/corrector.cr b/src/ameba/source/corrector.cr
new file mode 100644
index 00000000..ae3725e7
--- /dev/null
+++ b/src/ameba/source/corrector.cr
@@ -0,0 +1,57 @@
+require "./rewriter"
+
+class Ameba::Source
+  class Corrector
+    def self.correct(source : Source)
+      code = source.code
+      corrector = new(code)
+      source.issues.each(&.correct(corrector))
+      corrected_code = corrector.process
+      return if code == corrected_code
+
+      corrected_code
+    end
+
+    @line_sizes : Array(Int32)
+
+    def initialize(code : String)
+      @rewriter = Rewriter.new(code)
+      @line_sizes = code.lines(chomp: false).map(&.size)
+    end
+
+    alias SourceLocation = Crystal::Location | {Int32, Int32}
+
+    def replace(location : SourceLocation, end_location : SourceLocation, content)
+      @rewriter.replace(loc_to_pos(location), loc_to_pos(end_location) + 1, content)
+    end
+
+    def wrap(location : SourceLocation, end_location : SourceLocation, insert_before, insert_after)
+      @rewriter.wrap(loc_to_pos(location), loc_to_pos(end_location) + 1, insert_before, insert_after)
+    end
+
+    def remove(location : SourceLocation, end_location : SourceLocation)
+      @rewriter.remove(loc_to_pos(location), loc_to_pos(end_location) + 1)
+    end
+
+    def insert_before(location : SourceLocation, content)
+      @rewriter.insert_before(loc_to_pos(location), content)
+    end
+
+    def insert_after(location : SourceLocation, content)
+      @rewriter.insert_after(loc_to_pos(location) + 1, content)
+    end
+
+    private def loc_to_pos(location : SourceLocation)
+      if location.is_a?(Crystal::Location)
+        line, column = location.line_number, location.column_number
+      else
+        line, column = location
+      end
+      @line_sizes[0...line - 1].sum + (column - 1)
+    end
+
+    def process
+      @rewriter.process
+    end
+  end
+end
diff --git a/src/ameba/source/rewriter.cr b/src/ameba/source/rewriter.cr
new file mode 100644
index 00000000..62b5c65b
--- /dev/null
+++ b/src/ameba/source/rewriter.cr
@@ -0,0 +1,132 @@
+class Ameba::Source
+  # This class performs the heavy lifting in the source rewriting process.
+  # It schedules code updates to be performed in the correct order.
+  #
+  # For simple cases, the resulting source will be obvious.
+  #
+  # Examples for more complex cases follow. Assume these examples are acting on
+  # the source `puts(:hello, :world)`. The methods `#wrap`, `#remove`, etc.
+  # receive a range as the first two arguments; for clarity, examples below use
+  # English sentences and a string of raw code instead.
+  #
+  # ## Overlapping deletions:
+  #
+  # * remove `:hello, `
+  # * remove `, :world`
+  #
+  # The overlapping ranges are merged and `:hello, :world` will be removed.
+  #
+  # ## Multiple actions at the same end points:
+  #
+  # Results will always be independent of the order they were given.
+  # Exception: rewriting actions done on exactly the same range (covered next).
+  #
+  # Example:
+  #
+  # * replace `, ` by ` => `
+  # * wrap `:hello, :world` with `{` and `}`
+  # * replace `:world` with `:everybody`
+  # * wrap `:world` with `[`, `]`
+  #
+  # The resulting string will be `puts({:hello => [:everybody]})`
+  # and this result is independent of the order the instructions were given in.
+  #
+  # ## Multiple wraps on same range:
+  #
+  # * wrap `:hello` with `(` and `)`
+  # * wrap `:hello` with `[` and `]`
+  #
+  # The wraps are combined in order given and results would be `puts([(:hello)], :world)`.
+  #
+  # ## Multiple replacements on same range:
+  #
+  # * replace `:hello` by `:hi`, then
+  # * replace `:hello` by `:hey`
+  #
+  # The replacements are made in the order given, so the latter replacement
+  # supersedes the former and `:hello` will be replaced by `:hey`.
+  #
+  # ## Swallowed insertions:
+  #
+  # * wrap `world` by `__`, `__`
+  # * replace `:hello, :world` with `:hi`
+  #
+  # A containing replacement will swallow the contained rewriting actions
+  # and `:hello, :world` will be replaced by `:hi`.
+  #
+  # ## Implementation
+  #
+  # The updates are organized in a tree, according to the ranges they act on
+  # (where children are strictly contained by their parent).
+  class Rewriter
+    getter code
+
+    def initialize(@code : String)
+      @action_root = Rewriter::Action.new(0, code.size)
+    end
+
+    # Returns true if no (non trivial) update has been recorded
+    def empty?
+      @action_root.empty?
+    end
+
+    # Replaces the code of the given range with `content`.
+    def replace(begin_pos, end_pos, content)
+      combine(begin_pos, end_pos, replacement: content.to_s)
+    end
+
+    # Inserts the given strings before and after the given range.
+    def wrap(begin_pos, end_pos, insert_before, insert_after)
+      combine(begin_pos, end_pos, insert_before: insert_before.to_s, insert_after: insert_after.to_s)
+    end
+
+    # Shortcut for `replace(begin_pos, end_pos, "")`
+    def remove(begin_pos, end_pos)
+      replace(begin_pos, end_pos, "")
+    end
+
+    # Shortcut for `wrap(begin_pos, end_pos, content, nil)`
+    def insert_before(begin_pos, end_pos, content)
+      wrap(begin_pos, end_pos, content, nil)
+    end
+
+    # Shortcut for `wrap(begin_pos, end_pos, nil, content)`
+    def insert_after(begin_pos, end_pos, content)
+      wrap(begin_pos, end_pos, nil, content)
+    end
+
+    # Shortcut for `insert_before(pos, pos, content)`
+    def insert_before(pos, content)
+      insert_before(pos, pos, content)
+    end
+
+    # Shortcut for `insert_after(pos, pos, content)`
+    def insert_after(pos, content)
+      insert_after(pos, pos, content)
+    end
+
+    # Applies all scheduled changes and returns modified source as a new string.
+    def process
+      String.build do |io|
+        last_end = 0
+        @action_root.ordered_replacements.each do |begin_pos, end_pos, replacement|
+          io << code[last_end...begin_pos] << replacement
+          last_end = end_pos
+        end
+        io << code[last_end...code.size]
+      end
+    end
+
+    protected def combine(begin_pos, end_pos, **attributes)
+      check_range_validity(begin_pos, end_pos)
+      action = Rewriter::Action.new(begin_pos, end_pos, **attributes)
+      @action_root = @action_root.combine(action)
+    end
+
+    private def check_range_validity(begin_pos, end_pos)
+      if begin_pos < 0 || end_pos > code.size
+        raise IndexError.new("The range #{begin_pos}...#{end_pos} is outside the bounds of the source")
+      end
+    end
+  end
+end
diff --git a/src/ameba/source/rewriter/action.cr b/src/ameba/source/rewriter/action.cr
new file mode 100644
index 00000000..c33e607c
--- /dev/null
+++ b/src/ameba/source/rewriter/action.cr
@@ -0,0 +1,181 @@
+class Ameba::Source::Rewriter
+  # Actions are arranged in a tree and get combined so that:
+  # - children are strictly contained by their parent
+  # - siblings all disjoint from one another and ordered
+  # - only actions with `replacement == nil` may have children
+  class Action
+    getter begin_pos : Int32
+    getter end_pos : Int32
+    getter replacement : String?
+    getter insert_before : String
+    getter insert_after : String
+    protected getter children : Array(Action)
+
+    def initialize(@begin_pos,
+                   @end_pos,
+                   @insert_before = "",
+                   @replacement = nil,
+                   @insert_after = "",
+                   @children = [] of Action)
+    end
+
+    def combine(action)
+      return self if action.empty? # Ignore empty action
+
+      do_combine(action)
+    end
+
+    def empty?
+      replacement = @replacement
+      @insert_before.empty? &&
+        @insert_after.empty? &&
+        @children.empty? &&
+        (replacement.nil? || (replacement.empty? && @begin_pos == @end_pos))
+    end
+
+    def ordered_replacements
+      replacement = @replacement
+      reps = [] of {Int32, Int32, String}
+      reps << {@begin_pos, @begin_pos, @insert_before} unless @insert_before.empty?
+      reps << {@begin_pos, @end_pos, replacement} if replacement
+      reps.concat(@children.flat_map(&.ordered_replacements))
+      reps << {@end_pos, @end_pos, @insert_after} unless @insert_after.empty?
+      reps
+    end
+
+    def insertion?
+      replacement = @replacement
+      !@insert_before.empty? || !@insert_after.empty? || (replacement && !replacement.empty?)
+    end
+
+    protected def with(*,
+                       begin_pos = @begin_pos,
+                       end_pos = @end_pos,
+                       insert_before = @insert_before,
+                       replacement = @replacement,
+                       insert_after = @insert_after,
+                       children = @children)
+      children = [] of Action if replacement
+      self.class.new(begin_pos, end_pos, insert_before, replacement, insert_after, children)
+    end
+
+    # Assumes *action* is contained within `@begin_pos...@end_pos` and has no children
+    protected def do_combine(action)
+      if action.begin_pos == @begin_pos && action.end_pos == @end_pos
+        merge(action)
+      else
+        place_in_hierarchy(action)
+      end
+    end
+
+    protected def place_in_hierarchy(action)
+      family = analyse_hierarchy(action)
+      sibling_left, sibling_right = family[:sibling_left], family[:sibling_right]
+
+      if (fusible = family[:fusible])
+        child = family[:child]
+        child ||= [] of Action
+        fuse_deletions(action, fusible, sibling_left + child + sibling_right)
+      else
+        extra_sibling = if (parent = family[:parent])
+                          # action should be a descendant of one of the children
+                          parent.do_combine(action)
+                        elsif (child = family[:child])
+                          # or it should become the parent of some of the children,
+                          action.with(children: child).combine_children(action.children)
+                        else
+                          # or else it should become an additional child
+                          action
+                        end
+        self.with(children: sibling_left + [extra_sibling] + sibling_right)
+      end
+    end
+
+    # Assumes *more_children* all contained within `@begin_pos...@end_pos`
+    protected def combine_children(more_children)
+      more_children.reduce(self) do |parent, new_child|
+        parent.place_in_hierarchy(new_child)
+      end
+    end
+
+    protected def fuse_deletions(action, fusible, other_siblings)
+      without_fusible = self.with(children: other_siblings)
+      fusible = [action] + fusible
+      fused_begin_pos = fusible.map(&.begin_pos).min
+      fused_end_pos = fusible.map(&.end_pos).max
+      fused_deletion = action.with(begin_pos: fused_begin_pos, end_pos: fused_end_pos)
+      without_fusible.do_combine(fused_deletion)
+    end
+
+    # Similar to `@children.bsearch_index || size` except allows for a starting point
+    protected def bsearch_child_index(from = 0)
+      size = @children.size
+      (from...size).bsearch { |i| yield @children[i] } || size
+    end
+
+    # Returns the children in a hierarchy with respect to *action*:
+    #
+    # - `:sibling_left`, `:sibling_right` (for those that are disjoint from *action*)
+    # - `:parent` (in case one of our children contains *action*)
+    # - `:child` (in case *action* strictly contains some of our children)
+    # - `:fusible` (in case *action* overlaps some children but they can be fused in one deletion)
+    #
+    # In case a child has equal range to *action*, it is returned as `:parent`
+    #
+    # Reminder: an empty range 1...1 is considered disjoint from 1...10
+    protected def analyse_hierarchy(action) # ameba:disable Metrics/CyclomaticComplexity
+      # left_index is the index of the first child that isn't completely to the left of action
+      left_index = bsearch_child_index { |child| child.end_pos > action.begin_pos }
+      # right_index is the index of the first child that is completely on the right of action
+      start = left_index == 0 ? 0 : left_index - 1 # See "corner case" below for reason of -1
+      right_index = bsearch_child_index(start) { |child| child.begin_pos >= action.end_pos }
+      center = right_index - left_index
+      case center
+      when 0
+        # All children are disjoint from action, nothing else to do
+      when -1
+        # Corner case: if a child has empty range == action's range
+        # then it will appear to be both disjoint and to the left of action,
+        # as well as disjoint and to the right of action.
+        # Since ranges are equal, we return it as parent
+        left_index -= 1  # Fix indices, as otherwise this child would be
+        right_index += 1 # considered as a sibling (both left and right!)
+        parent = @children[left_index]
+      else
+        overlap_left = @children[left_index].begin_pos <=> action.begin_pos
+        overlap_right = @children[right_index - 1].end_pos <=> action.end_pos
+        raise "Unable to compare begin pos" if overlap_left.nil?
+        raise "Unable to compare end pos" if overlap_right.nil?
+
+        # For one child to be the parent of action, we must have:
+        if center == 1 && overlap_left <= 0 && overlap_right >= 0
+          parent = @children[left_index]
+        else
+          # Otherwise consider all non disjoint elements (center) to be contained...
+          contained = @children[left_index...right_index]
+          fusible = [] of Action
+          fusible << contained.shift if overlap_left < 0 # ... but check first and last one
+          fusible << contained.pop if overlap_right > 0  # ... for overlaps
+          fusible = nil if fusible.empty?
+        end
+      end
+
+      {
+        parent:        parent,
+        sibling_left:  @children[0...left_index],
+        sibling_right: @children[right_index...@children.size],
+        fusible:       fusible,
+        child:         contained,
+      }
+    end
+
+    # Assumes *action* has the exact same range and has no children
+    protected def merge(action)
+      self.with(
+        insert_before: "#{action.insert_before}#{insert_before}",
+        replacement: action.replacement || @replacement,
+        insert_after: "#{insert_after}#{action.insert_after}",
+      ).combine_children(action.children)
+    end
+  end
+end