ReplayGain optimizations/cleanup

src/backend/indexing.rs

@@ -1,13 +1,14 @@
 use std::{ffi::OsStr, fs::File, hash::Hasher, path::Path};
 
 use adler::Adler32;
-use lofty::{AudioFile, ItemKey, TaggedFileExt};
+use lofty::{AudioFile, TaggedFileExt};
 use miette::{miette, IntoDiagnostic, Result};
 use rayon::prelude::*;
 use sea_orm::{ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter, QuerySelect, Set};
 use symphonia::{
     core::{
-        formats::{FormatOptions, FormatReader},
+        errors::Error as SymphoniaError,
+        formats::{FormatOptions, FormatReader, Packet},
         io::{MediaSourceStream, MediaSourceStreamOptions},
         meta::{Limit, MetadataOptions},
         probe::Hint,
@@ -15,10 +16,10 @@ use symphonia::{
     default::get_probe,
 };
 use time::OffsetDateTime;
-use tracing::debug;
+use tracing::{debug, warn};
 use walkdir::{DirEntry, WalkDir};
 
-use crate::backend::replaygain::{format_gain, ReplayGain};
+use crate::backend::replaygain::{ReplayGainResult, ReplayGainState};
 
 use super::{
     config::Source,
@@ -49,14 +50,87 @@ fn get_packets(path: &Path) -> Result<Box<dyn FormatReader>> {
         .map(|v| v.format)
 }
 
-fn hash_packets(data: &mut Box<dyn FormatReader>) -> u64 {
+struct FormatReaderIter {
-    let mut adler = Adler32::new();
+    inner: Box<dyn FormatReader>,
+    error: Option<SymphoniaError>,
+    hash: Adler32,
+    rg: ReplayGainState,
+}
 
-    while let Ok(packet) = data.next_packet() {
+impl FormatReaderIter {
-        adler.write(&packet.data);
+    /// Initialize a new iterator over `FormatReader`.
+    /// Track ReplayGain values will be computed if `rg` is None.
+    fn new(
+        inner: Box<dyn FormatReader>,
+        rg: Option<ReplayGainResult>,
+    ) -> Result<Self, EleanorError> {
+        let track = inner
+            .default_track()
+            .ok_or_else(|| miette!("No default track was found"))?;
+        let params = &track.codec_params;
+
+        let (sample_rate, channels) = (params.sample_rate, params.channels);
+
+        // Only stereo is supported.
+        // In future we could try duplicating mono tracks and downmixing anything higher.
+        let rg = (!channels.is_some_and(|x| x.count() == 2))
+            .then_some(ReplayGainState::Failed(
+                miette!("Unsupported channel configuration: {:?}", channels).into(),
+            ))
+            .unwrap_or_else(|| ReplayGainState::new_with_override(sample_rate, track, rg));
+
+        Ok(Self {
+            inner,
+            error: None,
+            hash: Adler32::new(),
+            rg,
+        })
     }
 
-    adler.finish()
+    fn process(mut self) -> (u64, Result<ReplayGainResult, EleanorError>) {
+        // Loop over all packets
+        while let Some(packet) = (&mut self).next() {
+            // Hash the packet
+            self.hash.write(&packet.data);
+            // Copy into replaygain
+            self.rg.handle_packet(&packet);
+        }
+
+        let hash = self.hash.finish();
+
+        // Check for errors during iteration
+        if let Some(error) = self.error {
+            return (hash, Err(error.into()));
+        }
+
+        let rg = match self.rg {
+            ReplayGainState::Finished(rg_res) => Ok(rg_res),
+            ReplayGainState::Computing(rg) => Ok(rg.finish()),
+            ReplayGainState::Failed(e) => Err(e),
+        };
+
+        (hash, rg)
+    }
+}
+
+impl Iterator for &mut FormatReaderIter {
+    type Item = Packet;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let res_packet = self.inner.next_packet();
+        match res_packet {
+            Err(symphonia::core::errors::Error::IoError(ref packet_error))
+                if packet_error.kind() == std::io::ErrorKind::UnexpectedEof =>
+            {
+                None
+            }
+            Err(e) => {
+                self.error = Some(e);
+                None
+            }
+            Ok(packet) => Some(packet),
+        }
+    }
 }
 
 fn index_song(
@@ -65,6 +139,8 @@ fn index_song(
     force: bool,
     indexed_ts: OffsetDateTime,
 ) -> Result<Option<library::ActiveModel>, EleanorError> {
+    let filename = file.path().display().to_string();
+
     // Re-index previously indexed files
     if !force {
         let modified = file
@@ -75,7 +151,7 @@ fn index_song(
             .is_ok();
 
         if modified {
-            debug!("Skipping file {}", file.path().display());
+            debug!("Skipping file {filename}");
             return Ok(None);
         }
     }
@@ -87,42 +163,31 @@ fn index_song(
     let tags = tagged_file.primary_tag().or(tagged_file.first_tag());
     let properties = tagged_file.properties();
 
-    // Hash audio packets
+    // Hash audio packets and calculate replaygain
-    let mut packets = get_packets(file.path())?;
+    let (hash, rg) = FormatReaderIter::new(
-    let hash = hash_packets(&mut packets);
+        get_packets(file.path())?,
+        // Fall back to metadata
+        if let rg @ ReplayGainResult {
+            track_gain: Some(_),
+            track_peak: Some(_),
+            ..
+        } = ReplayGainResult::from(tags)
+        {
+            Some(rg)
+        } else {
+            None
+        },
+    )?
+    .process();
 
-    let rg_track_gain = tags
+    // Ignore replaygain errors
-        .and_then(|t| t.get_string(&ItemKey::ReplayGainTrackGain))
+    if let Err(e) = &rg {
-        .and_then(|v| format_gain(v).ok());
+        warn!(
-    let rg_track_peak = tags
+            track = filename,
-        .and_then(|t| t.get_string(&ItemKey::ReplayGainTrackPeak))
+            "Encountered error while calculating ReplayGain: {}", e
-        .and_then(|v| format_gain(v).ok());
+        );
-    let rg_album_gain = tags
-        .and_then(|t| t.get_string(&ItemKey::ReplayGainAlbumGain))
-        .and_then(|v| format_gain(v).ok());
-    let rg_album_peak = tags
-        .and_then(|t| t.get_string(&ItemKey::ReplayGainAlbumPeak))
-        .and_then(|v| format_gain(v).ok());
-
-    // Check for existing ReplayGain tags.
-    let mut rg = if let (Some(track_gain), Some(track_peak)) = (rg_track_gain, rg_track_peak) {
-        Ok(ReplayGain {
-            track_gain,
-            track_peak,
-            album_gain: None,
-            album_peak: None,
-        })
-    } else {
-        // Calculate replaygain values for the audio track
-        let mut packets = get_packets(file.path())?;
-        ReplayGain::try_calculate(&mut packets)
-    };
-
-    // Set album gain and peak, if present in metadata.
-    if let Ok(rg) = &mut rg {
-        rg.album_gain = rg_album_gain;
-        rg.album_peak = rg_album_peak;
     }
+    let rg = rg.ok();
 
     let song: library::ActiveModel = library::ActiveModel {
         path: Set(file
@@ -151,18 +216,10 @@ fn index_song(
         disc: Set(tags.and_then(lofty::Accessor::disk).map(|t| t as i32)),
         year: Set(tags.and_then(lofty::Accessor::year).map(|t| t as i32)),
         duration: Set(properties.duration().as_millis().try_into()?),
-        rg_track_gain: Set(rg.as_ref().map(|v| f64::from(v.track_gain)).ok()),
+        rg_track_gain: Set(rg.as_ref().and_then(|rg| rg.track_gain).map(f64::from)),
-        rg_track_peak: Set(rg.as_ref().map(|v| f64::from(v.track_peak)).ok()),
+        rg_track_peak: Set(rg.as_ref().and_then(|rg| rg.track_peak).map(f64::from)),
-        rg_album_gain: Set(rg
+        rg_album_gain: Set(rg.as_ref().and_then(|rg| rg.album_gain).map(f64::from)),
-            .as_ref()
+        rg_album_peak: Set(rg.as_ref().and_then(|rg| rg.album_peak).map(f64::from)),
-            .map(|v| v.album_gain.map(f64::from))
-            .ok()
-            .flatten()),
-        rg_album_peak: Set(rg
-            .as_ref()
-            .map(|v| v.album_peak.map(f64::from))
-            .ok()
-            .flatten()),
         ..Default::default()
     };
 

src/backend/replaygain.rs

@@ -1,105 +1,164 @@
+use lofty::{ItemKey, Tag};
 use miette::miette;
 
-use symphonia::core::{audio::SampleBuffer, codecs::DecoderOptions, formats::FormatReader};
+use symphonia::core::{
+    audio::SampleBuffer,
+    codecs::{Decoder, DecoderOptions},
+    formats::{Packet, Track},
+};
+use tracing::warn;
 
 use super::error::EleanorError;
 
-#[derive(Debug, Clone)]
+/// A wrapper around the data needed for ReplayGain calculation.
 pub struct ReplayGain {
-    pub track_gain: f32,
+    data: Vec<f32>,
-    pub track_peak: f32,
+    rg: replaygain::ReplayGain,
-    pub album_gain: Option<f32>,
+    decoder: Box<dyn Decoder>,
-    pub album_peak: Option<f32>,
+    track_id: u32,
+    sample_buf: Option<SampleBuffer<f32>>,
 }
 
 impl ReplayGain {
-    pub(crate) fn try_calculate(audio: &mut Box<dyn FormatReader>) -> Result<Self, EleanorError> {
+    pub fn init(sample_rate: usize, track: &Track) -> Result<Self, EleanorError> {
-        let track = audio
+        let rg = replaygain::ReplayGain::new(sample_rate)
-            .default_track()
+            .ok_or(miette!("Unsupported sample rate: {}", sample_rate))?;
-            .ok_or_else(|| miette!("No default track was found"))?;
 
         let params = &track.codec_params;
-
-        let (sample_rate, channels) = (params.sample_rate, params.channels);
-        let Some(sample_rate) = sample_rate else {
-            return Err(miette!("Sample rate must be known").into());
-        };
-
-        // Only stereo is supported.
-        if !channels.is_some_and(|x| x.count() == 2) {
-            return Err(miette!("Unsupported channel configuration: {:?}", channels).into());
-        }
-
-        // Only 44.1kHz and 48kHz are supported.
-        let Some(mut rg) = replaygain::ReplayGain::new(sample_rate as usize) else {
-            return Err(miette!("Unsupported sample rate: {:?}", sample_rate).into());
-        };
-
-        let mut decoder =
-            symphonia::default::get_codecs().make(params, &DecoderOptions::default())?;
-
         let track_id = track.id;
 
-        let mut samples: Vec<f32> = vec![];
+        let decoder = symphonia::default::get_codecs().make(params, &DecoderOptions::default())?;
-        let mut sample_buf = None;
-        loop {
-            let packet = match audio.next_packet() {
-                Ok(packet) => packet,
-                Err(symphonia::core::errors::Error::IoError(ref packet_error))
-                    if packet_error.kind() == std::io::ErrorKind::UnexpectedEof =>
-                {
-                    // End of audio stream
-                    break;
-                }
-                Err(e) => {
-                    return Err(e.into());
-                }
-            };
 
-            // Skip packets belonging to other audio tracks
+        Ok(Self {
-            if packet.track_id() != track_id {
+            rg,
-                continue;
+            data: Vec::new(),
-            }
+            decoder,
+            track_id,
+            sample_buf: None,
+        })
+    }
 
-            match decoder.decode(&packet) {
+    pub fn handle_packet(&mut self, packet: &Packet) -> Result<(), EleanorError> {
+        if packet.track_id() == self.track_id {
+            match self.decoder.decode(packet) {
                 Ok(buffer) => {
-                    if sample_buf.is_none() {
+                    if self.sample_buf.is_none() {
                         let spec = *buffer.spec();
                         let duration = buffer.capacity() as u64;
 
-                        sample_buf = Some(SampleBuffer::<f32>::new(duration, spec));
+                        self.sample_buf = Some(SampleBuffer::<f32>::new(duration, spec));
                     }
-                    if let Some(target) = &mut sample_buf {
+
+                    if let Some(target) = self.sample_buf.as_mut() {
                         target.copy_interleaved_ref(buffer);
                     }
                 }
                 Err(symphonia::core::errors::Error::DecodeError(_)) => (),
-                Err(_) => break,
+                Err(e) => return Err(e.into()),
             }
-
+            if let Some(buf) = self.sample_buf.as_mut() {
-            if let Some(buf) = &mut sample_buf {
+                self.data.extend(buf.samples());
-                samples.extend(buf.samples());
             }
         }
+        Ok(())
+    }
 
-        if samples.is_empty() {
+    pub fn finish(mut self) -> ReplayGainResult {
-            return Err(miette!("No samples were decoded from input audio").into());
+        let (track_gain, track_peak) = if self.data.is_empty() {
+            warn!("No samples were decoded from input audio");
+            (None, None)
+        } else {
+            self.rg.process_samples(&self.data);
+            let (gain, peak) = self.rg.finish();
+            (Some(gain), Some(peak))
         };
 
-        rg.process_samples(&samples);
+        ReplayGainResult {
-
-        let (track_gain, track_peak) = rg.finish();
-
-        Ok(Self {
             track_gain,
             track_peak,
             album_gain: None,
             album_peak: None,
-        })
+        }
     }
 }
 
-pub(crate) fn format_gain(gain: &str) -> Result<f32, EleanorError> {
+/// A representation of the possible calculation states
+pub enum ReplayGainState {
+    Finished(ReplayGainResult),
+    Computing(Box<ReplayGain>),
+    Failed(EleanorError),
+}
+
+impl ReplayGainState {
+    pub fn handle_packet(&mut self, packet: &Packet) {
+        if let ReplayGainState::Computing(rg) = self {
+            let result = rg.handle_packet(packet);
+
+            if let Err(e) = result {
+                *self = ReplayGainState::Failed(e);
+            }
+        }
+    }
+
+    pub fn new_with_override(
+        sample_rate: Option<u32>,
+        track: &Track,
+        r#override: Option<ReplayGainResult>,
+    ) -> Self {
+        let Some(sample_rate) = sample_rate else {
+            return Self::Failed(miette!("Sample rate must be known").into());
+        };
+
+        match r#override {
+            Some(rg_res) => Self::Finished(rg_res),
+            None => match ReplayGain::init(sample_rate as usize, track)
+                .map(Box::new)
+                .map(Self::Computing)
+            {
+                Ok(rg) => rg,
+                Err(e) => Self::Failed(e),
+            },
+        }
+    }
+}
+
+/// Finished ReplayGain calculation results.
+///
+/// Calculation errors may happen, but are non-critical
+/// and should be ignored by the player.
+#[derive(Debug, Clone)]
+pub struct ReplayGainResult {
+    pub track_gain: Option<f32>,
+    pub track_peak: Option<f32>,
+    pub album_gain: Option<f32>,
+    pub album_peak: Option<f32>,
+}
+
+impl From<Option<&Tag>> for ReplayGainResult {
+    fn from(tags: Option<&Tag>) -> Self {
+        let track_gain = tags
+            .and_then(|t| t.get_string(&ItemKey::ReplayGainTrackGain))
+            .and_then(|v| parse_gain(v).ok());
+        let track_peak = tags
+            .and_then(|t| t.get_string(&ItemKey::ReplayGainTrackPeak))
+            .and_then(|v| parse_gain(v).ok());
+        let album_gain = tags
+            .and_then(|t| t.get_string(&ItemKey::ReplayGainAlbumGain))
+            .and_then(|v| parse_gain(v).ok());
+        let album_peak = tags
+            .and_then(|t| t.get_string(&ItemKey::ReplayGainAlbumPeak))
+            .and_then(|v| parse_gain(v).ok());
+
+        Self {
+            track_gain,
+            track_peak,
+            album_gain,
+            album_peak,
+        }
+    }
+}
+
+pub fn parse_gain(gain: &str) -> Result<f32, EleanorError> {
     gain.chars()
         .filter(|c| c.is_numeric() || matches!(*c, '-' | '+' | '.'))
         .collect::<String>()
@@ -110,7 +169,7 @@ pub(crate) fn format_gain(gain: &str) -> Result<f32, EleanorError> {
 #[test]
 fn rg_parse_expected() {
     for (input, expected) in [("-8.97 dB", -8.97), ("12.75 dB", 12.75), ("0.00 dB", 0.0)] {
-        assert_eq!(format_gain(input).unwrap(), expected)
+        assert_eq!(parse_gain(input).unwrap(), expected)
     }
 }
 
@@ -124,13 +183,13 @@ fn rg_parse_malformed() {
         ("0.93dB", 0.93),
         ("24.12 deutscheBahn", 24.12),
     ] {
-        assert_eq!(format_gain(input).unwrap(), expected)
+        assert_eq!(parse_gain(input).unwrap(), expected)
     }
 }
 
 #[test]
 fn rg_parse_invalid() {
     for input in ["", "akjdfnhlkjfh", "🥺", "   DB "] {
-        assert!(format_gain(input).is_err())
+        assert!(parse_gain(input).is_err())
     }
 }