lib/mongoid/selectable.rb in mongoid-5.0.0 vs lib/mongoid/selectable.rb in mongoid-5.0.1

@@ -16,9 +16,44 @@
     #
     # @return [ Hash ] The document's selector.
     #
     # @since 1.0.0
     def atomic_selector
-      @atomic_selector ||= { "_id" => _root._id }.merge!(shard_key_selector)
+      @atomic_selector ||=
+        (embedded? ? embedded_atomic_selector : root_atomic_selector)
+    end
+
+    private
+
+    # Get the atomic selector for an embedded document.
+    #
+    # @api private
+    #
+    # @example Get the embedded atomic selector.
+    #   document.embedded_atomic_selector
+    #
+    # @return [ Hash ] The embedded document selector.
+    #
+    # @since 4.0.0
+    def embedded_atomic_selector
+      if persisted? && _id_changed?
+        _parent.atomic_selector
+      else
+        _parent.atomic_selector.merge("#{atomic_path}._id" => _id)
+      end
+    end
+
+    # Get the atomic selector for a root document.
+    #
+    # @api private
+    #
+    # @example Get the root atomic selector.
+    #   document.root_atomic_selector
+    #
+    # @return [ Hash ] The root document selector.
+    #
+    # @since 4.0.0
+    def root_atomic_selector
+      { "_id" => _id }.merge!(shard_key_selector)
     end
   end
 end