6 files changed, 94 insertions, 36 deletions

diff --git a/mm/huge_memory.c b/mm/huge_memory.c
index 010d32944d14..ce59ada09462 100644
--- a/mm/huge_memory.c
+++ b/mm/huge_memory.c
@@ -1375,13 +1375,14 @@ static void __split_huge_page(struct page *page,
                               struct anon_vma *anon_vma)
 {
         int mapcount, mapcount2;
+        pgoff_t pgoff = page->index << (PAGE_CACHE_SHIFT - PAGE_SHIFT);
         struct anon_vma_chain *avc;
 
         BUG_ON(!PageHead(page));
         BUG_ON(PageTail(page));
 
         mapcount = 0;
-        list_for_each_entry(avc, &anon_vma->head, same_anon_vma) {
+        anon_vma_interval_tree_foreach(avc, &anon_vma->rb_root, pgoff, pgoff) {
                 struct vm_area_struct *vma = avc->vma;
                 unsigned long addr = vma_address(page, vma);
                 BUG_ON(is_vma_temporary_stack(vma));
@@ -1407,7 +1408,7 @@ static void __split_huge_page(struct page *page,
         __split_huge_page_refcount(page);
 
         mapcount2 = 0;
-        list_for_each_entry(avc, &anon_vma->head, same_anon_vma) {
+        anon_vma_interval_tree_foreach(avc, &anon_vma->rb_root, pgoff, pgoff) {
                 struct vm_area_struct *vma = avc->vma;
                 unsigned long addr = vma_address(page, vma);
                 BUG_ON(is_vma_temporary_stack(vma));
diff --git a/mm/interval_tree.c b/mm/interval_tree.c
index 4ab7b9ec3a56..f7c72cd35e1d 100644
--- a/mm/interval_tree.c
+++ b/mm/interval_tree.c
@@ -8,6 +8,7 @@
 
 #include <linux/mm.h>
 #include <linux/fs.h>
+#include <linux/rmap.h>
 #include <linux/interval_tree_generic.h>
 
 static inline unsigned long vma_start_pgoff(struct vm_area_struct *v)
@@ -57,3 +58,16 @@ void vma_interval_tree_insert_after(struct vm_area_struct *node,
         rb_insert_augmented(&node->shared.linear.rb, root,
                             &vma_interval_tree_augment);
 }
+
+static inline unsigned long avc_start_pgoff(struct anon_vma_chain *avc)
+{
+        return vma_start_pgoff(avc->vma);
+}
+
+static inline unsigned long avc_last_pgoff(struct anon_vma_chain *avc)
+{
+        return vma_last_pgoff(avc->vma);
+}
+
+INTERVAL_TREE_DEFINE(struct anon_vma_chain, rb, unsigned long, rb_subtree_last,
+                     avc_start_pgoff, avc_last_pgoff,, anon_vma_interval_tree)
diff --git a/mm/ksm.c b/mm/ksm.c
index 9638620a7530..14ee5cf8a513 100644
--- a/mm/ksm.c
+++ b/mm/ksm.c
@@ -1618,7 +1618,8 @@ again:
                 struct vm_area_struct *vma;
 
                 anon_vma_lock(anon_vma);
-                list_for_each_entry(vmac, &anon_vma->head, same_anon_vma) {
+                anon_vma_interval_tree_foreach(vmac, &anon_vma->rb_root,
+                                               0, ULONG_MAX) {
                         vma = vmac->vma;
                         if (rmap_item->address < vma->vm_start ||
                             rmap_item->address >= vma->vm_end)
@@ -1671,7 +1672,8 @@ again:
                 struct vm_area_struct *vma;
 
                 anon_vma_lock(anon_vma);
-                list_for_each_entry(vmac, &anon_vma->head, same_anon_vma) {
+                anon_vma_interval_tree_foreach(vmac, &anon_vma->rb_root,
+                                               0, ULONG_MAX) {
                         vma = vmac->vma;
                         if (rmap_item->address < vma->vm_start ||
                             rmap_item->address >= vma->vm_end)
@@ -1723,7 +1725,8 @@ again:
                 struct vm_area_struct *vma;
 
                 anon_vma_lock(anon_vma);
-                list_for_each_entry(vmac, &anon_vma->head, same_anon_vma) {
+                anon_vma_interval_tree_foreach(vmac, &anon_vma->rb_root,
+                                               0, ULONG_MAX) {
                         vma = vmac->vma;
                         if (rmap_item->address < vma->vm_start ||
                             rmap_item->address >= vma->vm_end)
diff --git a/mm/memory-failure.c b/mm/memory-failure.c
index c38a6257d082..6c5899b9034a 100644
--- a/mm/memory-failure.c
+++ b/mm/memory-failure.c
@@ -400,18 +400,21 @@ static void collect_procs_anon(struct page *page, struct list_head *to_kill,
         struct vm_area_struct *vma;
         struct task_struct *tsk;
         struct anon_vma *av;
+        pgoff_t pgoff;
 
         av = page_lock_anon_vma(page);
         if (av == NULL) /* Not actually mapped anymore */
                 return;
 
+        pgoff = page->index << (PAGE_CACHE_SHIFT - PAGE_SHIFT);
         read_lock(&tasklist_lock);
         for_each_process (tsk) {
                 struct anon_vma_chain *vmac;
 
                 if (!task_early_kill(tsk))
                         continue;
-                list_for_each_entry(vmac, &av->head, same_anon_vma) {
+                anon_vma_interval_tree_foreach(vmac, &av->rb_root,
+                                               pgoff, pgoff) {
                         vma = vmac->vma;
                         if (!page_mapped_in_vma(page, vma))
                                 continue;
diff --git a/mm/mmap.c b/mm/mmap.c
index 66984aab7915..2e580ed79211 100644
--- a/mm/mmap.c
+++ b/mm/mmap.c
@@ -353,6 +353,38 @@ void validate_mm(struct mm_struct *mm)
 #define validate_mm(mm) do { } while (0)
 #endif
 
+/*
+ * vma has some anon_vma assigned, and is already inserted on that
+ * anon_vma's interval trees.
+ *
+ * Before updating the vma's vm_start / vm_end / vm_pgoff fields, the
+ * vma must be removed from the anon_vma's interval trees using
+ * anon_vma_interval_tree_pre_update_vma().
+ *
+ * After the update, the vma will be reinserted using
+ * anon_vma_interval_tree_post_update_vma().
+ *
+ * The entire update must be protected by exclusive mmap_sem and by
+ * the root anon_vma's mutex.
+ */
+static inline void
+anon_vma_interval_tree_pre_update_vma(struct vm_area_struct *vma)
+{
+        struct anon_vma_chain *avc;
+
+        list_for_each_entry(avc, &vma->anon_vma_chain, same_vma)
+                anon_vma_interval_tree_remove(avc, &avc->anon_vma->rb_root);
+}
+
+static inline void
+anon_vma_interval_tree_post_update_vma(struct vm_area_struct *vma)
+{
+        struct anon_vma_chain *avc;
+
+        list_for_each_entry(avc, &vma->anon_vma_chain, same_vma)
+                anon_vma_interval_tree_insert(avc, &avc->anon_vma->rb_root);
+}
+
 static int find_vma_links(struct mm_struct *mm, unsigned long addr,
                 unsigned long end, struct vm_area_struct **pprev,
                 struct rb_node ***rb_link, struct rb_node **rb_parent)
@@ -565,20 +597,17 @@ again:			remove_next = 1 + (end > next->vm_end);
 
         vma_adjust_trans_huge(vma, start, end, adjust_next);
 
-        /*
-         * When changing only vma->vm_end, we don't really need anon_vma
-         * lock. This is a fairly rare case by itself, but the anon_vma
-         * lock may be shared between many sibling processes.  Skipping
-         * the lock for brk adjustments makes a difference sometimes.
-         */
-        if (vma->anon_vma && (importer || start != vma->vm_start)) {
-                anon_vma = vma->anon_vma;
+        anon_vma = vma->anon_vma;
+        if (!anon_vma && adjust_next)
+                anon_vma = next->anon_vma;
+        if (anon_vma) {
                 VM_BUG_ON(adjust_next && next->anon_vma &&
                           anon_vma != next->anon_vma);
-        } else if (adjust_next && next->anon_vma)
-                anon_vma = next->anon_vma;
-        if (anon_vma)
                 anon_vma_lock(anon_vma);
+                anon_vma_interval_tree_pre_update_vma(vma);
+                if (adjust_next)
+                        anon_vma_interval_tree_pre_update_vma(next);
+        }
 
         if (root) {
                 flush_dcache_mmap_lock(mapping);
@@ -619,8 +648,12 @@ again:			remove_next = 1 + (end > next->vm_end);
                 __insert_vm_struct(mm, insert);
         }
 
-        if (anon_vma)
+        if (anon_vma) {
+                anon_vma_interval_tree_post_update_vma(vma);
+                if (adjust_next)
+                        anon_vma_interval_tree_post_update_vma(next);
                 anon_vma_unlock(anon_vma);
+        }
         if (mapping)
                 mutex_unlock(&mapping->i_mmap_mutex);
 
@@ -1748,7 +1781,9 @@ int expand_upwards(struct vm_area_struct *vma, unsigned long address)
                 if (vma->vm_pgoff + (size >> PAGE_SHIFT) >= vma->vm_pgoff) {
                         error = acct_stack_growth(vma, size, grow);
                         if (!error) {
+                                anon_vma_interval_tree_pre_update_vma(vma);
                                 vma->vm_end = address;
+                                anon_vma_interval_tree_post_update_vma(vma);
                                 perf_event_mmap(vma);
                         }
                 }
@@ -1798,8 +1833,10 @@ int expand_downwards(struct vm_area_struct *vma,
                 if (grow <= vma->vm_pgoff) {
                         error = acct_stack_growth(vma, size, grow);
                         if (!error) {
+                                anon_vma_interval_tree_pre_update_vma(vma);
                                 vma->vm_start = address;
                                 vma->vm_pgoff -= grow;
+                                anon_vma_interval_tree_post_update_vma(vma);
                                 perf_event_mmap(vma);
                         }
                 }
@@ -2515,7 +2552,7 @@ static DEFINE_MUTEX(mm_all_locks_mutex);
 
 static void vm_lock_anon_vma(struct mm_struct *mm, struct anon_vma *anon_vma)
 {
-        if (!test_bit(0, (unsigned long *) &anon_vma->root->head.next)) {
+        if (!test_bit(0, (unsigned long *) &anon_vma->root->rb_root.rb_node)) {
                 /*
                  * The LSB of head.next can't change from under us
                  * because we hold the mm_all_locks_mutex.
@@ -2531,7 +2568,7 @@ static void vm_lock_anon_vma(struct mm_struct *mm, struct anon_vma *anon_vma)
                  * anon_vma->root->mutex.
                  */
                 if (__test_and_set_bit(0, (unsigned long *)
-                                       &anon_vma->root->head.next))
+                                       &anon_vma->root->rb_root.rb_node))
                         BUG();
         }
 }
@@ -2572,7 +2609,7 @@ static void vm_lock_mapping(struct mm_struct *mm, struct address_space *mapping)
  * A single task can't take more than one mm_take_all_locks() in a row
  * or it would deadlock.
  *
- * The LSB in anon_vma->head.next and the AS_MM_ALL_LOCKS bitflag in
+ * The LSB in anon_vma->rb_root.rb_node and the AS_MM_ALL_LOCKS bitflag in
  * mapping->flags avoid to take the same lock twice, if more than one
  * vma in this mm is backed by the same anon_vma or address_space.
  *
@@ -2619,13 +2656,13 @@ out_unlock:
 
 static void vm_unlock_anon_vma(struct anon_vma *anon_vma)
 {
-        if (test_bit(0, (unsigned long *) &anon_vma->root->head.next)) {
+        if (test_bit(0, (unsigned long *) &anon_vma->root->rb_root.rb_node)) {
                 /*
                  * The LSB of head.next can't change to 0 from under
                  * us because we hold the mm_all_locks_mutex.
                  *
                  * We must however clear the bitflag before unlocking
-                 * the vma so the users using the anon_vma->head will
+                 * the vma so the users using the anon_vma->rb_root will
                  * never see our bitflag.
                  *
                  * No need of atomic instructions here, head.next
@@ -2633,7 +2670,7 @@ static void vm_unlock_anon_vma(struct anon_vma *anon_vma)
                  * anon_vma->root->mutex.
                  */
                 if (!__test_and_clear_bit(0, (unsigned long *)
-                                          &anon_vma->root->head.next))
+                                          &anon_vma->root->rb_root.rb_node))
                         BUG();
                 anon_vma_unlock(anon_vma);
         }
diff --git a/mm/rmap.c b/mm/rmap.c
index 8cbd62fde0f1..9c61bf387fd1 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -127,12 +127,7 @@ static void anon_vma_chain_link(struct vm_area_struct *vma,
         avc->vma = vma;
         avc->anon_vma = anon_vma;
         list_add(&avc->same_vma, &vma->anon_vma_chain);
-
-        /*
-         * It's critical to add new vmas to the tail of the anon_vma,
-         * see comment in huge_memory.c:__split_huge_page().
-         */
-        list_add_tail(&avc->same_anon_vma, &anon_vma->head);
+        anon_vma_interval_tree_insert(avc, &anon_vma->rb_root);
 }
 
 /**
@@ -336,13 +331,13 @@ void unlink_anon_vmas(struct vm_area_struct *vma)
                 struct anon_vma *anon_vma = avc->anon_vma;
 
                 root = lock_anon_vma_root(root, anon_vma);
-                list_del(&avc->same_anon_vma);
+                anon_vma_interval_tree_remove(avc, &anon_vma->rb_root);
 
                 /*
                  * Leave empty anon_vmas on the list - we'll need
                  * to free them outside the lock.
                  */
-                if (list_empty(&anon_vma->head))
+                if (RB_EMPTY_ROOT(&anon_vma->rb_root))
                         continue;
 
                 list_del(&avc->same_vma);
@@ -371,7 +366,7 @@ static void anon_vma_ctor(void *data)
 
         mutex_init(&anon_vma->mutex);
         atomic_set(&anon_vma->refcount, 0);
-        INIT_LIST_HEAD(&anon_vma->head);
+        anon_vma->rb_root = RB_ROOT;
 }
 
 void __init anon_vma_init(void)
@@ -724,6 +719,7 @@ static int page_referenced_anon(struct page *page,
 {
         unsigned int mapcount;
         struct anon_vma *anon_vma;
+        pgoff_t pgoff;
         struct anon_vma_chain *avc;
         int referenced = 0;
 
@@ -732,7 +728,8 @@ static int page_referenced_anon(struct page *page,
                 return referenced;
 
         mapcount = page_mapcount(page);
-        list_for_each_entry(avc, &anon_vma->head, same_anon_vma) {
+        pgoff = page->index << (PAGE_CACHE_SHIFT - PAGE_SHIFT);
+        anon_vma_interval_tree_foreach(avc, &anon_vma->rb_root, pgoff, pgoff) {
                 struct vm_area_struct *vma = avc->vma;
                 unsigned long address = vma_address(page, vma);
                 if (address == -EFAULT)
@@ -1445,6 +1442,7 @@ bool is_vma_temporary_stack(struct vm_area_struct *vma)
 static int try_to_unmap_anon(struct page *page, enum ttu_flags flags)
 {
         struct anon_vma *anon_vma;
+        pgoff_t pgoff;
         struct anon_vma_chain *avc;
         int ret = SWAP_AGAIN;
 
@@ -1452,7 +1450,8 @@ static int try_to_unmap_anon(struct page *page, enum ttu_flags flags)
         if (!anon_vma)
                 return ret;
 
-        list_for_each_entry(avc, &anon_vma->head, same_anon_vma) {
+        pgoff = page->index << (PAGE_CACHE_SHIFT - PAGE_SHIFT);
+        anon_vma_interval_tree_foreach(avc, &anon_vma->rb_root, pgoff, pgoff) {
                 struct vm_area_struct *vma = avc->vma;
                 unsigned long address;
 
@@ -1668,6 +1667,7 @@ static int rmap_walk_anon(struct page *page, int (*rmap_one)(struct page *,
                 struct vm_area_struct *, unsigned long, void *), void *arg)
 {
         struct anon_vma *anon_vma;
+        pgoff_t pgoff = page->index << (PAGE_CACHE_SHIFT - PAGE_SHIFT);
         struct anon_vma_chain *avc;
         int ret = SWAP_AGAIN;
 
@@ -1681,7 +1681,7 @@ static int rmap_walk_anon(struct page *page, int (*rmap_one)(struct page *,
         if (!anon_vma)
                 return ret;
         anon_vma_lock(anon_vma);
-        list_for_each_entry(avc, &anon_vma->head, same_anon_vma) {
+        anon_vma_interval_tree_foreach(avc, &anon_vma->rb_root, pgoff, pgoff) {
                 struct vm_area_struct *vma = avc->vma;
                 unsigned long address = vma_address(page, vma);
                 if (address == -EFAULT)