Centralize logic into fetch_sector_group

2024-11-23 12:54:05 -07:00 · 2024-11-23 12:54:05 -07:00 · d197b8e7c2
parent b8b06dcd5c
commit d197b8e7c2
6 changed files with 80 additions and 64 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@ -393,6 +393,15 @@ dependencies = [
 "serde",
 ]

+[[package]]
+name = "higher-kinded-types"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "561985554c8b8d4808605c90a5f1979cc6c31a5d20b78465cd59501233c6678e"
+dependencies = [
+ "never-say-never",
+]
+
 [[package]]
 name = "hybrid-array"
 version = "0.2.1"
@ -580,6 +589,12 @@ dependencies = [
 "adler2",
 ]

+[[package]]
+name = "never-say-never"
+version = "6.6.666"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cf5a574dadd7941adeaa71823ecba5e28331b8313fb2e1c6a5c7e5981ea53ad6"
+
 [[package]]
 name = "nod"
 version = "2.0.0-alpha.1"
@ -605,6 +620,7 @@ dependencies = [
 "memmap2",
 "miniz_oxide",
 "openssl",
+ "polonius-the-crab",
 "rand",
 "rayon",
 "sha1",
@ -751,6 +767,16 @@ version = "0.3.31"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "953ec861398dccce10c670dfeaf3ec4911ca479e9c02154b3a215178c5f566f2"

+[[package]]
+name = "polonius-the-crab"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e97ca2c89572ae41bbec1c99498251f87dd5a94e500c5ec19c382dd593dd5ce9"
+dependencies = [
+ "higher-kinded-types",
+ "never-say-never",
+]
+
 [[package]]
 name = "portable-atomic"
 version = "1.9.0"
--- a/nod/Cargo.toml
+++ b/nod/Cargo.toml
@ -45,6 +45,7 @@ md-5 = { workspace = true }
 memmap2 = "0.9"
 miniz_oxide = { version = "0.8", optional = true }
 openssl = { version = "0.10", optional = true }
+polonius-the-crab = "0.4"
 rand = "0.8"
 rayon = "1.10"
 sha1 = { workspace = true }
--- a/nod/src/disc/gcn.rs
+++ b/nod/src/disc/gcn.rs
@ -9,7 +9,7 @@ use zerocopy::FromBytes;

 use crate::{
    disc::{
-        preloader::{Preloader, SectorGroup, SectorGroupRequest},
+        preloader::{fetch_sector_group, Preloader, SectorGroup, SectorGroupRequest},
        ApploaderHeader, DiscHeader, DolHeader, PartitionHeader, BI2_SIZE, BOOT_SIZE,
        SECTOR_GROUP_SIZE, SECTOR_SIZE,
    },
@ -77,14 +77,9 @@ impl BufRead for PartitionReaderGC {
            mode: PartitionEncryption::Original,
        };

-        let sector_group = if matches!(&self.sector_group, Some(sector_group) if sector_group.request == request)
-        {
-            // We can improve this in Rust 2024 with `if_let_rescope`
-            // https://github.com/rust-lang/rust/issues/124085
-            self.sector_group.as_ref().unwrap()
-        } else {
-            self.sector_group.insert(self.preloader.fetch(request, max_groups)?)
-        };
+        // Load sector group
+        let (sector_group, _updated) =
+            fetch_sector_group(request, max_groups, &mut self.sector_group, &self.preloader)?;

        // Calculate the number of consecutive sectors in the group
        let group_sector = abs_sector - abs_group_sector;
@ -146,18 +141,9 @@ pub(crate) fn read_dol(
    let mut raw_dol: Vec<u8> =
        read_vec(reader, size_of::<DolHeader>()).context("Reading DOL header")?;
    let dol_header = DolHeader::ref_from_bytes(raw_dol.as_slice()).unwrap();
-    let dol_size = dol_header
-        .text_offs
-        .iter()
-        .zip(&dol_header.text_sizes)
+    let dol_size = (dol_header.text_offs.iter().zip(&dol_header.text_sizes))
+        .chain(dol_header.data_offs.iter().zip(&dol_header.data_sizes))
        .map(|(offs, size)| offs.get() + size.get())
-        .chain(
-            dol_header
-                .data_offs
-                .iter()
-                .zip(&dol_header.data_sizes)
-                .map(|(offs, size)| offs.get() + size.get()),
-        )
        .max()
        .unwrap_or(size_of::<DolHeader>() as u32);
    raw_dol.resize(dol_size as usize, 0);
@ -203,12 +189,10 @@ pub(crate) fn read_part_meta(
    let mut raw_apploader: Vec<u8> =
        read_vec(reader, size_of::<ApploaderHeader>()).context("Reading apploader header")?;
    let apploader_header = ApploaderHeader::ref_from_bytes(raw_apploader.as_slice()).unwrap();
-    raw_apploader.resize(
-        size_of::<ApploaderHeader>()
-            + apploader_header.size.get() as usize
-            + apploader_header.trailer_size.get() as usize,
-        0,
-    );
+    let apploader_size = size_of::<ApploaderHeader>()
+        + apploader_header.size.get() as usize
+        + apploader_header.trailer_size.get() as usize;
+    raw_apploader.resize(apploader_size, 0);
    reader
        .read_exact(&mut raw_apploader[size_of::<ApploaderHeader>()..])
        .context("Reading apploader")?;
--- a/nod/src/disc/preloader.rs
+++ b/nod/src/disc/preloader.rs
@ -11,6 +11,7 @@ use bytes::{Bytes, BytesMut};
 use crossbeam_channel::{Receiver, Sender};
 use crossbeam_utils::sync::WaitGroup;
 use lru::LruCache;
+use polonius_the_crab::{polonius, polonius_return};
 use simple_moving_average::{SingleSumSMA, SMA};
 use tracing::{debug, error, instrument, span, Level};
 use zerocopy::FromZeros;
@ -556,3 +557,22 @@ impl SectorGroupLoader {
        Ok((sector_bitmap, io_duration))
    }
 }
+
+/// Fetch a sector group from the cache or from the preloader.
+/// Returns a boolean indicating if the group was updated.
+pub fn fetch_sector_group<'a>(
+    request: SectorGroupRequest,
+    max_groups: u32,
+    mut cached: &'a mut Option<SectorGroup>,
+    preloader: &Preloader,
+) -> io::Result<(&'a SectorGroup, bool)> {
+    polonius!(|cached| -> io::Result<(&'polonius SectorGroup, bool)> {
+        if let Some(sector_group) = cached {
+            if sector_group.request == request {
+                polonius_return!(Ok((sector_group, false)));
+            }
+        }
+    });
+    let sector_group = preloader.fetch(request, max_groups)?;
+    Ok((cached.insert(sector_group), true))
+}
--- a/nod/src/disc/reader.rs
+++ b/nod/src/disc/reader.rs
@ -14,7 +14,9 @@ use crate::{
        direct::{DirectDiscReader, DirectDiscReaderMode},
        fst::{Fst, NodeKind},
        gcn::{read_fst, PartitionReaderGC},
-        preloader::{Preloader, SectorGroup, SectorGroupLoader, SectorGroupRequest},
+        preloader::{
+            fetch_sector_group, Preloader, SectorGroup, SectorGroupLoader, SectorGroupRequest,
+        },
        wii::{
            PartitionReaderWii, WiiPartEntry, WiiPartGroup, WiiPartitionHeader, REGION_OFFSET,
            REGION_SIZE, WII_PART_GROUP_OFF,
@ -32,6 +34,7 @@ use crate::{

 pub struct DiscReader {
    io: Box<dyn BlockReader>,
+    preloader: Arc<Preloader>,
    pos: u64,
    size: u64,
    mode: PartitionEncryption,
@ -39,7 +42,6 @@ pub struct DiscReader {
    partitions: Arc<[PartitionInfo]>,
    region: Option<[u8; REGION_SIZE]>,
    sector_group: Option<SectorGroup>,
-    preloader: Arc<Preloader>,
    alt_disc_header: Option<Arc<DiscHeader>>,
    alt_partitions: Option<Arc<[PartitionInfo]>>,
 }
@ -48,6 +50,7 @@ impl Clone for DiscReader {
    fn clone(&self) -> Self {
        Self {
            io: self.io.clone(),
+            preloader: self.preloader.clone(),
            pos: 0,
            size: self.size,
            mode: self.mode,
@ -55,7 +58,6 @@ impl Clone for DiscReader {
            partitions: self.partitions.clone(),
            region: self.region,
            sector_group: None,
-            preloader: self.preloader.clone(),
            alt_disc_header: self.alt_disc_header.clone(),
            alt_partitions: self.alt_partitions.clone(),
        }
@ -124,6 +126,7 @@ impl DiscReader {
        );
        Ok(Self {
            io,
+            preloader,
            pos: 0,
            size,
            mode: options.partition_encryption,
@ -131,7 +134,6 @@ impl DiscReader {
            partitions,
            region,
            sector_group: None,
-            preloader,
            alt_disc_header,
            alt_partitions,
        })
@ -246,14 +248,8 @@ impl DiscReader {
        };

        // Load sector group
-        let sector_group = if matches!(&self.sector_group, Some(sector_group) if sector_group.request == request)
-        {
-            // We can improve this in Rust 2024 with `if_let_rescope`
-            // https://github.com/rust-lang/rust/issues/124085
-            self.sector_group.as_ref().unwrap()
-        } else {
-            self.sector_group.insert(self.preloader.fetch(request, max_groups)?)
-        };
+        let (sector_group, _updated) =
+            fetch_sector_group(request, max_groups, &mut self.sector_group, &self.preloader)?;

        // Calculate the number of consecutive sectors in the group
        let group_sector = abs_sector - abs_group_sector;
@ -307,14 +303,8 @@ impl BufRead for DiscReader {
        };

        // Load sector group
-        let sector_group = if matches!(&self.sector_group, Some(sector_group) if sector_group.request == request)
-        {
-            // We can improve this in Rust 2024 with `if_let_rescope`
-            // https://github.com/rust-lang/rust/issues/124085
-            self.sector_group.as_ref().unwrap()
-        } else {
-            self.sector_group.insert(self.preloader.fetch(request, max_groups)?)
-        };
+        let (sector_group, _updated) =
+            fetch_sector_group(request, max_groups, &mut self.sector_group, &self.preloader)?;

        // Calculate the number of consecutive sectors in the group
        let group_sector = abs_sector - abs_group_sector;
--- a/nod/src/disc/wii.rs
+++ b/nod/src/disc/wii.rs
@ -15,7 +15,7 @@ use crate::{
    disc::{
        gcn::{read_part_meta, PartitionReaderGC},
        hashes::sha1_hash,
-        preloader::{Preloader, SectorGroup, SectorGroupRequest},
+        preloader::{fetch_sector_group, Preloader, SectorGroup, SectorGroupRequest},
        SECTOR_GROUP_SIZE, SECTOR_SIZE,
    },
    io::block::BlockReader,
@ -378,24 +378,19 @@ impl BufRead for PartitionReaderWii {
                PartitionEncryption::ForceDecryptedNoHashes
            },
        };
-        let sector_group = if matches!(&self.sector_group, Some(sector_group) if sector_group.request == request)
-        {
-            // We can improve this in Rust 2024 with `if_let_rescope`
-            // https://github.com/rust-lang/rust/issues/124085
-            self.sector_group.as_ref().unwrap()
-        } else {
-            let sector_group = self.preloader.fetch(request, max_groups)?;
-            if self.options.validate_hashes {
-                if let Some(h3_table) = self.meta.as_ref().and_then(|m| m.raw_h3_table.as_deref()) {
-                    verify_hashes(
-                        array_ref![sector_group.data, 0, SECTOR_GROUP_SIZE],
-                        group_idx,
-                        h3_table,
-                    )?;
-                }
+
+        // Load sector group
+        let (sector_group, updated) =
+            fetch_sector_group(request, max_groups, &mut self.sector_group, &self.preloader)?;
+        if updated && self.options.validate_hashes {
+            if let Some(h3_table) = self.meta.as_ref().and_then(|m| m.raw_h3_table.as_deref()) {
+                verify_hashes(
+                    array_ref![sector_group.data, 0, SECTOR_GROUP_SIZE],
+                    group_idx,
+                    h3_table,
+                )?;
            }
-            self.sector_group.insert(sector_group)
-        };
+        }

        // Read from sector group buffer
        let consecutive_sectors = sector_group.consecutive_sectors(group_sector);