journalmap/

journalmap.rs

use crate::JournalError;
use crate::Tag;
use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt::Debug;
use std::hash::Hash;
use std::ops::Range;

#[cfg(feature = "serde")]
use serde::Deserialize;
#[cfg(feature = "serde")]
use serde::Serialize;
#[cfg(feature = "serde")]
use serde::de::Deserialize as DeserializeTrait;
#[cfg(feature = "serde")]
use serde::de::Deserializer;
#[cfg(feature = "serde")]
use serde::ser::Serialize as SerializeTrait;
#[cfg(feature = "serde")]
use serde::ser::Serializer;

#[derive(Clone, Copy)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
struct Insert<K, V> {
    pub key: K,
    pub val: V,
    pub replacing: Option<V>,
}

impl<K: Eq + Hash + Clone, V: Clone> Insert<K, V> {
    pub fn forward(&self, map: &mut HashMap<K, V>) {
        map.insert(self.key.clone(), self.val.clone());
    }

    pub fn reverse(&self, map: &mut HashMap<K, V>) {
        match self.replacing.clone() {
            Some(val) => {
                map.insert(self.key.clone(), val);
            }
            None => {
                map.remove(&self.key);
            }
        }
    }
}

impl<K: Debug, V: Debug> std::fmt::Debug for Insert<K, V> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match &self.replacing {
            Some(replacing) => write!(f, "[>] {:?}:{:?}->{:?}", self.key, replacing, self.val),
            None => write!(f, "[+] {:?}:{:?}", self.key, self.val),
        }
    }
}

#[derive(Clone, Copy)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
struct Remove<K, V> {
    pub key: K,
    pub val: V,
}

impl<K: Eq + Hash + Clone, V: Clone> Remove<K, V> {
    pub fn forward(&self, map: &mut HashMap<K, V>) {
        map.remove(&self.key);
    }

    pub fn reverse(&self, map: &mut HashMap<K, V>) {
        map.insert(self.key.clone(), self.val.clone());
    }
}

impl<K: Debug, V: Debug> std::fmt::Debug for Remove<K, V> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "[-] '{:?}':'{:?}'", self.key, self.val)
    }
}

#[derive(Clone, Copy)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
enum JournalEntry<K, V, T> {
    Insert(Insert<K, V>),
    Remove(Remove<K, V>),
    Tag(Tag<T>),
}

impl<K: Clone + Hash + Eq, V: Clone, T: Hash + Eq + Clone> JournalEntry<K, V, T> {
    pub fn reverse(&self, map: &mut HashMap<K, V>) {
        match self {
            Self::Insert(entry_data) => entry_data.reverse(map),
            Self::Remove(entry_data) => entry_data.reverse(map),
            Self::Tag(_) => {}
        }
    }

    pub fn forward(&self, map: &mut HashMap<K, V>) {
        match self {
            Self::Insert(entry_data) => entry_data.forward(map),
            Self::Remove(entry_data) => entry_data.forward(map),
            Self::Tag(_) => {}
        }
    }
}

impl<K: Debug + Hash + Eq, V: Debug, T: Debug + Hash + Eq> std::fmt::Debug
    for JournalEntry<K, V, T>
{
    fn fmt(&self, f: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
        match self {
            JournalEntry::Insert(entry) => {
                write!(f, "{:?}", entry)
            }
            JournalEntry::Remove(entry) => {
                write!(f, "{:?}", entry)
            }
            JournalEntry::Tag(entry) => {
                write!(f, "{:?}", entry)
            }
        }
    }
}

/// Journaling HashMap. Each transaction is recorded to the journal and can be
/// undone and redone.
///
/// **NOTE:** Current implementation makes copies of all data, keys & tags.
/// Memory usage is at least doubled, and if you're using data types which
/// are large there can be a decent performance hit on write.
///
/// ## Usage
///
/// JournalMaps are designed to provide a simple interface to implement undo,
/// redo and version control. This is achieved via recording all transactions
/// in a journal and by adding "tags", which are basically just generic markers
/// that can be used as undo checkpoints or version identifiers.
///
/// ```rust
/// use journalmap::JournalMap;
///
/// fn main() {
///     let mut ages: JournalMap<&'static str, i32, _> = JournalMap::new();
///
///     // Always add a tag at the start
///     ages.append_tag("start").unwrap();
///
///     ages.insert("Sarah", 22);
///     ages.insert("John", 21);
///     ages.insert("Mariah", 45);
///
///     ages.append_tag("middle").unwrap();
///
///     ages.insert("Steve", 34);
///     ages.insert("Carol", 41);
///     ages.insert("Ellie", 21);
///     ages.insert("Mariah", 87);
///
///     ages.append_tag("end").unwrap();
///     
///     // Currently all the entries we have are present...
///     assert_eq!(ages.get(&"Sarah"), Some(&22));
///     assert_eq!(ages.get(&"John"), Some(&21));
///     assert_eq!(ages.get(&"Mariah"), Some(&87));
///     assert_eq!(ages.get(&"Steve"), Some(&34));
///     assert_eq!(ages.get(&"Carol"), Some(&41));
///     assert_eq!(ages.get(&"Ellie"), Some(&21));
///
///     // Now let's reverse it and watch our inserts be undone...
///     ages.reverse_to_tag(&"middle");
///
///     assert_eq!(ages.get(&"Sarah"), Some(&22));
///     assert_eq!(ages.get(&"John"), Some(&21));
///     assert_eq!(ages.get(&"Mariah"), Some(&45));
///     assert_eq!(ages.contains_key(&"Steve"), false);
///     assert_eq!(ages.contains_key(&"Carol"), false);
///     assert_eq!(ages.contains_key(&"Ellie"), false);
///
///     // Inserting a new entry truncates the future so we can no longer
///     // re-do the changes we've undone
///     ages.insert("Omar", 33);
///
///     ages.append_tag("end").unwrap();
///
///
///     assert_eq!(ages.get(&"Sarah"), Some(&22));
///     assert_eq!(ages.get(&"John"), Some(&21));
///     assert_eq!(ages.get(&"Mariah"), Some(&45));
///     assert_eq!(ages.get(&"Omar"), Some(&33));
///     assert_eq!(ages.contains_key(&"Steve"), false);
///     assert_eq!(ages.contains_key(&"Carol"), false);
///     assert_eq!(ages.contains_key(&"Ellie"), false);
/// }
/// ```
pub struct JournalMap<K: Eq + Hash, V, T: Eq + Hash> {
    journal: Vec<JournalEntry<K, V, T>>,
    // It may be better to store the indices with the tags, but I will need to
    // do benchmarking to figure this out. Storing tag indices may also over-
    // complicate things too.
    tags: HashSet<T>,
    map: HashMap<K, V>,
    current_tag_index: Option<usize>,
}

impl<K: Hash + Eq, V, T: Hash + Eq> Default for JournalMap<K, V, T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<K: Hash + Eq, V, T: Hash + Eq> JournalMap<K, V, T> {
    pub fn new() -> Self {
        Self {
            journal: Vec::new(),
            tags: HashSet::new(),
            map: HashMap::new(),
            current_tag_index: None,
        }
    }

    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            journal: Vec::with_capacity(capacity),
            tags: HashSet::with_capacity(capacity),
            map: HashMap::with_capacity(capacity),
            current_tag_index: None,
        }
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.map.len()
    }

    #[inline]
    pub fn journal_len(&self) -> usize {
        self.journal.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.map.is_empty()
    }

    pub fn current_tag(&self) -> Option<&T> {
        let Some(i_tag) = self.current_tag_index else {
            return None;
        };

        let JournalEntry::Tag(tag) = &self.journal[i_tag] else {
            panic!("Tag not at tag index");
        };

        Some(&tag.tag)
    }

    #[inline]
    pub fn as_hashmap(&self) -> &HashMap<K, V> {
        &self.map
    }

    /// Clears the map and the journal.
    pub fn clear_all(&mut self) {
        self.map.clear();
        self.journal.clear();
        self.tags.clear();
        self.current_tag_index = None;
    }

    /// Used to prepare for another journal entry. Truncates the journal past
    /// the current tag, and clears the current tag index.
    fn transaction_prep(&mut self) {
        let Some(i_tag) = self.current_tag_index.take() else {
            return;
        };

        let i_trunc = i_tag + 1;

        for entry in self.journal[i_trunc..].iter() {
            if let JournalEntry::Tag(tag) = entry {
                self.tags.remove(&tag.tag);
            }
        }

        self.journal.truncate(i_trunc);
    }

    /// Searches the journal in reverse from the current tag to find a tag
    /// that matches the match function, returning the range to apply the
    /// reverse function to. Also sets the current tag to the tag that matches
    /// the match function.
    fn i_tag_reverse_seek<MatchFn: Fn(&T) -> bool>(
        &mut self,
        match_fn: MatchFn,
    ) -> Option<Range<usize>> {
        let i_reverse_end = self.current_tag_index.unwrap_or(self.journal.len());

        let Some(i_tag_rev_rel) = self.journal[..i_reverse_end]
            .iter()
            .rev()
            .position(|entry| {
                if let JournalEntry::Tag(tag) = entry
                    && match_fn(&tag.tag)
                {
                    true
                } else {
                    false
                }
            })
        else {
            return None;
        };

        let i_tag = i_reverse_end - i_tag_rev_rel - 1;

        self.current_tag_index = Some(i_tag);

        let i_reverse_start = i_tag + 1;

        Some(i_reverse_start..i_reverse_end)
    }

    /// Searches the journal in from the current tag to find a tag
    /// that matches the match function, returning the range to apply the
    /// reverse function to. Also sets the current tag to the tag that matches
    /// the match function.
    fn i_tag_forward_seek<MatchFn: Fn(&T) -> bool>(
        &mut self,
        match_fn: MatchFn,
    ) -> Option<Range<usize>> {
        let Some(i_current_tag) = self.current_tag_index else {
            return None;
        };

        let i_forward_start = i_current_tag + 1;

        let Some(i_next_tag_rel) = self.journal[i_forward_start..].iter().position(|entry| {
            if let JournalEntry::Tag(tag) = entry
                && match_fn(&tag.tag)
            {
                true
            } else {
                false
            }
        }) else {
            return None;
        };

        let i_next_tag = i_forward_start + i_next_tag_rel;

        self.current_tag_index = Some(i_next_tag);

        let i_forward_end = i_next_tag;

        Some(i_forward_start..i_forward_end)
    }

    /// Retain tags that match the criteria in the retain closure.
    pub fn retain_tags<RetainFn: Fn(&T) -> bool>(&mut self, retain_fn: RetainFn) {
        self.current_tag_index = None;

        self.journal.retain(|entry| {
            let JournalEntry::Tag(tag) = entry else {
                return true;
            };

            let retain = retain_fn(&tag.tag);

            if !retain {
                self.tags.remove(&tag.tag);
            }

            retain
        });
    }
}

impl<K: Eq + Hash + Clone, V: Clone, T: Hash + Eq + Clone> JournalMap<K, V, T> {
    /// Create a JournalMap set from an iterator. You can optionally place a tag
    /// at the start of the journal with the `init_tag` argument.
    pub fn from_iter<Iter: Iterator<Item = (K, V)>>(iter: Iter, init_tag: Option<T>) -> Self {
        let map: HashMap<K, V> = HashMap::from_iter(iter);

        let capacity = map.capacity() + init_tag.is_some() as usize;

        let mut journal_map = Self {
            journal: Vec::with_capacity(capacity),
            tags: HashSet::with_capacity((capacity / 2).min(4)),
            map,
            current_tag_index: None,
        };

        if let Some(tag) = init_tag {
            journal_map.tags.insert(tag.clone());

            journal_map.journal.push(JournalEntry::Tag(Tag { tag }));
        }

        for (key, val) in journal_map.map.iter() {
            journal_map.journal.push(JournalEntry::Insert(Insert {
                key: key.clone(),
                val: val.clone(),
                replacing: None,
            }));
        }

        journal_map
    }

    /// **NOTE:** Inserting makes up to 3 clones; a clone of the key, a clone
    /// of the value, and a clone of the value being replaced if the key already
    /// exists and is assigned something.
    pub fn insert(&mut self, key: K, val: V) -> Option<V> {
        self.transaction_prep();

        let replacing = self.map.insert(key.clone(), val.clone());

        let entry = JournalEntry::Insert(Insert {
            key,
            val,
            replacing: replacing.clone(),
        });

        self.journal.push(entry);

        replacing
    }

    pub fn get(&self, key: &K) -> Option<&V> {
        self.map.get(key)
    }

    pub fn contains_key(&self, key: &K) -> bool {
        self.map.contains_key(key)
    }

    pub fn remove(&mut self, key: K) -> Option<V> {
        self.transaction_prep();

        let value_opt = self.map.remove(&key);

        if let Some(val) = value_opt.clone() {
            self.journal.push(JournalEntry::Remove(Remove { key, val }));
        }

        value_opt
    }

    /// **NOTE:** Does not clear journal. Performs remove operations
    /// on all items in the map, recording each. If you want to clear the map
    /// and the journal, use `clear_all()`
    pub fn clear(&mut self) {
        self.transaction_prep();

        for (key, val) in self.map.iter() {
            self.journal.push(JournalEntry::Remove(Remove {
                key: key.clone(),
                val: val.clone(),
            }));
        }

        self.map.clear();
    }

    /// Insert a tag at the current journal index. If the tag already exists,
    /// return an error.
    pub fn append_tag(&mut self, tag: T) -> Result<(), JournalError<T>> {
        if self.tags.contains(&tag) {
            return Err(JournalError::TagExists { tag });
        }
        self.current_tag_index = Some(self.journal.len());
        self.journal
            .push(JournalEntry::Tag(Tag { tag: tag.clone() }));
        self.tags.insert(tag);

        Ok(())
    }

    /// Replays the journal in reverse till the next tag is found. If no
    /// tag is found, nothing happens.
    pub fn reverse_to_next(&mut self) -> bool {
        let Some(range) = self.i_tag_reverse_seek(|_| true) else {
            return false;
        };

        if range.is_empty() {
            return false;
        }

        self.reverse_range(range);

        true
    }

    /// Replays the journal in reverse till the specified tag is found. If the
    /// tag isn't found, nothing happens.
    pub fn reverse_to_tag(&mut self, tag: &T) -> bool {
        let Some(range) = self.i_tag_reverse_seek(|other_tag| other_tag == tag) else {
            return false;
        };

        if range.is_empty() {
            return false;
        }

        self.reverse_range(range);

        true
    }

    /// Replays the journal till the next tag is found. If no
    /// tag is found, nothing happens.
    pub fn forward_to_next(&mut self) -> bool {
        let Some(range) = self.i_tag_forward_seek(|_| true) else {
            return false;
        };

        if range.is_empty() {
            return false;
        }

        self.forward_range(range);

        true
    }

    /// Replays the journal till the specified tag is found. If the
    /// tag isn't found, nothing happens.
    pub fn forward_to_tag(&mut self, tag: &T) -> bool {
        let Some(range) = self.i_tag_forward_seek(|other_tag| other_tag == tag) else {
            return false;
        };

        if range.is_empty() {
            return false;
        }

        self.forward_range(range);

        true
    }

    /// Replays the journal entries in the specified range in reverse
    fn reverse_range(&mut self, range: Range<usize>) {
        for entry in self.journal[range].iter().rev() {
            entry.reverse(&mut self.map);
        }
    }

    /// Replays the journal entries in the specified range
    fn forward_range(&mut self, range: Range<usize>) {
        for entry in self.journal[range].iter() {
            entry.forward(&mut self.map);
        }
    }

    /// Generates a journal struct based on the current state of the JournalMap.
    ///
    /// ### Flags:
    /// * `truncate_future` - remove all journal entries past the currently
    /// selected tag.
    pub fn generate_journal(&self, truncate_future: bool) -> JournalMapJournal<K, V, T> {
        let i_copy_start = 0;
        let i_copy_end = if truncate_future && let Some(i_tag) = self.current_tag_index {
            i_tag + 1
        } else {
            self.journal.len()
        };

        let journal_src = &self.journal[i_copy_start..i_copy_end];

        let journal = journal_src.to_vec();

        JournalMapJournal { journal }
    }

    /// Generated a JournalMap by replaying a JournalMapJournal from start to
    /// finish.
    pub fn from_journal(val: &JournalMapJournal<K, V, T>) -> Self {
        let capacity = val.journal.len();
        let map = HashMap::with_capacity(capacity);
        let tags = HashSet::with_capacity(capacity);
        let mut journal_map = Self {
            map,
            tags,
            journal: val.journal.clone(),
            current_tag_index: None,
        };

        for entry in journal_map.journal.iter() {
            entry.forward(&mut journal_map.map);

            if let JournalEntry::Tag(tag) = entry {
                journal_map.tags.insert(tag.tag.clone());
            }
        }

        journal_map
    }
}

impl<K: Debug + Eq + Hash, V: Debug, T: Debug + Eq + Hash> std::fmt::Debug for JournalMap<K, V, T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let alternate = f.alternate();
        write!(f, "{{")?;
        for (i_entry, entry) in self.journal.iter().enumerate() {
            if alternate {
                write!(f, "\n\t")?;
            }

            if let Some(i_tag) = self.current_tag_index
                && i_tag == i_entry
            {
                write!(f, "@ ")?;
            }

            write!(f, "{:?},", entry)?;
        }

        if alternate {
            writeln!(f)?;
        }

        write!(f, "}}")?;

        Ok(())
    }
}

/// An extracted journal from a JournalMap. Used for storage purposes: if you
/// want to serialize a JournalMap for later, you want to extract a
/// JournalMapJournal and then serialize said journal.
#[derive(Debug, Clone)]
pub struct JournalMapJournal<K: Hash + Eq, V, T: Hash + Eq> {
    journal: Vec<JournalEntry<K, V, T>>,
}

impl<K: Hash + Eq, V, T: Hash + Eq> JournalMapJournal<K, V, T> {
    #[inline]
    pub fn len(&self) -> usize {
        self.journal.len()
    }
}

#[cfg(feature = "serde")]
impl<K: SerializeTrait + Hash + Eq, V: SerializeTrait, T: SerializeTrait + Eq + Hash> SerializeTrait
    for JournalMapJournal<K, V, T>
{
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        self.journal.serialize(serializer)
    }
}

#[cfg(feature = "serde")]
impl<
    'de,
    K: DeserializeTrait<'de> + Hash + Eq,
    V: DeserializeTrait<'de>,
    T: DeserializeTrait<'de> + Eq + Hash,
> DeserializeTrait<'de> for JournalMapJournal<K, V, T>
{
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        Ok(Self {
            journal: Vec::deserialize(deserializer)?,
        })
    }
}