Help with some abstraction on my tree-walking code

code review
1

Hello Rust Community:

The question is a bit long.
Recently I am working on replicating the xv6 operation system using Rust. Here is the repository and the helper repository.

My question lies in the implementation of kvm.rs.
There are 3 functions that will walk through the page table.

  1. map_page
  2. unmap_page
  3. map_addr_recur

The common things in these function are that they traverse the page table, which can be viewed as a tree. However they will do different things during traversal.
I am wondering is it possible that I can create a shared struct to abstract the behavior of traversing the page table?

My initial thought is to create a struct, which has some common field like page table, virtual address and current level. It can have two function that user can set different behavior when encountering non-leaf page or leaf page:

struct<T> MapWalker {
    pagetable: &'a mut PageTable,
    va: VirtAddr,
    level: PageTableLevel,
    fn leaf() -> T,
    fn nonleaf() -> T,
}

By this I can implement the traversal behavior with:

impl MapWalker {
  fn walk(&mut self) {
    let index = self.va.get_index(level);
    let pte = &mut self.page_table[index];
    match self.level.next_level() {
          None => self.leaf(),
          Some(next_level) => self.nonleaf(),
    }
  }
}

However I have no idea how to design the function leaf and nonleaf. Since the arguments are different in map_page, unmap_page and map_addr.

  • In map_page, it needs the PhysAddr to and permission flag to set the page.
  • In unmap_page, it just needs a bool to indicate whether to free the page.
  • In map_addr_recur, it is going to return the PhysPage mapped to the VirtAddr.

Right now it is fine since the code is not much. However I wonder there is a good way to abstract the traversal behavior, and made my code cleaner.

2

Mutable and non-mutable is hard to abstract over, so you may want to consider separate methods for those.

As for the arguments, you can make it part of Self via generic parameter:

struct MapWalkerMut<'a, Extra> {
    pagetable: &'a mut PageTable,
    va: VirtAddr,
    level: PageTableLevel,
    extra: Extra,
}

And then rely on a trait implementation for the meat of the traversal:

pub trait PageTableVisitor {
    // n.b. requires nightly feature try_trait_v2
    // If that's no ok I'd probably just use Result
    // (You can use `.ok()` to get an `Option`)
    type Output: core::ops::Try;
    
    fn check_va(&mut self, va: VirtAddr) -> Self::Output;
    fn leaf(&mut self, pte: &mut PageTableEntry) -> Self::Output;
    fn nonleaf(&mut self, pte: &mut PageTableEntry) -> Self::Output;
}

And then implement on the MapWalkerMut:

impl<Extra: PageTableVisitor> MapWalkerMut<'_, Extra> {
    fn visit_mut(&mut self) -> Extra::Output {
        let _ = self.extra.check_va(self.va)?;
        let index = self.va.get_index(self.level);
        let pte = &mut self.pagetable[index];

        match self.level.next_level() {
            None => {
                self.extra.leaf(pte)
            }
            Some(next_level) => {
                let _ = self.extra.nonleaf(pte)?;

                let next_table = unsafe { &mut *(pte.addr() as *mut PageTable) };
                self.pagetable = next_table;
                self.level = next_level;
                self.visit_mut()
            }
        }
    }
}

If you want different return type without the Try trait, you'll have to work out something else; maybe as simple as another method to transform a Result, so you can still use ? most places.

    fn visit_mut(&mut self) -> Extra::Output {
        result = self._visit_mut(); // as above but with `Result` everywhere
        self.extra.finalize(result)
    }

Quick and dirty mock-up.

Edit: A (stable) version that doesn't use Try.

1 Like
3

Thanks quinedot:

I will check and try your code.
It comes to me that, is this kind of abstraction zero-cost? Since it is in kernel and I don't want any extra operation compared to what I implemented right now.

However I think with visitor pattern there will be some sort of function call overhead that is inevitable.

4

The code (visit_mut) will be monomorphized as if you wrote each implementation out. As for the function calls, you can mark the methods in your implementation of the trait as #[inline] or #[inline(always)] as a hint.

5

Just finish my code change. It works perfectly.
An extra question, you said that:

Mutable and non-mutable is hard to abstract over

Right now I am going to abstract out the map_addr_recur. So I will have to add another method visit in MapWalkerMut or should I just create another struct MapWalker which does not hold a mutable PageTable?

6

I would recommend the latter. You can only have one usable mutable borrow at a time (a better term would have been unique borrow), where as you can have as many shared borrowed as you like. So the difference between holding a &mut PageTable and a &PageTable can be quite large.

See also the standard library, where there's often a Iter and an IterMut, etc.

7

See, thank for your help :slight_smile:

8

Just for record:
The new implementation of the kvm is at here:
Map and Unmap page
Map Address

and the walker/visitor implementation

9

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.