WIP: [Tapioca Addon] Support gem RBI generation

.rubocop.yml

@@ -45,6 +45,8 @@ Sorbet/TrueSigil:
   Include:
     - "**/*.rb"
     - "**/*.rake"
+  Exclude:
+    - "lib/ruby_lsp/tapioca/server_addon.rb"
 
 Style/CaseEquality:
   Enabled: false

Gemfile

@@ -15,6 +15,8 @@ gem "irb"
 gem "rubocop-shopify"
 gem "rubocop-sorbet", ">= 0.4.1"
 gem "rubocop-rspec"
+gem "ruby-lsp", ">= 0.19.1", git: "https://github.com/Shopify/ruby-lsp.git", branch: "ar/shutdown-timestamp"
+gem "ruby-lsp-rails", ">= 0.3.18"
 
 group :deployment, :development do
   gem "rake"

Gemfile.lock

@@ -1,3 +1,14 @@
+GIT
+  remote: https://github.com/Shopify/ruby-lsp.git
+  revision: 06602886c19cad7f44101b143bdbba9a5fc40a0d
+  branch: ar/shutdown-timestamp
+  specs:
+    ruby-lsp (0.21.4)
+      language_server-protocol (~> 3.17.0)
+      prism (>= 1.2, < 2.0)
+      rbs (>= 3, < 4)
+      sorbet-runtime (>= 0.5.10782)
+
 GIT
   remote: https://github.com/csfrancis/cityhash.git
   revision: 3cfc7d01f333c01811d5e834f1495eaa29f87c36
@@ -273,6 +284,8 @@ GEM
     rbi (0.2.1)
       prism (~> 1.0)
       sorbet-runtime (>= 0.5.9204)
+    rbs (3.6.1)
+      logger
     rdoc (6.7.0)
       psych (>= 4.0.0)
     redis (5.0.8)
@@ -301,6 +314,8 @@ GEM
       rubocop (~> 1.51)
     rubocop-sorbet (0.8.6)
       rubocop (>= 1)
+    ruby-lsp-rails (0.3.25)
+      ruby-lsp (>= 0.21.2, < 0.22.0)
     ruby-progressbar (1.13.0)
     shopify-money (3.0.0)
     sidekiq (7.3.2)
@@ -385,6 +400,8 @@ DEPENDENCIES
   rubocop-rspec
   rubocop-shopify
   rubocop-sorbet (>= 0.4.1)
+  ruby-lsp (>= 0.19.1)!
+  ruby-lsp-rails (>= 0.3.18)
   shopify-money
   sidekiq
   smart_properties

lib/ruby_lsp/tapioca/addon.rb

@@ -0,0 +1,170 @@
+# typed: strict
+# frozen_string_literal: true
+
+RubyLsp::Addon.depend_on_ruby_lsp!(">= 0.20", "< 0.22")
+
+begin
+  # The Tapioca add-on depends on the Rails add-on to add a runtime component to the runtime server. We can allow the
+  # add-on to work outside of a Rails context in the future, but that may require Tapioca spawning its own runtime
+  # server
+  require "ruby-lsp-rails"
+rescue LoadError
+  return
+end
+
+require "zlib"
+
+module RubyLsp
+  module Tapioca
+    class Addon < ::RubyLsp::Addon
+      extend T::Sig
+
+      GEMFILE_LOCK_SNAPSHOT = "tmp/tapioca/.gemfile_lock_snapshot"
+      GIT_OPERATION_THRESHOLD = 15.0 # seconds
+
+      sig { void }
+      def initialize
+        super
+
+        @global_state = T.let(nil, T.nilable(RubyLsp::GlobalState))
+        @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])
+      end
+
+      sig { override.params(global_state: RubyLsp::GlobalState, outgoing_queue: Thread::Queue).void }
+      def activate(global_state, outgoing_queue)
+        @global_state = global_state
+        # TODO: Uncomment
+        # return unless @global_state.experimental_features
+
+        @index = @global_state.index
+        Thread.new do
+          # Get a handle to the Rails add-on's runtime client. The call to `rails_runner_client` will block this thread
+          # until the server has finished booting, but it will not block the main LSP. This has to happen inside of a
+          # thread
+          addon = T.cast(::RubyLsp::Addon.get("Ruby LSP Rails", ">= 0.3.18", "< 0.4"), ::RubyLsp::Rails::Addon)
+          @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__))
+
+          handle_gemfile_changes
+        rescue IncompatibleApiError
+          # The requested version for the Rails add-on no longer matches. We need to upgrade and fix the breaking
+          # changes
+        end
+      end
+
+      sig { override.void }
+      def deactivate
+      end
+
+      sig { override.returns(String) }
+      def name
+        "Tapioca"
+      end
+
+      sig { override.returns(String) }
+      def version
+        "0.1.0"
+      end
+
+      sig { params(changes: T::Array[{ uri: String, type: Integer }]).void }
+      def workspace_did_change_watched_files(changes)
+        # TODO: Uncomment
+        # return unless T.must(@global_state).experimental_features
+        return unless @rails_runner_client # Client is not ready
+
+        constants = changes.flat_map do |change|
+          path = URI(change[:uri]).to_standardized_path
+          next if path.end_with?("_test.rb", "_spec.rb")
+
+          case change[:type]
+          when Constant::FileChangeType::CREATED, Constant::FileChangeType::CHANGED
+            content = File.read(path)
+            current_checksum = Zlib.crc32(content).to_s
+
+            if change[:type] == Constant::FileChangeType::CHANGED && @file_checksums[path] == current_checksum
+              $stderr.puts "File has not changed. Skipping #{path}"
+              next
+            end
+
+            @file_checksums[path] = current_checksum
+          when Constant::FileChangeType::DELETED
+            @file_checksums.delete(path)
+          end
+
+          entries = T.must(@index).entries_for(path)
+          next unless entries
+
+          entries.filter_map do |entry|
+            entry.name if entry.class == RubyIndexer::Entry::Class || entry.class == RubyIndexer::Entry::Module
+          end
+        end.compact
+
+        return if constants.empty?
+
+        @rails_runner_client.trigger_reload
+        @rails_runner_client.delegate_notification(
+          server_addon_name: "Tapioca",
+          request_name: "dsl",
+          constants: constants,
+        )
+      end
+
+      private
+
+      sig { void }
+      def handle_gemfile_changes
+        return unless File.exist?(".git") && File.exist?(".ruby-lsp/shutdown-timestamp")
+
+        git_timestamp = fetch_last_git_operation_time
+        return unless git_timestamp
+
+        ruby_lsp_stop_time = Time.iso8601(File.read(".ruby-lsp/shutdown-timestamp"))
+
+        $stderr.puts("ruby_lsp_stop_time: #{ruby_lsp_stop_time}") # TODO: Remove
+        $stderr.puts("git_timestamp: #{git_timestamp}") # TODO: Remove
+
+        # If the Ruby LSP was stopped shortly after the last git checkout/pull operation, we don't need to regenerate
+        # RBIs since any Gemfile.lock changes were likely from version control, not from running bundle install
+        return if (ruby_lsp_stop_time - git_timestamp) <= GIT_OPERATION_THRESHOLD
+
+        process_gemfile_changes
+      end
+
+      sig { returns(T.nilable(Time)) }
+      def fetch_last_git_operation_time
+        git_reflog_output = %x(git reflog --date=iso | grep -E "checkout|pull" | head -n 1).strip
+        return if git_reflog_output.empty?
+
+        timestamp_string = T.must(git_reflog_output.match(/\{(.*)\}/))[1]
+        Time.iso8601(T.must(timestamp_string).sub(" ", "T").delete(" "))
+      end
+
+      sig { returns(T.nilable(Integer)) }
+      def process_gemfile_changes
+        current_lockfile = File.read("Gemfile.lock")
+        snapshot_lockfile = File.read(GEMFILE_LOCK_SNAPSHOT) if File.exist?(GEMFILE_LOCK_SNAPSHOT)
+
+        unless snapshot_lockfile
+          $stdout.puts("Creating initial Gemfile.lock snapshot at #{GEMFILE_LOCK_SNAPSHOT}")
+          FileUtils.mkdir_p(File.dirname(GEMFILE_LOCK_SNAPSHOT))
+          File.write(GEMFILE_LOCK_SNAPSHOT, current_lockfile)
+          return
+        end
+
+        return if current_lockfile == snapshot_lockfile
+
+        T.must(@rails_runner_client).delegate_notification(
+          server_addon_name: "Tapioca",
+          request_name: "gem",
+          snapshot_lockfile: snapshot_lockfile,
+          current_lockfile: current_lockfile,
+        )
+
+        File.write(GEMFILE_LOCK_SNAPSHOT, current_lockfile)
+      end
+    end
+  end
+end

lib/ruby_lsp/tapioca/server_addon.rb

@@ -0,0 +1,58 @@
+# typed: false
+# frozen_string_literal: true
+
+require "tapioca/internal"
+
+module RubyLsp
+  module Tapioca
+    class ServerAddon < ::RubyLsp::Rails::ServerAddon
+      def name
+        "Tapioca"
+      end
+
+      def execute(request, params)
+        case request
+        when "dsl"
+          dsl(params)
+        when "gem"
+          gem(params)
+        end
+      end
+
+      private
+
+      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)
+        snapshot_specs = parse_lockfile(params[:snapshot_lockfile])
+        current_specs = parse_lockfile(params[:current_lockfile])
+
+        removed_gems = snapshot_specs.keys - current_specs.keys
+        changed_gems = current_specs.select { |name, version| snapshot_specs[name] != version }.keys
+
+        return $stdout.puts("No gem changes detected") if removed_gems.empty? && changed_gems.empty?
+
+        if removed_gems.any?
+          $stdout.puts("Removing RBIs for deleted gems: #{removed_gems.join(", ")}")
+          FileUtils.rm_f(Dir.glob("sorbet/rbi/gems/{#{removed_gems.join(",")}}@*.rbi"))
+        end
+
+        if changed_gems.any?
+          $stdout.puts("Generating RBIs for changed gems: #{changed_gems.join(", ")}")
+
+          load("tapioca/cli.rb") # Reload the CLI to reset thor defaults between requests
+          ::Tapioca::Cli.start(["gem"] + changed_gems)
+        end
+      end
+
+      def parse_lockfile(content)
+        return {} if content.to_s.empty?
+
+        Bundler::LockfileParser.new(content).specs.to_h { |spec| [spec.name, spec.version.to_s] }
+      end
+    end
+  end
+end

lib/tapioca/cli.rb

@@ -144,6 +144,11 @@ def todo
       type: :hash,
       desc: "Options to pass to the DSL compilers",
       default: {}
+    option :lsp_addon,
+      type: :boolean,
+      desc: "Generate DSL RBIs from the LSP addon. Internal to tapioca and not intended for end-users",
+      default: false,
+      hide: true
     def dsl(*constant_or_paths)
       set_environment(options)
 
@@ -166,6 +171,7 @@ def dsl(*constant_or_paths)
         app_root: options[:app_root],
         halt_upon_load_error: options[:halt_upon_load_error],
         compiler_options: options[:compiler_options],
+        lsp_addon: options[:lsp_addon],
       }
 
       command = if options[:verify]

lib/tapioca/commands/abstract_dsl.rb

@@ -28,6 +28,7 @@ class AbstractDsl < CommandWithoutTracker
           app_root: String,
           halt_upon_load_error: T::Boolean,
           compiler_options: T::Hash[String, T.untyped],
+          lsp_addon: T::Boolean,
         ).void
       end
       def initialize(
@@ -47,7 +48,8 @@ def initialize(
         rbi_formatter: DEFAULT_RBI_FORMATTER,
         app_root: ".",
         halt_upon_load_error: true,
-        compiler_options: {}
+        compiler_options: {},
+        lsp_addon: false
       )
         @requested_constants = requested_constants
         @requested_paths = requested_paths
@@ -66,6 +68,7 @@ def initialize(
         @halt_upon_load_error = halt_upon_load_error
         @skip_constant = skip_constant
         @compiler_options = compiler_options
+        @lsp_addon = lsp_addon
 
         super()
       end
@@ -74,6 +77,10 @@ def initialize(
 
       sig { params(outpath: Pathname, quiet: T::Boolean).returns(T::Set[Pathname]) }
       def generate_dsl_rbi_files(outpath, quiet:)
+        if @lsp_addon
+          pipeline.active_compilers.each(&:reset_state)
+        end
+
         existing_rbi_files = existing_rbi_filenames(all_requested_constants)
 
         generated_files = pipeline.run do |constant, contents|
@@ -116,6 +123,7 @@ def load_application
           eager_load: @requested_constants.empty? && @requested_paths.empty?,
           app_root: @app_root,
           halt_upon_load_error: @halt_upon_load_error,
+          lsp_addon: @lsp_addon,
         )
       end
 
@@ -133,6 +141,7 @@ def create_pipeline
           skipped_constants: constantize(@skip_constant, ignore_missing: true),
           number_of_workers: @number_of_workers,
           compiler_options: @compiler_options,
+          lsp_addon: @lsp_addon,
         )
       end
 

lib/tapioca/dsl/compiler.rb

@@ -51,6 +51,13 @@ def requested_constants=(constants)
           @@requested_constants = constants # rubocop:disable Style/ClassVars
         end
 
+        sig { void }
+        def reset_state
+          @processable_constants = nil
+          @all_classes = nil
+          @all_modules = nil
+        end
+
         private
 
         sig do