}
 #endif /* CONFIG_BLOCK */
 
+static int expected_page_refs(struct page *page)
+{
+       int expected_count = 1;
+
+       /*
+        * Device public or private pages have an extra refcount as they are
+        * ZONE_DEVICE pages.
+        */
+       expected_count += is_device_private_page(page);
+       expected_count += is_device_public_page(page);
+       if (page_mapping(page))
+               expected_count += hpage_nr_pages(page) + page_has_private(page);
+
+       return expected_count;
+}
+
 /*
  * Replace the page in the mapping.
  *
        XA_STATE(xas, &mapping->i_pages, page_index(page));
        struct zone *oldzone, *newzone;
        int dirty;
-       int expected_count = 1 + extra_count;
-
-       /*
-        * Device public or private pages have an extra refcount as they are
-        * ZONE_DEVICE pages.
-        */
-       expected_count += is_device_private_page(page);
-       expected_count += is_device_public_page(page);
+       int expected_count = expected_page_refs(page) + extra_count;
 
        if (!mapping) {
                /* Anonymous page without mapping */
        newzone = page_zone(newpage);
 
        xas_lock_irq(&xas);
-
-       expected_count += hpage_nr_pages(page) + page_has_private(page);
        if (page_count(page) != expected_count || xas_load(&xas) != page) {
                xas_unlock_irq(&xas);
                return -EAGAIN;