feat: compact strains (#32)

* store zero strains compactly * added docs * added proptest * add feature to store strains compactly * update ci * rename CompactZerosVec to StrainsVec * add size_hint for iter * rename iter * minor adjustments * more minor adjustments
MaxOhn · Mar 31, 2024 · a577200 · a577200
1 parent 279e04c
commit a577200
Show file tree

Hide file tree

Showing 22 changed files with 552 additions and 84 deletions.
diff --git a/.github/workflows/check.yml b/.github/workflows/check.yml
@@ -1,4 +1,4 @@
-name: Clippy and Tests
+name: Clippy checks
 
 on:
   push:
@@ -14,7 +14,7 @@ jobs:
 
     steps:
       - name: Checkout project
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
 
       - name: Install stable toolchain
         uses: dtolnay/rust-toolchain@stable
@@ -29,34 +29,3 @@ jobs:
         run: |
           cargo install cargo-rdme
           cargo rdme --check --no-fail-on-warnings
-
-  tests:
-    name: Tests
-    runs-on: ubuntu-latest
-
-    steps:
-      - name: Checkout project
-        uses: actions/checkout@v3
-
-      - name: Install stable toolchain
-        uses: dtolnay/rust-toolchain@stable
-
-      - name: Cache dependencies
-        uses: Swatinem/rust-cache@v2
-
-      - name: Install nextest
-        uses: taiki-e/install-action@nextest
-
-      - name: Run tests on default features
-        run: cargo nextest run --no-default-features --no-fail-fast --failure-output=immediate-final
-
-      - name: Run tests with sync feature
-        run: >
-          cargo nextest run
-          --features sync
-          --filter-expr 'test(util::sync::tests::share_gradual_taiko)'
-          --filter-expr 'test(taiko::difficulty::gradual::tests::next_and_nth)'
-          --no-fail-fast --failure-output=immediate-final
-
-      - name: Run doctests
-        run: cargo test --no-default-features --doc
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -0,0 +1,101 @@
+name: Tests
+
+on:
+  push:
+    branches:
+      - main
+      - next
+  pull_request:
+
+jobs:
+  doc:
+    name: Doc tests
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout project
+        uses: actions/checkout@v4
+
+      - name: Install stable toolchain
+        uses: dtolnay/rust-toolchain@stable
+
+      - name: Cache dependencies
+        uses: Swatinem/rust-cache@v2
+
+      - name: Run doctests
+        run: >
+          cargo test
+          --doc
+          --no-default-features --features compact_strains
+
+  default:
+    name: Default tests
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout project
+        uses: actions/checkout@v4
+
+      - name: Install stable toolchain
+        uses: dtolnay/rust-toolchain@stable
+
+      - name: Cache dependencies
+        uses: Swatinem/rust-cache@v2
+
+      - name: Install nextest
+        uses: taiki-e/install-action@nextest
+
+      - name: Run all tests
+        run: >
+          cargo nextest run
+          --no-default-features --features compact_strains
+          --no-fail-fast --failure-output=immediate-final
+
+  sync:
+    name: Test sync feature
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout project
+        uses: actions/checkout@v4
+
+      - name: Install stable toolchain
+        uses: dtolnay/rust-toolchain@stable
+
+      - name: Cache dependencies
+        uses: Swatinem/rust-cache@v2
+
+      - name: Install nextest
+        uses: taiki-e/install-action@nextest
+
+      - name: Run specific tests
+        run: >
+          cargo nextest run
+          --features sync
+          --filter-expr 'test(util::sync::tests::share_gradual_taiko)'
+          --filter-expr 'test(taiko::difficulty::gradual::tests::next_and_nth)'
+          --no-fail-fast --failure-output=immediate-final
+
+  non_compact:
+    name: Test without compact_strains feature
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout project
+        uses: actions/checkout@v4
+
+      - name: Install stable toolchain
+        uses: dtolnay/rust-toolchain@stable
+
+      - name: Cache dependencies
+        uses: Swatinem/rust-cache@v2
+
+      - name: Install nextest
+        uses: taiki-e/install-action@nextest
+
+      - name: Run integration tests
+        run: >
+          cargo nextest run
+          --no-default-features
+          --test '*'
+          --no-fail-fast --failure-output=immediate-final
diff --git a/Cargo.toml b/Cargo.toml
@@ -11,7 +11,8 @@ description = "Difficulty and performance calculation for osu!"
 keywords = ["osu", "pp", "stars", "performance", "osu!"]
 
 [features]
-default = []
+default = ["compact_strains"]
+compact_strains = []
 sync = []
 tracing = ["rosu-map/tracing"]
 

diff --git a/README.md b/README.md
@@ -125,11 +125,12 @@ Calculating performances: Median: 44.13µs | Mean: 45.53µs
 
 ### Features
 
-| Flag      | Description | Dependencies
-| --------- | ----------- | ------------
-| `default` | No features |
-| `sync`    | Some gradual calculation types can only be shared across threads if this feature is enabled. This adds a performance penalty so only enable this if really needed. |
-| `tracing` | Any error encountered during beatmap decoding will be logged through `tracing::error`. If this feature is **not** enabled, errors will be ignored. | [`tracing`]
+| Flag              | Description                           | Dependencies
+| ----------------- | ------------------------------------- | ------------
+| `default`         | Enables the `compact_strains` feature |
+| `compact_strains` | Storing internal strain values in a plain Vec introduces an out-of-memory risk on maliciously long maps (see [/b/3739922](https://osu.ppy.sh/b/3739922)). This feature stores strains more compactly, but comes with a ~5% loss in performance. |
+| `sync`            | Some gradual calculation types can only be shared across threads if this feature is enabled. This adds a performance penalty so only enable this if really needed. |
+| `tracing`         | Any error encountered during beatmap decoding will be logged through `tracing::error`. If this feature is **not** enabled, errors will be ignored. | [`tracing`]
 
 ### Bindings
 

diff --git a/src/any/difficulty/skills.rs b/src/any/difficulty/skills.rs
@@ -1,3 +1,5 @@
+use crate::util::strains_vec::StrainsVec;
+
 pub fn strain_decay(ms: f64, strain_decay_base: f64) -> f64 {
     strain_decay_base.powf(ms / 1000.0)
 }
@@ -27,7 +29,7 @@ pub trait ISkill {
 pub struct StrainSkill {
     pub curr_section_peak: f64,
     pub curr_section_end: f64,
-    pub strain_peaks: Vec<f64>,
+    pub strain_peaks: StrainsVec,
 }
 
 impl Default for StrainSkill {
@@ -36,7 +38,7 @@ impl Default for StrainSkill {
             curr_section_peak: 0.0,
             curr_section_end: 0.0,
             // mean=386.81 | median=279
-            strain_peaks: Vec::with_capacity(256),
+            strain_peaks: StrainsVec::with_capacity(256),
         }
     }
 }
@@ -53,7 +55,7 @@ impl StrainSkill {
         self.curr_section_peak = initial_strain;
     }
 
-    pub fn get_curr_strain_peaks(self) -> Vec<f64> {
+    pub fn get_curr_strain_peaks(self) -> StrainsVec {
         let mut strain_peaks = self.strain_peaks;
         strain_peaks.push(self.curr_section_peak);
 
@@ -65,10 +67,8 @@ impl StrainSkill {
         let mut weight = 1.0;
 
         let mut peaks = self.get_curr_strain_peaks();
-        peaks.retain(|&strain| strain > 0.0);
-        peaks.sort_by(|a, b| b.total_cmp(a));
 
-        for strain in peaks {
+        for strain in peaks.sorted_non_zero_iter() {
             difficulty += strain * weight;
             weight *= decay_weight;
         }
@@ -95,7 +95,7 @@ impl StrainDecaySkill {
         self.inner.start_new_section_from(initial_strain);
     }
 
-    pub fn get_curr_strain_peaks(self) -> Vec<f64> {
+    pub fn get_curr_strain_peaks(self) -> StrainsVec {
         self.inner.get_curr_strain_peaks()
     }
 

diff --git a/src/catch/difficulty/skills/movement.rs b/src/catch/difficulty/skills/movement.rs
@@ -4,6 +4,7 @@ use crate::{
         skills::{strain_decay, ISkill, Skill, StrainDecaySkill},
     },
     catch::difficulty::object::CatchDifficultyObject,
+    util::strains_vec::StrainsVec,
 };
 
 const ABSOLUTE_PLAYER_POSITIONING_ERROR: f32 = 16.0;
@@ -107,7 +108,7 @@ impl Movement {
         dist_addition / weighted_strain_time
     }
 
-    pub fn get_curr_strain_peaks(self) -> Vec<f64> {
+    pub fn get_curr_strain_peaks(self) -> StrainsVec {
         self.inner.get_curr_strain_peaks()
     }
 

diff --git a/src/catch/strains.rs b/src/catch/strains.rs
@@ -20,6 +20,6 @@ pub fn strains(difficulty: &Difficulty, converted: &CatchBeatmap<'_>) -> CatchSt
     let DifficultyValues { movement, .. } = DifficultyValues::calculate(difficulty, converted);
 
     CatchStrains {
-        movement: movement.get_curr_strain_peaks(),
+        movement: movement.get_curr_strain_peaks().into_vec(),
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -121,11 +121,12 @@
 //!
 //! ## Features
 //!
-//! | Flag      | Description | Dependencies
-//! | --------- | ----------- | ------------
-//! | `default` | No features |
-//! | `sync`    | Some gradual calculation types can only be shared across threads if this feature is enabled. This adds a performance penalty so only enable this if really needed. |
-//! | `tracing` | Any error encountered during beatmap decoding will be logged through `tracing::error`. If this feature is **not** enabled, errors will be ignored. | [`tracing`]
+//! | Flag              | Description                           | Dependencies
+//! | ----------------- | ------------------------------------- | ------------
+//! | `default`         | Enables the `compact_strains` feature |
+//! | `compact_strains` | Storing internal strain values in a plain Vec introduces an out-of-memory risk on maliciously long maps (see [/b/3739922](https://osu.ppy.sh/b/3739922)). This feature stores strains more compactly, but comes with a ~5% loss in performance. |
+//! | `sync`            | Some gradual calculation types can only be shared across threads if this feature is enabled. This adds a performance penalty so only enable this if really needed. |
+//! | `tracing`         | Any error encountered during beatmap decoding will be logged through `tracing::error`. If this feature is **not** enabled, errors will be ignored. | [`tracing`]
 //!
 //! ## Bindings
 //!

diff --git a/src/mania/difficulty/skills/strain.rs b/src/mania/difficulty/skills/strain.rs
@@ -4,6 +4,7 @@ use crate::{
         skills::{strain_decay, ISkill, Skill, StrainDecaySkill},
     },
     mania::difficulty::object::ManiaDifficultyObject,
+    util::strains_vec::StrainsVec,
 };
 
 const INDIVIDUAL_DECAY_BASE: f64 = 0.125;
@@ -37,7 +38,7 @@ impl Strain {
         }
     }
 
-    pub fn get_curr_strain_peaks(self) -> Vec<f64> {
+    pub fn get_curr_strain_peaks(self) -> StrainsVec {
         self.inner.get_curr_strain_peaks()
     }
 

diff --git a/src/mania/strains.rs b/src/mania/strains.rs
@@ -20,6 +20,6 @@ pub fn strains(difficulty: &Difficulty, converted: &ManiaBeatmap<'_>) -> ManiaSt
     let values = DifficultyValues::calculate(difficulty, converted);
 
     ManiaStrains {
-        strains: values.strain.get_curr_strain_peaks(),
+        strains: values.strain.get_curr_strain_peaks().into_vec(),
     }
 }
diff --git a/src/osu/difficulty/skills/aim.rs b/src/osu/difficulty/skills/aim.rs
@@ -6,7 +6,7 @@ use crate::{
         skills::{strain_decay, ISkill, Skill},
     },
     osu::difficulty::object::OsuDifficultyObject,
-    util::float_ext::FloatExt,
+    util::{float_ext::FloatExt, strains_vec::StrainsVec},
 };
 
 use super::strain::OsuStrainSkill;
@@ -30,7 +30,7 @@ impl Aim {
         }
     }
 
-    pub fn get_curr_strain_peaks(self) -> Vec<f64> {
+    pub fn get_curr_strain_peaks(self) -> StrainsVec {
         self.inner.get_curr_strain_peaks()
     }
 

diff --git a/src/osu/difficulty/skills/flashlight.rs b/src/osu/difficulty/skills/flashlight.rs
@@ -6,7 +6,7 @@ use crate::{
         skills::{strain_decay, ISkill, Skill, StrainSkill},
     },
     osu::{difficulty::object::OsuDifficultyObject, object::OsuObjectKind},
-    util::mods::Mods,
+    util::{mods::Mods, strains_vec::StrainsVec},
 };
 
 use super::strain::OsuStrainSkill;
@@ -33,7 +33,7 @@ impl Flashlight {
         }
     }
 
-    pub fn get_curr_strain_peaks(self) -> Vec<f64> {
+    pub fn get_curr_strain_peaks(self) -> StrainsVec {
         self.inner.get_curr_strain_peaks()
     }
 
@@ -48,8 +48,7 @@ impl Flashlight {
     }
 
     fn static_difficulty_value(skill: StrainSkill) -> f64 {
-        skill.get_curr_strain_peaks().into_iter().sum::<f64>()
-            * OsuStrainSkill::DIFFICULTY_MULTIPLER
+        skill.get_curr_strain_peaks().sum() * OsuStrainSkill::DIFFICULTY_MULTIPLER
     }
 }
 

diff --git a/src/osu/difficulty/skills/speed.rs b/src/osu/difficulty/skills/speed.rs
@@ -6,6 +6,7 @@ use crate::{
         skills::{strain_decay, ISkill, Skill},
     },
     osu::difficulty::object::OsuDifficultyObject,
+    util::strains_vec::StrainsVec,
 };
 
 use super::strain::OsuStrainSkill;
@@ -37,7 +38,7 @@ impl Speed {
         }
     }
 
-    pub fn get_curr_strain_peaks(self) -> Vec<f64> {
+    pub fn get_curr_strain_peaks(self) -> StrainsVec {
         self.inner.get_curr_strain_peaks()
     }