vsrg 0.3.0

use crate::{
    notes::{
        Interval, NotesStorage, NotesStorageEditHook,
        generator::notes_storage::NoteState,
        relations::{AddRelatedSet, AssignPartnersToNote, PairCondition, RemoveFromRelatedSets},
    },
    rhythm::ClockTime,
};

/// Relates notes whose start times differ by at most some tolerance amount.
///
/// # Remarks
/// This relationship only operates in [`crate::rhythm::ClockTime`].
///
/// Times are rounded to the nearest 0.1ms before comparison (halfway values round away from 0).
/// Times passed to callback retain their original precision.
pub struct SameStartTime<S: NotesStorage> {
    tolerance: ClockTime,
    condition: Option<PairCondition<S>>,
    assign_related_notes: Option<AssignPartnersToNote<S>>,
    add_related_set: Option<AddRelatedSet<S>>,
    remove_notes_from_related_set: Option<RemoveFromRelatedSets<S>>,
    clear_related_sets: Option<fn(&mut S)>,

    /// List all notes that are in the timing range, to be filtered down for other uses later.
    scratch_in_range: Vec<S::NoteSnapshotEnum>,
    /// List notes that are to be passed to a callback.
    scratch_callback_arg: Vec<S::NoteSnapshotEnum>,
    /// List notes that are to be passed to remove_notes_from_related_set callback.
    scratch_set_callback_arg: Vec<S::NoteSnapshotEnum>,
}

impl<S: NotesStorage> SameStartTime<S> {
    /// Creates a relationship with no linking callbacks.
    ///
    /// # Remarks
    /// Continue configuring the relationship by calling its builder methods.
    pub const fn new() -> Self {
        Self {
            tolerance: ClockTime::from_duration(std::time::Duration::from_millis(1)),
            condition: None,
            assign_related_notes: None,
            add_related_set: None,
            remove_notes_from_related_set: None,
            clear_related_sets: None,
            scratch_in_range: vec![],
            scratch_callback_arg: vec![],
            scratch_set_callback_arg: vec![],
        }
    }

    /// Clear all relationship sets from the notes storage.
    ///
    /// This is optional, when not set, it calls [`Self::remove_notes_from_related_set`] for
    /// each note instead.
    ///
    /// Triggered after rebuilding notes storage.
    pub fn clear_related_sets(mut self, callback: fn(&mut S)) -> Self {
        self.clear_related_sets = Some(callback);
        self
    }

    /// Sets the inclusive timing tolerance.
    ///
    /// Times are rounded to the nearest 0.1ms before comparison (halfway values round away from 0).
    /// Times passed to callback retain their original precision.
    pub fn tolerance(mut self, value: ClockTime) -> Self {
        self.tolerance = value;
        self
    }

    /// Sets any additional pairing condition (Most common is restricting by note types).
    ///
    /// # Requirements
    /// - The condition must be *commutative*: pair(a, b) = pair(b, a).
    /// - The condition must be *transitive*: if pair(a, b) && pair(b, c), then pair (a, c).
    pub fn condition(mut self, value: PairCondition<S>) -> Self {
        self.condition = Some(value);
        self
    }

    /// Sets the callback that assigns each affected note with its current direct partners, excluding itself.
    ///
    /// # Remarks
    /// The note can copy any data from the partner notes as needed. For example, in cases where you
    /// need to draw a connection to the partner notes, then caching the partner's position in the
    /// note's runtime state may be a good idea.
    ///
    /// The slice contains the note's current direct matches, excluding the note itself. Passing an empty
    /// slice clears note of all partners.
    ///
    /// # Requirements
    /// - This operation must be *stateless*: how it writes should not depend on the current runtime
    ///   state of the note.
    /// - This operation must be *idempotent*: multiple applications of this callback on the same note
    ///   is the same as only one application.
    ///
    /// Effectively, it should sets the partners rather than appending the partners to the note.
    pub fn assign_related_notes(mut self, callback: AssignPartnersToNote<S>) -> Self {
        self.assign_related_notes = Some(callback);
        self
    }

    /// Sets the callback that Add a related set of notes to the notes storage.
    ///
    /// # Remarks
    /// The notes storage itself generated by [`crate::generate_notes_storage`] allows storing
    /// additional user-provided fields. For use cases where relationship state does not belong to any
    /// single note (e.g A connected line between notes of the same timing), then it may be stored in
    /// such user-provided field instead.
    ///
    /// Also configure [`Self::remove_notes_from_related_set`] to handle removals.
    ///
    /// # Requirements
    /// - When called with a set of notes, it should override any previously stored subset of the
    ///   provided set of notes. For example: if it currently stores {(A, B, C), (D, E), (A, D)}, and
    ///   the callback is called with (A, B, C, F), then it should now store {(A, B, C, F), (D, E),
    ///   (A, D)}
    /// - This means that the operation must be *idempotent*, since calling with a set that's already
    ///   stored should override the stored set.
    pub fn add_related_set(mut self, callback: AddRelatedSet<S>) -> Self {
        self.add_related_set = Some(callback);
        self
    }

    /// Sets the callback that removes a note from its related sets.
    ///
    /// # Remarks
    /// The notes storage itself generated by [`crate::generate_notes_storage`] allows storing
    /// additional user-provided fields. For use cases where relationship state does not belong to any
    /// single note (e.g A connected line between notes of the same timing), then it may be stored in
    /// such user-provided field instead.
    ///
    /// Also configure [`Self::add_related_set`] to handle additions.
    ///
    /// # Requirements
    /// - When called with a set of notes, it should those notes from all related sets stored in the
    ///   notes storage. Additionally, you may want to remove sets which now only contains one note.
    ///   For example: if it currently stores {(A, B, C), (D, E), (A, D)}, and the callback is called
    ///   with (C, D), then it should now store {(A, B)}
    /// - The operation must be *idempotent*: Removing the same notes twice is the same as removing
    ///   once.
    pub fn remove_notes_from_related_set(mut self, callback: RemoveFromRelatedSets<S>) -> Self {
        self.remove_notes_from_related_set = Some(callback);
        self
    }
}

fn is_valid_pair_candidate<S: NotesStorage>(
    note: &S::NoteSnapshotEnum,
    candidate: &S::NoteSnapshotEnum,
    tolerance: ClockTime,
    condition: Option<PairCondition<S>>,
) -> bool {
    let time = timing_ticks(note.clock_time_interval().start());
    let tolerance = timing_ticks(tolerance);
    !same_note::<S>(note, candidate)
        && (timing_ticks(candidate.clock_time_interval().start()) - time).abs() <= tolerance
        && condition.is_none_or(|condition| condition(note, candidate))
}

fn same_note<S: NotesStorage>(a: &S::NoteSnapshotEnum, b: &S::NoteSnapshotEnum) -> bool {
    a.id() == b.id() && std::mem::discriminant(a) == std::mem::discriminant(b)
}

fn timing_ticks(time: ClockTime) -> f64 {
    (time.seconds() * 10_000.0).round()
}

// Storage indexes use unrounded times. Include the rounding bins around both
// endpoints, then let the caller filter using rounded ticks. Round the query
// bounds outward as well so conversion back to seconds cannot exclude a match.
fn query_window(center: f64, radius: f64) -> Interval<ClockTime> {
    Interval::new(
        ClockTime::from_seconds(((center - radius - 1.0) / 10_000.0).next_down())
            .unwrap_or(ClockTime::ZERO - ClockTime::INF),
        ClockTime::from_seconds(((center + radius + 1.0) / 10_000.0).next_up())
            .unwrap_or(ClockTime::INF),
    )
}

impl<S: NotesStorage> Default for SameStartTime<S> {
    fn default() -> Self {
        Self::new()
    }
}

impl<S: NotesStorage> SameStartTime<S> {
    fn post_add_snapshot(&mut self, storage: &mut S, just_added: S::NoteSnapshotEnum) {
        self.scratch_in_range.clear();
        self.scratch_in_range
            .extend(storage.notes_overlapping_range(query_window(
                timing_ticks(just_added.clock_time_interval().start()),
                timing_ticks(self.tolerance) * 2.0,
            )));

        if let Some(assign) = self.assign_related_notes {
            for candidate in &self.scratch_in_range {
                if !same_note::<S>(&just_added, candidate)
                    && !is_valid_pair_candidate::<S>(
                        &just_added,
                        candidate,
                        self.tolerance,
                        self.condition,
                    )
                {
                    continue;
                }
                storage.query_notes_mut(
                    [(candidate.id(), S::NoteTypeEnum::from(candidate))].into_iter(),
                    |c| {
                        self.scratch_callback_arg.clear();
                        for n in &self.scratch_in_range {
                            if is_valid_pair_candidate::<S>(
                                candidate,
                                n,
                                self.tolerance,
                                self.condition,
                            ) {
                                self.scratch_callback_arg.push(n.clone());
                            }
                        }

                        assign(c, &self.scratch_callback_arg);
                    },
                );
            }
        }

        if let Some(activate) = self.add_related_set {
            self.scratch_in_range.retain(|candidate| {
                is_valid_pair_candidate::<S>(&just_added, candidate, self.tolerance, self.condition)
            });
            if !self.scratch_in_range.is_empty() {
                self.scratch_in_range.push(just_added);
                activate(storage, &self.scratch_in_range);
            }
        }
    }

    fn post_add_one(&mut self, storage: &mut S, note_type: S::NoteTypeEnum, note_id: u64) {
        let Some(just_added) = storage
            .query_notes_ref([(note_id, note_type)].into_iter())
            .next()
        else {
            return;
        };
        let just_added_snashot = just_added.to_snapshot();
        drop(just_added);
        self.post_add_snapshot(storage, just_added_snashot);
    }

    fn post_remove_one(&mut self, storage: &mut S, just_removed: &S::NoteSnapshotEnum) {
        let time = timing_ticks(just_removed.clock_time_interval().start());
        let tolerance = timing_ticks(self.tolerance);

        self.scratch_in_range.clear();
        self.scratch_in_range.extend(
            storage
                .notes_overlapping_range(query_window(time, tolerance * 2.0))
                .filter(|candidate| !same_note::<S>(just_removed, candidate)),
        );

        self.scratch_set_callback_arg.clear();
        if self.assign_related_notes.is_some() || self.remove_notes_from_related_set.is_some() {
            for candidate in &self.scratch_in_range {
                let candidate_time = timing_ticks(candidate.clock_time_interval().start());
                if candidate_time < time - tolerance || candidate_time > time + tolerance {
                    continue;
                }
                if !is_valid_pair_candidate::<S>(
                    just_removed,
                    candidate,
                    self.tolerance,
                    self.condition,
                ) {
                    continue;
                }

                self.scratch_callback_arg.clear();
                for n in &self.scratch_in_range {
                    if is_valid_pair_candidate::<S>(candidate, n, self.tolerance, self.condition) {
                        self.scratch_callback_arg.push(n.clone());
                    }
                }

                if self.scratch_callback_arg.is_empty() {
                    self.scratch_set_callback_arg.push(candidate.clone());
                }

                if let Some(assign) = self.assign_related_notes {
                    storage.query_notes_mut(
                        [(candidate.id(), S::NoteTypeEnum::from(candidate))].into_iter(),
                        |c| {
                            assign(c, &self.scratch_callback_arg);
                        },
                    );
                }
            }
        }

        if let Some(deactivate) = self.remove_notes_from_related_set {
            self.scratch_set_callback_arg.push(just_removed.clone());
            deactivate(storage, &self.scratch_set_callback_arg);
        }
    }

    fn post_update_one(
        &mut self,
        storage: &mut S,
        before: &S::NoteSnapshotEnum,
        after: &S::NoteSnapshotEnum,
    ) {
        self.post_remove_one(storage, before);
        self.post_add_snapshot(storage, after.clone());
    }
}

impl<S: NotesStorage> NotesStorageEditHook<S> for SameStartTime<S> {
    fn post_rebuild_clock_time(&mut self, storage: &mut S) {
        let notes: Vec<_> = storage.all_notes().collect();
        if let Some(clear) = self.clear_related_sets {
            clear(storage);
        } else if let Some(remove) = self.remove_notes_from_related_set {
            remove(storage, &notes);
        }
        for note in notes {
            self.post_add_one(storage, S::NoteTypeEnum::from(&note), note.id());
        }
    }

    fn post_add(&mut self, storage: &mut S, notes: &[(u64, S::NoteTypeEnum)]) {
        for &(id, ty) in notes {
            self.post_add_one(storage, ty, id);
        }
    }
    fn post_remove(&mut self, storage: &mut S, notes: &[S::NoteSnapshotEnum]) {
        for note in notes {
            self.post_remove_one(storage, note);
        }
    }
    fn post_update(
        &mut self,
        storage: &mut S,
        notes: &[(S::NoteSnapshotEnum, S::NoteSnapshotEnum)],
    ) {
        for (before, after) in notes {
            self.post_update_one(storage, before, after);
        }
    }
}

#[cfg(test)]
mod test {
    use std::collections::HashSet;
    use std::time::Duration;

    use proptest::prelude::*;

    use super::SameStartTime;
    use crate::notes::relations::test_common::{RelLong, RelShort, World, long_at, short_at};
    use crate::notes::{
        LongNoteTypeStorage, NotesStorage, NotesStorageEditHook, ShortNoteTypeStorage,
    };
    use crate::rhythm::ClockTime;

    /// Tolerance of the test notes storages, in milliseconds. A pair is same-timing iff
    /// their independently rounded starts differ by at most this, inclusively.
    const TOLERANCE_MS: f64 = 10.0;

    crate::generate_notes_storage! {
        enum TestNoteType,

        #[derive(PartialEq)]
        enum TestNote,

        #[derive(Debug, PartialEq)]
        enum TestNoteSnapshot,

        enum TestNoteRef,

        enum TestNoteMut,

        #[derive(Debug, Clone)]
        struct TestNotesStorage {
            note_type Short(RelShort) => shorts: ShortNoteTypeStorage<RelShort>,
            note_type Long(RelLong) => longs: LongNoteTypeStorage<RelLong>,
            field chords: Vec<HashSet<TestNoteId>>,
            edit_hook SameStartTime::<TestNotesStorage>::new()
                .tolerance(ClockTime::from_duration(Duration::from_millis(10)))
                .assign_related_notes(set_chord)
                .add_related_set(add_chord_to_field)
                .remove_notes_from_related_set(remove_chords_from_field),
        }
    }

    impl TestNotesStorage {
        fn health_check(&self) {
            self.shorts.health_check();
            self.longs.health_check();
        }
    }

    // Like `TestNotesStorage`, but only notes whose positions sum to 10 (a "mirror pair")
    // count as same-timing partners.
    crate::generate_notes_storage! {
        enum CondNoteType,

        #[derive(PartialEq)]
        enum CondNote,

        #[derive(Debug, PartialEq)]
        enum CondNoteSnapshot,

        enum CondNoteRef,

        enum CondNoteMut,

        #[derive(Debug, Clone)]
        struct CondNotesStorage {
            note_type Short(RelShort) => shorts: ShortNoteTypeStorage<RelShort>,
            note_type Long(RelLong) => longs: LongNoteTypeStorage<RelLong>,
            edit_hook SameStartTime::<CondNotesStorage>::new()
                .tolerance(ClockTime::from_duration(Duration::from_millis(10)))
                .condition(mirror_pos)
                .assign_related_notes(cond_set_chord)
        }
    }

    impl CondNotesStorage {
        fn health_check(&self) {
            self.shorts.health_check();
            self.longs.health_check();
        }
    }

    fn set_chord(mut state: TestNoteMut<'_>, matches: &[TestNoteSnapshot]) {
        let mut matching_ids: Vec<_> = matches
            .iter()
            .map(|note| (note.id(), matches!(note, TestNoteSnapshot::Short(_))))
            .collect();
        matching_ids.sort_unstable();
        assert!(matching_ids.windows(2).all(|pair| pair[0] != pair[1]));
        match &mut state {
            TestNoteMut::Short(s) => *s.runtime_data.matching_ids = matching_ids,
            TestNoteMut::Long(s) => *s.runtime_data.matching_ids = matching_ids,
        }
        assert!(matches.iter().all(|note| note.id() != state.id()

            || std::mem::discriminant(&note.to_owned_note())
                != std::mem::discriminant(&state.to_owned_note())));
        match state {
            TestNoteMut::Short(s) => *s.runtime_data.chord = !matches.is_empty(),
            TestNoteMut::Long(s) => *s.runtime_data.chord = !matches.is_empty(),
        }
    }

    fn add_chord_to_field(storage: &mut TestNotesStorage, notes: &[TestNoteSnapshot]) {
        let chord = notes
            .iter()
            .map(TestNoteId::from_snapshot)
            .collect::<HashSet<_>>();
        if storage
            .chords
            .iter()
            .any(|current| chord.is_subset(current))
        {
            return;
        }
        storage.chords.retain(|current| !current.is_subset(&chord));
        storage.chords.push(chord);
    }

    fn remove_chords_from_field(storage: &mut TestNotesStorage, notes: &[TestNoteSnapshot]) {
        let ids: Vec<TestNoteId> = notes.iter().map(TestNoteId::from_snapshot).collect();
        for cur_choord in storage.chords.iter_mut() {
            cur_choord.retain(|id| !ids.contains(id));
        }
        storage.chords.retain(|chord| chord.len() > 1);
    }

    fn mirror_pos(a: &CondNoteSnapshot, b: &CondNoteSnapshot) -> bool {
        let pos = |n: &CondNoteSnapshot| match n {
            CondNoteSnapshot::Short(s) => s.note.pos,
            CondNoteSnapshot::Long(s) => s.note.pos,
        };
        pos(a) + pos(b) == 10
    }

    fn cond_set_chord(state: CondNoteMut<'_>, matches: &[CondNoteSnapshot]) {
        assert!(matches.iter().all(|note| note.id() != state.id()

            || std::mem::discriminant(&note.to_owned_note())
                != std::mem::discriminant(&state.to_owned_note())));
        match state {
            CondNoteMut::Short(s) => *s.runtime_data.chord = !matches.is_empty(),
            CondNoteMut::Long(s) => *s.runtime_data.chord = !matches.is_empty(),
        }
    }

    /// Identifies a note across the two storages: ids are only unique within each
    /// note type's storage.
    #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
    enum TestNoteId {
        Short(u64),
        Long(u64),
    }

    impl TestNoteId {
        fn parts(self) -> (u64, TestNoteType) {
            match self {
                Self::Short(id) => (id, TestNoteType::Short),
                Self::Long(id) => (id, TestNoteType::Long),
            }
        }

        fn from_snapshot(note: &TestNoteSnapshot) -> Self {
            match note {
                TestNoteSnapshot::Short(n) => Self::Short(n.id),
                TestNoteSnapshot::Long(n) => Self::Long(n.id),
            }
        }
    }

    fn snote(ms: f64) -> TestNote {
        TestNote::Short(short_at(ms, 0))
    }

    fn lnote(start_ms: f64, len_ms: f64) -> TestNote {
        TestNote::Long(long_at(start_ms, len_ms, 0))
    }

    fn cnote(ms: f64, pos: u32) -> CondNote {
        CondNote::Short(short_at(ms, pos))
    }

    fn chord_flag(storage: &TestNotesStorage, id: TestNoteId) -> Option<bool> {
        storage
            .query_notes_ref(std::iter::once(id.parts()))
            .next()
            .map(|s| match s {
                TestNoteRef::Short(s) => *s.runtime_data.chord,
                TestNoteRef::Long(s) => *s.runtime_data.chord,
            })
    }

    fn chord_assigned(storage: &TestNotesStorage, id: TestNoteId) -> Vec<TestNoteId> {
        let state = storage
            .query_notes_ref([id.parts()].into_iter())
            .next()
            .unwrap();
        let ids = match state {
            TestNoteRef::Short(n) => n.runtime_data.matching_ids,
            TestNoteRef::Long(n) => n.runtime_data.matching_ids,
        };
        let mut result: Vec<_> = ids
            .iter()
            .map(|&(id, short)| {
                if short {
                    TestNoteId::Short(id)
                } else {
                    TestNoteId::Long(id)
                }
            })
            .collect();
        result.sort_unstable();
        result
    }

    fn cond_chord_flag(storage: &CondNotesStorage, id: u64) -> Option<bool> {
        storage
            .query_notes_ref(std::iter::once((id, CondNoteType::Short)))
            .next()
            .map(|s| match s {
                CondNoteRef::Short(s) => *s.runtime_data.chord,
                _ => false,
            })
    }

    fn build_notes_storage(world: &mut World, notes: Vec<TestNote>) -> TestNotesStorage {
        TestNotesStorage::with_notes(
            notes,
            &mut world.note_groups,
            &mut world.tempo_tracks,
            &mut world.scroll_tracks,
        )
    }

    fn add_note(world: &mut World, storage: &mut TestNotesStorage, note: TestNote) {
        storage.add_note(
            &mut world.note_groups,
            &mut world.tempo_tracks,
            &mut world.scroll_tracks,
            note,
        );
    }

    fn add_notes(world: &mut World, storage: &mut TestNotesStorage, notes: Vec<TestNote>) {
        storage.add_notes(
            &mut world.note_groups,
            &mut world.tempo_tracks,
            &mut world.scroll_tracks,
            notes,
        );
    }

    fn replace_note(world: &mut World, storage: &mut TestNotesStorage, id: u64, note: TestNote) {
        storage.replace_note(
            &mut world.note_groups,
            &mut world.tempo_tracks,
            &mut world.scroll_tracks,
            id,
            note,
        );
    }

    fn replace_notes(
        world: &mut World,
        storage: &mut TestNotesStorage,
        notes: Vec<(u64, TestNote)>,
    ) {
        storage.replace_notes(
            &mut world.note_groups,
            &mut world.tempo_tracks,
            &mut world.scroll_tracks,
            notes,
        );
    }

    fn setup(notes: Vec<TestNote>) -> (World, TestNotesStorage) {
        let mut world = World::new(0);
        let storage = build_notes_storage(&mut world, notes);
        (world, storage)
    }

    fn cond_setup(notes: Vec<CondNote>) -> (World, CondNotesStorage) {
        let mut world = World::new(0);
        let storage = CondNotesStorage::with_notes(
            notes,
            &mut world.note_groups,
            &mut world.tempo_tracks,
            &mut world.scroll_tracks,
        );
        (world, storage)
    }

    #[test]
    fn solo_note_stays_inactive() {
        let (_, storage) = setup(vec![snote(100.0)]);
        assert_eq!(chord_flag(&storage, TestNoteId::Short(1)), Some(false));
        assert!(storage.chords.is_empty(), "{:?}", storage.chords);
        storage.health_check();
    }

    #[test]
    fn pair_within_tolerance_activates_both() {
        let (_, storage) = setup(vec![snote(100.0), snote(105.0)]);
        assert_eq!(chord_flag(&storage, TestNoteId::Short(1)), Some(true));
        assert_eq!(
            chord_assigned(&storage, TestNoteId::Short(1)),
            vec![TestNoteId::Short(2)]
        );
        assert_eq!(chord_flag(&storage, TestNoteId::Short(2)), Some(true));
        assert_eq!(
            chord_assigned(&storage, TestNoteId::Short(2)),
            vec![TestNoteId::Short(1)]
        );
        assert_eq!(
            storage.chords,
            vec![
                [TestNoteId::Short(1), TestNoteId::Short(2)]
                    .into_iter()
                    .collect()
            ]
        );
        storage.health_check();
    }

    #[test]
    fn note_beyond_tolerance_stays_inactive() {
        let (_, storage) = setup(vec![snote(100.0), snote(130.0)]);
        assert_eq!(chord_flag(&storage, TestNoteId::Short(1)), Some(false));
        assert_eq!(chord_flag(&storage, TestNoteId::Short(2)), Some(false));
        storage.health_check();
    }

    #[test]
    fn tolerance_boundary_is_inclusive() {
        let (_, storage) = setup(vec![snote(100.0), snote(110.0)]);
        assert_eq!(chord_flag(&storage, TestNoteId::Short(1)), Some(true));
        assert_eq!(
            chord_assigned(&storage, TestNoteId::Short(1)),
            vec![TestNoteId::Short(2)]
        );
        assert_eq!(chord_flag(&storage, TestNoteId::Short(2)), Some(true));
        assert_eq!(
            chord_assigned(&storage, TestNoteId::Short(2)),
            vec![TestNoteId::Short(1)]
        );

        let (_, storage) = setup(vec![snote(100.0), snote(111.0)]);
        assert_eq!(chord_flag(&storage, TestNoteId::Short(1)), Some(false));
        assert_eq!(chord_flag(&storage, TestNoteId::Short(2)), Some(false));
    }

    #[test]
    fn rounded_timing_boundaries() {
        for (a, b, expected) in [
            (-60.0, -70.0, true),
            (100.0, 110.049, true),
            (100.0, 110.051, false),
            (-100.0, -110.049, true),
            (-100.0, -110.051, false),
            // Round each start, rather than their difference.
            (100.049, 110.051, false),
        ] {
            for notes in [
                vec![snote(a), lnote(b, 50.0)],
                vec![lnote(b, 50.0), snote(a)],
            ] {
                let (_, storage) = setup(notes);
                assert_eq!(chord_flag(&storage, TestNoteId::Short(1)), Some(expected));
                assert_eq!(chord_flag(&storage, TestNoteId::Long(1)), Some(expected));
                storage.health_check();
            }
        }
    }

    #[test]
    fn rounded_timing_removal_and_update() {
        let (mut world, mut storage) = setup(vec![snote(100.049), snote(110.0)]);
        replace_note(&mut world, &mut storage, 2, snote(110.049));
        assert_matches_naive(&storage);
        replace_note(&mut world, &mut storage, 2, snote(110.051));
        assert_matches_naive(&storage);

        // The outer partner lies beyond the unrounded double-tolerance window.
        let (_, mut storage) = setup(vec![snote(100.0), snote(110.049), snote(120.049)]);
        storage.remove_note(TestNoteType::Short, 1);
        assert_matches_naive(&storage);
        storage.remove_note(TestNoteType::Short, 3);
        assert_matches_naive(&storage);
        storage.health_check();
    }

    #[test]
    fn rounding_halfway_values_and_tolerance() {
        assert_eq!(
            super::timing_ticks(ClockTime::from_milis(0.05).unwrap()),
            1.0
        );
        assert_eq!(
            super::timing_ticks(ClockTime::from_milis(-0.05).unwrap()),
            -1.0
        );
        for (tolerance_ms, expected) in [(10.049, false), (10.051, true)] {
            let (_, mut storage) = setup(vec![snote(100.0), snote(110.1)]);
            let mut relationship = SameStartTime::<TestNotesStorage>::new()
                .tolerance(ClockTime::from_milis(tolerance_ms).unwrap())
                .assign_related_notes(set_chord);
            let note = storage
                .query_notes([(1, TestNoteType::Short)].into_iter())
                .next()
                .unwrap();
            relationship.post_add_snapshot(&mut storage, note);
            assert_eq!(chord_flag(&storage, TestNoteId::Short(1)), Some(expected));
            assert_eq!(chord_flag(&storage, TestNoteId::Short(2)), Some(expected));
        }
    }

    #[test]
    fn third_note_joins_cluster() {
        let (_, storage) = setup(vec![snote(100.0), snote(105.0), snote(108.0)]);
        for id in 1..=3 {
            assert_eq!(chord_flag(&storage, TestNoteId::Short(id)), Some(true));
        }
        assert_eq!(
            chord_assigned(&storage, TestNoteId::Short(1)),
            vec![TestNoteId::Short(2), TestNoteId::Short(3)]
        );
        assert_eq!(
            chord_assigned(&storage, TestNoteId::Short(2)),
            vec![TestNoteId::Short(1), TestNoteId::Short(3)]
        );
        assert_eq!(
            chord_assigned(&storage, TestNoteId::Short(3)),
            vec![TestNoteId::Short(1), TestNoteId::Short(2)]
        );
        storage.health_check();
    }

    #[test]
    fn mixed_types_form_chord() {
        let (_, storage) = setup(vec![snote(100.0), lnote(105.0, 50.0)]);
        assert_eq!(chord_flag(&storage, TestNoteId::Short(1)), Some(true));
        assert_eq!(
            chord_assigned(&storage, TestNoteId::Short(1)),
            vec![TestNoteId::Long(1)]
        );
        assert_eq!(chord_flag(&storage, TestNoteId::Long(1)), Some(true));
        assert_eq!(
            chord_assigned(&storage, TestNoteId::Long(1)),
            vec![TestNoteId::Short(1)]
        );
        storage.health_check();
    }

    #[test]
    fn adding_neighbor_preserves_outer_partners() {
        let (_, storage) = setup(vec![snote(100.0), snote(110.0), snote(90.0)]);
        assert_matches_naive(&storage);
    }

    #[test]
    fn storage_callbacks_work_without_note_assignment() {
        use crate::notes::NotesStorageEditHook;
        let (_, mut storage) = setup(vec![snote(100.0), snote(105.0)]);
        let removed = storage
            .query_notes([(2, TestNoteType::Short)].into_iter())
            .next()
            .unwrap();
        // Bypass the generated hook to exercise a relation with only set callbacks.
        storage.shorts.remove_notes([2].into_iter());
        let mut relation = SameStartTime::<TestNotesStorage>::new()
            .tolerance(ClockTime::from_milis(TOLERANCE_MS).unwrap())
            .remove_notes_from_related_set(|storage, notes| {
                assert_eq!(notes.len(), 2);
                remove_chords_from_field(storage, notes);
            });
        relation.post_remove(&mut storage, ::std::slice::from_ref(&removed));
        assert!(storage.chords.is_empty());
    }

    #[test]
    fn removing_one_of_pair_deactivates_survivor() {
        let (_, mut storage) = setup(vec![snote(100.0), snote(105.0)]);
        storage.remove_note(TestNoteType::Short, 2);
        assert_eq!(chord_flag(&storage, TestNoteId::Short(1)), Some(false));
        assert!(storage.chords.is_empty());
        storage.health_check();
    }

    #[test]
    fn removing_bridge_deactivates_both_isolated_neighbors() {
        let (_, mut storage) = setup(vec![snote(89.0), snote(100.0), snote(111.0)]);
        storage.remove_note(TestNoteType::Short, 2);

        assert_eq!(chord_flag(&storage, TestNoteId::Short(1)), Some(false));
        assert_eq!(chord_flag(&storage, TestNoteId::Short(3)), Some(false));
        assert!(storage.chords.is_empty());
    }

    #[test]
    fn removing_bridge_keeps_neighbors_with_outer_partners_active() {
        let (_, mut storage) = setup(vec![
            snote(80.0),
            snote(90.0),
            snote(100.0),
            snote(110.0),
            snote(120.0),
        ]);
        storage.remove_note(TestNoteType::Short, 3);

        for id in [1, 2, 4, 5] {
            assert_eq!(chord_flag(&storage, TestNoteId::Short(id)), Some(true));
        }
        assert!(
            storage.chords.contains(
                &[TestNoteId::Short(1), TestNoteId::Short(2)]
                    .into_iter()
                    .collect()
            )
        );
        assert!(
            storage.chords.contains(
                &[TestNoteId::Short(4), TestNoteId::Short(5)]
                    .into_iter()
                    .collect()
            )
        );
        assert!(
            storage
                .chords
                .iter()
                .all(|chord| !chord.contains(&TestNoteId::Short(3)))
        );
    }

    #[test]
    fn removing_one_of_triple_keeps_pair() {
        let (_, mut storage) = setup(vec![snote(100.0), snote(105.0), snote(108.0)]);
        storage.remove_note(TestNoteType::Short, 2);
        assert_eq!(chord_flag(&storage, TestNoteId::Short(1)), Some(true));
        assert_eq!(chord_flag(&storage, TestNoteId::Short(3)), Some(true));
        storage.health_check();
    }

    #[test]
    fn update_away_deactivates_both() {
        let (mut world, mut storage) = setup(vec![snote(100.0), snote(105.0)]);
        replace_note(&mut world, &mut storage, 2, snote(500.0));
        assert_eq!(chord_flag(&storage, TestNoteId::Short(1)), Some(false));
        assert_eq!(chord_flag(&storage, TestNoteId::Short(2)), Some(false));
        assert!(storage.chords.is_empty());
        storage.health_check();
    }

    #[test]
    fn update_into_cluster_activates() {
        let (mut world, mut storage) = setup(vec![snote(100.0), snote(500.0)]);
        replace_note(&mut world, &mut storage, 2, snote(105.0));
        assert_eq!(chord_flag(&storage, TestNoteId::Short(1)), Some(true));
        assert_eq!(chord_flag(&storage, TestNoteId::Short(2)), Some(true));
        storage.health_check();
    }

    #[test]
    fn update_moves_between_clusters() {
        let (mut world, mut storage) =
            setup(vec![snote(100.0), snote(105.0), snote(500.0), snote(505.0)]);
        // Move note 2 out of cluster A (100, 105) into cluster B (500, 505).
        replace_note(&mut world, &mut storage, 2, snote(502.0));
        // Cluster A lost its only pair: the remaining note deactivates.
        assert_eq!(chord_flag(&storage, TestNoteId::Short(1)), Some(false));
        // Cluster B now holds three notes: all active.
        assert_eq!(chord_flag(&storage, TestNoteId::Short(2)), Some(true));
        assert_eq!(chord_flag(&storage, TestNoteId::Short(3)), Some(true));
        assert_eq!(chord_flag(&storage, TestNoteId::Short(4)), Some(true));
        storage.health_check();
    }

    #[test]
    fn batch_add_forms_chord() {
        let (mut world, mut storage) = setup(vec![]);
        add_notes(
            &mut world,
            &mut storage,
            vec![snote(100.0), snote(105.0), snote(110.0)],
        );
        for id in 1..=3 {
            assert_eq!(chord_flag(&storage, TestNoteId::Short(id)), Some(true));
        }
        storage.health_check();
    }

    #[test]
    fn condition_must_accept_both_notes() {
        // Positions 3 + 7 == 10: the mirror condition accepts the pair.
        let (_, storage) = cond_setup(vec![cnote(100.0, 3), cnote(100.0, 7)]);
        assert_eq!(cond_chord_flag(&storage, 1), Some(true));
        assert_eq!(cond_chord_flag(&storage, 2), Some(true));

        // Positions 3 + 4 != 10: same timing, but no chord.
        let (_, storage) = cond_setup(vec![cnote(100.0, 3), cnote(100.0, 4)]);
        assert_eq!(cond_chord_flag(&storage, 1), Some(false));
        assert_eq!(cond_chord_flag(&storage, 2), Some(false));
    }

    #[test]
    fn condition_rechecked_on_update() {
        let (mut world, mut storage) = cond_setup(vec![cnote(100.0, 3), cnote(100.0, 7)]);
        assert_eq!(cond_chord_flag(&storage, 1), Some(true));
        // Breaking the mirror condition must deactivate both notes.
        storage.replace_note(
            &mut world.note_groups,
            &mut world.tempo_tracks,
            &mut world.scroll_tracks,
            2,
            cnote(100.0, 5),
        );
        assert_eq!(cond_chord_flag(&storage, 1), Some(false));
        assert_eq!(cond_chord_flag(&storage, 2), Some(false));
        storage.health_check();
    }

    /// Random edit sequences must keep the invariant
    /// `chord == the note has at least one same-timing partner`
    /// (and related sets cover all chord-active notes) after every edit,
    /// matching a naive recompute from final state.
    #[derive(Debug, Clone)]
    enum Op {
        Add(TestNote),
        AddBatch(Vec<TestNote>),
        Remove(TestNoteId),
        RemoveBatch(Vec<TestNoteId>),
        Replace(TestNoteId, TestNote),
        ReplaceBatch(Vec<(TestNoteId, TestNote)>),
    }

    /// Times on a 10ms grid, so tolerance windows actually overlap.
    fn arb_time() -> impl Strategy<Value = f64> {
        (-60..60i32).prop_map(|t| (t * 10) as f64)
    }

    fn arb_pos() -> impl Strategy<Value = u32> {
        0..10u32
    }

    fn arb_short() -> impl Strategy<Value = TestNote> {
        (arb_time(), arb_pos()).prop_map(|(t, pos)| TestNote::Short(short_at(t, pos)))
    }

    fn arb_long() -> impl Strategy<Value = TestNote> {
        (arb_time(), 0..40i32, arb_pos())
            .prop_map(|(s, len, pos)| TestNote::Long(long_at(s, (len * 10) as f64, pos)))
    }

    fn arb_note() -> impl Strategy<Value = TestNote> {
        prop_oneof![arb_short(), arb_long()]
    }

    fn arb_id() -> impl Strategy<Value = TestNoteId> {
        (1..=40u64, prop::bool::ANY).prop_map(|(id, is_short)| {
            if is_short {
                TestNoteId::Short(id)
            } else {
                TestNoteId::Long(id)
            }
        })
    }

    fn arb_op() -> impl Strategy<Value = Op> {
        prop_oneof![
            arb_note().prop_map(Op::Add),
            prop::collection::vec(arb_note(), 0..4).prop_map(Op::AddBatch),
            arb_id().prop_map(Op::Remove),
            prop::collection::vec(arb_id(), 0..4).prop_map(Op::RemoveBatch),
            (arb_id(), arb_note()).prop_map(|(id, note)| Op::Replace(id, note)),
            prop::collection::vec((arb_id(), arb_note()), 0..4).prop_map(Op::ReplaceBatch),
        ]
    }

    fn apply_op(world: &mut World, storage: &mut TestNotesStorage, op: Op) {
        match op {
            Op::Add(note) => add_note(world, storage, note),
            Op::AddBatch(notes) => add_notes(world, storage, notes),
            Op::Remove(id) => {
                let (id, ty) = id.parts();
                storage.remove_note(ty, id);
            }
            Op::RemoveBatch(ids) => {
                storage.remove_notes(ids.into_iter().map(|id| id.parts()).collect())
            }
            Op::Replace(id, note) => replace_note(world, storage, id.parts().0, note),
            Op::ReplaceBatch(pairs) => replace_notes(
                world,
                storage,
                pairs
                    .into_iter()
                    .map(|(id, note)| (id.parts().0, note))
                    .collect(),
            ),
        }
    }

    /// All notes of the notes storage with their id and rounded start time (in 0.1 ms ticks).
    fn all_starts(storage: &TestNotesStorage) -> Vec<(TestNoteId, f64)> {
        let mut notes: Vec<(TestNoteId, f64)> = storage
            .shorts
            .time_sorted_groups()
            .iter()
            .flat_map(|group| group.note_id().iter().zip(group.clock_time()))
            .map(|(&id, &t)| (TestNoteId::Short(id), (t.seconds() * 10_000.0).round()))
            .collect();
        notes.extend(
            storage
                .longs
                .time_sorted_groups()
                .iter()
                .flat_map(|group| group.note_id().iter().zip(group.clock_time()))
                .map(|(&id, &i)| {
                    (
                        TestNoteId::Long(id),
                        (i.start().seconds() * 10_000.0).round(),
                    )
                }),
        );
        notes
    }

    fn has_partner(notes: &[(TestNoteId, f64)], id: TestNoteId, t: f64) -> bool {
        notes
            .iter()
            .any(|(id2, t2)| *id2 != id && (t - t2).abs() <= (TOLERANCE_MS * 10.0).round())
    }

    fn assert_matches_naive(storage: &TestNotesStorage) {
        let notes = all_starts(storage);
        for (id, t) in &notes {
            let expected = has_partner(&notes, *id, *t);
            let mut expected_matches: Vec<_> = notes
                .iter()
                .filter(|(other, time)| {
                    other != id && (t - time).abs() <= (TOLERANCE_MS * 10.0).round()
                })
                .map(|(other, _)| match other {
                    TestNoteId::Short(id) => (*id, true),
                    TestNoteId::Long(id) => (*id, false),
                })
                .collect();
            expected_matches.sort_unstable();
            expected_matches.dedup();
            let state = storage
                .query_notes_ref([id.parts()].into_iter())
                .next()
                .unwrap();
            let actual = match state {
                TestNoteRef::Short(n) => n.runtime_data.matching_ids,
                TestNoteRef::Long(n) => n.runtime_data.matching_ids,
            };
            assert_eq!(
                actual, &expected_matches,
                "matching notes of {id:?} after edit"
            );
            assert_eq!(
                chord_flag(storage, *id),
                Some(expected),
                "chord flag of {id:?} after edit"
            );
        }

        // Sets can overlap; their union must cover exactly the active notes.
        assert!(storage.chords.iter().all(|chord| chord.len() > 1));
        let mut expected_highlight: Vec<TestNoteId> = notes
            .iter()
            .filter(|(id, t)| has_partner(&notes, *id, *t))
            .map(|(id, _)| *id)
            .collect();
        expected_highlight.sort_unstable();
        expected_highlight.dedup();

        let mut actual: Vec<_> = storage.chords.iter().flatten().copied().collect();
        actual.sort_unstable();
        actual.dedup();
        assert_eq!(actual, expected_highlight);
    }

    fn clear_relationships(storage: &mut TestNotesStorage) {
        for ty in [TestNoteType::Short, TestNoteType::Long] {
            storage.all_notes_mut_of_type(ty, |note| match note {
                TestNoteMut::Short(n) => {
                    n.runtime_data.matching_ids.clear();
                    *n.runtime_data.chord = false;
                }
                TestNoteMut::Long(n) => {
                    n.runtime_data.matching_ids.clear();
                    *n.runtime_data.chord = false;
                }
            });
        }
        storage.chords.clear();
    }

    #[test]
    fn rebuild_restores_cleared_relationships() {
        let (mut world, mut storage) = setup(vec![
            snote(0.0),
            lnote(5.0, 95.0),
            snote(100.0),
            snote(500.0),
        ]);
        clear_relationships(&mut storage);
        for _ in 0..2 {
            storage.rebuild_clock_time(
                &mut world.note_groups,
                &mut world.tempo_tracks,
                &mut world.scroll_tracks,
            );
            assert_matches_naive(&storage);
            storage.health_check();
        }
        let mut empty_world = World::new(0);
        let mut empty = TestNotesStorage::new();
        empty.rebuild_clock_time(
            &mut empty_world.note_groups,
            &mut empty_world.tempo_tracks,
            &mut empty_world.scroll_tracks,
        );
        assert_matches_naive(&empty);
    }

    #[test]
    fn rebuild_clears_auxiliary_sets_with_missing_ids() {
        let mut storage = TestNotesStorage::new();
        storage.chords.push(
            [TestNoteId::Short(999), TestNoteId::Long(999)]
                .into_iter()
                .collect(),
        );
        SameStartTime::<TestNotesStorage>::new()
            .clear_related_sets(|s| s.chords.clear())
            .post_rebuild_clock_time(&mut storage);
        assert!(storage.chords.is_empty());
    }

    proptest! {
        #[test]
        fn rebuild_matches_naive(notes in prop::collection::vec(arb_note(), 0..40)) {
            let (mut world, mut storage) = setup(notes);
            clear_relationships(&mut storage);
            storage.rebuild_clock_time(
                &mut world.note_groups,
                &mut world.tempo_tracks,
                &mut world.scroll_tracks,
            );
            assert_matches_naive(&storage);
            storage.health_check();
        }

        #[test]
        fn rounded_pairs_match_naive(
            a in -200_000i32..200_000,
            separation in -10_200i32..10_200,
        ) {
            // Microsecond inputs exercise both sides of each 0.1 ms rounding bin.
            let (_, storage) = setup(vec![
                snote(a as f64 / 1000.0),
                lnote((a + separation) as f64 / 1000.0, 50.0),
            ]);
            assert_matches_naive(&storage);
            storage.health_check();
        }

        #[test]
        fn arb_edits_match_naive(ops in prop::collection::vec(arb_op(), 0..20)) {
            let mut world = World::new(0);
            let mut storage = build_notes_storage(&mut world, vec![]);
            for op in ops {
                apply_op(&mut world, &mut storage, op);
                assert_matches_naive(&storage);
                storage.health_check();
            }
        }
    }
}