// Copyright 2019 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::borrow::Borrow;
use std::collections::BTreeMap;

/// A BTreeMap that supports 2 types of keys per value. All the usual restrictions and warnings for
/// `std::collections::BTreeMap` also apply to this struct. Additionally, there is a 1:1
/// relationship between the 2 key types. In other words, for each `K1` in the map, there is exactly
/// one `K2` in the map and vice versa.
pub struct MultikeyBTreeMap<K1, K2, V> {
    // We need to keep a copy of the second key in the main map so that we can remove entries using
    // just the main key. Otherwise we would require the caller to provide both keys when calling
    // `remove`.
    main: BTreeMap<K1, (K2, V)>,
    alt: BTreeMap<K2, K1>,
}

impl<K1, K2, V> MultikeyBTreeMap<K1, K2, V>
where
    K1: Clone + Ord,
    K2: Clone + Ord,
{
    /// Create a new empty MultikeyBTreeMap.
    pub fn new() -> Self {
        MultikeyBTreeMap {
            main: BTreeMap::default(),
            alt: BTreeMap::default(),
        }
    }

    /// Returns a reference to the value corresponding to the key.
    ///
    /// The key may be any borrowed form of `K1``, but the ordering on the borrowed form must match
    /// the ordering on `K1`.
    pub fn get<Q>(&self, key: &Q) -> Option<&V>
    where
        K1: Borrow<Q>,
        Q: Ord + ?Sized,
    {
        self.main.get(key).map(|(_, v)| v)
    }

    /// Returns a reference to the value corresponding to the alternate key.
    ///
    /// The key may be any borrowed form of the `K2``, but the ordering on the borrowed form must
    /// match the ordering on `K2`.
    ///
    /// Note that this method performs 2 lookups: one to get the main key and another to get the
    /// value associated with that key. For best performance callers should prefer the `get` method
    /// over this method whenever possible as `get` only needs to perform one lookup.
    pub fn get_alt<Q2>(&self, key: &Q2) -> Option<&V>
    where
        K2: Borrow<Q2>,
        Q2: Ord + ?Sized,
    {
        if let Some(k) = self.alt.get(key) {
            self.get(k)
        } else {
            None
        }
    }

    /// Inserts a new entry into the map with the given keys and value.
    ///
    /// Returns `None` if the map did not have an entry with `k1` or `k2` present. If exactly one
    /// key was present, then the value associated with that key is updated, the other key is
    /// removed, and the old value is returned. If **both** keys were present then the value
    /// associated with the main key is updated, the value associated with the alternate key is
    /// removed, and the old value associated with the main key is returned.
    pub fn insert(&mut self, k1: K1, k2: K2, v: V) -> Option<V> {
        let oldval = if let Some(oldkey) = self.alt.insert(k2.clone(), k1.clone()) {
            self.main.remove(&oldkey)
        } else {
            None
        };
        self.main
            .insert(k1, (k2.clone(), v))
            .or(oldval)
            .map(|(oldk2, v)| {
                if oldk2 != k2 {
                    self.alt.remove(&oldk2);
                }
                v
            })
    }

    /// Remove a key from the map, returning the value associated with that key if it was previously
    /// in the map.
    ///
    /// The key may be any borrowed form of `K1``, but the ordering on the borrowed form must match
    /// the ordering on `K1`.
    pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
    where
        K1: Borrow<Q>,
        Q: Ord + ?Sized,
    {
        self.main.remove(key).map(|(k2, v)| {
            self.alt.remove(&k2);
            v
        })
    }

    /// Clears the map, removing all values.
    pub fn clear(&mut self) {
        self.alt.clear();
        self.main.clear()
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn get() {
        let mut m = MultikeyBTreeMap::<u64, i64, u32>::new();

        let k1 = 0xc6c8_f5e0_b13e_ed40;
        let k2 = 0x1a04_ce4b_8329_14fe;
        let val = 0xf4e3_c360;
        assert!(m.insert(k1, k2, val).is_none());

        assert_eq!(*m.get(&k1).expect("failed to look up main key"), val);
        assert_eq!(*m.get_alt(&k2).expect("failed to look up alt key"), val);
    }

    #[test]
    fn update_main_key() {
        let mut m = MultikeyBTreeMap::<u64, i64, u32>::new();

        let k1 = 0xc6c8_f5e0_b13e_ed40;
        let k2 = 0x1a04_ce4b_8329_14fe;
        let val = 0xf4e3_c360;
        assert!(m.insert(k1, k2, val).is_none());

        let new_k1 = 0x3add_f8f8_c7c5_df5e;
        let val2 = 0x7389_f8a7;
        assert_eq!(
            m.insert(new_k1, k2, val2)
                .expect("failed to update main key"),
            val
        );

        assert!(m.get(&k1).is_none());
        assert_eq!(*m.get(&new_k1).expect("failed to look up main key"), val2);
        assert_eq!(*m.get_alt(&k2).expect("failed to look up alt key"), val2);
    }

    #[test]
    fn update_alt_key() {
        let mut m = MultikeyBTreeMap::<u64, i64, u32>::new();

        let k1 = 0xc6c8_f5e0_b13e_ed40;
        let k2 = 0x1a04_ce4b_8329_14fe;
        let val = 0xf4e3_c360;
        assert!(m.insert(k1, k2, val).is_none());

        let new_k2 = 0x6825_a60b_61ac_b333;
        let val2 = 0xbb14_8f2c;
        assert_eq!(
            m.insert(k1, new_k2, val2)
                .expect("failed to update alt key"),
            val
        );

        assert!(m.get_alt(&k2).is_none());
        assert_eq!(*m.get(&k1).expect("failed to look up main key"), val2);
        assert_eq!(
            *m.get_alt(&new_k2).expect("failed to look up alt key"),
            val2
        );
    }

    #[test]
    fn update_value() {
        let mut m = MultikeyBTreeMap::<u64, i64, u32>::new();

        let k1 = 0xc6c8_f5e0_b13e_ed40;
        let k2 = 0x1a04_ce4b_8329_14fe;
        let val = 0xf4e3_c360;
        assert!(m.insert(k1, k2, val).is_none());

        let val2 = 0xe42d_79ba;
        assert_eq!(
            m.insert(k1, k2, val2).expect("failed to update alt key"),
            val
        );

        assert_eq!(*m.get(&k1).expect("failed to look up main key"), val2);
        assert_eq!(*m.get_alt(&k2).expect("failed to look up alt key"), val2);
    }

    #[test]
    fn update_both_keys_main() {
        let mut m = MultikeyBTreeMap::<u64, i64, u32>::new();

        let k1 = 0xc6c8_f5e0_b13e_ed40;
        let k2 = 0x1a04_ce4b_8329_14fe;
        let val = 0xf4e3_c360;
        assert!(m.insert(k1, k2, val).is_none());

        let new_k1 = 0xc980_587a_24b3_ae30;
        let new_k2 = 0x2773_c5ee_8239_45a2;
        let val2 = 0x31f4_33f9;
        assert!(m.insert(new_k1, new_k2, val2).is_none());

        let val3 = 0x8da1_9cf7;
        assert_eq!(
            m.insert(k1, new_k2, val3)
                .expect("failed to update main key"),
            val
        );

        // Both new_k1 and k2 should now be gone from the map.
        assert!(m.get(&new_k1).is_none());
        assert!(m.get_alt(&k2).is_none());

        assert_eq!(*m.get(&k1).expect("failed to look up main key"), val3);
        assert_eq!(
            *m.get_alt(&new_k2).expect("failed to look up alt key"),
            val3
        );
    }

    #[test]
    fn update_both_keys_alt() {
        let mut m = MultikeyBTreeMap::<u64, i64, u32>::new();

        let k1 = 0xc6c8_f5e0_b13e_ed40;
        let k2 = 0x1a04_ce4b_8329_14fe;
        let val = 0xf4e3_c360;
        assert!(m.insert(k1, k2, val).is_none());

        let new_k1 = 0xc980_587a_24b3_ae30;
        let new_k2 = 0x2773_c5ee_8239_45a2;
        let val2 = 0x31f4_33f9;
        assert!(m.insert(new_k1, new_k2, val2).is_none());

        let val3 = 0x8da1_9cf7;
        assert_eq!(
            m.insert(new_k1, k2, val3)
                .expect("failed to update main key"),
            val2
        );

        // Both k1 and new_k2 should now be gone from the map.
        assert!(m.get(&k1).is_none());
        assert!(m.get_alt(&new_k2).is_none());

        assert_eq!(*m.get(&new_k1).expect("failed to look up main key"), val3);
        assert_eq!(*m.get_alt(&k2).expect("failed to look up alt key"), val3);
    }

    #[test]
    fn remove() {
        let mut m = MultikeyBTreeMap::<u64, i64, u32>::new();

        let k1 = 0xc6c8_f5e0_b13e_ed40;
        let k2 = 0x1a04_ce4b_8329_14fe;
        let val = 0xf4e3_c360;
        assert!(m.insert(k1, k2, val).is_none());

        assert_eq!(m.remove(&k1).expect("failed to remove entry"), val);
        assert!(m.get(&k1).is_none());
        assert!(m.get_alt(&k2).is_none());
    }
}