[WIP] [Tapioca Addon] Optimized gem RBI generation Exploration

lib/ruby_lsp/tapioca/addon.rb

@@ -13,6 +13,7 @@
 end
 
 require "zlib"
+require "open3"
 
 module RubyLsp
   module Tapioca
@@ -27,6 +28,7 @@ def initialize
         @rails_runner_client = T.let(nil, T.nilable(RubyLsp::Rails::RunnerClient))
         @index = T.let(nil, T.nilable(RubyIndexer::Index))
         @file_checksums = T.let({}, T::Hash[String, String])
+        @lockfile_diff = T.let(nil, T.nilable(String))
         @outgoing_queue = T.let(nil, T.nilable(Thread::Queue))
       end
 
@@ -45,6 +47,10 @@ def activate(global_state, outgoing_queue)
           @rails_runner_client = addon.rails_runner_client
           @outgoing_queue << Notification.window_log_message("Activating Tapioca add-on v#{version}")
           @rails_runner_client.register_server_addon(File.expand_path("server_addon.rb", __dir__))
+
+          if git_repo?
+            lockfile_changed? ? generate_gem_rbis : cleanup_orphaned_rbis
+          end
         rescue IncompatibleApiError
           # The requested version for the Rails add-on no longer matches. We need to upgrade and fix the breaking
           # changes
@@ -127,6 +133,56 @@ def file_updated?(change, path)
 
         false
       end
+
+      sig { returns(T.nilable(T::Boolean)) }
+      def git_repo?
+        _, status = Open3.capture2e("git rev-parse --is-inside-work-tree")
+
+        status.success?
+      end
+
+      sig { returns(T::Boolean) }
+      def lockfile_changed?
+        !fetch_lockfile_diff.empty?
+      end
+
+      sig { returns(String) }
+      def fetch_lockfile_diff
+        @lockfile_diff = %x(git diff HEAD Gemfile.lock).strip
+      end
+
+      sig { void }
+      def generate_gem_rbis
+        T.must(@rails_runner_client).delegate_notification(
+          server_addon_name: "Tapioca",
+          request_name: "gem",
+          diff: T.must(@lockfile_diff),
+        )
+      end
+
+      sig { void }
+      def cleanup_orphaned_rbis
+        untracked_files = %x(git ls-files --others --exclude-standard sorbet/rbi/gems/).lines.map(&:strip)
+        deleted_files = %x(git ls-files --deleted sorbet/rbi/gems/).lines.map(&:strip)
+
+        untracked_files.each do |file|
+          File.delete(file)
+
+          T.must(@outgoing_queue) << Notification.window_log_message(
+            "Deleted untracked RBI: #{file}",
+            type: Constant::MessageType::INFO,
+          )
+        end
+
+        deleted_files.each do |file|
+          %x(git checkout -- #{file})
+
+          T.must(@outgoing_queue) << Notification.window_log_message(
+            "Restored deleted RBI: #{file}",
+            type: Constant::MessageType::INFO,
+          )
+        end
+      end
     end
   end
 end

lib/ruby_lsp/tapioca/lockfile_diff_parser.rb

@@ -0,0 +1,43 @@
+# typed: true
+# frozen_string_literal: true
+
+module RubyLsp
+  module Tapioca
+    class LockfileDiffParser
+      GEM_NAME_PATTERN = /[\w\-]+/
+      DIFF_LINE_PATTERN = /[+-](.*#{GEM_NAME_PATTERN})\s*\(/
+      ADDED_LINE_PATTERN = /^\+.*#{GEM_NAME_PATTERN} \(.*\)/
+      REMOVED_LINE_PATTERN = /^-.*#{GEM_NAME_PATTERN} \(.*\)/
+
+      attr_reader :added_or_modified_gems
+      attr_reader :removed_gems
+
+      def initialize(diff_content)
+        @diff_content = diff_content.lines
+        @added_or_modified_gems = parse_added_or_modified_gems
+        @removed_gems = parse_removed_gems
+      end
+
+      private
+
+      def parse_added_or_modified_gems
+        @diff_content
+          .filter { |line| line.match?(ADDED_LINE_PATTERN) }
+          .map { |line| extract_gem(line) }
+          .uniq
+      end
+
+      def parse_removed_gems
+        @diff_content
+          .filter { |line| line.match?(REMOVED_LINE_PATTERN) }
+          .map { |line| extract_gem(line) }
+          .reject { |gem| @added_or_modified_gems.include?(gem) }
+          .uniq
+      end
+
+      def extract_gem(line)
+        line.match(DIFF_LINE_PATTERN)[1].strip
+      end
+    end
+  end
+end

lib/ruby_lsp/tapioca/server_addon.rb

@@ -2,6 +2,7 @@
 # frozen_string_literal: true
 
 require "tapioca/internal"
+require_relative "lockfile_diff_parser"
 
 module RubyLsp
   module Tapioca
@@ -16,6 +17,8 @@ def execute(request, params)
           fork do
             dsl(params)
           end
+        when "gem"
+          gem(params)
         end
       end
 
@@ -25,6 +28,24 @@ def dsl(params)
         load("tapioca/cli.rb") # Reload the CLI to reset thor defaults between requests
         ::Tapioca::Cli.start(["dsl", "--lsp_addon", "--workers=1"] + params[:constants])
       end
+
+      def gem(params)
+        gem_changes = LockfileDiffParser.new(params[:diff])
+
+        removed_gems = gem_changes.removed_gems
+        added_or_modified_gems = gem_changes.added_or_modified_gems
+
+        if added_or_modified_gems.any?
+          ::Tapioca::Cli.start([
+            "gem",
+            "--lsp_addon",
+            *added_or_modified_gems,
+          ])
+        elsif removed_gems.any?
+          FileUtils.rm_f(Dir.glob("sorbet/rbi/gems/{#{removed_gems.join(",")}}@*.rbi"))
+          $stdout.puts "Removed RBIs for: #{removed_gems.join(", ")}"
+        end
+      end
     end
   end
 end

lib/tapioca/cli.rb

@@ -266,6 +266,11 @@ def dsl(*constant_or_paths)
       type: :boolean,
       desc: "Halt upon a load error while loading the Rails application",
       default: true
+    option :lsp_addon,
+      type: :boolean,
+      desc: "Generate Gem RBIs from the LSP addon. Internal to tapioca and not intended for end-users",
+      default: false,
+      hide: true
     def gem(*gems)
       set_environment(options)
 
@@ -300,6 +305,7 @@ def gem(*gems)
         dsl_dir: options[:dsl_dir],
         rbi_formatter: rbi_formatter(options),
         halt_upon_load_error: options[:halt_upon_load_error],
+        lsp_addon: options[:lsp_addon],
       }
 
       command = if verify

lib/tapioca/commands/abstract_gem.rb

@@ -27,6 +27,7 @@ class AbstractGem < Command
           dsl_dir: String,
           rbi_formatter: RBIFormatter,
           halt_upon_load_error: T::Boolean,
+          lsp_addon: T.nilable(T::Boolean),
         ).void
       end
       def initialize(
@@ -45,7 +46,8 @@ def initialize(
         auto_strictness: true,
         dsl_dir: DEFAULT_DSL_DIR,
         rbi_formatter: DEFAULT_RBI_FORMATTER,
-        halt_upon_load_error: true
+        halt_upon_load_error: true,
+        lsp_addon: false
       )
         @gem_names = gem_names
         @exclude = exclude
@@ -59,6 +61,7 @@ def initialize(
         @auto_strictness = auto_strictness
         @dsl_dir = dsl_dir
         @rbi_formatter = rbi_formatter
+        @lsp_addon = lsp_addon
 
         super()
 
@@ -81,6 +84,8 @@ def gems_to_generate(gem_names)
           gem = @bundle.gem(gem_name)
 
           if gem.nil?
+            next say("Warning: Cannot find gem '#{gem_name}', skipping", :yellow) if @lsp_addon
+
             raise Thor::Error, set_color("Error: Cannot find gem '#{gem_name}'", :red)
           end
 
@@ -126,6 +131,7 @@ def compile_gem_rbi(gem)
             error_handler: ->(error) {
               say_error(error, :bold, :red)
             },
+            lsp_addon: T.must(@lsp_addon),
           ).compile
         end
 
@@ -192,6 +198,7 @@ def perform_additions
               postrequire: @postrequire,
               default_command: default_command(:require),
               halt_upon_load_error: @halt_upon_load_error,
+              lsp_addon: T.must(@lsp_addon),
             )
 
             Executor.new(gems, number_of_workers: @number_of_workers).run_in_parallel do |gem_name|

lib/tapioca/commands/gem_generate.rb

@@ -14,6 +14,7 @@ def execute
           postrequire: @postrequire,
           default_command: default_command(:require),
           halt_upon_load_error: @halt_upon_load_error,
+          lsp_addon: T.must(@lsp_addon),
         )
 
         gem_queue = gems_to_generate(@gem_names).reject { |gem| @exclude.include?(gem.name) }

lib/tapioca/gem/pipeline.rb

@@ -22,19 +22,22 @@ class Pipeline
           error_handler: T.proc.params(error: String).void,
           include_doc: T::Boolean,
           include_loc: T::Boolean,
+          lsp_addon: T::Boolean,
         ).void
       end
       def initialize(
         gem,
         error_handler:,
         include_doc: false,
-        include_loc: false
+        include_loc: false,
+        lsp_addon: false
       )
         @root = T.let(RBI::Tree.new, RBI::Tree)
         @gem = gem
         @seen = T.let(Set.new, T::Set[String])
         @alias_namespace = T.let(Set.new, T::Set[String])
         @error_handler = error_handler
+        @lsp_addon = lsp_addon
 
         @events = T.let([], T::Array[Gem::Event])
 

lib/tapioca/loaders/gem.rb

@@ -16,9 +16,12 @@ class << self
             postrequire: String,
             default_command: String,
             halt_upon_load_error: T::Boolean,
+            lsp_addon: T::Boolean,
           ).void
         end
-        def load_application(bundle:, prerequire:, postrequire:, default_command:, halt_upon_load_error:)
+        def load_application(bundle:, prerequire:, postrequire:, default_command:, halt_upon_load_error:, lsp_addon:)
+          return if lsp_addon
+
           loader = new(
             bundle: bundle,
             prerequire: prerequire,