address feedback from Chris

Author: babayet2

Cargo.lock

@@ -5293,7 +5293,9 @@ name = "page_table"
 version = "0.0.0"
 dependencies = [
  "bitfield-struct 0.11.0",
+ "static_assertions",
  "thiserror 2.0.16",
+ "tracing",
  "zerocopy 0.8.25",
 ]
 

tmk/tmk_vmm/src/load.rs

@@ -45,11 +45,10 @@ pub fn load_x86(
     let page_table_builder = page_table::x64::IdentityMapBuilder::new(
         page_table_base,
         page_table::IdentityMapSize::Size4Gb,
-    );
-    let page_tables = page_table_builder.build(
         page_table_work_buffer.as_mut_slice(),
         page_tables.as_mut_slice(),
     )?;
+    let page_tables = page_table_builder.build();
     loader
         .import_pages(
             page_table_base >> 12,

vm/loader/Cargo.toml

@@ -11,7 +11,7 @@ aarch64defs.workspace = true
 igvm.workspace = true
 loader_defs.workspace = true
 memory_range.workspace = true
-page_table.workspace = true
+page_table = {workspace = true, features = ["tracing"] }
 hvdef.workspace = true
 vm_topology.workspace = true
 x86defs.workspace = true

vm/loader/manifests/openhcl-x64-cvm-dev.json

@@ -57,4 +57,4 @@
             }
         }
     ]
-}
+}

vm/loader/page_table/Cargo.toml

@@ -9,6 +9,13 @@ rust-version.workspace = true
 [dependencies]
 bitfield-struct.workspace = true
 thiserror.workspace = true
+static_assertions.workspace = true
+tracing = { workspace = true, optional = true }
 zerocopy.workspace = true
+
+[features]
+tracing = ["dep:tracing", "std"]
+std = []
+
 [lints]
 workspace = true

vm/loader/page_table/src/aarch64.rs

@@ -437,6 +437,11 @@ impl<'a> Arm64PageTableSpace<'a> {
         debug_assert!(aligned(phys_table_start, Arm64PageSize::Small));
         debug_assert!(index < PAGE_SIZE_4K as usize / size_of::<Arm64PageTableEntry>());
 
+        #[cfg(feature = "tracing")]
+        tracing::debug!(
+            "Writing page table entry {entry:#016x}, index {index:#x}, table {phys_table_start:#x}"
+        );
+
         let pos = phys_table_start as usize - self.phys_page_table_root
             + index * size_of::<Arm64PageTableEntry>();
         self.space[pos..pos + 8].copy_from_slice(&entry.to_le_bytes());
@@ -689,6 +694,11 @@ pub fn build_identity_page_tables_aarch64(
         panic!("size not 2mb aligned");
     }
 
+    #[cfg(feature = "tracing")]
+    tracing::debug!(
+        "Creating Aarch64 page tables at {page_table_gpa:#x} mapping starting at {start_gpa:#x} of size {size} bytes"
+    );
+
     let mut page_tables =
         Arm64PageTableSpace::new(page_table_gpa as usize, page_table_space).unwrap();
     page_tables
@@ -705,12 +715,18 @@ pub fn build_identity_page_tables_aarch64(
 
     let used_space = page_tables.used_space();
 
+    #[cfg(feature = "tracing")]
+    {
+        tracing::debug!("Page tables use {used_space} bytes");
+        tracing::debug!("Page tables stats by level: {:?}", page_tables.lvl_stats());
+    }
+
     &page_table_space[0..used_space]
 }
 
 #[cfg(test)]
 mod tests {
-    extern crate std;
+    use std;
 
     use super::*;
     use std::vec;

vm/loader/page_table/src/lib.rs

@@ -3,7 +3,7 @@
 
 //! Methods to construct page tables.
 
-#![no_std]
+#![cfg_attr(not(feature = "std"), no_std)]
 #![expect(missing_docs)]
 #![forbid(unsafe_code)]
 
@@ -17,9 +17,29 @@ use thiserror::Error;
 pub enum Error {
     /// The PageTableBuilder bytes buffer does not match the size of the struct buffer
     #[error(
-        "PageTableBuilder bytes buffer size [{bytes_buf}] does not match the struct buffer size [{struct_buf}]"
+        "PageTableBuilder bytes buffer size {bytes_buf} does not match the struct buffer size [{struct_buf}]"
     )]
     BadBufferSize { bytes_buf: usize, struct_buf: usize },
+
+    /// The page table mapping size is not 2MB aligned
+    #[error("page table mapping size {0:#x} is not 2MB-aligned")]
+    SizeAlignment(u64),
+
+    /// The page table mapping size is not 2MB aligned
+    #[error("start_gpa {0:#x} is not 2MB aligned")]
+    StartAlignment(u64),
+
+    /// The page table mapping size is greater than 512GB
+    #[error("size {0:#x} is larger than 512GB")]
+    MappingSize(u64),
+
+    /// The page table mapping size is missing
+    #[error("the page table builder was invoked without a mapping size")]
+    MissingSize,
+
+    /// The page table builder is generating overlapping mappings
+    #[error("the page table builder was invoked without a mapping size")]
+    OverlappingMappings,
 }
 
 /// Size of the initial identity map