Reorganize dtb.rs

Author: Josen-B

modules/axhal/src/dtb.rs

@@ -1,50 +1,81 @@
 //! DTB (Device Tree Blob) related functionality.
-use fdt_parser::{Fdt, FdtHeader};
-
 use core::fmt::Write;
+use fdt_parser::{Fdt, FdtHeader};
 use lazyinit::LazyInit;
 
-static BOOTARGS_STR: LazyInit<heapless::String<256>> = LazyInit::new();
+/// Record address of the boot argument (DTB address).
+pub static BOOT_ARG: LazyInit<usize> = LazyInit::new();
+static CACHED_FDT: LazyInit<Fdt<'static>> = LazyInit::new();
+static BOOTARGS_STR: LazyInit<heapless::String<1024>> = LazyInit::new();
+
+/// Returns the boot argument.
+/// This is typically the device tree blob address passed from the bootloader.
+pub fn get_bootarg() -> usize {
+    *BOOT_ARG
+}
 
-/// Get the bootargs from the device tree.
-pub fn bootargs_message() -> Option<&'static str> {
-    let fdt_addr = crate::get_bootarg();
+/// Get the cached FDT or initialize it if not already done.
+fn get_fdt() -> Option<&'static Fdt<'static>> {
+    // Return cached FDT if available
+    if let Some(fdt) = CACHED_FDT.get() {
+        return Some(fdt);
+    }
 
+    // Parse and cache the FDT
+    let fdt_addr = get_bootarg();
     if fdt_addr == 0 {
         return None;
     }
 
     let virt_addr = crate::mem::phys_to_virt(crate::mem::PhysAddr::from(fdt_addr)).as_usize();
 
-    let fdt_header = unsafe {
+    let fdt = unsafe {
+        // First read the header to get the size
         let header_size = core::mem::size_of::<FdtHeader>();
-        let ptr = virt_addr as *const u8;
-        core::slice::from_raw_parts(ptr, header_size)
-    };
+        let header_ptr = virt_addr as *const u8;
+        let header_slice = core::slice::from_raw_parts(header_ptr, header_size);
 
-    let fdt_header = match FdtHeader::from_bytes(fdt_header) {
-        Ok(header) => header,
-        Err(_) => return None,
-    };
+        let fdt_header = match FdtHeader::from_bytes(header_slice) {
+            Ok(header) => header,
+            Err(_) => return None,
+        };
 
-    let fdt_bytes = unsafe {
-        let ptr = virt_addr as *const u8;
         let size = fdt_header.total_size() as usize;
-        core::slice::from_raw_parts(ptr, size)
-    };
+        let fdt_ptr = virt_addr as *const u8;
+        let fdt_slice = core::slice::from_raw_parts(fdt_ptr, size);
+
+        let fdt_slice_static = core::mem::transmute::<&[u8], &'static [u8]>(fdt_slice);
 
-    let fdt = match Fdt::from_bytes(fdt_bytes) {
-        Ok(fdt) => fdt,
-        Err(_) => return None,
+        match Fdt::from_bytes(fdt_slice_static) {
+            Ok(fdt) => fdt,
+            Err(_) => return None,
+        }
     };
 
+    // Store the FDT in the cache
+    CACHED_FDT.init_once(fdt);
+    CACHED_FDT.get()
+}
+
+/// Get the bootargs chosen from the device tree.
+pub fn get_chosen() -> Option<&'static str> {
+    // If bootargs are already cached, return them
+    if let Some(bootargs) = BOOTARGS_STR.get() {
+        return Some(bootargs.as_str());
+    }
+
+    // Get or initialize the cached FDT
+    let fdt = get_fdt()?;
+
     if let Some(chosen) = fdt.chosen() {
         if let Some(bootargs) = chosen.bootargs() {
             // Store bootargs in static variable
-            let mut bootargs_str = heapless::String::<256>::new();
+            let mut bootargs_str = heapless::String::<1024>::new();
             if write!(bootargs_str, "{}", bootargs).is_ok() {
                 BOOTARGS_STR.init_once(bootargs_str);
                 return Some(BOOTARGS_STR.as_str());
+            } else {
+                warn!("Failed to write bootargs");
             }
         }
     }

modules/axhal/src/lib.rs

@@ -109,19 +109,9 @@ pub use axplat::init::init_later;
 #[cfg(feature = "smp")]
 pub use axplat::init::{init_early_secondary, init_later_secondary};
 
-use lazyinit::LazyInit;
-
-static BOOT_ARG: LazyInit<usize> = LazyInit::new();
-
 /// Initializes the platform and boot argument.
 /// This function should be called as early as possible.
 pub fn init_early(cpu_id: usize, arg: usize) {
-    BOOT_ARG.init_once(arg);
+    dtb::BOOT_ARG.init_once(arg);
     axplat::init::init_early(cpu_id, arg);
 }
-
-/// Returns the boot argument.
-/// This is typically the device tree blob address passed from the bootloader.
-pub fn get_bootarg() -> usize {
-    *BOOT_ARG
-}