From a7170a4a44bc874c5b45499659c2307e59d2344a Mon Sep 17 00:00:00 2001
From: Philipp Oppermann <dev@phil-opp.com>
Date: Thu, 13 Apr 2017 19:11:31 +0200
Subject: [PATCH] Use type system tricks to make next_table methods safe

---
 src/memory/paging/table.rs | 46 ++++++++++++++++++++++++++++++++------
 1 file changed, 39 insertions(+), 7 deletions(-)

diff --git a/src/memory/paging/table.rs b/src/memory/paging/table.rs
index e5a4f14..949c35a 100644
--- a/src/memory/paging/table.rs
+++ b/src/memory/paging/table.rs
@@ -1,20 +1,24 @@
+use core::marker::PhantomData;
 use core::ops::{Index, IndexMut};
 use memory::paging::entry::*;
 use memory::paging::ENTRY_COUNT;
 
-pub const P4: *mut Table = 0xffffffff_fffff000 as *mut _;
+pub const P4: *mut Table<Level4> = 0xffffffff_fffff000 as *mut _;
 
-pub struct Table {
+pub struct Table<L: TableLevel> {
     entries: [Entry; ENTRY_COUNT],
+    level: PhantomData<L>,
 }
 
-impl Table {
+impl<L> Table<L> where L: TableLevel {
     pub fn zero(&mut self) {
         for entry in self.entries.iter_mut() {
             entry.set_unused();
         }
     }
+}
 
+impl<L> Table<L> where L: HierarchicalLevel {
     fn next_table_address(&self, index: usize) -> Option<usize> {
         let entry_flags = self[index].flags();
         if entry_flags.contains(PRESENT) && !entry_flags.contains(HUGE_PAGE) {
@@ -25,18 +29,18 @@ impl Table {
         }
     }
 
-    pub fn next_table(&self, index: usize) -> Option<&Table> {
+    pub fn next_table(&self, index: usize) -> Option<&Table<L::NextLevel>> {
         self.next_table_address(index)
             .map(|address| unsafe { &*(address as *const _) })
     }
 
-    pub fn next_table_mut(&mut self, index: usize) -> Option<&mut Table> {
+    pub fn next_table_mut(&mut self, index: usize) -> Option<&mut Table<L::NextLevel>> {
         self.next_table_address(index)
             .map(|address| unsafe { &mut *(address as *mut _) })
     }
 }
 
-impl Index<usize> for Table {
+impl<L> Index<usize> for Table<L> where L: TableLevel {
     type Output = Entry;
 
     fn index(&self, index: usize) -> &Entry {
@@ -44,8 +48,36 @@ impl Index<usize> for Table {
     }
 }
 
-impl IndexMut<usize> for Table {
+impl<L> IndexMut<usize> for Table<L> where L: TableLevel {
     fn index_mut(&mut self, index: usize) -> &mut Entry {
         &mut self.entries[index]
     }
 }
+
+pub trait TableLevel {}
+
+pub enum Level4 {}
+pub enum Level3 {}
+pub enum Level2 {}
+pub enum Level1 {}
+
+impl TableLevel for Level4 {}
+impl TableLevel for Level3 {}
+impl TableLevel for Level2 {}
+impl TableLevel for Level1 {}
+
+pub trait HierarchicalLevel: TableLevel {
+    type NextLevel: TableLevel;
+}
+
+impl HierarchicalLevel for Level4 {
+    type NextLevel = Level3;
+}
+
+impl HierarchicalLevel for Level3 {
+    type NextLevel = Level2;
+}
+
+impl HierarchicalLevel for Level2 {
+    type NextLevel = Level1;
+}