1 files changed, 15 insertions, 33 deletions

diff --git a/actionpack/lib/action_dispatch/middleware/flash.rb b/actionpack/lib/action_dispatch/middleware/flash.rb
index bc5b163931..cff0877030 100644
--- a/actionpack/lib/action_dispatch/middleware/flash.rb
+++ b/actionpack/lib/action_dispatch/middleware/flash.rb
@@ -78,7 +78,7 @@ module ActionDispatch
       include Enumerable
 
       def initialize #:nodoc:
-        @used    = Set.new
+        @discard = Set.new
         @closed  = false
         @flashes = {}
         @now     = nil
@@ -93,7 +93,7 @@ module ActionDispatch
       end
 
       def []=(k, v) #:nodoc:
-        keep(k)
+        @discard.delete k
         @flashes[k] = v
       end
 
@@ -102,7 +102,7 @@ module ActionDispatch
       end
 
       def update(h) #:nodoc:
-        h.keys.each { |k| keep(k) }
+        @discard.subtract h.keys
         @flashes.update h
         self
       end
@@ -116,6 +116,7 @@ module ActionDispatch
       end
 
       def delete(key)
+        @discard.delete key
         @flashes.delete key
         self
       end
@@ -129,6 +130,7 @@ module ActionDispatch
       end
 
       def clear
+        @discard.clear
         @flashes.clear
       end
 
@@ -139,7 +141,7 @@ module ActionDispatch
       alias :merge! :update
 
       def replace(h) #:nodoc:
-        @used = Set.new
+        @discard.clear
         @flashes.replace h
         self
       end
@@ -163,7 +165,8 @@ module ActionDispatch
       #    flash.keep            # keeps the entire flash
       #    flash.keep(:notice)   # keeps only the "notice" entry, the rest of the flash is discarded
       def keep(k = nil)
-        use(k, false)
+        @discard.subtract Array(k || keys)
+        k ? self[k] : self
       end
 
       # Marks the entire flash or a single flash entry to be discarded by the end of the current action:
@@ -171,24 +174,16 @@ module ActionDispatch
       #     flash.discard              # discard the entire flash at the end of the current action
       #     flash.discard(:warning)    # discard only the "warning" entry at the end of the current action
       def discard(k = nil)
-        use(k)
+        @discard.merge Array(k || keys)
+        k ? self[k] : self
       end
 
       # Mark for removal entries that were kept, and delete unkept ones.
       #
       # This method is called automatically by filters, so you generally don't need to care about it.
       def sweep #:nodoc:
-        keys.each do |k|
-          unless @used.include?(k)
-            @used << k
-          else
-            delete(k)
-            @used.delete(k)
-          end
-        end
-
-        # clean up after keys that could have been left over by calling reject! or shift on the flash
-        (@used - keys).each{ |k| @used.delete(k) }
+        @discard.each { |k| @flashes.delete k }
+        @discard.replace @flashes.keys
       end
 
       # Convenience accessor for flash[:alert]
@@ -212,22 +207,9 @@ module ActionDispatch
       end
 
       protected
-
-        def now_is_loaded?
-          !!@now
-        end
-
-        # Used internally by the <tt>keep</tt> and <tt>discard</tt> methods
-        #     use()               # marks the entire flash as used
-        #     use('msg')          # marks the "msg" entry as used
-        #     use(nil, false)     # marks the entire flash as unused (keeps it around for one more action)
-        #     use('msg', false)   # marks the "msg" entry as unused (keeps it around for one more action)
-        # Returns the single value for the key you asked to be marked (un)used or the FlashHash itself
-        # if no key is passed.
-        def use(key = nil, used = true)
-          Array(key || keys).each { |k| used ? @used << k : @used.delete(k) }
-          return key ? self[key] : self
-        end
+      def now_is_loaded?
+        @now
+      end
     end
 
     def initialize(app)