Make JournalStorage.may_omit_trials() reflect apply_discard

rustuna_pyo3/tests/storage_tests/test_storage.py

@@ -164,3 +164,31 @@ def test_get_n_trials_with_states_filter(storage: StorageProtocol) -> None:
         == 2
     )
     assert storage.get_n_trials(study._study_id, states=[]) == 0
+
+
+@pytest.mark.parametrize("storage_name", ["inmemory", "sqlite3", "journal-file"])
+@pytest.mark.parametrize("apply_discard", [False, True])
+def test_may_omit_trials_follows_apply_discard(
+    storage_name: str, apply_discard: bool
+) -> None:
+    with tempfile.TemporaryDirectory() as workdir:
+        storage: StorageProtocol
+        if storage_name == "inmemory":
+            storage = rustuna.storages.InMemoryStorage(apply_discard=apply_discard)
+        elif storage_name == "sqlite3":
+            storage = rustuna.storages.SQLite3Storage(
+                f"{workdir}/test.db", apply_discard=apply_discard
+            )
+        else:
+            storage = rustuna.storages.JournalFileStorage(
+                f"{workdir}/test.journal", apply_discard=apply_discard
+            )
+
+        assert storage.may_omit_trials() is apply_discard
+
+        study = storage.create_new_study(
+            "example study", [rustuna.study.StudyDirection.MINIMIZE]
+        )
+        trial_id = storage.create_new_trial(study.id)._trial_id
+        storage.discard_trials([trial_id])
+        assert storage.may_omit_trials() is apply_discard

rustuna_storage/src/journal/storage.rs

@@ -848,10 +848,7 @@ impl Storage for JournalStorage {
     }
 
     fn may_omit_trials(&self) -> bool {
-        self.replay
-            .trials_by_study
-            .values()
-            .any(|trials| trials.iter().any(Option::is_none))
+        self.replay.apply_discard
     }
 
     fn attr_format(&self) -> AttrFormat {
@@ -2588,6 +2585,29 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn may_omit_trials_follows_apply_discard_before_any_discard() -> Result<()> {
+        for apply_discard in [false, true] {
+            let backend = InMemoryJournalBackend {
+                logs: Arc::new(Mutex::new(Vec::new())),
+            };
+            let mut storage = JournalStorage::new_with_options(
+                Box::new(backend),
+                JournalStorageOptions {
+                    apply_discard,
+                    ..Default::default()
+                },
+            )?;
+            assert_eq!(storage.may_omit_trials(), apply_discard);
+
+            let study_id = storage.create_new_study("s", vec![Direction::Minimize])?.id;
+            let trial_id = storage.create_new_trial(study_id)?.id;
+            storage.discard_trials(&[trial_id])?;
+            assert_eq!(storage.may_omit_trials(), apply_discard);
+        }
+        Ok(())
+    }
+
     #[test]
     fn get_n_trials_counts_states_including_discarded_trials() -> Result<()> {
         let logs = Arc::new(Mutex::new(Vec::new()));