rbs.rake

namespace :rbs do
  task zeitwerk: :environment do
    namespace_keys = []

    # 他のRakeタスクと続けて実行すると既にクラスを読み込んでいるため、 `__namespace_dirs` が
    # 空になってしまうので、リロードしておく。
    Rails.autoloaders.main.reload

    # constantize すると `__namespace_dirs` が更新されるので、空になるまで繰り返す。
    until Rails.autoloaders.main.__namespace_dirs.empty?
      namespace_keys += Rails.autoloaders.main.__namespace_dirs.keys
      namespace_keys.each(&:constantize)
    end

    # Zeitwerkに生成された定数の一覧を取得する。
    namespaces = namespace_keys.map(&:constantize).select do |klass|
      filename, _line = Object.const_source_location(klass.to_s)
      filename.end_with?("lib/zeitwerk/loader/callbacks.rb")
    end

    # `Foo::Bar::Buz` を `{ Foo: { Bar: :Buz } }` の構造に変換する。
    # 処理を単純にするため、Object を先端に入れている。
    structured = { Object => {} }
    namespaces.each do |namespace|
      [namespace, *namespace.module_parents].reverse.inject(structured) do |result, klass|
        result[klass] ||= {}
      end
    end

    # Hashを再起的にRBSに変換する。
    to_rbs = ->(modules, indent) {
      next if modules.empty?

      modules.map {|klass, body|
        name = klass.name.split("::").last
        definitions = []
        definitions << if klass.instance_of?(Module)
                         "module #{name}"
                       elsif klass.instance_of?(Class) && klass.superclass
                         "class #{name} < #{klass.superclass}"
                       else
                         "class #{name}"
                       end

        definitions += [to_rbs[body, indent + 1], "end"]
        definitions.compact.join("\n").indent(indent)
      }.join("\n")
    }

    # RBSをファイルに書き込む。
    content = <<~RBS
      # Generated from lib/tasks/rbs.rake

      #{to_rbs[structured[Object], 0]}
    RBS
    path = Rails.root.join("sig/generated/zeitwerk.rbs")
    path.parent.mkdir unless path.parent.exist?
    path.write(content)
  end
end