c
	prev = NULL;
	vmi.vma = NULL; // Initialize vmi.vma
	while ((vma = vma_next(&vmi))) { // Use vma_next for manual iteration
		cond_resched();
		BUG_ON(!!vma->vm_userfaultfd_ctx.ctx ^
		       !!(vma->vm_flags & __VM_UFFD_FLAGS));
		if (vma->vm_userfaultfd_ctx.ctx != ctx) {
			prev = vma;
			continue;
		}
		new_flags = vma->vm_flags & ~__VM_UFFD_FLAGS;
		prev = vma_merge(&vmi, mm, prev, vma->vm_start, vma->vm_end,
				 new_flags, vma->anon_vma, vma->vm_file,
				 vma->vm_pgoff, vma_policy(vma),
				 vma->vm_userfaultfd_ctx, vma_get_anon_name(vma));
        if (vmi.vma == NULL) { // Check if the current VMA was removed
            // The VMA was merged/removed.  vma_next() already advanced.
            continue;
        } else {
           vma = vmi.vma; // Ensure vma points to a valid VMA
           // Existing code to work with the (potentially modified) vma.
           // ...
        }
    } // End of vma iteration loop
