strict_colorizer.rb

class Colorizer
  attr_reader :code, :ast, :lines, :line_indices
  def initialize(code)
    @code = code
    @ast = RubyVM::AbstractSyntaxTree.parse code
    @lines = code.lines
    @line_indices = [0]
    @lines.each { @line_indices << @line_indices.last + _1.size }
  end

  def nodes
    @nodes ||= [].tap do |nodes|
      traverse(ast) { nodes << _1 }
    end
  end

  TYPE_PRIORITY = %i[
    CONST
    LVAR
    CALL
    FCALL
    MATCH2
    STR
    LIT
    LIST
  ]

  def colorize
    adds = []
    deletes = []
    nodes.each do |node|
      range = text_range_of node
      next if range.begin == range.end
      item = [node.type, range.begin, range.end]
      (adds[range.begin] ||= []) << item
      (deletes[range.end] ||= []) << item
    end

    matches = []
    match_chars = code.chars.each_with_index.map do |c, i|
      if adds[i] || deletes[i]
        new_items = (adds[i] ||= []).sort_by(&:last).reverse
        matches = matches - (deletes[i] || []) + new_items
      end
      [matches, c]
    end

    match_chars.chunk(&:first).each do |matches, items|
      min_size = matches.map { _3 - _2 }.min
      types = matches.select { _3 - _2 == min_size }.map(&:first)
      type = (TYPE_PRIORITY & types).first || types.first
      text = items.map(&:last).join
      $> << colorize_node(type, text)
    end
  end

  module Terminal
    COLORS = {
      red: 31,
      green: 32,
      yellow: 33,
      blue: 34,
      magenta: 35,
      cyan: 36,
      gray: 37
    }
    BOLD = 1
    LIGHT = 2
    ITALIC = 3
    UNDERLINE = 4
    
    COLORS.each_key do |color|
      define_singleton_method color do |text, **args|
        self.text text, color, **args
      end
    end

    def self.text(text, color, bold: false, light: false, italic: false, underline: false)
      colors = [COLORS[color], (BOLD if bold), (LIGHT if light), (ITALIC if italic), (UNDERLINE if underline)].compact
      colors.empty? ? text : "\e[#{colors.join(';')}m#{text}\e[m"
    end
  end
  BIN_OP = %w[+ - * / ** % && || & | ^]
  OP = BIN_OP + BIN_OP.map { _1 + '=' } + %w[! != = == ~]

  def colorize_literal(text)
    case text[0]
    when '_'
      Terminal.cyan text, bold: true
    when '?', /\d/
      Terminal.blue text, bold: true
    when '"', "'", '%', '/'
      Terminal.red text
    when ':'
      Terminal.yellow text
    else
      text.end_with?(':') ? Terminal.magenta(text) : text
    end
  end

  KEYWORDS = %w[if while do end def until unless class module BEGIN END yield when case __END__ in for return next break super]

  def colorize_other(text)
    if text.include? '#'
      a, b = text.split('#', 2)
      return colorize_other(a) + Terminal.gray('#' + b)
    end
    return Terminal.magenta text if text.match? /: \z/
    if /\A(?<defpart>def *)(?<name>[^\s()]+)/ =~ text
      return Terminal.green(defpart) + Terminal.blue(name, bold: true) + text[defpart.size + name.size..]
    end
    text.scan(/[a-zA-Z0-9]+|[^a-zA-Z0-9]+/).map do |s|
      if KEYWORDS.include? s
        Terminal.green s
      elsif s =~ /\A[A-Z]/
        Terminal.blue s, bold: true, underline: true
      else
        s
      end
    end.join
  end

  def colorize_node(type, text)
    case type
    when :STR, :LIT
      colorize_literal text
    when :NIL, :TRUE, :FALSE, :SELF
      if %w[nil true false self].include? text
        Terminal.cyan text, bold: true
      else
        text
      end
    when :CONST
      Terminal.blue text, bold: true, underline: true
    when :GVAR
      Terminal.green text, bold: true
    when :LVAR, :CALL, :QCALL
      text
    when :CDECL, :COLON2
      text.gsub(/[A-Za-z0-9_]+/) { Terminal.blue _1, bold: true, underline: true }
    else
      colorize_other text
    end
  end

  def text_range_of(node)
    from = @line_indices[node.first_lineno - 1] + node.first_column
    to = @line_indices[node.last_lineno - 1] + node.last_column
    from...to
  end

  def text_of(node)
    code[text_range_of node]
  end

  def traverse(ast, &block)
    yield ast
    ast.children.grep(RubyVM::AbstractSyntaxTree::Node) do
      traverse(_1, &block)
    end
  end

  def self.colorize(code)
    new(code).colorize
  end
end

if ARGV.include? '-n'
  a = ARGV.dup
  i = a.index '-n'
  n = a[i + 1].to_i
  a[i, 2] = []
  a.each do
    code = (File.read(_1) + $/) * n
    Colorizer.new(code).colorize
  end
elsif ARGV.empty?
  Colorizer.new(STDIN.read).colorize
else
  ARGV.each do |file|
    Colorizer.new(File.read file).colorize
  end
end