textbringer 24 → 26

data/lib/textbringer/modes/ruby_mode.rb

@@ -1,4 +1,5 @@
-require "ripper"
+require "set"
+require "prism"
 
 module Textbringer
   CONFIG[:ruby_indent_level] = 2
@@ -10,81 +11,6 @@ module Textbringer
             (?:Gem|Rake|Cap|Thor|Vagrant|Guard|Pod)file)\z/ix
     self.interpreter_name_pattern = /ruby/i
 
-    define_syntax :comment, /
-      (?: \#.*(?:\\\n.*)*(?:\z|(?<!\\)\n) ) |
-      (?: ^=begin (?:.|\n)* (?> ^=end \b ) )
-    /x
-
-    define_syntax :keyword, /
-      (?<![$@.]) \b (?: (?:
-        class | module | def | undef | begin | rescue | ensure | end |
-        if | unless | then | elsif | else | case | when | while | until |
-        for | break | next | redo | retry | in | do | return | yield |
-        super | self | nil | true | false | and | or | not | alias
-      ) \b (?![!?]) | defined\? )
-    /x
-
-    define_syntax :string, /
-      (?: (?<! [a-zA-Z] ) \?
-              (:?
-                [^\\\s] |
-                \\ [0-7]{1,3} |
-                \\x [0-9a-fA-F]{2} |
-                \\u [0-9a-fA-F]{4} |
-                \\u \{ [0-9a-fA-F]+ \} |
-                \\C - . |
-                \\M - . |
-                \\ .
-              )
-      ) |
-      (?: %[qQrwWsiIx]?\{ (?: [^\\}] | \\ .  )* \} ) |
-      (?: %[qQrwWsiIx]?\( (?: [^\\)] | \\ .  )* \) ) |
-      (?: %[qQrwWsiIx]?\[ (?: [^\\\]] | \\ .  )* \] ) |
-      (?: %[qQrwWsiIx]?< (?: [^\\>] | \\ .  )* > ) |
-      (?:
-         %[qQrwWsiIx]?
-             (?<string_delimiter>[^{(\[<a-zA-Z0-9\s\u{0100}-\u{10ffff}])
-             (?: (?! \k<string_delimiter> ) [^\\] | \\ .  )*
-             \k<string_delimiter>
-      ) |
-      (?:
-        (?<! \$ )
-            " (?: [^\\"] | \\ .  )* "
-      ) |
-      (?:
-        (?<! \$ )
-            ' (?: [^\\'] | \\ .  )* '
-      ) |
-      (?:
-         (?<! [$.] | def | def \s )
-             ` (?: [^\\`] | \\ .  )* `
-      ) |
-      (?:
-        (?<=
-          ^ |
-          \b and | \b or | \b while | \b until | \b unless | \b if |
-          \b elsif | \b when | \b not | \b then | \b else |
-          [;~=!|&(,\[<>?:*+-]
-        ) \s*
-        \/ (?: [^\\\/] | \\ .  )* \/[iomxneus]*
-      ) |
-      (?:
-        (?<! class | class \s | [\]})"'.] | :: | \w )
-            <<[\-~]?(?<heredoc_quote>['"`]?)
-            (?<heredoc_terminator>
-              (?> [_a-zA-Z\u{0100}-\u{10ffff}]
-                  [_a-zA-Z0-9\u{0100}-\u{10ffff}]* )
-            )
-            \k<heredoc_quote>
-            (?> (?:.|\n)*? ^ [\ \t]* \k<heredoc_terminator> $ )
-      ) |
-      (?:
-        (?<! : ) :
-            [_a-zA-Z\u{0100}-\u{10ffff}]
-            [_a-zA-Z0-9\u{0100}-\u{10ffff}]*
-      )
-    /x
-
     def comment_start
       "#"
     end
@@ -93,17 +19,23 @@ module Textbringer
       super(buffer)
       @buffer[:indent_level] = CONFIG[:ruby_indent_level]
       @buffer[:indent_tabs_mode] = CONFIG[:ruby_indent_tabs_mode]
+      @prism_version = nil
+      @prism_tokens = nil
+      @prism_ast = nil
+      @prism_method_name_locs = nil
+      @literal_levels = nil
+      @literal_levels_version = nil
     end
 
     def forward_definition(n = number_prefix_arg || 1)
-      tokens = Ripper.lex(@buffer.to_s)
+      tokens = filter_prism_tokens_line_and_type
       @buffer.forward_line
       n.times do |i|
-        tokens = tokens.drop_while { |(l, _), e, t|
+        tokens = tokens.drop_while { |l, type|
           l < @buffer.current_line ||
-            e != :on_kw || /\A(?:class|module|def)\z/ !~ t
+            !DEFINITION_KEYWORDS.include?(type)
         }
-        (line,), = tokens.first
+        line, = tokens.first
         if line.nil?
           @buffer.end_of_buffer
           break
@@ -117,14 +49,14 @@ module Textbringer
     end
 
     def backward_definition(n = number_prefix_arg || 1)
-      tokens = Ripper.lex(@buffer.to_s).reverse
+      tokens = filter_prism_tokens_line_and_type.reverse
       @buffer.beginning_of_line
       n.times do |i|
-        tokens = tokens.drop_while { |(l, _), e, t|
+        tokens = tokens.drop_while { |l, type|
           l >= @buffer.current_line ||
-            e != :on_kw || /\A(?:class|module|def)\z/ !~ t
+            !DEFINITION_KEYWORDS.include?(type)
         }
-        (line,), = tokens.first
+        line, = tokens.first
         if line.nil?
           @buffer.beginning_of_buffer
           break
@@ -175,8 +107,167 @@ module Textbringer
       end
     end
 
+    def highlight(ctx)
+      ensure_prism_tokens
+      return unless @prism_tokens
+      ensure_method_name_locs
+      base_pos = ctx.buffer.point_min
+      hl_start = ctx.highlight_start
+      hl_end = ctx.highlight_end
+      in_symbol = false
+      after_class_or_module = false
+      @prism_tokens.each do |token_info|
+        token = token_info[0]
+        type = token.type
+        offset = token.location.start_offset
+        length = token.location.length
+        pos = base_pos + offset
+        pos_end = pos + length
+        break if pos >= hl_end
+        if pos_end > hl_start
+          face_name = PRISM_TOKEN_FACES[type]
+          if in_symbol
+            face_name = :string if face_name.nil? || face_name == :constant ||
+              face_name == :keyword || face_name == :operator
+          elsif @prism_method_name_locs.key?(offset)
+            face_name = :function_name
+          elsif face_name == :constant &&
+              (after_class_or_module || token.location.slice.match?(/\p{Lower}/))
+            face_name = :type
+          end
+          if face_name && (face = Face[face_name])
+            ctx.highlight(pos, pos_end, face)
+          end
+        end
+        in_symbol = type == :SYMBOL_BEGIN
+        after_class_or_module = (type == :KEYWORD_CLASS || type == :KEYWORD_MODULE) ||
+          (after_class_or_module && !(type == :NEWLINE || type == :SEMICOLON))
+      end
+    end
+
     private
 
+    PRISM_TOKEN_FACES = {
+      # Keywords
+      KEYWORD_ALIAS: :keyword, KEYWORD_AND: :keyword, KEYWORD_BEGIN: :keyword,
+      KEYWORD_BEGIN_UPCASE: :keyword, KEYWORD_BREAK: :keyword,
+      KEYWORD_CASE: :keyword, KEYWORD_CLASS: :keyword, KEYWORD_DEF: :keyword,
+      KEYWORD_DEFINED: :keyword, KEYWORD_DO: :keyword,
+      KEYWORD_DO_LOOP: :keyword, KEYWORD_ELSE: :keyword,
+      KEYWORD_ELSIF: :keyword, KEYWORD_END: :keyword,
+      KEYWORD_END_UPCASE: :keyword, KEYWORD_ENSURE: :keyword,
+      KEYWORD_FALSE: :builtin, KEYWORD_FOR: :keyword, KEYWORD_IF: :keyword,
+      KEYWORD_IF_MODIFIER: :keyword, KEYWORD_IN: :keyword,
+      KEYWORD_MODULE: :keyword, KEYWORD_NEXT: :keyword,
+      KEYWORD_NIL: :builtin,
+      KEYWORD_NOT: :keyword, KEYWORD_OR: :keyword, KEYWORD_REDO: :keyword,
+      KEYWORD_RESCUE: :keyword, KEYWORD_RESCUE_MODIFIER: :keyword,
+      KEYWORD_RETRY: :keyword, KEYWORD_RETURN: :keyword,
+      KEYWORD_SELF: :builtin, KEYWORD_SUPER: :builtin,
+      KEYWORD_THEN: :keyword, KEYWORD_TRUE: :builtin,
+      KEYWORD_UNDEF: :keyword,
+      KEYWORD_UNLESS: :keyword, KEYWORD_UNLESS_MODIFIER: :keyword,
+      KEYWORD_UNTIL: :keyword, KEYWORD_UNTIL_MODIFIER: :keyword,
+      KEYWORD_WHEN: :keyword, KEYWORD_WHILE: :keyword,
+      KEYWORD_WHILE_MODIFIER: :keyword, KEYWORD_YIELD: :keyword,
+      KEYWORD___FILE__: :builtin, KEYWORD___LINE__: :builtin,
+      KEYWORD___ENCODING__: :builtin,
+
+      # Comments
+      COMMENT: :comment, EMBDOC_BEGIN: :comment, EMBDOC_LINE: :comment,
+      EMBDOC_END: :comment,
+
+      # Strings and string-like
+      STRING_BEGIN: :string, STRING_CONTENT: :string, STRING_END: :string,
+      SYMBOL_BEGIN: :string, REGEXP_BEGIN: :string, REGEXP_END: :string,
+      PERCENT_LOWER_I: :string, PERCENT_UPPER_I: :string,
+      PERCENT_LOWER_W: :string, PERCENT_UPPER_W: :string,
+      PERCENT_LOWER_X: :string,
+      HEREDOC_START: :string, HEREDOC_END: :string,
+      LABEL: :property,
+
+      # Numbers
+      INTEGER: :number, FLOAT: :number,
+      INTEGER_RATIONAL: :number, FLOAT_RATIONAL: :number,
+      INTEGER_IMAGINARY: :number, FLOAT_IMAGINARY: :number,
+      INTEGER_RATIONAL_IMAGINARY: :number, FLOAT_RATIONAL_IMAGINARY: :number,
+
+      # Constants
+      CONSTANT: :constant,
+
+      # Variables
+      INSTANCE_VARIABLE: :variable, CLASS_VARIABLE: :variable,
+      GLOBAL_VARIABLE: :variable,
+
+      # Operators
+      PLUS: :operator, MINUS: :operator, STAR: :operator, SLASH: :operator,
+      PERCENT: :operator, STAR_STAR: :operator,
+      EQUAL: :operator, EQUAL_EQUAL: :operator, BANG_EQUAL: :operator,
+      LESS: :operator, GREATER: :operator,
+      LESS_EQUAL: :operator, GREATER_EQUAL: :operator,
+      LESS_EQUAL_GREATER: :operator, EQUAL_EQUAL_EQUAL: :operator,
+      EQUAL_TILDE: :operator, BANG_TILDE: :operator,
+      AMPERSAND_AMPERSAND: :operator, PIPE_PIPE: :operator,
+      BANG: :operator, TILDE: :operator,
+      LESS_LESS: :operator, GREATER_GREATER: :operator,
+      AMPERSAND: :operator, PIPE: :operator, CARET: :operator,
+      PLUS_EQUAL: :operator, MINUS_EQUAL: :operator,
+      STAR_EQUAL: :operator, SLASH_EQUAL: :operator,
+      PERCENT_EQUAL: :operator, STAR_STAR_EQUAL: :operator,
+      AMPERSAND_EQUAL: :operator, PIPE_EQUAL: :operator,
+      CARET_EQUAL: :operator,
+      AMPERSAND_AMPERSAND_EQUAL: :operator, PIPE_PIPE_EQUAL: :operator,
+      LESS_LESS_EQUAL: :operator, GREATER_GREATER_EQUAL: :operator,
+      DOT_DOT: :operator, DOT_DOT_DOT: :operator,
+      EQUAL_GREATER: :operator, UMINUS: :operator, UPLUS: :operator,
+      USTAR: :operator, USTAR_STAR: :operator, UAMPERSAND: :operator,
+
+      # Punctuation
+      DOT: :punctuation, COLON_COLON: :punctuation,
+      SEMICOLON: :punctuation, COMMA: :punctuation,
+      PARENTHESIS_LEFT: :punctuation, PARENTHESIS_RIGHT: :punctuation,
+      BRACKET_LEFT: :punctuation, BRACKET_LEFT_ARRAY: :punctuation,
+      BRACKET_RIGHT: :punctuation,
+      BRACE_LEFT: :punctuation, BRACE_RIGHT: :punctuation,
+      QUESTION_MARK: :punctuation, COLON: :punctuation,
+      LAMBDA_BEGIN: :punctuation,
+
+      # Method names (e.g. block_given?, is_a?)
+      METHOD_NAME: :function_name,
+    }.freeze
+
+    DEFINITION_KEYWORDS = %i[KEYWORD_CLASS KEYWORD_MODULE KEYWORD_DEF].to_set
+
+    LITERAL_BEGIN_TYPES = %i[STRING_BEGIN HEREDOC_START REGEXP_BEGIN EMBDOC_BEGIN].to_set
+    LITERAL_END_TYPES = %i[STRING_END HEREDOC_END REGEXP_END EMBDOC_END].to_set
+
+    CONTINUATION_OPERATOR_TYPES = %i[
+      PLUS MINUS STAR SLASH PERCENT STAR_STAR
+      EQUAL EQUAL_EQUAL BANG_EQUAL
+      LESS GREATER LESS_EQUAL GREATER_EQUAL LESS_EQUAL_GREATER
+      EQUAL_TILDE BANG_TILDE
+      AMPERSAND_AMPERSAND PIPE_PIPE
+      BANG TILDE
+      LESS_LESS GREATER_GREATER
+      AMPERSAND CARET
+      PLUS_EQUAL MINUS_EQUAL STAR_EQUAL SLASH_EQUAL PERCENT_EQUAL
+      STAR_STAR_EQUAL AMPERSAND_EQUAL PIPE_EQUAL CARET_EQUAL
+      AMPERSAND_AMPERSAND_EQUAL PIPE_PIPE_EQUAL
+      LESS_LESS_EQUAL GREATER_GREATER_EQUAL
+      EQUAL_GREATER
+    ].to_set
+
+    OPERATORS = %i(!= !~ =~ == === <=> > >= < <= & | ^ >> << - + % / * ** -@ +@ ~ ! [] []=)
+
+    BLOCK_END = {
+      EMBEXPR_BEGIN: :EMBEXPR_END,
+      BRACE_LEFT: :BRACE_RIGHT,
+      PARENTHESIS_LEFT: :PARENTHESIS_RIGHT,
+      BRACKET_LEFT: :BRACKET_RIGHT,
+      BRACKET_LEFT_ARRAY: :BRACKET_RIGHT,
+      LAMBDA_BEGIN: :BRACE_RIGHT,
+    }
+
     INDENT_BEG_RE = /^([ \t]*)(class|module|def|if|unless|case|while|until|for|begin)\b/
 
     def space_width(s)
@@ -187,8 +278,7 @@ module Textbringer
       loop do
         @buffer.re_search_backward(INDENT_BEG_RE)
         space = @buffer.match_string(1)
-        s = @buffer.substring(@buffer.point_min, @buffer.point)
-        if PartialLiteralAnalyzer.in_literal?(s)
+        if in_literal?(@buffer.point)
           next
         end
         return space_width(space)
@@ -199,19 +289,12 @@ module Textbringer
     end
 
     def lex(source)
-      line_count = source.count("\n")
-      s = source
-      lineno = 1
-      tokens = []
-      loop do
-        lexer = Ripper::Lexer.new(s, "-", lineno)
-        tokens.concat(lexer.lex)
-        last_line = tokens.dig(-1, 0, 0)
-        return tokens if last_line.nil? || last_line >= line_count
-        s = source.sub(/(.*\n?){#{last_line}}/, "")
-        return tokens if last_line + 1 <= lineno
-        lineno = last_line + 1
-      end
+      Prism.lex(source).value.filter_map { |token, _state|
+        type = token.type
+        next if type == :EOF
+        loc = token.location
+        [[loc.start_line, loc.start_column], type, token.value]
+      }
     end
 
     def calculate_indentation
@@ -227,19 +310,20 @@ module Textbringer
         start_line = @buffer.current_line
         tokens = lex(@buffer.substring(start_pos, bol_pos))
         _, event, text = tokens.last
-        if event == :on_nl
+        if event == :NEWLINE || event == :IGNORED_NEWLINE
           _, event, text = tokens[-2]
         end
-        if event == :on_tstring_beg ||
-            event == :on_heredoc_beg ||
-            event == :on_regexp_beg ||
-            (event == :on_regexp_end && text.size > 1) ||
-            event == :on_tstring_content
+        if event == :STRING_BEGIN ||
+            event == :HEREDOC_START ||
+            (event == :HEREDOC_END && text.empty?) ||
+            event == :REGEXP_BEGIN ||
+            event == :STRING_CONTENT ||
+            event == :HEREDOC_CONTENT
           return nil
         end
         i, extra_end_count = find_nearest_beginning_token(tokens)
         (line, column), event, = i ? tokens[i] : nil
-        if event == :on_lparen && tokens.dig(i + 1, 1) != :on_ignored_nl
+        if event == :PARENTHESIS_LEFT && tokens.dig(i + 1, 1) != :IGNORED_NEWLINE
           return column + 1
         end
         if line
@@ -268,101 +352,87 @@ module Textbringer
         if @buffer.looking_at?(/[ \t]*([}\])]|(end|else|elsif|when|in|rescue|ensure)\b)/)
           indentation -= @buffer[:indent_level]
         end
-        _, last_event, last_text = tokens.reverse_each.find { |_, e, _|
-          e != :on_sp && e != :on_nl && e != :on_ignored_nl
+        _, last_event, = tokens.reverse_each.find { |_, e, _|
+          e != :NEWLINE && e != :IGNORED_NEWLINE
         }
         if start_with_period ||
-            (last_event == :on_op && last_text != "|") ||
-            (last_event == :on_kw && /\A(and|or)\z/.match?(last_text)) ||
-            last_event == :on_period ||
-            (last_event == :on_comma && event != :on_lbrace &&
-             event != :on_lparen && event != :on_lbracket) ||
-            last_event == :on_label
+            CONTINUATION_OPERATOR_TYPES.include?(last_event) ||
+            last_event == :KEYWORD_AND || last_event == :KEYWORD_OR ||
+            last_event == :DOT ||
+            (last_event == :COMMA && event != :BRACE_LEFT &&
+             event != :PARENTHESIS_LEFT && event != :BRACKET_LEFT &&
+             event != :BRACKET_LEFT_ARRAY) ||
+            last_event == :LABEL
           indentation += @buffer[:indent_level]
         end
         indentation
       end
     end
 
-    BLOCK_END = {
-      '#{' => "}",
-      "{" => "}",
-      "(" => ")",
-      "[" => "]"
-    }
-
     def find_nearest_beginning_token(tokens)
       stack = []
       (tokens.size - 1).downto(0) do |i|
         (line, ), event, text = tokens[i]
         case event
-        when :on_kw
-          _, prev_event, _ = tokens[i - 1]
-          next if prev_event == :on_symbeg
-          case text
-          when "class", "module", "def", "if", "unless", "case",
-            "do", "for", "while", "until", "begin"
-            if /\A(if|unless|while|until)\z/.match?(text) &&
-                modifier?(tokens, i)
-              next
-            end
-            if text == "def" && endless_method_def?(tokens, i)
-              next
-            end
-            if stack.empty?
-              return i
-            end
-            if stack.last != "end"
-              raise EditorError, "#{@buffer.name}:#{line}: Unmatched #{text}"
-            end
-            stack.pop
-          when "end"
-            stack.push(text)
+        when :KEYWORD_CLASS, :KEYWORD_MODULE, :KEYWORD_DEF,
+          :KEYWORD_IF, :KEYWORD_UNLESS, :KEYWORD_CASE,
+          :KEYWORD_DO, :KEYWORD_DO_LOOP, :KEYWORD_FOR,
+          :KEYWORD_WHILE, :KEYWORD_UNTIL, :KEYWORD_BEGIN
+          if i > 0
+            _, prev_event, _ = tokens[i - 1]
+            next if prev_event == :SYMBOL_BEGIN
+          end
+          if event == :KEYWORD_DEF && endless_method_def?(tokens, i)
+            next
+          end
+          if stack.empty?
+            return i
+          end
+          if stack.last != :KEYWORD_END
+            raise EditorError, "#{@buffer.name}:#{line}: Unmatched #{text}"
           end
-        when :on_rbrace, :on_rparen, :on_rbracket, :on_embexpr_end
-          stack.push(text)
-        when :on_lbrace, :on_lparen, :on_lbracket, :on_tlambeg, :on_embexpr_beg
+          stack.pop
+        when :KEYWORD_END
+          if i > 0
+            _, prev_event, _ = tokens[i - 1]
+            next if prev_event == :SYMBOL_BEGIN
+          end
+          stack.push(:KEYWORD_END)
+        when :BRACE_RIGHT, :PARENTHESIS_RIGHT, :BRACKET_RIGHT, :EMBEXPR_END
+          stack.push(event)
+        when :BRACE_LEFT, :PARENTHESIS_LEFT, :BRACKET_LEFT,
+          :BRACKET_LEFT_ARRAY, :LAMBDA_BEGIN, :EMBEXPR_BEGIN
           if stack.empty?
             return i
           end
-          if stack.last != BLOCK_END[text]
+          if stack.last != BLOCK_END[event]
             raise EditorError, "#{@buffer.name}:#{line}: Unmatched #{text}"
           end
           stack.pop
         end
       end
-      return nil, stack.grep_v(/[)\]]/).size
-    end
-
-    def modifier?(tokens, i)
-      (line,), = tokens[i]
-      ts = tokens[0...i].reverse_each.take_while { |(l,_),| l == line }
-      t = ts.find { |_, e| e != :on_sp }
-      t && !(t[1] == :on_op && t[2] == "=")
+      return nil, stack.count { |t| t != :PARENTHESIS_RIGHT && t != :BRACKET_RIGHT }
     end
 
     def endless_method_def?(tokens, i)
       ts = tokens.drop(i + 1)
-      ts.shift while ts[0][1] == :on_sp
       _, event = ts.shift
-      return false if event != :on_ident
-      ts.shift while ts[0][1] == :on_sp
-      if ts[0][1] == :on_lparen
+      return false if event != :IDENTIFIER && event != :METHOD_NAME
+      if ts[0][1] == :PARENTHESIS_LEFT
         ts.shift
         count = 1
         while count > 0
           _, event = ts.shift
           return false if event.nil?
           case event
-          when :on_lparen
-            count +=1
-          when :on_rparen
-            count -=1
+          when :PARENTHESIS_LEFT
+            count += 1
+          when :PARENTHESIS_RIGHT
+            count -= 1
           end
         end
-        ts.shift while ts[0][1] == :on_sp
       end
-      ts[0][1] == :on_op && ts[0][2] == "="
+      ts[0][1] == :EQUAL
     rescue NoMethodError # no token
       return false
     end
@@ -395,31 +465,86 @@ module Textbringer
       nil
     end
 
-    class PartialLiteralAnalyzer < Ripper
-      def self.in_literal?(src)
-        new(src).in_literal?
-      end
+    def in_literal?(byte_offset)
+      ensure_prism_tokens
+      ensure_literal_levels
+      return false if @literal_levels.empty?
+      i = @literal_levels.bsearch_index { |offset, _| offset > byte_offset }
+      i = i ? i - 1 : @literal_levels.size - 1
+      return false if i < 0
+      @literal_levels[i][1] > 0
+    end
 
-      def in_literal?
-        @literal_level = 0
-        parse
-        @literal_level > 0
+    def filter_prism_tokens_line_and_type
+      ensure_prism_tokens
+      @prism_tokens.filter_map { |token, _state|
+        type = token.type
+        next if type == :EOF
+        [token.location.start_line, type]
+      }
+    end
+
+    def ensure_prism_tokens
+      return if @prism_version == @buffer.version
+      source = @buffer.to_s
+      if source.valid_encoding?
+        result = Prism.parse_lex(source)
+        @prism_ast, @prism_tokens = result.value
+      else
+        @prism_ast = nil
+        @prism_tokens = []
       end
+      @prism_method_name_locs = nil
+      @prism_version = @buffer.version
+      @literal_levels_version = nil
+    end
 
-      private
+    def ensure_method_name_locs
+      return if @prism_method_name_locs
+      @prism_method_name_locs = {}
+      return unless @prism_ast
+      collect_method_name_locs(@prism_ast)
+    end
 
-      %w(embdoc heredoc tstring regexp
-      symbols qsymbols words qwords).each do |name|
-        define_method("on_#{name}_beg") do |token|
-          @literal_level += 1
-        end
+    def collect_method_name_locs(node)
+      if node.type == :def_node
+        @prism_method_name_locs[node.name_loc.start_offset] = true
+      elsif node.type == :alias_method_node
+        add_alias_method_name_locs(node.new_name)
+        add_alias_method_name_locs(node.old_name)
+      elsif (node.type == :call_node &&
+             !(node.call_operator_loc.nil? && OPERATORS.include?(node.name)) &&          # exclude operators
+             !((node.call_operator_loc.nil? || node.call_operator_loc.slice == "::") &&  # exclude constants
+               /\A\p{Upper}/.match?(node.name))) ||
+          node.type == :call_operator_write_node ||
+          node.type == :call_and_write_node ||
+          node.type == :call_or_write_node
+        @prism_method_name_locs[node.message_loc.start_offset] = true
+      end
+      node.compact_child_nodes.each { |child| collect_method_name_locs(child) }
+    end
+
+    def add_alias_method_name_locs(node)
+      if node.type == :symbol_node && node.opening_loc.nil?
+        @prism_method_name_locs[node.value_loc.start_offset] = true
       end
+    end
 
-      %w(embdoc heredoc tstring regexp).each do |name|
-        define_method("on_#{name}_end") do |token|
-          @literal_level -= 1
+    def ensure_literal_levels
+      return if @literal_levels_version == @prism_version
+      level = 0
+      @literal_levels = []
+      @prism_tokens&.each do |token, _state|
+        type = token.type
+        if LITERAL_BEGIN_TYPES.include?(type)
+          level += 1
+        elsif LITERAL_END_TYPES.include?(type)
+          next if type == :HEREDOC_END && token.value.empty?
+          level -= 1
         end
+        @literal_levels << [token.location.start_offset, level]
       end
+      @literal_levels_version = @prism_version
     end
   end
 end