3 files changed, 24 insertions, 40 deletions

diff --git a/include/linux/page_cgroup.h b/include/linux/page_cgroup.h
index aaa60da8783c..1153095ee457 100644
--- a/include/linux/page_cgroup.h
+++ b/include/linux/page_cgroup.h
@@ -149,7 +149,7 @@ static inline void __init page_cgroup_init_flatmem(void)
 extern unsigned short swap_cgroup_cmpxchg(swp_entry_t ent,
                                         unsigned short old, unsigned short new);
 extern unsigned short swap_cgroup_record(swp_entry_t ent, unsigned short id);
-extern unsigned short lookup_swap_cgroup(swp_entry_t ent);
+extern unsigned short lookup_swap_cgroup_id(swp_entry_t ent);
 extern int swap_cgroup_swapon(int type, unsigned long max_pages);
 extern void swap_cgroup_swapoff(int type);
 #else
@@ -161,7 +161,7 @@ unsigned short swap_cgroup_record(swp_entry_t ent, unsigned short id)
 }
 
 static inline
-unsigned short lookup_swap_cgroup(swp_entry_t ent)
+unsigned short lookup_swap_cgroup_id(swp_entry_t ent)
 {
         return 0;
 }
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 71a9774e6ead..4c53e971749e 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -2474,7 +2474,7 @@ struct mem_cgroup *try_get_mem_cgroup_from_page(struct page *page)
                         memcg = NULL;
         } else if (PageSwapCache(page)) {
                 ent.val = page_private(page);
-                id = lookup_swap_cgroup(ent);
+                id = lookup_swap_cgroup_id(ent);
                 rcu_read_lock();
                 memcg = mem_cgroup_lookup(id);
                 if (memcg && !css_tryget(&memcg->css))
@@ -5264,7 +5264,7 @@ static int is_target_pte_for_mc(struct vm_area_struct *vma,
         }
         /* There is a swap entry and a page doesn't exist or isn't charged */
         if (ent.val && !ret &&
-                        css_id(&mc.from->css) == lookup_swap_cgroup(ent)) {
+                        css_id(&mc.from->css) == lookup_swap_cgroup_id(ent)) {
                 ret = MC_TARGET_SWAP;
                 if (target)
                         target->ent = ent;
diff --git a/mm/page_cgroup.c b/mm/page_cgroup.c
index e910524e5a08..b99d19edf89b 100644
--- a/mm/page_cgroup.c
+++ b/mm/page_cgroup.c
@@ -334,7 +334,6 @@ struct swap_cgroup {
         unsigned short          id;
 };
 #define SC_PER_PAGE     (PAGE_SIZE/sizeof(struct swap_cgroup))
-#define SC_POS_MASK     (SC_PER_PAGE - 1)
 
 /*
  * SwapCgroup implements "lookup" and "exchange" operations.
@@ -376,6 +375,21 @@ not_enough_page:
         return -ENOMEM;
 }
 
+static struct swap_cgroup *lookup_swap_cgroup(swp_entry_t ent,
+                                        struct swap_cgroup_ctrl **ctrlp)
+{
+        pgoff_t offset = swp_offset(ent);
+        struct swap_cgroup_ctrl *ctrl;
+        struct page *mappage;
+
+        ctrl = &swap_cgroup_ctrl[swp_type(ent)];
+        if (ctrlp)
+                *ctrlp = ctrl;
+
+        mappage = ctrl->map[offset / SC_PER_PAGE];
+        return page_address(mappage) + offset % SC_PER_PAGE;
+}
+
 /**
  * swap_cgroup_cmpxchg - cmpxchg mem_cgroup's id for this swp_entry.
  * @end: swap entry to be cmpxchged
@@ -388,21 +402,13 @@ not_enough_page:
 unsigned short swap_cgroup_cmpxchg(swp_entry_t ent,
                                         unsigned short old, unsigned short new)
 {
-        int type = swp_type(ent);
-        unsigned long offset = swp_offset(ent);
-        unsigned long idx = offset / SC_PER_PAGE;
-        unsigned long pos = offset & SC_POS_MASK;
         struct swap_cgroup_ctrl *ctrl;
-        struct page *mappage;
         struct swap_cgroup *sc;
         unsigned long flags;
         unsigned short retval;
 
-        ctrl = &swap_cgroup_ctrl[type];
+        sc = lookup_swap_cgroup(ent, &ctrl);
 
-        mappage = ctrl->map[idx];
-        sc = page_address(mappage);
-        sc += pos;
         spin_lock_irqsave(&ctrl->lock, flags);
         retval = sc->id;
         if (retval == old)
@@ -423,21 +429,13 @@ unsigned short swap_cgroup_cmpxchg(swp_entry_t ent,
  */
 unsigned short swap_cgroup_record(swp_entry_t ent, unsigned short id)
 {
-        int type = swp_type(ent);
-        unsigned long offset = swp_offset(ent);
-        unsigned long idx = offset / SC_PER_PAGE;
-        unsigned long pos = offset & SC_POS_MASK;
         struct swap_cgroup_ctrl *ctrl;
-        struct page *mappage;
         struct swap_cgroup *sc;
         unsigned short old;
         unsigned long flags;
 
-        ctrl = &swap_cgroup_ctrl[type];
+        sc = lookup_swap_cgroup(ent, &ctrl);
 
-        mappage = ctrl->map[idx];
-        sc = page_address(mappage);
-        sc += pos;
         spin_lock_irqsave(&ctrl->lock, flags);
         old = sc->id;
         sc->id = id;
@@ -447,28 +445,14 @@ unsigned short swap_cgroup_record(swp_entry_t ent, unsigned short id)
 }
 
 /**
- * lookup_swap_cgroup - lookup mem_cgroup tied to swap entry
+ * lookup_swap_cgroup_id - lookup mem_cgroup id tied to swap entry
  * @ent: swap entry to be looked up.
  *
  * Returns CSS ID of mem_cgroup at success. 0 at failure. (0 is invalid ID)
  */
-unsigned short lookup_swap_cgroup(swp_entry_t ent)
+unsigned short lookup_swap_cgroup_id(swp_entry_t ent)
 {
-        int type = swp_type(ent);
-        unsigned long offset = swp_offset(ent);
-        unsigned long idx = offset / SC_PER_PAGE;
-        unsigned long pos = offset & SC_POS_MASK;
-        struct swap_cgroup_ctrl *ctrl;
-        struct page *mappage;
-        struct swap_cgroup *sc;
-        unsigned short ret;
-
-        ctrl = &swap_cgroup_ctrl[type];
-        mappage = ctrl->map[idx];
-        sc = page_address(mappage);
-        sc += pos;
-        ret = sc->id;
-        return ret;
+        return lookup_swap_cgroup(ent, NULL)->id;
 }
 
 int swap_cgroup_swapon(int type, unsigned long max_pages)