syntax_tree 5.3.0 → 6.0.0

data/lib/syntax_tree/parser.rb

@@ -256,11 +256,37 @@ module SyntaxTree
       tokens[index] if index
     end
 
+    def find_token_between(type, left, right)
+      bounds = left.location.end_char...right.location.start_char
+      index =
+        tokens.rindex do |token|
+          char = token.location.start_char
+          break if char < bounds.begin
+
+          token.is_a?(type) && bounds.cover?(char)
+        end
+
+      tokens[index] if index
+    end
+
     def find_keyword(name)
       index = tokens.rindex { |token| token.is_a?(Kw) && (token.name == name) }
       tokens[index] if index
     end
 
+    def find_keyword_between(name, left, right)
+      bounds = left.end_char...right.start_char
+      index =
+        tokens.rindex do |token|
+          char = token.location.start_char
+          break if char < bounds.begin
+
+          token.is_a?(Kw) && (token.name == name) && bounds.cover?(char)
+        end
+
+      tokens[index] if index
+    end
+
     def find_operator(name)
       index = tokens.rindex { |token| token.is_a?(Op) && (token.name == name) }
       tokens[index] if index
@@ -348,6 +374,7 @@ module SyntaxTree
 
       start_char = find_next_statement_start(lbrace.location.end_char)
       statements.bind(
+        self,
         start_char,
         start_char - line_counts[lbrace.location.start_line - 1].start,
         rbrace.location.start_char,
@@ -386,6 +413,7 @@ module SyntaxTree
 
       start_char = find_next_statement_start(lbrace.location.end_char)
       statements.bind(
+        self,
         start_char,
         start_char - line_counts[lbrace.location.start_line - 1].start,
         rbrace.location.start_char,
@@ -640,13 +668,14 @@ module SyntaxTree
         stack.pop
       end
 
-      def visit_var_ref(node)
-        pins.shift
-        node.pin(stack[-2])
+      visit_methods do
+        def visit_var_ref(node)
+          node.pin(stack[-2], pins.shift)
+        end
       end
 
       def self.visit(node, tokens)
-        start_char = node.location.start_char
+        start_char = node.start_char
         allocated = []
 
         tokens.reverse_each do |token|
@@ -670,18 +699,22 @@ module SyntaxTree
     #     (nil | Array[untyped]) posts
     #   ) -> AryPtn
     def on_aryptn(constant, requireds, rest, posts)
-      parts = [constant, *requireds, rest, *posts].compact
+      lbracket = find_token(LBracket)
+      lbracket ||= find_token(LParen) if constant
 
-      # If there aren't any parts (no constant, no positional arguments), then
-      # we're matching an empty array. In this case, we're going to look for the
-      # left and right brackets explicitly. Otherwise, we'll just use the bounds
-      # of the various parts.
-      location =
-        if parts.empty?
-          consume_token(LBracket).location.to(consume_token(RBracket).location)
-        else
-          parts[0].location.to(parts[-1].location)
-        end
+      rbracket = find_token(RBracket)
+      rbracket ||= find_token(RParen) if constant
+
+      parts = [constant, lbracket, *requireds, rest, *posts, rbracket].compact
+
+      # The location is going to be determined by the first part to the last
+      # part. This includes potential brackets.
+      location = parts[0].location.to(parts[-1].location)
+
+      # Now that we have the location calculated, we can remove the brackets
+      # from the list of tokens.
+      tokens.delete(lbracket) if lbracket
+      tokens.delete(rbracket) if rbracket
 
       # If there is a plain *, then we're going to fix up the location of it
       # here because it currently doesn't have anything to use for its precise
@@ -820,6 +853,7 @@ module SyntaxTree
           end
 
         bodystmt.bind(
+          self,
           find_next_statement_start(keyword.location.end_char),
           keyword.location.end_column,
           end_location.end_char,
@@ -871,13 +905,34 @@ module SyntaxTree
     #   on_block_var: (Params params, (nil | Array[Ident]) locals) -> BlockVar
     def on_block_var(params, locals)
       index =
-        tokens.rindex do |node|
-          node.is_a?(Op) && %w[| ||].include?(node.value) &&
-            node.location.start_char < params.location.start_char
-        end
+        tokens.rindex { |node| node.is_a?(Op) && %w[| ||].include?(node.value) }
+
+      ending = tokens.delete_at(index)
+      beginning = ending.value == "||" ? ending : consume_operator(:|)
+
+      # If there are no parameters, then we didn't have anything to base the
+      # location information of off. Now that we have an opening of the
+      # block, we can correct this.
+      if params.empty?
+        start_line = params.location.start_line
+        start_char =
+          (
+            if beginning.value == "||"
+              beginning.location.start_char
+            else
+              find_next_statement_start(beginning.location.end_char)
+            end
+          )
 
-      beginning = tokens[index]
-      ending = tokens[-1]
+        location =
+          Location.fixed(
+            line: start_line,
+            char: start_char,
+            column: start_char - line_counts[start_line - 1].start
+          )
+
+        params = params.copy(location: location)
+      end
 
       BlockVar.new(
         params: params,
@@ -905,6 +960,14 @@ module SyntaxTree
     #     (nil | Ensure) ensure_clause
     #   ) -> BodyStmt
     def on_bodystmt(statements, rescue_clause, else_clause, ensure_clause)
+      # In certain versions of Ruby, the `statements` argument can be any node
+      # in the case that we're inside of an endless method definition. In this
+      # case we'll wrap it in a Statements node to be consistent.
+      unless statements.is_a?(Statements)
+        statements =
+          Statements.new(body: [statements], location: statements.location)
+      end
+
       parts = [statements, rescue_clause, else_clause, ensure_clause].compact
 
       BodyStmt.new(
@@ -929,6 +992,7 @@ module SyntaxTree
 
       start_char = find_next_statement_start(location.end_char)
       statements.bind(
+        self,
         start_char,
         start_char - line_counts[location.start_line - 1].start,
         rbrace.location.start_char,
@@ -1036,6 +1100,7 @@ module SyntaxTree
       start_char = find_next_statement_start(location.end_char)
 
       bodystmt.bind(
+        self,
         start_char,
         start_char - line_counts[location.start_line - 1].start,
         ending.location.start_char,
@@ -1154,13 +1219,23 @@ module SyntaxTree
     end
 
     # :call-seq:
-    #   on_const_path_field: (untyped parent, Const constant) -> ConstPathField
+    #   on_const_path_field: (untyped parent, Const constant) ->
+    #     ConstPathField | Field
     def on_const_path_field(parent, constant)
-      ConstPathField.new(
-        parent: parent,
-        constant: constant,
-        location: parent.location.to(constant.location)
-      )
+      if constant.is_a?(Const)
+        ConstPathField.new(
+          parent: parent,
+          constant: constant,
+          location: parent.location.to(constant.location)
+        )
+      else
+        Field.new(
+          parent: parent,
+          operator: consume_operator(:"::"),
+          name: constant,
+          location: parent.location.to(constant.location)
+        )
+      end
     end
 
     # :call-seq:
@@ -1235,6 +1310,7 @@ module SyntaxTree
         start_char = find_next_statement_start(params.location.end_char)
 
         bodystmt.bind(
+          self,
           start_char,
           start_char - line_counts[params.location.start_line - 1].start,
           ending.location.start_char,
@@ -1323,6 +1399,7 @@ module SyntaxTree
         start_char = find_next_statement_start(params.location.end_char)
 
         bodystmt.bind(
+          self,
           start_char,
           start_char - line_counts[params.location.start_line - 1].start,
           ending.location.start_char,
@@ -1362,6 +1439,7 @@ module SyntaxTree
       start_char = find_next_statement_start(location.end_char)
 
       bodystmt.bind(
+        self,
         start_char,
         start_char - line_counts[location.start_line - 1].start,
         ending.location.start_char,
@@ -1457,6 +1535,7 @@ module SyntaxTree
 
       start_char = find_next_statement_start(keyword.location.end_char)
       statements.bind(
+        self,
         start_char,
         start_char - line_counts[keyword.location.start_line - 1].start,
         ending.location.start_char,
@@ -1482,6 +1561,7 @@ module SyntaxTree
 
       start_char = find_next_statement_start(predicate.location.end_char)
       statements.bind(
+        self,
         start_char,
         start_char - line_counts[predicate.location.start_line - 1].start,
         ending.location.start_char,
@@ -1605,6 +1685,7 @@ module SyntaxTree
       ending = find_keyword(:end)
       start_char = find_next_statement_start(keyword.location.end_char)
       statements.bind(
+        self,
         start_char,
         start_char - line_counts[keyword.location.start_line - 1].start,
         ending.location.start_char,
@@ -1679,6 +1760,22 @@ module SyntaxTree
     #     VarField right
     #   ) -> FndPtn
     def on_fndptn(constant, left, values, right)
+      # The left and right of a find pattern are always going to be splats, so
+      # we're going to consume the * operators and use their location
+      # information to extend the location of the splats.
+      right, left =
+        [right, left].map do |node|
+          operator = consume_operator(:*)
+          location =
+            if node.value
+              operator.location.to(node.location)
+            else
+              operator.location
+            end
+
+          node.copy(location: location)
+        end
+
       # The opening of this find pattern is either going to be a left bracket, a
       # right left parenthesis, or the left splat. We're going to use this to
       # determine how to find the closing of the pattern, as well as determining
@@ -1719,21 +1816,20 @@ module SyntaxTree
       in_keyword = consume_keyword(:in)
       ending = consume_keyword(:end)
 
-      # Consume the do keyword if it exists so that it doesn't get confused for
-      # some other block
-      keyword = find_keyword(:do)
-      if keyword &&
-           keyword.location.start_char > collection.location.end_char &&
-           keyword.location.end_char < ending.location.start_char
-        tokens.delete(keyword)
-      end
+      delimiter =
+        find_keyword_between(:do, collection, ending) ||
+          find_token_between(Semicolon, collection, ending)
+
+      tokens.delete(delimiter) if delimiter
 
       start_char =
-        find_next_statement_start((keyword || collection).location.end_char)
+        find_next_statement_start((delimiter || collection).location.end_char)
+
       statements.bind(
+        self,
         start_char,
         start_char -
-          line_counts[(keyword || collection).location.end_line - 1].start,
+          line_counts[(delimiter || collection).location.end_line - 1].start,
         ending.location.start_char,
         ending.location.start_column
       )
@@ -1787,7 +1883,7 @@ module SyntaxTree
           line: lineno,
           char: char_pos,
           column: current_column,
-          size: value.size + 1
+          size: value.size
         )
 
       # Here we're going to artificially create an extra node type so that if
@@ -1822,7 +1918,7 @@ module SyntaxTree
           line: lineno,
           char: char_pos,
           column: current_column,
-          size: value.size + 1
+          size: value.size
         )
 
       heredoc_end = HeredocEnd.new(value: value.chomp, location: location)
@@ -1837,9 +1933,9 @@ module SyntaxTree
             start_line: heredoc.location.start_line,
             start_char: heredoc.location.start_char,
             start_column: heredoc.location.start_column,
-            end_line: lineno,
-            end_char: char_pos,
-            end_column: current_column
+            end_line: location.end_line,
+            end_char: location.end_char,
+            end_column: location.end_column
           )
       )
     end
@@ -1847,10 +1943,42 @@ module SyntaxTree
     # :call-seq:
     #   on_hshptn: (
     #     (nil | untyped) constant,
-    #     Array[[Label, untyped]] keywords,
+    #     Array[[Label | StringContent, untyped]] keywords,
     #     (nil | VarField) keyword_rest
     #   ) -> HshPtn
     def on_hshptn(constant, keywords, keyword_rest)
+      keywords =
+        (keywords || []).map do |(label, value)|
+          if label.is_a?(Label)
+            [label, value]
+          else
+            tstring_beg_index =
+              tokens.rindex do |token|
+                token.is_a?(TStringBeg) &&
+                  token.location.start_char < label.location.start_char
+              end
+
+            tstring_beg = tokens.delete_at(tstring_beg_index)
+
+            label_end_index =
+              tokens.rindex do |token|
+                token.is_a?(LabelEnd) &&
+                  token.location.start_char == label.location.end_char
+              end
+
+            label_end = tokens.delete_at(label_end_index)
+
+            [
+              DynaSymbol.new(
+                parts: label.parts,
+                quote: label_end.value[0],
+                location: tstring_beg.location.to(label_end.location)
+              ),
+              value
+            ]
+          end
+        end
+
       if keyword_rest
         # We're doing this to delete the token from the list so that it doesn't
         # confuse future patterns by thinking they have an extra ** on the end.
@@ -1863,7 +1991,7 @@ module SyntaxTree
         keyword_rest = VarField.new(value: nil, location: token.location)
       end
 
-      parts = [constant, *keywords&.flatten(1), keyword_rest].compact
+      parts = [constant, *keywords.flatten(1), keyword_rest].compact
 
       # If there's no constant, there may be braces, so we're going to look for
       # those to get our bounds.
@@ -1880,7 +2008,7 @@ module SyntaxTree
 
       HshPtn.new(
         constant: constant,
-        keywords: keywords || [],
+        keywords: keywords,
         keyword_rest: keyword_rest,
         location: parts[0].location.to(parts[-1].location)
       )
@@ -1911,8 +2039,14 @@ module SyntaxTree
       beginning = consume_keyword(:if)
       ending = consequent || consume_keyword(:end)
 
-      start_char = find_next_statement_start(predicate.location.end_char)
+      if (keyword = find_keyword_between(:then, predicate, ending))
+        tokens.delete(keyword)
+      end
+
+      start_char =
+        find_next_statement_start((keyword || predicate).location.end_char)
       statements.bind(
+        self,
         start_char,
         start_char - line_counts[predicate.location.end_line - 1].start,
         ending.location.start_char,
@@ -1946,7 +2080,7 @@ module SyntaxTree
       IfNode.new(
         predicate: predicate,
         statements:
-          Statements.new(self, body: [statement], location: statement.location),
+          Statements.new(body: [statement], location: statement.location),
         consequent: nil,
         location: statement.location.to(predicate.location)
       )
@@ -1995,8 +2129,10 @@ module SyntaxTree
         statements_start = token
       end
 
-      start_char = find_next_statement_start(statements_start.location.end_char)
+      start_char =
+        find_next_statement_start((token || statements_start).location.end_char)
       statements.bind(
+        self,
         start_char,
         start_char -
           line_counts[statements_start.location.start_line - 1].start,
@@ -2121,12 +2257,19 @@ module SyntaxTree
             token.location.start_char > beginning.location.start_char
         end
 
+      if braces
+        opening = consume_token(TLamBeg)
+        closing = consume_token(RBrace)
+      else
+        opening = consume_keyword(:do)
+        closing = consume_keyword(:end)
+      end
+
       # We need to do some special mapping here. Since ripper doesn't support
-      # capturing lambda var until 3.2, we need to normalize all of that here.
+      # capturing lambda vars, we need to normalize all of that here.
       params =
-        case params
-        when Paren
-          # In this case we've gotten to the <3.2 parentheses wrapping a set of
+        if params.is_a?(Paren)
+          # In this case we've gotten to the parentheses wrapping a set of
           # parameters case. Here we need to manually scan for lambda locals.
           range = (params.location.start_char + 1)...params.location.end_char
           locals = lambda_locals(source[range])
@@ -2148,27 +2291,31 @@ module SyntaxTree
 
           node.comments.concat(params.comments)
           node
-        when Params
-          # In this case we've gotten to the <3.2 plain set of parameters. In
-          # this case there cannot be lambda locals, so we will wrap the
-          # parameters into a lambda var that has no locals.
+        else
+          # If there are no parameters, then we didn't have anything to base the
+          # location information of off. Now that we have an opening of the
+          # block, we can correct this.
+          if params.empty?
+            opening_location = opening.location
+            location =
+              Location.fixed(
+                line: opening_location.start_line,
+                char: opening_location.start_char,
+                column: opening_location.start_column
+              )
+
+            params = params.copy(location: location)
+          end
+
+          # In this case we've gotten to the plain set of parameters. In this
+          # case there cannot be lambda locals, so we will wrap the parameters
+          # into a lambda var that has no locals.
           LambdaVar.new(params: params, locals: [], location: params.location)
-        when LambdaVar
-          # In this case we've gotten to 3.2+ lambda var. In this case we don't
-          # need to do anything and can just the value as given.
-          params
         end
 
-      if braces
-        opening = consume_token(TLamBeg)
-        closing = consume_token(RBrace)
-      else
-        opening = consume_keyword(:do)
-        closing = consume_keyword(:end)
-      end
-
       start_char = find_next_statement_start(opening.location.end_char)
       statements.bind(
+        self,
         start_char,
         start_char - line_counts[opening.location.end_line - 1].start,
         closing.location.start_char,
@@ -2353,23 +2500,30 @@ module SyntaxTree
 
     # :call-seq:
     #   on_method_add_block: (
-    #     (Call | Command | CommandCall) call,
+    #     (Break | Call | Command | CommandCall, Next) call,
     #     Block block
-    #   ) -> MethodAddBlock
+    #   ) -> Break | MethodAddBlock
     def on_method_add_block(call, block)
       location = call.location.to(block.location)
 
       case call
+      when Break, Next, ReturnNode
+        parts = call.arguments.parts
+
+        node = parts.pop
+        copied =
+          node.copy(block: block, location: node.location.to(block.location))
+
+        copied.comments.concat(call.comments)
+        parts << copied
+
+        call.copy(location: location)
       when Command, CommandCall
         node = call.copy(block: block, location: location)
         node.comments.concat(call.comments)
         node
       else
-        MethodAddBlock.new(
-          call: call,
-          block: block,
-          location: call.location.to(block.location)
-        )
+        MethodAddBlock.new(call: call, block: block, location: location)
       end
     end
 
@@ -2446,6 +2600,7 @@ module SyntaxTree
       start_char = find_next_statement_start(constant.location.end_char)
 
       bodystmt.bind(
+        self,
         start_char,
         start_char - line_counts[constant.location.start_line - 1].start,
         ending.location.start_char,
@@ -2592,19 +2747,40 @@ module SyntaxTree
       # have a `nil` for the value instead of a `false`.
       keywords&.map! { |(key, value)| [key, value || nil] }
 
-      parts = [
-        *requireds,
-        *optionals&.flatten(1),
-        rest,
-        *posts,
-        *keywords&.flatten(1),
-        (keyword_rest if keyword_rest != :nil),
-        (block if block != :&)
-      ].compact
+      # Here we're going to build up a list of all of the params so that we can
+      # determine our location information.
+      parts = []
+
+      requireds&.each { |required| parts << required.location }
+      optionals&.each do |(key, value)|
+        parts << key.location
+        parts << value.location if value
+      end
+
+      parts << rest.location if rest
+      posts&.each { |post| parts << post.location }
+
+      keywords&.each do |(key, value)|
+        parts << key.location
+        parts << value.location if value
+      end
+
+      if keyword_rest == :nil
+        # When we get a :nil here, it means that we have **nil syntax, which
+        # means this set of parameters accepts no more keyword arguments. In
+        # this case we need to go and find the location of these two tokens.
+        operator = consume_operator(:**)
+        parts << operator.location.to(consume_keyword(:nil).location)
+      elsif keyword_rest
+        parts << keyword_rest.location
+      end
+
+      parts << block.location if block && block != :&
+      parts = parts.compact
 
       location =
         if parts.any?
-          parts[0].location.to(parts[-1].location)
+          parts[0].to(parts[-1])
         else
           Location.fixed(line: lineno, char: char_pos, column: current_column)
         end
@@ -2701,7 +2877,7 @@ module SyntaxTree
         )
 
       statements.body << @__end__ if @__end__
-      statements.bind(0, 0, source.length, last_column)
+      statements.bind(self, 0, 0, source.length, last_column)
 
       program = Program.new(statements: statements, location: location)
       attach_comments(program, @comments)
@@ -3033,8 +3209,9 @@ module SyntaxTree
       exceptions = exceptions[0] if exceptions.is_a?(Array)
 
       last_node = variable || exceptions || keyword
-      start_char = find_next_statement_start(last_node.location.end_char)
+      start_char = find_next_statement_start(last_node.end_char)
       statements.bind(
+        self,
         start_char,
         start_char - line_counts[last_node.location.start_line - 1].start,
         char_pos,
@@ -3055,7 +3232,7 @@ module SyntaxTree
                 start_char: keyword.location.end_char + 1,
                 start_column: keyword.location.end_column + 1,
                 end_line: last_node.location.end_line,
-                end_char: last_node.location.end_char,
+                end_char: last_node.end_char,
                 end_column: last_node.location.end_column
               )
           )
@@ -3153,6 +3330,7 @@ module SyntaxTree
       start_char = find_next_statement_start(target.location.end_char)
 
       bodystmt.bind(
+        self,
         start_char,
         start_char - line_counts[target.location.start_line - 1].start,
         ending.location.start_char,
@@ -3166,9 +3344,29 @@ module SyntaxTree
       )
     end
 
-    # def on_semicolon(value)
-    #   value
-    # end
+    # Semicolons are tokens that get added to the token list but never get
+    # attached to the AST. Because of this they only need to track their
+    # associated location so they can be used for computing bounds.
+    class Semicolon
+      attr_reader :location
+
+      def initialize(location)
+        @location = location
+      end
+    end
+
+    # :call-seq:
+    #   on_semicolon: (String value) -> Semicolon
+    def on_semicolon(value)
+      tokens << Semicolon.new(
+        Location.token(
+          line: lineno,
+          char: char_pos,
+          column: current_column,
+          size: value.size
+        )
+      )
+    end
 
     # def on_sp(value)
     #   value
@@ -3186,18 +3384,13 @@ module SyntaxTree
           statements.location.to(statement.location)
         end
 
-      Statements.new(
-        self,
-        body: statements.body << statement,
-        location: location
-      )
+      Statements.new(body: statements.body << statement, location: location)
     end
 
     # :call-seq:
     #   on_stmts_new: () -> Statements
     def on_stmts_new
       Statements.new(
-        self,
         body: [],
         location:
           Location.fixed(line: lineno, char: char_pos, column: current_column)
@@ -3262,6 +3455,7 @@ module SyntaxTree
       embexpr_end = consume_token(EmbExprEnd)
 
       statements.bind(
+        self,
         embexpr_beg.location.end_char,
         embexpr_beg.location.end_column,
         embexpr_end.location.start_char,
@@ -3605,8 +3799,14 @@ module SyntaxTree
       beginning = consume_keyword(:unless)
       ending = consequent || consume_keyword(:end)
 
-      start_char = find_next_statement_start(predicate.location.end_char)
+      if (keyword = find_keyword_between(:then, predicate, ending))
+        tokens.delete(keyword)
+      end
+
+      start_char =
+        find_next_statement_start((keyword || predicate).location.end_char)
       statements.bind(
+        self,
         start_char,
         start_char - line_counts[predicate.location.end_line - 1].start,
         ending.location.start_char,
@@ -3629,7 +3829,7 @@ module SyntaxTree
       UnlessNode.new(
         predicate: predicate,
         statements:
-          Statements.new(self, body: [statement], location: statement.location),
+          Statements.new(body: [statement], location: statement.location),
         consequent: nil,
         location: statement.location.to(predicate.location)
       )
@@ -3641,17 +3841,18 @@ module SyntaxTree
       beginning = consume_keyword(:until)
       ending = consume_keyword(:end)
 
-      # Consume the do keyword if it exists so that it doesn't get confused for
-      # some other block
-      keyword = find_keyword(:do)
-      if keyword && keyword.location.start_char > predicate.location.end_char &&
-           keyword.location.end_char < ending.location.start_char
-        tokens.delete(keyword)
-      end
+      delimiter =
+        find_keyword_between(:do, predicate, statements) ||
+          find_token_between(Semicolon, predicate, statements)
+
+      tokens.delete(delimiter) if delimiter
 
       # Update the Statements location information
-      start_char = find_next_statement_start(predicate.location.end_char)
+      start_char =
+        find_next_statement_start((delimiter || predicate).location.end_char)
+
       statements.bind(
+        self,
         start_char,
         start_char - line_counts[predicate.location.end_line - 1].start,
         ending.location.start_char,
@@ -3673,7 +3874,7 @@ module SyntaxTree
       UntilNode.new(
         predicate: predicate,
         statements:
-          Statements.new(self, body: [statement], location: statement.location),
+          Statements.new(body: [statement], location: statement.location),
         location: statement.location.to(predicate.location)
       )
     end
@@ -3744,9 +3945,11 @@ module SyntaxTree
         statements_start = token
       end
 
-      start_char = find_next_statement_start(statements_start.location.end_char)
+      start_char =
+        find_next_statement_start((token || statements_start).location.end_char)
 
       statements.bind(
+        self,
         start_char,
         start_char -
           line_counts[statements_start.location.start_line - 1].start,
@@ -3768,17 +3971,18 @@ module SyntaxTree
       beginning = consume_keyword(:while)
       ending = consume_keyword(:end)
 
-      # Consume the do keyword if it exists so that it doesn't get confused for
-      # some other block
-      keyword = find_keyword(:do)
-      if keyword && keyword.location.start_char > predicate.location.end_char &&
-           keyword.location.end_char < ending.location.start_char
-        tokens.delete(keyword)
-      end
+      delimiter =
+        find_keyword_between(:do, predicate, statements) ||
+          find_token_between(Semicolon, predicate, statements)
+
+      tokens.delete(delimiter) if delimiter
 
       # Update the Statements location information
-      start_char = find_next_statement_start(predicate.location.end_char)
+      start_char =
+        find_next_statement_start((delimiter || predicate).location.end_char)
+
       statements.bind(
+        self,
         start_char,
         start_char - line_counts[predicate.location.end_line - 1].start,
         ending.location.start_char,
@@ -3800,7 +4004,7 @@ module SyntaxTree
       WhileNode.new(
         predicate: predicate,
         statements:
-          Statements.new(self, body: [statement], location: statement.location),
+          Statements.new(body: [statement], location: statement.location),
         location: statement.location.to(predicate.location)
       )
     end