Add RVZ packing support

nod/src/common.rs

@@ -2,8 +2,12 @@
 
 use std::{borrow::Cow, fmt, str::FromStr, sync::Arc};
 
+use zerocopy::FromBytes;
+
 use crate::{
-    disc::{wii::WiiPartitionHeader, DiscHeader, PartitionHeader, SECTOR_SIZE},
+    disc::{
+        fst::Fst, wii::WiiPartitionHeader, DiscHeader, PartitionHeader, BOOT_SIZE, SECTOR_SIZE,
+    },
     Error, Result,
 };
 
@@ -300,14 +304,14 @@ pub struct PartitionInfo {
     pub key: KeyBytes,
     /// The Wii partition header.
     pub header: Arc<WiiPartitionHeader>,
-    /// The disc header within the partition.
-    pub disc_header: Arc<DiscHeader>,
-    /// The partition header within the partition.
-    pub partition_header: Arc<PartitionHeader>,
     /// Whether the partition data is encrypted
     pub has_encryption: bool,
     /// Whether the partition data hashes are present
     pub has_hashes: bool,
+    /// Disc and partition header (boot.bin)
+    pub raw_boot: Arc<[u8; BOOT_SIZE]>,
+    /// File system table (fst.bin), or `None` if partition is invalid
+    pub raw_fst: Option<Arc<[u8]>>,
 }
 
 impl PartitionInfo {
@@ -322,4 +326,25 @@ impl PartitionInfo {
     pub fn data_contains_sector(&self, sector: u32) -> bool {
         sector >= self.data_start_sector && sector < self.data_end_sector
     }
+
+    /// A view into the disc header.
+    #[inline]
+    pub fn disc_header(&self) -> &DiscHeader {
+        DiscHeader::ref_from_bytes(&self.raw_boot[..size_of::<DiscHeader>()])
+            .expect("Invalid disc header alignment")
+    }
+
+    /// A view into the partition header.
+    #[inline]
+    pub fn partition_header(&self) -> &PartitionHeader {
+        PartitionHeader::ref_from_bytes(&self.raw_boot[size_of::<DiscHeader>()..])
+            .expect("Invalid partition header alignment")
+    }
+
+    /// A view into the file system table (FST).
+    #[inline]
+    pub fn fst(&self) -> Option<Fst> {
+        // FST has already been parsed, so we can safely unwrap
+        Some(Fst::new(self.raw_fst.as_deref()?).unwrap())
+    }
 }

nod/src/disc/gcn.rs

@@ -5,7 +5,7 @@ use std::{
     sync::Arc,
 };
 
-use zerocopy::FromBytes;
+use zerocopy::{FromBytes, FromZeros, IntoBytes};
 
 use crate::{
     disc::{
@@ -17,7 +17,7 @@ use crate::{
     read::{PartitionEncryption, PartitionMeta, PartitionReader},
     util::{
         impl_read_for_bufread,
-        read::{read_arc, read_arc_slice, read_vec},
+        read::{read_arc, read_arc_slice, read_from},
     },
     Result, ResultContext,
 };
@@ -138,17 +138,16 @@ pub(crate) fn read_dol(
     reader
         .seek(SeekFrom::Start(partition_header.dol_offset(is_wii)))
         .context("Seeking to DOL offset")?;
-    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_header: DolHeader = read_from(reader).context("Reading DOL header")?;
     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())
         .max()
         .unwrap_or(size_of::<DolHeader>() as u32);
-    raw_dol.resize(dol_size as usize, 0);
+    let mut raw_dol = <[u8]>::new_box_zeroed_with_elems(dol_size as usize)?;
+    raw_dol[..size_of::<DolHeader>()].copy_from_slice(dol_header.as_bytes());
     reader.read_exact(&mut raw_dol[size_of::<DolHeader>()..]).context("Reading DOL")?;
-    Ok(Arc::from(raw_dol.as_slice()))
+    Ok(Arc::from(raw_dol))
 }
 
 pub(crate) fn read_fst<R>(
@@ -173,6 +172,24 @@ where
     Ok(raw_fst)
 }
 
+pub(crate) fn read_apploader<R>(reader: &mut R) -> Result<Arc<[u8]>>
+where R: Read + Seek + ?Sized {
+    reader
+        .seek(SeekFrom::Start(BOOT_SIZE as u64 + BI2_SIZE as u64))
+        .context("Seeking to apploader offset")?;
+    let apploader_header: ApploaderHeader =
+        read_from(reader).context("Reading apploader header")?;
+    let apploader_size = size_of::<ApploaderHeader>()
+        + apploader_header.size.get() as usize
+        + apploader_header.trailer_size.get() as usize;
+    let mut raw_apploader = <[u8]>::new_box_zeroed_with_elems(apploader_size)?;
+    raw_apploader[..size_of::<ApploaderHeader>()].copy_from_slice(apploader_header.as_bytes());
+    reader
+        .read_exact(&mut raw_apploader[size_of::<ApploaderHeader>()..])
+        .context("Reading apploader")?;
+    Ok(Arc::from(raw_apploader))
+}
+
 pub(crate) fn read_part_meta(
     reader: &mut dyn PartitionReader,
     is_wii: bool,
@@ -186,17 +203,7 @@ pub(crate) fn read_part_meta(
     let raw_bi2: Arc<[u8; BI2_SIZE]> = read_arc(reader).context("Reading bi2.bin")?;
 
     // apploader.bin
-    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();
-    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")?;
-    let raw_apploader = Arc::from(raw_apploader.as_slice());
+    let raw_apploader = read_apploader(reader)?;
 
     // fst.bin
     let raw_fst = read_fst(reader, partition_header, is_wii)?;

nod/src/disc/hashes.rs

@@ -7,7 +7,7 @@ use crate::{
         wii::{HASHES_SIZE, SECTOR_DATA_SIZE},
         SECTOR_GROUP_SIZE, SECTOR_SIZE,
     },
-    util::{array_ref, array_ref_mut},
+    util::{array_ref, array_ref_mut, digest::sha1_hash},
 };
 
 /// Hashes for a single sector group (64 sectors).
@@ -73,20 +73,3 @@ pub fn hash_sector_group(sector_group: &[u8; SECTOR_GROUP_SIZE]) -> Box<GroupHas
     result.h3_hash = sha1_hash(result.h2_hashes.as_bytes());
     result
 }
-
-/// Hashes a byte slice with SHA-1.
-#[instrument(skip_all)]
-pub fn sha1_hash(buf: &[u8]) -> HashBytes {
-    #[cfg(feature = "openssl")]
-    {
-        // The one-shot openssl::sha::sha1 ends up being much slower
-        let mut hasher = openssl::sha::Sha1::new();
-        hasher.update(buf);
-        hasher.finish()
-    }
-    #[cfg(not(feature = "openssl"))]
-    {
-        use sha1::Digest;
-        HashBytes::from(sha1::Sha1::digest(buf))
-    }
-}

nod/src/disc/mod.rs

@@ -260,7 +260,7 @@ pub const DOL_MAX_TEXT_SECTIONS: usize = 7;
 pub const DOL_MAX_DATA_SECTIONS: usize = 11;
 
 /// Dolphin executable (DOL) header.
-#[derive(Debug, Clone, FromBytes, Immutable, KnownLayout)]
+#[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)]
 pub struct DolHeader {
     /// Text section offsets
     pub text_offs: [U32; DOL_MAX_TEXT_SECTIONS],

nod/src/disc/preloader.rs

@@ -415,7 +415,7 @@ impl SectorGroupLoader {
                 sector_data,
                 self.block_buf.as_mut(),
                 abs_sector,
-                &partition.disc_header,
+                partition.disc_header(),
                 Some(partition),
             )?;
             if !encrypted {

nod/src/disc/reader.rs

@@ -6,7 +6,7 @@ use std::{
 
 use bytes::Bytes;
 use tracing::warn;
-use zerocopy::IntoBytes;
+use zerocopy::{FromBytes, IntoBytes};
 
 use crate::{
     common::{PartitionInfo, PartitionKind},
@@ -21,7 +21,8 @@ use crate::{
             PartitionReaderWii, WiiPartEntry, WiiPartGroup, WiiPartitionHeader, REGION_OFFSET,
             REGION_SIZE, WII_PART_GROUP_OFF,
         },
-        DiscHeader, DL_DVD_SIZE, MINI_DVD_SIZE, SECTOR_GROUP_SIZE, SECTOR_SIZE, SL_DVD_SIZE,
+        DiscHeader, PartitionHeader, BOOT_SIZE, DL_DVD_SIZE, MINI_DVD_SIZE, SECTOR_GROUP_SIZE,
+        SECTOR_SIZE, SL_DVD_SIZE,
     },
     io::block::BlockReader,
     read::{DiscMeta, DiscOptions, PartitionEncryption, PartitionOptions, PartitionReader},
@@ -38,12 +39,22 @@ pub struct DiscReader {
     pos: u64,
     size: u64,
     mode: PartitionEncryption,
-    disc_header: Arc<DiscHeader>,
-    partitions: Arc<[PartitionInfo]>,
-    region: Option<[u8; REGION_SIZE]>,
-    sector_group: Option<SectorGroup>,
+    raw_boot: Arc<[u8; BOOT_SIZE]>,
     alt_disc_header: Option<Arc<DiscHeader>>,
+    disc_data: DiscReaderData,
+    sector_group: Option<SectorGroup>,
+}
+
+#[derive(Clone)]
+enum DiscReaderData {
+    GameCube {
+        raw_fst: Option<Arc<[u8]>>,
+    },
+    Wii {
+        partitions: Arc<[PartitionInfo]>,
         alt_partitions: Option<Arc<[PartitionInfo]>>,
+        region: [u8; REGION_SIZE],
+    },
 }
 
 impl Clone for DiscReader {
@@ -54,12 +65,10 @@ impl Clone for DiscReader {
             pos: 0,
             size: self.size,
             mode: self.mode,
-            disc_header: self.disc_header.clone(),
-            partitions: self.partitions.clone(),
-            region: self.region,
-            sector_group: None,
+            raw_boot: self.raw_boot.clone(),
             alt_disc_header: self.alt_disc_header.clone(),
-            alt_partitions: self.alt_partitions.clone(),
+            disc_data: self.disc_data.clone(),
+            sector_group: None,
         }
     }
 }
@@ -68,12 +77,14 @@ impl DiscReader {
     pub fn new(inner: Box<dyn BlockReader>, options: &DiscOptions) -> Result<Self> {
         let mut reader = DirectDiscReader::new(inner)?;
 
-        let disc_header: Arc<DiscHeader> = read_arc(&mut reader).context("Reading disc header")?;
+        let raw_boot: Arc<[u8; BOOT_SIZE]> =
+            read_arc(reader.as_mut()).context("Reading disc headers")?;
+        let disc_header = DiscHeader::ref_from_bytes(&raw_boot[..size_of::<DiscHeader>()])
+            .expect("Invalid disc header alignment");
+        let disc_header_arc = Arc::from(disc_header.clone());
+
         let mut alt_disc_header = None;
-        let mut region = None;
-        let mut partitions = Arc::<[PartitionInfo]>::default();
-        let mut alt_partitions = None;
-        if disc_header.is_wii() {
+        let disc_data = if disc_header.is_wii() {
             // Sanity check
             if disc_header.has_partition_encryption() && !disc_header.has_partition_hashes() {
                 return Err(Error::DiscFormat(
@@ -90,17 +101,22 @@ impl DiscReader {
 
             // Read region info
             reader.seek(SeekFrom::Start(REGION_OFFSET)).context("Seeking to region info")?;
-            region = Some(read_from(&mut reader).context("Reading region info")?);
+            let region: [u8; REGION_SIZE] =
+                read_from(&mut reader).context("Reading region info")?;
 
             // Read partition info
-            partitions = Arc::from(read_partition_info(&mut reader, disc_header.clone())?);
+            let partitions = Arc::<[PartitionInfo]>::from(read_partition_info(
+                &mut reader,
+                disc_header_arc.clone(),
+            )?);
+            let mut alt_partitions = None;
 
             // Update disc header with encryption mode
             if matches!(
                 options.partition_encryption,
                 PartitionEncryption::ForceDecrypted | PartitionEncryption::ForceEncrypted
             ) {
-                let mut disc_header = Box::new(disc_header.as_ref().clone());
+                let mut disc_header = Box::new(disc_header.clone());
                 let mut partitions = Box::<[PartitionInfo]>::from(partitions.as_ref());
                 disc_header.no_partition_encryption = match options.partition_encryption {
                     PartitionEncryption::ForceDecrypted => 1,
@@ -113,15 +129,23 @@ impl DiscReader {
                 alt_disc_header = Some(Arc::from(disc_header));
                 alt_partitions = Some(Arc::from(partitions));
             }
-        } else if !disc_header.is_gamecube() {
+
+            DiscReaderData::Wii { partitions, alt_partitions, region }
+        } else if disc_header.is_gamecube() {
+            DiscReaderData::GameCube { raw_fst: None }
+        } else {
             return Err(Error::DiscFormat("Invalid disc header".to_string()));
-        }
+        };
 
         // Calculate disc size
         let io = reader.into_inner();
-        let size = io.meta().disc_size.unwrap_or_else(|| guess_disc_size(&partitions));
+        let partitions = match &disc_data {
+            DiscReaderData::Wii { partitions, .. } => partitions,
+            _ => &Arc::default(),
+        };
+        let size = io.meta().disc_size.unwrap_or_else(|| guess_disc_size(partitions));
         let preloader = Preloader::new(
-            SectorGroupLoader::new(io.clone(), disc_header.clone(), partitions.clone()),
+            SectorGroupLoader::new(io.clone(), disc_header_arc, partitions.clone()),
             options.preloader_threads,
         );
         Ok(Self {
@@ -130,12 +154,10 @@ impl DiscReader {
             pos: 0,
             size,
             mode: options.partition_encryption,
-            disc_header,
-            partitions,
-            region,
+            raw_boot,
+            disc_data,
             sector_group: None,
             alt_disc_header,
-            alt_partitions,
         })
     }
 
@@ -150,15 +172,68 @@ impl DiscReader {
 
     #[inline]
     pub fn header(&self) -> &DiscHeader {
-        self.alt_disc_header.as_ref().unwrap_or(&self.disc_header)
+        self.alt_disc_header.as_deref().unwrap_or_else(|| {
+            DiscHeader::ref_from_bytes(&self.raw_boot[..size_of::<DiscHeader>()])
+                .expect("Invalid disc header alignment")
+        })
+    }
+
+    // #[inline]
+    // pub fn orig_header(&self) -> &DiscHeader {
+    //     DiscHeader::ref_from_bytes(&self.raw_boot[..size_of::<DiscHeader>()])
+    //         .expect("Invalid disc header alignment")
+    // }
+
+    #[inline]
+    pub fn region(&self) -> Option<&[u8; REGION_SIZE]> {
+        match &self.disc_data {
+            DiscReaderData::Wii { region, .. } => Some(region),
+            _ => None,
+        }
     }
 
     #[inline]
-    pub fn region(&self) -> Option<&[u8; REGION_SIZE]> { self.region.as_ref() }
+    pub fn partitions(&self) -> &[PartitionInfo] {
+        match &self.disc_data {
+            DiscReaderData::Wii { partitions, alt_partitions, .. } => {
+                alt_partitions.as_deref().unwrap_or(partitions)
+            }
+            _ => &[],
+        }
+    }
 
     #[inline]
-    pub fn partitions(&self) -> &[PartitionInfo] {
-        self.alt_partitions.as_deref().unwrap_or(&self.partitions)
+    pub fn orig_partitions(&self) -> &[PartitionInfo] {
+        match &self.disc_data {
+            DiscReaderData::Wii { partitions, .. } => partitions,
+            _ => &[],
+        }
+    }
+
+    #[inline]
+    pub fn partition_header(&self) -> Option<&PartitionHeader> {
+        match &self.disc_data {
+            DiscReaderData::GameCube { .. } => Some(
+                PartitionHeader::ref_from_bytes(
+                    &self.raw_boot[size_of::<DiscHeader>()
+                        ..size_of::<DiscHeader>() + size_of::<PartitionHeader>()],
+                )
+                .expect("Invalid partition header alignment"),
+            ),
+            _ => None,
+        }
+    }
+
+    /// A reference to the raw FST for GameCube discs.
+    /// For Wii discs, use the FST from the appropriate [PartitionInfo].
+    #[inline]
+    pub fn fst(&self) -> Option<Fst> {
+        match &self.disc_data {
+            DiscReaderData::GameCube { raw_fst } => {
+                raw_fst.as_deref().and_then(|v| Fst::new(v).ok())
+            }
+            _ => None,
+        }
     }
 
     #[inline]
@@ -170,7 +245,8 @@ impl DiscReader {
         index: usize,
         options: &PartitionOptions,
     ) -> Result<Box<dyn PartitionReader>> {
-        if self.disc_header.is_gamecube() {
+        match &self.disc_data {
+            DiscReaderData::GameCube { .. } => {
                 if index == 0 {
                     Ok(PartitionReaderGC::new(
                         self.io.clone(),
@@ -180,12 +256,21 @@ impl DiscReader {
                 } else {
                     Err(Error::DiscFormat("GameCube discs only have one partition".to_string()))
                 }
-        } else if let Some(part) = self.partitions.get(index) {
-            Ok(PartitionReaderWii::new(self.io.clone(), self.preloader.clone(), part, options)?)
+            }
+            DiscReaderData::Wii { partitions, .. } => {
+                if let Some(part) = partitions.get(index) {
+                    Ok(PartitionReaderWii::new(
+                        self.io.clone(),
+                        self.preloader.clone(),
+                        part,
+                        options,
+                    )?)
                 } else {
                     Err(Error::DiscFormat(format!("Partition {index} not found")))
                 }
             }
+        }
+    }
 
     /// Opens a new, decrypted partition read stream for the first partition matching
     /// the specified kind.
@@ -194,7 +279,8 @@ impl DiscReader {
         kind: PartitionKind,
         options: &PartitionOptions,
     ) -> Result<Box<dyn PartitionReader>> {
-        if self.disc_header.is_gamecube() {
+        match &self.disc_data {
+            DiscReaderData::GameCube { .. } => {
                 if kind == PartitionKind::Data {
                     Ok(PartitionReaderGC::new(
                         self.io.clone(),
@@ -204,12 +290,21 @@ impl DiscReader {
                 } else {
                     Err(Error::DiscFormat("GameCube discs only have a data partition".to_string()))
                 }
-        } else if let Some(part) = self.partitions.iter().find(|v| v.kind == kind) {
-            Ok(PartitionReaderWii::new(self.io.clone(), self.preloader.clone(), part, options)?)
+            }
+            DiscReaderData::Wii { partitions, .. } => {
+                if let Some(part) = partitions.iter().find(|v| v.kind == kind) {
+                    Ok(PartitionReaderWii::new(
+                        self.io.clone(),
+                        self.preloader.clone(),
+                        part,
+                        options,
+                    )?)
                 } else {
                     Err(Error::DiscFormat(format!("Partition type {kind} not found")))
                 }
             }
+        }
+    }
 
     pub fn fill_buf_internal(&mut self) -> io::Result<Bytes> {
         if self.pos >= self.size {
@@ -228,7 +323,7 @@ impl DiscReader {
         // Build sector group request
         let abs_sector = (self.pos / SECTOR_SIZE as u64) as u32;
         let (request, abs_group_sector, max_groups) = if let Some(partition) =
-            self.partitions.iter().find(|part| part.data_contains_sector(abs_sector))
+            self.orig_partitions().iter().find(|part| part.data_contains_sector(abs_sector))
         {
             let group_idx = (abs_sector - partition.data_start_sector) / 64;
             let abs_group_sector = partition.data_start_sector + group_idx * 64;
@@ -283,7 +378,7 @@ impl BufRead for DiscReader {
         // Build sector group request
         let abs_sector = (self.pos / SECTOR_SIZE as u64) as u32;
         let (request, abs_group_sector, max_groups) = if let Some(partition) =
-            self.partitions.iter().find(|part| part.data_contains_sector(abs_sector))
+            self.orig_partitions().iter().find(|part| part.data_contains_sector(abs_sector))
         {
             let group_idx = (abs_sector - partition.data_start_sector) / 64;
             let abs_group_sector = partition.data_start_sector + group_idx * 64;
@@ -395,11 +490,16 @@ fn read_partition_info(
                 data_start_sector,
                 key,
             });
-            let partition_disc_header: Arc<DiscHeader> =
-                read_arc(reader).context("Reading partition disc header")?;
-            let partition_header = read_arc(reader).context("Reading partition header")?;
-            if partition_disc_header.is_wii() {
-                let raw_fst = read_fst(reader, &partition_header, true)?;
+            let raw_boot: Arc<[u8; BOOT_SIZE]> =
+                read_arc(reader).context("Reading partition headers")?;
+            let partition_disc_header =
+                DiscHeader::ref_from_bytes(&raw_boot[..size_of::<DiscHeader>()])
+                    .expect("Invalid disc header alignment");
+            let partition_header =
+                PartitionHeader::ref_from_bytes(&raw_boot[size_of::<DiscHeader>()..])
+                    .expect("Invalid partition header alignment");
+            let raw_fst = if partition_disc_header.is_wii() {
+                let raw_fst = read_fst(reader, partition_header, true)?;
                 let fst = Fst::new(&raw_fst)?;
                 let max_fst_offset = fst
                     .nodes
@@ -420,9 +520,11 @@ fn read_partition_info(
                         )));
                     }
                 }
+                Some(raw_fst)
             } else {
                 warn!("Partition {group_idx}:{part_idx} is not valid");
-            }
+                None
+            };
             reader.reset(DirectDiscReaderMode::Raw);
 
             part_info.push(PartitionInfo {
@@ -433,10 +535,10 @@ fn read_partition_info(
                 data_end_sector,
                 key,
                 header,
-                disc_header: partition_disc_header,
-                partition_header,
                 has_encryption: disc_header.has_partition_encryption(),
                 has_hashes: disc_header.has_partition_hashes(),
+                raw_boot,
+                raw_fst,
             });
         }
     }

nod/src/disc/wii.rs

@@ -14,7 +14,6 @@ use crate::{
     common::{HashBytes, KeyBytes, PartitionInfo},
     disc::{
         gcn::{read_part_meta, PartitionReaderGC},
-        hashes::sha1_hash,
         preloader::{fetch_sector_group, Preloader, SectorGroup, SectorGroupRequest},
         SECTOR_GROUP_SIZE, SECTOR_SIZE,
     },
@@ -22,7 +21,9 @@ use crate::{
     read::{PartitionEncryption, PartitionMeta, PartitionOptions, PartitionReader},
     util::{
         aes::aes_cbc_decrypt,
-        array_ref, div_rem, impl_read_for_bufread,
+        array_ref,
+        digest::sha1_hash,
+        div_rem, impl_read_for_bufread,
         read::{read_arc, read_arc_slice},
         static_assert,
     },

nod/src/disc/writer.rs

@@ -253,6 +253,7 @@ pub(crate) fn check_block(
         if sector_data_iter(block).enumerate().all(|(i, sector_data)| {
             let sector_offset = partition_offset + i as u64 * SECTOR_DATA_SIZE as u64;
             lfg.check_sector_chunked(sector_data, disc_id, disc_num, sector_offset)
+                == sector_data.len()
         }) {
             return Ok(CheckBlockResult::Junk);
         }
@@ -260,7 +261,7 @@ pub(crate) fn check_block(
         if buf.iter().all(|&b| b == 0) {
             return Ok(CheckBlockResult::Zeroed);
         }
-        if lfg.check_sector_chunked(buf, disc_id, disc_num, input_position) {
+        if lfg.check_sector_chunked(buf, disc_id, disc_num, input_position) == buf.len() {
             return Ok(CheckBlockResult::Junk);
         }
     }

nod/src/io/gcz.rs

@@ -324,19 +324,15 @@ impl DiscWriter for DiscWriterGCZ {
             options.processor_threads,
             |block| {
                 // Update hashers
-                let disc_data_len = block.disc_data.len() as u64;
+                input_position += block.disc_data.len() as u64;
                 digest.send(block.disc_data);
 
                 // Update block map and hash
-                if block.meta.is_compressed {
-                    block_map[block.block_idx as usize] = data_position.into();
-                } else {
-                    block_map[block.block_idx as usize] = (data_position | (1 << 63)).into();
-                }
+                let uncompressed_bit = (!block.meta.is_compressed as u64) << 63;
+                block_map[block.block_idx as usize] = (data_position | uncompressed_bit).into();
                 block_hashes[block.block_idx as usize] = block.meta.block_hash.into();
 
                 // Write block data
-                input_position += disc_data_len;
                 data_position += block.block_data.len() as u64;
                 data_callback(block.block_data, input_position, disc_size)
                     .with_context(|| format!("Failed to write block {}", block.block_idx))?;

nod/src/io/split.rs

@@ -103,7 +103,6 @@ impl SplitFileReader {
 
     pub fn len(&self) -> u64 { self.files.last().map_or(0, |f| f.begin + f.size) }
 
-    #[instrument(name = "SplitFileReader::check_open_file", skip_all)]
     fn check_open_file(&mut self) -> io::Result<Option<&mut Split<BufReader<File>>>> {
         if self.open_file.is_none() || !self.open_file.as_ref().unwrap().contains(self.pos) {
             self.open_file = if let Some(split) = self.files.iter().find(|f| f.contains(self.pos)) {

nod/src/util/digest.rs

@@ -6,10 +6,33 @@ use digest::Digest;
 use tracing::instrument;
 
 use crate::{
+    common::HashBytes,
     io::nkit::NKitHeader,
     write::{DiscFinalization, ProcessOptions},
 };
 
+/// Hashes a byte slice with SHA-1.
+#[instrument(skip_all)]
+pub fn sha1_hash(buf: &[u8]) -> HashBytes {
+    #[cfg(feature = "openssl")]
+    {
+        // The one-shot openssl::sha::sha1 ends up being much slower
+        let mut hasher = openssl::sha::Sha1::new();
+        hasher.update(buf);
+        hasher.finish()
+    }
+    #[cfg(not(feature = "openssl"))]
+    {
+        use sha1::Digest;
+        HashBytes::from(sha1::Sha1::digest(buf))
+    }
+}
+
+/// Hashes a byte slice with XXH64.
+#[allow(unused_braces)] // https://github.com/rust-lang/rust/issues/116347
+#[instrument(skip_all)]
+pub fn xxh64_hash(buf: &[u8]) -> u64 { xxhash_rust::xxh64::xxh64(buf, 0) }
+
 pub type DigestThread = (Sender<Bytes>, JoinHandle<DigestResult>);
 
 pub fn digest_thread<H>() -> DigestThread
@@ -40,13 +63,13 @@ impl DigestManager {
         }
         if options.digest_md5 {
             #[cfg(feature = "openssl")]
-            threads.push(digest_thread::<ossl::HasherMD5>());
+            threads.push(digest_thread::<openssl_util::HasherMD5>());
             #[cfg(not(feature = "openssl"))]
             threads.push(digest_thread::<md5::Md5>());
         }
         if options.digest_sha1 {
             #[cfg(feature = "openssl")]
-            threads.push(digest_thread::<ossl::HasherSHA1>());
+            threads.push(digest_thread::<openssl_util::HasherSHA1>());
             #[cfg(not(feature = "openssl"))]
             threads.push(digest_thread::<sha1::Sha1>());
         }
@@ -156,7 +179,7 @@ impl Hasher for xxhash_rust::xxh64::Xxh64 {
 }
 
 #[cfg(feature = "openssl")]
-mod ossl {
+mod openssl_util {
     use tracing::instrument;
 
     use super::{DigestResult, Hasher};
@@ -208,7 +231,7 @@ mod ossl {
         }
 
         #[allow(unused_braces)] // https://github.com/rust-lang/rust/issues/116347
-        #[instrument(name = "ossl::HasherMD5::update", skip_all)]
+        #[instrument(name = "openssl_util::HasherMD5::update", skip_all)]
         fn update(&mut self, data: &[u8]) { self.hasher.update(data).unwrap() }
     }
 
@@ -222,7 +245,7 @@ mod ossl {
         }
 
         #[allow(unused_braces)] // https://github.com/rust-lang/rust/issues/116347
-        #[instrument(name = "ossl::HasherSHA1::update", skip_all)]
+        #[instrument(name = "openssl_util::HasherSHA1::update", skip_all)]
         fn update(&mut self, data: &[u8]) { self.hasher.update(data).unwrap() }
     }
 }

nod/src/util/lfg.rs

@@ -1,25 +1,31 @@
 //! Lagged Fibonacci generator for GC / Wii partition junk data.
 
 use std::{
-    cmp::min,
     io,
     io::{Read, Write},
 };
 
 use bytes::Buf;
+use tracing::instrument;
 use zerocopy::{transmute_ref, IntoBytes};
 
-use crate::disc::SECTOR_SIZE;
+use crate::{disc::SECTOR_SIZE, util::array_ref_mut};
 
 /// Value of `k` for the LFG.
 pub const LFG_K: usize = 521;
 
+/// Value of `k` for the LFG in bytes.
+pub const LFG_K_BYTES: usize = LFG_K * 4;
+
 /// Value of `j` for the LFG.
 pub const LFG_J: usize = 32;
 
 /// Number of 32-bit words in the seed.
 pub const SEED_SIZE: usize = 17;
 
+/// Size of the seed in bytes.
+pub const SEED_SIZE_BYTES: usize = SEED_SIZE * 4;
+
 /// Lagged Fibonacci generator for GC / Wii partition junk data.
 ///
 /// References (license CC0-1.0):
@@ -38,48 +44,68 @@ impl Default for LaggedFibonacci {
 impl LaggedFibonacci {
     fn init(&mut self) {
         for i in SEED_SIZE..LFG_K {
-            self.buffer[i] =
-                (self.buffer[i - 17] << 23) ^ (self.buffer[i - 16] >> 9) ^ self.buffer[i - 1];
+            self.buffer[i] = (self.buffer[i - SEED_SIZE] << 23)
+                ^ (self.buffer[i - SEED_SIZE + 1] >> 9)
+                ^ self.buffer[i - 1];
         }
         // Instead of doing the "shift by 18 instead of 16" oddity when actually outputting the data,
         // we can do the shifting (and byteswapping) at this point to make the output code simpler.
         for x in self.buffer.iter_mut() {
-            *x = ((*x & 0xFF00FFFF) | (*x >> 2 & 0x00FF0000)).swap_bytes();
+            *x = ((*x & 0xFF00FFFF) | (*x >> 2 & 0x00FF0000)).to_be();
         }
         for _ in 0..4 {
             self.forward();
         }
     }
 
-    /// Initializes the LFG with the standard seed for a given disc ID, disc number, and sector.
-    /// The partition offset is used to determine the sector and how many bytes to skip within the
-    /// sector.
-    pub fn init_with_seed(&mut self, disc_id: [u8; 4], disc_num: u8, partition_offset: u64) {
+    /// Generates the seed for GC / Wii partition junk data using the disc ID, disc number, and sector.
+    pub fn generate_seed(out: &mut [u32; SEED_SIZE], disc_id: [u8; 4], disc_num: u8, sector: u32) {
         let seed = u32::from_be_bytes([
             disc_id[2],
             disc_id[1],
             disc_id[3].wrapping_add(disc_id[2]),
             disc_id[0].wrapping_add(disc_id[1]),
         ]) ^ disc_num as u32;
-        let sector = (partition_offset / SECTOR_SIZE as u64) as u32;
-        let sector_offset = partition_offset % SECTOR_SIZE as u64;
         let mut n = seed.wrapping_mul(0x260BCD5) ^ sector.wrapping_mul(0x1EF29123);
-        for i in 0..SEED_SIZE {
-            let mut v = 0u32;
+        for v in &mut *out {
+            *v = 0u32;
             for _ in 0..LFG_J {
                 n = n.wrapping_mul(0x5D588B65).wrapping_add(1);
-                v = (v >> 1) | (n & 0x80000000);
+                *v = (*v >> 1) | (n & 0x80000000);
             }
-            self.buffer[i] = v;
         }
-        self.buffer[16] ^= self.buffer[0] >> 9 ^ self.buffer[16] << 23;
+        out[16] ^= out[0] >> 9 ^ out[16] << 23;
+    }
+
+    /// Same as [`generate_seed`], but ensures the resulting seed is big-endian.
+    pub fn generate_seed_be(
+        out: &mut [u32; SEED_SIZE],
+        disc_id: [u8; 4],
+        disc_num: u8,
+        sector: u32,
+    ) {
+        Self::generate_seed(out, disc_id, disc_num, sector);
+        for x in out.iter_mut() {
+            *x = x.to_be();
+        }
+    }
+
+    /// Initializes the LFG with the standard seed for a given disc ID, disc number, and sector.
+    /// The partition offset is used to determine the sector and how many bytes to skip within the
+    /// sector.
+    #[instrument(name = "LaggedFibonacci::init_with_seed", skip_all)]
+    pub fn init_with_seed(&mut self, disc_id: [u8; 4], disc_num: u8, partition_offset: u64) {
+        let sector = (partition_offset / SECTOR_SIZE as u64) as u32;
+        let sector_offset = (partition_offset % SECTOR_SIZE as u64) as usize;
+        Self::generate_seed(array_ref_mut![self.buffer, 0, SEED_SIZE], disc_id, disc_num, sector);
         self.position = 0;
         self.init();
-        self.skip(sector_offset as usize);
+        self.skip(sector_offset);
     }
 
     /// Initializes the LFG with the seed read from a reader. The seed is assumed to be big-endian.
     /// This is used for rebuilding junk data in WIA/RVZ files.
+    #[instrument(name = "LaggedFibonacci::init_with_reader", skip_all)]
     pub fn init_with_reader<R>(&mut self, reader: &mut R) -> io::Result<()>
     where R: Read + ?Sized {
         reader.read_exact(self.buffer[..SEED_SIZE].as_mut_bytes())?;
@@ -93,6 +119,7 @@ impl LaggedFibonacci {
 
     /// Initializes the LFG with the seed read from a [`Buf`]. The seed is assumed to be big-endian.
     /// This is used for rebuilding junk data in WIA/RVZ files.
+    #[instrument(name = "LaggedFibonacci::init_with_buf", skip_all)]
     pub fn init_with_buf(&mut self, reader: &mut impl Buf) -> io::Result<()> {
         let out = self.buffer[..SEED_SIZE].as_mut_bytes();
         if reader.remaining() < out.len() {
@@ -108,6 +135,9 @@ impl LaggedFibonacci {
     }
 
     /// Advances the LFG by one step.
+    // This gets vectorized and aggressively inlined, so it's better to
+    // keep it separate for code size and instruction cache pressure.
+    #[inline(never)]
     fn forward(&mut self) {
         for i in 0..LFG_J {
             self.buffer[i] ^= self.buffer[i + LFG_K - LFG_J];
@@ -120,61 +150,49 @@ impl LaggedFibonacci {
     /// Skips `n` bytes of junk data.
     pub fn skip(&mut self, n: usize) {
         self.position += n;
-        while self.position >= LFG_K * 4 {
+        while self.position >= LFG_K_BYTES {
             self.forward();
-            self.position -= LFG_K * 4;
+            self.position -= LFG_K_BYTES;
         }
     }
 
-    // pub fn backward(&mut self) {
-    //     for i in (LFG_J..LFG_K).rev() {
-    //         self.buffer[i] ^= self.buffer[i - LFG_J];
-    //     }
-    //     for i in (0..LFG_J).rev() {
-    //         self.buffer[i] ^= self.buffer[i + LFG_K - LFG_J];
-    //     }
-    // }
-
-    // pub fn get_seed(&mut self, seed: &mut [u8; SEED_SIZE]) {
-    //     for i in 0..SEED_SIZE {
-    //         seed[i] = self.buffer[i].to_be_bytes()[3];
-    //     }
-    // }
-
     /// Fills the buffer with junk data.
+    #[instrument(name = "LaggedFibonacci::fill", skip_all)]
     pub fn fill(&mut self, mut buf: &mut [u8]) {
         while !buf.is_empty() {
-            let len = min(buf.len(), LFG_K * 4 - self.position);
-            let bytes: &[u8; LFG_K * 4] = transmute_ref!(&self.buffer);
+            while self.position >= LFG_K_BYTES {
+                self.forward();
+                self.position -= LFG_K_BYTES;
+            }
+            let bytes: &[u8; LFG_K_BYTES] = transmute_ref!(&self.buffer);
+            let len = buf.len().min(LFG_K_BYTES - self.position);
             buf[..len].copy_from_slice(&bytes[self.position..self.position + len]);
             self.position += len;
             buf = &mut buf[len..];
-            if self.position == LFG_K * 4 {
-                self.forward();
-                self.position = 0;
-            }
         }
     }
 
     /// Writes junk data to the output stream.
+    #[instrument(name = "LaggedFibonacci::write", skip_all)]
     pub fn write<W>(&mut self, w: &mut W, mut len: u64) -> io::Result<()>
     where W: Write + ?Sized {
         while len > 0 {
-            let write_len = min(len, LFG_K as u64 * 4 - self.position as u64) as usize;
-            let bytes: &[u8; LFG_K * 4] = transmute_ref!(&self.buffer);
+            while self.position >= LFG_K_BYTES {
+                self.forward();
+                self.position -= LFG_K_BYTES;
+            }
+            let bytes: &[u8; LFG_K_BYTES] = transmute_ref!(&self.buffer);
+            let write_len = len.min((LFG_K_BYTES - self.position) as u64) as usize;
             w.write_all(&bytes[self.position..self.position + write_len])?;
             self.position += write_len;
             len -= write_len as u64;
-            if self.position == LFG_K * 4 {
-                self.forward();
-                self.position = 0;
-            }
         }
         Ok(())
     }
 
     /// The junk data on GC / Wii discs is reinitialized every 32KB. This functions handles the
     /// wrapping logic and reinitializes the LFG at sector boundaries.
+    #[instrument(name = "LaggedFibonacci::fill_sector_chunked", skip_all)]
     pub fn fill_sector_chunked(
         &mut self,
         mut buf: &mut [u8],
@@ -194,6 +212,7 @@ impl LaggedFibonacci {
 
     /// The junk data on GC / Wii discs is reinitialized every 32KB. This functions handles the
     /// wrapping logic and reinitializes the LFG at sector boundaries.
+    #[instrument(name = "LaggedFibonacci::write_sector_chunked", skip_all)]
     pub fn write_sector_chunked<W>(
         &mut self,
         w: &mut W,
@@ -215,31 +234,50 @@ impl LaggedFibonacci {
         Ok(())
     }
 
+    /// Checks if the data matches the junk data generated by the LFG, up to the first sector
+    /// boundary.
+    #[instrument(name = "LaggedFibonacci::check", skip_all)]
+    pub fn check(
+        &mut self,
+        buf: &[u8],
+        disc_id: [u8; 4],
+        disc_num: u8,
+        partition_offset: u64,
+    ) -> usize {
+        let mut lfg_buf = [0u8; SECTOR_SIZE];
+        self.init_with_seed(disc_id, disc_num, partition_offset);
+        let len = (SECTOR_SIZE - (partition_offset % SECTOR_SIZE as u64) as usize).min(buf.len());
+        self.fill(&mut lfg_buf[..len]);
+        buf[..len].iter().zip(&lfg_buf[..len]).take_while(|(a, b)| a == b).count()
+    }
+
     /// Checks if the data matches the junk data generated by the LFG. This function handles the
     /// wrapping logic and reinitializes the LFG at sector boundaries.
+    #[instrument(name = "LaggedFibonacci::check_sector_chunked", skip_all)]
     pub fn check_sector_chunked(
         &mut self,
         mut buf: &[u8],
         disc_id: [u8; 4],
         disc_num: u8,
         mut partition_offset: u64,
-    ) -> bool {
-        if buf.is_empty() {
-            return false;
-        }
+    ) -> usize {
         let mut lfg_buf = [0u8; SECTOR_SIZE];
+        let mut total_num_matching = 0;
         while !buf.is_empty() {
             self.init_with_seed(disc_id, disc_num, partition_offset);
             let len =
                 (SECTOR_SIZE - (partition_offset % SECTOR_SIZE as u64) as usize).min(buf.len());
             self.fill(&mut lfg_buf[..len]);
-            if buf[..len] != lfg_buf[..len] {
-                return false;
+            let num_matching =
+                buf[..len].iter().zip(&lfg_buf[..len]).take_while(|(a, b)| a == b).count();
+            total_num_matching += num_matching;
+            if num_matching != len {
+                break;
             }
             buf = &buf[len..];
             partition_offset += len as u64;
         }
-        true
+        total_num_matching
     }
 }
 

nodtool/src/cmd/gen.rs

@@ -673,7 +673,7 @@ fn check_junk_data(
             .fill_buf()
             .with_context(|| format!("Failed to read disc file at offset {}", offset))?;
         let read_len = (file_buf.len() as u64).min(remaining) as usize;
-        if !lfg.check_sector_chunked(&file_buf[..read_len], junk_id, disc_num, pos) {
+        if lfg.check_sector_chunked(&file_buf[..read_len], junk_id, disc_num, pos) != read_len {
             return Ok(false);
         }
 

nodtool/src/cmd/info.rs

@@ -92,12 +92,13 @@ fn info_file(path: &Path) -> nod::Result<()> {
             } else {
                 "N/A".to_string()
             };
-            println!("\tTitle: {}", info.disc_header.game_title_str());
-            println!("\tGame ID: {} ({})", info.disc_header.game_id_str(), title_id_str);
+            let part_disc_header = info.disc_header();
+            println!("\tTitle: {}", part_disc_header.game_title_str());
+            println!("\tGame ID: {} ({})", part_disc_header.game_id_str(), title_id_str);
             println!(
                 "\tDisc {}, Revision {}",
-                info.disc_header.disc_num + 1,
-                info.disc_header.disc_version
+                part_disc_header.disc_num + 1,
+                part_disc_header.disc_version
             );
         }
     } else if header.is_gamecube() {