1 files changed, 19 insertions, 46 deletions

diff --git a/actionpack/lib/action_dispatch/middleware/flash.rb b/actionpack/lib/action_dispatch/middleware/flash.rb
index e59404ef68..54f2669f43 100644
--- a/actionpack/lib/action_dispatch/middleware/flash.rb
+++ b/actionpack/lib/action_dispatch/middleware/flash.rb
@@ -4,7 +4,7 @@ module ActionDispatch
     # read a notice you put there or <tt>flash["notice"] = "hello"</tt>
     # to put a new one.
     def flash
-      @env[Flash::KEY] ||= (session["flash"] || Flash::FlashHash.new)
+      @env[Flash::KEY] ||= (session["flash"] || Flash::FlashHash.new).tap(&:sweep)
     end
   end
 
@@ -17,7 +17,7 @@ module ActionDispatch
   #     def create
   #       # save post
   #       flash[:notice] = "Post successfully created"
-  #       redirect_to posts_path(@post)
+  #       redirect_to @post
   #     end
   #
   #     def show
@@ -78,7 +78,7 @@ module ActionDispatch
       include Enumerable
 
       def initialize #:nodoc:
-        @used    = Set.new
+        @discard = Set.new
         @closed  = false
         @flashes = {}
         @now     = nil
@@ -93,8 +93,7 @@ module ActionDispatch
       end
 
       def []=(k, v) #:nodoc:
-        raise ClosedError, :flash if closed?
-        keep(k)
+        @discard.delete k
         @flashes[k] = v
       end
 
@@ -103,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
@@ -117,6 +116,7 @@ module ActionDispatch
       end
 
       def delete(key)
+        @discard.delete key
         @flashes.delete key
         self
       end
@@ -130,6 +130,7 @@ module ActionDispatch
       end
 
       def clear
+        @discard.clear
         @flashes.clear
       end
 
@@ -140,7 +141,7 @@ module ActionDispatch
       alias :merge! :update
 
       def replace(h) #:nodoc:
-        @used = Set.new
+        @discard.clear
         @flashes.replace h
         self
       end
@@ -159,16 +160,13 @@ module ActionDispatch
         @now ||= FlashNow.new(self)
       end
 
-      attr_reader :closed
-      alias :closed? :closed
-      def close!; @closed = true; end
-
       # Keeps either the entire current flash or a specific flash entry available for the next action:
       #
       #    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:
@@ -176,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]
@@ -217,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)
@@ -240,10 +217,6 @@ module ActionDispatch
     end
 
     def call(env)
-      if (session = env['rack.session']) && (flash = session['flash'])
-        flash.sweep
-      end
-
       @app.call(env)
     ensure
       session    = env['rack.session'] || {}
@@ -258,10 +231,10 @@ module ActionDispatch
         end
 
         env[KEY] = new_hash
-        new_hash.close!
       end
 
-      if session.key?('flash') && session['flash'].empty?
+      if (!session.respond_to?(:loaded?) || session.loaded?) && # (reset_session uses {}, which doesn't implement #loaded?)
+         session.key?('flash') && session['flash'].empty?
         session.delete('flash')
       end
     end