*
  * Returns: 0 on success
  */
-inline int vma_expand(struct ma_state *mas, struct vm_area_struct *vma,
+inline int vma_expand(struct vma_iterator *vmi, struct vm_area_struct *vma,
                      unsigned long start, unsigned long end, pgoff_t pgoff,
                      struct vm_area_struct *next)
 {
        /* Only handles expanding */
        VM_BUG_ON(vma->vm_start < start || vma->vm_end > end);
 
-       if (mas_preallocate(mas, vma, GFP_KERNEL))
+       if (vma_iter_prealloc(vmi, vma))
                goto nomem;
 
        vma_adjust_trans_huge(vma, start, end, 0);
        vma->vm_start = start;
        vma->vm_end = end;
        vma->vm_pgoff = pgoff;
-       /* Note: mas must be pointing to the expanding VMA */
-       vma_mas_store(vma, mas);
+       vma_iter_store(vmi, vma);
 
        if (file) {
                vma_interval_tree_insert(vma, root);
 
        /* Actually expand, if possible */
        if (vma &&
-           !vma_expand(&vmi.mas, vma, merge_start, merge_end, vm_pgoff, next)) {
+           !vma_expand(&vmi, vma, merge_start, merge_end, vm_pgoff, next)) {
                khugepaged_enter_vma(vma, vm_flags);
                goto expanded;
        }