c
prev = NULL;
for_each_vma(vmi, vma) {
    // ... (existing code up to the call to vma_merge)

    unsigned long start = vma->vm_start; // Keep track of the VMA start

    // Adjust the vma to remove the uffd flags. This doesn't change the vma list structure.
    vma_adjust(vma, vma->vm_start, vma->vm_end, vma->vm_pgoff,
              new_flags, vma->anon_vma, vma->vm_file);

    prev = vma_merge(&vmi, mm, prev, start, vma->vm_end, new_flags, 
                      vma->anon_vma, vma->vm_file, vma->vm_pgoff,
                      vma_policy(vma), NULL, vma_get_anon_name(vma));

    vma = find_vma(mm, start); // Find the VMA that now covers the start address.


    if (vma && vma->vm_start == start && vma->vm_userfaultfd_ctx.ctx == ctx) {
         // Correct VMA found, continue
         // ... rest of the loop ...
    }
}
