[vmbase] Update return type of layout functions to VirtualAddress

Test: atest vmbase_example.integration_test rialto_test
Test: m pvmfw_img
Change-Id: If33b0cc0bfffbe408908c382d0c077c92f3c197b
diff --git a/pvmfw/src/memory.rs b/pvmfw/src/memory.rs
index 5930ec9..27ab719 100644
--- a/pvmfw/src/memory.rs
+++ b/pvmfw/src/memory.rs
@@ -30,7 +30,7 @@
 
 /// Returns memory range reserved for the appended payload.
 pub fn appended_payload_range() -> Range<VirtualAddress> {
-    let start = align_up(layout::binary_end(), SIZE_4KB).unwrap();
+    let start = align_up(layout::binary_end().0, SIZE_4KB).unwrap();
     // pvmfw is contained in a 2MiB region so the payload can't be larger than the 2MiB alignment.
     let end = align_up(start, SIZE_2MB).unwrap();
     VirtualAddress(start)..VirtualAddress(end)
diff --git a/vmbase/example/src/layout.rs b/vmbase/example/src/layout.rs
index 4d14b1c..b032a30 100644
--- a/vmbase/example/src/layout.rs
+++ b/vmbase/example/src/layout.rs
@@ -29,14 +29,6 @@
     layout::stack_range(40 * PAGE_SIZE)
 }
 
-fn data_load_address() -> VirtualAddress {
-    VirtualAddress(layout::data_load_address())
-}
-
-fn binary_end() -> VirtualAddress {
-    VirtualAddress(layout::binary_end())
-}
-
 pub fn print_addresses() {
     let dtb = layout::dtb_range();
     info!("dtb:        {}..{} ({} bytes)", dtb.start, dtb.end, dtb.end - dtb.start);
@@ -44,14 +36,14 @@
     info!("text:       {}..{} ({} bytes)", text.start, text.end, text.end - text.start);
     let rodata = layout::rodata_range();
     info!("rodata:     {}..{} ({} bytes)", rodata.start, rodata.end, rodata.end - rodata.start);
-    info!("binary end: {}", binary_end());
+    info!("binary end: {}", layout::binary_end());
     let data = layout::data_range();
     info!(
         "data:       {}..{} ({} bytes, loaded at {})",
         data.start,
         data.end,
         data.end - data.start,
-        data_load_address(),
+        layout::data_load_address(),
     );
     let bss = layout::bss_range();
     info!("bss:        {}..{} ({} bytes)", bss.start, bss.end, bss.end - bss.start);
diff --git a/vmbase/src/layout/mod.rs b/vmbase/src/layout/mod.rs
index ffa29e7..f7e8170 100644
--- a/vmbase/src/layout/mod.rs
+++ b/vmbase/src/layout/mod.rs
@@ -31,7 +31,8 @@
     ($symbol:ident) => {{
         // SAFETY: We're just getting the address of an extern static symbol provided by the linker,
         // not dereferencing it.
-        unsafe { addr_of!($crate::linker::$symbol) as usize }
+        let addr = unsafe { addr_of!($crate::linker::$symbol) as usize };
+        VirtualAddress(addr)
     }};
 }
 
@@ -42,7 +43,7 @@
         let start = linker_addr!($begin);
         let end = linker_addr!($end);
 
-        VirtualAddress(start)..VirtualAddress(end)
+        start..end
     }};
 }
 
@@ -74,10 +75,10 @@
 /// Writable data region for the stack.
 pub fn stack_range(stack_size: usize) -> Range<VirtualAddress> {
     let end = linker_addr!(init_stack_pointer);
-    let start = end.checked_sub(stack_size).unwrap();
+    let start = VirtualAddress(end.0.checked_sub(stack_size).unwrap());
     assert!(start >= linker_addr!(stack_limit));
 
-    VirtualAddress(start)..VirtualAddress(end)
+    start..end
 }
 
 /// All writable sections, excluding the stack.
@@ -93,12 +94,12 @@
 }
 
 /// Read-write data (original).
-pub fn data_load_address() -> usize {
+pub fn data_load_address() -> VirtualAddress {
     linker_addr!(data_lma)
 }
 
 /// End of the binary image.
-pub fn binary_end() -> usize {
+pub fn binary_end() -> VirtualAddress {
     linker_addr!(bin_end)
 }