steep 1.0.2 → 1.1.0.pre.1

data/lib/steep/type_inference/multiple_assignment.rb

@@ -0,0 +1,189 @@
+module Steep
+  module TypeInference
+    class MultipleAssignment
+      Assignments = _ = Struct.new(:rhs_type, :optional, :leading_assignments, :trailing_assignments, :splat_assignment, keyword_init: true) do
+        # @implements Assignments
+
+        def each(&block)
+          if block
+            leading_assignments.each(&block)
+            if sp = splat_assignment
+              yield sp
+            end
+            trailing_assignments.each(&block)
+          else
+            enum_for :each
+          end
+        end
+      end
+
+      def expand(mlhs, rhs_type, optional)
+        lhss = mlhs.children
+
+        case rhs_type
+        when AST::Types::Tuple
+          expand_tuple(lhss.dup, rhs_type, rhs_type.types.dup, optional)
+        when AST::Types::Name::Instance
+          if AST::Builtin::Array.instance_type?(rhs_type)
+            expand_array(lhss.dup, rhs_type, optional)
+          end
+        when AST::Types::Any
+          expand_any(lhss, rhs_type, AST::Builtin.any_type, optional)
+        end
+      end
+
+      def expand_tuple(lhss, rhs_type, tuples, optional)
+        # @type var leading_assignments: Array[node_type_pair]
+        leading_assignments = []
+        # @type var trailing_assignments: Array[node_type_pair]
+        trailing_assignments = []
+        # @type var splat_assignment: node_type_pair?
+        splat_assignment = nil
+
+        while !lhss.empty?
+          first = lhss.first or raise
+
+          case
+          when first.type == :splat
+            break
+          else
+            leading_assignments << [first, tuples.first || AST::Builtin.nil_type]
+            lhss.shift
+            tuples.shift
+          end
+        end
+
+        while !lhss.empty?
+          last = lhss.last or raise
+
+          case
+          when last.type == :splat
+            break
+          else
+            trailing_assignments << [last, tuples.last || AST::Builtin.nil_type]
+            lhss.pop
+            tuples.pop
+          end
+        end
+
+        case lhss.size
+        when 0
+          # nop
+        when 1
+          splat_assignment = [lhss.first || raise, AST::Types::Tuple.new(types: tuples)]
+        else
+          raise
+        end
+
+        Assignments.new(
+          rhs_type: rhs_type,
+          optional: optional,
+          leading_assignments: leading_assignments,
+          trailing_assignments: trailing_assignments,
+          splat_assignment: splat_assignment
+        )
+      end
+
+      def expand_array(lhss, rhs_type, optional)
+        element_type = rhs_type.args[0] or raise
+
+        # @type var leading_assignments: Array[node_type_pair]
+        leading_assignments = []
+        # @type var trailing_assignments: Array[node_type_pair]
+        trailing_assignments = []
+        # @type var splat_assignment: node_type_pair?
+        splat_assignment = nil
+
+        while !lhss.empty?
+          first = lhss.first or raise
+
+          case
+          when first.type == :splat
+            break
+          else
+            leading_assignments << [first, AST::Builtin.optional(element_type)]
+            lhss.shift
+          end
+        end
+
+        while !lhss.empty?
+          last = lhss.last or raise
+
+          case
+          when last.type == :splat
+            break
+          else
+            trailing_assignments << [last, AST::Builtin.optional(element_type)]
+            lhss.pop
+          end
+        end
+
+        case lhss.size
+        when 0
+          # nop
+        when 1
+          splat_assignment = [
+            lhss.first || raise,
+            AST::Builtin::Array.instance_type(element_type)
+          ]
+        else
+          raise
+        end
+
+        Assignments.new(
+          rhs_type: rhs_type,
+          optional: optional,
+          leading_assignments: leading_assignments,
+          trailing_assignments: trailing_assignments,
+          splat_assignment: splat_assignment
+        )
+      end
+
+      def expand_any(nodes, rhs_type, element_type, optional)
+        # @type var leading_assignments: Array[node_type_pair]
+        leading_assignments = []
+        # @type var trailing_assignments: Array[node_type_pair]
+        trailing_assignments = []
+        # @type var splat_assignment: node_type_pair?
+        splat_assignment = nil
+
+        array = leading_assignments
+
+        nodes.each do |node|
+          case node.type
+          when :splat
+            splat_assignment = [node, AST::Builtin::Array.instance_type(element_type)]
+            array = trailing_assignments
+          else
+            array << [node, element_type]
+          end
+        end
+
+        Assignments.new(
+          rhs_type: rhs_type,
+          optional: optional,
+          leading_assignments: leading_assignments,
+          trailing_assignments: trailing_assignments,
+          splat_assignment: splat_assignment
+        )
+      end
+
+      def hint_for_mlhs(mlhs, env)
+        case mlhs.type
+        when :mlhs
+          types = mlhs.children.map do |node|
+            hint_for_mlhs(node, env) or return
+          end
+          AST::Types::Tuple.new(types: types)
+        when :lvasgn, :ivasgn, :gvasgn
+          name = mlhs.children[0]
+          env[name] || AST::Builtin.any_type
+        when :splat
+          return
+        else
+          return
+        end
+      end
+    end
+  end
+end

data/lib/steep/type_inference/type_env.rb

@@ -1,169 +1,320 @@
 module Steep
   module TypeInference
     class TypeEnv
-      attr_reader :subtyping
-      attr_reader :const_types
-      attr_reader :gvar_types
-      attr_reader :ivar_types
-      attr_reader :const_env
-
-      def initialize(subtyping:, const_env:)
-        @subtyping = subtyping
-        @const_types = {}
-        @gvar_types = {}
-        @ivar_types = {}
-        @const_env = const_env
-      end
-
-      def initialize_copy(other)
-        @subtyping = other.subtyping
-        @const_types = other.const_types.dup
-        @gvar_types = other.gvar_types.dup
-        @ivar_types = other.ivar_types.dup
-        @const_env = other.const_env
-      end
-
-      def self.build(annotations:, signatures:, subtyping:, const_env:)
-        new(subtyping: subtyping, const_env: const_env).tap do |env|
-          annotations.ivar_types.each do |name, type|
-            env.set(ivar: name, type: type)
-          end
-          annotations.const_types.each do |name, type|
-            env.set(const: name, type: type)
-          end
-          signatures.global_decls.each do |name, entry|
-            type = entry.decl.type
-            env.set(gvar: name, type: subtyping.factory.type(type))
+      include NodeHelper
+
+      attr_reader :local_variable_types
+      attr_reader :instance_variable_types, :global_types, :constant_types
+      attr_reader :constant_env
+      attr_reader :pure_method_calls
+
+      def to_s
+        array = []
+
+        local_variable_types.each do |name, entry|
+          if enforced_type = entry[1]
+            array << "#{name}: #{entry[0].to_s} <#{enforced_type.to_s}>"
+          else
+            array << "#{name}: #{entry[0].to_s}"
           end
         end
-      end
 
-      def with_annotations(ivar_types: {}, const_types: {}, gvar_types: {}, self_type:, instance_type:, class_type:, &block)
-        dup.tap do |env|
-          merge!(original_env: env.ivar_types, override_env: ivar_types, self_type: self_type, instance_type: instance_type, class_type: class_type, &block)
-          merge!(original_env: env.gvar_types, override_env: gvar_types, self_type: self_type, instance_type: instance_type, class_type: class_type, &block)
+        instance_variable_types.each do |name, type|
+          array << "#{name}: #{type.to_s}"
+        end
 
-          const_types.each do |name, annotated_type|
-            env.const_types[name] = annotated_type
-          end
+        global_types.each do |name, type|
+          array << "#{name}: #{type.to_s}"
+        end
+
+        constant_types.each do |name, type|
+          array << "#{name}: #{type.to_s}"
         end
+
+        pure_method_calls.each do |node, pair|
+          call, type = pair
+          array << "`#{node.loc.expression.source.lines[0]}`: #{type || call.return_type}"
+        end
+
+        "{ #{array.join(", ")} }"
+      end
+
+      def initialize(constant_env, local_variable_types: {}, instance_variable_types: {}, global_types: {}, constant_types: {}, pure_method_calls: {})
+        @constant_env = constant_env
+        @local_variable_types = local_variable_types
+        @instance_variable_types = instance_variable_types
+        @global_types = global_types
+        @constant_types = constant_types
+        @pure_method_calls = pure_method_calls
+
+        @pure_node_descendants = {}
+      end
+
+      def update(local_variable_types: self.local_variable_types, instance_variable_types: self.instance_variable_types, global_types: self.global_types, constant_types: self.constant_types, pure_method_calls: self.pure_method_calls)
+        TypeEnv.new(
+          constant_env,
+          local_variable_types: local_variable_types,
+          instance_variable_types: instance_variable_types,
+          global_types: global_types,
+          constant_types: constant_types,
+          pure_method_calls: pure_method_calls
+        )
+      end
+
+      def merge(local_variable_types: {}, instance_variable_types: {}, global_types: {}, constant_types: {}, pure_method_calls: {})
+        local_variable_types = self.local_variable_types.merge(local_variable_types)
+        instance_variable_types = self.instance_variable_types.merge(instance_variable_types)
+        global_types = self.global_types.merge(global_types)
+        constant_types = self.constant_types.merge(constant_types)
+        pure_method_calls = self.pure_method_calls.merge(pure_method_calls)
+
+        TypeEnv.new(
+          constant_env,
+          local_variable_types: local_variable_types,
+          instance_variable_types: instance_variable_types,
+          global_types:  global_types,
+          constant_types: constant_types,
+          pure_method_calls: pure_method_calls
+        )
       end
 
-      # @type method assert: (const: TypeName) { () -> void } -> AST::Type
-      #                    | (gvar: Symbol) { () -> void } -> AST::Type
-      #                    | (ivar: Symbol) { () -> void } -> AST::Type
-      def get(const: nil, gvar: nil, ivar: nil)
-        case
-        when const
-          if const_types.key?(const)
-            const_types[const]
+      def [](name)
+        case name
+        when Symbol
+          case
+          when local_variable_name?(name)
+            local_variable_types[name]&.[](0)
+          when instance_variable_name?(name)
+            instance_variable_types[name]
+          when global_name?(name)
+            global_types[name]
           else
-            yield
-            AST::Types::Any.new
+            raise "Unexpected variable name: #{name}"
           end
-        else
-          lookup_dictionary(ivar: ivar, gvar: gvar) do |var_name, dictionary|
-            if dictionary.key?(var_name)
-              dictionary[var_name]
-            else
-              yield
-              AST::Types::Any.new
+        when Parser::AST::Node
+          case name.type
+          when :lvar
+            self[name.children[0]]
+          when :send
+            if (call, type = pure_method_calls[name])
+              type || call.return_type
             end
           end
         end
       end
 
-      def merge_const_types(env)
-        env.const_types.each do |name, type|
-          const_types[name] = type unless const_types.key?(name)
+      def enforced_type(name)
+        local_variable_types[name]&.[](1)
+      end
+
+      def assign_local_variables(assignments)
+        local_variable_types = {}
+        invalidated_nodes = Set[]
+
+        assignments.each do |name, new_type|
+          local_variable_types[name] = [new_type, enforced_type(name)]
+          invalidated_nodes.merge(invalidated_pure_nodes(::Parser::AST::Node.new(:lvar, [name])))
         end
+
+        invalidation = pure_node_invalidation(invalidated_nodes)
+
+        merge(
+          local_variable_types: local_variable_types,
+          pure_method_calls: invalidation
+        )
       end
 
-      def set(const: nil, gvar: nil, ivar: nil, type:)
-        case
-        when const
-          const_types[const] = type
-        else
-          lookup_dictionary(ivar: ivar, gvar: gvar) do |var_name, dictionary|
-            dictionary[var_name] = type
-          end
+      def assign_local_variable(name, var_type, enforced_type)
+        merge(
+          local_variable_types: { name => [enforced_type || var_type, enforced_type] },
+          pure_method_calls: pure_node_invalidation(invalidated_pure_nodes(::Parser::AST::Node.new(:lvar, [name])))
+        )
+      end
+
+      def refine_types(local_variable_types: {}, pure_call_types: {})
+        local_variable_updates = {}
+
+        local_variable_types.each do |name, type|
+          local_variable_updates[name] = [type, enforced_type(name)]
+        end
+
+        invalidated_nodes = Set.new(pure_call_types.each_key)
+        local_variable_types.each_key do |name|
+          invalidated_nodes.merge(invalidated_pure_nodes(Parser::AST::Node.new(:lvar, [name])))
+        end
+
+        pure_call_updates = pure_node_invalidation(invalidated_nodes)
+
+        pure_call_types.each do |node, type|
+          call, _ = pure_call_updates[node]
+          pure_call_updates[node] = [call, type]
         end
+
+        merge(local_variable_types: local_variable_updates, pure_method_calls: pure_call_updates)
       end
 
-      # @type method assign: (const: TypeName, type: AST::Type) { (Subtyping::Result::Failure | nil) -> void } -> AST::Type
-      #                    | (gvar: Symbol, type: AST::Type) { (Subtyping::Result::Failure | nil) -> void } -> AST::Type
-      #                    | (ivar: Symbol, type: AST::Type) { (Subtyping::Result::Failure | nil) -> void } -> AST::Type
-      def assign(gvar: nil, ivar: nil, type:, self_type:, instance_type:, class_type:, &block)
-        lookup_dictionary(ivar: ivar, gvar: gvar) do |var_name, dictionary|
-          if dictionary.key?(var_name)
-            assert_assign(
-              var_type: dictionary[var_name],
-              lhs_type: type,
-              self_type: self_type,
-              instance_type: instance_type,
-              class_type: class_type,
-              &block
-            )
+      def constant(arg1, arg2)
+        if arg1.is_a?(RBS::TypeName) && arg2.is_a?(Symbol)
+          constant_env.resolve_child(arg1, arg2)
+        elsif arg1.is_a?(Symbol)
+          if arg2
+            constant_env.toplevel(arg1)
           else
-            yield nil
-            AST::Types::Any.new
+            constant_env.resolve(arg1)
           end
         end
       end
 
-      def lookup_dictionary(ivar:, gvar:)
-        case
-        when ivar
-          yield ivar, ivar_types
-        when gvar
-          yield gvar, gvar_types
+      def annotated_constant(name)
+        constant_types[name]
+      end
+
+      def pin_local_variables(names)
+        names = Set.new(names) if names
+
+        local_variable_types.each.with_object({}) do |pair, hash|
+          name, entry = pair
+
+          if names.nil? || names.include?(name)
+            type, enforced_type = entry
+            unless enforced_type
+              hash[name] = [type, type]
+            end
+          end
         end
       end
 
-      def assert_assign(var_type:, lhs_type:, self_type:, instance_type:, class_type:)
-        return var_type if var_type == lhs_type
+      def unpin_local_variables(names)
+        names = Set.new(names) if names
 
-        var_type = subtyping.expand_alias(var_type)
-        lhs_type = subtyping.expand_alias(lhs_type)
+        local_var_types = local_variable_types.each.with_object({}) do |pair, hash|
+          name, entry = pair
+
+          if names.nil? || names.include?(name)
+            type, _ = entry
+            hash[name] = [type, nil]
+          end
+        end
 
-        relation = Subtyping::Relation.new(sub_type: lhs_type, super_type: var_type)
-        constraints = Subtyping::Constraints.new(unknowns: Set.new)
+        merge(local_variable_types: local_var_types)
+      end
+
+      def subst(s)
+        update(
+          local_variable_types: local_variable_types.transform_values do |entry|
+            # @type block: local_variable_entry
+
+            type, enforced_type = entry
+            [
+              type.subst(s),
+              enforced_type&.yield_self {|ty| ty.subst(s) }
+            ]
+          end
+        )
+      end
 
-        subtyping.check(relation, self_type: self_type, constraints: constraints, instance_type: instance_type, class_type: class_type).else do |result|
-          yield result
+      def join(*envs)
+        # @type var all_lvar_types: Hash[Symbol, Array[AST::Types::t]]
+        all_lvar_types = envs.each_with_object({}) do |env, hash|
+          env.local_variable_types.each_key do |name|
+            hash[name] = []
+          end
+        end
+
+        envs.each do |env|
+          all_lvar_types.each_key do |name|
+            all_lvar_types[name] << (env[name] || AST::Builtin.nil_type)
+          end
+        end
+
+        assignments =
+          all_lvar_types
+            .transform_values {|types| AST::Types::Union.build(types: types) }
+            .reject {|var, type| self[var] == type }
+
+        common_pure_nodes = envs
+          .map {|env| Set.new(env.pure_method_calls.each_key) }
+          .inject(Set.new(pure_method_calls.each_key)) {|s1, s2| s1.intersection(s2) }
+
+        pure_call_updates = common_pure_nodes.each_with_object({}) do |node, hash|
+          pairs = envs.map {|env| env.pure_method_calls[node] }
+          refined_type = AST::Types::Union.build(types: pairs.map {|pair| pair[1] || pair[0].return_type })
+          call, _ = (pure_method_calls[node] or raise)
+
+          hash[node] = [call, refined_type]
         end
 
-        var_type
+        assign_local_variables(assignments).merge(pure_method_calls: pure_call_updates)
       end
 
-      def merge!(original_env:, override_env:, self_type:, instance_type:, class_type:, &block)
-        original_env.merge!(override_env) do |name, original_type, override_type|
-          assert_annotation(
-            name,
-            annotated_type: override_type,
-            original_type: original_type,
-            self_type: self_type,
-            instance_type: instance_type,
-            class_type: class_type,
-            &block
-          )
+      def add_pure_call(node, call, type)
+        if (c, _ = pure_method_calls[node]) && c == call
+          return self
+        end
+
+        update =
+          pure_node_invalidation(invalidated_pure_nodes(node))
+            .merge!({ node => [call, type] })
+
+        merge(pure_method_calls: update)
+      end
+
+      def replace_pure_call_type(node, type)
+        if (call, _ = pure_method_calls[node])
+          calls = pure_method_calls.dup
+          calls[node] = [call, type]
+          update(pure_method_calls: calls)
+        else
+          raise
         end
       end
 
-      def assert_annotation(name, annotated_type:, original_type:, self_type:, instance_type:, class_type:)
-        return annotated_type if annotated_type == original_type
+      def invalidate_pure_node(node)
+        merge(pure_method_calls: pure_node_invalidation(invalidated_pure_nodes(node)))
+      end
+
+      def pure_node_invalidation(invalidated_nodes)
+        # @type var invalidation: Hash[Parser::AST::Node, [MethodCall::Typed, AST::Types::t?]]
+        invalidation = {}
+
+        invalidated_nodes.each do |node|
+          if (call, _ = pure_method_calls[node])
+            invalidation[node] = [call, nil]
+          end
+        end
 
-        annotated_type = subtyping.expand_alias(annotated_type)
-        original_type = subtyping.expand_alias(original_type)
+        invalidation
+      end
 
-        relation = Subtyping::Relation.new(sub_type: annotated_type, super_type: original_type)
-        constraints = Subtyping::Constraints.new(unknowns: Set.new)
+      def invalidated_pure_nodes(invalidated_node)
+        invalidated_nodes = Set[]
 
-        subtyping.check(relation, constraints: constraints, self_type: self_type, instance_type: instance_type, class_type: class_type).else do |result|
-          yield name, relation, result
+        pure_method_calls.each_key do |pure_node|
+          descendants = @pure_node_descendants[pure_node] ||= each_descendant_node(pure_node).to_set
+          if descendants.member?(invalidated_node)
+            invalidated_nodes << pure_node
+          end
         end
 
-        annotated_type
+        invalidated_nodes
+      end
+
+      def local_variable_name?(name)
+        name.start_with?(/[a-z_]/)
+      end
+
+      def instance_variable_name?(name)
+        name.start_with?(/@[^@]/)
+      end
+
+      def global_name?(name)
+        name.start_with?('$')
+      end
+
+      def inspect
+        s = "#<%s:%#018x " % [self.class, object_id]
+        s << instance_variables.map(&:to_s).sort.map {|name| "#{name}=..." }.join(", ")
+        s + ">"
       end
     end
   end