Reorganize dtb.rs

Author: Josen-B

modules/axhal/src/dtb.rs

@@ -1,15 +1,30 @@
 //! DTB (Device Tree Blob) related functionality.
-use fdt_parser::{Fdt, FdtHeader};
+extern crate alloc;
 
+use alloc::boxed::Box;
 use core::fmt::Write;
+use fdt_parser::{Fdt, FdtHeader};
 use lazyinit::LazyInit;
 
-static BOOTARGS_STR: LazyInit<heapless::String<256>> = LazyInit::new();
+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();
 
-/// Get the bootargs from the device tree.
-pub fn bootargs_message() -> Option<&'static str> {
-    let fdt_addr = crate::get_bootarg();
+/// 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 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;
     }
@@ -33,18 +48,38 @@ pub fn bootargs_message() -> Option<&'static str> {
         core::slice::from_raw_parts(ptr, size)
     };
 
-    let fdt = match Fdt::from_bytes(fdt_bytes) {
-        Ok(fdt) => fdt,
-        Err(_) => return None,
-    };
+    // Leak the FDT bytes to create a 'static reference
+    let fdt_bytes_static = Box::leak(fdt_bytes.to_vec().into_boxed_slice());
+
+    match Fdt::from_bytes(fdt_bytes_static) {
+        Ok(fdt) => {
+            // Store the FDT in the cache
+            CACHED_FDT.init_once(fdt);
+            CACHED_FDT.get()
+        }
+        Err(_) => None,
+    }
+}
+
+/// 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
-}