if (vma &&
            (!vma->anon_vma || list_is_singular(&vma->anon_vma_chain)) &&
            ((vma->vm_flags & ~VM_SOFTDIRTY) == flags)) {
-               mas->index = vma->vm_start;
-               mas->last = addr + len - 1;
-               vma_adjust_trans_huge(vma, addr, addr + len, 0);
+               mas_set_range(mas, vma->vm_start, addr + len - 1);
+               if (mas_preallocate(mas, vma, GFP_KERNEL))
+                       return -ENOMEM;
+
+               vma_adjust_trans_huge(vma, vma->vm_start, addr + len, 0);
                if (vma->anon_vma) {
                        anon_vma_lock_write(vma->anon_vma);
                        anon_vma_interval_tree_pre_update_vma(vma);
                }
                vma->vm_end = addr + len;
                vma->vm_flags |= VM_SOFTDIRTY;
-               if (mas_store_gfp(mas, vma, GFP_KERNEL))
-                       goto mas_expand_failed;
+               mas_store_prealloc(mas, vma);
 
                if (vma->anon_vma) {
                        anon_vma_interval_tree_post_update_vma(vma);
 vma_alloc_fail:
        vm_unacct_memory(len >> PAGE_SHIFT);
        return -ENOMEM;
-
-mas_expand_failed:
-       if (vma->anon_vma) {
-               anon_vma_interval_tree_post_update_vma(vma);
-               anon_vma_unlock_write(vma->anon_vma);
-       }
-       return -ENOMEM;
 }
 
 int vm_brk_flags(unsigned long addr, unsigned long request, unsigned long flags)