1 files changed, 4 insertions, 9 deletions

diff --git a/drivers/xen/balloon.c b/drivers/xen/balloon.c
index a0fb5eac407c..8dc7109d61b7 100644
--- a/drivers/xen/balloon.c
+++ b/drivers/xen/balloon.c
@@ -44,13 +44,15 @@
 #include <linux/list.h>
 #include <linux/sysdev.h>
 
-#include <asm/xen/hypervisor.h>
 #include <asm/page.h>
 #include <asm/pgalloc.h>
 #include <asm/pgtable.h>
 #include <asm/uaccess.h>
 #include <asm/tlb.h>
 
+#include <asm/xen/hypervisor.h>
+#include <asm/xen/hypercall.h>
+#include <xen/interface/xen.h>
 #include <xen/interface/memory.h>
 #include <xen/xenbus.h>
 #include <xen/features.h>
@@ -122,14 +124,7 @@ static struct timer_list balloon_timer;
 static void scrub_page(struct page *page)
 {
 #ifdef CONFIG_XEN_SCRUB_PAGES
-        if (PageHighMem(page)) {
-                void *v = kmap(page);
-                clear_page(v);
-                kunmap(v);
-        } else {
-                void *v = page_address(page);
-                clear_page(v);
-        }
+        clear_highpage(page);
 #endif
 }