infer_type.rb

require 'set'

class InferType
  def self.run(&block)
    infer_type = InferType.new
    infer_type.start
    block.call
  ensure
    infer_type.finish
  end

  def initialize
    @records = {}
  end

  def start
    @trace = TracePoint.new(:call, &method(:on_call))
    @trace.enable
  end

  def finish
    @trace.disable
    report
  end

  private

  def on_call(tp)
    meth = find_method(tp)
    if meth
      meth.parameters.each do |arg_type, arg_name|
        if arg_name
          sig = build_signature(tp, meth)
          @records[sig] ||= {}
          @records[sig][arg_name] ||= Set.new
          @records[sig][arg_name].add(tp.binding.local_variable_get(arg_name).class)
        end
      end
    end
  end

  def find_method(tp)
    tp.defined_class.instance_method(tp.method_id)
  rescue NameError
    nil
  end

  def build_signature(tp, meth)
    sig =
      if tp.defined_class.singleton_class?
        attached = tp.defined_class.inspect.slice(/#<Class:(.+)>/, 1)
        "#{attached}.#{tp.method_id}"
      else
        "#{tp.defined_class}##{tp.method_id}"
      end
    loc = meth.source_location
    if loc
      "#{sig} #{loc[0]}:#{loc[1]}"
    else
      sig
    end
  end

  def report
    re = Regexp.new(ENV.fetch('INFER_TYPE_TARGET', '.'))
    @records.each do |signature, method_info|
      if re === signature
        puts signature
        method_info.each do |arg_name, klasses|
          t = infer_type(klasses.to_a)
          puts "  @param #{arg_name} [#{t}]"
        end
      end
    end
  end

  def infer_type(klasses)
    if klasses.size == 1 && klasses.member?(NilClass)
      return 'nil'
    end

    nullable = klasses.member?(NilClass)
    klasses.delete(NilClass)

    if klasses.size == 2 && klasses.member?(TrueClass) && klasses.member?(FalseClass)
      return annotate_nullable('Boolean', nullable)
    end

    common = klasses.drop(1).inject(klasses[0], &method(:common_type))
    annotate_nullable(common, nullable)
  end

  def annotate_nullable(t, nullable)
    if nullable
      "#{t}, nil"
    else
      t
    end
  end

  def common_type(t1, t2)
    top = BasicObject
    parents(t1).zip(parents(t2)) do |x, y|
      unless x.equal?(y)
        return top
      end
      top = x
    end
    top
  end

  def parents(t)
    klasses = []
    until t.equal?(BasicObject)
      klasses << t
      t = t.superclass
    end
    klasses << BasicObject
    klasses.reverse
  end
end

if $0 == __FILE__
  path = ARGV.shift
  InferType.run { load path }
end

# require_relative '../infer_type'
# RSpec.configure do |config|
#   config.before(:suite) { @infer_type = InferType.new; @infer_type.start }
#   config.after(:suite) { @infer_type.finish }
# end