pour 0.2.1

use super::*;
use num_traits::Zero;
use std::fmt::{self, Debug, Formatter};

/// The log of the branching factor for a radix trie
const LOG_BRANCHING_FACTOR: u8 = 2;

/// The branching factor for a radix trie
const BRANCHING_FACTOR: usize = 1 << LOG_BRANCHING_FACTOR;

/// The bin mask
const BIN_MASK: u8 = BRANCHING_FACTOR as u8 - 1;

/// Round a depth to a bin depth
pub fn round_depth<D>(depth: D) -> D
where
    D: PrimInt,
{
    let lbf: D = num_traits::NumCast::from(LOG_BRANCHING_FACTOR).unwrap();
    (depth / lbf) * lbf
}

/// Get the bin in a byte at a depth
pub fn byte_bin(byte: u8, depth: usize) -> usize {
    let byte_depth = round_depth(depth % 8);
    let bin_mask = BIN_MASK.wrapping_shl(byte_depth as u32);
    let bin_bits = byte & bin_mask;
    bin_bits.wrapping_shr(byte_depth as u32) as usize
}

/// Get the bin in a pattern at a depth
pub fn pattern_bin<P, D>(pattern: P, depth: D) -> usize
where
    P: Pattern<D>,
    D: Copy + AsPrimitive<usize>,
{
    let byte = pattern.byte(depth);
    byte_bin(byte, depth.as_())
}

/// An `IdMap`'s internal data.
#[derive(Clone)]
pub struct InnerMap<K: RadixKey, V: Clone> {
    pattern: K::PatternType,
    depth: K::DepthType,
    len: usize,
    branches: [Branch<K, V>; BRANCHING_FACTOR],
}

impl<K: RadixKey + Debug, V: Clone + Debug> Debug for InnerMap<K, V> {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        fmt.debug_struct("InnerMap")
            .field("len", &self.len)
            .field("branches", &self.branches)
            .finish()
    }
}

impl<K: RadixKey, V: Clone> InnerMap<K, V> {
    /// The size of the map this `InnerMap` represents
    pub fn len(&self) -> usize {
        self.len
    }
    /// The depth of this `InnerMap`
    pub fn depth(&self) -> usize {
        self.depth.to_usize().unwrap_or(usize::MAX)
    }
    /// Whether this `InnerMap` is a root `InnerMap`, i.e. can be placed directly into a `Set`
    pub fn is_root(&self) -> bool {
        K::pattern_no(self.depth) == 0
    }
    /// Whether this `InnerMap` is empty
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }
    /// Create an empty `InnerMap`
    fn empty() -> InnerMap<K, V> {
        InnerMap {
            pattern: K::PatternType::default(),
            depth: K::DepthType::zero(),
            len: 0,
            branches: InnerMap::empty_branches(),
        }
    }
    /// Get the branches for an empty `InnerMap`
    fn empty_branches() -> [Branch<K, V>; BRANCHING_FACTOR] {
        Default::default()
    }
    /// Create an `InnerMap` representing a singleton
    pub fn singleton(key: K, value: V) -> InnerMap<K, V> {
        let depth = K::DepthType::zero();
        let pattern = key.pattern(depth);
        let bin = pattern_bin(pattern, depth);
        let mut branches = InnerMap::empty_branches();
        branches[bin] = Branch::Mapping(key, value);
        InnerMap {
            pattern,
            depth,
            len: 1,
            branches,
        }
    }
    /// Create a consed, leveled `InnerMap` holding two key-value pairs, guaranteed to have different keys
    fn double_in<C: ConsCtx<K, V>>(
        first: (K, V),
        second: (K, V),
        depth: K::DepthType,
        ctx: &mut C,
    ) -> Arc<InnerMap<K, V>> {
        let first_pattern = first.0.pattern(depth);
        let second_pattern = second.0.pattern(depth);
        let diff = first_pattern.diff(second_pattern);
        if diff.as_() >= K::PatternType::MAX_BITS {
            unimplemented!("Level++")
        } else {
            let depth = round_depth(diff);
            let len = 2;
            let mut branches = InnerMap::empty_branches();
            branches[pattern_bin(first_pattern, depth)] = Branch::Mapping(first.0, first.1);
            branches[pattern_bin(second_pattern, depth)] = Branch::Mapping(second.0, second.1);
            ctx.cons(InnerMap {
                depth,
                len,
                branches,
                pattern: first_pattern,
            })
        }
    }
    /// Convert an `InnerMap` into an `IdMap` within a given context
    pub(super) fn into_idmap_in<C: ConsCtx<K, V>>(self, ctx: &mut C) -> IdMap<K, V> {
        if !self.is_root() {
            unimplemented!("Non root InnerMap conversion")
        }
        IdMap(self.cons_in(ctx))
    }
    /// Cons a tree in a given context
    fn cons_in<C: ConsCtx<K, V>>(self, ctx: &mut C) -> Option<Arc<InnerMap<K, V>>> {
        let premap = self.into_premap();
        premap.into_inner_in(ctx)
    }
    /// Convert a tree to a branch in a given context
    fn into_branch_in<C: ConsCtx<K, V>>(self, ctx: &mut C) -> Branch<K, V> {
        let premap = self.into_premap();
        premap.into_branch_in(ctx)
    }
    /// Convert a rooted `InnerMap` into a `PreMap<T>`
    fn into_premap(mut self) -> PreMap<K, V> {
        let level = K::pattern_no(self.depth);
        let mut last_tree = &mut None;
        let mut last_map = 0;
        let mut no_maps = 0;
        let mut tree_len = 0;
        let mut no_trees = 0;
        let depth = self.depth();
        for (ix, branch) in self.branches.iter_mut().enumerate() {
            match branch {
                Branch::Tree(t) => {
                    if let Some(tree) = t {
                        tree_len += tree.len();
                        no_trees += 1;
                        let tree_depth = tree.depth();
                        debug_assert!(
                            tree_depth > depth || tree_depth == 0,
                            "0 < Subtree #{} depth {} <= InnerMap depth {}",
                            ix,
                            tree_depth,
                            depth
                        );
                        last_tree = t;
                    }
                }
                Branch::Mapping(_, _) => {
                    no_maps += 1;
                    last_map = ix;
                }
            }
        }
        debug_assert_eq!(
            tree_len + no_maps,
            self.len,
            "Tree length = {} + no_maps = {} != self.len = {}",
            tree_len,
            no_maps,
            self.len
        );
        if last_tree.is_none() {
            debug_assert_eq!(no_trees, 0, "Last tree is none, but there are trees")
        } else {
            debug_assert_ne!(no_trees, 0, "There are no trees, but the last tree is Some")
        }
        if no_trees == BRANCHING_FACTOR {
            debug_assert_eq!(no_maps, 0, "There are no maps, since everything's a tree")
        }
        if no_maps == BRANCHING_FACTOR {
            debug_assert_eq!(no_trees, 0, "There are no trees, since everything's a map")
        }
        match (last_tree, no_maps, no_trees) {
            (None, 0, 0) => PreMap::Empty,
            (None, 1, 0) => PreMap::Mapping(self, last_map),
            (branch @ Some(_), 0, 1) => {
                let tree = branch.as_ref().unwrap();
                if K::pattern_no(tree.depth) <= level {
                    let mut root = None;
                    std::mem::swap(branch, &mut root);
                    PreMap::Root(root.unwrap())
                } else {
                    PreMap::Ready(self)
                }
            }
            _ => PreMap::Ready(self),
        }
    }
    /// Mutate this `InnerMap`, given a `this` pointer to itself, returning a result and an optional new `InnerMap<K, V>`.
    /// The returned `InnerMap` is always on the same level (i.e. `pattern_no`)
    ///
    /// Returns `None` if no changes were performed.
    pub(super) fn mutate<B, M, R, C>(
        &self,
        this: &Arc<InnerMap<K, V>>,
        key: B,
        mutator: M,
        ctx: &mut C,
    ) -> (Option<InnerMap<K, V>>, R)
    where
        B: Borrow<K>,
        M: FnOnce(B, Option<&V>) -> (Mutation<K, V>, R),
        C: ConsCtx<K, V>,
    {
        let bkey = key.borrow();
        let pattern = bkey.pattern(self.depth);
        let residual = self.depth % K::PatternType::max_bits();
        let diff = self.pattern.diff(pattern);
        if diff < residual {
            // The key cannot be in this bin, so handle that
            // Note we *know* that this implies the key can be on this level!
            let (mutation, result) = mutator(key, None);
            let (key, value) = match mutation {
                Mutation::Null => return (None, result),
                Mutation::Remove => return (None, result),
                Mutation::Update(_) => return (None, result),
                Mutation::Insert(key, value) => (key, value),
            };
            let depth = round_depth(self.depth - residual + diff);
            let this_bin = pattern_bin(self.pattern, depth);
            let other_bin = pattern_bin(pattern, depth);
            let mut branches = Self::empty_branches();
            branches[this_bin] = Branch::Tree(Some(this.clone()));
            branches[other_bin] = Branch::Mapping(key, value);
            let inner = InnerMap {
                len: self.len + 1,
                depth,
                branches,
                pattern,
            };
            return (Some(inner), result);
        }
        // The key may be in a sub-bin
        let bin = pattern_bin(pattern, self.depth);
        let (mutated_branch, result) =
            self.branches[bin].mutate(this, key, mutator, self.depth, ctx);
        if let Some(mutated_branch) = mutated_branch {
            let len = self.len - self.branches[bin].len() + mutated_branch.len();
            let mut branches = Self::empty_branches();
            for (ix, (new_branch, old_branch)) in
                branches.iter_mut().zip(self.branches.iter()).enumerate()
            {
                if ix != bin {
                    *new_branch = old_branch.clone_in(ctx)
                }
            }
            branches[bin] = mutated_branch;
            let inner = InnerMap {
                pattern: self.pattern,
                depth: self.depth,
                len,
                branches,
            };
            (Some(inner), result)
        } else {
            (None, result)
        }
    }
    /// Get an element in this `InnerMap`
    pub(super) fn get(&self, key: &K) -> Option<&V> {
        let pattern = key.pattern(self.depth);
        let depth: usize = self.depth.as_();
        let residual = depth % K::PatternType::MAX_BITS;
        let diff: usize = self.pattern.diff(pattern).as_();
        if diff < residual {
            return None;
        }
        let bin = pattern_bin(pattern, self.depth);
        match &self.branches[bin] {
            Branch::Mapping(entry_key, value) if entry_key == key => Some(value),
            Branch::Tree(Some(tree)) => tree.get(key),
            _ => None,
        }
    }
    /// Mutate the values of an `InnerMap` in a given context, yielding a new `InnerMap` on the same level if anything changed
    pub(super) fn mutate_vals_in<M, C>(
        &self,
        mutator: &mut M,
        ctx: &mut C,
    ) -> Option<InnerMap<K, V>>
    where
        M: UnaryMutator<K, V>,
        C: ConsCtx<K, V>,
    {
        let mut branches = InnerMap::empty_branches();
        let mut mutated_branches = [false; BRANCHING_FACTOR];
        // Mutate branches
        for ((new_branch, is_mutated), old_branch) in branches
            .iter_mut()
            .zip(mutated_branches.iter_mut())
            .zip(self.branches.iter())
        {
            if let Some(branch) = old_branch.mutate_vals_in(mutator, ctx) {
                *new_branch = branch;
                *is_mutated = true;
            }
        }
        // If no branches were mutated, early exit with no changes
        if mutated_branches.iter().all(|x| !x) {
            return None;
        }
        // Clone over unchanged branches
        for ((new_branch, is_mutated), old_branch) in branches
            .iter_mut()
            .zip(mutated_branches.iter())
            .zip(self.branches.iter())
        {
            if !*is_mutated {
                *new_branch = old_branch.clone_in(ctx)
            }
        }
        // Build a new `InnerMap` from the mutated branches
        let len = branches.iter().map(|branch| branch.len()).sum();
        let inner = InnerMap {
            branches,
            len,
            pattern: self.pattern,
            depth: self.depth,
        };
        Some(inner)
    }
    /// Join-mutate two maps by applying a binary mutator to their key intersection and unary
    /// mutators to their left and right intersection.
    pub(super) fn join_mutate_in<IM, LM, RM, C>(
        &self,
        other: &Self,
        intersection_mutator: &mut IM,
        left_mutator: &mut LM,
        right_mutator: &mut RM,
        ctx: &mut C,
    ) -> BinaryResult<InnerMap<K, V>>
    where
        IM: BinaryMutator<K, V>,
        LM: UnaryMutator<K, V>,
        RM: UnaryMutator<K, V>,
        C: ConsCtx<K, V>,
    {
        if self as *const _ == other as *const _ {
            match intersection_mutator.kind() {
                BinaryMutatorKind::Nilpotent => return New(InnerMap::empty()),
                BinaryMutatorKind::Idempotent => return Ambi,
                BinaryMutatorKind::Left => return Ambi,
                BinaryMutatorKind::Right => return Ambi,
                BinaryMutatorKind::Ambi => return Ambi,
                BinaryMutatorKind::General => {}
            }
        }
        use BinaryResult::*;
        let mut branches = Self::empty_branches();
        let mut sources = [New(()); BRANCHING_FACTOR];
        let mut left_count = 0;
        let mut right_count = 0;
        let mut ambi_count = 0;
        let mut new_count = 0;
        let editor = branches.iter_mut().zip(sources.iter_mut());
        let inputs = self.branches.iter().zip(other.branches.iter());
        // Perform branch-wise merges
        for ((branch, source), (left, right)) in editor.zip(inputs) {
            match left.join_mutate_in(
                right,
                intersection_mutator,
                left_mutator,
                right_mutator,
                self.depth,
                ctx,
            ) {
                New(new_branch) => {
                    *branch = new_branch;
                    new_count += 1;
                }
                Left => {
                    left_count += 1;
                    *source = Left
                }
                Right => {
                    right_count += 1;
                    *source = Right
                }
                Ambi => {
                    ambi_count += 1;
                    *source = Ambi
                }
            }
        }
        // Short circuit left and right returns
        if new_count == 0 {
            match (left_count, right_count, ambi_count) {
                (0, 0, _) => {
                    debug_assert_eq!(ambi_count, BRANCHING_FACTOR);
                    return Ambi;
                }
                (left_count, 0, ambi_count) => {
                    debug_assert_eq!(left_count + ambi_count, BRANCHING_FACTOR);
                    return Left;
                }
                (0, right_count, ambi_count) => {
                    debug_assert_eq!(right_count + ambi_count, BRANCHING_FACTOR);
                    return Right;
                }
                _ => {}
            }
        }

        // Copy over left and right sources
        let editor = branches.iter_mut().zip(sources.iter_mut());
        let inputs = self.branches.iter().zip(other.branches.iter());
        for ((branch, source), (left, right)) in editor.zip(inputs) {
            match source {
                New(()) => {}
                Ambi => *branch = left.clone_in(ctx),
                Left => *branch = left.clone_in(ctx),
                Right => *branch = right.clone_in(ctx),
            }
        }
        // Build a new `InnerMap` from the fused branches
        let len = branches.iter().map(|branch| branch.len()).sum();
        let inner = InnerMap {
            branches,
            len,
            pattern: self.pattern,
            depth: self.depth,
        };
        New(inner)
    }
}

impl<K: RadixKey, V: Clone + Eq> InnerMap<K, V> {
    /// Check whether two `InnerMap`s are recursively equal, i.e. contain the same data versus point to the same data
    pub fn rec_eq(&self, other: &InnerMap<K, V>) -> bool {
        if self as *const _ == other as *const _ {
            return true;
        }
        self.depth == other.depth
            && self.len == other.len
            && self
                .branches
                .iter()
                .zip(other.branches.iter())
                .all(|(this, other)| this.rec_eq(other))
    }
    /// Check whether this map is a submap of another. A map is considered to be a submap of itself.
    ///
    /// If `cons` is true, this map is assumed to be hash-consed with the other
    pub(super) fn is_submap(&self, other: &InnerMap<K, V>, cons: bool) -> bool {
        if self as *const _ as *const u8 == other as *const _ as *const u8 {
            // This is correct since both sides implement Eq
            return true;
        }
        if self.len > other.len {
            // Length heuristic
            return false;
        }
        self.branches
            .iter()
            .zip(other.branches.iter())
            .all(|(left, right)| left.is_submap(right, cons))
    }
    /// Check whether this map's key set is a subset of another map's keyset.
    pub(super) fn domain_is_subset<U: Clone + Eq>(&self, other: &InnerMap<K, U>) -> bool {
        if self as *const _ as *const u8 == other as *const _ as *const u8 {
            // This is correct since both sides implement Eq
            return true;
        }
        if self.len > other.len {
            return false;
        }
        self.branches
            .iter()
            .zip(other.branches.iter())
            .all(|(left, right)| left.domain_is_subset(right))
    }
    /// Check whether this map's domains intersect
    pub(super) fn domains_intersect<U: Clone + Eq>(&self, other: &InnerMap<K, U>) -> bool {
        if self as *const _ as *const u8 == other as *const _ as *const u8 {
            // This is correct since both sides implement Eq
            return true;
        }
        self.branches
            .iter()
            .zip(other.branches.iter())
            .any(|(left, right)| left.domains_intersect(right))
    }
}

/// A borrowed iterator over an `IdMap`
#[derive(Debug, Clone)]
pub struct IdMapIter<'a, K: RadixKey, V: Clone> {
    inner_stack: Vec<(&'a InnerMap<K, V>, usize)>,
}

impl<'a, K: RadixKey, V: Clone> IdMapIter<'a, K, V> {
    /// Create a new, empty iterator
    pub fn empty() -> IdMapIter<'a, K, V> {
        IdMapIter {
            inner_stack: Vec::new(),
        }
    }
    /// Add a new root to a given iterator
    pub(super) fn root(&mut self, root: &'a InnerMap<K, V>) {
        self.inner_stack.push((root, 0))
    }
}

impl<'a, K: RadixKey, V: Clone> Iterator for IdMapIter<'a, K, V> {
    type Item = (&'a K, &'a V);
    fn next(&mut self) -> Option<(&'a K, &'a V)> {
        loop {
            let (top, ix) = self.inner_stack.last_mut()?;
            if *ix >= BRANCHING_FACTOR {
                self.inner_stack.pop();
                continue;
            }
            match &top.branches[*ix] {
                Branch::Mapping(key, value) => {
                    *ix += 1;
                    return Some((key, value));
                }
                Branch::Tree(None) => {
                    *ix += 1;
                }
                Branch::Tree(Some(tree)) => {
                    *ix += 1;
                    self.inner_stack.push((&*tree, 0));
                }
            }
        }
    }
}

/// An iterator over an `IdMap`
#[derive(Debug, Clone)]
pub struct IdMapIntoIter<K: RadixKey, V: Clone> {
    inner_stack: Vec<(Arc<InnerMap<K, V>>, usize)>,
}

impl<K: RadixKey, V: Clone> IdMapIntoIter<K, V> {
    /// Create a new, empty iterator
    pub fn empty() -> IdMapIntoIter<K, V> {
        IdMapIntoIter {
            inner_stack: Vec::new(),
        }
    }
    /// Add a new root to a given iterator
    pub(super) fn root(&mut self, root: Arc<InnerMap<K, V>>) {
        self.inner_stack.push((root, 0))
    }
}

impl<K: RadixKey, V: Clone> Iterator for IdMapIntoIter<K, V> {
    type Item = (K, V);
    fn next(&mut self) -> Option<(K, V)> {
        loop {
            let (top, ix) = self.inner_stack.last_mut()?;
            if *ix >= BRANCHING_FACTOR {
                self.inner_stack.pop();
                continue;
            }
            match top.branches[*ix].clone() {
                Branch::Mapping(key, value) => {
                    *ix += 1;
                    return Some((key, value));
                }
                Branch::Tree(None) => {
                    *ix += 1;
                }
                Branch::Tree(Some(tree)) => {
                    *ix += 1;
                    self.inner_stack.push((tree, 0));
                }
            }
        }
    }
}

/// An `InnerMap` converted to a "pre-map"
#[derive(Debug, Clone)]
enum PreMap<K: RadixKey, V: Clone> {
    /// An empty map
    Empty,
    /// A single non-null root tree entry
    Root(Arc<InnerMap<K, V>>),
    /// A single mapping
    Mapping(InnerMap<K, V>, usize),
    /// A map ready for direct conversion
    Ready(InnerMap<K, V>),
}

impl<K: RadixKey, V: Clone> PreMap<K, V> {
    pub fn into_inner_in<C: ConsCtx<K, V>>(self, ctx: &mut C) -> Option<Arc<InnerMap<K, V>>> {
        use PreMap::*;
        match self {
            Empty => None,
            Root(root) => Some(root),
            Ready(map) | Mapping(map, _) => Some(ctx.cons(map)),
        }
    }
    pub fn into_branch_in<C: ConsCtx<K, V>>(self, ctx: &mut C) -> Branch<K, V> {
        use Branch::*;
        use PreMap::{Mapping, *};
        match self {
            Empty => Tree(None),
            Root(root) => Tree(Some(root)),
            Mapping(mut map, ix) => {
                let mut result = Tree(None);
                std::mem::swap(&mut result, &mut map.branches[ix]);
                result
            }
            Ready(map) => Tree(Some(ctx.cons(map))),
        }
    }
}

/// An branch in a layer of an `IdMap`, which is either empty, a mapping or a recursive tree
#[derive(Debug, Clone)]
enum Branch<K: RadixKey, V: Clone> {
    /// A single mapping
    Mapping(K, V),
    /// A tree of recursive mappings
    Tree(Option<Arc<InnerMap<K, V>>>),
}

impl<K: RadixKey, V: Clone> Default for Branch<K, V> {
    fn default() -> Branch<K, V> {
        Branch::Tree(None)
    }
}

impl<K: RadixKey, V: Clone + Eq> Branch<K, V> {
    /// Check whether two branches are recursively equal
    pub fn rec_eq(&self, other: &Branch<K, V>) -> bool {
        use Branch::*;
        match (self, other) {
            (Mapping(this_key, this_value), Mapping(other_key, other_value)) => {
                this_key == other_key && this_value == other_value
            }
            (Tree(None), Tree(None)) => true,
            (Tree(Some(this)), Tree(Some(other))) => this.rec_eq(other),
            _ => false,
        }
    }
    /// Check whether this branch is a submap of another. A map is considered to be a submap of itself.
    ///
    /// If `cons` is true, this map is assumed to be hash-consed with the other
    fn is_submap(&self, other: &Branch<K, V>, cons: bool) -> bool {
        use Branch::*;
        match (self, other) {
            (Tree(None), _) => true,
            (_, Tree(None)) => false,
            (Tree(Some(this)), Tree(Some(other))) => {
                if Arc::ptr_eq(this, other) {
                    true
                } else if cons && this.len() == other.len() {
                    false
                } else {
                    this.is_submap(other, cons)
                }
            }
            (Tree(Some(this)), Mapping(key, value)) => {
                if this.len() != 1 {
                    false
                } else {
                    this.get(key) == Some(value)
                }
            }
            (Mapping(key, value), Tree(Some(other))) => other.get(key) == Some(value),
            (Mapping(left_key, left_value), Mapping(right_key, right_value)) => {
                left_key == right_key && left_value == right_value
            }
        }
    }
    /// Check whether this branch's domain is a subset of another's.
    fn domain_is_subset<U: Clone + Eq>(&self, other: &Branch<K, U>) -> bool {
        use Branch::*;
        match (self, other) {
            (Tree(None), _) => true,
            (_, Tree(None)) => false,
            (Tree(Some(this)), Tree(Some(other))) => {
                if Arc::as_ptr(this) as *const u8 == Arc::as_ptr(other) as *const u8 {
                    true
                } else {
                    this.domain_is_subset(other)
                }
            }
            (Tree(Some(this)), Mapping(key, _)) => {
                if this.len() != 1 {
                    false
                } else {
                    this.get(key).is_some()
                }
            }
            (Mapping(key, _), Tree(Some(other))) => other.get(key).is_some(),
            (Mapping(left_key, _), Mapping(right_key, _)) => left_key == right_key,
        }
    }
    /// Check whether this branch's domain intersects another's
    fn domains_intersect<U: Clone + Eq>(&self, other: &Branch<K, U>) -> bool {
        use Branch::*;
        match (self, other) {
            (Tree(None), _) => false,
            (_, Tree(None)) => false,
            (Tree(Some(this)), Tree(Some(other))) => this.domains_intersect(other),
            (Tree(Some(this)), Mapping(key, _)) => this.get(key).is_some(),
            (Mapping(key, _), Tree(Some(other))) => other.get(key).is_some(),
            (Mapping(left_key, _), Mapping(right_key, _)) => left_key == right_key,
        }
    }
}

impl<K: RadixKey, V: Clone> Branch<K, V> {
    /// Clone a branch in a given context, consing appropriately
    pub fn clone_in<C: ConsCtx<K, V>>(&self, ctx: &mut C) -> Branch<K, V> {
        use Branch::*;
        match self {
            Mapping(key, value) => Mapping(key.clone(), value.clone()),
            Tree(None) => Tree(None),
            Tree(Some(tree)) => Tree(Some(ctx.cons_recursive(tree))),
        }
    }
    /// Mutate this `Branch`, given a `this` pointer to itself, returning a result and an optional new `Branch<K, V>`.
    ///
    /// Returns `None` if no changes were performed.
    pub(super) fn mutate<B, M, R, C>(
        &self,
        this: &Arc<InnerMap<K, V>>,
        key: B,
        mutator: M,
        depth: K::DepthType,
        ctx: &mut C,
    ) -> (Option<Branch<K, V>>, R)
    where
        B: Borrow<K>,
        M: FnOnce(B, Option<&V>) -> (Mutation<K, V>, R),
        C: ConsCtx<K, V>,
    {
        use Branch::*;
        let bkey = key.borrow();
        match self {
            Mapping(entry_key, value) if entry_key == bkey => {
                let (mutation, result) = mutator(key, Some(value));
                let mutated = match mutation {
                    Mutation::Null => None,
                    Mutation::Remove => Some(Tree(None)),
                    Mutation::Insert(_, _) => None,
                    Mutation::Update(value) => Some(Mapping(entry_key.clone(), value)),
                };
                (mutated, result)
            }
            Mapping(entry_key, entry_value) => {
                let (mutation, result) = mutator(key, None);
                let mutated = match mutation {
                    Mutation::Null => None,
                    Mutation::Remove => None,
                    Mutation::Insert(key, value) => Some(Tree(Some(InnerMap::double_in(
                        (key, value),
                        (entry_key.clone(), entry_value.clone()),
                        depth,
                        ctx,
                    )))),
                    Mutation::Update(_) => None,
                };
                (mutated, result)
            }
            Tree(None) => {
                let (mutation, result) = mutator(key, None);
                let mutated = match mutation {
                    Mutation::Null => None,
                    Mutation::Remove => None,
                    Mutation::Insert(key, value) => Some(Mapping(key, value)),
                    Mutation::Update(_) => None,
                };
                (mutated, result)
            }
            Tree(Some(tree)) => {
                let (try_tree, result) = tree.mutate(this, key, mutator, ctx);
                let mutated = if let Some(tree) = try_tree {
                    Some(tree.into_branch_in(ctx))
                } else {
                    None
                };
                (mutated, result)
            }
        }
    }
    /// Mutate the values stored in a branch, returning `Some` if anything changed
    fn mutate_vals_in<M, C>(&self, mutator: &mut M, ctx: &mut C) -> Option<Branch<K, V>>
    where
        M: UnaryMutator<K, V>,
        C: ConsCtx<K, V>,
    {
        use Branch::*;
        use UnaryResult::*;
        match mutator.kind() {
            UnaryMutatorKind::Null => return None,
            UnaryMutatorKind::Delete => {
                return if self.is_empty() {
                    None
                } else {
                    Some(Tree(None))
                }
            }
            _ => {}
        }
        match self {
            Mapping(key, value) => match mutator.mutate(key, value) {
                New(None) => Some(Tree(None)),
                New(Some(value)) => Some(Mapping(key.clone(), value)),
                Old => None,
            },
            Tree(Some(tree)) => {
                let mutated_tree = tree.mutate_vals_in(mutator, ctx)?;
                Some(mutated_tree.into_branch_in(ctx))
            }
            Tree(None) => None,
        }
    }
    /// Join-mutate two branches by applying a binary mutator to their key intersection and unary
    /// mutators to their left and right intersection.
    pub(super) fn join_mutate_in<IM, LM, RM, C>(
        &self,
        other: &Self,
        intersection_mutator: &mut IM,
        left_mutator: &mut LM,
        right_mutator: &mut RM,
        depth: K::DepthType,
        ctx: &mut C,
    ) -> BinaryResult<Branch<K, V>>
    where
        IM: BinaryMutator<K, V>,
        LM: UnaryMutator<K, V>,
        RM: UnaryMutator<K, V>,
        C: ConsCtx<K, V>,
    {
        use BinaryResult::*;
        use Branch::*;
        match (self, other) {
            (this, Tree(None)) => BinaryResult::or_left(this.mutate_vals_in(left_mutator, ctx)),
            (Tree(None), other) => BinaryResult::or_right(other.mutate_vals_in(right_mutator, ctx)),
            (Mapping(lkey, lvalue), Mapping(rkey, rvalue)) => {
                if lkey == rkey {
                    match intersection_mutator.kind() {
                        BinaryMutatorKind::Left => Left,
                        BinaryMutatorKind::Right => Right,
                        BinaryMutatorKind::Ambi => Ambi,
                        _ => intersection_mutator
                            .mutate_bin(lkey, lvalue, rvalue)
                            .map(|value| {
                                if let Some(value) = value {
                                    Mapping(lkey.clone(), value)
                                } else {
                                    Tree(None)
                                }
                            }),
                    }
                } else {
                    use UnaryResult::{New as Nw, Old};
                    let left = left_mutator.mutate(lkey, lvalue);
                    let right = right_mutator.mutate(rkey, rvalue);
                    let (left, right) = match (left, right) {
                        (Old, Nw(None)) => return Left,
                        (Nw(None), Old) => return Right,
                        (Nw(None), Nw(None)) => return New(Tree(None)),
                        (Nw(Some(left)), Nw(None)) => return New(Mapping(lkey.clone(), left)),
                        (Nw(None), Nw(Some(right))) => return New(Mapping(rkey.clone(), right)),
                        (Nw(Some(left)), Nw(Some(right))) => (left, right),
                        (Nw(Some(left)), Old) => (left, rvalue.clone()),
                        (Old, Nw(Some(right))) => (lvalue.clone(), right),
                        (Old, Old) => (lvalue.clone(), rvalue.clone()),
                    };
                    let inner = InnerMap::double_in(
                        (lkey.clone(), left),
                        (rkey.clone(), right),
                        depth,
                        ctx,
                    );
                    New(Tree(Some(inner)))
                }
            }
            (Tree(Some(this)), Tree(Some(other))) => this
                .join_mutate_in(
                    other,
                    intersection_mutator,
                    left_mutator,
                    right_mutator,
                    ctx,
                )
                .map(|tree| tree.into_branch_in(ctx)),
            (Tree(Some(this)), Mapping(key, value)) => Branch::join_mutate_mapping(
                (key, value),
                this,
                Swapped::ref_cast_mut(intersection_mutator),
                // NOTE: right and left are swapped here, then we swap sides again
                right_mutator,
                left_mutator,
                ctx,
            )
            .swap_sides(),
            (Mapping(key, value), Tree(Some(other))) => Branch::join_mutate_mapping(
                (key, value),
                other,
                intersection_mutator,
                left_mutator,
                right_mutator,
                ctx,
            ),
        }
    }
    /// A helper to perform an intersection-mapped mutation on a mapping
    pub fn join_mutate_mapping<IM, LM, RM, C>(
        left: (&K, &V),
        right: &Arc<InnerMap<K, V>>,
        intersection_mutator: &mut IM,
        left_mutator: &mut LM,
        right_mutator: &mut RM,
        ctx: &mut C,
    ) -> BinaryResult<Branch<K, V>>
    where
        IM: BinaryMutator<K, V>,
        LM: UnaryMutator<K, V>,
        RM: UnaryMutator<K, V>,
        C: ConsCtx<K, V>,
    {
        use BinaryResult::*;
        use Branch::*;
        use UnaryResult::{New as Nw, Old};
        let key = left.0;
        let left_value = left.1;
        match right_mutator.kind() {
            UnaryMutatorKind::Null => {
                // The right tree is always left alone except for the target, so insert into it unchanged.
                let (mutated_tree, ()) = right.mutate(
                    right,
                    key,
                    |key, right_value| {
                        use Mutation::*;
                        let mutation = if let Some(right_value) = right_value {
                            match intersection_mutator.mutate_bin(key, left_value, right_value) {
                                New(Some(value)) => Update(value),
                                New(None) => Remove,
                                Left => Update(left_value.clone()),
                                Right | Ambi => Null,
                            }
                        } else {
                            match left_mutator.mutate(key, left_value) {
                                Old => Insert(key.clone(), left_value.clone()),
                                Nw(Some(value)) => Insert(key.clone(), value),
                                Nw(None) => Null,
                            }
                        };
                        (mutation, ())
                    },
                    ctx,
                );
                if let Some(mutated_tree) = mutated_tree {
                    New(mutated_tree.into_branch_in(ctx))
                } else {
                    Right
                }
            }
            UnaryMutatorKind::Delete => {
                // The right tree is always deleted except for the target, so extract the target value and operate on that
                let right_value = right.get(key);
                if let Some(right_value) = right_value {
                    match intersection_mutator.mutate_bin(key, left_value, right_value) {
                        // The left is alone and unchanged
                        Left | Ambi => Left,
                        // Only the right remains, so return it as a new mapping
                        Right => New(Mapping(key.clone(), right_value.clone())),
                        // A new value has been created, so return it as a new mapping
                        New(Some(value)) => New(Mapping(key.clone(), value)),
                        // Everything has been deleted, so return the empty tree
                        New(None) => New(Tree(None)),
                    }
                } else {
                    match left_mutator.mutate(key, left_value) {
                        // The left is alone and unchanged
                        Old => Left,
                        // The left has been modified, so return it as a new mapping
                        Nw(Some(left)) => New(Mapping(key.clone(), left)),
                        // Everything has been deleted, so return the empty tree
                        Nw(None) => New(Tree(None)),
                    }
                }
            }
            UnaryMutatorKind::General => {
                //TODO: new recursive subroutine needed here for maximum efficiency
                let mut left_mutated = false;
                let mut mutator = FilterMap::annotated(|entry_key, right_value| {
                    if entry_key == key {
                        left_mutated = true;
                        match intersection_mutator.mutate_bin(key, left_value, right_value) {
                            Ambi => Old,
                            Right => Old,
                            Left => Nw(Some(right_value.clone())),
                            New(n) => Nw(n),
                        }
                    } else {
                        right_mutator.mutate(entry_key, right_value)
                    }
                });
                let mutated_tree = right.mutate_vals_in(&mut mutator, ctx);
                let uninserted_tree = match (mutated_tree, left_mutated) {
                    (Some(tree), true) => return New(tree.into_branch_in(ctx)),
                    (Some(tree), false) => match tree.cons_in(ctx) {
                        Some(tree) => tree,
                        None => return Left,
                    },
                    (None, true) => return New(Tree(None)),
                    (None, false) => return Left,
                };
                let (inserted_tree, _) = uninserted_tree.mutate(
                    &uninserted_tree,
                    key,
                    |key, _value| (Mutation::Insert(key.clone(), left_value.clone()), ()),
                    ctx,
                );
                New(inserted_tree
                    .expect("If the left was not inserted, it can't be in the tree")
                    .into_branch_in(ctx))
            }
        }
    }
    /// Get whether this `Branch` is empty
    pub fn is_empty(&self) -> bool {
        matches!(self, Branch::Tree(None))
    }
    /// The number of mappings this `Branch` represents
    pub fn len(&self) -> usize {
        use Branch::*;
        match self {
            Mapping(_, _) => 1,
            Tree(None) => 0,
            Tree(Some(t)) => t.len(),
        }
    }
}

impl<K: RadixKey + Hash, V: Clone + Hash> Hash for Branch<K, V> {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        match self {
            Branch::Mapping(key, value) => {
                // To make this `NonNull`, unlike the pointer
                0x34u8.hash(hasher);
                key.hash(hasher);
                value.hash(hasher)
            }
            Branch::Tree(None) => std::ptr::null::<*const InnerMap<K, V>>().hash(hasher),
            Branch::Tree(Some(tree)) => Arc::as_ptr(tree).hash(hasher),
        }
    }
}

impl<K: RadixKey, V: Clone + Eq> PartialEq for Branch<K, V> {
    fn eq(&self, other: &Branch<K, V>) -> bool {
        use Branch::*;
        match (self, other) {
            (Mapping(key1, value1), Mapping(key2, value2)) => key1 == key2 && value1 == value2,
            (Tree(None), Tree(None)) => true,
            (Tree(Some(this)), Tree(Some(other))) => Arc::ptr_eq(this, other),
            _ => false,
        }
    }
}

impl<K: RadixKey + Eq, V: Clone + Eq> Eq for Branch<K, V> {}

impl<K: RadixKey + PartialOrd, V: Clone + Eq + PartialOrd> PartialOrd for Branch<K, V> {
    fn partial_cmp(&self, other: &Branch<K, V>) -> Option<Ordering> {
        use Branch::*;
        use Ordering::*;
        match (self, other) {
            (Mapping(key1, value1), Mapping(key2, value2)) => {
                (key1, value1).partial_cmp(&(key2, value2))
            }
            (Mapping(_, _), Tree(None)) => Some(Less),
            (Mapping(_, _), Tree(Some(_))) => Some(Less),
            (Tree(None), Mapping(_, _)) => Some(Greater),
            (Tree(None), Tree(None)) => Some(Equal),
            (Tree(None), Tree(Some(_))) => Some(Less),
            (Tree(Some(_)), Mapping(_, _)) => Some(Greater),
            (Tree(Some(_)), Tree(None)) => Some(Greater),
            (Tree(Some(tree1)), Tree(Some(tree2))) => {
                Arc::as_ptr(tree1).partial_cmp(&Arc::as_ptr(&tree2))
            }
        }
    }
}

impl<K: RadixKey + Eq + Ord, V: Clone + Eq + Ord> Ord for Branch<K, V> {
    fn cmp(&self, other: &Branch<K, V>) -> Ordering {
        use Branch::*;
        use Ordering::*;
        match (self, other) {
            (Mapping(key1, value1), Mapping(key2, value2)) => (key1, value1).cmp(&(key2, value2)),
            (Mapping(_, _), Tree(None)) => Less,
            (Mapping(_, _), Tree(Some(_))) => Less,
            (Tree(None), Mapping(_, _)) => Greater,
            (Tree(None), Tree(None)) => Equal,
            (Tree(None), Tree(Some(_))) => Less,
            (Tree(Some(_)), Mapping(_, _)) => Greater,
            (Tree(Some(_)), Tree(None)) => Greater,
            (Tree(Some(tree1)), Tree(Some(tree2))) => Arc::as_ptr(tree1).cmp(&Arc::as_ptr(&tree2)),
        }
    }
}

impl<K: RadixKey, V: Clone + Eq> PartialEq for InnerMap<K, V> {
    fn eq(&self, other: &InnerMap<K, V>) -> bool {
        self.branches == other.branches
    }
}

impl<K: RadixKey, V: Clone + Eq> Eq for InnerMap<K, V> {}

impl<K: RadixKey + PartialOrd, V: Clone + Eq + PartialOrd> PartialOrd for InnerMap<K, V> {
    fn partial_cmp(&self, other: &InnerMap<K, V>) -> Option<Ordering> {
        self.branches.partial_cmp(&other.branches)
    }
}

impl<K: RadixKey + Eq + Ord, V: Clone + Eq + Ord> Ord for InnerMap<K, V> {
    fn cmp(&self, other: &InnerMap<K, V>) -> Ordering {
        self.branches.cmp(&other.branches)
    }
}

impl<K: RadixKey + Hash, V: Clone + Hash> Hash for InnerMap<K, V> {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.branches.hash(hasher);
    }
}

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

    #[test]
    fn basic_double_in_test() {
        let double = InnerMap::double_in((0b110, 3), (0b011, 2), 0, &mut ());
        assert_eq!(double.len(), 2);
        assert_eq!(double.depth(), 0);
        let double = InnerMap::double_in((0b111, 3), (0b011, 2), 0, &mut ());
        assert_eq!(double.len(), 2);
        assert_eq!(double.depth(), 2);
        let double = InnerMap::double_in((0b111, 3), (0b1111, 2), 0, &mut ());
        assert_eq!(double.len(), 2);
        assert_eq!(double.depth(), 2);
        let double = InnerMap::double_in((0b1111, 3), (0b11111, 2), 0, &mut ());
        assert_eq!(double.len(), 2);
        assert_eq!(double.depth(), 4);
    }

    #[test]
    fn empty_inner_consing() {
        let empty = InnerMap::<u64, u64>::empty();
        assert_eq!(empty.len(), 0);
        assert!(empty.is_empty());
        assert_eq!(empty.clone().cons_in(&mut ()), None);
        assert_eq!(empty.into_idmap_in(&mut ()), IdMap::EMPTY);
    }

    #[test]
    fn basic_branch_ordering() {
        use Branch::*;
        let sorted_branches = [
            Mapping(3, 5),
            Mapping(3, 7),
            Mapping(4, 5),
            Mapping(4, 7),
            Tree(None),
            Tree(Some(InnerMap::double_in((3, 5), (4, 7), 0, &mut ()))),
        ];
        for (i, l) in sorted_branches.iter().enumerate() {
            for (j, r) in sorted_branches.iter().enumerate() {
                assert_eq!(
                    i.cmp(&j),
                    l.cmp(r),
                    "Comparison of {:?}@{} {:?}@{} wrong",
                    l,
                    i,
                    r,
                    j
                );
                assert_eq!(
                    i.partial_cmp(&j),
                    l.partial_cmp(r),
                    "Partial comparison of {:?}@{} {:?}@{} wrong",
                    l,
                    i,
                    r,
                    j
                )
            }
        }
    }

    #[test]
    fn basic_branch_joins() {
        use Branch::*;
        let map1 = Mapping(3, 6);
        let map2 = Mapping(3, 5);
        let map3 = Mapping(4, 13);
        assert_eq!(
            map1.join_mutate_in(
                &map2,
                &mut LeftMutator,
                &mut NullMutator,
                &mut NullMutator,
                0,
                &mut ()
            ),
            BinaryResult::Left
        );
        assert_eq!(
            map1.join_mutate_in(
                &map2,
                &mut RightMutator,
                &mut NullMutator,
                &mut NullMutator,
                0,
                &mut ()
            ),
            BinaryResult::Right
        );
        assert_eq!(
            map1.join_mutate_in(
                &map2,
                &mut AmbiMutator,
                &mut NullMutator,
                &mut NullMutator,
                0,
                &mut ()
            ),
            BinaryResult::Ambi
        );
        assert_eq!(
            map1.join_mutate_in(
                &map3,
                &mut LeftMutator,
                &mut DeleteMutator,
                &mut NullMutator,
                0,
                &mut ()
            ),
            BinaryResult::Right
        );
        assert_eq!(
            map1.join_mutate_in(
                &map3,
                &mut LeftMutator,
                &mut NullMutator,
                &mut DeleteMutator,
                0,
                &mut ()
            ),
            BinaryResult::Left
        );
        assert_eq!(
            map1.join_mutate_in(
                &map3,
                &mut LeftMutator,
                &mut DeleteMutator,
                &mut DeleteMutator,
                0,
                &mut ()
            ),
            BinaryResult::New(Tree(None))
        );
        assert!(map1
            .join_mutate_in(
                &map3,
                &mut LeftMutator,
                &mut NullMutator,
                &mut NullMutator,
                0,
                &mut ()
            )
            .unwrap()
            .rec_eq(&Tree(Some(InnerMap::double_in(
                (3, 6),
                (4, 13),
                0,
                &mut ()
            )))));
    }
}