diff --git a/src/kernel_space.rs b/src/kernel_space.rs
index cc04541..68296e9 100644
--- a/src/kernel_space.rs
+++ b/src/kernel_space.rs
@@ -8,15 +8,29 @@
 use alloc::sync::Arc;
 use spin::Mutex;
 
+use crate::KernelAddressTranslator;
+
 use super::{
     address::{PhysicalAddress, VirtualAddress, VirtualAddressRange},
     page_pool::{Page, PagePool},
     MemoryAccessRights, RegimeVaRange, TranslationGranule, TranslationRegime, Xlat, XlatError,
 };
 
+struct KernelAddressTranslatorIdentity;
+
+impl KernelAddressTranslator for KernelAddressTranslatorIdentity {
+    fn kernel_to_pa(va: VirtualAddress) -> PhysicalAddress {
+        PhysicalAddress(va.0 & 0x0000_000f_ffff_ffff)
+    }
+
+    fn pa_to_kernel(pa: PhysicalAddress) -> VirtualAddress {
+        VirtualAddress(pa.0 | 0xffff_fff0_0000_0000)
+    }
+}
+
 #[derive(Clone)]
 pub struct KernelSpace {
-    xlat: Arc<Mutex<Xlat<36>>>,
+    xlat: Arc<Mutex<Xlat<KernelAddressTranslatorIdentity, 36>>>,
 }
 
 /// # Kernel space memory mapping
@@ -154,16 +168,6 @@
     }
 
     /// Kernel virtual address to physical address
-    #[cfg(not(test))]
-    pub const fn kernel_to_pa(kernel_address: u64) -> u64 {
-        kernel_address & 0x0000_000f_ffff_ffff
-    }
-    /// Physical address to kernel virtual address
-    #[cfg(not(test))]
-    pub const fn pa_to_kernel(pa: u64) -> u64 {
-        // TODO: make this consts assert_eq!(pa & 0xffff_fff0_0000_0000, 0);
-        pa | 0xffff_fff0_0000_0000
-    }
 
     // Do not use any mapping in test build
     #[cfg(test)]