rbs_optimizer.rb

# frozen_string_literal: true

#
# RBS loading optimizer
#
# Usage: require 'rbs_optimizer' in your .irbrc
#

require 'rbs'
require 'rbs/cli'

module RBS
  class EnvironmentLoader
    attr_accessor :resolved, :cached_sources

    def latest_modified_time
      @_latest_modified_time ||= begin
        mtimes = []
        each_dir do |source, dir|
          skip_hidden = !source.is_a?(Pathname)
          mtimes << RBS::FileFinder.each_file(dir, skip_hidden:, immediate: true).map { |f| File.mtime(f) }.max
        end
        mtimes.max
      end
    end

    def load(env:)
      # @type var loaded: Array[[AST::Declarations::t, Pathname, source]]
      loaded = []
      cached = cached_sources || []

      each_signature do |source, path, buffer, decls, dirs|
        decls.each do |decl|
          decl.resolved = cached.include?(source)
          loaded << [decl, path, source]
        end
        env.add_signature(buffer:, directives: dirs, decls:)
      end

      loaded
    end
  end

  module AST
    module Declarations
      class Base
        attr_accessor :resolved
      end
    end
  end

  class Environment
    attr_accessor :loader

    def self.from_loader(loader, cache: true)
      new_loader = cache ? cached_loader(loader) : loader
      new.tap do |env|
        env.loader = new_loader
        new_loader.load(env:)
      end
    end

    def self.cached_loader(loader)
      if system_cached?(loader.latest_modified_time)
        fully_cached_loader
      else
        write_collection_cache unless collection_cached?

        collection_loader.dirs.each { |dir| loader.dirs.delete(dir) }
        collection_loader.libs.each { |lib| loader.libs.delete(lib) }

        loader.instance_eval { @core_root = nil }
        loader.dirs.unshift collection_cache_path
        loader.cached_sources = [collection_cache_path]
        loader
      end
    end

    def self.fully_cached_loader
      loader = RBS::EnvironmentLoader.new(core_root: nil)
      loader.cached_sources = [system_cache_path]
      loader.add(path: system_cache_path)
      loader
    end

    def self.collection_loader
      loader = RBS::CLI::LibraryOptions.new.loader
      loader.instance_eval { @core_root = nil }
      loader
    end

    def self.system_cached?(mtime)
      system_cache_path.exist? && system_cache_path.mtime == mtime
    end

    def self.collection_cached?
      mtime = collection_loader.latest_modified_time
      collection_cache_path.exist? && collection_cache_path.mtime == mtime
    end

    def self.system_cache_path
      @_system_cache_path ||= cached_path('system.rbs')
    end

    def self.collection_cache_path
      @_collection_cache_path ||= cached_path('collection.rbs')
    end

    def self.cached_path(filename)
      Pathname(ENV['XDG_CACHE_HOME'] || File.expand_path('~/.cache')).join('rbs', filename)
    end

    def self.write_collection_cache
      loader = collection_loader
      env = RBS::Environment.from_loader(loader, cache: false)
      env.resolve_type_names
      write_cache(collection_cache_path, env, loader)
    end

    def self.write_cache(filename, env, loader)
      mtime = loader.latest_modified_time
      cache_path = cached_path(filename)
      return if cache_path.exist? && cache_path.mtime == mtime

      cache_path.parent.mkpath
      cache_path.open('wt') do |f|
        RBS::Writer.new(out: f).write(env.declarations)
      end
      cache_path.utime(mtime, mtime)
    end

    alias original_resolve_type_names resolve_type_names

    def resolve_type_names(only: nil, cache_filename: 'system.rbs')
      original_resolve_type_names(only:)

      self.class.write_cache(cache_filename, self, loader)
    end

    alias original_resolve_declaration resolve_declaration

    def resolve_declaration(resolver, map, decl, outer:, prefix:)
      return decl if decl.resolved

      original_resolve_declaration(resolver, map, decl, outer:, prefix:)
    end
  end
end