Support extend

lib/typeprof/core/ast.rb

@@ -283,6 +283,8 @@ def self.create_node(raw_node, lenv, use_result = true, allow_meta = false)
           case raw_node.name
           when :include
             return IncludeMetaNode.new(raw_node, lenv)
+          when :extend
+            return ExtendMetaNode.new(raw_node, lenv)
           when :attr_reader
             return AttrReaderMetaNode.new(raw_node, lenv)
           when :attr_writer
@@ -460,6 +462,7 @@ def self.create_rbs_member(raw_decl, lenv)
       when RBS::AST::Members::Prepend
         SigPrependNode.new(raw_decl, lenv)
       when RBS::AST::Members::Extend
+        SigExtendNode.new(raw_decl, lenv)
       when RBS::AST::Members::Public
       when RBS::AST::Members::Private
       when RBS::AST::Members::Alias

lib/typeprof/core/ast/meta.rb

@@ -46,6 +46,52 @@ def install0(genv)
       end
     end
 
+    class ExtendMetaNode < Node
+      def initialize(raw_node, lenv)
+        super(raw_node, lenv)
+        # TODO: error for splat
+        @args = raw_node.arguments.arguments.map do |raw_arg|
+          next if raw_arg.is_a?(Prism::SplatNode)
+          lenv.use_strict_const_scope do
+            AST.create_node(raw_arg, lenv)
+          end
+        end.compact
+        # TODO: error for non-LIT
+        # TODO: fine-grained hover
+      end
+
+      attr_reader :args
+
+      def subnodes = { args: }
+
+      def define0(genv)
+        mod = genv.resolve_cpath(@lenv.cref.cpath)
+        @args.each do |arg|
+          arg.define(genv)
+          if arg.static_ret
+            arg.static_ret.followers << mod
+            mod.add_extend_def(genv, arg)
+          end
+        end
+      end
+
+      def undefine0(genv)
+        mod = genv.resolve_cpath(@lenv.cref.cpath)
+        @args.each do |arg|
+          if arg.static_ret
+            mod.remove_extend_def(genv, arg)
+          end
+          arg.undefine(genv)
+        end
+        super(genv)
+      end
+
+      def install0(genv)
+        @args.each {|arg| arg.install(genv) }
+        Source.new
+      end
+    end
+
     class AttrReaderMetaNode < Node
       def initialize(raw_node, lenv)
         super(raw_node, lenv)

lib/typeprof/core/ast/sig_decl.rb

@@ -318,6 +318,55 @@ def install0(genv)
       end
     end
 
+    class SigExtendNode < Node
+      def initialize(raw_decl, lenv)
+        super(raw_decl, lenv)
+        name = raw_decl.name
+        @cpath = name.namespace.path + [name.name]
+        @toplevel = name.namespace.absolute?
+        @args = raw_decl.args.map {|arg| AST.create_rbs_type(arg, lenv) }
+      end
+
+      attr_reader :cpath, :toplevel, :args
+      def subnodes = { args: }
+      def attrs = { cpath:, toplevel: }
+
+      def define0(genv)
+        @args.each {|arg| arg.define(genv) }
+        const_reads = []
+        const_read = BaseConstRead.new(genv, @cpath.first, @toplevel ? CRef::Toplevel : @lenv.cref, true)
+        const_reads << const_read
+        @cpath[1..].each do |cname|
+          const_read = ScopedConstRead.new(cname, const_read, true)
+          const_reads << const_read
+        end
+        mod = genv.resolve_cpath(@lenv.cref.cpath)
+        const_read.followers << mod
+        mod.add_extend_decl(genv, self)
+        const_reads
+      end
+
+      def define_copy(genv)
+        mod = genv.resolve_cpath(@lenv.cref.cpath)
+        mod.add_extend_decl(genv, self)
+        mod.remove_extend_decl(genv, @prev_node)
+        super(genv)
+      end
+
+      def undefine0(genv)
+        mod = genv.resolve_cpath(@lenv.cref.cpath)
+        mod.remove_extend_decl(genv, self)
+        @static_ret.each do |const_read|
+          const_read.destroy(genv)
+        end
+        @args.each {|arg| arg.undefine(genv) }
+      end
+
+      def install0(genv)
+        Source.new
+      end
+    end
+
     class SigAliasNode < Node
       def initialize(raw_decl, lenv)
         super(raw_decl, lenv)

lib/typeprof/core/env.rb

@@ -80,7 +80,7 @@ def each_superclass(mod, singleton, &blk)
         # TODO: prepended modules
         yield mod, singleton
         if singleton
-          # TODO: extended modules
+          each_extended_module(mod, &blk)
         else
           each_included_module(mod, &blk)
         end
@@ -95,6 +95,15 @@ def each_included_module(mod, &blk)
       end
     end
 
+    def each_extended_module(mod, &blk)
+      # An extended module contributes its instance methods (singleton = false)
+      # as the receiver's singleton methods.
+      mod.extended_modules.each do |_ext_decl, ext_mod|
+        yield ext_mod, false
+        each_included_module(ext_mod, &blk)
+      end
+    end
+
     def get_superclass(singleton, mod)
       super_mod = mod.superclass
       if super_mod

lib/typeprof/core/env/module_entity.rb

@@ -9,6 +9,8 @@ def initialize(cpath, outer_module = self)
       @include_defs = Set.empty
       @prepend_decls = []
       @prepend_defs = []
+      @extend_decls = Set.empty
+      @extend_defs = Set.empty
 
       # `class Foo = Bar` / `module Foo = Bar` declarations attached to this entity.
       # Maps an alias decl to the target ModuleEntity at the time of registration.
@@ -23,6 +25,7 @@ def initialize(cpath, outer_module = self)
       @self_types = {}
       @included_modules = {}
       @prepended_modules = {}
+      @extended_modules = {}
       @basic_object = @cpath == [:BasicObject]
 
       # child modules (subclasses and all modules that include me)
@@ -57,6 +60,7 @@ def initialize(cpath, outer_module = self)
     attr_reader :self_types
     attr_reader :included_modules
     attr_reader :prepended_modules
+    attr_reader :extended_modules
     attr_reader :child_modules
 
     attr_reader :superclass_type_args
@@ -219,6 +223,26 @@ def remove_include_def(genv, node)
       genv.add_static_eval_queue(:parent_modules_changed, self)
     end
 
+    def add_extend_decl(genv, node)
+      @extend_decls << node
+      genv.add_static_eval_queue(:parent_modules_changed, self)
+    end
+
+    def remove_extend_decl(genv, node)
+      @extend_decls.delete(node) || raise
+      genv.add_static_eval_queue(:parent_modules_changed, self)
+    end
+
+    def add_extend_def(genv, node)
+      @extend_defs << node
+      genv.add_static_eval_queue(:parent_modules_changed, self)
+    end
+
+    def remove_extend_def(genv, node)
+      @extend_defs.delete(node) || raise
+      genv.add_static_eval_queue(:parent_modules_changed, self)
+    end
+
     def add_prepend_decl(genv, node)
       @prepend_decls << node
       genv.add_static_eval_queue(:parent_modules_changed, self)
@@ -396,6 +420,30 @@ def on_parent_modules_changed(genv)
           any_updated = true
         end
       end
+      @extend_decls.each do |edecl|
+        new_parent_cpath = edecl.static_ret.last.cpath
+        new_parent, updated = update_parent(genv, edecl, @extended_modules[edecl], new_parent_cpath)
+        if updated
+          if new_parent
+            @extended_modules[edecl] = new_parent
+          else
+            @extended_modules.delete(edecl) || raise
+          end
+          any_updated = true
+        end
+      end
+      @extend_defs.each do |edef|
+        new_parent_cpath = edef.static_ret ? edef.static_ret.cpath : nil
+        new_parent, updated = update_parent(genv, edef, @extended_modules[edef], new_parent_cpath)
+        if updated
+          if new_parent
+            @extended_modules[edef] = new_parent
+          else
+            @extended_modules.delete(edef) || raise
+          end
+          any_updated = true
+        end
+      end
       @included_modules.delete_if do |origin, old_mod|
         if @include_decls.include?(origin) || @include_defs.include?(origin)
           false
@@ -414,6 +462,15 @@ def on_parent_modules_changed(genv)
           true
         end
       end
+      @extended_modules.delete_if do |origin, old_mod|
+        if @extend_decls.include?(origin) || @extend_defs.include?(origin)
+          false
+        else
+          _new_parent, updated = update_parent(genv, origin, old_mod, nil)
+          any_updated ||= updated
+          true
+        end
+      end
 
       if any_updated
         @subclass_checks.each do |mcall_box|

lib/typeprof/core/graph/box.rb

@@ -1158,7 +1158,10 @@ def resolve(genv, changes, &blk)
           skip = false
 
           if ty.is_a?(Type::Singleton)
-            # TODO: extended modules
+            # Check extended modules (their instance methods are singleton methods here)
+            break if resolve_extended_modules(genv, changes, base_ty_env, ty, mid) do |me, ty, mid|
+              yield me, ty, mid, orig_ty
+            end
           else
             # Finally check included modules
             break if resolve_included_modules(genv, changes, base_ty_env, ty, mid) do |me, ty, mid|
@@ -1258,6 +1261,38 @@ def resolve_included_modules(genv, changes, base_ty_env, ty, mid, &blk)
       found
     end
 
+    def resolve_extended_modules(genv, changes, base_ty_env, ty, mid, &blk)
+      found = false
+
+      alias_limit = 0
+      # An extended module's instance methods are resolved as the receiver's
+      # singleton methods, so look them up with singleton = false.
+      ty.mod.extended_modules.each do |ext_decl, ext_mod|
+        if ext_decl.is_a?(AST::SigExtendNode) && ext_mod.type_params
+          ext_ty = genv.get_instance_type(ext_mod, ext_decl.args, changes, base_ty_env, ty)
+        else
+          type_params = ext_mod.type_params.map { Source.new() } # TODO: better support
+          ext_ty = Type::Instance.new(genv, ext_mod, type_params)
+        end
+
+        me = ext_ty.mod.get_method(false, mid)
+        changes.add_depended_method_entity(me) if changes
+        if !me.aliases.empty?
+          mid = me.aliases.values.first
+          alias_limit += 1
+          redo if alias_limit < 5
+        end
+        if me.exist?
+          found = true
+          yield me, ext_ty, mid
+        else
+          # The extended module may itself include other modules.
+          found ||= resolve_included_modules(genv, changes, base_ty_env, ext_ty, mid, &blk)
+        end
+      end
+      found
+    end
+
     def resolve_subclasses(genv, changes)
       # TODO: This does not follow new subclasses
       @recv.each_type do |ty|

lib/typeprof/core/service.rb

@@ -497,6 +497,11 @@ def dump_declarations(path)
                   out << "  " * stack.size + "include #{ inc_mod.show_cpath }"
                 end
               end
+              mod.extended_modules.each do |ext_def, ext_mod|
+                if ext_def.is_a?(AST::ConstantReadNode) && ext_def.lenv.path == path
+                  out << "  " * stack.size + "extend #{ ext_mod.show_cpath }"
+                end
+              end
               # Output method definitions from meta nodes (StructNewNode etc.)
               node.boxes(:mdef) do |mdef|
                 out << "  " * stack.size + "def #{ mdef.singleton ? "self." : "" }#{ mdef.mid }: " + mdef.show(@options[:output_parameter_names])

scenario/incremental/extend.rb

@@ -0,0 +1,40 @@
+## update
+module Helpers
+  def shout(s) = s.upcase
+end
+
+class App
+end
+
+App.shout("hi")
+
+## assert
+module Helpers
+  def shout: (untyped) -> untyped
+end
+class App
+end
+
+## diagnostics
+(8,4)-(8,9): undefined method: singleton(App)#shout
+
+## update
+module Helpers
+  def shout(s) = s.upcase
+end
+
+class App
+  extend Helpers
+end
+
+App.shout("hi")
+
+## assert
+module Helpers
+  def shout: (String) -> String
+end
+class App
+  extend Helpers
+end
+
+## diagnostics

scenario/rbs/extend.rb

@@ -0,0 +1,19 @@
+## update: test.rbs
+module Helpers
+  def shout: (String) -> String
+end
+class App
+  extend Helpers
+end
+
+## update: test.rb
+def check
+  App.shout("hi")
+end
+
+## assert: test.rb
+class Object
+  def check: -> String
+end
+
+## diagnostics: test.rb