vsrg 0.6.0

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

/// Relates notes when time endpoints differ by at most some tolerance amount.
///
/// # Remarks
/// This relationship is a superset of [`super::same_start_time::SameStartTime`] that also notes when
/// any start/start, start/end, end/start, or end/end pair differs by at most the inclusive tolerance.
///
/// 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 MatchingTimeEndpoint<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)>,

    scratch_snapshot: Vec<S::NoteSnapshotEnum>,
    scratch_candidates: Vec<S::NoteSnapshotEnum>,
    scratch_matches: Vec<S::NoteSnapshotEnum>,

    // Capture every partner batch before callbacks can mutate note state.
    scratch_partners: Vec<S::NoteSnapshotEnum>,
    scratch_partner_ranges: Vec<std::ops::Range<usize>>,
}

impl<S: NotesStorage> MatchingTimeEndpoint<S> {
    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_snapshot: vec![],
            scratch_candidates: vec![],
            scratch_matches: vec![],
            scratch_partners: vec![],
            scratch_partner_ranges: 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 matching_notes(
        storage: &S,
        note: &S::NoteSnapshotEnum,
        tolerance: ClockTime,
        condition: Option<PairCondition<S>>,
        matches: &mut Vec<S::NoteSnapshotEnum>,
    ) {
        matches.clear();
        collect_candidates(storage, note, timing_ticks(tolerance), matches);
        matches.retain(|candidate| {
            !same_note::<S>(note, candidate)
                && condition.is_none_or(|condition| condition(note, candidate))
        });
    }

    fn post_add_snapshot(&mut self, storage: &mut S, just_added: S::NoteSnapshotEnum) {
        Self::matching_notes(
            storage,
            &just_added,
            self.tolerance,
            self.condition,
            &mut self.scratch_snapshot,
        );

        self.scratch_snapshot.push(just_added);
        if let Some(activate) = self.assign_related_notes {
            self.scratch_partners.clear();
            self.scratch_partner_ranges.clear();
            for note in &self.scratch_snapshot {
                Self::matching_notes(
                    storage,
                    note,
                    self.tolerance,
                    self.condition,
                    &mut self.scratch_matches,
                );
                let start = self.scratch_partners.len();
                self.scratch_partners.append(&mut self.scratch_matches);
                self.scratch_partner_ranges
                    .push(start..self.scratch_partners.len());
            }

            for (note, range) in self
                .scratch_snapshot
                .iter()
                .zip(&self.scratch_partner_ranges)
            {
                let partners = &self.scratch_partners[range.clone()];
                storage.query_notes_mut(
                    [(note.id(), S::NoteTypeEnum::from(note))].into_iter(),
                    |state| activate(state, partners),
                );
            }
        }

        if self.scratch_snapshot.len() > 1
            && let Some(activate) = self.add_related_set
        {
            activate(storage, &self.scratch_snapshot);
        }
    }
}

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

// Query both narrow endpoint windows.
fn collect_candidates<S: NotesStorage>(
    storage: &S,
    note: &S::NoteSnapshotEnum,
    tolerance: f64,
    out: &mut Vec<S::NoteSnapshotEnum>,
) {
    for endpoint in [
        note.clock_time_interval().start(),
        note.clock_time_interval().end(),
    ] {
        let time = timing_ticks(endpoint);
        for candidate in storage.notes_overlapping_range(query_window(time, tolerance)) {
            if [
                candidate.clock_time_interval().start(),
                candidate.clock_time_interval().end(),
            ]
            .into_iter()
            .any(|t| (timing_ticks(t) - time).abs() <= tolerance)
                && !out
                    .iter()
                    .any(|existing| same_note::<S>(existing, &candidate))
            {
                out.push(candidate);
            }
        }
    }
}

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

impl<S: NotesStorage> MatchingTimeEndpoint<S> {
    fn post_add_one(&mut self, storage: &mut S, note_type: S::NoteTypeEnum, note_id: u64) {
        let just_added = storage
            .query_notes_ref([(note_id, note_type)].into_iter())
            .next()
            .map(|note| note.to_snapshot());
        if let Some(just_added) = just_added {
            self.post_add_snapshot(storage, just_added);
        }
    }

    fn post_remove_one(&mut self, storage: &mut S, just_removed: &S::NoteSnapshotEnum) {
        Self::matching_notes(
            storage,
            just_removed,
            self.tolerance,
            self.condition,
            &mut self.scratch_candidates,
        );

        self.scratch_snapshot.clear();
        self.scratch_partners.clear();
        self.scratch_partner_ranges.clear();

        for candidate in &self.scratch_candidates {
            Self::matching_notes(
                storage,
                candidate,
                self.tolerance,
                self.condition,
                &mut self.scratch_matches,
            );
            self.scratch_matches
                .retain(|partner| !same_note::<S>(just_removed, partner));
            let start = self.scratch_partners.len();
            self.scratch_partners.append(&mut self.scratch_matches);
            self.scratch_partner_ranges
                .push(start..self.scratch_partners.len());
        }

        for (candidate, range) in self
            .scratch_candidates
            .iter()
            .zip(&self.scratch_partner_ranges)
        {
            let partners = &self.scratch_partners[range.clone()];
            if partners.is_empty() {
                self.scratch_snapshot.push(candidate.clone());
            }
            if let Some(assign) = self.assign_related_notes {
                storage.query_notes_mut(
                    [(candidate.id(), S::NoteTypeEnum::from(candidate))].into_iter(),
                    |state| assign(state, partners),
                );
            }
        }

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

    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 MatchingTimeEndpoint<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);
        }
    }
}

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),
    )
}

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

    use proptest::prelude::*;

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

    /// Inclusive endpoint tolerance of the test notes storages, in milliseconds.
    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 highlight: Vec<HashSet<TestNoteId>>,
            edit_hook MatchingTimeEndpoint::<TestNotesStorage>::new()
                .tolerance(ClockTime::from_duration(Duration::from_millis(10)))
                .assign_related_notes(set_chord)
                .add_related_set(add_highlight)
                .remove_notes_from_related_set(remove_highlight),
        }
    }

    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 MatchingTimeEndpoint::<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_highlight(storage: &mut TestNotesStorage, notes: &[TestNoteSnapshot]) {
        let mut ids: Vec<_> = notes.iter().map(TestNoteId::from_snapshot).collect();
        ids.sort_unstable();
        ids.dedup();
        assert_eq!(ids.len(), notes.len(), "duplicate activation snapshot");
        let set: HashSet<_> = ids.into_iter().collect();
        assert!(set.len() > 1);
        if storage
            .highlight
            .iter()
            .any(|current| set.is_subset(current))
        {
            return;
        }
        storage.highlight.retain(|current| !current.is_subset(&set));
        storage.highlight.push(set);
    }

    fn remove_highlight(storage: &mut TestNotesStorage, notes: &[TestNoteSnapshot]) {
        let ids: HashSet<_> = notes.iter().map(TestNoteId::from_snapshot).collect();
        assert_eq!(ids.len(), notes.len(), "duplicate removal snapshot");
        for set in &mut storage.highlight {
            set.retain(|id| !ids.contains(id));
        }
        storage.highlight.retain(|set| set.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 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,
            &world.note_groups,
            &world.tempo_tracks,
            &world.scroll_tracks,
        )
    }

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

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

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

    fn replace_notes(
        world: &mut World,
        storage: &mut TestNotesStorage,
        notes: Vec<(u64, TestNote)>,
    ) {
        storage.replace_notes(
            &world.note_groups,
            &world.tempo_tracks,
            &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 world = World::new(0);
        let storage = CondNotesStorage::with_notes(
            notes,
            &world.note_groups,
            &world.tempo_tracks,
            &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.highlight.is_empty(), "{:?}", storage.highlight);
        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_flag(&storage, TestNoteId::Short(2)), Some(true));
        let mut highlight: Vec<_> = storage.highlight.iter().flatten().copied().collect();
        highlight.sort();
        assert_eq!(highlight, vec![TestNoteId::Short(1), TestNoteId::Short(2)]);
        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_flag(&storage, TestNoteId::Short(2)), Some(true));

        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 = MatchingTimeEndpoint::<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));
        }
        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_flag(&storage, TestNoteId::Long(1)), Some(true));
        storage.health_check();
    }

    #[test]
    fn edits_preserve_partners_at_the_other_endpoint() {
        let (mut world, mut storage) = setup(vec![lnote(100.0, 900.0), snote(1000.0)]);
        add_note(&mut world, &mut storage, snote(100.0));
        assert_matches_naive(&storage);
        storage.remove_note(TestNoteType::Short, 2);
        assert_matches_naive(&storage);
    }

    #[test]
    fn isolated_add_assigns_empty_partners() {
        let (_, mut storage) = setup(vec![snote(100.0)]);
        let note = storage
            .query_notes([(1, TestNoteType::Short)].into_iter())
            .next()
            .unwrap();
        storage.query_notes_mut([(1, TestNoteType::Short)].into_iter(), |state| {
            if let TestNoteMut::Short(n) = state {
                *n.runtime_data.chord = true;
            }
        });
        MatchingTimeEndpoint::<TestNotesStorage>::new()
            .assign_related_notes(set_chord)
            .post_add_snapshot(&mut storage, note);
        assert_eq!(chord_flag(&storage, TestNoteId::Short(1)), Some(false));
    }

    #[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();
        storage.shorts.remove_notes([2].into_iter());
        MatchingTimeEndpoint::<TestNotesStorage>::new()
            .tolerance(ClockTime::from_milis(TOLERANCE_MS).unwrap())
            .remove_notes_from_related_set(|storage, notes| {
                assert_eq!(notes.len(), 2);
                remove_highlight(storage, notes);
            })
            .post_remove(&mut storage, ::std::slice::from_ref(&removed));
        assert!(storage.highlight.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.highlight.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.highlight.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
                    .highlight
                    .iter()
                    .any(|set| set.contains(&TestNoteId::Short(id)))
            );
        }
        assert!(
            !storage
                .highlight
                .iter()
                .any(|set| set.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.highlight.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 (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(
            &world.note_groups,
            &world.tempo_tracks,
            &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 endpoints with their note id and rounded 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()))
                .flat_map(|(&id, &i)| {
                    [
                        (
                            TestNoteId::Long(id),
                            (i.start().seconds() * 10_000.0).round(),
                        ),
                        (TestNoteId::Long(id), (i.end().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, _) in &notes {
            let expected = notes
                .iter()
                .any(|(other_id, endpoint)| other_id == id && has_partner(&notes, *id, *endpoint));
            let mut expected_matches: Vec<_> = notes
                .iter()
                .filter(|(other, time)| {
                    other != id
                        && notes.iter().any(|(own, endpoint)| {
                            own == id && (endpoint - 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"
            );
        }

        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();

        assert!(storage.highlight.iter().all(|set| set.len() > 1));
        let mut actual: Vec<_> = storage.highlight.iter().flatten().copied().collect();
        actual.sort_unstable();
        actual.dedup();
        assert_eq!(actual, expected_highlight);
    }

    #[test]
    fn endpoint_combinations_and_spanning_nonmatch() {
        for (notes, expected) in [
            (vec![lnote(0.0, 100.0), lnote(100.0, 100.0)], true),
            (vec![lnote(100.0, 100.0), lnote(0.0, 100.0)], true),
            (vec![lnote(0.0, 200.0), lnote(100.0, 100.0)], true),
            (vec![lnote(0.0, 200.0), lnote(0.0, 200.0)], true),
            (vec![lnote(0.0, 200.0), lnote(50.0, 100.0)], false),
            (vec![lnote(0.0, 0.0), lnote(0.0, 0.0)], true),
        ] {
            let (_, storage) = setup(notes);
            for id in [1, 2] {
                assert_eq!(chord_flag(&storage, TestNoteId::Long(id)), Some(expected));
            }
            assert_matches_naive(&storage);
        }
    }

    #[test]
    fn removal_preserves_partner_at_far_endpoint() {
        let (_, mut storage) = setup(vec![snote(0.0), lnote(0.0, 1000.0), snote(1000.0)]);
        storage.remove_note(TestNoteType::Short, 1);
        assert_eq!(chord_flag(&storage, TestNoteId::Long(1)), Some(true));
        assert_matches_naive(&storage);
        storage.remove_note(TestNoteType::Short, 2);
        assert_eq!(chord_flag(&storage, TestNoteId::Long(1)), Some(false));
        assert_matches_naive(&storage);
    }

    #[test]
    fn endpoint_rounding_and_end_only_update() {
        let (mut world, mut storage) = setup(vec![lnote(0.0, 100.049), snote(110.049)]);
        assert_eq!(chord_flag(&storage, TestNoteId::Long(1)), Some(true));
        replace_note(&mut world, &mut storage, 1, lnote(0.0, 99.949));
        assert_eq!(chord_flag(&storage, TestNoteId::Long(1)), Some(false));
        assert_matches_naive(&storage);
        replace_note(&mut world, &mut storage, 1, lnote(0.0, 110.049));
        assert_matches_naive(&storage);
        replace_note(&mut world, &mut storage, 1, lnote(1000.0, 100.0));
        assert_eq!(chord_flag(&storage, TestNoteId::Long(1)), Some(false));
        assert_matches_naive(&storage);
    }

    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.highlight.clear();
    }

    #[test]
    fn rebuild_restores_cleared_relationships() {
        let (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(
                &world.note_groups,
                &world.tempo_tracks,
                &world.scroll_tracks,
            );
            assert_matches_naive(&storage);
            storage.health_check();
        }
        let empty_world = World::new(0);
        let mut empty = TestNotesStorage::new();
        empty.rebuild_clock_time(
            &empty_world.note_groups,
            &empty_world.tempo_tracks,
            &empty_world.scroll_tracks,
        );
        assert_matches_naive(&empty);
    }

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

    proptest! {
        #[test]
        fn rebuild_matches_naive(notes in prop::collection::vec(arb_note(), 0..40)) {
            let (world, mut storage) = setup(notes);
            clear_relationships(&mut storage);
            storage.rebuild_clock_time(
                &world.note_groups,
                &world.tempo_tracks,
                &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();
            }
        }
    }
}