<commit_before>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::marker;
use std::fmt::Debug;
use std::marker::PhantomData;
use snapshot_vec as sv;

#[cfg(test)]
mod tests;

/// This trait is implemented by any type that can serve as a type
/// variable. We call such variables *unification keys*. For example,
/// this trait is implemented by `IntVid`, which represents integral
/// variables.
///
/// Each key type has an associated value type `V`. For example, for
/// `IntVid`, this is `Option<IntVarValue>`, representing some
/// (possibly not yet known) sort of integer.
///
/// Clients are expected to provide implementations of this trait; you
/// can see some examples in the `test` module.
pub trait UnifyKey: Copy + Clone + Debug + PartialEq {
    type Value: Clone + PartialEq + Debug;

    fn index(&self) -> u32;

    fn from_index(u: u32) -> Self;

    fn tag(k: Option<Self>) -> &'static str;
}

/// This trait is implemented for unify values that can be
/// combined. This relation should be a monoid.
pub trait Combine {
    fn combine(&self, other: &Self) -> Self;
}

impl Combine for () {
    fn combine(&self, _other: &()) {}
}

/// Value of a unification key. We implement Tarjan's union-find
/// algorithm: when two keys are unified, one of them is converted
/// into a "redirect" pointing at the other. These redirects form a
/// DAG: the roots of the DAG (nodes that are not redirected) are each
/// associated with a value of type `V` and a rank. The rank is used
/// to keep the DAG relatively balanced, which helps keep the running
/// time of the algorithm under control. For more information, see
/// <http://en.wikipedia.org/wiki/Disjoint-set_data_structure>.
#[derive(PartialEq,Clone,Debug)]
pub struct VarValue<K: UnifyKey> {
    parent: K, // if equal to self, this is a root
    value: K::Value, // value assigned (only relevant to root)
    rank: u32, // max depth (only relevant to root)
}

/// Table of unification keys and their values.
pub struct UnificationTable<K: UnifyKey> {
    /// Indicates the current value of each key.
    values: sv::SnapshotVec<Delegate<K>>,
}

/// At any time, users may snapshot a unification table.  The changes
/// made during the snapshot may either be *committed* or *rolled back*.
pub struct Snapshot<K: UnifyKey> {
    // Link snapshot to the key type `K` of the table.
    marker: marker::PhantomData<K>,
    snapshot: sv::Snapshot,
}

#[derive(Copy, Clone)]
struct Delegate<K>(PhantomData<K>);

impl<K: UnifyKey> VarValue<K> {
    fn new_var(key: K, value: K::Value) -> VarValue<K> {
        VarValue::new(key, value, 0)
    }

    fn new(parent: K, value: K::Value, rank: u32) -> VarValue<K> {
        VarValue {
            parent: parent, // this is a root
            value: value,
            rank: rank,
        }
    }

    fn redirect(self, to: K) -> VarValue<K> {
        VarValue { parent: to, ..self }
    }

    fn root(self, rank: u32, value: K::Value) -> VarValue<K> {
        VarValue {
            rank: rank,
            value: value,
            ..self
        }
    }

    /// Returns the key of this node. Only valid if this is a root
    /// node, which you yourself must ensure.
    fn key(&self) -> K {
        self.parent
    }

    fn parent(&self, self_key: K) -> Option<K> {
        self.if_not_self(self.parent, self_key)
    }

    fn if_not_self(&self, key: K, self_key: K) -> Option<K> {
        if key == self_key { None } else { Some(key) }
    }
}

// We can't use V:LatticeValue, much as I would like to,
// because frequently the pattern is that V=Option<U> for some
// other type parameter U, and we have no way to say
// Option<U>:LatticeValue.

impl<K: UnifyKey> UnificationTable<K> {
    pub fn new() -> UnificationTable<K> {
        UnificationTable { values: sv::SnapshotVec::new() }
    }

    /// Starts a new snapshot. Each snapshot must be either
    /// rolled back or committed in a "LIFO" (stack) order.
    pub fn snapshot(&mut self) -> Snapshot<K> {
        Snapshot {
            marker: marker::PhantomData::<K>,
            snapshot: self.values.start_snapshot(),
        }
    }

    /// Reverses all changes since the last snapshot. Also
    /// removes any keys that have been created since then.
    pub fn rollback_to(&mut self, snapshot: Snapshot<K>) {
        debug!("{}: rollback_to()", UnifyKey::tag(None::<K>));
        self.values.rollback_to(snapshot.snapshot);
    }

    /// Commits all changes since the last snapshot. Of course, they
    /// can still be undone if there is a snapshot further out.
    pub fn commit(&mut self, snapshot: Snapshot<K>) {
        debug!("{}: commit()", UnifyKey::tag(None::<K>));
        self.values.commit(snapshot.snapshot);
    }

    pub fn new_key(&mut self, value: K::Value) -> K {
        let len = self.values.len();
        let key: K = UnifyKey::from_index(len as u32);
        self.values.push(VarValue::new_var(key, value));
        debug!("{}: created new key: {:?}", UnifyKey::tag(None::<K>), key);
        key
    }

    /// Find the root node for `vid`. This uses the standard
    /// union-find algorithm with path compression:
    /// <http://en.wikipedia.org/wiki/Disjoint-set_data_structure>.
    ///
    /// NB. This is a building-block operation and you would probably
    /// prefer to call `probe` below.
    fn get(&mut self, vid: K) -> VarValue<K> {
        let index = vid.index() as usize;
        let mut value: VarValue<K> = self.values.get(index).clone();
        match value.parent(vid) {
            Some(redirect) => {
                let root: VarValue<K> = self.get(redirect);
                if root.key() != redirect {
                    // Path compression
                    value.parent = root.key();
                    self.values.set(index, value);
                }
                root
            }
            None => value,
        }
    }

    fn is_root(&self, key: K) -> bool {
        let index = key.index() as usize;
        self.values.get(index).parent(key).is_none()
    }

    /// Sets the value for `vid` to `new_value`. `vid` MUST be a root
    /// node! This is an internal operation used to impl other things.
    fn set(&mut self, key: K, new_value: VarValue<K>) {
        assert!(self.is_root(key));

        debug!("Updating variable {:?} to {:?}", key, new_value);

        let index = key.index() as usize;
        self.values.set(index, new_value);
    }

    /// Either redirects `node_a` to `node_b` or vice versa, depending
    /// on the relative rank. The value associated with the new root
    /// will be `new_value`.
    ///
    /// NB: This is the "union" operation of "union-find". It is
    /// really more of a building block. If the values associated with
    /// your key are non-trivial, you would probably prefer to call
    /// `unify_var_var` below.
    fn unify(&mut self, root_a: VarValue<K>, root_b: VarValue<K>, new_value: K::Value) -> K {
        debug!("unify(root_a(id={:?}, rank={:?}), root_b(id={:?}, rank={:?}))",
               root_a.key(),
               root_a.rank,
               root_b.key(),
               root_b.rank);

        if root_a.rank > root_b.rank {
            // a has greater rank, so a should become b's parent,
            // i.e., b should redirect to a.
            self.redirect_root(root_a.rank, root_b, root_a, new_value)
        } else if root_a.rank < root_b.rank {
            // b has greater rank, so a should redirect to b.
            self.redirect_root(root_b.rank, root_a, root_b, new_value)
        } else {
            // If equal, redirect one to the other and increment the
            // other's rank.
            self.redirect_root(root_a.rank + 1, root_a, root_b, new_value)
        }
    }

    fn redirect_root(&mut self,
                     new_rank: u32,
                     old_root: VarValue<K>,
                     new_root: VarValue<K>,
                     new_value: K::Value)
                     -> K {
        let old_root_key = old_root.key();
        let new_root_key = new_root.key();
        self.set(old_root_key, old_root.redirect(new_root_key));
        self.set(new_root_key, new_root.root(new_rank, new_value));
        new_root_key
    }
}

impl<K: UnifyKey> sv::SnapshotVecDelegate for Delegate<K> {
    type Value = VarValue<K>;
    type Undo = ();

    fn reverse(_: &mut Vec<VarValue<K>>, _: ()) {}
}

// # Base union-find algorithm, where we are just making sets

impl<'tcx, K: UnifyKey> UnificationTable<K>
    where K::Value: Combine
{
    pub fn union(&mut self, a_id: K, b_id: K) -> K {
        let node_a = self.get(a_id);
        let node_b = self.get(b_id);
        let a_id = node_a.key();
        let b_id = node_b.key();
        if a_id != b_id {
            let new_value = node_a.value.combine(&node_b.value);
            self.unify(node_a, node_b, new_value)
        } else {
            a_id
        }
    }

    pub fn find(&mut self, id: K) -> K {
        self.get(id).key()
    }

    pub fn find_value(&mut self, id: K) -> K::Value {
        self.get(id).value
    }

    pub fn unioned(&mut self, a_id: K, b_id: K) -> bool {
        self.find(a_id) == self.find(b_id)
    }
}

// # Non-subtyping unification
//
// Code to handle keys which carry a value, like ints,
// floats---anything that doesn't have a subtyping relationship we
// need to worry about.

impl<'tcx, K, V> UnificationTable<K>
    where K: UnifyKey<Value = Option<V>>,
          V: Clone + PartialEq + Debug
{
    pub fn unify_var_var(&mut self, a_id: K, b_id: K) -> Result<K, (V, V)> {
        let node_a = self.get(a_id);
        let node_b = self.get(b_id);
        let a_id = node_a.key();
        let b_id = node_b.key();

        if a_id == b_id {
            return Ok(a_id);
        }

        let combined = {
            match (&node_a.value, &node_b.value) {
                (&None, &None) => None,
                (&Some(ref v), &None) |
                (&None, &Some(ref v)) => Some(v.clone()),
                (&Some(ref v1), &Some(ref v2)) => {
                    if *v1 != *v2 {
                        return Err((v1.clone(), v2.clone()));
                    }
                    Some(v1.clone())
                }
            }
        };

        Ok(self.unify(node_a, node_b, combined))
    }

    /// Sets the value of the key `a_id` to `b`. Because simple keys do not have any subtyping
    /// relationships, if `a_id` already has a value, it must be the same as `b`.
    pub fn unify_var_value(&mut self, a_id: K, b: V) -> Result<(), (V, V)> {
        let mut node_a = self.get(a_id);

        match node_a.value {
            None => {
                node_a.value = Some(b);
                self.set(node_a.key(), node_a);
                Ok(())
            }

            Some(ref a_t) => {
                if *a_t == b {
                    Ok(())
                } else {
                    Err((a_t.clone(), b))
                }
            }
        }
    }

    pub fn has_value(&mut self, id: K) -> bool {
        self.get(id).value.is_some()
    }

    pub fn probe(&mut self, a_id: K) -> Option<V> {
        self.get(a_id).value.clone()
    }

    pub fn unsolved_variables(&mut self) -> Vec<K> {
        self.values
            .iter()
            .filter_map(|vv| {
                if vv.value.is_some() {
                    None
                } else {
                    Some(vv.key())
                }
            })
            .collect()
    }
}
<commit_msg>Rollup merge of #37848 - nnethercote:UnificationTable-probe, r=arielb1<commit_after>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::marker;
use std::fmt::Debug;
use std::marker::PhantomData;
use snapshot_vec as sv;

#[cfg(test)]
mod tests;

/// This trait is implemented by any type that can serve as a type
/// variable. We call such variables *unification keys*. For example,
/// this trait is implemented by `IntVid`, which represents integral
/// variables.
///
/// Each key type has an associated value type `V`. For example, for
/// `IntVid`, this is `Option<IntVarValue>`, representing some
/// (possibly not yet known) sort of integer.
///
/// Clients are expected to provide implementations of this trait; you
/// can see some examples in the `test` module.
pub trait UnifyKey: Copy + Clone + Debug + PartialEq {
    type Value: Clone + PartialEq + Debug;

    fn index(&self) -> u32;

    fn from_index(u: u32) -> Self;

    fn tag(k: Option<Self>) -> &'static str;
}

/// This trait is implemented for unify values that can be
/// combined. This relation should be a monoid.
pub trait Combine {
    fn combine(&self, other: &Self) -> Self;
}

impl Combine for () {
    fn combine(&self, _other: &()) {}
}

/// Value of a unification key. We implement Tarjan's union-find
/// algorithm: when two keys are unified, one of them is converted
/// into a "redirect" pointing at the other. These redirects form a
/// DAG: the roots of the DAG (nodes that are not redirected) are each
/// associated with a value of type `V` and a rank. The rank is used
/// to keep the DAG relatively balanced, which helps keep the running
/// time of the algorithm under control. For more information, see
/// <http://en.wikipedia.org/wiki/Disjoint-set_data_structure>.
#[derive(PartialEq,Clone,Debug)]
pub struct VarValue<K: UnifyKey> {
    parent: K, // if equal to self, this is a root
    value: K::Value, // value assigned (only relevant to root)
    rank: u32, // max depth (only relevant to root)
}

/// Table of unification keys and their values.
pub struct UnificationTable<K: UnifyKey> {
    /// Indicates the current value of each key.
    values: sv::SnapshotVec<Delegate<K>>,
}

/// At any time, users may snapshot a unification table.  The changes
/// made during the snapshot may either be *committed* or *rolled back*.
pub struct Snapshot<K: UnifyKey> {
    // Link snapshot to the key type `K` of the table.
    marker: marker::PhantomData<K>,
    snapshot: sv::Snapshot,
}

#[derive(Copy, Clone)]
struct Delegate<K>(PhantomData<K>);

impl<K: UnifyKey> VarValue<K> {
    fn new_var(key: K, value: K::Value) -> VarValue<K> {
        VarValue::new(key, value, 0)
    }

    fn new(parent: K, value: K::Value, rank: u32) -> VarValue<K> {
        VarValue {
            parent: parent, // this is a root
            value: value,
            rank: rank,
        }
    }

    fn redirect(self, to: K) -> VarValue<K> {
        VarValue { parent: to, ..self }
    }

    fn root(self, rank: u32, value: K::Value) -> VarValue<K> {
        VarValue {
            rank: rank,
            value: value,
            ..self
        }
    }

    /// Returns the key of this node. Only valid if this is a root
    /// node, which you yourself must ensure.
    fn key(&self) -> K {
        self.parent
    }

    fn parent(&self, self_key: K) -> Option<K> {
        self.if_not_self(self.parent, self_key)
    }

    fn if_not_self(&self, key: K, self_key: K) -> Option<K> {
        if key == self_key { None } else { Some(key) }
    }
}

// We can't use V:LatticeValue, much as I would like to,
// because frequently the pattern is that V=Option<U> for some
// other type parameter U, and we have no way to say
// Option<U>:LatticeValue.

impl<K: UnifyKey> UnificationTable<K> {
    pub fn new() -> UnificationTable<K> {
        UnificationTable { values: sv::SnapshotVec::new() }
    }

    /// Starts a new snapshot. Each snapshot must be either
    /// rolled back or committed in a "LIFO" (stack) order.
    pub fn snapshot(&mut self) -> Snapshot<K> {
        Snapshot {
            marker: marker::PhantomData::<K>,
            snapshot: self.values.start_snapshot(),
        }
    }

    /// Reverses all changes since the last snapshot. Also
    /// removes any keys that have been created since then.
    pub fn rollback_to(&mut self, snapshot: Snapshot<K>) {
        debug!("{}: rollback_to()", UnifyKey::tag(None::<K>));
        self.values.rollback_to(snapshot.snapshot);
    }

    /// Commits all changes since the last snapshot. Of course, they
    /// can still be undone if there is a snapshot further out.
    pub fn commit(&mut self, snapshot: Snapshot<K>) {
        debug!("{}: commit()", UnifyKey::tag(None::<K>));
        self.values.commit(snapshot.snapshot);
    }

    pub fn new_key(&mut self, value: K::Value) -> K {
        let len = self.values.len();
        let key: K = UnifyKey::from_index(len as u32);
        self.values.push(VarValue::new_var(key, value));
        debug!("{}: created new key: {:?}", UnifyKey::tag(None::<K>), key);
        key
    }

    /// Find the root node for `vid`. This uses the standard
    /// union-find algorithm with path compression:
    /// <http://en.wikipedia.org/wiki/Disjoint-set_data_structure>.
    ///
    /// NB. This is a building-block operation and you would probably
    /// prefer to call `probe` below.
    fn get(&mut self, vid: K) -> VarValue<K> {
        let index = vid.index() as usize;
        let mut value: VarValue<K> = self.values.get(index).clone();
        match value.parent(vid) {
            Some(redirect) => {
                let root: VarValue<K> = self.get(redirect);
                if root.key() != redirect {
                    // Path compression
                    value.parent = root.key();
                    self.values.set(index, value);
                }
                root
            }
            None => value,
        }
    }

    fn is_root(&self, key: K) -> bool {
        let index = key.index() as usize;
        self.values.get(index).parent(key).is_none()
    }

    /// Sets the value for `vid` to `new_value`. `vid` MUST be a root
    /// node! This is an internal operation used to impl other things.
    fn set(&mut self, key: K, new_value: VarValue<K>) {
        assert!(self.is_root(key));

        debug!("Updating variable {:?} to {:?}", key, new_value);

        let index = key.index() as usize;
        self.values.set(index, new_value);
    }

    /// Either redirects `node_a` to `node_b` or vice versa, depending
    /// on the relative rank. The value associated with the new root
    /// will be `new_value`.
    ///
    /// NB: This is the "union" operation of "union-find". It is
    /// really more of a building block. If the values associated with
    /// your key are non-trivial, you would probably prefer to call
    /// `unify_var_var` below.
    fn unify(&mut self, root_a: VarValue<K>, root_b: VarValue<K>, new_value: K::Value) -> K {
        debug!("unify(root_a(id={:?}, rank={:?}), root_b(id={:?}, rank={:?}))",
               root_a.key(),
               root_a.rank,
               root_b.key(),
               root_b.rank);

        if root_a.rank > root_b.rank {
            // a has greater rank, so a should become b's parent,
            // i.e., b should redirect to a.
            self.redirect_root(root_a.rank, root_b, root_a, new_value)
        } else if root_a.rank < root_b.rank {
            // b has greater rank, so a should redirect to b.
            self.redirect_root(root_b.rank, root_a, root_b, new_value)
        } else {
            // If equal, redirect one to the other and increment the
            // other's rank.
            self.redirect_root(root_a.rank + 1, root_a, root_b, new_value)
        }
    }

    fn redirect_root(&mut self,
                     new_rank: u32,
                     old_root: VarValue<K>,
                     new_root: VarValue<K>,
                     new_value: K::Value)
                     -> K {
        let old_root_key = old_root.key();
        let new_root_key = new_root.key();
        self.set(old_root_key, old_root.redirect(new_root_key));
        self.set(new_root_key, new_root.root(new_rank, new_value));
        new_root_key
    }
}

impl<K: UnifyKey> sv::SnapshotVecDelegate for Delegate<K> {
    type Value = VarValue<K>;
    type Undo = ();

    fn reverse(_: &mut Vec<VarValue<K>>, _: ()) {}
}

// # Base union-find algorithm, where we are just making sets

impl<'tcx, K: UnifyKey> UnificationTable<K>
    where K::Value: Combine
{
    pub fn union(&mut self, a_id: K, b_id: K) -> K {
        let node_a = self.get(a_id);
        let node_b = self.get(b_id);
        let a_id = node_a.key();
        let b_id = node_b.key();
        if a_id != b_id {
            let new_value = node_a.value.combine(&node_b.value);
            self.unify(node_a, node_b, new_value)
        } else {
            a_id
        }
    }

    pub fn find(&mut self, id: K) -> K {
        self.get(id).key()
    }

    pub fn find_value(&mut self, id: K) -> K::Value {
        self.get(id).value
    }

    pub fn unioned(&mut self, a_id: K, b_id: K) -> bool {
        self.find(a_id) == self.find(b_id)
    }
}

// # Non-subtyping unification
//
// Code to handle keys which carry a value, like ints,
// floats---anything that doesn't have a subtyping relationship we
// need to worry about.

impl<'tcx, K, V> UnificationTable<K>
    where K: UnifyKey<Value = Option<V>>,
          V: Clone + PartialEq + Debug
{
    pub fn unify_var_var(&mut self, a_id: K, b_id: K) -> Result<K, (V, V)> {
        let node_a = self.get(a_id);
        let node_b = self.get(b_id);
        let a_id = node_a.key();
        let b_id = node_b.key();

        if a_id == b_id {
            return Ok(a_id);
        }

        let combined = {
            match (&node_a.value, &node_b.value) {
                (&None, &None) => None,
                (&Some(ref v), &None) |
                (&None, &Some(ref v)) => Some(v.clone()),
                (&Some(ref v1), &Some(ref v2)) => {
                    if *v1 != *v2 {
                        return Err((v1.clone(), v2.clone()));
                    }
                    Some(v1.clone())
                }
            }
        };

        Ok(self.unify(node_a, node_b, combined))
    }

    /// Sets the value of the key `a_id` to `b`. Because simple keys do not have any subtyping
    /// relationships, if `a_id` already has a value, it must be the same as `b`.
    pub fn unify_var_value(&mut self, a_id: K, b: V) -> Result<(), (V, V)> {
        let mut node_a = self.get(a_id);

        match node_a.value {
            None => {
                node_a.value = Some(b);
                self.set(node_a.key(), node_a);
                Ok(())
            }

            Some(ref a_t) => {
                if *a_t == b {
                    Ok(())
                } else {
                    Err((a_t.clone(), b))
                }
            }
        }
    }

    pub fn has_value(&mut self, id: K) -> bool {
        self.get(id).value.is_some()
    }

    pub fn probe(&mut self, a_id: K) -> Option<V> {
        self.get(a_id).value
    }

    pub fn unsolved_variables(&mut self) -> Vec<K> {
        self.values
            .iter()
            .filter_map(|vv| {
                if vv.value.is_some() {
                    None
                } else {
                    Some(vv.key())
                }
            })
            .collect()
    }
}
<|endoftext|>


<commit_before>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::marker;
use std::fmt::Debug;
use std::marker::PhantomData;
use snapshot_vec as sv;

#[cfg(test)]
mod tests;

/// This trait is implemented by any type that can serve as a type
/// variable. We call such variables *unification keys*. For example,
/// this trait is implemented by `IntVid`, which represents integral
/// variables.
///
/// Each key type has an associated value type `V`. For example, for
/// `IntVid`, this is `Option<IntVarValue>`, representing some
/// (possibly not yet known) sort of integer.
///
/// Clients are expected to provide implementations of this trait; you
/// can see some examples in the `test` module.
pub trait UnifyKey : Copy + Clone + Debug + PartialEq {
    type Value: Clone + PartialEq + Debug;

    fn index(&self) -> u32;

    fn from_index(u: u32) -> Self;

    fn tag(k: Option<Self>) -> &'static str;
}

/// This trait is implemented for unify values that can be
/// combined. This relation should be a monoid.
pub trait Combine {
    fn combine(&self, other: &Self) -> Self;
}

impl Combine for () {
    fn combine(&self, _other: &()) {}
}

/// Value of a unification key. We implement Tarjan's union-find
/// algorithm: when two keys are unified, one of them is converted
/// into a "redirect" pointing at the other. These redirects form a
/// DAG: the roots of the DAG (nodes that are not redirected) are each
/// associated with a value of type `V` and a rank. The rank is used
/// to keep the DAG relatively balanced, which helps keep the running
/// time of the algorithm under control. For more information, see
/// <http://en.wikipedia.org/wiki/Disjoint-set_data_structure>.
#[derive(PartialEq,Clone,Debug)]
pub struct VarValue<K: UnifyKey> {
    parent: K, // if equal to self, this is a root
    value: K::Value, // value assigned (only relevant to root)
    rank: u32, // max depth (only relevant to root)
}

/// Table of unification keys and their values.
pub struct UnificationTable<K: UnifyKey> {
    /// Indicates the current value of each key.
    values: sv::SnapshotVec<Delegate<K>>,
}

/// At any time, users may snapshot a unification table.  The changes
/// made during the snapshot may either be *committed* or *rolled back*.
pub struct Snapshot<K: UnifyKey> {
    // Link snapshot to the key type `K` of the table.
    marker: marker::PhantomData<K>,
    snapshot: sv::Snapshot,
}

#[derive(Copy, Clone)]
struct Delegate<K>(PhantomData<K>);

impl<K: UnifyKey> VarValue<K> {
    fn new_var(key: K, value: K::Value) -> VarValue<K> {
        VarValue::new(key, value, 0)
    }

    fn new(parent: K, value: K::Value, rank: u32) -> VarValue<K> {
        VarValue {
            parent: parent, // this is a root
            value: value,
            rank: rank,
        }
    }

    fn redirect(self, to: K) -> VarValue<K> {
        VarValue { parent: to, ..self }
    }

    fn root(self, rank: u32, value: K::Value) -> VarValue<K> {
        VarValue {
            rank: rank,
            value: value,
            ..self
        }
    }

    /// Returns the key of this node. Only valid if this is a root
    /// node, which you yourself must ensure.
    fn key(&self) -> K {
        self.parent
    }

    fn parent(&self, self_key: K) -> Option<K> {
        self.if_not_self(self.parent, self_key)
    }

    fn if_not_self(&self, key: K, self_key: K) -> Option<K> {
        if key == self_key {
            None
        } else {
            Some(key)
        }
    }
}

// We can't use V:LatticeValue, much as I would like to,
// because frequently the pattern is that V=Option<U> for some
// other type parameter U, and we have no way to say
// Option<U>:LatticeValue.

impl<K: UnifyKey> UnificationTable<K> {
    pub fn new() -> UnificationTable<K> {
        UnificationTable { values: sv::SnapshotVec::new() }
    }

    /// Starts a new snapshot. Each snapshot must be either
    /// rolled back or committed in a "LIFO" (stack) order.
    pub fn snapshot(&mut self) -> Snapshot<K> {
        Snapshot {
            marker: marker::PhantomData::<K>,
            snapshot: self.values.start_snapshot(),
        }
    }

    /// Reverses all changes since the last snapshot. Also
    /// removes any keys that have been created since then.
    pub fn rollback_to(&mut self, snapshot: Snapshot<K>) {
        debug!("{}: rollback_to()", UnifyKey::tag(None::<K>));
        self.values.rollback_to(snapshot.snapshot);
    }

    /// Commits all changes since the last snapshot. Of course, they
    /// can still be undone if there is a snapshot further out.
    pub fn commit(&mut self, snapshot: Snapshot<K>) {
        debug!("{}: commit()", UnifyKey::tag(None::<K>));
        self.values.commit(snapshot.snapshot);
    }

    pub fn new_key(&mut self, value: K::Value) -> K {
        let len = self.values.len();
        let key: K = UnifyKey::from_index(len as u32);
        self.values.push(VarValue::new_var(key, value));
        debug!("{}: created new key: {:?}", UnifyKey::tag(None::<K>), key);
        key
    }

    /// Find the root node for `vid`. This uses the standard
    /// union-find algorithm with path compression:
    /// <http://en.wikipedia.org/wiki/Disjoint-set_data_structure>.
    ///
    /// NB. This is a building-block operation and you would probably
    /// prefer to call `probe` below.
    fn get(&mut self, vid: K) -> VarValue<K> {
        let index = vid.index() as usize;
        let mut value: VarValue<K> = self.values.get(index).clone();
        match value.parent(vid) {
            Some(redirect) => {
                let root: VarValue<K> = self.get(redirect);
                if root.key() != redirect {
                    // Path compression
                    value.parent = root.key();
                    self.values.set(index, value);
                }
                root
            }
            None => value,
        }
    }

    fn is_root(&self, key: K) -> bool {
        let index = key.index() as usize;
        self.values.get(index).parent(key).is_none()
    }

    /// Sets the value for `vid` to `new_value`. `vid` MUST be a root
    /// node! This is an internal operation used to impl other things.
    fn set(&mut self, key: K, new_value: VarValue<K>) {
        assert!(self.is_root(key));

        debug!("Updating variable {:?} to {:?}", key, new_value);

        let index = key.index() as usize;
        self.values.set(index, new_value);
    }

    /// Either redirects `node_a` to `node_b` or vice versa, depending
    /// on the relative rank. The value associated with the new root
    /// will be `new_value`.
    ///
    /// NB: This is the "union" operation of "union-find". It is
    /// really more of a building block. If the values associated with
    /// your key are non-trivial, you would probably prefer to call
    /// `unify_var_var` below.
    fn unify(&mut self, root_a: VarValue<K>, root_b: VarValue<K>, new_value: K::Value) -> K {
        debug!("unify(root_a(id={:?}, rank={:?}), root_b(id={:?}, rank={:?}))",
               root_a.key(),
               root_a.rank,
               root_b.key(),
               root_b.rank);

        if root_a.rank > root_b.rank {
            // a has greater rank, so a should become b's parent,
            // i.e., b should redirect to a.
            self.redirect_root(root_a.rank, root_b, root_a, new_value)
        } else if root_a.rank < root_b.rank {
            // b has greater rank, so a should redirect to b.
            self.redirect_root(root_b.rank, root_a, root_b, new_value)
        } else {
            // If equal, redirect one to the other and increment the
            // other's rank.
            self.redirect_root(root_a.rank + 1, root_a, root_b, new_value)
        }
    }

    fn redirect_root(&mut self,
                     new_rank: u32,
                     old_root: VarValue<K>,
                     new_root: VarValue<K>,
                     new_value: K::Value) -> K {
        let old_root_key = old_root.key();
        let new_root_key = new_root.key();
        self.set(old_root_key, old_root.redirect(new_root_key));
        self.set(new_root_key, new_root.root(new_rank, new_value));
        new_root_key
    }
}

impl<K: UnifyKey> sv::SnapshotVecDelegate for Delegate<K> {
    type Value = VarValue<K>;
    type Undo = ();

    fn reverse(_: &mut Vec<VarValue<K>>, _: ()) {}
}

// # Base union-find algorithm, where we are just making sets

impl<'tcx, K: UnifyKey> UnificationTable<K>
    where K::Value: Combine
{
    pub fn union(&mut self, a_id: K, b_id: K) -> K {
        let node_a = self.get(a_id);
        let node_b = self.get(b_id);
        let a_id = node_a.key();
        let b_id = node_b.key();
        if a_id != b_id {
            let new_value = node_a.value.combine(&node_b.value);
            self.unify(node_a, node_b, new_value)
        } else {
            a_id
        }
    }

    pub fn find(&mut self, id: K) -> K {
        self.get(id).key()
    }

    pub fn find_value(&mut self, id: K) -> K::Value {
        self.get(id).value
    }

    pub fn unioned(&mut self, a_id: K, b_id: K) -> bool {
        self.find(a_id) == self.find(b_id)
    }
}

// # Non-subtyping unification
//
// Code to handle keys which carry a value, like ints,
// floats---anything that doesn't have a subtyping relationship we
// need to worry about.

impl<'tcx, K, V> UnificationTable<K>
    where K: UnifyKey<Value = Option<V>>,
          V: Clone + PartialEq + Debug
{
    pub fn unify_var_var(&mut self, a_id: K, b_id: K) -> Result<K, (V, V)> {
        let node_a = self.get(a_id);
        let node_b = self.get(b_id);
        let a_id = node_a.key();
        let b_id = node_b.key();

        if a_id == b_id {
            return Ok(a_id);
        }

        let combined = {
            match (&node_a.value, &node_b.value) {
                (&None, &None) => None,
                (&Some(ref v), &None) | (&None, &Some(ref v)) => Some(v.clone()),
                (&Some(ref v1), &Some(ref v2)) => {
                    if *v1 != *v2 {
                        return Err((v1.clone(), v2.clone()));
                    }
                    Some(v1.clone())
                }
            }
        };

        Ok(self.unify(node_a, node_b, combined))
    }

    /// Sets the value of the key `a_id` to `b`. Because simple keys do not have any subtyping
    /// relationships, if `a_id` already has a value, it must be the same as `b`.
    pub fn unify_var_value(&mut self, a_id: K, b: V) -> Result<(), (V, V)> {
        let mut node_a = self.get(a_id);

        match node_a.value {
            None => {
                node_a.value = Some(b);
                self.set(node_a.key(), node_a);
                Ok(())
            }

            Some(ref a_t) => {
                if *a_t == b {
                    Ok(())
                } else {
                    Err((a_t.clone(), b))
                }
            }
        }
    }

    pub fn has_value(&mut self, id: K) -> bool {
        self.get(id).value.is_some()
    }

    pub fn probe(&mut self, a_id: K) -> Option<V> {
        self.get(a_id).value.clone()
    }

    pub fn unsolved_variables(&mut self) -> Vec<K> {
        self.values
            .iter()
            .filter_map(|vv| {
                if vv.value.is_some() {
                    None
                } else {
                    Some(vv.key())
                }
            })
            .collect()
    }
}
<commit_msg>run rustfmt on unify folder<commit_after>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::marker;
use std::fmt::Debug;
use std::marker::PhantomData;
use snapshot_vec as sv;

#[cfg(test)]
mod tests;

/// This trait is implemented by any type that can serve as a type
/// variable. We call such variables *unification keys*. For example,
/// this trait is implemented by `IntVid`, which represents integral
/// variables.
///
/// Each key type has an associated value type `V`. For example, for
/// `IntVid`, this is `Option<IntVarValue>`, representing some
/// (possibly not yet known) sort of integer.
///
/// Clients are expected to provide implementations of this trait; you
/// can see some examples in the `test` module.
pub trait UnifyKey: Copy + Clone + Debug + PartialEq {
    type Value: Clone + PartialEq + Debug;

    fn index(&self) -> u32;

    fn from_index(u: u32) -> Self;

    fn tag(k: Option<Self>) -> &'static str;
}

/// This trait is implemented for unify values that can be
/// combined. This relation should be a monoid.
pub trait Combine {
    fn combine(&self, other: &Self) -> Self;
}

impl Combine for () {
    fn combine(&self, _other: &()) {}
}

/// Value of a unification key. We implement Tarjan's union-find
/// algorithm: when two keys are unified, one of them is converted
/// into a "redirect" pointing at the other. These redirects form a
/// DAG: the roots of the DAG (nodes that are not redirected) are each
/// associated with a value of type `V` and a rank. The rank is used
/// to keep the DAG relatively balanced, which helps keep the running
/// time of the algorithm under control. For more information, see
/// <http://en.wikipedia.org/wiki/Disjoint-set_data_structure>.
#[derive(PartialEq,Clone,Debug)]
pub struct VarValue<K: UnifyKey> {
    parent: K, // if equal to self, this is a root
    value: K::Value, // value assigned (only relevant to root)
    rank: u32, // max depth (only relevant to root)
}

/// Table of unification keys and their values.
pub struct UnificationTable<K: UnifyKey> {
    /// Indicates the current value of each key.
    values: sv::SnapshotVec<Delegate<K>>,
}

/// At any time, users may snapshot a unification table.  The changes
/// made during the snapshot may either be *committed* or *rolled back*.
pub struct Snapshot<K: UnifyKey> {
    // Link snapshot to the key type `K` of the table.
    marker: marker::PhantomData<K>,
    snapshot: sv::Snapshot,
}

#[derive(Copy, Clone)]
struct Delegate<K>(PhantomData<K>);

impl<K: UnifyKey> VarValue<K> {
    fn new_var(key: K, value: K::Value) -> VarValue<K> {
        VarValue::new(key, value, 0)
    }

    fn new(parent: K, value: K::Value, rank: u32) -> VarValue<K> {
        VarValue {
            parent: parent, // this is a root
            value: value,
            rank: rank,
        }
    }

    fn redirect(self, to: K) -> VarValue<K> {
        VarValue { parent: to, ..self }
    }

    fn root(self, rank: u32, value: K::Value) -> VarValue<K> {
        VarValue {
            rank: rank,
            value: value,
            ..self
        }
    }

    /// Returns the key of this node. Only valid if this is a root
    /// node, which you yourself must ensure.
    fn key(&self) -> K {
        self.parent
    }

    fn parent(&self, self_key: K) -> Option<K> {
        self.if_not_self(self.parent, self_key)
    }

    fn if_not_self(&self, key: K, self_key: K) -> Option<K> {
        if key == self_key { None } else { Some(key) }
    }
}

// We can't use V:LatticeValue, much as I would like to,
// because frequently the pattern is that V=Option<U> for some
// other type parameter U, and we have no way to say
// Option<U>:LatticeValue.

impl<K: UnifyKey> UnificationTable<K> {
    pub fn new() -> UnificationTable<K> {
        UnificationTable { values: sv::SnapshotVec::new() }
    }

    /// Starts a new snapshot. Each snapshot must be either
    /// rolled back or committed in a "LIFO" (stack) order.
    pub fn snapshot(&mut self) -> Snapshot<K> {
        Snapshot {
            marker: marker::PhantomData::<K>,
            snapshot: self.values.start_snapshot(),
        }
    }

    /// Reverses all changes since the last snapshot. Also
    /// removes any keys that have been created since then.
    pub fn rollback_to(&mut self, snapshot: Snapshot<K>) {
        debug!("{}: rollback_to()", UnifyKey::tag(None::<K>));
        self.values.rollback_to(snapshot.snapshot);
    }

    /// Commits all changes since the last snapshot. Of course, they
    /// can still be undone if there is a snapshot further out.
    pub fn commit(&mut self, snapshot: Snapshot<K>) {
        debug!("{}: commit()", UnifyKey::tag(None::<K>));
        self.values.commit(snapshot.snapshot);
    }

    pub fn new_key(&mut self, value: K::Value) -> K {
        let len = self.values.len();
        let key: K = UnifyKey::from_index(len as u32);
        self.values.push(VarValue::new_var(key, value));
        debug!("{}: created new key: {:?}", UnifyKey::tag(None::<K>), key);
        key
    }

    /// Find the root node for `vid`. This uses the standard
    /// union-find algorithm with path compression:
    /// <http://en.wikipedia.org/wiki/Disjoint-set_data_structure>.
    ///
    /// NB. This is a building-block operation and you would probably
    /// prefer to call `probe` below.
    fn get(&mut self, vid: K) -> VarValue<K> {
        let index = vid.index() as usize;
        let mut value: VarValue<K> = self.values.get(index).clone();
        match value.parent(vid) {
            Some(redirect) => {
                let root: VarValue<K> = self.get(redirect);
                if root.key() != redirect {
                    // Path compression
                    value.parent = root.key();
                    self.values.set(index, value);
                }
                root
            }
            None => value,
        }
    }

    fn is_root(&self, key: K) -> bool {
        let index = key.index() as usize;
        self.values.get(index).parent(key).is_none()
    }

    /// Sets the value for `vid` to `new_value`. `vid` MUST be a root
    /// node! This is an internal operation used to impl other things.
    fn set(&mut self, key: K, new_value: VarValue<K>) {
        assert!(self.is_root(key));

        debug!("Updating variable {:?} to {:?}", key, new_value);

        let index = key.index() as usize;
        self.values.set(index, new_value);
    }

    /// Either redirects `node_a` to `node_b` or vice versa, depending
    /// on the relative rank. The value associated with the new root
    /// will be `new_value`.
    ///
    /// NB: This is the "union" operation of "union-find". It is
    /// really more of a building block. If the values associated with
    /// your key are non-trivial, you would probably prefer to call
    /// `unify_var_var` below.
    fn unify(&mut self, root_a: VarValue<K>, root_b: VarValue<K>, new_value: K::Value) -> K {
        debug!("unify(root_a(id={:?}, rank={:?}), root_b(id={:?}, rank={:?}))",
               root_a.key(),
               root_a.rank,
               root_b.key(),
               root_b.rank);

        if root_a.rank > root_b.rank {
            // a has greater rank, so a should become b's parent,
            // i.e., b should redirect to a.
            self.redirect_root(root_a.rank, root_b, root_a, new_value)
        } else if root_a.rank < root_b.rank {
            // b has greater rank, so a should redirect to b.
            self.redirect_root(root_b.rank, root_a, root_b, new_value)
        } else {
            // If equal, redirect one to the other and increment the
            // other's rank.
            self.redirect_root(root_a.rank + 1, root_a, root_b, new_value)
        }
    }

    fn redirect_root(&mut self,
                     new_rank: u32,
                     old_root: VarValue<K>,
                     new_root: VarValue<K>,
                     new_value: K::Value)
                     -> K {
        let old_root_key = old_root.key();
        let new_root_key = new_root.key();
        self.set(old_root_key, old_root.redirect(new_root_key));
        self.set(new_root_key, new_root.root(new_rank, new_value));
        new_root_key
    }
}

impl<K: UnifyKey> sv::SnapshotVecDelegate for Delegate<K> {
    type Value = VarValue<K>;
    type Undo = ();

    fn reverse(_: &mut Vec<VarValue<K>>, _: ()) {}
}

// # Base union-find algorithm, where we are just making sets

impl<'tcx, K: UnifyKey> UnificationTable<K>
    where K::Value: Combine
{
    pub fn union(&mut self, a_id: K, b_id: K) -> K {
        let node_a = self.get(a_id);
        let node_b = self.get(b_id);
        let a_id = node_a.key();
        let b_id = node_b.key();
        if a_id != b_id {
            let new_value = node_a.value.combine(&node_b.value);
            self.unify(node_a, node_b, new_value)
        } else {
            a_id
        }
    }

    pub fn find(&mut self, id: K) -> K {
        self.get(id).key()
    }

    pub fn find_value(&mut self, id: K) -> K::Value {
        self.get(id).value
    }

    pub fn unioned(&mut self, a_id: K, b_id: K) -> bool {
        self.find(a_id) == self.find(b_id)
    }
}

// # Non-subtyping unification
//
// Code to handle keys which carry a value, like ints,
// floats---anything that doesn't have a subtyping relationship we
// need to worry about.

impl<'tcx, K, V> UnificationTable<K>
    where K: UnifyKey<Value = Option<V>>,
          V: Clone + PartialEq + Debug
{
    pub fn unify_var_var(&mut self, a_id: K, b_id: K) -> Result<K, (V, V)> {
        let node_a = self.get(a_id);
        let node_b = self.get(b_id);
        let a_id = node_a.key();
        let b_id = node_b.key();

        if a_id == b_id {
            return Ok(a_id);
        }

        let combined = {
            match (&node_a.value, &node_b.value) {
                (&None, &None) => None,
                (&Some(ref v), &None) |
                (&None, &Some(ref v)) => Some(v.clone()),
                (&Some(ref v1), &Some(ref v2)) => {
                    if *v1 != *v2 {
                        return Err((v1.clone(), v2.clone()));
                    }
                    Some(v1.clone())
                }
            }
        };

        Ok(self.unify(node_a, node_b, combined))
    }

    /// Sets the value of the key `a_id` to `b`. Because simple keys do not have any subtyping
    /// relationships, if `a_id` already has a value, it must be the same as `b`.
    pub fn unify_var_value(&mut self, a_id: K, b: V) -> Result<(), (V, V)> {
        let mut node_a = self.get(a_id);

        match node_a.value {
            None => {
                node_a.value = Some(b);
                self.set(node_a.key(), node_a);
                Ok(())
            }

            Some(ref a_t) => {
                if *a_t == b {
                    Ok(())
                } else {
                    Err((a_t.clone(), b))
                }
            }
        }
    }

    pub fn has_value(&mut self, id: K) -> bool {
        self.get(id).value.is_some()
    }

    pub fn probe(&mut self, a_id: K) -> Option<V> {
        self.get(a_id).value.clone()
    }

    pub fn unsolved_variables(&mut self) -> Vec<K> {
        self.values
            .iter()
            .filter_map(|vv| {
                if vv.value.is_some() {
                    None
                } else {
                    Some(vv.key())
                }
            })
            .collect()
    }
}
<|endoftext|>


<commit_before>mod render;

#[cfg(test)]
mod tests;

use std::iter;

use either::Either;
use hir::{HasSource, Semantics};
use ide_db::{
    base_db::FileRange,
    defs::{Definition, IdentClass},
    helpers::{pick_best_token, FamousDefs},
    FxIndexSet, RootDatabase,
};
use itertools::Itertools;
use syntax::{ast, match_ast, AstNode, SyntaxKind::*, SyntaxNode, SyntaxToken, T};

use crate::{
    doc_links::token_as_doc_comment,
    markup::Markup,
    runnables::{runnable_fn, runnable_mod},
    FileId, FilePosition, NavigationTarget, RangeInfo, Runnable, TryToNav,
};
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct HoverConfig {
    pub links_in_hover: bool,
    pub documentation: Option<HoverDocFormat>,
}

impl HoverConfig {
    fn markdown(&self) -> bool {
        matches!(self.documentation, Some(HoverDocFormat::Markdown))
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum HoverDocFormat {
    Markdown,
    PlainText,
}

#[derive(Debug, Clone)]
pub enum HoverAction {
    Runnable(Runnable),
    Implementation(FilePosition),
    Reference(FilePosition),
    GoToType(Vec<HoverGotoTypeData>),
}

impl HoverAction {
    fn goto_type_from_targets(db: &RootDatabase, targets: Vec<hir::ModuleDef>) -> Self {
        let targets = targets
            .into_iter()
            .filter_map(|it| {
                Some(HoverGotoTypeData {
                    mod_path: render::path(
                        db,
                        it.module(db)?,
                        it.name(db).map(|name| name.to_string()),
                    ),
                    nav: it.try_to_nav(db)?,
                })
            })
            .collect();
        HoverAction::GoToType(targets)
    }
}

#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub struct HoverGotoTypeData {
    pub mod_path: String,
    pub nav: NavigationTarget,
}

/// Contains the results when hovering over an item
#[derive(Debug, Default)]
pub struct HoverResult {
    pub markup: Markup,
    pub actions: Vec<HoverAction>,
}

// Feature: Hover
//
// Shows additional information, like the type of an expression or the documentation for a definition when "focusing" code.
// Focusing is usually hovering with a mouse, but can also be triggered with a shortcut.
//
// image::https://user-images.githubusercontent.com/48062697/113020658-b5f98b80-917a-11eb-9f88-3dbc27320c95.gif[]
pub(crate) fn hover(
    db: &RootDatabase,
    FileRange { file_id, range }: FileRange,
    config: &HoverConfig,
) -> Option<RangeInfo<HoverResult>> {
    let sema = &hir::Semantics::new(db);
    let file = sema.parse(file_id).syntax().clone();

    let offset = if !range.is_empty() {
        return hover_ranged(&file, range, sema, config);
    } else {
        range.start()
    };

    let original_token = pick_best_token(file.token_at_offset(offset), |kind| match kind {
        IDENT | INT_NUMBER | LIFETIME_IDENT | T![self] | T![super] | T![crate] => 3,
        T!['('] | T![')'] => 2,
        kind if kind.is_trivia() => 0,
        _ => 1,
    })?;

    if let Some(doc_comment) = token_as_doc_comment(&original_token) {
        cov_mark::hit!(no_highlight_on_comment_hover);
        return doc_comment.get_definition_with_descend_at(sema, offset, |def, node, range| {
            let res = hover_for_definition(sema, file_id, def, &node, config)?;
            Some(RangeInfo::new(range, res))
        });
    }

    let descended = sema.descend_into_macros(original_token.clone());

    // FIXME: Definition should include known lints and the like instead of having this special case here
    if let Some(res) = descended.iter().find_map(|token| {
        let attr = token.ancestors().find_map(ast::Attr::cast)?;
        render::try_for_lint(&attr, token)
    }) {
        return Some(RangeInfo::new(original_token.text_range(), res));
    }

    let result = descended
        .iter()
        .filter_map(|token| {
            let node = token.parent()?;
            let class = IdentClass::classify_token(sema, token)?;
            Some(class.definitions().into_iter().zip(iter::once(node).cycle()))
        })
        .flatten()
        .unique_by(|&(def, _)| def)
        .filter_map(|(def, node)| hover_for_definition(sema, file_id, def, &node, config))
        .reduce(|mut acc: HoverResult, HoverResult { markup, actions }| {
            acc.actions.extend(actions);
            acc.markup = Markup::from(format!("{}\n---\n{}", acc.markup, markup));
            acc
        });

    if result.is_none() {
        // fallbacks, show keywords or types
        if let Some(res) = render::keyword(sema, config, &original_token) {
            return Some(RangeInfo::new(original_token.text_range(), res));
        }
        let res = descended
            .iter()
            .find_map(|token| hover_type_fallback(sema, config, token, &original_token));
        if let res @ Some(_) = res {
            return res;
        }
    }
    result.map(|mut res: HoverResult| {
        res.actions = dedupe_or_merge_hover_actions(res.actions);
        RangeInfo::new(original_token.text_range(), res)
    })
}

pub(crate) fn hover_for_definition(
    sema: &Semantics<RootDatabase>,
    file_id: FileId,
    definition: Definition,
    node: &SyntaxNode,
    config: &HoverConfig,
) -> Option<HoverResult> {
    let famous_defs = match &definition {
        Definition::BuiltinType(_) => Some(FamousDefs(sema, sema.scope(node).krate())),
        _ => None,
    };
    if let Some(markup) = render::definition(sema.db, definition, famous_defs.as_ref(), config) {
        let mut res = HoverResult::default();
        res.markup = render::process_markup(sema.db, definition, &markup, config);
        if let Some(action) = show_implementations_action(sema.db, definition) {
            res.actions.push(action);
        }

        if let Some(action) = show_fn_references_action(sema.db, definition) {
            res.actions.push(action);
        }

        if let Some(action) = runnable_action(sema, definition, file_id) {
            res.actions.push(action);
        }

        if let Some(action) = goto_type_action_for_def(sema.db, definition) {
            res.actions.push(action);
        }
        return Some(res);
    }
    None
}

fn hover_ranged(
    file: &SyntaxNode,
    range: syntax::TextRange,
    sema: &Semantics<RootDatabase>,
    config: &HoverConfig,
) -> Option<RangeInfo<HoverResult>> {
    let expr_or_pat = file.covering_element(range).ancestors().find_map(|it| {
        match_ast! {
            match it {
                ast::Expr(expr) => Some(Either::Left(expr)),
                ast::Pat(pat) => Some(Either::Right(pat)),
                _ => None,
            }
        }
    })?;
    let res = match &expr_or_pat {
        Either::Left(ast::Expr::TryExpr(try_expr)) => render::try_expr(sema, config, try_expr),
        Either::Left(ast::Expr::PrefixExpr(prefix_expr))
            if prefix_expr.op_kind() == Some(ast::UnaryOp::Deref) =>
        {
            render::deref_expr(sema, config, prefix_expr)
        }
        _ => None,
    };
    let res = res.or_else(|| render::type_info(sema, config, &expr_or_pat));
    res.map(|it| {
        let range = match expr_or_pat {
            Either::Left(it) => it.syntax().text_range(),
            Either::Right(it) => it.syntax().text_range(),
        };
        RangeInfo::new(range, it)
    })
}

fn hover_type_fallback(
    sema: &Semantics<RootDatabase>,
    config: &HoverConfig,
    token: &SyntaxToken,
    original_token: &SyntaxToken,
) -> Option<RangeInfo<HoverResult>> {
    let node = token
        .ancestors()
        .take_while(|it| !ast::Item::can_cast(it.kind()))
        .find(|n| ast::Expr::can_cast(n.kind()) || ast::Pat::can_cast(n.kind()))?;

    let expr_or_pat = match_ast! {
        match node {
            ast::Expr(it) => Either::Left(it),
            ast::Pat(it) => Either::Right(it),
            // If this node is a MACRO_CALL, it means that `descend_into_macros_many` failed to resolve.
            // (e.g expanding a builtin macro). So we give up here.
            ast::MacroCall(_it) => return None,
            _ => return None,
        }
    };

    let res = render::type_info(sema, config, &expr_or_pat)?;
    let range = sema
        .original_range_opt(&node)
        .map(|frange| frange.range)
        .unwrap_or_else(|| original_token.text_range());
    Some(RangeInfo::new(range, res))
}

fn show_implementations_action(db: &RootDatabase, def: Definition) -> Option<HoverAction> {
    fn to_action(nav_target: NavigationTarget) -> HoverAction {
        HoverAction::Implementation(FilePosition {
            file_id: nav_target.file_id,
            offset: nav_target.focus_or_full_range().start(),
        })
    }

    let adt = match def {
        Definition::Trait(it) => return it.try_to_nav(db).map(to_action),
        Definition::Adt(it) => Some(it),
        Definition::SelfType(it) => it.self_ty(db).as_adt(),
        _ => None,
    }?;
    adt.try_to_nav(db).map(to_action)
}

fn show_fn_references_action(db: &RootDatabase, def: Definition) -> Option<HoverAction> {
    match def {
        Definition::Function(it) => it.try_to_nav(db).map(|nav_target| {
            HoverAction::Reference(FilePosition {
                file_id: nav_target.file_id,
                offset: nav_target.focus_or_full_range().start(),
            })
        }),
        _ => None,
    }
}

fn runnable_action(
    sema: &hir::Semantics<RootDatabase>,
    def: Definition,
    file_id: FileId,
) -> Option<HoverAction> {
    match def {
        Definition::Module(it) => runnable_mod(sema, it).map(HoverAction::Runnable),
        Definition::Function(func) => {
            let src = func.source(sema.db)?;
            if src.file_id != file_id.into() {
                cov_mark::hit!(hover_macro_generated_struct_fn_doc_comment);
                cov_mark::hit!(hover_macro_generated_struct_fn_doc_attr);
                return None;
            }

            runnable_fn(sema, func).map(HoverAction::Runnable)
        }
        _ => None,
    }
}

fn goto_type_action_for_def(db: &RootDatabase, def: Definition) -> Option<HoverAction> {
    let mut targets: Vec<hir::ModuleDef> = Vec::new();
    let mut push_new_def = |item: hir::ModuleDef| {
        if !targets.contains(&item) {
            targets.push(item);
        }
    };

    if let Definition::GenericParam(hir::GenericParam::TypeParam(it)) = def {
        it.trait_bounds(db).into_iter().for_each(|it| push_new_def(it.into()));
    } else {
        let ty = match def {
            Definition::Local(it) => it.ty(db),
            Definition::GenericParam(hir::GenericParam::ConstParam(it)) => it.ty(db),
            Definition::Field(field) => field.ty(db),
            Definition::Function(function) => function.ret_type(db),
            _ => return None,
        };

        walk_and_push_ty(db, &ty, &mut push_new_def);
    }

    Some(HoverAction::goto_type_from_targets(db, targets))
}

fn walk_and_push_ty(
    db: &RootDatabase,
    ty: &hir::Type,
    push_new_def: &mut dyn FnMut(hir::ModuleDef),
) {
    ty.walk(db, |t| {
        if let Some(adt) = t.as_adt() {
            push_new_def(adt.into());
        } else if let Some(trait_) = t.as_dyn_trait() {
            push_new_def(trait_.into());
        } else if let Some(traits) = t.as_impl_traits(db) {
            traits.for_each(|it| push_new_def(it.into()));
        } else if let Some(trait_) = t.as_associated_type_parent_trait(db) {
            push_new_def(trait_.into());
        }
    });
}

fn dedupe_or_merge_hover_actions(actions: Vec<HoverAction>) -> Vec<HoverAction> {
    let mut deduped_actions = Vec::with_capacity(actions.len());
    let mut go_to_type_targets = FxIndexSet::default();

    let mut seen_implementation = false;
    let mut seen_reference = false;
    let mut seen_runnable = false;
    for action in actions {
        match action {
            HoverAction::GoToType(targets) => {
                go_to_type_targets.extend(targets);
            }
            HoverAction::Implementation(..) => {
                if !seen_implementation {
                    seen_implementation = true;
                    deduped_actions.push(action);
                }
            }
            HoverAction::Reference(..) => {
                if !seen_reference {
                    seen_reference = true;
                    deduped_actions.push(action);
                }
            }
            HoverAction::Runnable(..) => {
                if !seen_runnable {
                    seen_runnable = true;
                    deduped_actions.push(action);
                }
            }
        };
    }

    if !go_to_type_targets.is_empty() {
        deduped_actions.push(HoverAction::GoToType(go_to_type_targets.into_iter().collect()));
    }

    deduped_actions
}
<commit_msg>fix: keyword hover works on non-keyword tokens if expanded to keyword<commit_after>mod render;

#[cfg(test)]
mod tests;

use std::iter;

use either::Either;
use hir::{HasSource, Semantics};
use ide_db::{
    base_db::FileRange,
    defs::{Definition, IdentClass},
    helpers::{pick_best_token, FamousDefs},
    FxIndexSet, RootDatabase,
};
use itertools::Itertools;
use syntax::{ast, match_ast, AstNode, SyntaxKind::*, SyntaxNode, SyntaxToken, T};

use crate::{
    doc_links::token_as_doc_comment,
    markup::Markup,
    runnables::{runnable_fn, runnable_mod},
    FileId, FilePosition, NavigationTarget, RangeInfo, Runnable, TryToNav,
};
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct HoverConfig {
    pub links_in_hover: bool,
    pub documentation: Option<HoverDocFormat>,
}

impl HoverConfig {
    fn markdown(&self) -> bool {
        matches!(self.documentation, Some(HoverDocFormat::Markdown))
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum HoverDocFormat {
    Markdown,
    PlainText,
}

#[derive(Debug, Clone)]
pub enum HoverAction {
    Runnable(Runnable),
    Implementation(FilePosition),
    Reference(FilePosition),
    GoToType(Vec<HoverGotoTypeData>),
}

impl HoverAction {
    fn goto_type_from_targets(db: &RootDatabase, targets: Vec<hir::ModuleDef>) -> Self {
        let targets = targets
            .into_iter()
            .filter_map(|it| {
                Some(HoverGotoTypeData {
                    mod_path: render::path(
                        db,
                        it.module(db)?,
                        it.name(db).map(|name| name.to_string()),
                    ),
                    nav: it.try_to_nav(db)?,
                })
            })
            .collect();
        HoverAction::GoToType(targets)
    }
}

#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub struct HoverGotoTypeData {
    pub mod_path: String,
    pub nav: NavigationTarget,
}

/// Contains the results when hovering over an item
#[derive(Debug, Default)]
pub struct HoverResult {
    pub markup: Markup,
    pub actions: Vec<HoverAction>,
}

// Feature: Hover
//
// Shows additional information, like the type of an expression or the documentation for a definition when "focusing" code.
// Focusing is usually hovering with a mouse, but can also be triggered with a shortcut.
//
// image::https://user-images.githubusercontent.com/48062697/113020658-b5f98b80-917a-11eb-9f88-3dbc27320c95.gif[]
pub(crate) fn hover(
    db: &RootDatabase,
    FileRange { file_id, range }: FileRange,
    config: &HoverConfig,
) -> Option<RangeInfo<HoverResult>> {
    let sema = &hir::Semantics::new(db);
    let file = sema.parse(file_id).syntax().clone();

    if !range.is_empty() {
        return hover_ranged(&file, range, sema, config);
    }
    let offset = range.start();

    let original_token = pick_best_token(file.token_at_offset(offset), |kind| match kind {
        IDENT | INT_NUMBER | LIFETIME_IDENT | T![self] | T![super] | T![crate] => 3,
        T!['('] | T![')'] => 2,
        kind if kind.is_trivia() => 0,
        _ => 1,
    })?;

    if let Some(doc_comment) = token_as_doc_comment(&original_token) {
        cov_mark::hit!(no_highlight_on_comment_hover);
        return doc_comment.get_definition_with_descend_at(sema, offset, |def, node, range| {
            let res = hover_for_definition(sema, file_id, def, &node, config)?;
            Some(RangeInfo::new(range, res))
        });
    }

    let descended = sema.descend_into_macros(original_token.clone());

    // FIXME: Definition should include known lints and the like instead of having this special case here
    let hovered_lint = descended.iter().find_map(|token| {
        let attr = token.ancestors().find_map(ast::Attr::cast)?;
        render::try_for_lint(&attr, token)
    });
    if let Some(res) = hovered_lint {
        return Some(RangeInfo::new(original_token.text_range(), res));
    }

    let result = descended
        .iter()
        .filter_map(|token| {
            let node = token.parent()?;
            let class = IdentClass::classify_token(sema, token)?;
            Some(class.definitions().into_iter().zip(iter::once(node).cycle()))
        })
        .flatten()
        .unique_by(|&(def, _)| def)
        .filter_map(|(def, node)| hover_for_definition(sema, file_id, def, &node, config))
        .reduce(|mut acc: HoverResult, HoverResult { markup, actions }| {
            acc.actions.extend(actions);
            acc.markup = Markup::from(format!("{}\n---\n{}", acc.markup, markup));
            acc
        });

    if result.is_none() {
        // fallbacks, show keywords or types

        let res = descended.iter().find_map(|token| render::keyword(sema, config, &token));
        if let Some(res) = res {
            return Some(RangeInfo::new(original_token.text_range(), res));
        }
        let res = descended
            .iter()
            .find_map(|token| hover_type_fallback(sema, config, token, &original_token));
        if let res @ Some(_) = res {
            return res;
        }
    }
    result.map(|mut res: HoverResult| {
        res.actions = dedupe_or_merge_hover_actions(res.actions);
        RangeInfo::new(original_token.text_range(), res)
    })
}

pub(crate) fn hover_for_definition(
    sema: &Semantics<RootDatabase>,
    file_id: FileId,
    definition: Definition,
    node: &SyntaxNode,
    config: &HoverConfig,
) -> Option<HoverResult> {
    let famous_defs = match &definition {
        Definition::BuiltinType(_) => Some(FamousDefs(sema, sema.scope(node).krate())),
        _ => None,
    };
    if let Some(markup) = render::definition(sema.db, definition, famous_defs.as_ref(), config) {
        let mut res = HoverResult::default();
        res.markup = render::process_markup(sema.db, definition, &markup, config);
        if let Some(action) = show_implementations_action(sema.db, definition) {
            res.actions.push(action);
        }

        if let Some(action) = show_fn_references_action(sema.db, definition) {
            res.actions.push(action);
        }

        if let Some(action) = runnable_action(sema, definition, file_id) {
            res.actions.push(action);
        }

        if let Some(action) = goto_type_action_for_def(sema.db, definition) {
            res.actions.push(action);
        }
        return Some(res);
    }
    None
}

fn hover_ranged(
    file: &SyntaxNode,
    range: syntax::TextRange,
    sema: &Semantics<RootDatabase>,
    config: &HoverConfig,
) -> Option<RangeInfo<HoverResult>> {
    // FIXME: make this work in attributes
    let expr_or_pat = file.covering_element(range).ancestors().find_map(|it| {
        match_ast! {
            match it {
                ast::Expr(expr) => Some(Either::Left(expr)),
                ast::Pat(pat) => Some(Either::Right(pat)),
                _ => None,
            }
        }
    })?;
    let res = match &expr_or_pat {
        Either::Left(ast::Expr::TryExpr(try_expr)) => render::try_expr(sema, config, try_expr),
        Either::Left(ast::Expr::PrefixExpr(prefix_expr))
            if prefix_expr.op_kind() == Some(ast::UnaryOp::Deref) =>
        {
            render::deref_expr(sema, config, prefix_expr)
        }
        _ => None,
    };
    let res = res.or_else(|| render::type_info(sema, config, &expr_or_pat));
    res.map(|it| {
        let range = match expr_or_pat {
            Either::Left(it) => it.syntax().text_range(),
            Either::Right(it) => it.syntax().text_range(),
        };
        RangeInfo::new(range, it)
    })
}

fn hover_type_fallback(
    sema: &Semantics<RootDatabase>,
    config: &HoverConfig,
    token: &SyntaxToken,
    original_token: &SyntaxToken,
) -> Option<RangeInfo<HoverResult>> {
    let node = token
        .ancestors()
        .take_while(|it| !ast::Item::can_cast(it.kind()))
        .find(|n| ast::Expr::can_cast(n.kind()) || ast::Pat::can_cast(n.kind()))?;

    let expr_or_pat = match_ast! {
        match node {
            ast::Expr(it) => Either::Left(it),
            ast::Pat(it) => Either::Right(it),
            // If this node is a MACRO_CALL, it means that `descend_into_macros_many` failed to resolve.
            // (e.g expanding a builtin macro). So we give up here.
            ast::MacroCall(_it) => return None,
            _ => return None,
        }
    };

    let res = render::type_info(sema, config, &expr_or_pat)?;
    let range = sema
        .original_range_opt(&node)
        .map(|frange| frange.range)
        .unwrap_or_else(|| original_token.text_range());
    Some(RangeInfo::new(range, res))
}

fn show_implementations_action(db: &RootDatabase, def: Definition) -> Option<HoverAction> {
    fn to_action(nav_target: NavigationTarget) -> HoverAction {
        HoverAction::Implementation(FilePosition {
            file_id: nav_target.file_id,
            offset: nav_target.focus_or_full_range().start(),
        })
    }

    let adt = match def {
        Definition::Trait(it) => return it.try_to_nav(db).map(to_action),
        Definition::Adt(it) => Some(it),
        Definition::SelfType(it) => it.self_ty(db).as_adt(),
        _ => None,
    }?;
    adt.try_to_nav(db).map(to_action)
}

fn show_fn_references_action(db: &RootDatabase, def: Definition) -> Option<HoverAction> {
    match def {
        Definition::Function(it) => it.try_to_nav(db).map(|nav_target| {
            HoverAction::Reference(FilePosition {
                file_id: nav_target.file_id,
                offset: nav_target.focus_or_full_range().start(),
            })
        }),
        _ => None,
    }
}

fn runnable_action(
    sema: &hir::Semantics<RootDatabase>,
    def: Definition,
    file_id: FileId,
) -> Option<HoverAction> {
    match def {
        Definition::Module(it) => runnable_mod(sema, it).map(HoverAction::Runnable),
        Definition::Function(func) => {
            let src = func.source(sema.db)?;
            if src.file_id != file_id.into() {
                cov_mark::hit!(hover_macro_generated_struct_fn_doc_comment);
                cov_mark::hit!(hover_macro_generated_struct_fn_doc_attr);
                return None;
            }

            runnable_fn(sema, func).map(HoverAction::Runnable)
        }
        _ => None,
    }
}

fn goto_type_action_for_def(db: &RootDatabase, def: Definition) -> Option<HoverAction> {
    let mut targets: Vec<hir::ModuleDef> = Vec::new();
    let mut push_new_def = |item: hir::ModuleDef| {
        if !targets.contains(&item) {
            targets.push(item);
        }
    };

    if let Definition::GenericParam(hir::GenericParam::TypeParam(it)) = def {
        it.trait_bounds(db).into_iter().for_each(|it| push_new_def(it.into()));
    } else {
        let ty = match def {
            Definition::Local(it) => it.ty(db),
            Definition::GenericParam(hir::GenericParam::ConstParam(it)) => it.ty(db),
            Definition::Field(field) => field.ty(db),
            Definition::Function(function) => function.ret_type(db),
            _ => return None,
        };

        walk_and_push_ty(db, &ty, &mut push_new_def);
    }

    Some(HoverAction::goto_type_from_targets(db, targets))
}

fn walk_and_push_ty(
    db: &RootDatabase,
    ty: &hir::Type,
    push_new_def: &mut dyn FnMut(hir::ModuleDef),
) {
    ty.walk(db, |t| {
        if let Some(adt) = t.as_adt() {
            push_new_def(adt.into());
        } else if let Some(trait_) = t.as_dyn_trait() {
            push_new_def(trait_.into());
        } else if let Some(traits) = t.as_impl_traits(db) {
            traits.for_each(|it| push_new_def(it.into()));
        } else if let Some(trait_) = t.as_associated_type_parent_trait(db) {
            push_new_def(trait_.into());
        }
    });
}

fn dedupe_or_merge_hover_actions(actions: Vec<HoverAction>) -> Vec<HoverAction> {
    let mut deduped_actions = Vec::with_capacity(actions.len());
    let mut go_to_type_targets = FxIndexSet::default();

    let mut seen_implementation = false;
    let mut seen_reference = false;
    let mut seen_runnable = false;
    for action in actions {
        match action {
            HoverAction::GoToType(targets) => {
                go_to_type_targets.extend(targets);
            }
            HoverAction::Implementation(..) => {
                if !seen_implementation {
                    seen_implementation = true;
                    deduped_actions.push(action);
                }
            }
            HoverAction::Reference(..) => {
                if !seen_reference {
                    seen_reference = true;
                    deduped_actions.push(action);
                }
            }
            HoverAction::Runnable(..) => {
                if !seen_runnable {
                    seen_runnable = true;
                    deduped_actions.push(action);
                }
            }
        };
    }

    if !go_to_type_targets.is_empty() {
        deduped_actions.push(HoverAction::GoToType(go_to_type_targets.into_iter().collect()));
    }

    deduped_actions
}
<|endoftext|>


<commit_before>use super::*;

use std::{collections::BTreeSet, ops::Bound};

#[derive(Debug, Clone)]
pub(crate) struct View<'a> {
    pub(crate) pid: PageId,
    pub(crate) lo: &'a IVec,
    pub(crate) hi: &'a IVec,
    pub(crate) is_index: bool,
    pub(crate) next: Option<PageId>,
    pub(crate) ptr: TreePtr<'a>,
    frags: Vec<&'a Frag>,
    base_offset: usize,
    pub(crate) base_data: &'a Data,
    pub(crate) merging_child: Option<PageId>,
    pub(crate) merging: bool,
}

impl<'a> View<'a> {
    pub(crate) fn new(
        pid: PageId,
        ptr: TreePtr<'a>,
        frags: Vec<&'a Frag>,
    ) -> View<'a> {
        let mut view = View {
            pid,
            ptr,
            frags,
            lo: unsafe { std::mem::uninitialized() },
            hi: unsafe { std::mem::uninitialized() },
            is_index: false,
            base_offset: usize::max_value(),
            base_data: unsafe { std::mem::uninitialized() },
            next: None,
            merging_child: None,
            merging: false,
        };

        let mut merge_confirmed = false;

        for (offset, frag) in view.frags.iter().enumerate() {
            match frag {
                Frag::Base(node) => {
                    if view.base_offset == usize::max_value() {
                        // hi and next may be changed via a
                        // parent split if we re-add that frag
                        view.hi = &node.hi;
                        view.next = node.next;
                    }
                    view.lo = &node.lo;
                    view.is_index = node.data.index_ref().is_some();
                    view.base_offset = offset;
                    view.base_data = &node.data;
                    break;
                }
                Frag::ParentMergeIntention(pid) => {
                    if merge_confirmed {
                        merge_confirmed = false;
                    } else {
                        assert!(view.merging_child.is_none());
                        view.merging_child = Some(*pid);
                    }
                }
                Frag::ParentMergeConfirm => {
                    assert!(!merge_confirmed);
                    merge_confirmed = true;
                }
                Frag::ChildMergeCap => {
                    assert!(!view.merging);
                    assert_eq!(offset, 0);
                    view.merging = true;
                }
                _ => {}
            }
        }

        assert_ne!(
            view.base_offset,
            usize::max_value(),
            "view was never initialized with a base"
        );

        view
    }

    pub(crate) fn contains_upper_bound(&self, bound: &Bound<IVec>) -> bool {
        match bound {
            Bound::Unbounded => self.hi.is_empty(),
            Bound::Included(bound) => self.hi > bound || self.hi.is_empty(),
            Bound::Excluded(bound) => self.hi >= bound || self.hi.is_empty(),
        }
    }

    pub(crate) fn contains_lower_bound(&self, bound: &Bound<IVec>) -> bool {
        match bound {
            Bound::Unbounded => self.lo.is_empty(),
            Bound::Included(bound) | Bound::Excluded(bound) => self.lo <= bound,
        }
    }

    fn keys(&self) -> BTreeSet<&IVec> {
        let mut keys: BTreeSet<&IVec> = self
            .base_data
            .leaf_ref()
            .unwrap()
            .iter()
            .map(|(k, _v)| k)
            .collect();

        for offset in (0..self.base_offset).rev() {
            match self.frags[offset] {
                Frag::Set(k, _) => {
                    keys.insert(k);
                }
                Frag::Del(k) => {
                    keys.remove(k);
                }
                Frag::Merge(k, _) => {
                    keys.insert(k);
                }
                Frag::Base(_) => {
                    panic!(
                        "somehow hit 2 base nodes while \
                         searching for a successor"
                    );
                }
                Frag::ChildMergeCap => {}
                Frag::ParentMergeIntention(_) | Frag::ParentMergeConfirm => {
                    panic!(
                        "somehow hit parent merge \
                         frags while searching for a \
                         successor"
                    )
                }
            }
        }

        keys
    }

    pub(crate) fn successor(
        &self,
        bound: &Bound<IVec>,
        config: &Config,
    ) -> Option<(IVec, IVec)> {
        assert!(!self.is_index);

        // This encoding happens this way because
        // keys cannot be lower than the node's lo key.
        let predecessor_key = match bound {
            Bound::Unbounded => prefix_encode(self.lo, self.lo),
            Bound::Included(b) => {
                let max = std::cmp::max(b, self.lo);
                prefix_encode(self.lo, max)
            }
            Bound::Excluded(b) => {
                let max = std::cmp::max(b, self.lo);
                prefix_encode(self.lo, max)
            }
        };

        let keys = self.keys();

        let successor_keys = keys.range(predecessor_key..);

        for encoded_key in successor_keys {
            let decoded_key = prefix_decode(self.lo, &encoded_key);

            if let Bound::Excluded(e) = bound {
                if &*e == &decoded_key {
                    // skip this excluded key
                    continue;
                }
            }

            // try to get this key until it works
            if let Some(value) = self.leaf_value_for_key(&decoded_key, config) {
                return Some((IVec::from(decoded_key), value));
            }
        }

        None
    }

    pub(crate) fn predecessor(
        &self,
        bound: &Bound<IVec>,
        config: &Config,
    ) -> Option<(IVec, IVec)> {
        assert!(!self.is_index);

        // This encoding happens this way because
        // the rightmost (unbounded) node has
        // a hi key represented by the empty slice
        let successor_key = match bound {
            Bound::Unbounded => {
                if self.hi.is_empty() {
                    prefix_encode(self.lo, &[255; 1024 * 1024])
                } else {
                    prefix_encode(self.lo, self.hi)
                }
            }
            Bound::Included(b) => {
                let min = if self.hi.is_empty() {
                    b
                } else {
                    std::cmp::min(b, self.hi)
                };
                prefix_encode(self.lo, min)
            }
            Bound::Excluded(b) => {
                let min = if self.hi.is_empty() {
                    b
                } else {
                    std::cmp::min(b, self.hi)
                };
                prefix_encode(self.lo, min)
            }
        };

        let keys = self.keys();

        let predecessor_keys = keys.range(..=successor_key).rev();

        for encoded_key in predecessor_keys {
            let decoded_key = prefix_decode(self.lo, &encoded_key);

            if let Bound::Excluded(e) = bound {
                if &*e == &decoded_key {
                    // skip this excluded key
                    continue;
                }
            }

            // try to get this key until it works
            if let Some(value) = self.leaf_value_for_key(&decoded_key, config) {
                return Some((IVec::from(decoded_key), value));
            }
        }

        None
    }

    pub(crate) fn is_free(&self) -> bool {
        self.frags.is_empty()
    }

    pub(crate) fn leaf_value_for_key(
        &self,
        key: &[u8],
        config: &Config,
    ) -> Option<IVec> {
        assert!(!self.is_index);

        let mut merge_base = None;
        let mut merges = vec![];

        for frag in self.frags[..self.base_offset + 1].iter() {
            match frag {
                Frag::Set(k, val) if self.key_eq(k, key) => {
                    if merges.is_empty() {
                        return Some(val.clone());
                    } else {
                        merge_base = Some(val);
                        break;
                    }
                }
                Frag::Del(k) if self.key_eq(k, key) => return None,
                Frag::Merge(k, val) if self.key_eq(k, key) => merges.push(val),
                Frag::Base(node) => {
                    let data = &node.data;
                    let items =
                        data.leaf_ref().expect("last_node should be a leaf");
                    let search = items
                        .binary_search_by(|&(ref k, ref _v)| {
                            prefix_cmp_encoded(k, key.as_ref(), &node.lo)
                        })
                        .ok();

                    let val = search.map(|idx| &items[idx].1);
                    if merges.is_empty() {
                        return val.cloned();
                    } else {
                        merge_base = val;
                    }
                }
                _ => {}
            }
        }

        if merges.is_empty() {
            None
        } else {
            let merge_fn_ptr = config
                .merge_operator
                .expect("must have a merge operator set");

            unsafe {
                let merge_fn: MergeOperator = std::mem::transmute(merge_fn_ptr);

                let mut ret = merge_fn(
                    key,
                    merge_base.map(|iv| &**iv),
                    &merges.pop().unwrap(),
                );
                                       ;
                for merge in merges.into_iter().rev() {
                    if let Some(v) = ret {
                        ret = merge_fn(key, Some(&*v), merge);
                    } else {
                        ret = merge_fn(key, None, merge);
                    }
                }

                ret.map(IVec::from)
            }
        }
    }

    #[inline]
    fn key_eq(&self, encoded: &[u8], not_encoded: &[u8]) -> bool {
        prefix_cmp_encoded(encoded, not_encoded, self.lo)
            == std::cmp::Ordering::Equal
    }

    pub(crate) fn index_next_node(&self, key: &[u8]) -> PageId {
        assert!(self.is_index);

        for frag in self.frags[..self.base_offset + 1].iter() {
            match frag {
                Frag::Set(..) => unimplemented!(),
                Frag::Del(..) => unimplemented!(),
                Frag::Merge(..) => unimplemented!(),
                Frag::Base(node) => {
                    let data = &node.data;
                    let items =
                        data.index_ref().expect("last_node should be a leaf");
                    let search =
                        binary_search_lub(items, |&(ref k, ref _v)| {
                            prefix_cmp_encoded(k, key.as_ref(), &node.lo)
                        });

                    // This might be none if ord is Less and we're
                    // searching for the empty key
                    let index = search.expect("failed to traverse index");

                    return items[index].1;
                }
                Frag::ParentMergeIntention(_)
                | Frag::ParentMergeConfirm
                | Frag::ChildMergeCap => {
                    // nothing to do for these frags
                }
            }
        }
        panic!("no index found")
    }

    pub(crate) fn should_split(&self, max_sz: u64) -> bool {
        let children = self.base_data.len();
        children > 2
            && self.size_in_bytes() > max_sz
            && self.merging_child.is_none()
            && !self.merging
    }

    pub(crate) fn should_merge(&self, min_sz: u64) -> bool {
        self.size_in_bytes() < min_sz
            && self.merging_child.is_none()
            && !self.merging
    }

    pub(crate) fn can_merge_child(&self) -> bool {
        self.merging_child.is_none() && !self.merging
    }

    pub(crate) fn compact(&self, config: &Config) -> Node {
        let mut lhs = self.frags[self.base_offset].unwrap_base().clone();
        for offset in (0..self.base_offset).rev() {
            let frag = self.frags[offset];
            lhs.apply(frag, config.merge_operator);
        }
        lhs
    }

    pub(crate) fn split(&self, config: &Config) -> (Node, Node) {
        let mut lhs = self.compact(config);
        let rhs = lhs.split();

        lhs.data.drop_gte(&rhs.lo, &lhs.lo);
        lhs.hi = rhs.lo.clone();

        // intentionally make this the end to make
        // any issues pop out with setting it
        // correctly after the split.
        lhs.next = None;

        (lhs, rhs)
    }

    #[inline]
    pub(crate) fn size_in_bytes(&self) -> u64 {
        // TODO needs to better account for
        // sizes that don't actually fall under
        // a merge threshold once we support one.
        self.frags[..self.base_offset + 1]
            .iter()
            .map(|f| f.size_in_bytes())
            .sum()
    }
}
<commit_msg>Better account for prefix encoding while iterating<commit_after>use super::*;

use std::{collections::HashSet, ops::Bound};

#[derive(Debug, Clone)]
pub(crate) struct View<'a> {
    pub(crate) pid: PageId,
    pub(crate) lo: &'a IVec,
    pub(crate) hi: &'a IVec,
    pub(crate) is_index: bool,
    pub(crate) next: Option<PageId>,
    pub(crate) ptr: TreePtr<'a>,
    frags: Vec<&'a Frag>,
    base_offset: usize,
    pub(crate) base_data: &'a Data,
    pub(crate) merging_child: Option<PageId>,
    pub(crate) merging: bool,
}

impl<'a> View<'a> {
    pub(crate) fn new(
        pid: PageId,
        ptr: TreePtr<'a>,
        frags: Vec<&'a Frag>,
    ) -> View<'a> {
        let mut view = View {
            pid,
            ptr,
            frags,
            lo: unsafe { std::mem::uninitialized() },
            hi: unsafe { std::mem::uninitialized() },
            is_index: false,
            base_offset: usize::max_value(),
            base_data: unsafe { std::mem::uninitialized() },
            next: None,
            merging_child: None,
            merging: false,
        };

        let mut merge_confirmed = false;

        for (offset, frag) in view.frags.iter().enumerate() {
            match frag {
                Frag::Base(node) => {
                    if view.base_offset == usize::max_value() {
                        // hi and next may be changed via a
                        // parent split if we re-add that frag
                        view.hi = &node.hi;
                        view.next = node.next;
                    }
                    view.lo = &node.lo;
                    view.is_index = node.data.index_ref().is_some();
                    view.base_offset = offset;
                    view.base_data = &node.data;
                    break;
                }
                Frag::ParentMergeIntention(pid) => {
                    if merge_confirmed {
                        merge_confirmed = false;
                    } else {
                        assert!(view.merging_child.is_none());
                        view.merging_child = Some(*pid);
                    }
                }
                Frag::ParentMergeConfirm => {
                    assert!(!merge_confirmed);
                    merge_confirmed = true;
                }
                Frag::ChildMergeCap => {
                    assert!(!view.merging);
                    assert_eq!(offset, 0);
                    view.merging = true;
                }
                _ => {}
            }
        }

        assert_ne!(
            view.base_offset,
            usize::max_value(),
            "view was never initialized with a base"
        );

        view
    }

    pub(crate) fn contains_upper_bound(&self, bound: &Bound<IVec>) -> bool {
        match bound {
            Bound::Unbounded => self.hi.is_empty(),
            Bound::Included(bound) => self.hi > bound || self.hi.is_empty(),
            Bound::Excluded(bound) => self.hi >= bound || self.hi.is_empty(),
        }
    }

    pub(crate) fn contains_lower_bound(&self, bound: &Bound<IVec>) -> bool {
        match bound {
            Bound::Unbounded => self.lo.is_empty(),
            Bound::Included(bound) | Bound::Excluded(bound) => self.lo <= bound,
        }
    }

    fn keys(&self) -> Vec<&IVec> {
        let mut keys: HashSet<&IVec> = self
            .base_data
            .leaf_ref()
            .unwrap()
            .iter()
            .map(|(k, _v)| k)
            .collect();

        for offset in (0..self.base_offset).rev() {
            match self.frags[offset] {
                Frag::Set(k, _) => {
                    keys.insert(k);
                }
                Frag::Del(k) => {
                    keys.remove(k);
                }
                Frag::Merge(k, _) => {
                    keys.insert(k);
                }
                Frag::Base(_) => {
                    panic!(
                        "somehow hit 2 base nodes while \
                         searching for a successor"
                    );
                }
                Frag::ChildMergeCap => {}
                Frag::ParentMergeIntention(_) | Frag::ParentMergeConfirm => {
                    panic!(
                        "somehow hit parent merge \
                         frags while searching for a \
                         successor"
                    )
                }
            }
        }

        let mut keys: Vec<_> = keys.into_iter().collect();
        keys.sort_by(|a, b| prefix_cmp(&a, &b));

        keys
    }

    pub(crate) fn successor(
        &self,
        bound: &Bound<IVec>,
        config: &Config,
    ) -> Option<(IVec, IVec)> {
        assert!(!self.is_index);

        // This encoding happens this way because
        // keys cannot be lower than the node's lo key.
        let predecessor_key = match bound {
            Bound::Unbounded => prefix_encode(self.lo, self.lo),
            Bound::Included(b) => {
                let max = std::cmp::max(b, self.lo);
                prefix_encode(self.lo, max)
            }
            Bound::Excluded(b) => {
                let max = std::cmp::max(b, self.lo);
                prefix_encode(self.lo, max)
            }
        };

        let keys = self.keys();

        let successor_keys = keys.iter().filter(|k| {
            let ord = prefix_cmp(k, &predecessor_key);

            ord == std::cmp::Ordering::Greater
                || ord == std::cmp::Ordering::Equal
        });

        for encoded_key in successor_keys {
            let decoded_key = prefix_decode(self.lo, &encoded_key);

            if let Bound::Excluded(e) = bound {
                if &*e == &decoded_key {
                    // skip this excluded key
                    continue;
                }
            }

            // try to get this key until it works
            if let Some(value) = self.leaf_value_for_key(&decoded_key, config) {
                return Some((IVec::from(decoded_key), value));
            }
        }

        None
    }

    pub(crate) fn predecessor(
        &self,
        bound: &Bound<IVec>,
        config: &Config,
    ) -> Option<(IVec, IVec)> {
        assert!(!self.is_index);

        // This encoding happens this way because
        // the rightmost (unbounded) node has
        // a hi key represented by the empty slice
        let successor_key = match bound {
            Bound::Unbounded => {
                if self.hi.is_empty() {
                    prefix_encode(self.lo, &[255; 1024 * 1024])
                } else {
                    prefix_encode(self.lo, self.hi)
                }
            }
            Bound::Included(b) => {
                let min = if self.hi.is_empty() {
                    b
                } else {
                    std::cmp::min(b, self.hi)
                };
                prefix_encode(self.lo, min)
            }
            Bound::Excluded(b) => {
                let min = if self.hi.is_empty() {
                    b
                } else {
                    std::cmp::min(b, self.hi)
                };
                prefix_encode(self.lo, min)
            }
        };

        let keys = self.keys();

        let predecessor_keys = keys.iter().filter(|k| {
            let ord = prefix_cmp(k, &successor_key);

            ord == std::cmp::Ordering::Less || ord == std::cmp::Ordering::Equal
        });

        for encoded_key in predecessor_keys {
            let decoded_key = prefix_decode(self.lo, &encoded_key);

            if let Bound::Excluded(e) = bound {
                if &*e == &decoded_key {
                    // skip this excluded key
                    continue;
                }
            }

            // try to get this key until it works
            if let Some(value) = self.leaf_value_for_key(&decoded_key, config) {
                return Some((IVec::from(decoded_key), value));
            }
        }

        None
    }

    pub(crate) fn is_free(&self) -> bool {
        self.frags.is_empty()
    }

    pub(crate) fn leaf_value_for_key(
        &self,
        key: &[u8],
        config: &Config,
    ) -> Option<IVec> {
        assert!(!self.is_index);

        let mut merge_base = None;
        let mut merges = vec![];

        for frag in self.frags[..self.base_offset + 1].iter() {
            match frag {
                Frag::Set(k, val) if self.key_eq(k, key) => {
                    if merges.is_empty() {
                        return Some(val.clone());
                    } else {
                        merge_base = Some(val);
                        break;
                    }
                }
                Frag::Del(k) if self.key_eq(k, key) => return None,
                Frag::Merge(k, val) if self.key_eq(k, key) => merges.push(val),
                Frag::Base(node) => {
                    let data = &node.data;
                    let items =
                        data.leaf_ref().expect("last_node should be a leaf");
                    let search = items
                        .binary_search_by(|&(ref k, ref _v)| {
                            prefix_cmp_encoded(k, key.as_ref(), &node.lo)
                        })
                        .ok();

                    let val = search.map(|idx| &items[idx].1);
                    if merges.is_empty() {
                        return val.cloned();
                    } else {
                        merge_base = val;
                    }
                }
                _ => {}
            }
        }

        if merges.is_empty() {
            None
        } else {
            let merge_fn_ptr = config
                .merge_operator
                .expect("must have a merge operator set");

            unsafe {
                let merge_fn: MergeOperator = std::mem::transmute(merge_fn_ptr);

                let mut ret = merge_fn(
                    key,
                    merge_base.map(|iv| &**iv),
                    &merges.pop().unwrap(),
                );
                                       ;
                for merge in merges.into_iter().rev() {
                    if let Some(v) = ret {
                        ret = merge_fn(key, Some(&*v), merge);
                    } else {
                        ret = merge_fn(key, None, merge);
                    }
                }

                ret.map(IVec::from)
            }
        }
    }

    #[inline]
    fn key_eq(&self, encoded: &[u8], not_encoded: &[u8]) -> bool {
        prefix_cmp_encoded(encoded, not_encoded, self.lo)
            == std::cmp::Ordering::Equal
    }

    pub(crate) fn index_next_node(&self, key: &[u8]) -> PageId {
        assert!(self.is_index);

        for frag in self.frags[..self.base_offset + 1].iter() {
            match frag {
                Frag::Set(..) => unimplemented!(),
                Frag::Del(..) => unimplemented!(),
                Frag::Merge(..) => unimplemented!(),
                Frag::Base(node) => {
                    let data = &node.data;
                    let items =
                        data.index_ref().expect("last_node should be a leaf");
                    let search =
                        binary_search_lub(items, |&(ref k, ref _v)| {
                            prefix_cmp_encoded(k, key.as_ref(), &node.lo)
                        });

                    // This might be none if ord is Less and we're
                    // searching for the empty key
                    let index = search.expect("failed to traverse index");

                    return items[index].1;
                }
                Frag::ParentMergeIntention(_)
                | Frag::ParentMergeConfirm
                | Frag::ChildMergeCap => {
                    // nothing to do for these frags
                }
            }
        }
        panic!("no index found")
    }

    pub(crate) fn should_split(&self, max_sz: u64) -> bool {
        let children = self.base_data.len();
        children > 2
            && self.size_in_bytes() > max_sz
            && self.merging_child.is_none()
            && !self.merging
    }

    pub(crate) fn should_merge(&self, min_sz: u64) -> bool {
        self.size_in_bytes() < min_sz
            && self.merging_child.is_none()
            && !self.merging
    }

    pub(crate) fn can_merge_child(&self) -> bool {
        self.merging_child.is_none() && !self.merging
    }

    pub(crate) fn compact(&self, config: &Config) -> Node {
        let mut lhs = self.frags[self.base_offset].unwrap_base().clone();
        for offset in (0..self.base_offset).rev() {
            let frag = self.frags[offset];
            lhs.apply(frag, config.merge_operator);
        }
        lhs
    }

    pub(crate) fn split(&self, config: &Config) -> (Node, Node) {
        let mut lhs = self.compact(config);
        let rhs = lhs.split();

        lhs.data.drop_gte(&rhs.lo, &lhs.lo);
        lhs.hi = rhs.lo.clone();

        // intentionally make this the end to make
        // any issues pop out with setting it
        // correctly after the split.
        lhs.next = None;

        (lhs, rhs)
    }

    #[inline]
    pub(crate) fn size_in_bytes(&self) -> u64 {
        // TODO needs to better account for
        // sizes that don't actually fall under
        // a merge threshold once we support one.
        self.frags[..self.base_offset + 1]
            .iter()
            .map(|f| f.size_in_bytes())
            .sum()
    }
}
<|endoftext|>