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
 }
 
 static inline
-unsigned short lookup_swap_cgroup(swp_entry_t ent)
+unsigned short lookup_swap_cgroup_id(swp_entry_t ent)
 {
        return 0;
 }
 
        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.
        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
 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)
  */
 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;
 }
 
 /**
- * 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)