diff --git a/src/rust/cpu.rs b/src/rust/cpu.rs
index cf76419c..a76c6666 100644
--- a/src/rust/cpu.rs
+++ b/src/rust/cpu.rs
@@ -1,16 +1,11 @@
 // TODO: Make this an instance, so we can plug in a fake cpu
 
+use cpu2;
 use page::Page;
 use state_flags::CachedStateFlags;
 
 mod unsafe_cpu {
     extern "C" {
-        pub fn tlb_set_has_code(physical_page: u32, has_code: bool);
-        pub fn read8(addr: u32) -> u8;
-        pub fn read16(addr: u32) -> u16;
-        pub fn read32s(addr: u32) -> u32;
-        pub fn check_tlb_invariants();
-
         pub fn codegen_finalize(
             wasm_table_index: u16,
             phys_addr: u32,
@@ -29,15 +24,15 @@ pub enum BitSize {
     DWORD,
 }
 
-pub fn read8(addr: u32) -> u8 { unsafe { unsafe_cpu::read8(addr) } }
-pub fn read16(addr: u32) -> u16 { unsafe { unsafe_cpu::read16(addr) } }
-pub fn read32(addr: u32) -> u32 { unsafe { unsafe_cpu::read32s(addr) } }
+pub fn read8(addr: u32) -> u8 { unsafe { cpu2::memory::read8(addr) as u8 } }
+pub fn read16(addr: u32) -> u16 { unsafe { cpu2::memory::read16(addr) as u16 } }
+pub fn read32(addr: u32) -> u32 { unsafe { cpu2::memory::read32s(addr) as u32 } }
 
 pub fn tlb_set_has_code(physical_page: Page, has_code: bool) {
-    unsafe { unsafe_cpu::tlb_set_has_code(physical_page.to_u32(), has_code) }
+    unsafe { cpu2::cpu::tlb_set_has_code(physical_page, has_code) }
 }
 
-pub fn check_tlb_invariants() { unsafe { unsafe_cpu::check_tlb_invariants() } }
+pub fn check_tlb_invariants() { unsafe { cpu2::cpu::check_tlb_invariants() } }
 
 pub fn codegen_finalize(
     wasm_table_index: u16,
diff --git a/src/rust/cpu2/cpu.rs b/src/rust/cpu2/cpu.rs
index 86ec59ba..12d15067 100644
--- a/src/rust/cpu2/cpu.rs
+++ b/src/rust/cpu2/cpu.rs
@@ -1034,8 +1034,8 @@ pub unsafe fn translate_address_write(address: i32) -> OrPageFault<u32> {
 }
 
 #[no_mangle]
-pub unsafe fn tlb_set_has_code(physical_page: u32, has_code: bool) {
-    dbg_assert!(physical_page < (1 << 20) as u32);
+pub unsafe fn tlb_set_has_code(physical_page: Page, has_code: bool) {
+    let physical_page = physical_page.to_u32();
     for i in 0..valid_tlb_entries_count {
         let page: i32 = valid_tlb_entries[i as usize];
         let entry: i32 = *tlb_data.offset(page as isize);