dec_mm_counter(mm, MM_ANONPAGES);
                else
                        dec_mm_counter(mm, MM_FILEPAGES);
+       } else if (IS_ENABLED(CONFIG_MIGRATION) && (flags & TTU_MIGRATION)) {
+               swp_entry_t entry;
+               pte_t swp_pte;
+               /*
+                * Store the pfn of the page in a special migration
+                * pte. do_swap_page() will wait until the migration
+                * pte is removed and then restart fault handling.
+                */
+               entry = make_migration_entry(page, pte_write(pteval));
+               swp_pte = swp_entry_to_pte(entry);
+               if (pte_soft_dirty(pteval))
+                       swp_pte = pte_swp_mksoft_dirty(swp_pte);
+               set_pte_at(mm, address, pte, swp_pte);
        } else if (PageAnon(page)) {
                swp_entry_t entry = { .val = page_private(page) };
                pte_t swp_pte;
-
-               if (PageSwapCache(page)) {
-                       /*
-                        * Store the swap location in the pte.
-                        * See handle_pte_fault() ...
-                        */
-                       if (swap_duplicate(entry) < 0) {
-                               set_pte_at(mm, address, pte, pteval);
-                               ret = SWAP_FAIL;
-                               goto out_unmap;
-                       }
-                       if (list_empty(&mm->mmlist)) {
-                               spin_lock(&mmlist_lock);
-                               if (list_empty(&mm->mmlist))
-                                       list_add(&mm->mmlist, &init_mm.mmlist);
-                               spin_unlock(&mmlist_lock);
-                       }
-                       dec_mm_counter(mm, MM_ANONPAGES);
-                       inc_mm_counter(mm, MM_SWAPENTS);
-               } else if (IS_ENABLED(CONFIG_MIGRATION)) {
-                       /*
-                        * Store the pfn of the page in a special migration
-                        * pte. do_swap_page() will wait until the migration
-                        * pte is removed and then restart fault handling.
-                        */
-                       BUG_ON(!(flags & TTU_MIGRATION));
-                       entry = make_migration_entry(page, pte_write(pteval));
+               /*
+                * Store the swap location in the pte.
+                * See handle_pte_fault() ...
+                */
+               VM_BUG_ON_PAGE(!PageSwapCache(page), page);
+               if (swap_duplicate(entry) < 0) {
+                       set_pte_at(mm, address, pte, pteval);
+                       ret = SWAP_FAIL;
+                       goto out_unmap;
+               }
+               if (list_empty(&mm->mmlist)) {
+                       spin_lock(&mmlist_lock);
+                       if (list_empty(&mm->mmlist))
+                               list_add(&mm->mmlist, &init_mm.mmlist);
+                       spin_unlock(&mmlist_lock);
                }
+               dec_mm_counter(mm, MM_ANONPAGES);
+               inc_mm_counter(mm, MM_SWAPENTS);
                swp_pte = swp_entry_to_pte(entry);
                if (pte_soft_dirty(pteval))
                        swp_pte = pte_swp_mksoft_dirty(swp_pte);
                set_pte_at(mm, address, pte, swp_pte);
-       } else if (IS_ENABLED(CONFIG_MIGRATION) &&
-                  (flags & TTU_MIGRATION)) {
-               /* Establish migration entry for a file page */
-               swp_entry_t entry;
-               entry = make_migration_entry(page, pte_write(pteval));
-               set_pte_at(mm, address, pte, swp_entry_to_pte(entry));
        } else
                dec_mm_counter(mm, MM_FILEPAGES);