1 files changed, 63 insertions, 33 deletions

diff --git a/lib/bundler/index.rb b/lib/bundler/index.rb
index b8c599f63a..df46facc88 100644
--- a/lib/bundler/index.rb
+++ b/lib/bundler/index.rb
@@ -10,8 +10,8 @@ module Bundler
       i
     end
 
-    attr_reader :specs, :all_specs, :sources
-    protected :specs, :all_specs
+    attr_reader :specs, :duplicates, :sources
+    protected :specs, :duplicates
 
     RUBY = "ruby"
     NULL = "\0"
@@ -19,21 +19,21 @@ module Bundler
     def initialize
       @sources = []
       @cache = {}
-      @specs = Hash.new {|h, k| h[k] = {} }
-      @all_specs = Hash.new {|h, k| h[k] = EMPTY_SEARCH }
+      @specs = {}
+      @duplicates = {}
     end
 
     def initialize_copy(o)
       @sources = o.sources.dup
       @cache = {}
-      @specs = Hash.new {|h, k| h[k] = {} }
-      @all_specs = Hash.new {|h, k| h[k] = EMPTY_SEARCH }
+      @specs = {}
+      @duplicates = {}
 
       o.specs.each do |name, hash|
         @specs[name] = hash.dup
       end
-      o.all_specs.each do |name, array|
-        @all_specs[name] = array.dup
+      o.duplicates.each do |name, array|
+        @duplicates[name] = array.dup
       end
     end
 
@@ -46,12 +46,11 @@ module Bundler
       true
     end
 
-    def search_all(name)
-      all_matches = local_search(name) + @all_specs[name]
-      @sources.each do |source|
-        all_matches.concat(source.search_all(name))
-      end
-      all_matches
+    def search_all(name, &blk)
+      return enum_for(:search_all, name) unless blk
+      specs_by_name(name).each(&blk)
+      @duplicates[name]&.each(&blk)
+      @sources.each {|source| source.search_all(name, &blk) }
     end
 
     # Search this index's specs, and any source indexes that this index knows
@@ -61,11 +60,14 @@ module Bundler
       return results unless @sources.any?
 
       @sources.each do |source|
-        results.concat(source.search(query))
+        results = safe_concat(results, source.search(query))
       end
-      results.uniq(&:full_name)
+      results.uniq!(&:full_name) unless results.empty? # avoid modifying frozen EMPTY_SEARCH
+      results
     end
 
+    alias_method :[], :search
+
     def local_search(query)
       case query
       when Gem::Specification, RemoteSpecification, LazySpecification, EndpointSpecification then search_by_spec(query)
@@ -76,12 +78,10 @@ module Bundler
       end
     end
 
-    alias_method :[], :search
-
-    def <<(spec)
-      @specs[spec.name][spec.full_name] = spec
-      spec
+    def add(spec)
+      (@specs[spec.name] ||= {}).store(spec.full_name, spec)
     end
+    alias_method :<<, :add
 
     def each(&blk)
       return enum_for(:each) unless blk
@@ -115,15 +115,25 @@ module Bundler
       names.uniq
     end
 
-    def use(other, override_dupes = false)
+    # Combines indexes proritizing existing specs, like `Hash#reverse_merge!`
+    # Duplicate specs found in `other` are stored in `@duplicates`.
+    def use(other)
       return unless other
-      other.each do |s|
-        if (dupes = search_by_spec(s)) && !dupes.empty?
-          # safe to << since it's a new array when it has contents
-          @all_specs[s.name] = dupes << s
-          next unless override_dupes
+      other.each do |spec|
+        exist?(spec) ? add_duplicate(spec) : add(spec)
+      end
+      self
+    end
+
+    # Combines indexes proritizing specs from `other`, like `Hash#merge!`
+    # Duplicate specs found in `self` are saved in `@duplicates`.
+    def merge!(other)
+      return unless other
+      other.each do |spec|
+        if existing = find_by_spec(spec)
+          add_duplicate(existing)
         end
-        self << s
+        add spec
       end
       self
     end
@@ -135,8 +145,7 @@ module Bundler
     end
 
     # Whether all the specs in self are in other
-    # TODO: rename to #include?
-    def ==(other)
+    def subset?(other)
       all? do |spec|
         other_spec = other[spec].first
         other_spec && dependencies_eql?(spec, other_spec) && spec.source == other_spec.source
@@ -157,19 +166,40 @@ module Bundler
 
     private
 
+    def safe_concat(a, b)
+      return a if b.empty?
+      return b if a.empty?
+      a.concat(b)
+    end
+
+    def add_duplicate(spec)
+      (@duplicates[spec.name] ||= []) << spec
+    end
+
     def specs_by_name_and_version(name, version)
-      specs_by_name(name).select {|spec| spec.version == version }
+      results = @specs[name]&.values
+      return EMPTY_SEARCH unless results
+      results.select! {|spec| spec.version == version }
+      results
     end
 
     def specs_by_name(name)
-      @specs[name].values
+      @specs[name]&.values || EMPTY_SEARCH
     end
 
     EMPTY_SEARCH = [].freeze
 
     def search_by_spec(spec)
-      spec = @specs[spec.name][spec.full_name]
+      spec = find_by_spec(spec)
       spec ? [spec] : EMPTY_SEARCH
     end
+
+    def find_by_spec(spec)
+      @specs[spec.name]&.fetch(spec.full_name, nil)
+    end
+
+    def exist?(spec)
+      @specs[spec.name]&.key?(spec.full_name)
+    end
   end
 end