1use core::ops::{Deref, DerefMut, Range};
2
3use ax_memory_addr::MemoryAddr;
4
5use crate::{
6 FrameAllocator, PageTableEntry, PagingError, PagingResult, PhysAddr, PteConfigOf, TableMeta,
7 VirtAddr,
8 frame::Frame,
9 map::{MapConfig, MapRecursiveConfig, UnmapConfig, UnmapRecursiveConfig},
10 walk::{PageTableWalker, WalkConfig},
11};
12
13const TARGETED_FLUSH_LIMIT: usize = 32;
14
15pub struct PageTable<T: TableMeta, A: FrameAllocator> {
16 inner: PageTableRef<T, A>,
17 #[cfg(feature = "copy-from")]
18 borrowed_root_entries: Option<Range<usize>>,
19}
20
21impl<T: TableMeta, A: FrameAllocator> PageTable<T, A> {
22 pub const VALID_BITS: usize = Frame::<T, A>::PT_VALID_BITS;
23
24 pub fn new(allocator: A) -> PagingResult<Self> {
26 let inner = unsafe { PageTableRef::new(allocator) }?;
27 Ok(Self {
28 inner,
29 #[cfg(feature = "copy-from")]
30 borrowed_root_entries: None,
31 })
32 }
33
34 pub const fn root_paddr(&self) -> PhysAddr {
35 self.inner.root.paddr
36 }
37
38 pub fn clone_missing_root_entries_from(
52 &mut self,
53 other: &PageTableRef<T, A>,
54 start_vaddr: VirtAddr,
55 size: usize,
56 ) -> PagingResult {
57 let Some(entries) = Self::root_entry_range(start_vaddr, size)? else {
58 return Ok(());
59 };
60
61 let root_level = Frame::<T, A>::PT_LEVEL;
62 let mut changed = false;
63 for index in entries {
64 changed |= self
65 .inner
66 .root
67 .clone_entry_from(&other.root, index, root_level)?;
68 }
69 if changed {
70 T::flush(None);
71 }
72 Ok(())
73 }
74
75 #[cfg(feature = "copy-from")]
86 pub unsafe fn share_root_entries_from(
87 &mut self,
88 other: &Self,
89 start_vaddr: VirtAddr,
90 size: usize,
91 ) -> PagingResult {
92 if size == 0 {
93 return Ok(());
94 }
95 if self.borrowed_root_entries.is_some() {
96 return Err(PagingError::hierarchy_error(
97 "Page table already contains shared root entries",
98 ));
99 }
100
101 let Some(entries) = Self::root_entry_range(start_vaddr, size)? else {
102 return Ok(());
103 };
104 let root_level = Frame::<T, A>::PT_LEVEL;
105
106 for index in entries.clone() {
107 self.inner.root.dealloc_entry_recursive(index, root_level);
108 self.inner.root.as_slice_mut()[index] = other.inner.root.as_slice()[index];
109 }
110 self.borrowed_root_entries = Some(entries);
111 T::flush(None);
112 Ok(())
113 }
114
115 fn root_entry_range(start_vaddr: VirtAddr, size: usize) -> PagingResult<Option<Range<usize>>> {
116 if size == 0 {
117 return Ok(None);
118 }
119 let end_vaddr = start_vaddr
120 .as_usize()
121 .checked_add(size)
122 .ok_or_else(|| PagingError::address_overflow("root_entry_range"))?;
123 let root_level = Frame::<T, A>::PT_LEVEL;
124 let start_index = Frame::<T, A>::virt_to_index(start_vaddr, root_level);
125 let end_index =
126 Frame::<T, A>::virt_to_index(VirtAddr::from_usize(end_vaddr - 1), root_level) + 1;
127 if start_index >= end_index {
128 return Err(PagingError::invalid_range(
129 "Range must be contiguous in the root page table",
130 ));
131 }
132 Ok(Some(start_index..end_index))
133 }
134
135 #[cfg(feature = "copy-from")]
136 fn detach_borrowed_root_entries(&mut self) {
137 let Some(entries) = self.borrowed_root_entries.take() else {
138 return;
139 };
140 for index in entries {
141 self.inner.root.as_slice_mut()[index].clear();
142 }
143 }
144}
145
146impl<T: TableMeta, A: FrameAllocator> Drop for PageTable<T, A> {
147 fn drop(&mut self) {
148 #[cfg(feature = "copy-from")]
149 self.detach_borrowed_root_entries();
150 unsafe {
151 self.deallocate();
153 }
154 }
155}
156
157impl<T: TableMeta, A: FrameAllocator> Deref for PageTable<T, A> {
158 type Target = PageTableRef<T, A>;
159
160 fn deref(&self) -> &Self::Target {
161 &self.inner
162 }
163}
164
165impl<T: TableMeta, A: FrameAllocator> DerefMut for PageTable<T, A> {
166 fn deref_mut(&mut self) -> &mut Self::Target {
167 &mut self.inner
168 }
169}
170
171#[derive(Clone, Copy)]
172pub struct PageTableRef<T: TableMeta, A: FrameAllocator> {
173 pub root: Frame<T, A>,
174}
175
176impl<T: TableMeta, A: FrameAllocator> core::fmt::Debug for PageTableRef<T, A>
177where
178 T::P: core::fmt::Debug,
179{
180 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
181 f.debug_struct("PageTable")
182 .field(
183 "root_paddr",
184 &format_args!("{:#x}", self.root.paddr.as_usize()),
185 )
186 .field("table_levels", &T::LEVEL_BITS.len())
187 .field("max_block_level", &T::MAX_BLOCK_LEVEL)
188 .field("page_size", &format_args!("{:#x}", T::PAGE_SIZE))
189 .finish()
190 }
191}
192
193impl<T: TableMeta, A: FrameAllocator> PageTableRef<T, A> {
194 pub unsafe fn new(allocator: A) -> PagingResult<Self> {
200 let root = Frame::new_root(allocator)?;
201 Ok(Self { root })
202 }
203
204 pub fn from_paddr(paddr: PhysAddr, allocator: A) -> Self {
205 let root = Frame::from_root_paddr(paddr, allocator);
206 Self { root }
207 }
208
209 pub fn map_page(
211 &mut self,
212 vaddr: VirtAddr,
213 paddr: PhysAddr,
214 page_size: usize,
215 config: PteConfigOf<T>,
216 ) -> PagingResult {
217 let Some(level) = Frame::<T, A>::level_for_page_size(page_size) else {
218 return Err(PagingError::invalid_size(
219 "Page size is not represented by the page-table levels",
220 ));
221 };
222 if level > 1 && level > T::MAX_BLOCK_LEVEL {
223 return Err(PagingError::invalid_size(
224 "Page size exceeds the architecture's block-mapping level",
225 ));
226 }
227 self.map(&MapConfig {
228 vaddr: vaddr.align_down(page_size),
229 paddr: paddr.align_down(page_size),
230 size: page_size,
231 pte: config,
232 allow_huge: level > 1,
233 flush: true,
234 })
235 }
236
237 pub fn map_region(
241 &mut self,
242 start_vaddr: VirtAddr,
243 get_paddr: impl Fn(VirtAddr) -> PhysAddr,
244 size: usize,
245 config: PteConfigOf<T>,
246 allow_huge: bool,
247 ) -> PagingResult {
248 if size == 0 {
249 return Err(PagingError::invalid_size("Region size cannot be zero"));
250 }
251 if !start_vaddr.as_usize().is_multiple_of(T::PAGE_SIZE)
252 || !size.is_multiple_of(T::PAGE_SIZE)
253 {
254 return Err(PagingError::alignment_error(
255 "Region start and size must be base-page aligned",
256 ));
257 }
258 start_vaddr.as_usize().checked_add(size).ok_or_else(|| {
259 PagingError::address_overflow("Virtual address overflow in map_region")
260 })?;
261 self.validate_address_width(start_vaddr, size, "map_region")?;
262
263 let mut offset = 0;
264 let mut flush_addrs = heapless::Vec::<VirtAddr, TARGETED_FLUSH_LIMIT>::new();
265 let mut full_flush = false;
266 let result = loop {
267 if offset >= size {
268 break Ok(());
269 }
270 let vaddr = start_vaddr + offset;
271 let paddr = get_paddr(vaddr);
272 let remaining = size - offset;
273 let page_size = largest_page_size::<T, A>(vaddr, paddr, remaining, allow_huge);
274 if let Err(err) = self.map(&MapConfig {
275 vaddr: vaddr.align_down(page_size),
276 paddr: paddr.align_down(page_size),
277 size: page_size,
278 pte: config,
279 allow_huge: page_size > T::PAGE_SIZE,
280 flush: false,
281 }) {
282 let rollback_result = if offset == 0 {
283 Ok(())
284 } else {
285 self.unmap_with_config(&UnmapConfig {
286 start_vaddr,
287 size: offset,
288 flush: false,
289 })
290 };
291 break match rollback_result {
292 Ok(()) => Err(err),
293 Err(rollback_err) => Err(rollback_err),
294 };
295 }
296 if !full_flush && flush_addrs.push(vaddr).is_err() {
297 full_flush = true;
298 flush_addrs.clear();
299 }
300 offset += page_size;
301 };
302
303 if full_flush {
304 T::flush(None);
305 } else {
306 for vaddr in flush_addrs {
307 T::flush(Some(vaddr));
308 }
309 }
310 result
311 }
312
313 pub fn unmap_page(
315 &mut self,
316 vaddr: VirtAddr,
317 ) -> PagingResult<(PhysAddr, PteConfigOf<T>, usize)> {
318 let (pte, level) = self
319 .root
320 .find_occupied_leaf(vaddr, Frame::<T, A>::PT_LEVEL)?;
321 let page_size = Frame::<T, A>::level_size(level);
322 let is_dir = level > 1;
323 let paddr = pte.paddr(is_dir);
324 let config = pte.config(is_dir);
325 self.unmap_with_config(&UnmapConfig {
326 start_vaddr: vaddr.align_down(page_size),
327 size: page_size,
328 flush: true,
329 })?;
330 Ok((paddr, config, page_size))
331 }
332
333 pub fn protect_page(&mut self, vaddr: VirtAddr, config: PteConfigOf<T>) -> PagingResult<usize> {
335 let page_size = self
336 .root
337 .protect_recursive(vaddr, config, Frame::<T, A>::PT_LEVEL)?;
338 T::flush(Some(vaddr));
339 Ok(page_size)
340 }
341
342 pub fn protect_region(
344 &mut self,
345 start_vaddr: VirtAddr,
346 size: usize,
347 config: PteConfigOf<T>,
348 ) -> PagingResult {
349 let end = start_vaddr
350 .as_usize()
351 .checked_add(size)
352 .ok_or_else(|| PagingError::address_overflow("protect_region"))?;
353 if size == 0 {
354 return Ok(());
355 }
356
357 self.root
360 .split_leaf_for_boundary(start_vaddr, start_vaddr, Frame::<T, A>::PT_LEVEL)?;
361 self.root.split_leaf_for_boundary(
362 VirtAddr::from_usize(end - 1),
363 VirtAddr::from_usize(end),
364 Frame::<T, A>::PT_LEVEL,
365 )?;
366
367 let mut vaddr = start_vaddr;
368 while vaddr.as_usize() < end {
369 match self.protect_page(vaddr, config) {
370 Ok(page_size) => vaddr += page_size,
371 Err(PagingError::NotMapped) => vaddr += T::PAGE_SIZE,
372 Err(err) => return Err(err),
373 }
374 }
375 Ok(())
376 }
377
378 pub fn remap_page(
380 &mut self,
381 vaddr: VirtAddr,
382 paddr: PhysAddr,
383 config: PteConfigOf<T>,
384 ) -> PagingResult<usize> {
385 let page_size = self
386 .root
387 .remap_recursive(vaddr, paddr, config, Frame::<T, A>::PT_LEVEL)?;
388 T::flush(Some(vaddr));
389 Ok(page_size)
390 }
391
392 pub fn query(&self, vaddr: VirtAddr) -> PagingResult<(PhysAddr, PteConfigOf<T>, usize)> {
394 let (paddr, pte, level) = self.translate_with_level(vaddr)?;
395 Ok((
396 paddr,
397 pte.config(level > 1),
398 Frame::<T, A>::level_size(level),
399 ))
400 }
401
402 pub fn map(&mut self, config: &MapConfig<PteConfigOf<T>>) -> PagingResult {
404 self.validate_map_config(config)?;
406
407 if config.vaddr.as_usize().checked_add(config.size).is_none()
409 || config.paddr.as_usize().checked_add(config.size).is_none()
410 {
411 return Err(PagingError::address_overflow(
412 "Virtual or physical address overflow",
413 ));
414 }
415 self.validate_address_width(config.vaddr, config.size, "map")?;
416
417 self.root.map_range_recursive(MapRecursiveConfig {
418 start_vaddr: config.vaddr,
419 start_paddr: config.paddr,
420 end_vaddr: config.vaddr + config.size,
421 level: Frame::<T, A>::PT_LEVEL,
422 allow_huge: config.allow_huge,
423 flush: config.flush,
424 pte_template: config.pte,
425 })?;
426
427 Ok(())
428 }
429
430 pub fn unmap(&mut self, start_vaddr: VirtAddr, size: usize) -> PagingResult<()> {
446 self.validate_unmap_params(start_vaddr, size)?;
448
449 let end_vaddr: VirtAddr = match start_vaddr.as_usize().checked_add(size) {
451 Some(end) => VirtAddr::from_usize(end),
452 None => {
453 return Err(PagingError::address_overflow(
454 "Virtual address overflow in unmap",
455 ));
456 }
457 };
458 self.validate_address_width(start_vaddr, size, "unmap")?;
459
460 self.root.unmap_range_recursive(UnmapRecursiveConfig {
461 start_vaddr,
462 end_vaddr,
463 level: Frame::<T, A>::PT_LEVEL,
464 flush: true, })?;
466
467 Ok(())
468 }
469
470 pub fn unmap_with_config(&mut self, config: &UnmapConfig) -> PagingResult<()> {
472 self.validate_unmap_params(config.start_vaddr, config.size)?;
473
474 let end_vaddr = match config.start_vaddr.as_usize().checked_add(config.size) {
475 Some(end) => VirtAddr::from_usize(end),
476 None => {
477 return Err(PagingError::address_overflow(
478 "Virtual address overflow in unmap_with_config",
479 ));
480 }
481 };
482 self.validate_address_width(config.start_vaddr, config.size, "unmap_with_config")?;
483
484 self.root.unmap_range_recursive(UnmapRecursiveConfig {
485 start_vaddr: config.start_vaddr,
486 end_vaddr,
487 level: Frame::<T, A>::PT_LEVEL,
488 flush: config.flush,
489 })?;
490
491 Ok(())
492 }
493
494 fn validate_unmap_params(&self, start_vaddr: VirtAddr, size: usize) -> PagingResult<()> {
496 if size == 0 {
497 return Err(PagingError::invalid_size("Size cannot be zero in unmap"));
498 }
499
500 if !start_vaddr.as_usize().is_multiple_of(T::PAGE_SIZE) {
502 return Err(PagingError::alignment_error(
503 "Start virtual address not page aligned in unmap",
504 ));
505 }
506
507 if !size.is_multiple_of(T::PAGE_SIZE) {
509 return Err(PagingError::alignment_error(
510 "Size not page aligned in unmap",
511 ));
512 }
513
514 Ok(())
515 }
516
517 pub fn walk_all(&self, config: WalkConfig) -> PageTableWalker<'_, T, A> {
519 PageTableWalker::new(self, config)
520 }
521
522 pub fn walk(
523 &self,
524 start_vaddr: VirtAddr,
525 end_vaddr: VirtAddr,
526 ) -> impl Iterator<Item = crate::walk::PteInfo<T::P>> + '_ {
527 let config = WalkConfig {
528 start_vaddr,
529 end_vaddr,
530 };
531 PageTableWalker::new(self, config).filter(|p| p.pte.present())
532 }
533
534 pub fn walk_valid(&self) -> impl Iterator<Item = crate::walk::PteInfo<T::P>> + '_ {
536 self.walk(0.into(), usize::MAX.into())
537 .filter(|p| p.pte.present() && p.is_final_mapping)
538 }
539
540 fn validate_map_config(&self, config: &MapConfig<PteConfigOf<T>>) -> PagingResult {
542 if config.size == 0 {
543 return Err(PagingError::invalid_size("Size cannot be zero"));
544 }
545
546 if !config.vaddr.as_usize().is_multiple_of(T::PAGE_SIZE) {
548 return Err(PagingError::alignment_error(
549 "Virtual address not page aligned",
550 ));
551 }
552
553 if !config.paddr.as_usize().is_multiple_of(T::PAGE_SIZE) {
554 return Err(PagingError::alignment_error(
555 "Physical address not page aligned",
556 ));
557 }
558
559 Ok(())
560 }
561
562 fn validate_address_width(
563 &self,
564 start_vaddr: VirtAddr,
565 size: usize,
566 operation: &'static str,
567 ) -> PagingResult<()> {
568 if !T::STRICT_ADDRESS_WIDTH {
569 return Ok(());
570 }
571 let Some(end) = start_vaddr.as_usize().checked_add(size) else {
572 return Err(PagingError::address_overflow(
573 "Virtual address range overflow",
574 ));
575 };
576 let last = end.saturating_sub(1);
577 if !Self::is_addr_in_width(start_vaddr.as_usize()) || !Self::is_addr_in_width(last) {
578 return Err(PagingError::address_overflow(operation));
579 }
580 Ok(())
581 }
582
583 pub const fn page_size() -> usize {
584 T::PAGE_SIZE
585 }
586
587 pub const fn table_levels() -> usize {
588 T::LEVEL_BITS.len()
589 }
590
591 pub const fn valid_bits() -> usize {
592 Frame::<T, A>::PT_VALID_BITS
593 }
594
595 fn is_addr_in_width(addr: usize) -> bool {
596 let valid_bits = Self::valid_bits();
597 if valid_bits >= usize::BITS as usize {
598 return true;
599 }
600 addr < (1usize << valid_bits)
601 }
602
603 pub unsafe fn destroy(mut self) {
616 self.root.deallocate_recursive(Frame::<T, A>::PT_LEVEL);
617 }
618
619 pub unsafe fn deallocate(&mut self) {
634 self.root.deallocate_recursive(Frame::<T, A>::PT_LEVEL);
635 }
636
637 pub fn deallocate_range(&mut self, start_vaddr: VirtAddr, end_vaddr: VirtAddr) -> PagingResult {
642 if start_vaddr >= end_vaddr {
643 return Err(PagingError::invalid_range(
644 "Start address must be less than end address",
645 ));
646 }
647
648 Ok(())
655 }
656
657 pub fn translate(&self, vaddr: VirtAddr) -> PagingResult<(PhysAddr, T::P)> {
669 self.translate_with_level(vaddr)
670 .map(|(phys_addr, pte, _)| (phys_addr, pte))
671 }
672
673 pub fn translate_with_level(&self, vaddr: VirtAddr) -> PagingResult<(PhysAddr, T::P, usize)> {
675 if T::STRICT_ADDRESS_WIDTH && !Self::is_addr_in_width(vaddr.as_usize()) {
676 return Err(PagingError::address_overflow("translate"));
677 }
678
679 let (pte, level) = self
680 .root
681 .translate_recursive_with_level(vaddr, Frame::<T, A>::PT_LEVEL)?;
682
683 let is_huge = pte.huge(level > 1);
684 let pte_paddr = pte.paddr(level > 1);
685
686 let (phys_addr, _) = if is_huge {
688 let level_size = Frame::<T, A>::level_size(level);
690 let offset_in_page = vaddr.as_usize() % level_size;
691 (
692 PhysAddr::from_usize(pte_paddr.as_usize() + offset_in_page),
693 level_size,
694 )
695 } else {
696 let offset_in_page = vaddr.as_usize() % T::PAGE_SIZE;
698 (
699 PhysAddr::from_usize(pte_paddr.as_usize() + offset_in_page),
700 T::PAGE_SIZE,
701 )
702 };
703
704 Ok((phys_addr, pte, level))
705 }
706
707 pub fn translate_phys(&self, vaddr: VirtAddr) -> PagingResult<PhysAddr> {
716 let (p, _) = self.translate(vaddr)?;
717 Ok(p)
718 }
719
720 pub fn is_mapped(&self, vaddr: VirtAddr) -> bool {
731 self.translate(vaddr).is_ok()
732 }
733
734 pub fn root_paddr(&self) -> crate::PhysAddr {
736 self.root.paddr
737 }
738}
739
740fn largest_page_size<T: TableMeta, A: FrameAllocator>(
741 vaddr: VirtAddr,
742 paddr: PhysAddr,
743 remaining: usize,
744 allow_huge: bool,
745) -> usize {
746 if allow_huge {
747 let max_level = Frame::<T, A>::PT_LEVEL.min(T::MAX_BLOCK_LEVEL);
748 for level in (2..=max_level).rev() {
749 let page_size = Frame::<T, A>::level_size(level);
750 if vaddr.is_aligned(page_size) && paddr.is_aligned(page_size) && remaining >= page_size
751 {
752 return page_size;
753 }
754 }
755 }
756 T::PAGE_SIZE
757}