rigortype 0.1.8 → 0.1.10

data/lib/rigor/inference/def_return_typer.rb

@@ -0,0 +1,98 @@
+# frozen_string_literal: true
+
+require "prism"
+
+require_relative "../type"
+
+module Rigor
+  module Inference
+    # Returns the inferred return type of a `Prism::DefNode`, or nil
+    # when no type can be derived (empty body, scope-lookup miss,
+    # or any failure during inference — caller surfaces "no
+    # annotation" on a nil).
+    #
+    # The inferred type is the union of:
+    #
+    # - the body's last-statement type, and
+    # - the type of every explicit `return value` reachable in the
+    #   body. Nested `def` / lambda / block bodies are return
+    #   barriers — their `return`s do not bubble up to the enclosing
+    #   method.
+    #
+    # Extracted from `Rigor::SigGen::Generator#infer_return_type` so
+    # `LineTypeCollector` (`rigor annotate`'s def-line annotator)
+    # and the sig-generator share one source of truth.
+    module DefReturnTyper
+      RETURN_BARRIER_NODES = [Prism::DefNode, Prism::LambdaNode, Prism::BlockNode].freeze
+      private_constant :RETURN_BARRIER_NODES
+
+      module_function
+
+      def call(def_node, scope_index)
+        body = def_node.body
+        return nil if body.nil?
+
+        last = body_last_expression(body)
+        return nil if last.nil?
+
+        inner_scope = scope_index[last] || scope_index[body] || scope_index[def_node]
+        return nil if inner_scope.nil?
+
+        last_type = safe_type_of(inner_scope, last)
+        return nil if last_type.nil?
+
+        union_with_explicit_returns(body, last_type, scope_index)
+      rescue StandardError
+        nil
+      end
+
+      def body_last_expression(body)
+        case body
+        when Prism::StatementsNode then body.body.last
+        when Prism::BeginNode then body_last_expression(body.statements)
+        else body
+        end
+      end
+
+      def union_with_explicit_returns(body, last_type, scope_index)
+        return_types = []
+        collect_return_types(body, scope_index, return_types)
+        return last_type if return_types.empty?
+
+        Type::Combinator.union(last_type, *return_types)
+      end
+
+      def collect_return_types(node, scope_index, out)
+        return unless node.is_a?(Prism::Node)
+        return if RETURN_BARRIER_NODES.any? { |klass| node.is_a?(klass) }
+
+        type_return_node(node, scope_index, out) if node.is_a?(Prism::ReturnNode)
+        node.compact_child_nodes.each { |c| collect_return_types(c, scope_index, out) }
+      end
+
+      def type_return_node(return_node, scope_index, out)
+        args = return_node.arguments&.arguments || []
+        if args.empty?
+          out << Type::Combinator.constant_of(nil)
+          return
+        end
+
+        scope = scope_index[return_node] || scope_index[args.first]
+        return if scope.nil?
+        # `return a, b` packs into a Tuple at runtime; the MVP only
+        # handles the single-value form. Multi-arg returns
+        # contribute no type to keep the implementation focused.
+        return unless args.size == 1
+
+        type = safe_type_of(scope, args.first)
+        out << type unless type.nil?
+      end
+
+      def safe_type_of(scope, node)
+        scope.type_of(node)
+      rescue StandardError
+        nil
+      end
+    end
+  end
+end

data/lib/rigor/inference/expression_typer.rb

@@ -8,6 +8,7 @@ require_relative "block_parameter_binder"
 require_relative "fallback"
 require_relative "macro_block_self_type"
 require_relative "method_dispatcher"
+require_relative "narrowing"
 
 module Rigor
   module Inference
@@ -660,16 +661,26 @@ module Rigor
       end
 
       # Returns `:truthy`, `:falsey`, or `nil` for an arbitrary
-      # predicate expression. Only `Type::Constant` answers
-      # decisively — `Union[true, false]`, `Nominal[bool]`, and
-      # `Dynamic[T]` keep both branches live.
+      # predicate expression under three-valued logic. Uses the
+      # same {Narrowing} probe as `StatementEvaluator#eval_if`:
+      # the predicate is truthy when its falsey fragment is `Bot`,
+      # falsey when its truthy fragment is `Bot`. So
+      # `Nominal[Integer]` (always truthy in Ruby), `Constant[nil]`,
+      # and `Constant[false]` fold one branch; `Union[true, false]`,
+      # `Dynamic[T]`, and `Top` keep both branches live.
       def constant_predicate_polarity(predicate)
         return nil if predicate.nil?
 
         type = type_of(predicate)
-        return nil unless type.is_a?(Type::Constant)
+        return nil if type.nil? || type.is_a?(Type::Bot)
 
-        type.value ? :truthy : :falsey
+        truthy_bot = Narrowing.narrow_truthy(type).is_a?(Type::Bot)
+        falsey_bot = Narrowing.narrow_falsey(type).is_a?(Type::Bot)
+
+        return :falsey if truthy_bot && !falsey_bot
+        return :truthy if !truthy_bot && falsey_bot
+
+        nil
       end
 
       def type_of_else(node)
@@ -712,15 +723,135 @@ module Rigor
         type.value ? :truthy : :falsey
       end
 
+      # Three-valued evaluation of `case predicate when pattern`
+      # dispatch. For each `when` clause we ask: under static types,
+      # does `pattern === predicate` definitely match (`:yes`),
+      # definitely not match (`:no`), or possibly match (`:maybe`)?
+      # Walking in source order:
+      #
+      # - `:yes` — this branch fires, subsequent branches are
+      #   unreachable. Result = union(prior `:maybe` branches, this
+      #   `:yes` branch).
+      # - `:no`  — branch dropped.
+      # - `:maybe` — branch is a candidate, continue.
+      #
+      # If no `:yes` was reached, the else clause (or `Constant[nil]`
+      # when absent) is added to the candidate set.
+      #
+      # The `case ... in` pattern-matching form (`CaseMatchNode`) and
+      # the predicate-less form (`case; when c1; ...`) bypass the
+      # `===` analysis: pattern matching has richer semantics, and a
+      # predicate-less `case` reduces to a `if c1; ...; elsif c2`
+      # chain that statement-level narrowing already handles.
       def type_of_case(node)
-        branch_types = node.conditions.map { |branch| type_of(branch) }
-        else_type =
-          if node.else_clause
-            type_of(node.else_clause)
-          else
-            Type::Combinator.constant_of(nil)
+        return type_of_case_simple_union(node) if node.is_a?(Prism::CaseMatchNode) || node.predicate.nil?
+
+        subject_type = type_of(node.predicate)
+        candidates = []
+        reached_yes = false
+
+        node.conditions.each do |when_node|
+          case case_when_branch_certainty(subject_type, when_node)
+          when :yes
+            candidates << type_of(when_node)
+            reached_yes = true
+            break
+          when :maybe
+            candidates << type_of(when_node)
+            # :no — drop the branch
           end
-        Type::Combinator.union(*branch_types, else_type)
+        end
+
+        candidates << type_of_case_else(node) unless reached_yes
+        Type::Combinator.union(*candidates)
+      end
+
+      def type_of_case_simple_union(node)
+        branch_types = node.conditions.map { |branch| type_of(branch) }
+        Type::Combinator.union(*branch_types, type_of_case_else(node))
+      end
+
+      def type_of_case_else(node)
+        return Type::Combinator.constant_of(nil) if node.else_clause.nil?
+
+        type_of(node.else_clause)
+      end
+
+      # Combines per-pattern certainty across a `when` clause's
+      # conditions (`when a, b, c` ≡ `a === s || b === s || c === s`).
+      # `:yes` if any pattern is `:yes`; `:no` if all are `:no`;
+      # `:maybe` otherwise.
+      def case_when_branch_certainty(subject_type, when_node)
+        return :maybe unless when_node.respond_to?(:conditions)
+
+        results = when_node.conditions.map { |c| case_when_pattern_certainty(subject_type, c) }
+        return :maybe if results.empty?
+        return :yes if results.include?(:yes)
+        return :no if results.all?(:no)
+
+        :maybe
+      end
+
+      # Static three-valued certainty for `pattern === subject`.
+      # Specialises two pattern shapes:
+      #
+      # - **Class / Module reference** (`Integer`, `Foo::Bar`):
+      #   reduce to `subject.is_a?(class)` via
+      #   `Narrowing.narrow_class` / `narrow_not_class`. A Bot
+      #   truthy fragment means no inhabitant matches (`:no`); a
+      #   Bot falsey fragment means every inhabitant matches
+      #   (`:yes`).
+      # - **Value-equality literal** (numeric / String / Symbol /
+      #   true / false / nil) against a `Constant[c]` subject:
+      #   the static comparison `pattern_value === c` is exact.
+      #   Other subject carriers stay `:maybe` because the
+      #   runtime value isn't pinned.
+      #
+      # Other pattern shapes (Range, Regexp, custom `===`) stay
+      # `:maybe` — the existing union fallback handles them.
+      def case_when_pattern_certainty(subject_type, pattern_node)
+        class_name = build_constant_path_name(pattern_node)
+        return class_pattern_certainty(subject_type, class_name) if class_name
+
+        literal = literal_pattern_value(pattern_node)
+        return literal_pattern_certainty(subject_type, literal[:value]) if literal
+
+        :maybe
+      end
+
+      def class_pattern_certainty(subject_type, class_name)
+        env = scope.environment
+        truthy_bot = Narrowing.narrow_class(subject_type, class_name, environment: env).is_a?(Type::Bot)
+        falsey_bot = Narrowing.narrow_not_class(subject_type, class_name, environment: env).is_a?(Type::Bot)
+
+        return :no if truthy_bot && !falsey_bot
+        return :yes if !truthy_bot && falsey_bot
+
+        :maybe
+      end
+
+      VALUE_EQUALITY_CLASSES = [Integer, Float, Rational, Complex, String, Symbol,
+                                TrueClass, FalseClass, NilClass].freeze
+      private_constant :VALUE_EQUALITY_CLASSES
+
+      # Returns `{ value: v }` when `pattern_node` types to a
+      # `Constant[v]` of a value-equality-safe class (so `===`
+      # reduces to `==`), else nil. Wrapped in a hash so a literal
+      # `nil` / `false` value doesn't collide with the "no literal"
+      # signal.
+      def literal_pattern_value(pattern_node)
+        type = type_of(pattern_node)
+        return nil unless type.is_a?(Type::Constant)
+        return nil unless VALUE_EQUALITY_CLASSES.any? { |klass| type.value.is_a?(klass) }
+
+        { value: type.value }
+      end
+
+      def literal_pattern_certainty(subject_type, pattern_value)
+        return :maybe unless subject_type.is_a?(Type::Constant)
+        return :maybe unless VALUE_EQUALITY_CLASSES.any? { |klass| subject_type.value.is_a?(klass) }
+
+        pattern_value == subject_type.value ? :yes : :no
       end
 
       # `when` clauses for `case` and `in` clauses for `case ... in` have
@@ -1056,6 +1187,9 @@ module Rigor
         per_element = try_per_element_block_fold(node, receiver)
         return per_element if per_element
 
+        hash_transform = try_hash_shape_block_fold(node, receiver)
+        return hash_transform if hash_transform
+
         result = MethodDispatcher.dispatch(
           receiver_type: receiver,
           method_name: node.name,
@@ -1455,10 +1589,17 @@ module Rigor
       # the dispatch chain untouched.
       PER_ELEMENT_TUPLE_METHODS = Set[
         :map, :collect, :filter_map, :flat_map,
+        :select, :filter, :reject,
         :find, :detect, :find_index, :index
       ].freeze
       private_constant :PER_ELEMENT_TUPLE_METHODS
 
+      HASH_SHAPE_TRANSFORM_METHODS = Set[
+        :transform_keys, :transform_keys!,
+        :transform_values, :transform_values!
+      ].freeze
+      private_constant :HASH_SHAPE_TRANSFORM_METHODS
+
       # Cardinality cap for per-element block fold over
       # finite-bound `Constant<Range>` receivers. Walking
       # `(1..1_000_000).map { … }` element-wise would balloon
@@ -1475,15 +1616,50 @@ module Rigor
         element_types = per_element_elements_of(receiver_type)
         return nil if element_types.nil? || element_types.empty?
 
-        block_node = call_node.block
-        return nil unless block_node.is_a?(Prism::BlockNode)
+        per_position = per_element_block_results(call_node.block, element_types)
+        return nil if per_position.nil? || per_position.any?(&:nil?)
 
-        per_position = element_types.map do |element_type|
-          type_block_body_with_param(block_node, [element_type])
+        assemble_per_element_result(call_node.name, per_position, element_types)
+      end
+
+      # Evaluates the call's block once per receiver element.
+      # Two block shapes are supported:
+      #
+      # - `Prism::BlockNode` — a full `do … end` / `{ … }` block;
+      #   the body is re-typed per position with the element
+      #   bound to the block parameter.
+      # - `Prism::BlockArgumentNode` wrapping a `SymbolNode` —
+      #   the `&:predicate` shorthand; the symbol is dispatched
+      #   as a zero-arg method on each element type.
+      #
+      # Any other shape (`&proc_local`, `&method(:foo)`, no
+      # block) returns `nil` so the fold declines.
+      def per_element_block_results(block, element_types)
+        case block
+        when Prism::BlockNode
+          element_types.map { |element_type| type_block_body_with_param(block, [element_type]) }
+        when Prism::BlockArgumentNode
+          per_element_symbol_results(block, element_types)
         end
-        return nil if per_position.any?(&:nil?)
+      end
 
-        assemble_per_element_result(call_node.name, per_position, element_types)
+      def per_element_symbol_results(block_arg, element_types)
+        expression = block_arg.expression
+        return nil unless expression.is_a?(Prism::SymbolNode)
+
+        method_name = expression.unescaped.to_sym
+        element_types.map do |element_type|
+          MethodDispatcher.dispatch(
+            receiver_type: element_type,
+            method_name: method_name,
+            arg_types: [],
+            block_type: nil,
+            environment: scope.environment,
+            scope: scope
+          )
+        end
+      rescue StandardError
+        nil
       end
 
       # Returns the per-position element types for a finite,
@@ -1532,11 +1708,37 @@ module Rigor
         when :map, :collect then Type::Combinator.tuple_of(*per_position)
         when :filter_map then assemble_filter_map_result(per_position)
         when :flat_map then assemble_flat_map_result(per_position)
+        when :select, :filter
+          assemble_filter_result(per_position, element_types, keep_on_truthy: true)
+        when :reject
+          assemble_filter_result(per_position, element_types, keep_on_truthy: false)
         when :find, :detect then assemble_find_result(per_position, element_types)
         when :find_index, :index then assemble_find_index_result(per_position)
         end
       end
 
+      # `select` / `filter` / `reject`: keeps each receiver
+      # element whose per-position predicate result folds to a
+      # decisive `Constant` — Ruby-truthy for `select` / `filter`,
+      # Ruby-falsey for `reject`. The surviving elements assemble
+      # into a `Tuple`, strictly tighter than the RBS-projected
+      # `Array[Elem]`.
+      #
+      # Folds tightly only when EVERY position is a `Constant`:
+      # a single non-`Constant` position leaves the result
+      # cardinality unknown (the element might or might not
+      # survive), so the dispatcher declines and the RBS tier
+      # widens to `Array[Elem]`. `[].select` style empty results
+      # are sound — an empty `Tuple` is the empty-array carrier.
+      def assemble_filter_result(per_position, element_types, keep_on_truthy:)
+        return nil unless per_position.all?(Type::Constant)
+
+        kept = element_types.each_index.filter_map do |index|
+          element_types[index] if truthy_constant?(per_position[index]) == keep_on_truthy
+        end
+        Type::Combinator.tuple_of(*kept)
+      end
+
       # `filter_map` folds tightly only when every per-position
       # result is a `Constant`: positions whose value is `nil`
       # or `false` drop, the rest survive in declaration order.
@@ -1611,6 +1813,94 @@ module Rigor
         type.is_a?(Type::Constant) && type.value && type.value != false
       end
 
+      # Per-pair block fold for `HashShape#transform_keys` and
+      # `HashShape#transform_values` (and their bang variants).
+      #
+      # When the receiver is a closed `HashShape` with no optional
+      # keys, applies the call's block (a `Prism::BlockNode` or
+      # `Prism::BlockArgumentNode`) to each key/value pair
+      # independently and assembles a new `HashShape`:
+      #
+      # - `transform_values` / `transform_values!`: re-types
+      #   each VALUE by binding it to the block parameter; keys
+      #   are preserved unchanged.
+      # - `transform_keys` / `transform_keys!`: re-types each
+      #   KEY by wrapping it in `Constant[k]` and passing it to
+      #   the block; values are preserved unchanged. The result
+      #   key must be a `Constant[Symbol | String]` — otherwise
+      #   the tier declines (the new key cannot be used as a
+      #   static HashShape index). Collisions (two old keys
+      #   mapping to the same new key) also decline.
+      #
+      # Returns `nil` on any decline so the dispatcher falls
+      # through to `RbsDispatch` and gets the widened `Hash[K, V]`
+      # answer.
+      def try_hash_shape_block_fold(call_node, receiver_type)
+        return nil unless HASH_SHAPE_TRANSFORM_METHODS.include?(call_node.name)
+        return nil unless receiver_type.is_a?(Type::HashShape)
+        return nil unless receiver_type.closed?
+        return nil unless receiver_type.optional_keys.empty?
+
+        block_arg = call_node.block
+        return nil if block_arg.nil?
+
+        if %i[transform_values transform_values!].include?(call_node.name)
+          fold_hash_shape_transform_values(receiver_type, block_arg)
+        else
+          fold_hash_shape_transform_keys(receiver_type, block_arg)
+        end
+      end
+
+      def fold_hash_shape_transform_values(shape, block_arg)
+        new_pairs = {}
+        shape.pairs.each do |key, value|
+          new_value = apply_hash_block(block_arg, value)
+          return nil if new_value.nil?
+
+          new_pairs[key] = new_value
+        end
+        Type::Combinator.hash_shape_of(new_pairs)
+      end
+
+      def fold_hash_shape_transform_keys(shape, block_arg)
+        new_pairs = {}
+        shape.pairs.each do |key, value|
+          key_type = Type::Combinator.constant_of(key)
+          new_key_type = apply_hash_block(block_arg, key_type)
+          return nil unless new_key_type.is_a?(Type::Constant)
+
+          new_key = new_key_type.value
+          return nil unless new_key.is_a?(Symbol) || new_key.is_a?(String)
+          return nil if new_pairs.key?(new_key)
+
+          new_pairs[new_key] = value
+        end
+        Type::Combinator.hash_shape_of(new_pairs)
+      end
+
+      # Applies a single-argument block (either a full BlockNode
+      # or a `&:symbol` BlockArgumentNode) to `param_type` and
+      # returns the resulting type, or `nil` on failure.
+      def apply_hash_block(block_arg, param_type)
+        case block_arg
+        when Prism::BlockNode
+          type_block_body_with_param(block_arg, [param_type])
+        when Prism::BlockArgumentNode
+          expression = block_arg.expression
+          return nil unless expression.is_a?(Prism::SymbolNode)
+
+          MethodDispatcher.dispatch(
+            receiver_type: param_type,
+            method_name: expression.unescaped.to_sym,
+            arg_types: [],
+            block_type: nil,
+            environment: scope.environment,
+            call_node: block_arg,
+            scope: scope
+          )
+        end
+      end
+
       def type_block_body_with_param(block_node, expected_param_types)
         bindings = BlockParameterBinder.new(expected_param_types: expected_param_types).bind(block_node)
         block_scope = bindings.reduce(scope) { |acc, (name, type)| acc.with_local(name, type) }

data/lib/rigor/inference/method_dispatcher/cgi_folding.rb

@@ -0,0 +1,109 @@
+# frozen_string_literal: true
+
+require "cgi/util"
+require_relative "../../type"
+
+module Rigor
+  module Inference
+    module MethodDispatcher
+      # Folds `CGI` module-function calls on statically known
+      # string constants.
+      #
+      # `CGI.escapeHTML` / `CGI.unescapeHTML` and related methods
+      # are pure, deterministic functions over their string inputs.
+      # When the argument is a `Constant[String]`, the analyzer can
+      # evaluate the call at inference time and return the concrete
+      # `Constant[String]` result.
+      #
+      # === Supported methods
+      #
+      # * `escapeHTML(str)` / `escape_html(str)` / `h(str)` —
+      #   HTML-escape. Returns `Constant[String]`.
+      # * `unescapeHTML(str)` / `unescape_html(str)` —
+      #   HTML-unescape. Returns `Constant[String]`.
+      # * `escape(str)` / `unescape(str)` —
+      #   URL-encode / decode (`application/x-www-form-urlencoded`).
+      #   Returns `Constant[String]`.
+      # * `escapeURIComponent(str)` / `escape_uri_component(str)`,
+      #   `unescapeURIComponent(str)` / `unescape_uri_component(str)` —
+      #   URI-component percent-encode / decode. Returns `Constant[String]`.
+      # * `escapeElement(str, *elements)` / `escape_element(str, *elements)`,
+      #   `unescapeElement(str, *elements)` / `unescape_element(str, *elements)` —
+      #   element-level escape / unescape (first arg is the string,
+      #   remaining args are element names). Returns `Constant[String]`.
+      #
+      # === Non-constant / unsupported cases
+      #
+      # Returns `nil` (deferring to the next dispatcher tier) when:
+      # - the receiver is not `Singleton[CGI]`,
+      # - the first argument is not a `Constant[String]`,
+      # - the method is not in the supported set.
+      module CGIFolding
+        CGI_HTML_ESCAPE_METHODS = Set[:escapeHTML, :escape_html, :h].freeze
+        CGI_HTML_UNESCAPE_METHODS = Set[:unescapeHTML, :unescape_html].freeze
+        CGI_URL_ESCAPE_METHODS = Set[:escape, :unescape].freeze
+        CGI_URI_ESCAPE_METHODS = Set[
+          :escapeURIComponent, :escape_uri_component,
+          :unescapeURIComponent, :unescape_uri_component
+        ].freeze
+        CGI_ELEMENT_ESCAPE_METHODS = Set[
+          :escapeElement, :escape_element,
+          :unescapeElement, :unescape_element
+        ].freeze
+        CGI_ALL_ESCAPE_METHODS = (
+          CGI_HTML_ESCAPE_METHODS | CGI_HTML_UNESCAPE_METHODS |
+          CGI_URL_ESCAPE_METHODS | CGI_URI_ESCAPE_METHODS |
+          CGI_ELEMENT_ESCAPE_METHODS
+        ).freeze
+
+        private_constant :CGI_HTML_ESCAPE_METHODS, :CGI_HTML_UNESCAPE_METHODS,
+                         :CGI_URL_ESCAPE_METHODS, :CGI_URI_ESCAPE_METHODS,
+                         :CGI_ELEMENT_ESCAPE_METHODS, :CGI_ALL_ESCAPE_METHODS
+
+        module_function
+
+        # @return [Rigor::Type, nil] folded result, or nil to defer.
+        def try_dispatch(receiver:, method_name:, args:)
+          return nil unless dispatch_target?(receiver)
+          return nil unless CGI_ALL_ESCAPE_METHODS.include?(method_name)
+
+          fold_cgi_call(method_name, args)
+        end
+
+        def dispatch_target?(receiver)
+          receiver.is_a?(Type::Singleton) && receiver.class_name == "CGI"
+        end
+
+        def fold_cgi_call(method_name, args)
+          return nil if args.empty?
+          return nil unless args.first.is_a?(Type::Constant) && args.first.value.is_a?(String)
+
+          str = args.first.value
+
+          if CGI_ELEMENT_ESCAPE_METHODS.include?(method_name)
+            fold_cgi_element(method_name, str, args.drop(1))
+          else
+            Type::Combinator.constant_of(CGI.public_send(method_name, str))
+          end
+        rescue StandardError
+          nil
+        end
+
+        # `CGI.escapeElement(str, "elem1", "elem2", ...)` — element-
+        # level escape / unescape. The remaining args after the first
+        # must be `Constant[String]` element names.
+        def fold_cgi_element(method_name, str, element_args)
+          elements = element_args.map do |arg|
+            return nil unless arg.is_a?(Type::Constant) && arg.value.is_a?(String)
+
+            arg.value
+          end
+
+          Type::Combinator.constant_of(CGI.public_send(method_name, str, *elements))
+        rescue StandardError
+          nil
+        end
+      end
+    end
+  end
+end