Merge pull request #302 from glennsarti/can-i-lex

(GH-306) Add a syntax aware code folding provider

Author: jpogran

lib/puppet-languageserver/manifest/folding_provider.rb

@@ -0,0 +1,217 @@
+# frozen_string_literal: true
+
+require 'puppet-languageserver/puppet_lexer_helper'
+require 'lsp/lsp'
+
+module PuppetLanguageServer
+  module Manifest
+    class FoldingProvider
+      class << self
+        def instance
+          @instance ||= new
+        end
+
+        def supported?
+          # Folding is only supported on Puppet 6.3.0 and above
+          # Requires - https://github.com/puppetlabs/puppet/commit/6d375ab4d735779031d49ab8631bd9d161a9c3e3
+          @supported ||= Gem::Version.new(Puppet.version) >= Gem::Version.new('6.3.0')
+        end
+      end
+
+      REGION_NONE = nil
+      REGION_COMMENT = 'comment'
+      REGION_REGION = 'region'
+
+      def start_region?(text)
+        !(text =~ %r{^#\s*region\b}).nil?
+      end
+
+      def end_region?(text)
+        !(text =~ %r{^#\s*endregion\b}).nil?
+      end
+
+      def folding_ranges(tokens, show_last_line = false)
+        return nil unless self.class.supported?
+        ranges = {}
+
+        brace_stack = []
+        brack_stack = []
+        comment_stack = []
+
+        index = 0
+        until index > tokens.length - 1
+          case tokens[index][0]
+          # Find comments
+          when :TOKEN_COMMENT
+            if block_comment?(index, tokens)
+              comment = tokens[index][1].locator.extract_text(tokens[index][1].offset, tokens[index][1].length)
+              if start_region?(comment) # rubocop:disable Metrics/BlockNesting
+                comment_stack.push(tokens[index][1])
+              elsif end_region?(comment) && !comment_stack.empty? # rubocop:disable Metrics/BlockNesting
+                add_range!(create_range_span_tokens(comment_stack.pop, tokens[index][1], REGION_REGION), ranges)
+              else
+                index = process_block_comment!(index, tokens, ranges)
+              end
+            end
+
+          # Find old style comments /* -> */
+          when :TOKEN_MLCOMMENT
+            add_range!(create_range_whole_token(tokens[index][1], REGION_COMMENT), ranges)
+
+          # Find matching braces { -> } and select brace ?{ -> }
+          when :LBRACE, :SELBRACE
+            brace_stack.push(tokens[index][1])
+          when :RBRACE
+            add_range!(create_range_span_tokens(brace_stack.pop, tokens[index][1], REGION_NONE), ranges) unless brace_stack.empty?
+
+          # Find matching braces [ -> ], list and index
+          when :LISTSTART, :LBRACK
+            brack_stack.push(tokens[index][1])
+          when :RBRACK
+            add_range!(create_range_span_tokens(brack_stack.pop, tokens[index][1], REGION_NONE), ranges) unless brack_stack.empty?
+
+          # Find matching Heredoc and heredoc sublocations
+          when :HEREDOC
+            # Need to check if the next token is :SUBLOCATE
+            if index < tokens.length - 2 && tokens[index + 1][0] == :SUBLOCATE # rubocop:disable Style/IfUnlessModifier
+              add_range!(create_range_heredoc(tokens[index][1], tokens[index + 1][1], REGION_NONE), ranges)
+            end
+          end
+
+          index += 1
+        end
+
+        # If we are showing the last line then decrement the EndLine by one, if possible
+        if show_last_line
+          ranges.values.each do |range|
+            range.endLine = [range.startLine, range.endLine - 1].max
+            range.endCharacter = 0 # We don't know where the previous line actually ends so set it to zero
+          end
+        end
+
+        ranges.values
+      end
+
+      private
+
+      # region Internal Helper methods to call locator methods on Locators or SubLocators
+      def line_for_offset(token, offset = nil)
+        locator_method_with_offset(token, :line_for_offset, offset || token.offset)
+      end
+
+      def pos_on_line(token, offset = nil)
+        locator_method_with_offset(token, :pos_on_line, offset || token.offset)
+      end
+
+      def locator_method_with_offset(token, method_name, offset)
+        if token.locator.is_a?(Puppet::Pops::Parser::Locator::SubLocator)
+          global_offset, = token.locator.to_global(offset, token.length)
+          token.locator.locator.send(method_name, global_offset)
+        else
+          token.locator.send(method_name, offset)
+        end
+      end
+
+      def extract_text(token)
+        if token.locator.is_a?(Puppet::Pops::Parser::Locator::SubLocator)
+          global_offset, global_length = token.locator.to_global(token.offset, token.length)
+          token.locator.locator.extract_text(global_offset, global_length)
+        else
+          token.locator.extract_text(token.offset, token.length)
+        end
+      end
+      # endregion
+
+      # Return nil if not valid range
+      def create_range_span_tokens(start_token, end_token, kind)
+        start_line = line_for_offset(start_token) - 1
+        end_line = line_for_offset(end_token) - 1
+        return nil if start_line == end_line
+        LSP::FoldingRange.new({
+                                'startLine'      => start_line,
+                                'startCharacter' => pos_on_line(start_token) - 1,
+                                'endLine'        => end_line,
+                                'endCharacter'   => pos_on_line(end_token, end_token.offset + end_token.length) - 1,
+                                'kind'           => kind
+                              })
+      end
+
+      # Return nil if not valid range
+      def create_range_whole_token(token, kind)
+        start_line = line_for_offset(token) - 1
+        end_line = line_for_offset(token, token.offset + token.length) - 1
+        return nil if start_line == end_line
+        LSP::FoldingRange.new({
+                                'startLine'      => start_line,
+                                'startCharacter' => pos_on_line(token) - 1,
+                                'endLine'        => end_line,
+                                'endCharacter'   => pos_on_line(token, token.offset + token.length) - 1,
+                                'kind'           => kind
+                              })
+      end
+
+      # Return nil if not valid range
+      def create_range_heredoc(heredoc_token, subloc_token, kind)
+        start_line = line_for_offset(heredoc_token) - 1
+        # The lexer does not output the end heredoc_token. Instead we
+        # use the heredoc sublocator endline and add one
+        end_line = line_for_offset(heredoc_token, heredoc_token.offset + heredoc_token.length + subloc_token.length)
+        return nil if start_line == end_line
+        LSP::FoldingRange.new({
+                                'startLine'      => start_line,
+                                'startCharacter' => pos_on_line(heredoc_token) - 1,
+                                'endLine'        => end_line,
+                                # We don't know where the end token for the Heredoc is, so just assume it's at the start of the line
+                                'endCharacter'   => 0,
+                                'kind'           => kind
+                              })
+      end
+
+      # Adds a FoldingReference to the list and enforces ordering rules e.g. Only one fold per start line
+      def add_range!(range, ranges)
+        # Make sure the arguments are correct
+        return nil if range.nil? || ranges.nil?
+
+        # Ignore the range if there is an existing one which is bigger
+        return nil unless ranges[range.startLine].nil? || ranges[range.startLine].endLine < range.endLine
+        ranges[range.startLine] = range
+        nil
+      end
+
+      # Returns new index position
+      def process_block_comment!(index, tokens, ranges)
+        start_index = index
+        line_num = line_for_offset(tokens[index][1])
+        while index < tokens.length - 2
+          break unless tokens[index + 1][0] == :TOKEN_COMMENT
+          next_line = line_for_offset(tokens[index + 1][1])
+          # Tokens must be on contiguous lines
+          break unless next_line == line_num + 1
+          # Must not be a region comment
+          comment = extract_text(tokens[index + 1][1])
+          break if start_region?(comment) || end_region?(comment)
+          # It's a block comment
+          line_num = next_line
+          index += 1
+        end
+
+        return index if start_index == index
+
+        add_range!(create_range_span_tokens(tokens[start_index][1], tokens[index][1], REGION_COMMENT), ranges)
+        index
+      end
+
+      def block_comment?(index, tokens)
+        # Has to be a comment token
+        return false unless tokens[index][0] == :TOKEN_COMMENT
+        # If it's the first token then it has to be at the start of a line
+        return true if index.zero?
+        # It has to be the first token on this line
+        this_token_line = line_for_offset(tokens[index][1])
+        prev_token_line = line_for_offset(tokens[index - 1][1])
+
+        this_token_line != prev_token_line
+      end
+    end
+  end
+end

lib/puppet-languageserver/message_handler.rb

@@ -29,9 +29,12 @@ def request_initialize(_, json_rpc_message)
       PuppetLanguageServer.log_message(:debug, 'Received initialize method')
 
       language_client.parse_lsp_initialize!(json_rpc_message.params)
+      static_folding_provider = !language_client.client_capability('textDocument', 'foldingRange', 'dynamicRegistration') &&
+                                PuppetLanguageServer::ServerCapabilites.folding_provider_supported?
       # Setup static registrations if dynamic registration is not available
       info = {
-        :documentOnTypeFormattingProvider => !language_client.client_capability('textDocument', 'onTypeFormatting', 'dynamicRegistration')
+        :documentOnTypeFormattingProvider => !language_client.client_capability('textDocument', 'onTypeFormatting', 'dynamicRegistration'),
+        :foldingRangeProvider             => static_folding_provider
       }
 
       # Configure the document store
@@ -162,6 +165,20 @@ def request_textdocument_completion(_, json_rpc_message)
       LSP::CompletionList.new('isIncomplete' => false, 'items' => [])
     end
 
+    def request_textdocument_foldingrange(_, json_rpc_message)
+      return nil unless language_client.folding_range
+      file_uri = json_rpc_message.params['textDocument']['uri']
+      case documents.document_type(file_uri)
+      when :manifest
+        PuppetLanguageServer::Manifest::FoldingProvider.instance.folding_ranges(documents.document_tokens(file_uri))
+      else
+        raise "Unable to provide folding ranages on #{file_uri}"
+      end
+    rescue StandardError => e
+      PuppetLanguageServer.log_message(:error, "(textDocument/foldingRange) #{e}")
+      nil
+    end
+
     def request_completionitem_resolve(_, json_rpc_message)
       PuppetLanguageServer::Manifest::CompletionProvider.resolve(session_state, LSP::CompletionItem.new(json_rpc_message.params))
     rescue StandardError => e
@@ -285,6 +302,7 @@ def notification_initialized(_, _json_rpc_message)
           "Unable to use Puppet version '#{server_options[:puppet_version]}' as it is not available. Using version '#{Puppet.version}' instead."
         )
       end
+
       # Register for workspace setting changes if it's supported
       if language_client.client_capability('workspace', 'didChangeConfiguration', 'dynamicRegistration') == true
         language_client.register_capability('workspace/didChangeConfiguration')

lib/puppet-languageserver/providers.rb

@@ -5,6 +5,7 @@
   manifest/completion_provider
   manifest/definition_provider
   manifest/document_symbol_provider
+  manifest/folding_provider
   manifest/format_on_type_provider
   manifest/signature_provider
   manifest/validation_provider

lib/puppet-languageserver/puppet_lexer_helper.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module Puppet
+  module Pops
+    module Parser
+      # This Lexer adds code to create comment tokens.
+      # The default lexer just throws them out
+      # Ref - https://github.com/puppetlabs/puppet-specifications/blob/master/language/lexical_structure.md#comments
+      class Lexer2WithComments < Puppet::Pops::Parser::Lexer2
+        # The PATTERN_COMMENT in lexer2 also consumes the trailing \r in the token and
+        # we don't want that.
+        PATTERN_COMMENT_NO_WS = %r{#[^\r\n]*}.freeze
+
+        TOKEN_COMMENT = [:COMMENT, '#', 1].freeze
+        TOKEN_MLCOMMENT = [:MLCOMMENT, nil, 0].freeze
+
+        def initialize
+          super
+
+          # Remove the selector for line comments so we can add our own
+          @new_selector = @selector.reject { |k, _v| k == '#' }
+
+          # Add code to scan line comments
+          @new_selector['#'] = lambda {
+            scn = @scanner
+            before = scn.pos
+            value = scn.scan(PATTERN_COMMENT_NO_WS)
+
+            if value
+              emit_completed([:TOKEN_COMMENT, value[1..-1].freeze, scn.pos - before], before)
+            else
+              # It's probably not possible to EVER get here ... but just incase
+              emit(TOKEN_COMMENT, before)
+            end
+          }.freeze
+
+          # Add code to scan multi-line comments
+          old_lambda = @new_selector['/']
+          @new_selector['/'] = lambda {
+            scn = @scanner
+            la = scn.peek(2)
+            if la[1] == '*'
+              before = scn.pos
+              value = scn.scan(PATTERN_MLCOMMENT)
+              return emit_completed([:TOKEN_MLCOMMENT, value[2..-3].freeze, scn.pos - before], before) if value
+            end
+            old_lambda.call
+          }.freeze
+          @new_selector.freeze
+          @selector = @new_selector
+        end
+      end
+    end
+  end
+end

lib/puppet-languageserver/server_capabilities.rb

@@ -4,6 +4,10 @@
 
 module PuppetLanguageServer
   module ServerCapabilites
+    def self.folding_provider_supported?
+      @folding_provider ||= PuppetLanguageServer::Manifest::FoldingProvider.supported?
+    end
+
     def self.capabilities(options = {})
       # https://github.com/Microsoft/language-server-protocol/blob/master/protocol.md#initialize-request
 
@@ -22,6 +26,7 @@ def self.capabilities(options = {})
         }
       }
       value['documentOnTypeFormattingProvider'] = document_on_type_formatting_options if options[:documentOnTypeFormattingProvider]
+      value['foldingRangeProvider'] = folding_range_provider_options if options[:foldingRangeProvider]
       value
     end
 
@@ -31,6 +36,10 @@ def self.document_on_type_formatting_options
       }
     end
 
+    def self.folding_range_provider_options
+      true
+    end
+
     def self.no_capabilities
       # Any empty hash denotes no capabilities at all
       {