iptrie 0.11.1

//! Generic prefix trie set structures
use crate::prefix::*;
use crate::trie::lctrie::LevelCompressedTrie;
use crate::trie::patricia::RadixTrie;
use std::num::NonZeroUsize;

#[cfg(feature = "graphviz")]
use crate::graphviz::DotWriter;

#[cfg(feature = "graphviz")]
use std::fmt::Display;

/// A set of Ip prefixes based on a radix binary trie
#[derive(Clone)]
pub struct RTrieSet<P: IpPrefix>(pub(crate) RadixTrie<P, ()>);

/// A set of Ip prefixes based on a level-compressed trie
pub struct LCTrieSet<P: IpPrefix>(pub(crate) LevelCompressedTrie<P, ()>);

impl<P: IpRootPrefix> RTrieSet<P> {
    /// Creates a new set which contains the root prefix.
    #[inline]
    pub fn new() -> Self {
        Self::with_capacity(1000)
    }

    /// Creates a new set with an initial capacity.
    ///
    /// The returned set already contains the root prefix.
    #[inline]
    pub fn with_capacity(capacity: usize) -> Self {
        Self(RadixTrie::new((), capacity))
    }
}

impl<P: IpPrefix> RTrieSet<P> {
    /// Returns the size of the set.
    ///
    /// Notice that it never equals zero since the top prefix is
    /// always present in the set.
    #[inline]
    #[allow(clippy::len_without_is_empty)]
    pub fn len(&self) -> NonZeroUsize {
        self.0.len()
    }

    /// Compress this Patricia trie in a LC-Trie.
    ///
    /// For lookup algorithms, a Patricia trie performs unit bit checking and LC-Trie
    /// performs multi bits checking. So the last one is more performant but it
    /// cannot be modified (no insertion or removal operations are provided).
    #[inline]
    pub fn compress(self) -> LCTrieSet<P> {
        LCTrieSet(LevelCompressedTrie::new(self.0))
    }

    #[inline]
    pub fn shrink_to_fit(&mut self) {
        self.0.shrink_to_fit()
    }

    /// Inserts a new element in the set.
    ///
    /// If the specified element already exists in the set, `false` is returned.
    ///
    /// # Example
    /// ```
    /// #  use iptrie::*;
    /// use std::net::Ipv4Addr;
    ///
    /// let addr = Ipv4Addr::new(1,1,1,1);
    /// let mut trie = Ipv4RTrieSet::new();
    ///
    /// let ip20 = Ipv4Prefix::new(addr, 20).unwrap();
    /// let ip22 = Ipv4Prefix::new(addr, 22).unwrap();
    ///
    /// assert_eq!( trie.insert(ip20), true);
    /// assert_eq!( trie.insert(ip22), true);
    /// assert_eq!( trie.insert(ip20), false);
    /// ```
    #[inline]
    pub fn insert(&mut self, k: P) -> bool {
        self.0.insert(k, ()).is_none()
    }

    /// Checks if an element is present (exact match).
    ///
    /// # Example
    /// ```
    /// #  use iptrie::*;
    /// use std::net::Ipv4Addr;
    ///
    /// let addr = Ipv4Addr::new(1,1,1,1);
    ///
    /// let ip20 = Ipv4Prefix::new(addr, 20).unwrap();
    /// let ip22 = Ipv4Prefix::new(addr, 22).unwrap();
    /// let ip24 = Ipv4Prefix::new(addr, 24).unwrap();
    ///
    /// let trie = Ipv4RTrieSet::from_iter([ip20,ip24]);
    ///
    /// assert_eq!( trie.contains(&ip20), true);
    /// assert_eq!( trie.contains(&ip22), false);
    /// assert_eq!( trie.contains(&ip24), true);
    /// ```
    #[inline]
    pub fn contains<Q>(&self, k: &Q) -> bool
    where
        Q: IpPrefix<Addr = P::Addr>,
        P: IpPrefixCovering<Q>,
    {
        self.0.get(k).is_some()
    }

    /// Removes a previously inserted prefix (exact match).
    ///
    /// Returns `false` is the element was not present in the set
    /// and `true` if the removal is effective.
    ///
    /// # Example
    /// ```
    /// #  use iptrie::*;
    /// use std::net::Ipv4Addr;
    ///
    /// let addr = Ipv4Addr::new(1,1,1,1);
    ///
    /// let ip20 = Ipv4Prefix::new(addr, 20).unwrap();
    /// let ip22 = Ipv4Prefix::new(addr, 22).unwrap();
    ///
    /// let mut trie = Ipv4RTrieSet::from_iter([ip20,ip22]);
    ///
    /// assert_eq!( trie.contains(&ip20), true);
    /// assert_eq!(trie.remove(&ip20), true);
    /// assert_eq!(trie.remove(&ip20), false);
    ///
    /// assert_eq!( trie.contains(&ip22), true);
    /// assert_eq!( trie.contains(&ip20), false);
    /// ```
    #[inline]
    pub fn remove<Q>(&mut self, k: &Q) -> bool
    where
        Q: IpPrefix<Addr = P::Addr>,
        P: IpPrefixCovering<Q>,
    {
        self.0.remove(k).is_some()
    }

    /// Replace an existing prefix.
    ///
    /// Adds a prefix to the set, replacing the existing one, if any (exact match performed).
    /// Returns the replaced value.
    ///
    /// # Example
    /// ```
    /// # use iptrie::*;
    /// # use iptrie::set::*;
    /// use std::net::Ipv4Addr;
    /// use ipnet::Ipv4Net;
    /// let mut trie = RTrieSet::new();
    ///
    /// let addr1 = Ipv4Addr::new(1,1,1,1);
    /// let addr2 = Ipv4Addr::new(1,1,1,2);
    /// let ip20 = Ipv4Net::new(addr1, 20).unwrap();
    /// let ip20b = Ipv4Net::new(addr2, 20).unwrap();
    ///
    /// trie.insert(ip20);
    /// assert_eq!(trie.get(&ip20).unwrap().to_string(), "1.1.1.1/20".to_string());
    ///
    /// assert_eq!(trie.insert(ip20b), false);
    /// assert_eq!(trie.get(&ip20).unwrap().to_string(), "1.1.1.1/20".to_string());
    ///
    /// assert_eq!(trie.replace(ip20b), Some(ip20));
    /// assert_eq!(trie.get(&ip20).unwrap().to_string(), "1.1.1.2/20".to_string());
    /// ```
    #[inline]
    pub fn replace(&mut self, k: P) -> Option<P> {
        self.0.replace(k, ()).map(|(k, _)| k)
    }

    /// Gets the value associated with an exact match of the key.
    ///
    /// To access to the longest prefix match, use [`Self::lookup`].
    ///
    /// # Example
    /// ```
    /// # use iptrie::*;
    /// # use iptrie::set::*;
    /// use std::net::Ipv4Addr;
    /// use ipnet::Ipv4Net;
    /// let mut trie = RTrieSet::new();
    ///
    /// let addr1 = Ipv4Addr::new(1,1,1,1);
    /// let addr2 = Ipv4Addr::new(1,1,1,2);
    /// let ip20 = Ipv4Net::new(addr1, 20).unwrap();
    /// let ip20b = Ipv4Net::new(addr2, 20).unwrap();
    ///
    /// trie.insert(ip20);
    /// assert_eq!(trie.get(&ip20).unwrap().to_string(), "1.1.1.1/20".to_string());
    ///
    /// trie.insert(ip20b);
    /// assert_eq!(trie.get(&ip20).unwrap().to_string(), "1.1.1.1/20".to_string());
    /// ```
    #[inline]
    pub fn get<Q>(&self, k: &Q) -> Option<&P>
    where
        Q: IpPrefix<Addr = P::Addr>,
        P: IpPrefixCovering<Q>,
    {
        self.0.get(k).map(|(k, _)| k)
    }

    /// Gets the longest prefix which matches the given key.
    ///
    /// As the top prefix always matches, it never fails.
    ///
    /// To access to the exact prefix match, use [`Self::get`].
    ///
    /// # Example
    /// ```
    /// # use iptrie::*;
    /// use std::net::Ipv4Addr;
    /// let mut trie = Ipv4RTrieSet::new();
    ///
    /// let addr = Ipv4Addr::new(1,1,1,1);
    /// let ip20 = Ipv4Prefix::new(addr, 20).unwrap();
    /// let ip22 = Ipv4Prefix::new(addr, 22).unwrap();
    /// let ip24 = Ipv4Prefix::new(addr, 24).unwrap();
    ///
    /// trie.insert(ip20);
    /// trie.insert(ip24);
    ///
    /// assert_eq!( trie.lookup(&ip20), &ip20);
    /// assert_eq!( trie.lookup(&ip22), &ip20);
    /// assert_eq!( trie.lookup(&ip24), &ip24);
    ///
    /// assert_eq!( trie.lookup(&addr), &ip24);
    /// ```
    #[inline]
    pub fn lookup<Q>(&self, k: &Q) -> &P
    where
        Q: IpPrefix<Addr = P::Addr>,
        P: IpPrefixCovering<Q>,
    {
        self.0.lookup(k).0
    }

    /// Iterates over all the prefixes of this set.
    #[inline]
    pub fn iter(&self) -> impl Iterator<Item = &P> + '_ {
        self.0.leaves.0.iter().map(|(k, _)| k)
    }

    #[inline]
    pub fn info(&self) {
        self.0.info()
    }
}

impl<P: IpRootPrefix> Default for RTrieSet<P> {
    #[inline]
    fn default() -> Self {
        Self::new()
    }
}

impl<P: IpPrefix> Extend<P> for RTrieSet<P> {
    fn extend<I: IntoIterator<Item = P>>(&mut self, iter: I) {
        iter.into_iter().for_each(|item| {
            self.insert(item);
        })
    }
}

impl<P: IpRootPrefix> FromIterator<P> for RTrieSet<P> {
    fn from_iter<I: IntoIterator<Item = P>>(iter: I) -> Self {
        let mut trieset = Self::default();
        trieset.extend(iter);
        trieset
    }
}

impl<P: IpPrefix> LCTrieSet<P> {
    /// Returns the size of the set.
    ///
    /// Notice that it never equals zero since the top prefix is
    /// always present in the set.
    #[inline]
    #[allow(clippy::len_without_is_empty)]
    pub fn len(&self) -> NonZeroUsize {
        self.0.len()
    }

    #[inline]
    pub fn info(&self) {
        self.0.info()
    }

    /// Checks if an element is present (exact match).
    ///
    /// # Example
    /// ```
    /// #  use iptrie::*;
    /// use std::net::Ipv4Addr;
    ///
    /// let addr = Ipv4Addr::new(1,1,1,1);
    ///
    /// let ip20 = Ipv4Prefix::new(addr, 20).unwrap();
    /// let ip22 = Ipv4Prefix::new(addr, 22).unwrap();
    /// let ip24 = Ipv4Prefix::new(addr, 24).unwrap();
    ///
    /// let trie = Ipv4RTrieSet::from_iter([ip20,ip24]);
    /// let lctrie = trie.compress();
    ///
    /// assert_eq!( lctrie.contains(&ip20), true);
    /// assert_eq!( lctrie.contains(&ip22), false);
    /// assert_eq!( lctrie.contains(&ip24), true);
    /// ```
    #[inline]
    pub fn contains<Q>(&self, k: &Q) -> bool
    where
        Q: IpPrefix<Addr = P::Addr>,
        P: IpPrefixCovering<Q>,
    {
        self.0.get(k).is_some()
    }

    /// Gets the value associated with an exact match of the key.
    ///
    /// To access to the longest prefix match, use [`Self::lookup`].
    ///
    /// # Example
    /// ```
    /// # use iptrie::*;
    /// # use iptrie::set::*;
    /// use std::net::Ipv4Addr;
    /// let mut trie = RTrieSet::new();
    ///
    /// let addr = Ipv4Addr::new(1,1,1,1);
    /// let ip20 = Ipv4Prefix::new(addr, 20).unwrap();
    /// let ip22 = Ipv4Prefix::new(addr, 22).unwrap();
    /// let ip24 = Ipv4Prefix::new(addr, 24).unwrap();
    ///
    /// trie.insert(ip24);
    /// trie.insert(ip20);
    ///
    /// let lctrie = trie.compress();
    ///
    /// assert!( lctrie.get(&ip24).is_some());
    /// assert!( lctrie.get(&ip22).is_none());
    /// assert!( lctrie.get(&ip20).is_some());
    /// ```
    #[inline]
    pub fn get<Q>(&self, k: &Q) -> Option<&P>
    where
        Q: IpPrefix<Addr = P::Addr>,
        P: IpPrefixCovering<Q>,
    {
        self.0.get(k).map(|(k, _)| k)
    }

    /// Gets the longest prefix which matches the given key.
    ///
    /// As the top prefix always matches, it never fails.
    ///
    /// To access to the exact prefix match, use [`Self::get`].
    ///
    /// # Example
    /// ```
    /// # use iptrie::*;
    /// use std::net::Ipv4Addr;
    /// let mut trie = Ipv4RTrieSet::new();
    ///
    /// let addr = Ipv4Addr::new(1,1,1,1);
    /// let ip20 = Ipv4Prefix::new(addr, 20).unwrap();
    /// let ip22 = Ipv4Prefix::new(addr, 22).unwrap();
    /// let ip24 = Ipv4Prefix::new(addr, 24).unwrap();
    ///
    /// trie.insert(ip20);
    /// trie.insert(ip24);
    ///
    /// let lctrie = trie.compress();
    ///
    /// assert_eq!( lctrie.lookup(&ip20), &ip20);
    /// assert_eq!( lctrie.lookup(&ip22), &ip20);
    /// assert_eq!( lctrie.lookup(&ip24), &ip24);
    ///
    /// assert_eq!( lctrie.lookup(&addr), &ip24);
    /// ```
    #[inline]
    pub fn lookup<Q>(&self, k: &Q) -> &P
    where
        Q: IpPrefix<Addr = P::Addr>,
        P: IpPrefixCovering<Q>,
    {
        self.0.lookup(k).0
    }

    /// Iterates over all the prefixes of this set.
    #[inline]
    pub fn iter(&self) -> impl Iterator<Item = &P> + '_ {
        self.0.leaves.0.iter().map(|(k, _)| k)
    }
}

impl<P: IpRootPrefix> FromIterator<P> for LCTrieSet<P> {
    fn from_iter<I: IntoIterator<Item = P>>(iter: I) -> Self {
        RTrieSet::from_iter(iter).compress()
    }
}

#[cfg(feature = "graphviz")]
impl<P: IpPrefix + Display> DotWriter for RTrieSet<P> {
    fn write_dot(&self, dot: &mut dyn std::io::Write) -> std::io::Result<()> {
        self.0.write_dot(dot)
    }
}

#[cfg(feature = "graphviz")]
impl<P: IpPrefix + Display> DotWriter for LCTrieSet<P> {
    fn write_dot(&self, dot: &mut dyn std::io::Write) -> std::io::Result<()> {
        self.0.write_dot(dot)
    }
}