Merge pull request #84 from ruby-syntax-tree/lambda-locals

Handle lambda-local variables

Author: kddnewton

lib/syntax_tree/node.rb

@@ -5900,7 +5900,7 @@ def deconstruct_keys(_keys)
   #     ->(value) { value * 2 }
   #
   class Lambda < Node
-    # [Params | Paren] the parameter declaration for this lambda
+    # [LambdaVar | Paren] the parameter declaration for this lambda
     attr_reader :params
 
     # [BodyStmt | Statements] the expressions to be executed in this lambda
@@ -5955,24 +5955,100 @@ def format(q)
                 node.is_a?(Command) || node.is_a?(CommandCall)
               end
 
-            q.text(force_parens ? "{" : "do")
-            q.indent do
+            if force_parens
+              q.text("{")
+
+              unless statements.empty?
+                q.indent do
+                  q.breakable
+                  q.format(statements)
+                end
+                q.breakable
+              end
+
+              q.text("}")
+            else
+              q.text("do")
+
+              unless statements.empty?
+                q.indent do
+                  q.breakable
+                  q.format(statements)
+                end
+              end
+
               q.breakable
-              q.format(statements)
+              q.text("end")
             end
-
-            q.breakable
-            q.text(force_parens ? "}" : "end")
           end
           .if_flat do
-            q.text("{ ")
-            q.format(statements)
-            q.text(" }")
+            q.text("{")
+
+            unless statements.empty?
+              q.text(" ")
+              q.format(statements)
+              q.text(" ")
+            end
+
+            q.text("}")
           end
       end
     end
   end
 
+  # LambdaVar represents the parameters being declared for a lambda. Effectively
+  # this node is everything contained within the parentheses. This includes all
+  # of the various parameter types, as well as block-local variable
+  # declarations.
+  #
+  #     -> (positional, optional = value, keyword:, &block; local) do
+  #     end
+  #
+  class LambdaVar < Node
+    # [Params] the parameters being declared with the block
+    attr_reader :params
+
+    # [Array[ Ident ]] the list of block-local variable declarations
+    attr_reader :locals
+
+    # [Array[ Comment | EmbDoc ]] the comments attached to this node
+    attr_reader :comments
+
+    def initialize(params:, locals:, location:, comments: [])
+      @params = params
+      @locals = locals
+      @location = location
+      @comments = comments
+    end
+
+    def accept(visitor)
+      visitor.visit_lambda_var(self)
+    end
+
+    def child_nodes
+      [params, *locals]
+    end
+
+    alias deconstruct child_nodes
+
+    def deconstruct_keys(_keys)
+      { params: params, locals: locals, location: location, comments: comments }
+    end
+
+    def empty?
+      params.empty? && locals.empty?
+    end
+
+    def format(q)
+      q.format(params)
+
+      if locals.any?
+        q.text("; ")
+        q.seplist(locals, -> { q.text(", ") }) { |local| q.format(local) }
+      end
+    end
+  end
+
   # LBrace represents the use of a left brace, i.e., {.
   class LBrace < Node
     # [String] the left brace

lib/syntax_tree/parser.rb

@@ -1940,6 +1940,41 @@ def on_lambda(params, statements)
             token.location.start_char > beginning.location.start_char
         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.
+      params =
+        case params
+        in Paren[contents: Params]
+          # In this case we've gotten to the <3.2 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])
+
+          location = params.contents.location
+          location = location.to(locals.last.location) if locals.any?
+
+          Paren.new(
+            lparen: params.lparen,
+            contents:
+              LambdaVar.new(
+                params: params.contents,
+                locals: locals,
+                location: location
+              ),
+            location: params.location,
+            comments: params.comments
+          )
+        in 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.
+          LambdaVar.new(params: params, locals: [], location: params.location)
+        in 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 = find_token(TLamBeg)
         closing = find_token(RBrace)
@@ -1962,6 +1997,84 @@ def on_lambda(params, statements)
       )
     end
 
+    # :call-seq:
+    #   on_lambda_var: (Params params, Array[ Ident ] locals) -> LambdaVar
+    def on_lambda_var(params, locals)
+      location = params.location
+      location = location.to(locals.last.location) if locals.any?
+
+      LambdaVar.new(params: params, locals: locals || [], location: location)
+    end
+
+    # Ripper doesn't support capturing lambda local variables until 3.2. To
+    # mitigate this, we have to parse that code for ourselves. We use the range
+    # from the parentheses to find where we _should_ be looking. Then we check
+    # if the resulting tokens match a pattern that we determine means that the
+    # declaration has block-local variables. Once it does, we parse those out
+    # and convert them into Ident nodes.
+    def lambda_locals(source)
+      tokens = Ripper.lex(source)
+
+      # First, check that we have a semi-colon. If we do, then we can start to
+      # parse the tokens _after_ the semicolon.
+      index = tokens.rindex { |token| token[1] == :on_semicolon }
+      return [] unless index
+
+      # Next, map over the tokens and convert them into Ident nodes. Bail out
+      # midway through if we encounter a token we didn't expect. Basically we're
+      # making our own mini-parser here. To do that we'll walk through a small
+      # state machine:
+      #
+      #     ┌────────┐               ┌────────┐                ┌─────────┐
+      #     │        │               │        │                │┌───────┐│
+      # ──> │  item  │ ─── ident ──> │  next  │ ─── rparen ──> ││ final ││
+      #     │        │ <── comma ─── │        │                │└───────┘│
+      #     └────────┘               └────────┘                └─────────┘
+      #        │  ^                     │  ^
+      #        └──┘                     └──┘
+      #    ignored_nl, sp              nl, sp
+      #
+      state = :item
+      transitions = {
+        item: {
+          on_ignored_nl: :item,
+          on_sp: :item,
+          on_ident: :next
+        },
+        next: {
+          on_nl: :next,
+          on_sp: :next,
+          on_comma: :item,
+          on_rparen: :final
+        },
+        final: {
+        }
+      }
+
+      tokens[(index + 1)..].each_with_object([]) do |token, locals|
+        (lineno, column), type, value, = token
+
+        # Make the state transition for the parser. This is going to raise a
+        # KeyError if we don't have a transition for the current state and type.
+        # But that shouldn't actually be possible because ripper would have
+        # found a syntax error by then.
+        state = transitions[state].fetch(type)
+
+        # If we hit an identifier, then add it to our list.
+        next if type != :on_ident
+
+        location =
+          Location.token(
+            line: lineno,
+            char: line_counts[lineno - 1][column],
+            column: column,
+            size: value.size
+          )
+
+        locals << Ident.new(value: value, location: location)
+      end
+    end
+
     # :call-seq:
     #   on_lbrace: (String value) -> LBrace
     def on_lbrace(value)

lib/syntax_tree/visitor.rb

@@ -301,6 +301,9 @@ def visit_child_nodes(node)
     # Visit a Lambda node.
     alias visit_lambda visit_child_nodes
 
+    # Visit a LambdaVar node.
+    alias visit_lambda_var visit_child_nodes
+
     # Visit a LBrace node.
     alias visit_lbrace visit_child_nodes
 

lib/syntax_tree/visitor/field_visitor.rb

@@ -586,6 +586,14 @@ def visit_lambda(node)
         end
       end
 
+      def visit_lambda_var(node)
+        node(node, "lambda_var") do
+          field("params", node.params)
+          list("locals", node.locals) if node.locals.any?
+          comments(node)
+        end
+      end
+
       def visit_lbrace(node)
         visit_token(node, "lbrace")
       end

test/fixtures/lambda.rb

@@ -1,4 +1,6 @@
 %
+-> {}
+%
 -> { foo }
 %
 ->(foo, bar) { baz }
@@ -40,3 +42,27 @@
 -> { -> foo do bar end.baz }.qux
 -
 -> { ->(foo) { bar }.baz }.qux
+%
+->(;a) {}
+-
+->(; a) {}
+%
+->(; a) {}
+%
+->(; a,b) {}
+-
+->(; a, b) {}
+%
+->(; a, b) {}
+%
+->(;
+a
+) {}
+-
+->(; a) {}
+%
+->(; a , 
+b
+) {}
+-
+->(; a, b) {}