Remove nil from compact element types

lib/solargraph/api_map.rb

@@ -991,6 +991,7 @@ def create_resolved_alias_pin alias_pin, original
         visibility: original.visibility,
         signatures: original.signatures.map(&:clone).freeze,
         attribute: original.attribute?,
+        removes_nil: original.removes_nil?,
         generics: original.generics.clone,
         return_type: original.return_type,
         source: :resolve_method_alias

lib/solargraph/pin/method.rb

@@ -12,6 +12,9 @@ class Method < Callable
 
       attr_writer :signatures
 
+      # @return [Boolean]
+      attr_writer :removes_nil
+
       # @return [Parser::AST::Node]
       attr_reader :node
 
@@ -22,10 +25,11 @@ class Method < Callable
       # @param attribute [Boolean]
       # @param signatures [::Array<Signature>, nil]
       # @param anon_splat [Boolean]
+      # @param removes_nil [Boolean] declared by RBS %a{solargraph:removes-nil}
       # @param context [ComplexType, ComplexType::UniqueType, nil]
       # @param [Hash{Symbol => Object}] splat
       def initialize visibility: :public, explicit: true, block: :undefined, node: nil, attribute: false, signatures: nil, anon_splat: false,
-                     context: nil, **splat
+                     removes_nil: false, context: nil, **splat
         super(**splat)
         @visibility = visibility
         @explicit = explicit
@@ -34,6 +38,7 @@ def initialize visibility: :public, explicit: true, block: :undefined, node: nil
         @attribute = attribute
         @signatures = signatures
         @anon_splat = anon_splat
+        @removes_nil = removes_nil
         @context = context if context
       end
 
@@ -68,6 +73,7 @@ def combine_with other, attrs = {}
           parameters: parameters,
           signatures: sigs,
           anon_splat: assert_same(other, :anon_splat?),
+          removes_nil: removes_nil? || other.removes_nil?,
           return_type: nil # pulled from signatures on first call
         }.merge(attrs)
         super(other, new_attrs)
@@ -381,6 +387,14 @@ def attribute?
         @attribute
       end
 
+      # True when the method returns a collection with its receiver's nil
+      # elements (Hash: nil values) removed (RBS %a{solargraph:removes-nil})
+      #
+      # @return [Boolean]
+      def removes_nil?
+        @removes_nil
+      end
+
       # @parm other [self]
       def nearly? other
         super &&

lib/solargraph/pin/search.rb

@@ -44,7 +44,6 @@ def do_query
              .compact
              # @param a [self]
              # @param b [self]
-             # @sg-ignore https://github.com/castwide/solargraph/pull/1050
              .sort { |a, b| b.match <=> a.match }
              .map(&:pin)
       end

lib/solargraph/rbs_map/conversions.rb

@@ -551,6 +551,7 @@ def method_def_to_pin decl, closure, context
         generics = decl.overloads.map(&:method_type).map do |method_type|
           type_parameter_names method_type
         end
+        removes_nil = removes_nil?(decl)
 
         if decl.instance?
           name = decl.name.to_s
@@ -565,6 +566,7 @@ def method_def_to_pin decl, closure, context
             signatures: [],
             generics: generics,
             visibility: visibility,
+            removes_nil: removes_nil,
             source: :rbs
           )
           pin.signatures.concat method_def_to_sigs(decl, pin)
@@ -587,12 +589,29 @@ def method_def_to_pin decl, closure, context
           scope: final_scope,
           signatures: [],
           generics: generics,
+          removes_nil: removes_nil,
           source: :rbs
         )
         pin.signatures.concat method_def_to_sigs(decl, pin)
         pins.push pin
       end
 
+      # %a{solargraph:removes-nil} on the method, or on every overload
+      #
+      # @param decl [RBS::AST::Members::MethodDefinition]
+      # @return [Boolean]
+      def removes_nil? decl
+        return true if removes_nil_annotation?(decl.annotations)
+
+        !decl.overloads.empty? && decl.overloads.all? { |overload| removes_nil_annotation?(overload.annotations) }
+      end
+
+      # @param annotations [Array<RBS::AST::Annotation>]
+      # @return [Boolean]
+      def removes_nil_annotation? annotations
+        annotations.any? { |annotation| annotation.string == 'solargraph:removes-nil' }
+      end
+
       # @param decl [RBS::AST::Members::MethodDefinition]
       # @param pin [Pin::Method]
       # @return [void]

lib/solargraph/rbs_map/core.rb

@@ -4,6 +4,9 @@ module Solargraph
   module RbsMap
     class Core < Base
       FILLS_DIRECTORY = File.expand_path(File.join(File.dirname(__FILE__), '..', '..', '..', 'rbs', 'fills'))
+      # Every .rbs file in it is loaded with the fills, its annotations
+      # copied onto core methods (see #annotate)
+      ANNOTATIONS_DIRECTORY = File.expand_path(File.join(File.dirname(__FILE__), '..', '..', '..', 'rbs', 'annotations'))
 
       def pins
         @pins ||= generate_pins
@@ -27,8 +30,14 @@ def generate_pins
         # Avoid RBS::DuplicatedDeclarationError by loading in a different EnvironmentLoader
         fill_loader = RBS::EnvironmentLoader.new(core_root: nil, repository: RBS::Repository.new(no_stdlib: false))
         fill_loader.add(path: Pathname(FILLS_DIRECTORY))
+        fill_loader.add(path: Pathname(ANNOTATIONS_DIRECTORY))
         fill_conversions = Conversions.new(loader: fill_loader)
-        new_pins.concat fill_conversions.pins
+        annotation_pins, fill_pins = fill_conversions.pins.partition do |pin|
+          pin.type_location&.filename&.start_with?(ANNOTATIONS_DIRECTORY)
+        end
+        new_pins.concat fill_pins
+
+        annotate new_pins, annotation_pins
 
         # add some overrides
         new_pins.concat RbsMap::CoreFills::ALL
@@ -39,6 +48,20 @@ def generate_pins
         # File.write cache_file, serial, mode: 'wb'
         processed
       end
+
+      # Copies the annotations that ANNOTATIONS_DIRECTORY declares, e.g. on
+      # 'def compact: ...' in a reopened Array, onto the core methods,
+      # keeping their core signatures
+      #
+      # @param pins [Array<Pin::Base>]
+      # @param annotation_pins [Array<Pin::Base>] pins from ANNOTATIONS_DIRECTORY
+      # @return [void]
+      def annotate pins, annotation_pins
+        removes_nil_paths = annotation_pins.select { |pin| pin.is_a?(Pin::Method) && pin.removes_nil? }.to_set(&:path)
+        pins.each do |pin|
+          pin.removes_nil = true if pin.is_a?(Pin::Method) && removes_nil_paths.include?(pin.path)
+        end
+      end
     end
   end
 end

lib/solargraph/source/chain/call.rb

@@ -212,6 +212,7 @@ def inferred_pins pins, api_map, name_pin, locals
               break if type.defined?
             end
             p = p.with_single_signature(new_signature_pin) unless new_signature_pin.nil?
+            type = without_nil_elements(type) if p.removes_nil? && type.defined?
             next p.proxy(type) if type.defined?
             if !p.macros.empty?
               result = process_macro(p, api_map, name_pin.context, locals)
@@ -242,6 +243,23 @@ def inferred_pins pins, api_map, name_pin, locals
           end
         end
 
+        # @param type [ComplexType] the return type of a method annotated
+        #   %a{solargraph:removes-nil}
+        # @return [ComplexType]
+        def without_nil_elements type
+          ComplexType.new(type.items.map { |unique_type| without_nil_subtypes(unique_type) })
+        end
+
+        # @param type [ComplexType::UniqueType]
+        # @return [ComplexType::UniqueType]
+        def without_nil_subtypes type
+          non_nil_items = type.subtypes.map do |subtype|
+            subtype.items.reject { |item| item.nil_type? || item.name == 'NilClass' }
+          end
+          subtypes = non_nil_items.reject(&:empty?).map { |items| ComplexType.new(items) }
+          subtypes.empty? ? type : type.recreate(new_subtypes: subtypes)
+        end
+
         # @param api_map [ApiMap]
         # @param method_pin [Pin::Method]
         # @return [ComplexType, nil]

rbs/annotations/removes_nil.rbs

@@ -0,0 +1,24 @@
+# Adds %a{solargraph:removes-nil} to core methods.
+#
+# Each 'def name: ...' adds only its annotations; RbsMap::Core copies them
+# onto the core method and keeps the core signatures, and uses nothing
+# else from this file.
+#
+# %a{solargraph:removes-nil} drops nil from the element types
+# (Hash: value types) of the returned collection, which a signature like
+# '() -> ::Array[Elem]' cannot express.
+
+class Array[unchecked out Elem]
+  %a{solargraph:removes-nil}
+  def compact: ...
+end
+
+module Enumerable[unchecked out Elem]
+  %a{solargraph:removes-nil}
+  def compact: ...
+end
+
+class Hash[unchecked out K, unchecked out V]
+  %a{solargraph:removes-nil}
+  def compact: ...
+end

spec/rbs_map/core_spec.rb

@@ -120,4 +120,17 @@ class Foo
     clip = api_map.clip_at('test.rb', [2, 6])
     expect(clip.infer.to_s).to eq('String')
   end
+
+  it 'copies %a{solargraph:removes-nil} onto core compact methods' do
+    store = Solargraph::ApiMap::Store.new(core.pins)
+    flags = %w[Array#compact Enumerable#compact Hash#compact Array#compact!].to_h do |path|
+      [path, store.get_path_pins(path).first.removes_nil?]
+    end
+    expect(flags).to eq('Array#compact' => true, 'Enumerable#compact' => true, 'Hash#compact' => true, 'Array#compact!' => false)
+  end
+
+  it 'keeps the core signatures of annotated methods' do
+    store = Solargraph::ApiMap::Store.new(core.pins)
+    expect(store.get_path_pins('Array#compact').first.signatures).not_to be_empty
+  end
 end