rbs_macros 0.1.0

data/lib/rbs_macros/environment.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+require "prism"
+
+module RbsMacros
+  # An environment for the Ruby program being analyzed.
+  class Environment
+    attr_reader :rbs, :object_class, :decls
+
+    def initialize
+      @rbs = RBS::Environment.new
+      @object_class = MetaClass.new(self, "Object", is_class: true)
+      @decls = []
+      @exact_handlers = {}
+    end
+
+    def register_handler(name, handler)
+      (@exact_handlers[name.to_sym] ||= []) << handler
+    end
+
+    def invoke(params)
+      handlers = @exact_handlers[params.name]
+      handlers&.each do |handler|
+        handler.(params)
+      end
+    end
+
+    def meta_eval_ruby(code, filename:)
+      result = Prism.parse(code, filepath: "#{filename}.rb")
+      raise ArgumentError, "Parse error: #{result.errors}" if result.failure?
+
+      ExecCtx.new(
+        env: self,
+        filename:,
+        self: nil, # TODO
+        cref: @object_class,
+        cref_dynamic: @object_class,
+        locals: {}
+      ).eval_node(result.value)
+    end
+
+    def add_decl(decl, mod:, file:)
+      @decls << DeclarationEntry.new(declaration: decl, mod:, file:)
+    end
+
+    HandlerParams = _ = Data.define(:env, :filename, :receiver, :name, :positional, :keyword, :block) # rubocop:disable Naming/ConstantName
+
+    DeclarationEntry = _ = Data.define(:declaration, :mod, :file) # rubocop:disable Naming/ConstantName
+  end
+end

data/lib/rbs_macros/exec_ctx.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+module RbsMacros
+  ExecCtx = _ = Data.define(:env, :filename, :self, :cref, :cref_dynamic, :locals) # rubocop:disable Naming/ConstantName
+
+  # Context, including self, class, and local variables
+  class ExecCtx
+    def eval_node(node)
+      case node
+      when nil
+        # do nothing
+      when Prism::CallNode
+        recv = \
+          if node.receiver
+            eval_node(node.receiver)
+          else
+            self.self
+          end
+        positional = [] # : Array[Object]
+        keyword = {} # : Hash[Object, Object]
+        node.arguments&.arguments&.each do |arg|
+          positional << eval_node(arg)
+        end
+        env.invoke(
+          Environment::HandlerParams.new(
+            env:,
+            filename:,
+            receiver: recv,
+            name: node.name,
+            positional:,
+            keyword:,
+            block: nil
+          )
+        )
+      when Prism::ClassNode
+        klass = cref.define_module(node.name)
+        klass.class!
+        with(
+          self: klass,
+          cref: klass,
+          cref_dynamic: klass,
+          locals: {}
+        ).eval_node(node.body)
+      when Prism::ModuleNode
+        mod = cref.define_module(node.name)
+        mod.module!
+        with(
+          self: mod,
+          cref: mod,
+          cref_dynamic: mod,
+          locals: {}
+        ).eval_node(node.body)
+      when Prism::ProgramNode
+        eval_node(node.statements)
+      when Prism::StatementsNode
+        node.body.each { |stmt| eval_node(stmt) }
+      when Prism::StringNode
+        node.unescaped.dup.freeze
+      when Prism::SymbolNode
+        node.unescaped.to_sym
+      else
+        $stderr.puts "Dismissing node: #{node.inspect}" # rubocop:disable Style/StderrPuts
+      end
+    end
+  end
+end

data/lib/rbs_macros/library_registry.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+require "forwardable"
+
+module RbsMacros
+  # RbsMacros allow you to define a reusable set of macro definitions
+  # called a library.
+  # This is usually registered to the global singleton of LibraryRegistry
+  # like:
+  #
+  #   # my_library/rbs_macros.rb
+  #   RbsMacros::LibraryRegistry.register_macros("my_library/rbs_macros") do |macros|
+  #     macros << MyMacro1
+  #     macros << MyMacro2
+  #   end
+  class LibraryRegistry
+    extend SingleForwardable
+    def_single_delegators :@global, :register_macros
+
+    def initialize
+      @libraries = {}
+    end
+
+    def register_macros(name, macros = [], &block)
+      a = @libraries.fetch(name) { |k| @libraries[k] = [] }
+      a.push(*macros)
+      block&.(a)
+      nil
+    end
+
+    def lookup_macros(name, soft_fail: false)
+      unless @libraries.key?(name)
+        require_library(name, soft_fail:)
+        raise ArgumentError, "Unknown library: #{name}" if !@libraries.key?(name) && !soft_fail
+      end
+
+      @libraries[name] || []
+    end
+
+    def require_library(name, soft_fail: false)
+      # To be implemented by subclasses
+      raise ArgumentError, "Unknown library: #{name}" unless soft_fail
+    end
+
+    @global = LibraryRegistry.new
+    def @global.require_library(name, soft_fail: false)
+      require name
+    rescue LoadError
+      raise unless soft_fail
+    end
+  end
+end

data/lib/rbs_macros/macro.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+module RbsMacros
+  # Base class for macro implementations.
+  # Macros react to method invocations in Ruby code (usually in module/class bodies)
+  # and generate RBS declarations for them.
+  class Macro
+    def setup(env) # rubocop:disable Lint/UnusedMethodArgument
+      raise NoMethodError, "Not implemented: #{self.class}#setup"
+    end
+  end
+end

data/lib/rbs_macros/macros/forwardable.rb

@@ -0,0 +1,197 @@
+# frozen_string_literal: true
+
+require "rbs_macros"
+
+module RbsMacros
+  module Macros
+    # TODO: resolve duplication between ForwardableMacros and SingleForwardableMacros
+    # TODO: fallback to untyped even when something is wrong
+
+    # Implements macros for the `Forwardable` module.
+    class ForwardableMacros < Macro
+      def setup(env)
+        env.register_handler(:def_delegator, method(:meta_def_delegator))
+        env.register_handler(:def_instance_delegator, method(:meta_def_delegator))
+
+        env.register_handler(:def_delegators, method(:meta_def_delegators))
+        env.register_handler(:def_instance_delegators, method(:meta_def_delegators))
+      end
+
+      def meta_def_delegator(params)
+        recv = params.receiver
+        return unless recv.is_a?(MetaModule)
+
+        accessor = params.positional[0]
+        method = params.positional[1]
+        ali = params.positional[2] || method
+        return unless accessor.is_a?(Symbol) || accessor.is_a?(String)
+        return unless method.is_a?(Symbol) || method.is_a?(String)
+        return unless ali.is_a?(Symbol) || ali.is_a?(String)
+
+        builder = RBS::DefinitionBuilder.new(env: params.env.rbs)
+        self_instance = builder.build_instance(recv.rbs_type)
+        accessor_type =
+          case accessor.to_s
+          when /\A@/
+            ivar = self_instance.instance_variables[accessor.to_sym]
+            return unless ivar
+
+            ivar.type
+          else
+            # TODO
+            return
+          end
+        accessor_instance = widened_instance(accessor_type, builder:)
+        return unless accessor_instance
+
+        meth = accessor_instance[0].methods[method.to_sym]&.sub(accessor_instance[1])
+        return unless meth
+
+        params.env.add_decl(
+          RBS::AST::Members::MethodDefinition.new(
+            name: ali.to_sym,
+            kind: :instance,
+            overloads:
+              meth.defs.map do |typedef|
+                RBS::AST::Members::MethodDefinition::Overload.new(
+                  method_type: typedef.type,
+                  annotations: []
+                )
+              end,
+            annotations: [],
+            location: nil,
+            comment: nil,
+            overloading: false,
+            visibility: nil
+          ),
+          mod: recv,
+          file: params.filename
+        )
+      end
+
+      def meta_def_delegators(params)
+        recv = params.receiver
+        return unless recv.is_a?(MetaModule)
+
+        accessor = params.positional[0]
+        methods = params.positional[1..] || []
+        methods.each do |method|
+          params.env.invoke(
+            params.with(
+              name: :def_instance_delegator,
+              positional: [accessor, method],
+              keyword: {},
+              block: nil
+            )
+          )
+        end
+      end
+
+      private
+
+      def widened_instance(type, builder:)
+        case type
+        when RBS::Types::ClassInstance
+          # Using public_send because tapp_subst is declared as private although defined as public.
+          [builder.build_instance(type.name), builder.public_send(:tapp_subst, type.name, type.args)]
+        end
+      end
+    end
+
+    # Implements macros for the `SingleForwardable` module.
+    class SingleForwardableMacros < Macro
+      def setup(env)
+        env.register_handler(:def_delegator, method(:meta_def_delegator))
+        env.register_handler(:def_single_delegator, method(:meta_def_delegator))
+
+        env.register_handler(:def_delegators, method(:meta_def_delegators))
+        env.register_handler(:def_single_delegators, method(:meta_def_delegators))
+      end
+
+      def meta_def_delegator(params)
+        recv = params.receiver
+        return unless recv.is_a?(MetaModule)
+
+        accessor = params.positional[0]
+        method = params.positional[1]
+        ali = params.positional[2] || method
+        return unless accessor.is_a?(Symbol) || accessor.is_a?(String)
+        return unless method.is_a?(Symbol) || method.is_a?(String)
+        return unless ali.is_a?(Symbol) || ali.is_a?(String)
+
+        builder = RBS::DefinitionBuilder.new(env: params.env.rbs)
+        singleton = builder.build_singleton(recv.rbs_type)
+        accessor_type =
+          case accessor.to_s
+          when /\A@/
+            ivar = singleton.instance_variables[accessor.to_sym]
+            return unless ivar
+
+            ivar.type
+          else
+            # TODO
+            return
+          end
+        accessor_instance = widened_instance(accessor_type, builder:)
+        return unless accessor_instance
+
+        meth = accessor_instance[0].methods[method.to_sym]&.sub(accessor_instance[1])
+        return unless meth
+
+        params.env.add_decl(
+          RBS::AST::Members::MethodDefinition.new(
+            name: ali.to_sym,
+            kind: :singleton,
+            overloads:
+              meth.defs.map do |typedef|
+                RBS::AST::Members::MethodDefinition::Overload.new(
+                  method_type: typedef.type,
+                  annotations: []
+                )
+              end,
+            annotations: [],
+            location: nil,
+            comment: nil,
+            overloading: false,
+            visibility: nil
+          ),
+          mod: recv,
+          file: params.filename
+        )
+      end
+
+      def meta_def_delegators(params)
+        recv = params.receiver
+        return unless recv.is_a?(MetaModule)
+
+        accessor = params.positional[0]
+        methods = params.positional[1..] || []
+        methods.each do |method|
+          params.env.invoke(
+            params.with(
+              name: :def_single_delegator,
+              positional: [accessor, method],
+              keyword: {},
+              block: nil
+            )
+          )
+        end
+      end
+
+      private
+
+      def widened_instance(type, builder:)
+        case type
+        when RBS::Types::ClassInstance
+          # Using public_send because tapp_subst is declared as private although defined as public.
+          [builder.build_instance(type.name), builder.public_send(:tapp_subst, type.name, type.args)]
+        end
+      end
+    end
+  end
+end
+
+RbsMacros::LibraryRegistry.register_macros("rbs_macros/macros/forwardable") do |macros|
+  macros << RbsMacros::Macros::ForwardableMacros
+  macros << RbsMacros::Macros::SingleForwardableMacros
+end

data/lib/rbs_macros/meta_module.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+module RbsMacros
+  # Designates a module in a Ruby program being analyzed.
+  class MetaModule
+    attr_reader :env, :name, :is_class, :superclass
+
+    def initialize(env, name, is_class: nil, superclass: nil)
+      @env = env
+      @name = name
+      @is_class = is_class
+      @superclass = superclass
+      @constants = {}
+    end
+
+    def rbs_type
+      return @rbs_type if defined?(@rbs_type)
+
+      *ns, name = (@name || raise("Anonymous module given")).split("::")
+      @rbs_type = RBS::TypeName.new(
+        name: (name || raise("Anonymous module gien")).to_sym,
+        namespace: RBS::Namespace.new(
+          path: ns.map(&:to_sym),
+          absolute: true
+        )
+      )
+    end
+
+    def class!
+      @is_class = true if @is_class.nil?
+    end
+
+    def module!
+      @is_class = false if @is_class.nil?
+    end
+
+    def define_module(name)
+      @constants[name] ||= MetaModule.new(@env, child_module_name(name.to_s))
+    end
+
+    def meta_const_set(name, value)
+      @constants[name] = value
+    end
+
+    def meta_const_get(name)
+      @constants[name]
+    end
+
+    def meta_constants
+      @constants.keys
+    end
+
+    private
+
+    def child_module_name(child_name)
+      if child_name && (name && name != "Object")
+        "#{name}::#{child_name}"
+      elsif child_name
+        child_name
+      end
+    end
+  end
+
+  MetaClass = MetaModule
+end

data/lib/rbs_macros/project.rb

@@ -0,0 +1,110 @@
+# frozen_string_literal: true
+
+require "fileutils"
+require "pathname"
+
+module RbsMacros
+  # Refers to the filesystem RbsMacros operates on.
+  class AbstractProject
+    # rubocop:disable Lint/UnusedMethodArgument
+    def glob(
+      ext:,
+      include:,
+      exclude:,
+      &block
+    )
+      raise NoMethodError, "Not implemented: #{self.class}#each_rbs"
+    end
+
+    def write(path, content)
+      raise NoMethodError, "Not implemented: #{self.class}#write"
+    end
+
+    def read(path)
+      raise NoMethodError, "Not implemented: #{self.class}#read"
+    end
+    # rubocop:enable Lint/UnusedMethodArgument
+  end
+
+  # A project based on real FS.
+  class Project < AbstractProject
+    attr_accessor :base_dir
+
+    def initialize(base_dir: Pathname(Dir.pwd))
+      super()
+      @base_dir = base_dir
+    end
+
+    def glob(
+      ext:,
+      include:,
+      exclude:,
+      &block
+    )
+      return enum_for(:glob, ext:, include:, exclude:) unless block
+
+      loaded = Set.new # : Set[String]
+      include.each do |incl_dir|
+        Dir.glob(
+          "#{incl_dir}/**/*#{ext}",
+          base: @base_dir
+        ).sort.each do |path|
+          next unless File.file?(@base_dir + path)
+          next if loaded.include?(path)
+
+          loaded << path
+          is_excluded = exclude.any? do |excl_dir|
+            "#{path}/".start_with?("#{excl_dir}/")
+          end
+          block.(path) unless is_excluded
+        end
+      end
+    end
+
+    def write(path, content)
+      full_path = @base_dir + path
+      FileUtils.mkdir_p(full_path.dirname)
+      File.write(full_path.to_s, content)
+    end
+
+    def read(path)
+      File.read((@base_dir + path).to_s)
+    end
+  end
+
+  # An in-memory project.
+  class FakeProject < AbstractProject
+    def initialize
+      super
+      @files = {}
+    end
+
+    def glob(
+      ext:,
+      include:,
+      exclude:,
+      &block
+    )
+      return enum_for(:glob, ext:, include:, exclude:) unless block
+
+      @files.each_key do |path|
+        has_ext = path.end_with?(ext)
+        is_incl = include.any? do |incl_dir|
+          "#{path}/".start_with?("#{incl_dir}/")
+        end
+        is_excl = exclude.any? do |excl_dir|
+          "#{path}/".start_with?("#{excl_dir}/")
+        end
+        block.(path) if has_ext && is_incl && !is_excl
+      end
+    end
+
+    def write(path, content)
+      @files[path] = content
+    end
+
+    def read(path)
+      @files[path] || raise(Errno::ENOENT, path)
+    end
+  end
+end

data/lib/rbs_macros/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module RbsMacros
+  VERSION = "0.1.0"
+end

data/lib/rbs_macros.rb

@@ -0,0 +1,137 @@
+# frozen_string_literal: true
+
+require_relative "rbs_macros/version"
+require_relative "rbs_macros/config"
+require_relative "rbs_macros/environment"
+require_relative "rbs_macros/exec_ctx"
+require_relative "rbs_macros/library_registry"
+require_relative "rbs_macros/macro"
+require_relative "rbs_macros/meta_module"
+require_relative "rbs_macros/project"
+
+require "stringio"
+require "rbs"
+
+# RbsMacros is a utility that looks for metaprogramming-related
+# method invocation in your Ruby code and generates RBS files for them.
+module RbsMacros
+  def self.run(&block)
+    config = Config.new
+    block&.(config)
+
+    env = Environment.new
+    config.loader.load(env: env.rbs) if config.use_loader
+    config.macros.each do |macro|
+      macro.setup(env)
+    end
+
+    config.project.glob(ext: ".rbs", include: config.sigs, exclude: [config.output_dir]) do |filename|
+      source = config.project.read(filename)
+      buffer = RBS::Buffer.new(name: filename, content: source)
+      _, directives, decls = RBS::Parser.parse_signature(buffer)
+      env.rbs.add_signature(buffer:, directives:, decls:)
+    end
+    # TODO: streamline this private method invocation
+    env.instance_variable_set(:@rbs, env.rbs.resolve_type_names)
+    config.project.glob(ext: ".rb", include: config.load_dirs, exclude: []) do |filename|
+      relative_filename = filename
+      config.load_dirs.each do |load_dir|
+        if filename.start_with?("#{load_dir}/")
+          relative_filename = filename.delete_prefix("#{load_dir}/")
+          break
+        end
+      end
+      source = config.project.read(filename)
+      env.meta_eval_ruby(source, filename: relative_filename.sub(/\.rb\z/, ""))
+    end
+
+    files = {} # : Hash[String, Array[RBS::AST::Declarations::t]]
+    env.decls.each do |entry|
+      file_decls = (files[entry.file] ||= [])
+      current_mod = env.object_class
+      container = nil # : (RBS::AST::Declarations::Module | RBS::AST::Declarations::Class)?
+      (entry.mod.name || "").split("::").each do |name|
+        inner_mod = current_mod.meta_const_get(name.to_sym)
+        raise "Not found: #{current_mod.name}::#{name}" unless inner_mod
+        raise "Not a module: #{current_mod.name}::#{name}" unless inner_mod.is_a?(MetaModule)
+
+        current_mod = inner_mod
+
+        current_decls = container&.members || file_decls
+        container = nil
+        current_decls.each do |decl|
+          if (decl.is_a?(RBS::AST::Declarations::Class) || decl.is_a?(RBS::AST::Declarations::Module)) \
+            && decl.name.name == name.to_sym
+            container = decl
+          end
+        end
+        next if container
+
+        if inner_mod.is_class
+          container = c = RBS::AST::Declarations::Class.new(
+            name: RBS::TypeName.new(name: name.to_sym, namespace: RBS::Namespace.empty),
+            type_params: [],
+            super_class: nil,
+            members: [],
+            annotations: [],
+            location: nil,
+            comment: nil
+          )
+          current_decls << c
+        else
+          container = m = RBS::AST::Declarations::Module.new(
+            name: RBS::TypeName.new(name: name.to_sym, namespace: RBS::Namespace.empty),
+            type_params: [],
+            members: [],
+            self_types: [],
+            annotations: [],
+            location: nil,
+            comment: nil
+          )
+          current_decls << m
+        end
+      end
+      if container
+        container.members << entry.declaration
+      else
+        d = entry.declaration
+        case d
+        when RBS::AST::Declarations::Class,
+            RBS::AST::Declarations::Module,
+            RBS::AST::Declarations::Interface,
+            RBS::AST::Declarations::Constant,
+            RBS::AST::Declarations::Global,
+            RBS::AST::Declarations::TypeAlias,
+            RBS::AST::Declarations::ClassAlias,
+            RBS::AST::Declarations::ModuleAlias
+          file_decls << d
+        else
+          raise "Not allowed here: #{d.class}"
+        end
+      end
+    end
+
+    decls = [] # : Array[RBS::AST::Declarations::t]
+    env.object_class.meta_constants.each do |name|
+      value = env.object_class.meta_const_get(name)
+      next unless value.is_a?(MetaModule)
+
+      decls << RBS::AST::Declarations::Module.new(
+        name: RBS::TypeName.new(name:, namespace: RBS::Namespace.empty),
+        type_params: [],
+        members: [],
+        self_types: [],
+        annotations: [],
+        location: nil,
+        comment: nil
+      )
+    end
+
+    files.each do |filename, file_decls|
+      out = StringIO.new(+"", "w")
+      writer = RBS::Writer.new(out:)
+      writer.write file_decls
+      config.project.write("#{config.output_dir}/#{filename}.rbs", out.string)
+    end
+  end
+end