lib/mongoid/contextual/mongo.rb in mongoid-3.0.5 vs lib/mongoid/contextual/mongo.rb in mongoid-3.0.6

@@ -127,14 +127,13 @@
       # @return [ Enumerator ] The enumerator.
       #
       # @since 3.0.0
       def each(&block)
         if block_given?
-          reset_length
           selecting do
             documents_for_iteration.each do |doc|
-              yield_and_increment(doc, &block)
+              yield_document(doc, &block)
             end
             @cache_loaded = true
             eager_loadable? ? docs : self
           end
         else
@@ -570,38 +569,10 @@
       # @since 3.0.0
       def eager_loadable?
         !eager_loaded && !criteria.inclusions.empty?
       end
 
-      # Increment the length of the results.
-      #
-      # @api private
-      #
-      # @example Increment the length.
-      #   context.increment_length
-      #
-      # @return [ Integer ] The new length
-      #
-      # @since 3.0.0
-      def increment_length
-        @length += 1
-      end
-
-      # Reset the length to zero. This happens once before iteration.
-      #
-      # @api private
-      #
-      # @example Reset the length.
-      #   context.reset_length
-      #
-      # @return [ Integer ] zero.
-      #
-      # @since 3.0.0
-      def reset_length
-        @length = 0
-      end
-
       # Apply all the optional criterion.
       #
       # @example Apply the options.
       #   context.apply_options
       #
@@ -625,16 +596,15 @@
       # @return [ Object ] The yielded value.
       #
       # @since 2.4.4
       def selecting
         begin
-          unless criteria.options[:fields].blank?
-            Threaded.selection = criteria.options[:fields]
-          end
+          fields = criteria.options[:fields]
+          Threaded.set_selection(criteria.object_id, fields) unless fields.blank?
           yield
         ensure
-          Threaded.selection = nil
+          Threaded.set_selection(criteria.object_id, nil)
         end
       end
 
       # If the provided document exists, eager load it's dependencies or return
       # nil.
@@ -648,31 +618,31 @@
       #
       # @since 3.0.0
       def with_eager_loading(document)
         selecting do
           return nil unless document
-          doc = Factory.from_db(klass, document)
+          doc = Factory.from_db(klass, document, criteria.object_id)
           eager_load([ doc ]) if eager_loadable?
           doc
         end
       end
 
-      # Yield to the document and increment the length.
+      # Yield to the document.
       #
       # @api private
       #
-      # @example Yield and increment.
-      #   context.yield_and_increment(doc) do |doc|
+      # @example Yield the document.
+      #   context.yield_document(doc) do |doc|
       #     ...
       #   end
       #
       # @param [ Document ] document The document to yield to.
       #
       # @since 3.0.0
-      def yield_and_increment(document, &block)
-        doc = document.respond_to?(:_id) ? document : Factory.from_db(klass, document)
+      def yield_document(document, &block)
+        doc = document.respond_to?(:_id) ?
+          document : Factory.from_db(klass, document, criteria.object_id)
         yield(doc)
-        increment_length
         documents.push(doc) if cacheable?
       end
     end
   end
 end