Bug fix for get_disks()

In rare circumstances we could be omitting drives from the list

Cargo.lock

@@ -467,7 +467,6 @@ dependencies = [
  "json5",
  "lazy_static",
  "libc",
- "once_cell",
  "pretty_assertions",
  "rand",
  "ratatui",

Cargo.toml

@@ -15,5 +15,5 @@
 
 [workspace]
 members = ["core", "boot_diags", "deja_vu", "pe_menu"]
-default-members = ["boot_diags", "deja_vu", "pe_menu"]
+default-members = ["core", "boot_diags", "deja_vu", "pe_menu"]
 resolver = "2"

core/Cargo.toml

@@ -40,7 +40,6 @@ human-panic = "2.0.1"
 json5 = "0.4.1"
 lazy_static = "1.5.0"
 libc = "0.2.158"
-once_cell = "1.20.2"
 pretty_assertions = "1.4.0"
 rand = "0.9.0"
 ratatui = { version = "0.29.0", features = ["serde", "macros"] }

core/src/lib.rs

@@ -23,4 +23,5 @@ pub mod logging;
 pub mod state;
 pub mod system;
 pub mod tasks;
+pub mod tests;
 pub mod tui;

core/src/system/disk.rs

@@ -17,12 +17,12 @@ use serde::{Deserialize, Serialize};
 use std::{
     fmt,
     process::{Command, Stdio},
+    sync::OnceLock,
     thread::sleep,
     time::Duration,
 };
 use tracing::info;
 
-use once_cell::sync::Lazy;
 use regex::Regex;
 
 use crate::system::diskpart;
@@ -385,10 +385,11 @@ pub fn bytes_to_string(size: u64) -> String {
 ///
 /// Will panic if s is not simliar to 32B, 64MB, etc...
 pub fn string_to_bytes(s: &str) -> u64 {
-    static RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"(\d+)\s+(\w+)B").unwrap());
+    static RE_BYTES: OnceLock<Regex> = OnceLock::new();
+    let re_bytes = RE_BYTES.get_or_init(|| Regex::new(r"(\d+)\s+(\w+)B").unwrap());
     let base: u64 = 1024;
     let mut size: u64 = 0;
-    for (_, [size_str, suffix]) in RE.captures_iter(s).map(|c| c.extract()) {
+    for (_, [size_str, suffix]) in re_bytes.captures_iter(s).map(|c| c.extract()) {
         let x: u64 = size_str.parse().unwrap();
         size += x;
         match suffix {

core/src/system/diskpart.rs

@@ -19,9 +19,9 @@ use std::{
     fs::File,
     io::Write,
     process::{Command, Output, Stdio},
+    sync::OnceLock,
 };
 
-use once_cell::sync::Lazy;
 use regex::Regex;
 use tempfile::tempdir;
 use tracing::{info, warn};
@@ -32,13 +32,88 @@ use crate::system::disk::{
 
 static DEFAULT_MAX_DISKS: usize = 8;
 
+struct RegexList {
+    re_detail_all_disks: OnceLock<Regex>,
+    re_detail_disk: OnceLock<Regex>,
+    re_detail_partition: OnceLock<Regex>,
+    re_disk_numbers: OnceLock<Regex>,
+    re_list_disk: OnceLock<Regex>,
+    re_list_partition: OnceLock<Regex>,
+    re_list_volumes: OnceLock<Regex>,
+    re_split_all_disks: OnceLock<Regex>,
+    re_uuid: OnceLock<Regex>,
+}
+
+impl RegexList {
+    pub fn detail_all_disks(&self) -> &Regex {
+        self.re_detail_all_disks.get_or_init(|| {
+            Regex::new(r"(?s)Disk (\d+) is now the selected disk.*?\r?\n\s*\r?\n(.*)").unwrap()
+        })
+    }
+
+    pub fn detail_disk(&self) -> &Regex {
+        self.re_detail_disk.get_or_init(|| {
+            Regex::new(r"(.*?)\r?\nDisk ID\s*:\s+(.*?)\r?\nType\s*:\s+(.*?)\r?\n").unwrap()
+        })
+    }
+
+    pub fn detail_partition(&self) -> &Regex {
+        self.re_detail_partition
+            .get_or_init(|| Regex::new(r"Partition (\d+)\r?\nType\s*: (\S+)(\r?\n.*){5}\s*(Volume.*\r?\n.*\r?\n|There is no volume)(.*)").unwrap())
+    }
+
+    pub fn disk_numbers(&self) -> &Regex {
+        self.re_disk_numbers
+            .get_or_init(|| Regex::new(r"\s+Disk\s+(\d+).*\n.*\n.*\nDisk ID:").unwrap())
+    }
+
+    pub fn list_disk(&self) -> &Regex {
+        self.re_list_disk
+            .get_or_init(|| Regex::new(r"Disk\s+(\d+)\s+(\w+)\s+(\d+\s+\w+B)").unwrap())
+    }
+
+    pub fn list_partition(&self) -> &Regex {
+        self.re_list_partition
+            .get_or_init(|| Regex::new(r"Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+B)").unwrap())
+    }
+
+    pub fn split_all_disks(&self) -> &Regex {
+        self.re_split_all_disks
+            .get_or_init(|| Regex::new(r"Disk \d+ is now the selected disk").unwrap())
+    }
+
+    pub fn list_volumes(&self) -> &Regex {
+        self.re_list_volumes.get_or_init(|| {
+            //   Volume ###  Ltr  Label        Fs     Type        Size     Status     Info
+            //   ----------  ---  -----------  -----  ----------  -------  ---------  --------
+            // * Volume 1     S   ESP          FAT32  Partition    100 MB  Healthy    Hidden
+            Regex::new(r"..Volume (\d.{2})  (.{3})  (.{11})  (.{5})").unwrap()
+        })
+    }
+
+    pub fn uuid(&self) -> &Regex {
+        self.re_uuid.get_or_init(|| {
+            Regex::new(r"^\{[0-9A-F]{8}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{12}\}$")
+                .unwrap()
+        })
+    }
+}
+
+static REGEXES: RegexList = RegexList {
+    re_detail_all_disks: OnceLock::new(),
+    re_detail_disk: OnceLock::new(),
+    re_detail_partition: OnceLock::new(),
+    re_disk_numbers: OnceLock::new(),
+    re_list_disk: OnceLock::new(),
+    re_list_partition: OnceLock::new(),
+    re_list_volumes: OnceLock::new(),
+    re_split_all_disks: OnceLock::new(),
+    re_uuid: OnceLock::new(),
+};
+
 pub fn get_disk_details(disk_id: usize, disk_size: u64, disk_details: Option<&str>) -> Disk {
-    static RE_DETAILS: Lazy<Regex> = Lazy::new(|| {
-        Regex::new(r"(.*?)\r?\nDisk ID\s*:\s+(.*?)\r?\nType\s*:\s+(.*?)\r?\n").unwrap()
-    });
-    static RE_UUID: Lazy<Regex> = Lazy::new(|| {
-        Regex::new(r"^\{[0-9A-F]{8}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{12}\}$").unwrap()
-    });
+    let re_details = REGEXES.detail_disk();
+    let re_uuid = REGEXES.uuid();
     let mut disk = Disk {
         id: disk_id,
         size: disk_size,
@@ -56,11 +131,11 @@ pub fn get_disk_details(disk_id: usize, disk_size: u64, disk_details: Option<&st
 
     // Parse details
     for (_, [model, part_type, conn_type]) in
-        RE_DETAILS.captures_iter(&details).map(|c| c.extract())
+        re_details.captures_iter(&details).map(|c| c.extract())
     {
         disk.model = String::from(model);
         disk.conn_type = String::from(conn_type);
-        if RE_UUID.is_match(part_type) {
+        if re_uuid.is_match(part_type) {
             disk.part_type = PartitionTableType::Guid;
         } else {
             disk.part_type = PartitionTableType::Legacy;
@@ -77,8 +152,7 @@ pub fn get_partition_details(
     disk_details: Option<&str>,
     part_details: Option<&str>,
 ) -> Vec<Partition> {
-    static RE_LIS: Lazy<Regex> =
-        Lazy::new(|| Regex::new(r"Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+B)").unwrap());
+    let re_list_partitions = REGEXES.list_partition();
     let mut parts = Vec::new();
 
     // List partition
@@ -89,7 +163,10 @@ pub fn get_partition_details(
         let script = format!("select disk {disk_id}\r\nlist partition");
         contents = run_script(&script);
     };
-    for (_, [number, size]) in RE_LIS.captures_iter(&contents).map(|c| c.extract()) {
+    for (_, [number, size]) in re_list_partitions
+        .captures_iter(&contents)
+        .map(|c| c.extract())
+    {
         let part_num = number.parse().unwrap();
         if part_num != 0 {
             // part_num == 0 is reserved for extended partition "containers" so we can exclude them
@@ -120,7 +197,7 @@ pub fn get_partition_details(
         part_contents = String::from(details);
     } else {
         part_contents = run_script(script.join("\r\n").as_str());
-    };
+    }
     parse_partition_details(&mut parts, &part_contents);
 
     // Done
@@ -149,13 +226,12 @@ pub fn build_dest_format_script(disk_id: usize, part_type: &PartitionTableType)
 
 #[must_use]
 pub fn build_get_disk_script(disk_nums: Option<Vec<&str>>) -> String {
-    let capacity = DEFAULT_MAX_DISKS * 3 + 1;
     let script: String;
 
     // Get disk and partition details
     if let Some(disks) = disk_nums {
         // (Slower case)
-        let mut script_parts = Vec::with_capacity(capacity);
+        let mut script_parts = Vec::with_capacity(DEFAULT_MAX_DISKS * 3 + 1);
 
         // Get list of disks
         script_parts.push(String::from("list disk"));
@@ -171,7 +247,7 @@ pub fn build_get_disk_script(disk_nums: Option<Vec<&str>>) -> String {
         script = script_parts.join("\n");
     } else {
         // (Best case)
-        let mut script_parts = Vec::with_capacity(capacity);
+        let mut script_parts = Vec::with_capacity(DEFAULT_MAX_DISKS + 1);
 
         // Get list of disks
         script_parts.push("list disk");
@@ -181,9 +257,9 @@ pub fn build_get_disk_script(disk_nums: Option<Vec<&str>>) -> String {
         script_parts.push("detail disk");
         script_parts.push("list partition");
 
-        // Limit to 8 disks (if there's more the manual "worst" case will be used)
+        // Limit to DEFAULT_MAX_DISKS (if there's more the manual "worst" case will be used)
         let mut i = 0;
-        while i < 8 {
+        while i < DEFAULT_MAX_DISKS {
             script_parts.push("select disk next");
             script_parts.push("detail disk");
             script_parts.push("list partition");
@@ -198,11 +274,8 @@ pub fn build_get_disk_script(disk_nums: Option<Vec<&str>>) -> String {
 }
 
 pub fn get_disks() -> Vec<Disk> {
-    static RE_DIS_DET: Lazy<Regex> = Lazy::new(|| {
-        Regex::new(r"(?s)Disk (\d+) is now the selected disk.*?\r?\n\s*\r?\n(.*)").unwrap()
-    });
-    static RE_DIS_LIS: Lazy<Regex> =
-        Lazy::new(|| Regex::new(r"Disk\s+(\d+)\s+(\w+)\s+(\d+\s+\w+B)").unwrap());
+    let re_detail_all_disks = REGEXES.detail_all_disks();
+    let re_list_disk = REGEXES.list_disk();
     let mut contents: String;
     let mut output;
     let mut script: String;
@@ -213,7 +286,7 @@ pub fn get_disks() -> Vec<Disk> {
     contents = String::from_utf8_lossy(&output.stdout).to_string();
     if let Some(return_code) = output.status.code() {
         let disk_nums = parse_disk_numbers(&contents);
-        if return_code != 0 && !disk_nums.is_empty() {
+        if return_code != 0 && !disk_nums.is_empty() && disk_nums.len() != DEFAULT_MAX_DISKS {
             // The base assumptions were correct! skipping fallback method
             //
             // Since the return_code was not zero, and at least one disk was detected, that
@@ -244,7 +317,7 @@ pub fn get_disks() -> Vec<Disk> {
     //       i.e. 0, 1, 3, 4
     //       For instance, this can happen if a drive is disconnected after startup
     let mut disks_map: HashMap<&str, Disk> = HashMap::with_capacity(DEFAULT_MAX_DISKS);
-    for (_, [number, _status, size]) in RE_DIS_LIS
+    for (_, [number, _status, size]) in re_list_disk
         .captures_iter(dp_sections.remove(0)) // This is the "list disk" section
         .map(|c| c.extract())
     {
@@ -261,7 +334,10 @@ pub fn get_disks() -> Vec<Disk> {
     // Add Disk details
     let mut disks_raw: Vec<Disk> = Vec::with_capacity(DEFAULT_MAX_DISKS);
     for section in dp_sections {
-        for (_, [id, details]) in RE_DIS_DET.captures_iter(section).map(|c| c.extract()) {
+        for (_, [id, details]) in re_detail_all_disks
+            .captures_iter(section)
+            .map(|c| c.extract())
+        {
             if let Some(disk) = disks_map.remove(id) {
                 // We remove the disk from the HashMap because we're moving it to the Vec
                 let mut disk = get_disk_details(disk.id, disk.size, Some(details));
@@ -281,31 +357,19 @@ pub fn parse_disk_numbers(contents: &str) -> Vec<&str> {
     //
     //Red Hat VirtIO SCSI Disk Device
     //Disk ID: {E9CE8DFA-46B2-43C1-99BB-850C661CEE6B}
-    static RE: Lazy<Regex> =
-        Lazy::new(|| Regex::new(r"\s+Disk\s+(\d+).*\n.*\n.*\nDisk ID:").unwrap());
-
+    let re_disk_numbers = REGEXES.disk_numbers();
     let mut disk_nums = Vec::new();
-    for (_, [number]) in RE.captures_iter(contents).map(|c| c.extract()) {
+    for (_, [number]) in re_disk_numbers.captures_iter(contents).map(|c| c.extract()) {
         disk_nums.push(number);
     }
     disk_nums
 }
 
 pub fn parse_partition_details(parts: &mut [Partition], contents: &str) {
-    static RE_PAR: Lazy<Regex> = Lazy::new(|| {
-        Regex::new(
-            r"Partition (\d+)\r?\nType\s*: (\S+)(\r?\n.*){5}\s*(Volume.*\r?\n.*\r?\n|There is no volume)(.*)",
-        )
-        .unwrap()
-    });
-    static RE_VOL: Lazy<Regex> = Lazy::new(|| {
-        //   Volume ###  Ltr  Label        Fs     Type        Size     Status     Info
-        //   ----------  ---  -----------  -----  ----------  -------  ---------  --------
-        // * Volume 1     S   ESP          FAT32  Partition    100 MB  Healthy    Hidden
-        Regex::new(r"..Volume (\d.{2})  (.{3})  (.{11})  (.{5})").unwrap()
-    });
+    let re_detail_partition = REGEXES.detail_partition();
+    let re_list_volume = REGEXES.list_volumes();
 
-    for (part_index, (_, [_part_id, part_type, _, _vol_header, vol_line])) in RE_PAR
+    for (part_index, (_, [_part_id, part_type, _, _vol_header, vol_line])) in re_detail_partition
         .captures_iter(contents)
         .map(|c| c.extract())
         .enumerate()
@@ -316,7 +380,7 @@ pub fn parse_partition_details(parts: &mut [Partition], contents: &str) {
 
             // Volume info
             for (_, [_id, letter, label, fs_type]) in
-                RE_VOL.captures_iter(vol_line).map(|c| c.extract())
+                re_list_volume.captures_iter(vol_line).map(|c| c.extract())
             {
                 part.label = String::from(label.trim());
                 part.letter = String::from(letter.trim());
@@ -361,12 +425,11 @@ pub fn run_script(script: &str) -> String {
 
 pub fn split_diskpart_disk_output(contents: &str) -> Vec<&str> {
     // NOTE: A simple split isn't helpful since we want to include the matching lines
-    static RE: Lazy<Regex> =
-        Lazy::new(|| Regex::new(r"Disk \d+ is now the selected disk").unwrap());
+    let re_split_all_disks = REGEXES.split_all_disks();
     let mut sections = Vec::new();
     let mut starts: Vec<usize> = vec![0];
     let mut ends: Vec<usize> = Vec::new();
-    let _: Vec<_> = RE
+    let _: Vec<_> = re_split_all_disks
         .find_iter(contents)
         .map(|m| {
             ends.push(m.start() - 1);