*/
 
 #include <linux/mm.h>
+#include <linux/mm_inline.h>
 #include <linux/hugetlb.h>
 #include <linux/shm.h>
 #include <linux/ksm.h>
 #include <linux/mmu_notifier.h>
 #include <linux/uaccess.h>
 #include <linux/userfaultfd_k.h>
+#include <linux/mempolicy.h>
 
 #include <asm/cacheflush.h>
 #include <asm/tlb.h>
                /* can we just expand the current mapping? */
                if (vma_expandable(vma, new_len - old_len)) {
                        long pages = (new_len - old_len) >> PAGE_SHIFT;
+                       unsigned long extension_start = addr + old_len;
+                       unsigned long extension_end = addr + new_len;
+                       pgoff_t extension_pgoff = vma->vm_pgoff + (old_len >> PAGE_SHIFT);
 
                        if (vma->vm_flags & VM_ACCOUNT) {
                                if (security_vm_enough_memory_mm(mm, pages)) {
                                }
                        }
 
-                       if (vma_adjust(vma, vma->vm_start, addr + new_len,
-                                      vma->vm_pgoff, NULL)) {
+                       /*
+                        * Function vma_merge() is called on the extension we are adding to
+                        * the already existing vma, vma_merge() will merge this extension with
+                        * the already existing vma (expand operation itself) and possibly also
+                        * with the next vma if it becomes adjacent to the expanded vma and
+                        * otherwise compatible.
+                        */
+                       vma = vma_merge(mm, vma, extension_start, extension_end,
+                                       vma->vm_flags, vma->anon_vma, vma->vm_file,
+                                       extension_pgoff, vma_policy(vma),
+                                       vma->vm_userfaultfd_ctx, anon_vma_name(vma));
+                       if (!vma) {
                                vm_unacct_memory(pages);
                                ret = -ENOMEM;
                                goto out;
 
        return addr;
 }
 
+/*
+ * This test validates that merge is called when expanding a mapping.
+ * Mapping containing three pages is created, middle page is unmapped
+ * and then the mapping containing the first page is expanded so that
+ * it fills the created hole. The two parts should merge creating
+ * single mapping with three pages.
+ */
+static void mremap_expand_merge(unsigned long page_size)
+{
+       char *test_name = "mremap expand merge";
+       FILE *fp;
+       char *line = NULL;
+       size_t len = 0;
+       bool success = false;
+       char *start = mmap(NULL, 3 * page_size, PROT_READ | PROT_WRITE,
+                          MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+
+       munmap(start + page_size, page_size);
+       mremap(start, page_size, 2 * page_size, 0);
+
+       fp = fopen("/proc/self/maps", "r");
+       if (fp == NULL) {
+               ksft_test_result_fail("%s\n", test_name);
+               return;
+       }
+
+       while (getline(&line, &len, fp) != -1) {
+               char *first = strtok(line, "- ");
+               void *first_val = (void *)strtol(first, NULL, 16);
+               char *second = strtok(NULL, "- ");
+               void *second_val = (void *) strtol(second, NULL, 16);
+
+               if (first_val == start && second_val == start + 3 * page_size) {
+                       success = true;
+                       break;
+               }
+       }
+       if (success)
+               ksft_test_result_pass("%s\n", test_name);
+       else
+               ksft_test_result_fail("%s\n", test_name);
+       fclose(fp);
+}
+
 /*
  * Returns the start address of the mapping on success, else returns
  * NULL on failure.
        int i, run_perf_tests;
        unsigned int threshold_mb = VALIDATION_DEFAULT_THRESHOLD;
        unsigned int pattern_seed;
+       int num_expand_tests = 1;
        struct test test_cases[MAX_TEST];
        struct test perf_test_cases[MAX_PERF_TEST];
        int page_size;
                                (threshold_mb * _1MB >= _1GB);
 
        ksft_set_plan(ARRAY_SIZE(test_cases) + (run_perf_tests ?
-                     ARRAY_SIZE(perf_test_cases) : 0));
+                     ARRAY_SIZE(perf_test_cases) : 0) + num_expand_tests);
 
        for (i = 0; i < ARRAY_SIZE(test_cases); i++)
                run_mremap_test_case(test_cases[i], &failures, threshold_mb,
                                     pattern_seed);
 
+       mremap_expand_merge(page_size);
+
        if (run_perf_tests) {
                ksft_print_msg("\n%s\n",
                 "mremap HAVE_MOVE_PMD/PUD optimization time comparison for 1GB region:");