vsrg 0.6.0

use crate::notes::relations::{AssignParterToNote, AssignPartnersToNote};
use crate::notes::{NoteState, NotesStorage, NotesStorageEditHook};

// For future users:
// - It's not immediately clear whether storing parent-child relationship in an auxiliary
//   field in notes storage is useful. Therefore we skip it for now to simplify the API.
// - These relationship cannot represent relationship between notes of the same type,
//   or circular relationship, since those would cause an infinite loop when removing the
//   notes. A separate dedicated relationship is necessary.
//
// If you need these features, then let us know by opening an issue.

/// Reads the parent of a note.
///
/// # Requirements
/// Returns [`Some`] with the parent's id if the note declares itself a child
/// of that note, or [`None`] if it has no parent.
pub type ParentIdOf<S> = fn(&<S as NotesStorage>::NoteEnum) -> Option<u64>;

/// Remove any parent linking from a child note.
///
/// # 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.
pub type DetachNote<S> = fn(<S as NotesStorage>::NoteMutEnum<'_>);

/// What happens to a child note whose parent is removed.
#[derive(Debug, Clone, Copy)]
pub enum OrphanPolicy<S: NotesStorage> {
    /// Detach the child notes from its parent, without removing the note.
    ///
    /// # Remarks
    /// You must provide a callback to detach the note.
    ///
    /// # 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.
    Detach(DetachNote<S>),

    /// Remove the orphaned children from the notes storage.
    Remove,
}

/// Parent-child relationship.
///
/// # Remarks
/// This assumes each child note stores the id of its parent note (and not the reverse where parent
/// stores ids of its children notes), and wants to cache data from its parent notes into its
/// runtime state. For example: tap notes belonging to a parent curve caching the parent curve's position
/// at its time.
pub struct ParentChild<S: NotesStorage> {
    parent_note_type: S::NoteTypeEnum,
    child_note_type: S::NoteTypeEnum,

    parent_id_of: ParentIdOf<S>,
    link_parent_to_child: Option<AssignParterToNote<S>>,
    link_children_to_parent: Option<AssignPartnersToNote<S>>,
    orphan_policy: OrphanPolicy<S>,

    scratch_children: Vec<S::NoteSnapshotEnum>,
}

impl<S: NotesStorage> ParentChild<S> {
    /// Creates a relationship with no linking callbacks.
    ///
    /// # Remarks
    /// Continue configuring the relationship by calling its builder methods.
    ///
    /// Removing a parent removes its children by default; use [Self::orphan_policy] to configure
    /// detachment instead.
    pub const fn new(
        parent_note_type: S::NoteTypeEnum,
        child_note_type: S::NoteTypeEnum,
        parent_id_of: ParentIdOf<S>,
    ) -> Self {
        Self {
            parent_note_type,
            child_note_type,
            parent_id_of,
            link_parent_to_child: None,
            link_children_to_parent: None,
            orphan_policy: OrphanPolicy::Remove,
            scratch_children: Vec::new(),
        }
    }

    /// Assign the parent to the child.
    ///
    /// # Remarks
    /// The note can copy any data from the partner note as needed. For example: copying the
    /// parent curve's position to the child note.
    ///
    /// # 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 parent rather than appending the parent to the note.
    pub const fn link_parent_to_child(mut self, callback: AssignParterToNote<S>) -> Self {
        self.link_parent_to_child = Some(callback);
        self
    }

    /// Assign the children to the parent.
    ///
    /// # Remarks
    /// The note can copy any data from the partner notes as needed.
    /// The slice contains the all children notes whose stored parent id is the parent note.
    /// Empty slice means it no longer has any children.
    ///
    /// # 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 children rather than appending the children to the note.
    pub const fn link_children_to_parent(mut self, callback: AssignPartnersToNote<S>) -> Self {
        self.link_children_to_parent = Some(callback);
        self
    }

    /// Sets whether removing a parent removes or detaches its children.
    pub const fn orphan_policy(mut self, value: OrphanPolicy<S>) -> Self {
        self.orphan_policy = value;
        self
    }
}
impl<S: NotesStorage> ParentChild<S> {
    fn post_add_one(&mut self, storage: &mut S, note_type: S::NoteTypeEnum, note_id: u64) {
        if note_type != self.child_note_type {
            return;
        }

        let Some(just_added) = storage
            .query_notes_ref([(note_id, note_type)].into_iter())
            .map(|x| S::NoteEnum::from(&x))
            .next()
        else {
            return;
        };

        let Some(parent_id) = (self.parent_id_of)(&just_added) else {
            return;
        };
        drop(just_added);

        let Some(parent) = storage
            .query_notes([(parent_id, self.parent_note_type)].into_iter())
            .next()
        else {
            return;
        };

        if let Some(link) = self.link_parent_to_child {
            storage.query_notes_mut([(note_id, note_type)].into_iter(), |note| {
                link(note, parent.clone());
            })
        }

        if let Some(link) = self.link_children_to_parent {
            self.scratch_children.clear();
            self.scratch_children.extend(
                storage
                    .all_notes_ref_of_type(self.child_note_type)
                    .filter_map(|n| {
                        (self.parent_id_of)(&S::NoteEnum::from(&n))
                            .is_some_and(|i| parent_id == i)
                            .then_some(n.to_snapshot())
                    }),
            );
            storage.query_notes_mut([(parent_id, self.parent_note_type)].into_iter(), |parent| {
                link(parent, &self.scratch_children)
            });
        }
    }

    fn pre_remove_one(
        &mut self,
        storage: &mut S,
        note_type: <S as NotesStorage>::NoteTypeEnum,
        note_id: u64,
        removal_batch: &[S::NoteSnapshotEnum],
    ) {
        if note_type != self.child_note_type {
            return;
        }

        let Some(link) = self.link_children_to_parent else {
            return;
        };

        let Some(removing) = storage
            .query_notes_ref([(note_id, note_type)].into_iter())
            .map(|x| S::NoteEnum::from(&x))
            .next()
        else {
            return;
        };

        let Some(parent_id) = (self.parent_id_of)(&removing) else {
            return;
        };
        drop(removing);

        self.scratch_children.clear();
        self.scratch_children.extend(
            storage
                .all_notes_ref_of_type(self.child_note_type)
                .filter_map(|n| {
                    let matching_parent =
                        (self.parent_id_of)(&S::NoteEnum::from(&n)).is_some_and(|i| parent_id == i);
                    let is_removing = removal_batch.iter().any(|removed| {
                        removed.id() == n.id()
                            && S::NoteTypeEnum::from(removed) == self.child_note_type
                    });
                    (matching_parent && !is_removing).then(|| n.to_snapshot())
                }),
        );
        storage.query_notes_mut([(parent_id, self.parent_note_type)].into_iter(), |parent| {
            link(parent, &self.scratch_children)
        });
    }

    fn post_remove_one(&mut self, storage: &mut S, just_removed: &S::NoteSnapshotEnum) {
        if S::NoteTypeEnum::from(just_removed) != self.parent_note_type {
            return;
        }

        match self.orphan_policy {
            OrphanPolicy::Remove => {
                // remove_notes consumes its input, so this buffer cannot be reused.
                let children = storage
                    .all_notes_ref_of_type(self.child_note_type)
                    .filter_map(|n| {
                        let id = n.id();
                        (self.parent_id_of)(&S::NoteEnum::from(&n))
                            .is_some_and(|id| id == just_removed.id())
                            .then_some((id, self.child_note_type))
                    })
                    .collect::<Vec<_>>();
                storage.remove_notes(children);
            }

            OrphanPolicy::Detach(detach) => {
                storage.all_notes_mut_of_type(self.child_note_type, |candidate| {
                    if (self.parent_id_of)(&(&candidate).into())
                        .is_some_and(|id| id == just_removed.id())
                    {
                        detach(candidate);
                    }
                })
            }
        }
    }

    fn pre_update_one(
        &mut self,
        storage: &mut S,
        replacing_id: u64,
        replacing_note: &<S as NotesStorage>::NoteEnum,
        replacements: &[(u64, S::NoteEnum)],
    ) {
        // If a note has changed its parent id, then we need to notify the previous parent.
        let Some(link) = self.link_children_to_parent else {
            return;
        };

        let note_type = S::NoteTypeEnum::from(replacing_note);
        if note_type != self.child_note_type {
            return;
        }

        let Some(note_ref) = storage
            .query_notes_ref([(replacing_id, note_type)].into_iter())
            .map(|x| S::NoteEnum::from(&x))
            .next()
        else {
            return;
        };

        let Some(current_parent_id) = (self.parent_id_of)(&note_ref) else {
            return;
        };
        drop(note_ref);

        let new_parent_id = (self.parent_id_of)(replacing_note);
        if new_parent_id == Some(current_parent_id) {
            return;
        }

        self.scratch_children.clear();
        self.scratch_children.extend(
            storage
                .all_notes_ref_of_type(self.child_note_type)
                .filter_map(|n| {
                    let matching_parent = (self.parent_id_of)(&S::NoteEnum::from(&n))
                        .is_some_and(|i| current_parent_id == i);
                    let will_change_parent = replacements.iter().any(|(id, replacement)| {
                        *id == n.id()
                            && S::NoteTypeEnum::from(replacement) == self.child_note_type
                            && (self.parent_id_of)(replacement) != Some(current_parent_id)
                    });
                    (matching_parent && !will_change_parent).then(|| n.to_snapshot())
                }),
        );
        storage.query_notes_mut(
            [(current_parent_id, self.parent_note_type)].into_iter(),
            |parent| link(parent, &self.scratch_children),
        );
    }

    fn post_update_one(
        &mut self,
        storage: &mut S,
        _before: &S::NoteSnapshotEnum,
        after: &S::NoteSnapshotEnum,
    ) {
        let note_type = S::NoteTypeEnum::from(after);
        if note_type == self.child_note_type {
            // Unconditionally link child to the new parent, which is the same as adding the note.
            // Unlinking is already handled in pre_update
            self.post_add_one(storage, note_type, after.id());
        } else if note_type == self.parent_note_type {
            // Scan all notes of children note type, for any notes whose parent id matches the
            // updated note. Then link them
            let Some(link) = self.link_parent_to_child else {
                return;
            };
            let Some(parent) = storage
                .query_notes([(after.id(), note_type)].into_iter())
                .next()
            else {
                return;
            };

            let children = storage
                .all_notes_ref_of_type(self.child_note_type)
                .filter_map(|child| {
                    let child_note = S::NoteEnum::from(&child);
                    ((self.parent_id_of)(&child_note) == Some(after.id())).then_some(child.id())
                })
                .collect::<Vec<_>>();

            for child_id in children {
                storage.query_notes_mut([(child_id, self.child_note_type)].into_iter(), |child| {
                    link(child, parent.clone())
                });
            }
        }
    }
}

impl<S: NotesStorage> NotesStorageEditHook<S> for ParentChild<S> {
    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 pre_remove(&mut self, storage: &mut S, notes: &[S::NoteSnapshotEnum]) {
        for note in notes {
            self.pre_remove_one(storage, S::NoteTypeEnum::from(note), note.id(), notes);
        }
    }
    fn post_remove(&mut self, storage: &mut S, notes: &[S::NoteSnapshotEnum]) {
        for note in notes {
            self.post_remove_one(storage, note);
        }
    }
    fn pre_update(&mut self, storage: &mut S, notes: &mut [(u64, S::NoteEnum)]) {
        for (id, note) in notes.iter() {
            self.pre_update_one(storage, *id, note, notes);
        }
    }
    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 super::*;
    use crate::notes::relations::test_common::{RelLong, RelShort, World, long_at, short_at};
    use crate::notes::{LongNoteTypeStorage, ShortNoteTypeStorage};

    crate::generate_notes_storage! {
        enum Kind,
        enum Note,
        enum Snapshot,
        enum NoteRef,
        enum NoteMut,
        struct Storage {
            note_type Child(RelShort) => children: ShortNoteTypeStorage<RelShort>,
            note_type Parent(RelLong) => parents: LongNoteTypeStorage<RelLong>,
            edit_hook ParentChild::<Storage>::new(Kind::Parent, Kind::Child, |note| match note {
                    Note::Child(child) => Some(child.pos as u64),
                    _ => None,
                }).link_parent_to_child(|child, parent| {
                    if let (NoteMut::Child(child), Snapshot::Parent(parent)) = (child, parent) {
                        *child.runtime_data.parent_id = Some(parent.id());
                        *child.runtime_data.matching_ids = vec![(parent.note.pos as u64, false)];
                    }
                }).link_children_to_parent(|parent, children| {
                    if let NoteMut::Parent(parent) = parent {
                        *parent.runtime_data.matching_ids = children.iter().map(|n| (n.id(), false)).collect();
                    }
                }),
        }
    }

    fn fixture(world: &mut World) -> Storage {
        Storage::with_notes(
            vec![
                Note::Parent(long_at(0.0, 10.0, 0)),
                Note::Parent(long_at(20.0, 10.0, 0)),
                Note::Child(short_at(1.0, 1)),
                Note::Child(short_at(2.0, 1)),
            ],
            &world.note_groups,
            &world.tempo_tracks,
            &world.scroll_tracks,
        )
    }

    fn children(storage: &Storage, id: u64) -> Vec<u64> {
        match storage
            .query_notes_ref([(id, Kind::Parent)].into_iter())
            .next()
            .unwrap()
        {
            NoteRef::Parent(parent) => parent
                .runtime_data
                .matching_ids
                .iter()
                .map(|(id, _)| *id)
                .collect(),
            _ => unreachable!(),
        }
    }

    #[test]
    fn removing_children_together_excludes_the_entire_batch() {
        let mut world = World::new(0);
        let mut storage = fixture(&mut world);
        assert_eq!(children(&storage, 1), vec![1, 2]);
        storage.remove_notes(vec![(1, Kind::Child), (2, Kind::Child)]);
        assert!(children(&storage, 1).is_empty());
    }

    #[test]
    fn moving_children_together_refreshes_both_parents() {
        let mut world = World::new(0);
        let mut storage = fixture(&mut world);
        storage.replace_notes(
            &world.note_groups,
            &world.tempo_tracks,
            &world.scroll_tracks,
            vec![
                (1, Note::Child(short_at(1.0, 2))),
                (2, Note::Child(short_at(2.0, 2))),
            ],
        );
        assert!(children(&storage, 1).is_empty());
        assert_eq!(children(&storage, 2), vec![1, 2]);
    }

    #[test]
    fn removing_parent_and_child_together_cascades_remaining_children() {
        let mut world = World::new(0);
        let mut storage = fixture(&mut world);
        storage.remove_notes(vec![(1, Kind::Parent), (1, Kind::Child)]);
        assert_eq!(storage.all_notes_of_type(Kind::Child).count(), 0);
        assert_eq!(storage.all_notes_of_type(Kind::Parent).count(), 1);
    }

    fn add(world: &mut World, storage: &mut Storage, note: Note) {
        storage.add_note(
            &world.note_groups,
            &world.tempo_tracks,
            &world.scroll_tracks,
            note,
        );
    }

    fn child_parent(storage: &Storage, id: u64) -> (Option<u64>, Vec<(u64, bool)>) {
        match storage
            .query_notes_ref([(id, Kind::Child)].into_iter())
            .next()
            .unwrap()
        {
            NoteRef::Child(child) => (
                *child.runtime_data.parent_id,
                child.runtime_data.matching_ids.clone(),
            ),
            _ => unreachable!(),
        }
    }

    #[test]
    fn adding_child_copies_the_selected_parents_position() {
        let mut world = World::new(0);
        let mut storage = Storage::new();
        add(
            &mut world,
            &mut storage,
            Note::Parent(long_at(0.0, 10.0, 17)),
        );
        add(
            &mut world,
            &mut storage,
            Note::Parent(long_at(20.0, 10.0, 42)),
        );
        add(&mut world, &mut storage, Note::Child(short_at(1.0, 2)));
        assert_eq!(child_parent(&storage, 1), (Some(2), vec![(42, false)]));
        assert!(children(&storage, 1).is_empty());
        assert_eq!(children(&storage, 2), vec![1]);
    }

    #[test]
    fn replacing_child_parent_refreshes_its_position_and_both_parents() {
        let mut world = World::new(0);
        let mut storage = Storage::new();
        add(
            &mut world,
            &mut storage,
            Note::Parent(long_at(0.0, 10.0, 17)),
        );
        add(
            &mut world,
            &mut storage,
            Note::Parent(long_at(20.0, 10.0, 42)),
        );
        add(&mut world, &mut storage, Note::Child(short_at(1.0, 1)));
        assert_eq!(child_parent(&storage, 1), (Some(1), vec![(17, false)]));
        storage.replace_note(
            &world.note_groups,
            &world.tempo_tracks,
            &world.scroll_tracks,
            1,
            Note::Child(short_at(1.0, 2)),
        );
        assert_eq!(child_parent(&storage, 1), (Some(2), vec![(42, false)]));
        assert!(children(&storage, 1).is_empty());
        assert_eq!(children(&storage, 2), vec![1]);
    }

    #[test]
    fn replacing_parent_children_refreshes() {
        let mut world = World::new(0);
        let mut storage = Storage::new();
        add(
            &mut world,
            &mut storage,
            Note::Parent(long_at(0.0, 10.0, 17)),
        );
        add(&mut world, &mut storage, Note::Child(short_at(1.0, 1)));
        assert_eq!(child_parent(&storage, 1), (Some(1), vec![(17, false)]));
        storage.replace_note(
            &world.note_groups,
            &world.tempo_tracks,
            &world.scroll_tracks,
            1,
            Note::Parent(long_at(0.0, 10.0, 42)),
        );
        assert_eq!(child_parent(&storage, 1), (Some(1), vec![(42, false)]));
        assert_eq!(children(&storage, 1), vec![1]);
    }

    #[test]
    fn removing_parent_cascades_only_its_children() {
        let mut world = World::new(0);
        let mut storage = fixture(&mut world);
        add(&mut world, &mut storage, Note::Child(short_at(3.0, 2)));
        storage.remove_note(Kind::Parent, 1);
        assert!(
            storage
                .query_notes([(1, Kind::Child), (2, Kind::Child)].into_iter())
                .next()
                .is_none()
        );
        assert_eq!(child_parent(&storage, 3).0, Some(2));
        assert_eq!(children(&storage, 2), vec![3]);
    }

    #[test]
    fn removing_one_child_preserves_the_other_link() {
        let mut world = World::new(0);
        let mut storage = fixture(&mut world);
        storage.remove_note(Kind::Child, 1);
        assert!(
            storage
                .query_notes([(1, Kind::Child)].into_iter())
                .next()
                .is_none()
        );
        assert_eq!(children(&storage, 1), vec![2]);
        assert_eq!(child_parent(&storage, 2).0, Some(1));
        storage.remove_note(Kind::Child, 2);
        assert!(children(&storage, 1).is_empty());
    }

    mod detach {
        use super::*;
        crate::generate_notes_storage! {
            enum Kind,
            enum Note,
            enum Snapshot,
            enum NoteRef,
            enum NoteMut,
            struct Storage {
                note_type Child(RelShort) => children: ShortNoteTypeStorage<RelShort>,
                note_type Parent(RelLong) => parents: LongNoteTypeStorage<RelLong>,
                edit_hook ParentChild::<Storage>::new(Kind::Parent, Kind::Child, |note| match note {
                        Note::Child(child) if child.pos != 0 => Some(child.pos as u64),
                        _ => None,
                    }).link_parent_to_child(|child, parent| {
                        if let NoteMut::Child(child) = child {
                            *child.runtime_data.parent_id = Some(parent.id());
                        }
                    }).orphan_policy(OrphanPolicy::Detach(|child| {
                        if let NoteMut::Child(child) = child {
                            // Detachment clears runtime linkage; authored parent IDs are immutable here.
                            *child.runtime_data.parent_id = None;
                        }
                    })),
            }
        }

        #[test]
        fn removing_parent_detaches_children_and_preserves_unrelated_links() {
            let world = World::new(0);
            let mut storage = Storage::with_notes(
                vec![
                    Note::Parent(long_at(0.0, 10.0, 17)),
                    Note::Parent(long_at(20.0, 10.0, 42)),
                    Note::Child(short_at(1.0, 1)),
                    Note::Child(short_at(2.0, 2)),
                ],
                &world.note_groups,
                &world.tempo_tracks,
                &world.scroll_tracks,
            );
            storage.remove_note(Kind::Parent, 1);
            assert_eq!(storage.all_notes_of_type(Kind::Child).count(), 2);
            for (id, expected) in [(1, None), (2, Some(2))] {
                match storage
                    .query_notes_ref([(id, Kind::Child)].into_iter())
                    .next()
                    .unwrap()
                {
                    NoteRef::Child(child) => {
                        assert_eq!(*child.runtime_data.parent_id, expected);
                        assert_eq!(*child.value_data.0 as u64, id);
                    }
                    _ => unreachable!(),
                }
            }
        }
    }

    fn assert_matches_scan(storage: &Storage) {
        for snapshot in storage.all_notes_of_type(Kind::Parent) {
            let parent = snapshot.id();
            let position = match snapshot {
                Snapshot::Parent(n) => n.note.pos,
                _ => unreachable!(),
            };
            let mut expected: Vec<_> = storage
                .all_notes_of_type(Kind::Child)
                .filter_map(|n| match n {
                    Snapshot::Child(n) if n.note.pos as u64 == parent => Some(n.id()),
                    _ => None,
                })
                .collect();
            expected.sort_unstable();
            let mut actual = children(storage, parent);
            actual.sort_unstable();
            assert_eq!(actual, expected.clone());
            for id in expected {
                assert_eq!(
                    child_parent(storage, id),
                    (Some(parent), vec![(position as u64, false)])
                );
            }
        }
    }

    use proptest::prelude::*;

    proptest! {
        #[test]
        fn random_child_edits_match_manual_scan(
            ops in prop::collection::vec((0u8..6, 0usize..20, 1u32..=2), 1..40)
        ) {
            let mut world = World::new(0);
            let mut storage = fixture(&mut world);
            for (op, index, parent_index) in ops {
                let parents: Vec<_> = storage.all_notes_of_type(Kind::Parent).map(|n| n.id()).collect();
                if op == 5 || parents.is_empty() {
                    add(&mut world, &mut storage, Note::Parent(long_at(index as f64, 10.0, index as u32)));
                } else if op == 4 {
                    let parent = parents[index % parents.len()];
                    let removed_children = children(&storage, parent);
                    storage.remove_note(Kind::Parent, parent);
                    for id in removed_children {
                        prop_assert!(storage.query_notes([(id, Kind::Child)].into_iter()).next().is_none());
                    }
                } else {
                let parent = parents[parent_index as usize % parents.len()] as u32;
                let ids: Vec<_> = storage.all_notes_of_type(Kind::Child).map(|n| n.id()).collect();
                if op == 0 || ids.is_empty() {
                    add(&mut world, &mut storage, Note::Child(short_at(index as f64, parent)));
                } else {
                    let id = ids[index % ids.len()];
                    if op == 1 {
                        storage.remove_note(Kind::Child, id);
                    } else if op == 2 {
                        storage.replace_note(
                            &world.note_groups,
                            &world.tempo_tracks,
                            &world.scroll_tracks,
                            id,
                            Note::Child(short_at(index as f64, parent))
                        );
                    } else {
                        storage.remove_notes(ids.into_iter().take(2).map(|id| (id, Kind::Child)).collect());
                    }
                }
                }
                assert_matches_scan(&storage);
                storage.children.health_check();
                storage.parents.health_check();
            }
        }
    }
}