/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

pub trait AbstractDomain: Clone + Eq {
    fn bottom() -> Self;
    fn top() -> Self;
    fn is_bottom(&self) -> bool;
    fn is_top(&self) -> bool;
    fn leq(&self, rhs: &Self) -> bool;

    fn join(mut self, rhs: Self) -> Self {
        self.join_with(rhs);
        self
    }

    fn meet(mut self, rhs: Self) -> Self {
        self.meet_with(rhs);
        self
    }

    fn widen(mut self, rhs: Self) -> Self {
        self.widen_with(rhs);
        self
    }

    fn narrow(mut self, rhs: Self) -> Self {
        self.narrow_with(rhs);
        self
    }

    fn join_with(&mut self, rhs: Self);
    fn meet_with(&mut self, rhs: Self);
    fn widen_with(&mut self, rhs: Self);
    fn narrow_with(&mut self, rhs: Self);
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::borrow::Cow;
use std::hash::Hash;

use im::HashMap;

use crate::datatype::bitvec::BitVec;
use crate::datatype::AbstractDomain;
use crate::datatype::PatriciaTreeMap;

/*
 * An abstract environment is a type of abstract domain that maps the variables
 * of a program to elements of a common abstract domain. For example, to perform
 * range analysis one can use an abstract environment that maps variable names
 * to intervals:
 *
 *   {"x" -> [-1, 1], "i" -> [0, 10], ...}
 *
 * Another example is descriptive type analysis for Dex code, where one computes
 * the set of all possible Java classes a register can hold a reference to at
 * any point in the code:
 *
 *  {"v0" -> {android.app.Fragment, java.lang.Object}, "v1" -> {...}, ...}
 *
 * This type of domain is commonly used for nonrelational (also called
 * attribute-independent) analyses that do not track relationships among
 * program variables. Please note that by definition of an abstract
 * environment, if the value _|_ appears in a variable binding, then no valid
 * execution state can ever be represented by this abstract environment. Hence,
 * assigning _|_ to a variable is equivalent to setting the entire environment
 * to _|_.
 */

pub trait AbstractEnvironment<V, D: AbstractDomain>: AbstractDomain {
    type ContainerType;
    fn bindings(&self) -> Option<&Self::ContainerType>;
    fn into_bindings(self) -> Option<Self::ContainerType>;

    fn len(&self) -> usize;
    fn is_empty(&self) -> bool;
    fn get(&self, variable: &V) -> Cow<'_, D>;
    fn set(&mut self, variable: V, domain: D);
    fn update(&mut self, variable: &V, op: impl FnOnce(&mut D));
}

#[derive(Clone, PartialEq, Eq, Debug)]
pub enum HashMapAbstractEnvironment<V: Clone + Eq + Hash, D: AbstractDomain> {
    Value(HashMap<V, D>),
    Bottom,
}

impl<V, D> AbstractEnvironment<V, D> for HashMapAbstractEnvironment<V, D>
where
    V: Clone + Eq + Hash,
    D: AbstractDomain,
{
    type ContainerType = HashMap<V, D>;

    fn bindings(&self) -> Option<&HashMap<V, D>> {
        match self {
            Self::Value(ref map) => Some(map),
            _ => None,
        }
    }

    fn into_bindings(self) -> Option<HashMap<V, D>> {
        match self {
            Self::Value(map) => Some(map),
            _ => None,
        }
    }

    fn len(&self) -> usize {
        self.bindings()
            .expect("Bottom doesn't have a length!")
            .len()
    }

    fn is_empty(&self) -> bool {
        match self {
            Self::Value(map) => map.is_empty(),
            Self::Bottom => true,
        }
    }

    fn get(&self, variable: &V) -> Cow<'_, D> {
        let map = match self {
            Self::Value(map) => map,
            Self::Bottom => return Cow::Owned(D::bottom()),
        };

        match map.get(variable) {
            Some(domain) => Cow::Borrowed(domain),
            None => Cow::Owned(D::top()),
        }
    }

    fn set(&mut self, variable: V, domain: D) {
        if let Self::Value(map) = self {
            if domain.is_top() {
                map.remove(&variable);
            } else if domain.is_bottom() {
                *self = Self::Bottom;
            } else {
                map.insert(variable, domain);
            }
        }
    }

    fn update(&mut self, variable: &V, op: impl FnOnce(&mut D)) {
        if let Self::Value(map) = self {
            match map.get_mut(variable) {
                Some(explicit_value) => {
                    op(explicit_value);
                    if explicit_value.is_top() {
                        // Use implicit binding
                        map.remove(variable);
                    } else if explicit_value.is_bottom() {
                        *self = Self::Bottom;
                    }
                }
                None => {
                    let mut domain = D::top();
                    op(&mut domain);
                    if domain.is_top() {
                        // Do nothing. Continue to use implicit binding.
                    } else if domain.is_bottom() {
                        *self = Self::Bottom;
                    } else {
                        map.insert(variable.clone(), domain);
                    }
                }
            }
        }
    }
}

impl<V, D> AbstractDomain for HashMapAbstractEnvironment<V, D>
where
    V: Clone + Eq + Hash,
    D: AbstractDomain,
{
    fn bottom() -> Self {
        Self::Bottom
    }

    fn top() -> Self {
        Self::Value(HashMap::new())
    }

    fn is_bottom(&self) -> bool {
        matches!(self, Self::Bottom)
    }

    fn is_top(&self) -> bool {
        match self {
            Self::Value(map) => map.is_empty(),
            _ => false,
        }
    }

    fn leq(&self, rhs: &Self) -> bool {
        use HashMapAbstractEnvironment::*;
        match (self, rhs) {
            (Value(l_map), Value(r_map)) => {
                if l_map.len() < r_map.len() {
                    // Perf optimization
                    false
                } else {
                    for (l_k, l_v) in l_map.iter() {
                        if let Some(r_v) = r_map.get(l_k) {
                            if !l_v.leq(r_v) {
                                return false;
                            }
                        }
                    }

                    for (r_k, _) in r_map.iter() {
                        if l_map.get(r_k).is_none() {
                            return false;
                        }
                    }
                    true
                }
            }
            (Bottom, _) => true,
            (_, Bottom) => self.is_bottom(),
        }
    }

    fn join_with(&mut self, rhs: Self) {
        Self::join_like_operation(self, rhs, |d1, d2| d1.join_with(d2));
    }

    fn meet_with(&mut self, rhs: Self) {
        Self::meet_like_operation(self, rhs, |d1, d2| d1.meet_with(d2));
    }

    fn widen_with(&mut self, rhs: Self) {
        Self::join_like_operation(self, rhs, |d1, d2| d1.widen_with(d2));
    }

    fn narrow_with(&mut self, rhs: Self) {
        Self::meet_like_operation(self, rhs, |d1, d2| d1.narrow_with(d2));
    }
}

impl<V, D> HashMapAbstractEnvironment<V, D>
where
    V: Clone + Eq + Hash,
    D: AbstractDomain,
{
    fn join_like_operation(lhs: &mut Self, rhs: Self, operation: impl Fn(&mut D, D)) {
        use HashMapAbstractEnvironment::*;

        match (lhs, rhs) {
            (Value(l_map), Value(ref mut r_map)) => {
                l_map.retain(|l_k, _| r_map.contains_key(l_k));
                for (l_k, l_v) in l_map.iter_mut() {
                    let r_v = r_map.remove(l_k).unwrap();
                    operation(l_v, r_v);
                }
                l_map.retain(|_, l_v| !l_v.is_top());
            }
            (lhs @ Bottom, rhs) => *lhs = rhs,
            (_, Bottom) => {}
        }
    }

    fn meet_like_operation(lhs: &mut Self, rhs: Self, operation: impl Fn(&mut D, D)) {
        use HashMapAbstractEnvironment::*;

        match (&mut *lhs, rhs) {
            (Value(l_map), Value(r_map)) => {
                for (r_k, r_v) in r_map.into_iter() {
                    if let Some(l_v) = l_map.get_mut(&r_k) {
                        operation(l_v, r_v);
                        // l_v wasn't top. A meet-like operation should not make it top.
                        assert!(!l_v.is_top());
                        if l_v.is_bottom() {
                            *lhs = Bottom;
                            return;
                        }
                    } else {
                        // The value is Top, we just insert the other value (Top is the
                        // identity for meet-like operations).
                        l_map.insert(r_k, r_v);
                    }
                }
            }
            (Bottom, _) => {}
            (_, rhs @ Bottom) => {
                *lhs = rhs;
            }
        }
    }
}

#[derive(Debug)]
pub enum PatriciaTreeMapAbstractEnvironment<V: Into<BitVec> + Clone, D: Sized + Eq + AbstractDomain>
{
    Value(PatriciaTreeMap<V, D>), // Use empty map value as top
    Bottom,
}

impl<V, D> Clone for PatriciaTreeMapAbstractEnvironment<V, D>
where
    V: Into<BitVec> + Clone,
    D: Sized + AbstractDomain,
{
    fn clone(&self) -> Self {
        match self {
            Self::Value(map) => Self::Value(map.clone()),
            Self::Bottom => Self::Bottom,
        }
    }
}

impl<V, D> PartialEq for PatriciaTreeMapAbstractEnvironment<V, D>
where
    V: Into<BitVec> + Clone,
    D: Sized + AbstractDomain,
{
    fn eq(&self, rhs: &Self) -> bool {
        match (self, rhs) {
            (Self::Value(l_map), Self::Value(r_map)) => l_map == r_map,
            (Self::Bottom, Self::Bottom) => true,
            (_, _) => false,
        }
    }
}

impl<V, D> Eq for PatriciaTreeMapAbstractEnvironment<V, D>
where
    V: Into<BitVec> + Clone,
    D: Sized + AbstractDomain,
{
}

impl<V, D> AbstractEnvironment<V, D> for PatriciaTreeMapAbstractEnvironment<V, D>
where
    V: Into<BitVec> + Clone,
    D: Sized + AbstractDomain,
{
    type ContainerType = PatriciaTreeMap<V, D>;

    fn bindings(&self) -> Option<&PatriciaTreeMap<V, D>> {
        match self {
            Self::Value(ref map) => Some(map),
            _ => None,
        }
    }

    fn into_bindings(self) -> Option<PatriciaTreeMap<V, D>> {
        match self {
            Self::Value(map) => Some(map),
            _ => None,
        }
    }

    fn len(&self) -> usize {
        self.bindings()
            .expect("Bottom doesn't have a length!")
            .len()
    }

    fn is_empty(&self) -> bool {
        match self {
            Self::Value(map) => map.is_empty(),
            Self::Bottom => true,
        }
    }

    fn get(&self, variable: &V) -> Cow<'_, D> {
        let map = match self {
            Self::Value(map) => map,
            Self::Bottom => return Cow::Owned(D::bottom()),
        };

        match map.get(variable.clone()) {
            Some(domain) => Cow::Borrowed(domain),
            None => Cow::Owned(D::top()),
        }
    }

    fn set(&mut self, variable: V, domain: D) {
        if let Self::Value(map) = self {
            if domain.is_top() {
                map.remove(variable);
            } else if domain.is_bottom() {
                *self = Self::Bottom;
            } else {
                map.upsert(variable, domain);
            }
        }
    }

    fn update(&mut self, variable: &V, op: impl FnOnce(&mut D)) {
        let map = match self {
            Self::Bottom => return,
            Self::Value(ref mut map) => map,
        };

        let mut update_domain = match map.get(variable.clone()) {
            Some(domain) => domain.clone(),
            None => D::top(),
        };

        op(&mut update_domain);

        self.set(variable.clone(), update_domain);
    }
}

impl<V, D> AbstractDomain for PatriciaTreeMapAbstractEnvironment<V, D>
where
    V: Into<BitVec> + Clone,
    D: Sized + AbstractDomain,
{
    fn bottom() -> Self {
        Self::Bottom
    }

    fn top() -> Self {
        Self::Value(PatriciaTreeMap::new())
    }

    fn is_bottom(&self) -> bool {
        matches!(self, Self::Bottom)
    }

    fn is_top(&self) -> bool {
        match self {
            Self::Value(map) => map.is_empty(),
            _ => false,
        }
    }

    fn leq(&self, rhs: &Self) -> bool {
        use PatriciaTreeMapAbstractEnvironment::*;

        match (self, rhs) {
            (Value(self_map), Value(other_map)) => self_map.leq(other_map, &D::top()),
            (Bottom, _) => true,
            (_, Bottom) => self.is_bottom(),
        }
    }

    fn join_with(&mut self, rhs: Self) {
        Self::join_like_operation(self, rhs, |d1, d2| d1.join_with(d2));
    }

    fn meet_with(&mut self, rhs: Self) {
        Self::meet_like_operation(self, rhs, |d1, d2| d1.meet_with(d2));
    }

    fn widen_with(&mut self, rhs: Self) {
        Self::join_like_operation(self, rhs, |d1, d2| d1.widen_with(d2));
    }

    fn narrow_with(&mut self, rhs: Self) {
        Self::meet_like_operation(self, rhs, |d1, d2| d1.narrow_with(d2));
    }
}

impl<V, D> PatriciaTreeMapAbstractEnvironment<V, D>
where
    V: Into<BitVec> + Clone,
    D: Sized + AbstractDomain,
{
    fn join_like_operation(lhs: &mut Self, rhs: Self, operation: impl Fn(&mut D, D)) {
        use PatriciaTreeMapAbstractEnvironment::*;

        match (lhs, rhs) {
            (Value(lmap), Value(rmap)) => {
                lmap.intersect_with(&rmap, |s, t| {
                    let mut s = s.clone();
                    operation(&mut s, t.clone());
                    s
                });
            }
            (lhs @ Bottom, rhs) => *lhs = rhs,
            (_, Bottom) => {}
        }
    }

    fn meet_like_operation(lhs: &mut Self, rhs: Self, operation: impl Fn(&mut D, D)) {
        use PatriciaTreeMapAbstractEnvironment::*;

        match (lhs, rhs) {
            (Value(lmap), Value(rmap)) => {
                lmap.union_with(&rmap, |s, t| {
                    let mut s = s.clone();
                    operation(&mut s, t.clone());
                    s
                });
            }
            (Bottom, _) => {}
            (lhs, rhs @ Bottom) => {
                *lhs = rhs;
            }
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::borrow::Cow;
use std::hash::Hash;

use im::HashMap;

use crate::datatype::bitvec::BitVec;
use crate::datatype::AbstractDomain;
use crate::datatype::PatriciaTreeMap;

/*
 * A partition is a mapping from a set of labels to elements in an abstract
 * domain. It denotes a union of properties. A partition is Bottom iff all its
 * bindings are set to Bottom, and it is Top iff all its bindings are set to
 * Top.
 *
 * All lattice operations are applied componentwise.
 */
pub trait AbstractPartition<L, D: AbstractDomain>: AbstractDomain {
    type ContainerType;
    fn bindings(&self) -> Option<&Self::ContainerType>;
    fn into_bindings(self) -> Option<Self::ContainerType>;

    fn len(&self) -> usize;
    fn is_empty(&self) -> bool;
    fn get(&self, label: &L) -> Cow<'_, D>;
    fn set(&mut self, label: L, domain: D);
    fn update(&mut self, label: &L, op: impl FnOnce(&mut D));
}

/*
 * In order to minimize the size of the hashtable, we do not explicitly
 * represent bindings to Bottom.
 *
 * This implementation differs slightly from the textbook definition of a
 * partition: our Top partition cannot have its labels re-bound to anything
 * other than Top. I.e. for all labels L and domains D,
 *
 *   HashMapAbstractPartition::top().set(L, D) == HashMapAbstractPartition::top()
 *
 * This makes for a much simpler implementation.
 */
#[derive(Clone, PartialEq, Eq, Debug)]
pub enum HashMapAbstractPartition<L: Clone + Eq + Hash, D: AbstractDomain> {
    Top,
    Value(HashMap<L, D>), // Use empty map value as bottom
}

impl<L, D> AbstractPartition<L, D> for HashMapAbstractPartition<L, D>
where
    L: Clone + Eq + Hash,
    D: AbstractDomain,
{
    type ContainerType = HashMap<L, D>;

    fn bindings(&self) -> Option<&Self::ContainerType> {
        match self {
            Self::Value(ref map) => Some(map),
            _ => None,
        }
    }

    fn into_bindings(self) -> Option<Self::ContainerType> {
        match self {
            Self::Value(map) => Some(map),
            _ => None,
        }
    }

    fn len(&self) -> usize {
        self.bindings()
            .expect("Top abstract domain doesn't have a length!")
            .len()
    }

    fn is_empty(&self) -> bool {
        match self {
            Self::Top => false,
            Self::Value(map) => map.is_empty(),
        }
    }

    fn get(&self, label: &L) -> Cow<'_, D> {
        let map = match self {
            Self::Top => return Cow::Owned(D::top()),
            Self::Value(map) => map,
        };

        match map.get(label) {
            Some(domain) => Cow::Borrowed(domain),
            None => Cow::Owned(D::bottom()),
        }
    }

    fn set(&mut self, label: L, domain: D) {
        let map = match self {
            Self::Top => return,
            Self::Value(ref mut map) => map,
        };

        // Save some memory by implicitly storing bottom.
        if domain.is_bottom() {
            map.remove(&label);
        } else {
            map.insert(label, domain);
        }
    }

    fn update(&mut self, label: &L, op: impl FnOnce(&mut D)) {
        let map = match self {
            Self::Top => return,
            Self::Value(ref mut map) => map,
        };

        match map.get_mut(label) {
            Some(domain) => op(domain),
            None => {
                let mut temp = D::bottom();
                op(&mut temp);
                if !temp.is_bottom() {
                    map.insert(label.clone(), temp);
                }
            }
        }
    }
}

impl<L, D> AbstractDomain for HashMapAbstractPartition<L, D>
where
    L: Clone + Eq + Hash,
    D: AbstractDomain,
{
    fn bottom() -> Self {
        Self::Value(HashMap::new())
    }

    fn top() -> Self {
        Self::Top
    }

    fn is_bottom(&self) -> bool {
        match self {
            Self::Value(map) => map.is_empty(),
            _ => false,
        }
    }

    fn is_top(&self) -> bool {
        matches!(self, Self::Top)
    }

    fn leq(&self, rhs: &Self) -> bool {
        use HashMapAbstractPartition::*;

        match (self, rhs) {
            (Top, _) => rhs.is_top(),
            (_, Top) => true,
            (Value(self_map), Value(other_map)) => {
                if self_map.len() > other_map.len() {
                    // Perf optimization
                    false
                } else {
                    for (k, v) in self_map.iter() {
                        match other_map.get(k) {
                            Some(rd) => {
                                if !v.leq(rd) {
                                    return false;
                                }
                            }
                            None => return false,
                        }
                    }
                    true
                }
            }
        }
    }

    fn join_with(&mut self, rhs: Self) {
        Self::join_like_operation(self, rhs, |d1, d2| d1.join_with(d2));
    }

    fn meet_with(&mut self, rhs: Self) {
        Self::meet_like_operation(self, rhs, |d1, d2| d1.meet_with(d2));
    }

    fn widen_with(&mut self, rhs: Self) {
        Self::join_like_operation(self, rhs, |d1, d2| d1.widen_with(d2));
    }

    fn narrow_with(&mut self, rhs: Self) {
        Self::meet_like_operation(self, rhs, |d1, d2| d1.narrow_with(d2));
    }
}

impl<L, D> HashMapAbstractPartition<L, D>
where
    L: Clone + Eq + Hash,
    D: AbstractDomain,
{
    fn join_like_operation(lhs: &mut Self, rhs: Self, operation: impl Fn(&mut D, D)) {
        use HashMapAbstractPartition::*;

        match (lhs, rhs) {
            (Top, _) => {}
            (lhs, Top) => {
                *lhs = Top;
            }
            (Value(l_map), Value(r_map)) => {
                for (r_k, r_v) in r_map.into_iter() {
                    if let Some(l_v) = l_map.get_mut(&r_k) {
                        operation(l_v, r_v);
                        // l_v wasn't bottom. A join-like operation should not make it bottom.
                        assert!(!l_v.is_bottom());
                    } else {
                        // The value is Bottom, we just insert the other value (Bottom is the
                        // identity for join-like operations).
                        l_map.insert(r_k, r_v);
                    }
                }
            }
        }
    }

    fn meet_like_operation(lhs: &mut Self, rhs: Self, operation: impl Fn(&mut D, D)) {
        use HashMapAbstractPartition::*;

        match (lhs, rhs) {
            (lhs @ Top, rhs) => {
                *lhs = rhs;
            }
            (_, Top) => {}
            (Value(l_map), Value(mut r_map)) => {
                l_map.retain(|l_k, _| r_map.contains_key(l_k));
                for (l_k, l_v) in l_map.iter_mut() {
                    let r_v = r_map.remove(l_k).unwrap();
                    operation(l_v, r_v);
                }

                l_map.retain(|_, l_v| !l_v.is_bottom());
            }
        }
    }
}

#[derive(Debug)]
pub enum PatriciaTreeMapAbstractPartition<L: Into<BitVec> + Clone, D: Sized + Eq + AbstractDomain> {
    Top,
    Value(PatriciaTreeMap<L, D>), // Use empty map value as bottom
}

impl<L, D> Clone for PatriciaTreeMapAbstractPartition<L, D>
where
    L: Into<BitVec> + Clone,
    D: Sized + AbstractDomain,
{
    fn clone(&self) -> Self {
        use PatriciaTreeMapAbstractPartition::*;

        match self {
            Top => Top,
            Value(map) => Value(map.clone()),
        }
    }
}

impl<L, D> PartialEq for PatriciaTreeMapAbstractPartition<L, D>
where
    L: Into<BitVec> + Clone,
    D: Sized + AbstractDomain,
{
    fn eq(&self, rhs: &Self) -> bool {
        use PatriciaTreeMapAbstractPartition::*;

        match (self, rhs) {
            (Top, Top) => true,
            (Value(l_map), Value(r_map)) => l_map == r_map,
            (_, _) => false,
        }
    }
}

impl<L, D> Eq for PatriciaTreeMapAbstractPartition<L, D>
where
    L: Into<BitVec> + Clone,
    D: Sized + AbstractDomain,
{
}

impl<L, D> AbstractPartition<L, D> for PatriciaTreeMapAbstractPartition<L, D>
where
    L: Into<BitVec> + Clone,
    D: Sized + AbstractDomain,
{
    type ContainerType = PatriciaTreeMap<L, D>;

    fn bindings(&self) -> Option<&Self::ContainerType> {
        match self {
            Self::Value(ref map) => Some(map),
            _ => None,
        }
    }

    fn into_bindings(self) -> Option<Self::ContainerType> {
        match self {
            Self::Value(map) => Some(map),
            _ => None,
        }
    }

    fn len(&self) -> usize {
        self.bindings()
            .expect("Top abstract domain doesn't have a length!")
            .len()
    }

    fn is_empty(&self) -> bool {
        match self {
            Self::Top => false,
            Self::Value(map) => map.is_empty(),
        }
    }

    fn get(&self, label: &L) -> Cow<'_, D> {
        let map = match self {
            Self::Top => return Cow::Owned(D::top()),
            Self::Value(map) => map,
        };

        match map.get(label.clone()) {
            Some(domain) => Cow::Borrowed(domain),
            None => Cow::Owned(D::bottom()),
        }
    }

    fn set(&mut self, label: L, domain: D) {
        let map = match self {
            Self::Top => return,
            Self::Value(ref mut map) => map,
        };

        // Save some memory by implicitly storing bottom.
        if domain.is_bottom() {
            map.remove(label);
        } else {
            map.upsert(label, domain);
        }
    }

    fn update(&mut self, label: &L, op: impl FnOnce(&mut D)) {
        let map = match self {
            Self::Top => return,
            Self::Value(ref mut map) => map,
        };

        let mut update_domain = match map.get(label.clone()) {
            Some(domain) => domain.clone(),
            None => D::bottom(),
        };

        op(&mut update_domain);

        self.set(label.clone(), update_domain);
    }
}

impl<L, D> AbstractDomain for PatriciaTreeMapAbstractPartition<L, D>
where
    L: Into<BitVec> + Clone,
    D: Sized + AbstractDomain,
{
    fn bottom() -> Self {
        Self::Value(PatriciaTreeMap::new())
    }

    fn top() -> Self {
        Self::Top
    }

    fn is_bottom(&self) -> bool {
        match self {
            Self::Value(map) => map.is_empty(),
            _ => false,
        }
    }

    fn is_top(&self) -> bool {
        matches!(self, Self::Top)
    }

    fn leq(&self, rhs: &Self) -> bool {
        use PatriciaTreeMapAbstractPartition::*;

        match (self, rhs) {
            (Top, _) => rhs.is_top(),
            (_, Top) => true,
            (Value(self_map), Value(other_map)) => self_map.leq(other_map, &D::bottom()),
        }
    }

    fn join_with(&mut self, rhs: Self) {
        Self::join_like_operation(self, rhs, |d1, d2| d1.join_with(d2));
    }

    fn meet_with(&mut self, rhs: Self) {
        Self::meet_like_operation(self, rhs, |d1, d2| d1.meet_with(d2));
    }

    fn widen_with(&mut self, rhs: Self) {
        Self::join_like_operation(self, rhs, |d1, d2| d1.widen_with(d2));
    }

    fn narrow_with(&mut self, rhs: Self) {
        Self::meet_like_operation(self, rhs, |d1, d2| d1.narrow_with(d2));
    }
}

impl<L, D> PatriciaTreeMapAbstractPartition<L, D>
where
    L: Into<BitVec> + Clone,
    D: Sized + AbstractDomain,
{
    fn join_like_operation(lhs: &mut Self, rhs: Self, operation: impl Fn(&mut D, D)) {
        use PatriciaTreeMapAbstractPartition::*;

        match (lhs, rhs) {
            (Top, _) => {}
            (lhs, rhs @ Top) => {
                *lhs = rhs;
            }
            (Value(lmap), Value(rmap)) => {
                lmap.union_with(&rmap, |s, t| {
                    let mut s = s.clone();
                    operation(&mut s, t.clone());
                    s
                });
            }
        }
    }

    fn meet_like_operation(lhs: &mut Self, rhs: Self, operation: impl Fn(&mut D, D)) {
        use PatriciaTreeMapAbstractPartition::*;

        match (lhs, rhs) {
            (lhs @ Top, rhs) => {
                *lhs = rhs;
            }
            (_, Top) => {}
            (Value(lmap), Value(rmap)) => {
                lmap.intersect_with(&rmap, |s, t| {
                    let mut s = s.clone();
                    operation(&mut s, t.clone());
                    s
                });
            }
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::fmt::Debug;
use std::fmt::Formatter;
use std::fmt::Result;
use std::string::ToString;

const POINTER_SIZE_IN_BITS: usize = std::mem::size_of::<usize>() * 8;

union PointerOrBits {
    pointer: *mut usize, // TODO: When len > pointer size, allocate and store pointer here.
    bits: usize,
}

pub struct BitVec {
    actual_storage: PointerOrBits,
    len: usize,
}

impl Debug for BitVec {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
        let bits_str = format!("{:#b}", unsafe { self.actual_storage.bits });
        f.debug_struct("BitVec")
            .field("actual_storage", &bits_str)
            .field("len", &self.len)
            .finish()
    }
}

fn make_mask(len: usize) -> usize {
    // avoid integer overflow
    assert!(len <= POINTER_SIZE_IN_BITS);
    if len == POINTER_SIZE_IN_BITS {
        usize::MAX
    } else {
        (1 << len) - 1
    }
}

impl BitVec {
    pub fn new() -> Self {
        BitVec {
            actual_storage: PointerOrBits { bits: 0 },
            len: 0,
        }
    }

    pub fn from_int(v: usize) -> Self {
        BitVec {
            actual_storage: PointerOrBits { bits: v },
            len: POINTER_SIZE_IN_BITS,
        }
    }

    pub fn from_int_with_len(v: usize, len: usize) -> Self {
        assert!(len <= POINTER_SIZE_IN_BITS);

        let mask = make_mask(len);
        BitVec {
            actual_storage: PointerOrBits { bits: v & mask },
            len,
        }
    }

    pub fn get(&self, idx: usize) -> bool {
        assert!(idx < self.len, "idx: {}, len: {}", idx, self.len);

        let storage_unit;
        let intraunit_offset;

        if self.len <= POINTER_SIZE_IN_BITS {
            // Actual vector is stored in where it's supposed to be a pointer.
            storage_unit = unsafe { self.actual_storage.bits };
            intraunit_offset = idx;
        } else {
            todo!("Long bitvec not implemented");
        }

        let mask: usize = 1 << intraunit_offset;
        (storage_unit & mask) != 0
    }

    fn push(&mut self, val: bool) {
        let storage_unit: &mut usize;
        let intraunit_offset: usize;

        if self.len < POINTER_SIZE_IN_BITS {
            storage_unit = unsafe { &mut self.actual_storage.bits };
            intraunit_offset = self.len;
        } else {
            todo!("Long bitvec not implemented");
        }

        if val {
            *storage_unit |= 1 << intraunit_offset;
        } else {
            *storage_unit &= !(1 << intraunit_offset);
        }

        self.len += 1;
    }

    pub fn len(&self) -> usize {
        self.len
    }

    pub fn begins_with(&self, prefix: &BitVec) -> bool {
        if self.len() < prefix.len() {
            return false;
        }

        if self.len() <= POINTER_SIZE_IN_BITS {
            let compare_mask = make_mask(prefix.len());
            let compare_left = unsafe { self.actual_storage.bits & compare_mask };
            let compare_right = unsafe { prefix.actual_storage.bits & compare_mask };
            compare_left == compare_right
        } else {
            todo!("Long bitvec not implemented");
        }
    }

    pub fn common_prefix(v1: &BitVec, v2: &BitVec) -> BitVec {
        let min_len = std::cmp::min(v1.len(), v2.len());

        if min_len > POINTER_SIZE_IN_BITS {
            todo!("Long bitvec not implemented");
        }

        let cmp1 = unsafe { v1.actual_storage.bits };
        let cmp2 = unsafe { v2.actual_storage.bits };

        let mask = make_mask(min_len);

        let mut same_bits = (!(cmp1 ^ cmp2)) & mask;
        let mut count_prefix: usize = 0;

        while same_bits % 2 == 1 {
            same_bits >>= 1;
            count_prefix += 1;
        }

        assert!(
            count_prefix <= min_len,
            "count_prefix: {}, min_len: {}",
            count_prefix,
            min_len
        );

        BitVec::from_int_with_len(cmp1, count_prefix)
    }

    pub fn back(&self) -> bool {
        self.get(self.len() - 1)
    }
}

impl Default for BitVec {
    fn default() -> Self {
        Self::new()
    }
}

impl Clone for BitVec {
    fn clone(&self) -> Self {
        if self.len <= POINTER_SIZE_IN_BITS {
            Self {
                actual_storage: PointerOrBits {
                    bits: unsafe { self.actual_storage.bits },
                },
                len: self.len,
            }
        } else {
            todo!("Long bitvec not implemented");
        }
    }
}

impl ToString for BitVec {
    fn to_string(&self) -> String {
        let bits = unsafe { self.actual_storage.bits };
        format!("({} bits {:#b} aka {})", self.len, bits, bits)
    }
}

impl PartialEq for BitVec {
    fn eq(&self, other: &Self) -> bool {
        if self.len != other.len {
            return false;
        }

        if self.len > POINTER_SIZE_IN_BITS {
            todo!("Long bitvec not implemented");
        }

        let compare_left = unsafe { self.actual_storage.bits };
        let compare_right = unsafe { other.actual_storage.bits };
        compare_left == compare_right
    }
}

impl Eq for BitVec {}

macro_rules! impl_bitvec_convert_for_integral_type {
    ( $ ($type:ty), * ) => {
        $(
            impl From<$type> for BitVec {
                fn from(u: $type) -> BitVec {
                    const SIZE: usize = std::mem::size_of::<$type>() * 8;
                    BitVec::from_int_with_len(u as usize, SIZE)
                }
            }

            impl From<&BitVec> for $type {
                fn from(bv: &BitVec) -> $type {
                    const LIMIT: usize = std::mem::size_of::<$type>() * 8;
                    if bv.len() > LIMIT {
                        panic!("BitVec too long for $type")
                    } else {
                        let bits = unsafe { bv.actual_storage.bits };
                        (bits & make_mask(bv.len())) as $type
                    }
                }
            }
        )*
    };
}

impl_bitvec_convert_for_integral_type!(u8, u16, u32, u64, i8, i16, i32, i64, usize);

#[cfg(test)]
mod tests {

    use crate::datatype::bitvec::*;

    #[test]
    fn test_make_mask() {
        assert_eq!(make_mask(0), 0b0);
        assert_eq!(make_mask(1), 0b1);
        assert_eq!(make_mask(2), 0b11);
        assert_eq!(make_mask(3), 0b111);
        assert_eq!(make_mask(42), 0x3FFFFFFFFFF);
    }

    #[test]
    fn test_common_prefix() {
        let bv1 = BitVec::from_int_with_len(0b1010010, 8);
        let bv2 = BitVec::from_int_with_len(0b1110010, 8);
        let bv3 = BitVec::from_int_with_len(0b10010, 5);

        assert!(
            BitVec::common_prefix(&bv1, &bv2) == bv3,
            "bv1: {}, bv2: {}, bv3: {}",
            bv1.to_string(),
            bv2.to_string(),
            bv3.to_string()
        );
    }

    #[test]
    fn test_create_bitvec() {
        let bv: BitVec = BitVec::from(10);
        assert!(bv == bv);

        let bv8: BitVec = BitVec::from_int_with_len(0xFEFE, 8);
        let bv12: BitVec = BitVec::from_int_with_len(0xABFEFE, 12);

        assert!(bv12.begins_with(&bv8));
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::hash::Hash;

use im::HashSet;

use super::powerset::SetAbstractDomainOps;
use super::powerset::SetElementOps;

impl<T: Eq + Hash + Clone> SetAbstractDomainOps for HashSet<T> {
    fn is_subset(&self, other: &Self) -> bool {
        self.is_subset(other)
    }

    fn intersection_with(&mut self, other: &Self) {
        self.retain(|elem| other.contains(elem));
    }

    fn union_with(&mut self, other: Self) {
        other.into_iter().for_each(|elem| {
            self.insert(elem);
        })
    }
}

impl<T: Eq + Hash + Clone> SetElementOps for HashSet<T> {
    type Element = T;

    fn add_element(&mut self, e: T) {
        self.insert(e);
    }

    fn remove_element(&mut self, e: &T) {
        self.remove(e);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use super::abstract_domain::AbstractDomain;

#[derive(PartialEq, Eq, Clone, Debug)]
pub enum LiftedDomain<D: AbstractDomain> {
    Lifted(D),
    Bottom,
}

macro_rules! get_lowered {
    ( $self: ident ) => {
        match $self {
            Self::Lifted(underlying) => underlying,
            Self::Bottom => panic!("It's bottom and cannot be lowered!"),
        }
    };
}

impl<D: AbstractDomain> LiftedDomain<D> {
    fn lifted(underlying: D) -> Self {
        Self::Lifted(underlying)
    }

    fn is_lifted(&self) -> bool {
        matches!(self, Self::Lifted(_))
    }

    fn lowered(&self) -> &D {
        get_lowered!(self)
    }

    fn lowered_mut(&mut self) -> &mut D {
        get_lowered!(self)
    }

    fn into_lowered(self) -> D {
        get_lowered!(self)
    }
}

impl<D: AbstractDomain> AbstractDomain for LiftedDomain<D> {
    fn bottom() -> Self {
        Self::Bottom
    }

    fn top() -> Self {
        Self::Lifted(D::top())
    }

    fn is_bottom(&self) -> bool {
        matches!(self, Self::Bottom)
    }

    fn is_top(&self) -> bool {
        matches!(self, Self::Lifted(s) if s.is_top())
    }

    fn leq(&self, rhs: &Self) -> bool {
        match (self, rhs) {
            (Self::Bottom, _) => true,
            (_, Self::Bottom) => false,
            (Self::Lifted(s), Self::Lifted(t)) => s.leq(t),
        }
    }

    fn join_with(&mut self, rhs: Self) {
        match (self, rhs) {
            (x @ Self::Bottom, y) => *x = y,
            (_, Self::Bottom) => {}
            (Self::Lifted(s), Self::Lifted(t)) => s.join_with(t),
        }
    }

    fn meet_with(&mut self, rhs: Self) {
        match (self, rhs) {
            (Self::Bottom, _) => {}
            (x, y @ Self::Bottom) => *x = y,
            (Self::Lifted(s), Self::Lifted(t)) => s.meet_with(t),
        }
    }

    fn widen_with(&mut self, rhs: Self) {
        match (self, rhs) {
            (x @ Self::Bottom, y) => *x = y,
            (_, Self::Bottom) => {}
            (Self::Lifted(s), Self::Lifted(t)) => s.widen_with(t),
        }
    }

    fn narrow_with(&mut self, rhs: Self) {
        match (self, rhs) {
            (Self::Bottom, _) => {}
            (x, y @ Self::Bottom) => *x = y,
            (Self::Lifted(s), Self::Lifted(t)) => s.narrow_with(t),
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::datatype::lifted_domain::*;
    use crate::datatype::powerset::*;

    type UnderlyingDomain = HashSetAbstractDomain<i64>;
    type Domain = LiftedDomain<UnderlyingDomain>;

    #[test]
    fn test_basic_lift_ops() {
        let underlying_bottom = UnderlyingDomain::bottom();
        let underlying_top = UnderlyingDomain::top();
        let underlying_value1: UnderlyingDomain = vec![1, 2].into_iter().collect();
        let underlying_value2: UnderlyingDomain = vec![1, 2, 3].into_iter().collect();
        let underlying_value3: UnderlyingDomain = vec![1, 2, 4].into_iter().collect();
        let underlying_value4: UnderlyingDomain = vec![3, 4, 5].into_iter().collect();

        let bottom = Domain::bottom();
        let top = Domain::top();
        let mut lifted_bottom = Domain::lifted(underlying_bottom);
        let mut lifted_top = Domain::lifted(underlying_top);
        let mut lifted_value1 = Domain::lifted(underlying_value1);
        let mut lifted_value2 = Domain::lifted(underlying_value2);
        let lifted_value3 = Domain::lifted(underlying_value3);
        let lifted_value4 = Domain::lifted(underlying_value4);

        // Domain::Bottom is equal to Domain::bottom()
        assert!(bottom.leq(&Domain::Bottom));
        assert!(Domain::Bottom.leq(&bottom));
        assert_eq!(Domain::Bottom, bottom);

        // The lifted bottom is strictly greater than the bottom
        assert!(bottom.leq(&lifted_bottom));
        assert!(!lifted_bottom.leq(&bottom));
        assert_ne!(bottom, lifted_bottom);

        // The lifted top is equal to Domain::top()
        assert!(top.leq(&lifted_top));
        assert!(lifted_top.leq(&top));
        assert_eq!(top, lifted_top);

        // Check the ordering among the lifted values
        assert!(lifted_bottom.leq(&lifted_value1));
        assert!(lifted_value1.leq(&lifted_value2));
        assert!(lifted_value1.leq(&lifted_value3));

        assert!(!lifted_value2.leq(&lifted_value3));
        assert!(!lifted_value3.leq(&lifted_value2));

        assert!(lifted_value2.leq(&lifted_top));
        assert!(lifted_value3.leq(&lifted_top));

        // Test is_lifted
        assert!(lifted_bottom.is_lifted());
        assert!(lifted_top.is_lifted());
        assert!(lifted_value1.is_lifted());
        assert!(!bottom.is_lifted());

        // Test lowered
        assert_eq!(*lifted_bottom.lowered(), UnderlyingDomain::bottom());
        assert_eq!(*lifted_top.lowered(), UnderlyingDomain::top());
        assert_eq!(*lifted_value1.lowered(), vec![1, 2].into_iter().collect());

        // Test lowered_mut
        let underlying_val1_mut: &mut UnderlyingDomain = lifted_value1.lowered_mut();
        underlying_val1_mut.add_element(5);
        assert_eq!(*underlying_val1_mut, vec![1, 2, 5].into_iter().collect());

        // Test into_lowered
        let mut underlying_val1_into: UnderlyingDomain = lifted_value1.into_lowered();
        underlying_val1_into.add_element(6);
        assert_eq!(underlying_val1_into, vec![1, 2, 5, 6].into_iter().collect());
        // lifted_value1 cannot be used anymore after into_lowered

        // Test join_with
        lifted_bottom.join_with(Domain::top());
        assert_eq!(lifted_bottom, Domain::top());

        lifted_value2.join_with(lifted_value3);
        assert_eq!(
            *lifted_value2.lowered(),
            vec![1, 2, 3, 4].into_iter().collect()
        );

        // Test meet_with
        lifted_top.meet_with(Domain::bottom());
        assert_eq!(lifted_top, Domain::bottom());

        lifted_value2.meet_with(lifted_value4);
        assert_eq!(*lifted_value2.lowered(), vec![3, 4].into_iter().collect());
    }

    #[test]
    #[should_panic]
    fn test_bottom_lowered() {
        Domain::Bottom.lowered();
    }

    #[test]
    #[should_panic]
    fn test_bottom_lowered_mut() {
        Domain::Bottom.lowered_mut();
    }

    #[test]
    #[should_panic]
    fn test_bottom_into_lowered() {
        Domain::Bottom.into_lowered();
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

mod abstract_domain;
mod abstract_environment;
mod abstract_partition;
mod bitvec;
mod hash_set_impl;
mod lifted_domain;
mod patricia_tree_impl;
mod patricia_tree_map;
mod patricia_tree_set;
mod powerset;

pub use abstract_domain::*;
pub use abstract_environment::*;
pub use abstract_partition::*;
pub use hash_set_impl::*;
pub use lifted_domain::*;
pub use patricia_tree_map::*;
pub use patricia_tree_set::*;
pub use powerset::*;

extern crate sparta_proc_macros;
pub use sparta_proc_macros::DisjointUnion;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::rc::Rc;
use std::string::ToString;

use crate::datatype::bitvec::BitVec;
use crate::datatype::AbstractDomain;

/// This structure implements a map of integer/pointer keys to (possibly empty)
/// values. It's based on the following paper:
///
/// C. Okasaki, A. Gill. Fast Mergeable Integer Maps. In Workshop on ML (1998).

#[derive(Debug)]
enum Node<V> {
    Leaf {
        key: BitVec,
        value: V,
    },
    // Rust doesn't have higher kinded types for Rc.
    // Ideally we should reuse this enum for an Arc variant of it.
    Branch {
        prefix: BitVec,
        left: Rc<Node<V>>,
        right: Rc<Node<V>>,
    },
}

impl<V: Eq> PartialEq for Node<V> {
    fn eq(&self, other: &Self) -> bool {
        use Node::*;

        match (self, other) {
            (
                Leaf {
                    key: l_key,
                    value: l_value,
                },
                Leaf {
                    key: r_key,
                    value: r_value,
                },
            ) => l_key == r_key && l_value == r_value,
            (
                Branch {
                    prefix: l_prefix,
                    left: l_left,
                    right: l_right,
                },
                Branch {
                    prefix: r_prefix,
                    left: r_left,
                    right: r_right,
                },
            ) => {
                (l_prefix == r_prefix)
                    && ((Rc::ptr_eq(l_left, r_left) || l_left == r_left)
                        && (Rc::ptr_eq(l_right, r_right) || l_right == r_right))
            }
            (_, _) => false,
        }
    }
}

impl<V> ToString for Node<V> {
    fn to_string(&self) -> String {
        use Node::*;
        match self {
            Leaf { key, value: _ } => format!("(Leaf {})", key.to_string()),
            Branch {
                prefix,
                left,
                right,
            } => format!(
                "(Branch prefix: {} Left: {} Right: {})",
                prefix.to_string(),
                left.to_string(),
                right.to_string()
            ),
        }
    }
}

impl<V> Node<V> {
    /// Core algorithm for node insert, update, and removal.
    /// Returns: updated tree.
    ///
    /// `op` will be called in two separate occasions.
    /// - When node with `key` is found. Then `op` will be called with a `Some` value containing
    ///   the matching node. The entire matching subtree will be replaced by return value of `op`.
    ///   If `op` returned `None`, the entire subtree will be removed.
    /// - When node with `key` is not found. Then `op` will be called with a `None` value. The
    ///   return value of `op` is emplaced to the tree. If the return value of `op` is `None`, a
    ///   value equivalent to the original tree `maybe_node` is returned.
    fn update_node_by_key<F>(
        maybe_node: Option<Rc<Node<V>>>,
        key: &BitVec,
        op: F,
    ) -> Option<Rc<Node<V>>>
    where
        F: FnOnce(Option<Rc<Node<V>>>) -> Option<Rc<Node<V>>>,
    {
        use Node::*;

        if let Some(ref node) = maybe_node {
            match node.as_ref() {
                Leaf {
                    key: node_key,
                    value: _,
                } => {
                    if node_key == key {
                        op(maybe_node)
                    } else {
                        let maybe_new_node = op(None);
                        match maybe_new_node {
                            Some(new_node) => {
                                Some(Rc::new(Node::make_branch(new_node, node.clone())))
                            }
                            None => maybe_node,
                        }
                    }
                }
                Branch {
                    prefix,
                    left,
                    right,
                } => {
                    if key.begins_with(prefix) {
                        let branching_bit = key.get(prefix.len());
                        if !branching_bit {
                            let maybe_new_left =
                                Self::update_node_by_key(Some(left.clone()), key, op);
                            match maybe_new_left {
                                Some(new_left) => {
                                    // Possible optimization: If `new_left` ptr_eq to `left`, do nothing.
                                    Some(Rc::new(Node::make_branch(new_left, right.clone())))
                                }
                                None => Some(right.clone()),
                            }
                        } else {
                            let maybe_new_right =
                                Self::update_node_by_key(Some(right.clone()), key, op);
                            match maybe_new_right {
                                Some(new_right) => {
                                    Some(Rc::new(Node::make_branch(left.clone(), new_right)))
                                }
                                None => Some(left.clone()),
                            }
                        }
                    } else {
                        // Branch differs, create new branch like how you'd do with another leaf.
                        match op(None) {
                            Some(new_node) => {
                                Some(Rc::new(Node::make_branch(new_node, node.clone())))
                            }
                            None => maybe_node,
                        }
                    }
                }
            }
        } else {
            op(None)
        }
    }

    fn find_node_by_key<'a>(
        maybe_node: Option<&'a Rc<Node<V>>>,
        lookup_key: &BitVec,
    ) -> Option<&'a Rc<Node<V>>> {
        use Node::*;
        if let Some(node) = maybe_node {
            match node.as_ref() {
                Leaf { ref key, value: _ } => {
                    if key == lookup_key {
                        Some(node)
                    } else {
                        None
                    }
                }
                Branch {
                    ref prefix,
                    ref left,
                    ref right,
                } => {
                    if prefix.len() < lookup_key.len() {
                        if !lookup_key.get(prefix.len()) {
                            Self::find_node_by_key(Some(left), lookup_key)
                        } else {
                            Self::find_node_by_key(Some(right), lookup_key)
                        }
                    } else if prefix == lookup_key {
                        Some(node)
                    } else {
                        None
                    }
                }
            }
        } else {
            None
        }
    }

    fn find_leaf_by_key<'a>(
        maybe_node: Option<&'a Rc<Node<V>>>,
        lookup_key: &BitVec,
    ) -> Option<&'a Rc<Node<V>>> {
        if let Some(found_node) = Self::find_node_by_key(maybe_node, lookup_key) {
            return match found_node.as_ref() {
                Node::Leaf { key: _, value: _ } => Some(found_node),
                _ => None,
            };
        }
        None
    }

    fn contains_leaf_with_key(maybe_node: Option<&Rc<Node<V>>>, lookup_key: &BitVec) -> bool {
        Self::find_leaf_by_key(maybe_node, lookup_key).is_some()
    }

    fn key_or_prefix(&self) -> &BitVec {
        use Node::*;
        match self {
            Leaf { key, value: _ } => key,
            Branch {
                prefix,
                left: _,
                right: _,
            } => prefix,
        }
    }

    fn make_branch(one: Rc<Self>, other: Rc<Self>) -> Self {
        let v1 = one.key_or_prefix();
        let v2 = other.key_or_prefix();
        assert!(v1 != v2);
        let common = BitVec::common_prefix(v1, v2);
        let branching_bit = common.len();

        let b1 = v1.get(branching_bit);
        let b2 = v2.get(branching_bit);
        assert!(b1 != b2);

        let left;
        let right;

        if !b1 {
            left = one;
            right = other;
        } else {
            left = other;
            right = one;
        }

        Node::Branch {
            prefix: common,
            left,
            right,
        }
    }

    fn combine_leaves_by_key(
        node: Rc<Node<V>>,
        key: &BitVec,
        other: Rc<Node<V>>,
        leaf_combine: &impl Fn(Rc<Node<V>>, Rc<Node<V>>) -> Option<Rc<Node<V>>>,
    ) -> Rc<Node<V>> {
        let updated = Self::update_node_by_key(Some(node), key, move |leaf| match leaf {
            Some(leaf) => leaf_combine(leaf, other),
            None => Some(other),
        });
        updated.unwrap() // Leaf combine should not make deletions.
    }

    /// Merge two trees. Combined tree should contain all keys from s and t.
    /// If duplicate keys are found, two nodes are passed to `leaf_combine` shall be called
    /// with values from s on the left hand side and values from t on the right hand side.
    fn merge_trees(
        s: &Rc<Node<V>>,
        t: &Rc<Node<V>>,
        leaf_combine: &impl Fn(Rc<Node<V>>, Rc<Node<V>>) -> Option<Rc<Node<V>>>,
    ) -> Rc<Node<V>> {
        use Node::*;

        if Rc::ptr_eq(s, t) {
            // Quickly checking if the two trees are identical to allow union
            // operation to complete in sublinear time when operands share some structures.
            return s.clone();
        }

        match (s.as_ref(), t.as_ref()) {
            // We check if t is a leaf first, before s. Because the leaf combine
            // operator may not be commutative when s and t are both leaves.
            (
                _,
                Leaf {
                    key: t_key,
                    value: _,
                },
            ) => {
                // Insert t into s, where s may or may not be a leaf.
                // If t has the same key as one of the element in s, t is rhs of the combine operator.
                Self::combine_leaves_by_key(s.clone(), t_key, t.clone(), leaf_combine)
            }
            (
                Leaf {
                    key: s_key,
                    value: _,
                },
                _,
            ) => {
                // Insert s into t
                Self::combine_leaves_by_key(t.clone(), s_key, s.clone(), leaf_combine)
            }
            (
                Branch {
                    prefix: s_prefix,
                    left: s_left,
                    right: s_right,
                },
                Branch {
                    prefix: t_prefix,
                    left: t_left,
                    right: t_right,
                },
            ) => {
                if s_prefix == t_prefix {
                    // The two trees have the same prefix. We just merge the subtrees.
                    let new_left = Self::merge_trees(s_left, t_left, leaf_combine);
                    let new_right = Self::merge_trees(s_right, t_right, leaf_combine);

                    if Rc::ptr_eq(&new_left, s_left) && Rc::ptr_eq(&new_right, s_right) {
                        s.clone()
                    } else if Rc::ptr_eq(&new_left, t_left) && Rc::ptr_eq(&new_right, t_right) {
                        t.clone()
                    } else {
                        Rc::new(Node::make_branch(new_left, new_right))
                    }
                } else if t_prefix.begins_with(s_prefix) {
                    let branching_bit = t_prefix.get(s_prefix.len());
                    if !branching_bit {
                        let new_left = Self::merge_trees(s_left, t, leaf_combine);
                        if Rc::ptr_eq(s_left, &new_left) {
                            s.clone()
                        } else {
                            Rc::new(Node::make_branch(new_left, s_right.clone()))
                        }
                    } else {
                        let new_right = Self::merge_trees(s_right, t, leaf_combine);
                        if Rc::ptr_eq(s_right, &new_right) {
                            s.clone()
                        } else {
                            Rc::new(Node::make_branch(s_left.clone(), new_right))
                        }
                    }
                } else if s_prefix.begins_with(t_prefix) {
                    let branching_bit = s_prefix.get(t_prefix.len());
                    if !branching_bit {
                        let new_left = Self::merge_trees(s, t_left, leaf_combine);
                        if Rc::ptr_eq(t_left, &new_left) {
                            t.clone()
                        } else {
                            Rc::new(Node::make_branch(new_left, t_right.clone()))
                        }
                    } else {
                        let new_right = Self::merge_trees(s, t_right, leaf_combine);
                        if Rc::ptr_eq(t_right, &new_right) {
                            t.clone()
                        } else {
                            Rc::new(Node::make_branch(t_left.clone(), new_right))
                        }
                    }
                } else {
                    // The prefixes disagree.
                    Rc::new(Node::make_branch(s.clone(), t.clone()))
                }
            }
        }
    }

    /// Intersection of two trees. Combined tree should contain all keys that are found from both
    /// s and t. If duplicate keys are found, two nodes are passed to `leaf_combine` shall be
    /// called with values from s on the left hand side and values from t on the right hand side.
    fn intersect_trees(
        s: &Rc<Node<V>>,
        t: &Rc<Node<V>>,
        leaf_combine: &impl Fn(Rc<Node<V>>, Rc<Node<V>>) -> Option<Rc<Node<V>>>,
    ) -> Option<Rc<Node<V>>> {
        use Node::*;

        if Rc::ptr_eq(s, t) {
            // This conditions allows the inclusion test to run in sublinear time
            // when comparing Patricia trees that share some structure.
            return Some(s.clone());
        }

        match (s.as_ref(), t.as_ref()) {
            (Leaf { key, value: _ }, _) => match Self::find_leaf_by_key(Some(t), key) {
                // Keep leaves from s on the left.
                Some(t_leaf) => leaf_combine(s.clone(), t_leaf.clone()),
                None => None,
            },
            (_, Leaf { key, value: _ }) => match Self::find_leaf_by_key(Some(s), key) {
                // Keep leaves from s on the left.
                Some(s_leaf) => leaf_combine(s_leaf.clone(), t.clone()),
                None => None,
            },
            (
                Branch {
                    prefix: s_prefix,
                    left: s_left,
                    right: s_right,
                },
                Branch {
                    prefix: t_prefix,
                    left: t_left,
                    right: t_right,
                },
            ) => {
                if s_prefix == t_prefix {
                    let new_left = Self::intersect_trees(s_left, t_left, leaf_combine);
                    let new_right = Self::intersect_trees(s_right, t_right, leaf_combine);
                    match (new_left, new_right) {
                        (left, None) => left,
                        (None, right) => right,
                        (Some(left), Some(right)) => Some(Rc::new(Self::make_branch(left, right))),
                    }
                } else if t_prefix.begins_with(s_prefix) {
                    let branching_bit = t_prefix.get(s_prefix.len());
                    Self::intersect_trees(
                        if !branching_bit { s_left } else { s_right },
                        t,
                        leaf_combine,
                    )
                } else if s_prefix.begins_with(t_prefix) {
                    let branching_bit = s_prefix.get(t_prefix.len());
                    Self::intersect_trees(
                        s,
                        if !branching_bit { t_left } else { t_right },
                        leaf_combine,
                    )
                } else {
                    None
                }
            }
        }
    }

    /// Returns true if s is a subset of t.
    fn is_tree_subset_of(s: &Rc<Node<V>>, t: &Rc<Node<V>>) -> bool {
        use Node::*;

        if Rc::ptr_eq(s, t) {
            // This conditions allows the inclusion test to run in sublinear time
            // when comparing Patricia trees that share some structure.
            return true;
        }

        match (s.as_ref(), t.as_ref()) {
            (Leaf { key, value: _ }, _) => Self::contains_leaf_with_key(Some(t), key),
            (_, Leaf { key: _, value: _ }) => false,
            (
                Branch {
                    prefix: s_prefix,
                    left: s_left,
                    right: s_right,
                },
                Branch {
                    prefix: t_prefix,
                    left: t_left,
                    right: t_right,
                },
            ) => {
                if s_prefix == t_prefix {
                    Self::is_tree_subset_of(s_left, t_left)
                        && Self::is_tree_subset_of(s_right, t_right)
                } else if s_prefix.begins_with(t_prefix) {
                    assert!(s_prefix.len() > t_prefix.len());
                    let branching_bit = s_prefix.get(t_prefix.len());
                    if !branching_bit {
                        Self::is_tree_subset_of(s_left, t_left)
                            && Self::is_tree_subset_of(s_right, t_left)
                    } else {
                        Self::is_tree_subset_of(s_left, t_right)
                            && Self::is_tree_subset_of(s_right, t_right)
                    }
                } else {
                    false
                }
            }
        }
    }
}

impl<D: AbstractDomain> Node<D> {
    fn is_tree_leq(s: Option<&Rc<Node<D>>>, t: Option<&Rc<Node<D>>>, implicit_value: &D) -> bool {
        match (s, t) {
            (None, None) => true,
            (None, _) => implicit_value.is_bottom(),
            (_, None) => implicit_value.is_top(),
            (Some(s), Some(t)) => Self::is_tree_leq_impl(s, t, implicit_value),
        }
    }

    fn is_tree_leq_impl(s: &Rc<Node<D>>, t: &Rc<Node<D>>, implicit_value: &D) -> bool {
        use Node::*;

        if Rc::ptr_eq(s, t) {
            return true;
        }

        match (s.as_ref(), t.as_ref()) {
            (
                Leaf {
                    key: s_key,
                    value: s_value,
                },
                Leaf {
                    key: t_key,
                    value: t_value,
                },
            ) => s_key == t_key && s_value.leq(t_value),
            (
                Leaf {
                    key: s_key,
                    value: s_value,
                },
                _,
            ) => {
                if implicit_value.is_top() {
                    false
                } else {
                    match Self::find_leaf_by_key(Some(t), s_key) {
                        Some(rc) => match rc.as_ref() {
                            Leaf {
                                key: _,
                                value: t_value,
                            } => s_value.leq(t_value),
                            _ => unreachable!(),
                        },
                        None => false,
                    }
                }
            }
            (
                _,
                Leaf {
                    key: t_key,
                    value: t_value,
                },
            ) => {
                if implicit_value.is_bottom() {
                    false
                } else {
                    match Self::find_leaf_by_key(Some(s), t_key) {
                        Some(rc) => match rc.as_ref() {
                            Leaf {
                                key: _,
                                value: s_value,
                            } => s_value.leq(t_value),
                            _ => unreachable!(),
                        },
                        None => false,
                    }
                }
            }
            (
                Branch {
                    prefix: s_prefix,
                    left: s_left,
                    right: s_right,
                },
                Branch {
                    prefix: t_prefix,
                    left: t_left,
                    right: t_right,
                },
            ) => {
                if s_prefix == t_prefix {
                    // The two trees have the same prefix, compare each subtrees.
                    Self::is_tree_leq_impl(s_left, t_left, implicit_value)
                        && Self::is_tree_leq_impl(s_right, t_right, implicit_value)
                } else if s_prefix.begins_with(t_prefix) {
                    // The tree t only contains bindings present in a subtree of s, and s has bindings not present in t.
                    let branching_bit = s_prefix.get(t_prefix.len());
                    implicit_value.is_top()
                        && Self::is_tree_leq_impl(
                            if !branching_bit { s_left } else { s_right },
                            t,
                            implicit_value,
                        )
                } else if t_prefix.begins_with(s_prefix) {
                    let branching_bit = t_prefix.get(s_prefix.len());
                    implicit_value.is_bottom()
                        && Self::is_tree_leq_impl(
                            s,
                            if !branching_bit { t_left } else { t_right },
                            implicit_value,
                        )
                } else {
                    false
                }
            }
        }
    }
}

// Create an interface that gives the user a "mutable" illusion of an immutable data structure.
#[derive(Debug)]
pub(crate) struct PatriciaTree<V> {
    root: Option<Rc<Node<V>>>,
}

impl<V> PatriciaTree<V> {
    pub(crate) fn new() -> Self {
        Self { root: None }
    }

    pub(crate) fn clear(&mut self) {
        self.root = None;
    }

    pub(crate) fn is_empty(&self) -> bool {
        matches!(self.root, None)
    }

    // Not a very fast operation.
    pub(crate) fn len(&self) -> usize {
        self.iter().count()
    }

    fn apply_root_operation<F>(&mut self, op: F)
    where
        F: FnOnce(Option<Rc<Node<V>>>) -> Option<Rc<Node<V>>>,
    {
        let mut temp_root = None;
        std::mem::swap(&mut self.root, &mut temp_root);
        self.root = op(temp_root);
    }

    pub(crate) fn insert(&mut self, key: BitVec, value: V) {
        let new_leaf = Rc::new(Node::Leaf {
            key: key.clone(),
            value,
        });
        let node_op = move |_| Some(new_leaf);
        let root_op = move |root| Node::update_node_by_key(root, &key, node_op);
        self.apply_root_operation(root_op);
    }

    pub(crate) fn contains_key(&self, key: &BitVec) -> bool {
        !matches!(self.get(key), None)
    }

    pub(crate) fn get(&self, key: &BitVec) -> Option<&V> {
        use Node::*;
        let node = Node::find_leaf_by_key(self.root.as_ref(), key);
        match node {
            Some(leaf_node) => match leaf_node.as_ref() {
                Leaf { key: _, ref value } => Some(value),
                _ => panic!("Did not correctly get a leaf!"),
            },
            None => None,
        }
    }

    pub(crate) fn remove(&mut self, key: &BitVec) {
        let root_op = move |root| Node::update_node_by_key(root, key, |_| None);
        self.apply_root_operation(root_op);
    }

    pub(crate) fn iter(&self) -> PatriciaTreePostOrderIterator<V> {
        PatriciaTreePostOrderIterator::<V>::from_tree(self)
    }

    fn get_leaf_combine_with_value_op_semantics(
        value_op_on_duplicate_key: impl Fn(&V, &V) -> V,
    ) -> impl Fn(Rc<Node<V>>, Rc<Node<V>>) -> Option<Rc<Node<V>>> {
        use Node::*;

        move |one_leaf: Rc<Node<V>>, other_leaf: Rc<Node<V>>| match (
            one_leaf.as_ref(),
            other_leaf.as_ref(),
        ) {
            (
                Leaf {
                    key,
                    value: ref l_value,
                },
                Leaf {
                    key: _,
                    value: ref r_value,
                },
            ) => Some(Rc::new(Leaf {
                key: key.clone(),
                value: value_op_on_duplicate_key(l_value, r_value),
            })),
            _ => panic!("leaf_combine should only be called on leaves!"),
        }
    }

    pub(crate) fn union_with(
        &mut self,
        other: &Self,
        value_op_on_duplicate_key: impl Fn(&V, &V) -> V,
    ) {
        match (self.root.as_ref(), other.root.as_ref()) {
            (None, _) => self.root = other.root.clone(),
            (Some(_), None) => {}
            (Some(self_node), Some(other_node)) => {
                self.root = Some(Node::merge_trees(
                    self_node,
                    other_node,
                    &Self::get_leaf_combine_with_value_op_semantics(value_op_on_duplicate_key),
                ));
            }
        }
    }

    pub(crate) fn intersect_with(
        &mut self,
        other: &Self,
        value_op_on_duplicate_key: impl Fn(&V, &V) -> V,
    ) {
        match (self.root.as_ref(), other.root.as_ref()) {
            (None, _) => {}
            (Some(_), None) => {
                self.root = None;
            }
            (Some(self_node), Some(other_node)) => {
                self.root = Node::intersect_trees(
                    self_node,
                    other_node,
                    &Self::get_leaf_combine_with_value_op_semantics(value_op_on_duplicate_key),
                );
            }
        }
    }

    pub(crate) fn subset_of(&self, other: &Self) -> bool {
        match (self.root.as_ref(), other.root.as_ref()) {
            (None, _) => true,
            (_, None) => false,
            (Some(s), Some(t)) => Node::is_tree_subset_of(s, t),
        }
    }
}

impl<D: AbstractDomain> PatriciaTree<D> {
    // Leq operation here facilitates the abstract partition and abstract environment leqs.
    // Both structures have implicit values that they use to encode top and bottom values
    // efficiently, which changes the semantics of the leq operation.
    pub(crate) fn leq(&self, other: &Self, implicit_value: &D) -> bool {
        Node::is_tree_leq(self.root.as_ref(), other.root.as_ref(), implicit_value)
    }
}

impl<V> Clone for PatriciaTree<V> {
    fn clone(&self) -> Self {
        // Cloning a patricia tree is just cloning the shared reference to the root node.
        Self {
            root: self.root.clone(),
        }
    }
}

impl<V: Eq> PartialEq for PatriciaTree<V> {
    fn eq(&self, other: &Self) -> bool {
        match (self.root.as_ref(), other.root.as_ref()) {
            (None, None) => true,
            (Some(ref self_node), Some(ref other_node)) => {
                Rc::ptr_eq(self_node, other_node) || (self_node == other_node)
            }
            (_, _) => false,
        }
    }
}

impl<V: Eq> Eq for PatriciaTree<V> {}

#[derive(Debug)]
pub(crate) struct PatriciaTreePostOrderIterator<'a, V> {
    branch_stack: Vec<&'a Node<V>>,
    current: Option<&'a Node<V>>,
}

impl<'a, V> PatriciaTreePostOrderIterator<'a, V> {
    pub(crate) fn from_tree(tree: &'a PatriciaTree<V>) -> Self {
        let mut ret = Self {
            branch_stack: vec![],
            current: None,
        };

        match tree.root {
            Some(ref node) => ret.next_leaf(node),
            None => (),
        };

        ret
    }

    fn next_leaf(&mut self, subtree: &'a Rc<Node<V>>) {
        let mut node = subtree.as_ref();

        while let Node::Branch {
            prefix: _,
            left,
            right: _,
        } = node
        {
            self.branch_stack.push(node);
            node = left.as_ref();
        }

        // node is a leaf now.
        self.current = Some(node);
    }

    fn next_node(&mut self) -> Option<&'a Node<V>> {
        let ret = self.current;
        self.current = None;

        if let Some(br) = self.branch_stack.pop() {
            match br {
                Node::Branch {
                    prefix: _,
                    left: _,
                    ref right,
                } => self.next_leaf(right),
                _ => panic!("Malformed Patricia Tree Iterator"),
            }
        }

        ret
    }

    fn into_tuple(node: Option<&Node<V>>) -> Option<(&BitVec, &V)> {
        match node {
            Some(leaf) => match leaf {
                Node::Leaf { ref key, ref value } => Some((key, value)),
                _ => panic!("Malformed Patricia Tree Iterator"),
            },
            None => None,
        }
    }
}

impl<'a, V> Iterator for PatriciaTreePostOrderIterator<'a, V> {
    type Item = (&'a BitVec, &'a V);

    fn next(&mut self) -> Option<Self::Item> {
        Self::into_tuple(self.next_node())
    }
}

#[cfg(test)]
mod tests {
    use crate::datatype::patricia_tree_impl::*;

    #[test]
    fn test_basic_insertion() {
        let mut map: PatriciaTree<usize> = PatriciaTree::new();
        map.insert(1.into(), 111);
        map.insert(22.into(), 222);
        map.insert(42.into(), 444);
        map.insert(42.into(), 444);
        map.insert(42.into(), 444);
        map.insert(13.into(), 1313);

        assert!(map.contains_key(&1.into()));
        assert!(map.contains_key(&22.into()));
        assert!(map.contains_key(&42.into()));
        assert!(!map.contains_key(&2.into()));
        assert!(!map.contains_key(&3.into()));

        assert_eq!(map.len(), 4);

        let mut map2 = map.clone();

        assert!(map2 == map);

        map2.insert(55.into(), 555);
        assert_eq!(map.len(), 4);
        assert_eq!(map2.len(), 5);

        assert!(map2 != map);

        map2.remove(&55.into());
        assert_eq!(map2.len(), 4);
        assert!(map2 == map);

        map2.remove(&1.into());
        assert!(map.contains_key(&1.into()));
        assert!(!map2.contains_key(&1.into()));

        map2.remove(&1.into());
        assert_eq!(map2.len(), 3);
        assert!(!map2.contains_key(&1.into()));

        map2.remove(&22.into());
        assert_eq!(map2.len(), 2);
        assert!(!map2.contains_key(&22.into()));

        map2.remove(&13.into());
        assert_eq!(map2.len(), 1);
        assert!(!map2.contains_key(&13.into()));

        map2.remove(&42.into());
        assert_eq!(map2.len(), 0);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::iter::FromIterator;
use std::iter::Iterator;
use std::marker::PhantomData;

use crate::datatype::bitvec::BitVec;
use crate::datatype::patricia_tree_impl::PatriciaTree;
use crate::datatype::patricia_tree_impl::PatriciaTreePostOrderIterator;
use crate::datatype::AbstractDomain;

// Interface structs for PatriciaTreeMap. Does not require V to impl Clone.
#[derive(Debug)]
pub struct PatriciaTreeMap<K: Into<BitVec>, V> {
    storage: PatriciaTree<V>,
    _key_type_phantom: PhantomData<K>,
}

impl<K: Into<BitVec>, V> PatriciaTreeMap<K, V> {
    pub fn new() -> Self {
        Self {
            storage: PatriciaTree::<V>::new(),
            _key_type_phantom: Default::default(),
        }
    }

    pub fn clear(&mut self) {
        self.storage.clear();
    }

    pub fn is_empty(&self) -> bool {
        self.storage.is_empty()
    }

    // Not a very fast operation.
    pub fn len(&self) -> usize {
        self.storage.len()
    }

    pub fn upsert(&mut self, key: K, value: V) {
        self.storage.insert(key.into(), value)
    }

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

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

    pub fn remove(&mut self, key: K) {
        self.storage.remove(&key.into())
    }

    pub fn iter(&self) -> PatriciaTreeMapIterator<'_, K, V> {
        self.storage.iter().into()
    }
}

impl<K: Into<BitVec>, V> Default for PatriciaTreeMap<K, V> {
    fn default() -> Self {
        Self::new()
    }
}

impl<K: Into<BitVec>, V: Eq> PartialEq for PatriciaTreeMap<K, V> {
    fn eq(&self, other: &Self) -> bool {
        self.storage.eq(&other.storage)
    }
}

impl<K: Into<BitVec>, V: Eq> Eq for PatriciaTreeMap<K, V> {}

impl<K: Into<BitVec>, V> Clone for PatriciaTreeMap<K, V> {
    fn clone(&self) -> Self {
        Self {
            storage: self.storage.clone(),
            ..Default::default()
        }
    }
}

#[derive(Debug)]
pub struct PatriciaTreeMapIterator<'a, K: Into<BitVec>, V> {
    iter_impl: PatriciaTreePostOrderIterator<'a, V>,
    _key_type_phantom: PhantomData<K>,
}

impl<'a, K: Into<BitVec>, V> From<PatriciaTreePostOrderIterator<'a, V>>
    for PatriciaTreeMapIterator<'a, K, V>
{
    fn from(iter_impl: PatriciaTreePostOrderIterator<'a, V>) -> Self {
        Self {
            iter_impl,
            _key_type_phantom: Default::default(),
        }
    }
}

impl<'a, K: 'a + Into<BitVec> + From<&'a BitVec>, V> Iterator
    for PatriciaTreeMapIterator<'a, K, V>
{
    type Item = (K, &'a V);

    fn next(&mut self) -> Option<Self::Item> {
        let wrapped = self.iter_impl.next();
        wrapped.map(|(key, value)| (key.into(), value))
    }
}

impl<K: Into<BitVec>, V> FromIterator<(K, V)> for PatriciaTreeMap<K, V> {
    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
        let mut ret: PatriciaTreeMap<K, V> = PatriciaTreeMap::new();
        for (k, v) in iter {
            ret.upsert(k, v);
        }
        ret
    }
}

impl<'a, K: Into<BitVec> + From<&'a BitVec>, V> IntoIterator for &'a PatriciaTreeMap<K, V> {
    type Item = (K, &'a V);
    type IntoIter = PatriciaTreeMapIterator<'a, K, V>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

// Specialized Leq for abstract partitions and environments when D is an abstract domain.
impl<K: Into<BitVec>, D: AbstractDomain> PatriciaTreeMap<K, D> {
    pub(crate) fn leq(&self, other: &Self, implicit_value: &D) -> bool {
        self.storage.leq(&other.storage, implicit_value)
    }

    pub(crate) fn union_with(
        &mut self,
        other: &Self,
        value_op_on_duplicate_key: impl Fn(&D, &D) -> D,
    ) {
        self.storage
            .union_with(&other.storage, value_op_on_duplicate_key)
    }

    pub(crate) fn intersect_with(
        &mut self,
        other: &Self,
        value_op_on_duplicate_key: impl Fn(&D, &D) -> D,
    ) {
        self.storage
            .intersect_with(&other.storage, value_op_on_duplicate_key)
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use crate::datatype::PatriciaTreeMap;

    #[test]
    fn test_map_operations() {
        let num_vals: usize = 10000;

        let mut items: HashMap<usize, String> = HashMap::new();

        for i in 0..num_vals {
            items.insert(i, format!("{}", i));
        }

        let map: PatriciaTreeMap<_, _> = items.clone().into_iter().collect();

        assert_eq!(map.len(), 10000);
        assert_eq!(*map.get(15).unwrap(), "15".to_string());
        assert_eq!(*map.get(9999).unwrap(), "9999".to_string());
        assert_eq!(map.get(10000), None);

        assert!(map == map);
        let mut map2 = map.clone();
        assert!(map == map2);
        map2.upsert(0xabcdef12, "0xabcdef12".to_owned());
        assert!(map != map2);
        map2.remove(0xabcdef12);
        assert!(map == map2)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::iter::FromIterator;
use std::iter::Iterator;
use std::marker::PhantomData;

use super::powerset::SetAbstractDomainOps;
use super::powerset::SetElementOps;
use crate::datatype::bitvec::BitVec;
use crate::datatype::patricia_tree_impl::PatriciaTree;
use crate::datatype::patricia_tree_impl::PatriciaTreePostOrderIterator;

// Interface structs for PatriciaTree
#[derive(Debug)]
pub struct PatriciaTreeSet<K: Into<BitVec>> {
    storage: PatriciaTree<()>,
    _key_type_phantom: PhantomData<K>,
}

impl<K: Into<BitVec>> PatriciaTreeSet<K> {
    pub fn new() -> Self {
        Self {
            storage: PatriciaTree::<()>::new(),
            _key_type_phantom: Default::default(),
        }
    }

    pub fn clear(&mut self) {
        self.storage.clear();
    }

    pub fn is_empty(&self) -> bool {
        self.storage.is_empty()
    }

    // Not a very fast operation.
    pub fn len(&self) -> usize {
        self.storage.len()
    }

    pub fn insert(&mut self, key: K) {
        self.storage.insert(key.into(), ())
    }

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

    pub fn remove(&mut self, key: K) {
        self.storage.remove(&key.into())
    }

    pub fn iter(&self) -> PatriciaTreeSetIterator<'_, K> {
        self.storage.iter().into()
    }
}

impl<K: Into<BitVec>> Default for PatriciaTreeSet<K> {
    fn default() -> Self {
        Self::new()
    }
}

impl<K: Into<BitVec>> Clone for PatriciaTreeSet<K> {
    fn clone(&self) -> Self {
        Self {
            storage: self.storage.clone(),
            ..Default::default()
        }
    }
}

impl<K: Into<BitVec>> PartialEq for PatriciaTreeSet<K> {
    fn eq(&self, other: &Self) -> bool {
        self.storage.eq(&other.storage)
    }
}

impl<K: Into<BitVec>> Eq for PatriciaTreeSet<K> {}

#[derive(Debug)]
pub struct PatriciaTreeSetIterator<'a, K: Into<BitVec>> {
    iter_impl: PatriciaTreePostOrderIterator<'a, ()>,
    _key_type_phantom: PhantomData<K>,
}

impl<'a, K: Into<BitVec>> From<PatriciaTreePostOrderIterator<'a, ()>>
    for PatriciaTreeSetIterator<'a, K>
{
    fn from(iter_impl: PatriciaTreePostOrderIterator<'a, ()>) -> Self {
        Self {
            iter_impl,
            _key_type_phantom: Default::default(),
        }
    }
}

impl<'a, K: 'a + Into<BitVec> + From<&'a BitVec>> Iterator for PatriciaTreeSetIterator<'a, K> {
    type Item = K;

    fn next(&mut self) -> Option<Self::Item> {
        let wrapped = self.iter_impl.next();
        wrapped.map(|(key, _)| key.into())
    }
}

impl<K: Into<BitVec>> FromIterator<K> for PatriciaTreeSet<K> {
    fn from_iter<I: IntoIterator<Item = K>>(iter: I) -> Self {
        let mut ret: PatriciaTreeSet<K> = PatriciaTreeSet::new();
        for item in iter {
            ret.insert(item);
        }
        ret
    }
}

impl<'a, K: Into<BitVec> + From<&'a BitVec>> IntoIterator for &'a PatriciaTreeSet<K> {
    type Item = K;
    type IntoIter = PatriciaTreeSetIterator<'a, K>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

impl<K> SetAbstractDomainOps for PatriciaTreeSet<K>
where
    K: Into<BitVec> + Clone,
{
    fn is_subset(&self, other: &Self) -> bool {
        self.storage.subset_of(&other.storage)
    }

    fn intersection_with(&mut self, other: &Self) {
        self.storage.intersect_with(&other.storage, |_, _| ())
    }

    fn union_with(&mut self, other: Self) {
        self.storage.union_with(&other.storage, |_, _| ())
    }
}

impl<K: Into<BitVec> + Clone> SetElementOps for PatriciaTreeSet<K> {
    type Element = K;

    fn add_element(&mut self, e: K) {
        self.insert(e);
    }

    fn remove_element(&mut self, e: &K) {
        self.remove(e.clone());
    }
}

impl<K, const N: usize> From<[K; N]> for PatriciaTreeSet<K>
where
    K: Into<BitVec> + Clone,
{
    fn from(arr: [K; N]) -> Self {
        arr.into_iter().collect()
    }
}

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

    use rand::Rng;

    use crate::datatype::powerset::SetAbstractDomainOps;
    use crate::datatype::PatriciaTreeSet;

    #[test]
    fn test_set_operations() {
        let num_vals: u32 = 10000;

        let items: HashSet<u32> = (0..num_vals).collect();
        let mut set: PatriciaTreeSet<u32> = items.iter().cloned().collect();
        let out_items: HashSet<u32> = set.iter().collect();

        assert_eq!(items, out_items);
        assert_eq!(set.len(), 10000);
        assert!(!set.contains(10001));

        set.insert(200); // Nothing bad happens
        set.insert(10200); // Nothing bad happens

        assert_eq!(set.len(), 10001);

        assert!(set == set);
        let mut set2 = set.clone();
        assert!(set == set2);
        set2.insert(0xabcdef12);
        assert!(set != set2);
        set2.remove(0xabcdef12);
        assert!(set == set2)
    }

    #[test]
    fn test_robustness_with_rng() {
        let mut rng = rand::thread_rng();

        let mut items: HashSet<u32> = HashSet::new();
        let mut set: PatriciaTreeSet<u32> = PatriciaTreeSet::new();

        let num_vals = 10000;
        for _ in 0..num_vals {
            let i: u32 = rng.gen();
            items.insert(i);
            set.insert(i);
        }

        let out_items: HashSet<u32> = set.iter().collect();

        for item in items.iter() {
            assert!(set.contains(*item));
        }

        assert_eq!(items, out_items);
    }

    #[test]
    fn test_union_operation_empty() {
        let mut set1 = PatriciaTreeSet::<u32>::from([1, 2, 3]);
        let set2 = PatriciaTreeSet::<u32>::from([]);
        set1.union_with(set2);

        let actual: HashSet<_> = set1.iter().collect();
        let expected = HashSet::<u32>::from([1, 2, 3]);

        assert_eq!(actual, expected);
    }

    #[test]
    fn test_simple_union_operation() {
        let mut set1 = PatriciaTreeSet::<u32>::from([1, 2, 3, 64, 99]);
        let set2 = PatriciaTreeSet::<u32>::from([4, 5, 6, 65, 102]);
        set1.union_with(set2);

        let actual: HashSet<_> = set1.iter().collect();
        let expected = HashSet::<u32>::from([1, 2, 3, 4, 5, 6, 64, 65, 99, 102]);

        assert_eq!(actual, expected);
    }

    #[test]
    fn test_union_operation_with_dup() {
        let mut set1 = PatriciaTreeSet::<u32>::from([1, 2, 3, 4, 5]);
        let set2 = PatriciaTreeSet::<u32>::from([4, 5, 6, 7, 8]);
        set1.union_with(set2);

        let actual: HashSet<_> = set1.iter().collect();
        let expected = HashSet::<u32>::from([1, 2, 3, 4, 5, 6, 7, 8]);

        assert_eq!(actual, expected);
    }

    #[test]
    fn test_is_subset() {
        let set1 = PatriciaTreeSet::<u32>::from([]);
        let set2 = PatriciaTreeSet::<u32>::from([]);
        assert!(set1.is_subset(&set2));
        assert!(set2.is_subset(&set1));

        let set1 = PatriciaTreeSet::<u32>::from([1, 2]);
        let set2 = PatriciaTreeSet::<u32>::from([1, 2, 3]);
        assert!(set1.is_subset(&set2));

        let set1 = PatriciaTreeSet::<u32>::from([203, 345, 324]);
        let mut set2 = PatriciaTreeSet::<u32>::from([203, 234]);
        assert!(!set1.is_subset(&set2));

        set2.union_with(set1.clone());
        assert!(set1.is_subset(&set2));

        let set1 = PatriciaTreeSet::<u32>::from([10, 20, 30]);
        let mut set2 = PatriciaTreeSet::<u32>::from([1, 2, 3, 4, 5]);
        assert!(!set1.is_subset(&set2));

        set2.union_with(set1.clone());
        assert!(set1.is_subset(&set2));
    }

    #[test]
    fn test_intersect_operation_empty() {
        let mut set1 = PatriciaTreeSet::<u32>::from([1, 2, 3]);
        let set2 = PatriciaTreeSet::<u32>::from([]);
        set1.intersection_with(&set2);

        let actual: HashSet<_> = set1.iter().collect();
        let expected = HashSet::<u32>::from([]);

        assert_eq!(actual, expected);

        let mut set1 = PatriciaTreeSet::<u32>::from([]);
        let set2 = PatriciaTreeSet::<u32>::from([]);
        set1.intersection_with(&set2);

        let actual: HashSet<_> = set1.iter().collect();
        let expected = HashSet::<u32>::from([]);

        assert_eq!(actual, expected);
    }

    #[test]
    fn test_simple_intersect_operation() {
        let mut set1 = PatriciaTreeSet::<u32>::from([0, 1, 2, 3, 4]);
        let set2 = PatriciaTreeSet::<u32>::from([2, 3, 4, 5, 6]);
        set1.intersection_with(&set2);

        let actual: HashSet<_> = set1.iter().collect();
        let expected = HashSet::<u32>::from([2, 3, 4]);

        assert_eq!(actual, expected);

        let mut set1 = PatriciaTreeSet::<u32>::from([0, 1, 2]);
        let set2 = PatriciaTreeSet::<u32>::from([3, 4, 5, 6]);
        set1.intersection_with(&set2);

        let actual: HashSet<_> = set1.iter().collect();
        let expected = HashSet::<u32>::from([]);

        assert_eq!(actual, expected);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::iter::FromIterator;

use im::HashSet;

use super::abstract_domain::AbstractDomain;
use crate::datatype::PatriciaTreeSet;

pub trait SetAbstractDomainOps: Clone + Eq {
    fn is_subset(&self, other: &Self) -> bool;
    fn intersection_with(&mut self, other: &Self);
    fn union_with(&mut self, other: Self);
}

pub trait SetElementOps {
    type Element;

    fn add_element(&mut self, e: Self::Element);
    fn remove_element(&mut self, e: &Self::Element);
}

#[derive(PartialEq, Eq, Clone, Debug)]
pub enum PowersetLattice<S: SetAbstractDomainOps> {
    Top,
    Value(S),
    Bottom,
}

impl<S: SetAbstractDomainOps> PowersetLattice<S> {
    pub fn value_from_set(set: S) -> Self {
        Self::Value(set)
    }
}

impl<'a, S: SetAbstractDomainOps + SetElementOps + 'a> PowersetLattice<S>
where
    &'a S: IntoIterator,
{
    pub fn add_element(&mut self, e: S::Element) {
        if let Self::Value(powerset) = self {
            powerset.add_element(e);
        }
    }

    pub fn add_elements<I: IntoIterator<Item = S::Element>>(&mut self, elements: I) {
        if let Self::Value(powerset) = self {
            for e in elements {
                powerset.add_element(e);
            }
        }
    }

    pub fn remove_element(&mut self, e: &S::Element) {
        if let Self::Value(powerset) = self {
            powerset.remove_element(e);
        }
    }

    pub fn remove_elements<'b, I: IntoIterator<Item = &'b S::Element>>(&mut self, elements: I)
    where
        S::Element: 'b,
    {
        if let Self::Value(powerset) = self {
            for e in elements {
                powerset.remove_element(e);
            }
        }
    }

    pub fn elements(&'a self) -> impl Iterator<Item = <&'a S as IntoIterator>::Item> {
        // NOTE: this is a workaround to make an empty iter, since we don't
        // know the actual type of S::ElementIter, we cannot create an empty
        // iter with the same type.
        let res = match self {
            Self::Value(powerset) => Some(powerset.into_iter()),
            _ => None,
        };
        res.into_iter().flatten()
    }

    pub fn set(&self) -> &S {
        match self {
            Self::Value(powerset) => powerset,
            _ => panic!("set called on Top or Bottom value!"),
        }
    }

    pub fn into_set(self) -> S {
        match self {
            Self::Value(powerset) => powerset,
            _ => panic!("into_set called on Top or Bottom value!"),
        }
    }
}

impl<S: SetAbstractDomainOps> AbstractDomain for PowersetLattice<S> {
    fn bottom() -> Self {
        Self::Bottom
    }

    fn top() -> Self {
        Self::Top
    }

    fn is_bottom(&self) -> bool {
        matches!(self, Self::Bottom)
    }

    fn is_top(&self) -> bool {
        matches!(self, Self::Top)
    }

    fn leq(&self, rhs: &Self) -> bool {
        match self {
            Self::Top => rhs.is_top(),
            Self::Value(ref s) => match rhs {
                Self::Top => true,
                Self::Value(ref t) => s.is_subset(t),
                Self::Bottom => false,
            },
            Self::Bottom => true,
        }
    }

    fn join_with(&mut self, rhs: Self) {
        match self {
            Self::Top => {}
            Self::Value(ref mut s) => match rhs {
                Self::Top => *self = rhs,
                Self::Value(t) => {
                    s.union_with(t);
                }
                Self::Bottom => {}
            },
            Self::Bottom => *self = rhs,
        };
    }

    fn meet_with(&mut self, rhs: Self) {
        match self {
            Self::Top => *self = rhs,
            Self::Value(ref mut s) => match rhs {
                Self::Top => {}
                Self::Value(ref t) => {
                    s.intersection_with(t);
                }
                Self::Bottom => *self = rhs,
            },
            Self::Bottom => {}
        };
    }

    fn widen_with(&mut self, _rhs: Self) {
        *self = Self::top();
    }

    fn narrow_with(&mut self, _rhs: Self) {
        *self = Self::bottom();
    }
}

impl<S, A> FromIterator<A> for PowersetLattice<S>
where
    S: SetAbstractDomainOps + FromIterator<A>,
{
    fn from_iter<T: IntoIterator<Item = A>>(iter: T) -> Self {
        Self::value_from_set(S::from_iter(iter))
    }
}

pub type PatriciaTreeSetAbstractDomain<T> = PowersetLattice<PatriciaTreeSet<T>>;
pub type HashSetAbstractDomain<T> = PowersetLattice<HashSet<T>>;

#[cfg(test)]
mod tests {
    use crate::datatype::abstract_domain::AbstractDomain;
    use crate::datatype::HashSetAbstractDomain;
    use crate::datatype::PatriciaTreeSetAbstractDomain;

    macro_rules! test_powerset_impl {
        ($powerset_type:ty) => {
            type IntPowerset = $powerset_type;

            let top = IntPowerset::Top;
            let value1: IntPowerset = vec![1, 2, 3, 4, 5].into_iter().collect();
            let value2: IntPowerset = vec![3, 4, 5, 6, 7].into_iter().collect();

            assert!(top.is_top());
            assert!(!value1.is_top());
            assert!(value1.leq(&IntPowerset::Top));
            assert!(IntPowerset::Bottom.leq(&IntPowerset::Top));

            assert!(!value1.leq(&value2));
            assert!(!value2.leq(&value1));

            let mut elements: Vec<_> = value1.elements().map(|e| e.to_owned()).collect();
            elements.sort();
            assert_eq!(elements, vec![1, 2, 3, 4, 5]);
            let mut elements: Vec<_> = top.elements().map(|e| e.to_owned()).collect();
            elements.sort();
            assert_eq!(elements, vec![]);

            let joined = value1.clone().join(value2.clone());
            let expected_joined: IntPowerset = vec![1, 2, 3, 4, 5, 6, 7].into_iter().collect();

            assert_eq!(joined, expected_joined);

            let met = value1.meet(value2);
            let expected_met: IntPowerset = vec![3, 4, 5].into_iter().collect();
            assert_eq!(met, expected_met);
        };
    }

    #[test]
    fn test_hash_powerset() {
        test_powerset_impl!(HashSetAbstractDomain<i64>);
    }

    #[test]
    fn test_patriciatree_powerset() {
        test_powerset_impl!(PatriciaTreeSetAbstractDomain<i64>);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

mod abstract_environment_test {
    use sparta::datatype::AbstractDomain;
    use sparta::datatype::AbstractEnvironment;
    use sparta::datatype::HashMapAbstractEnvironment;
    use sparta::datatype::HashSetAbstractDomain;
    use sparta::datatype::PatriciaTreeMapAbstractEnvironment;

    type Domain = HashSetAbstractDomain<String>;

    fn build_domain<const N: usize>(arr: [&str; N]) -> Domain {
        arr.iter().map(|str| str.to_string()).collect()
    }

    #[test]
    fn test_hmae_lattice_operations() {
        type Environment = HashMapAbstractEnvironment<String, Domain>;

        assert!(Environment::top().leq(&Environment::top()));
        assert!(!Environment::top().leq(&Environment::bottom()));
        assert!(Environment::bottom().leq(&Environment::top()));
        assert!(Environment::bottom().leq(&Environment::bottom()));

        assert!(Environment::bottom() == Environment::bottom());
        assert!(Environment::top() == Environment::top());
        assert!(Environment::bottom() != Environment::top());

        let mut e1 = Environment::top();
        e1.set("v1".to_owned(), build_domain(["a", "b"]));
        e1.set("v2".to_owned(), build_domain(["c"]));
        e1.set("v3".to_owned(), build_domain(["d", "e", "f"]));
        e1.set("v4".to_owned(), build_domain(["a", "f"]));
        assert_eq!(e1.len(), 4);

        let mut e2 = Environment::top();
        e2.set("v0".to_owned(), build_domain(["c", "f"]));
        e2.set("v2".to_owned(), build_domain(["c", "d"]));
        e2.set("v3".to_owned(), build_domain(["d", "e", "g", "h"]));
        assert_eq!(e2.len(), 3);

        let mut e3 = Environment::top();
        e3.set("v0".to_owned(), build_domain(["c", "d"]));
        e3.set("v2".to_owned(), Domain::bottom());
        e3.set("v3".to_owned(), build_domain(["a", "f", "g"]));
        assert!(e3.is_bottom());

        assert!(Environment::bottom().leq(&e1));
        assert!(!e1.leq(&Environment::bottom()));
        assert!(!Environment::top().leq(&e1));
        assert!(e1.leq(&Environment::top()));
        assert!(!e1.leq(&e2));
        assert!(!e2.leq(&e1));

        assert!(e1 == e1);
        assert!(e1 != e2);

        let join = e1.clone().join(e2.clone());
        assert!(e1.leq(&join));
        assert!(e2.leq(&join));
        assert_eq!(join.len(), 2);
        assert!(*join.get(&("v2".to_string())) == build_domain(["c", "d"]));
        assert!(*join.get(&("v3".to_string())) == build_domain(["d", "e", "f", "g", "h"]));

        assert!(*join.get(&("v4".to_string())) == Domain::top());

        assert!(e1.clone().join(Environment::top()).is_top());
        assert!(e1.clone().join(Environment::bottom()) == (e1));

        let meet = e1.clone().meet(e2.clone());
        assert!(meet.leq(&e1));
        assert!(meet.leq(&e2));
        assert_eq!(meet.len(), 5);

        assert!(*meet.get(&("v2".to_string())) == build_domain(["c"]));
        assert!(*meet.get(&("v3".to_string())) == build_domain(["d", "e"]));

        assert!(e1.clone().meet(Environment::bottom()).is_bottom());
        assert!(e1.clone().meet(Environment::top()) == e1);
    }

    #[test]
    fn test_ptmae_lattice_operations() {
        type Environment = PatriciaTreeMapAbstractEnvironment<u32, Domain>;

        assert!(Environment::top().leq(&Environment::top()));
        assert!(!Environment::top().leq(&Environment::bottom()));
        assert!(Environment::bottom().leq(&Environment::top()));
        assert!(Environment::bottom().leq(&Environment::bottom()));

        assert!(Environment::bottom() == Environment::bottom());
        assert!(Environment::top() == Environment::top());
        assert!(Environment::bottom() != Environment::top());

        let mut e1 = Environment::top();
        e1.set(1, build_domain(["a", "b"]));
        e1.set(2, build_domain(["c"]));
        e1.set(3, build_domain(["d", "e", "f"]));
        e1.set(4, build_domain(["a", "f"]));
        assert_eq!(e1.len(), 4);

        let mut e2 = Environment::top();
        e2.set(0, build_domain(["c", "f"]));
        e2.set(2, build_domain(["c", "d"]));
        e2.set(3, build_domain(["d", "e", "g", "h"]));
        assert_eq!(e2.len(), 3);

        let mut e3 = Environment::top();
        e3.set(0, build_domain(["c", "d"]));
        e3.set(2, Domain::bottom());
        e3.set(3, build_domain(["a", "f", "g"]));
        assert!(e3.is_bottom());

        assert!(Environment::bottom().leq(&e1));
        assert!(!e1.leq(&Environment::bottom()));
        assert!(!Environment::top().leq(&e1));
        assert!(e1.leq(&Environment::top()));
        assert!(!e1.leq(&e2));
        assert!(!e2.leq(&e1));

        assert!(e1 == e1);
        assert!(e1 != e2);

        let join = e1.clone().join(e2.clone());
        assert!(e1.leq(&join));
        assert!(e2.leq(&join));
        assert_eq!(join.len(), 2);
        assert!(!join.is_bottom());
        assert!(*join.get(&2) == build_domain(["c", "d"]));
        assert!(*join.get(&3) == build_domain(["d", "e", "f", "g", "h"]));

        assert!(*join.get(&4) == Domain::top());

        assert!(e1.clone().join(Environment::top()).is_top());
        assert!(e1.clone().join(Environment::bottom()) == (e1));

        let meet = e1.clone().meet(e2.clone());
        assert!(meet.leq(&e1));
        assert!(meet.leq(&e2));
        assert_eq!(meet.len(), 5);

        assert!(*meet.get(&2) == build_domain(["c"]));
        assert!(*meet.get(&3) == build_domain(["d", "e"]));

        assert!(e1.clone().meet(Environment::bottom()).is_bottom());
        assert!(e1.clone().meet(Environment::top()) == e1);
    }
}