<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 #37287 - srinivasreddy:unify, r=eddyb<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>// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

use super::*;
use std::fmt::Write;
use std::ops::Deref;
use std::str::FromStr;

/// Sized, heap-allocated string type guaranteed to contain a well-formed [IETF-RFC3986] URI
/// or [network path](enum.UriType.html#variant.NetworkPath).
///
/// The unsized counterpart is [`Uri`](crate::Uri).
///
/// This type implements [`std::ops::Deref<Uri>`], so you can also use all of the
/// methods from [`Uri`] on this type.
///
/// [IETF-RFC3986]: https://tools.ietf.org/html/rfc3986
#[derive(Clone, Eq, Hash)]
pub struct UriBuf(pub(super) UriRefBuf);

impl_uri_buf_traits!(UriBuf, Uri);

impl Deref for UriBuf {
    type Target = Uri;

    fn deref(&self) -> &Self::Target {
        self.as_uri()
    }
}

impl AsRef<Uri> for UriBuf {
    fn as_ref(&self) -> &Uri {
        self.as_uri()
    }
}

impl From<&Uri> for UriBuf {
    fn from(x: &Uri) -> Self {
        x.to_uri_buf()
    }
}

impl FromStr for UriBuf {
    type Err = ParseError;

    fn from_str(input: &str) -> Result<Self, Self::Err> {
        Self::from_str(input)
    }
}
impl std::fmt::Display for UriBuf {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::result::Result<(), std::fmt::Error> {
        self.write_to(f)
    }
}

/// # Unsafe Methods
///
/// `UriBuf` needs some unsafe methods in order to function properly. This section is where
/// they are all located.
impl UriBuf {
    /// Unchecked version of [`UriBuf::from_string`].
    ///
    /// # Safety
    ///
    /// This method is marked as unsafe because it allows you to construct a `UriBuf` with
    /// a value that is not a well-formed URI reference.
    #[inline(always)]
    pub unsafe fn from_string_unchecked(s: String) -> UriBuf {
        UriBuf(UriRefBuf::from_string_unchecked(s))
    }
}

impl UriBuf {
    /// Creates a new `UriBuf` from *unescaped* component values.
    pub fn new<Sch, Hos, Pat, Que, Frg>(
        scheme: Sch,
        host: Hos,
        port: Option<u16>,
        path: Pat,
        query: Option<Que>,
        fragment: Option<Frg>,
    ) -> UriBuf
    where
        Sch: Into<String>,
        Hos: AsRef<str>,
        Pat: AsRef<str>,
        Que: AsRef<str>,
        Frg: AsRef<str>,
    {
        let mut ret: String = Self::from_scheme_host_port(scheme, host, port).into();

        let mut path = path.as_ref();

        if path.starts_with('/') {
            path = &path[1..];
        }

        let path_segment_iter =
            path.split('/')
                .filter_map(|seg| if seg == "." { None } else { Some(seg) });

        for seg in path_segment_iter {
            ret.push('/');
            ret.extend(seg.escape_uri());
        }

        if let Some(query) = query {
            let mut first = true;
            ret.push('?');
            for seg in query.as_ref().split(|c| c == '&' || c == ';') {
                if first {
                    first = false;
                } else {
                    ret.push('&');
                }
                ret.extend(seg.escape_uri().for_query());
            }
        }

        if let Some(fragment) = fragment {
            ret.push('#');
            ret.extend(fragment.as_ref().escape_uri().for_fragment());
        }

        unsafe { Self::from_string_unchecked(ret) }
    }

    /// Constructs a `UriBuf` from a scheme and authority.
    ///
    /// The authority should not be percent encoded. If the given scheme contains invalid
    /// characters, this method will panic.
    ///
    /// # Example
    ///
    /// ```
    /// use async_coap_uri::prelude::*;
    /// let authority = "user@[2001:0db8:85a3::1%en2]:8080";
    ///
    /// let uri_buf = UriBuf::from_scheme_authority("http", authority);
    ///
    /// assert_eq!(uri_buf, uri!("http://user@[2001:0db8:85a3::1%25en2]:8080"));
    ///
    /// ```
    pub fn from_scheme_authority<Sch, Aut>(scheme: Sch, authority: Aut) -> UriBuf
    where
        Sch: Into<String>,
        Aut: AsRef<str>,
    {
        let mut ret = scheme.into();

        // Make sure that the scheme is proper. This is likely a string constant,
        // so we go ahead and panic if we detect that it is busted.
        assert_eq!(
            ret.find(|c: char| !(c.is_ascii_alphanumeric() || c == '+' || c == '-' || c == '.')),
            None,
            "Scheme contains invalid characters: {:?}",
            ret
        );

        ret.push_str("://");

        ret.extend(authority.as_ref().escape_uri().for_authority());

        unsafe { Self::from_string_unchecked(ret) }
    }

    /// Constructs a network path from a host and a relative reference.
    ///
    /// # Example
    ///
    /// ```
    /// use async_coap_uri::prelude::*;
    /// let host = "example.com";
    /// let rel_ref = rel_ref!("./foobar?q");
    ///
    /// let uri_buf = UriBuf::from_host_rel_ref(host, rel_ref);
    ///
    /// assert_eq!(uri_buf, uri!("//example.com/foobar?q"));
    ///
    /// ```
    pub fn from_host_rel_ref<Hos, RR>(host: Hos, rel_ref: RR) -> UriBuf
    where
        Hos: AsRef<str>,
        RR: AsRef<RelRef>,
    {
        let host = host.as_ref();
        let rel_ref = rel_ref
            .as_ref()
            .trim_leading_dot_slashes()
            .trim_leading_slashes();

        // UNWRAP-SAFETY: This is safe because we are fully
        // escaping the host and we already know rel_ref to
        // be well-formed.
        uri_format!("//{}/{}", host.escape_uri().full(), rel_ref).unwrap()
    }

    /// Constructs a `UriBuf` from a scheme, host and an optional port number.
    ///
    /// The host should not be percent encoded. If the given scheme contains invalid
    /// characters, this method will panic.
    pub fn from_scheme_host_port<Sch, Hos>(scheme: Sch, host: Hos, port: Option<u16>) -> UriBuf
    where
        Sch: Into<String>,
        Hos: AsRef<str>,
    {
        let mut ret = scheme.into();

        // Make sure that the scheme is proper. This is likely a string constant,
        // so we go ahead and panic if we detect that it is busted.
        assert_eq!(
            ret.find(|c: char| !(c.is_ascii_alphanumeric() || c == '+' || c == '-' || c == '.')),
            None,
            "Scheme contains invalid characters: {:?}",
            ret
        );

        ret.push_str("://");

        let mut host = host.as_ref();

        // Trim enclosing brackets.
        if host.starts_with('[') && host.ends_with(']') {
            host = &host[1..host.len()];
        }

        if host.find(':').is_some() {
            ret.push('[');
            ret.extend(host.escape_uri());
            ret.push(']');
        }

        if let Some(port) = port {
            write!(ret, ":{}", port).unwrap();
        }

        unsafe { Self::from_string_unchecked(ret) }
    }

    /// Attempts to create a new [`UriBuf`] from a string slice.
    pub fn from_str<S: AsRef<str>>(s: S) -> Result<UriBuf, ParseError> {
        let s = s.as_ref();
        let components = UriRawComponents::from_str(s)?;

        if components.uri_type().can_borrow_as_uri() {
            Ok(unsafe { Self::from_string_unchecked(s.to_string()) })
        } else {
            Err(ParseError::new("Missing scheme or authority", None))
        }
    }

    /// Attempts to create a new [`UriBuf`] from a [`String`].
    pub fn from_string(s: String) -> Result<UriBuf, ParseError> {
        let components = UriRawComponents::from_str(s.as_str())?;

        if components.uri_type().can_borrow_as_uri() {
            Ok(unsafe { Self::from_string_unchecked(s) })
        } else {
            Err(ParseError::new("Missing scheme or authority", None))
        }
    }

    /// Attempts to create a new [`UriBuf`] from a `UriRef` slice.
    pub fn from_uri<S: AsRef<UriRef>>(s: S) -> Option<UriBuf> {
        if s.as_ref().uri_type().can_borrow_as_uri() {
            Some(UriBuf(s.as_ref().to_uri_ref_buf()))
        } else {
            None
        }
    }
}

impl UriBuf {
    /// Borrows a [`Uri`] slice containing this URI.
    #[inline(always)]
    pub fn as_uri(&self) -> &Uri {
        unsafe { Uri::from_str_unchecked(self.as_str()) }
    }
}

/// # Manipulation
impl UriBuf {
    /// Using this URI as the base, performs "relative resolution" to the given instance
    /// implementing [`AnyUriRef`], updating the content of this `UriBuf` with the result.
    pub fn resolve<T: AnyUriRef + ?Sized>(&mut self, dest: &T) -> Result<(), ResolveError> {
        self.0.resolve(dest)
    }

    /// Replaces the path, query, and fragment with that from `rel`.
    pub fn replace_path(&mut self, rel: &RelRef) {
        self.0.replace_path(rel)
    }
}

inherits_uri_ref_buf!(UriBuf);

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

    #[test]
    fn test_from_str() {
        assert_eq!(
            <UriBuf as FromStr>::from_str("https://www.google.com/"),
            UriBuf::from_str("https://www.google.com/")
        );
    }
}
<commit_msg>async-coap-uri: impl TryFrom for UriBuf<commit_after>// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

use super::*;
use std::convert::TryFrom;
use std::fmt::Write;
use std::ops::Deref;
use std::str::FromStr;

/// Sized, heap-allocated string type guaranteed to contain a well-formed [IETF-RFC3986] URI
/// or [network path](enum.UriType.html#variant.NetworkPath).
///
/// The unsized counterpart is [`Uri`](crate::Uri).
///
/// This type implements [`std::ops::Deref<Uri>`], so you can also use all of the
/// methods from [`Uri`] on this type.
///
/// [IETF-RFC3986]: https://tools.ietf.org/html/rfc3986
#[derive(Clone, Eq, Hash)]
pub struct UriBuf(pub(super) UriRefBuf);

impl_uri_buf_traits!(UriBuf, Uri);

impl Deref for UriBuf {
    type Target = Uri;

    fn deref(&self) -> &Self::Target {
        self.as_uri()
    }
}

impl AsRef<Uri> for UriBuf {
    fn as_ref(&self) -> &Uri {
        self.as_uri()
    }
}

impl From<&Uri> for UriBuf {
    fn from(x: &Uri) -> Self {
        x.to_uri_buf()
    }
}

impl FromStr for UriBuf {
    type Err = ParseError;

    fn from_str(input: &str) -> Result<Self, Self::Err> {
        Self::from_str(input)
    }
}

impl TryFrom<&str> for UriBuf {
    type Error = ParseError;

    fn try_from(value: &str) -> Result<Self, Self::Error> {
        Self::from_str(value)
    }
}

impl TryFrom<String> for UriBuf {
    type Error = ParseError;

    fn try_from(value: String) -> Result<Self, Self::Error> {
        Self::from_string(value)
    }
}

impl<'a> TryFrom<&'a String> for UriBuf {
    type Error = <Self as TryFrom<&'a str>>::Error;

    fn try_from(value: &'a String) -> Result<Self, Self::Error> {
        <Self as TryFrom<&'a str>>::try_from(value.as_str())
    }
}

impl std::fmt::Display for UriBuf {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::result::Result<(), std::fmt::Error> {
        self.write_to(f)
    }
}

/// # Unsafe Methods
///
/// `UriBuf` needs some unsafe methods in order to function properly. This section is where
/// they are all located.
impl UriBuf {
    /// Unchecked version of [`UriBuf::from_string`].
    ///
    /// # Safety
    ///
    /// This method is marked as unsafe because it allows you to construct a `UriBuf` with
    /// a value that is not a well-formed URI reference.
    #[inline(always)]
    pub unsafe fn from_string_unchecked(s: String) -> UriBuf {
        UriBuf(UriRefBuf::from_string_unchecked(s))
    }
}

impl UriBuf {
    /// Creates a new `UriBuf` from *unescaped* component values.
    pub fn new<Sch, Hos, Pat, Que, Frg>(
        scheme: Sch,
        host: Hos,
        port: Option<u16>,
        path: Pat,
        query: Option<Que>,
        fragment: Option<Frg>,
    ) -> UriBuf
    where
        Sch: Into<String>,
        Hos: AsRef<str>,
        Pat: AsRef<str>,
        Que: AsRef<str>,
        Frg: AsRef<str>,
    {
        let mut ret: String = Self::from_scheme_host_port(scheme, host, port).into();

        let mut path = path.as_ref();

        if path.starts_with('/') {
            path = &path[1..];
        }

        let path_segment_iter =
            path.split('/')
                .filter_map(|seg| if seg == "." { None } else { Some(seg) });

        for seg in path_segment_iter {
            ret.push('/');
            ret.extend(seg.escape_uri());
        }

        if let Some(query) = query {
            let mut first = true;
            ret.push('?');
            for seg in query.as_ref().split(|c| c == '&' || c == ';') {
                if first {
                    first = false;
                } else {
                    ret.push('&');
                }
                ret.extend(seg.escape_uri().for_query());
            }
        }

        if let Some(fragment) = fragment {
            ret.push('#');
            ret.extend(fragment.as_ref().escape_uri().for_fragment());
        }

        unsafe { Self::from_string_unchecked(ret) }
    }

    /// Constructs a `UriBuf` from a scheme and authority.
    ///
    /// The authority should not be percent encoded. If the given scheme contains invalid
    /// characters, this method will panic.
    ///
    /// # Example
    ///
    /// ```
    /// use async_coap_uri::prelude::*;
    /// let authority = "user@[2001:0db8:85a3::1%en2]:8080";
    ///
    /// let uri_buf = UriBuf::from_scheme_authority("http", authority);
    ///
    /// assert_eq!(uri_buf, uri!("http://user@[2001:0db8:85a3::1%25en2]:8080"));
    ///
    /// ```
    pub fn from_scheme_authority<Sch, Aut>(scheme: Sch, authority: Aut) -> UriBuf
    where
        Sch: Into<String>,
        Aut: AsRef<str>,
    {
        let mut ret = scheme.into();

        // Make sure that the scheme is proper. This is likely a string constant,
        // so we go ahead and panic if we detect that it is busted.
        assert_eq!(
            ret.find(|c: char| !(c.is_ascii_alphanumeric() || c == '+' || c == '-' || c == '.')),
            None,
            "Scheme contains invalid characters: {:?}",
            ret
        );

        ret.push_str("://");

        ret.extend(authority.as_ref().escape_uri().for_authority());

        unsafe { Self::from_string_unchecked(ret) }
    }

    /// Constructs a network path from a host and a relative reference.
    ///
    /// # Example
    ///
    /// ```
    /// use async_coap_uri::prelude::*;
    /// let host = "example.com";
    /// let rel_ref = rel_ref!("./foobar?q");
    ///
    /// let uri_buf = UriBuf::from_host_rel_ref(host, rel_ref);
    ///
    /// assert_eq!(uri_buf, uri!("//example.com/foobar?q"));
    ///
    /// ```
    pub fn from_host_rel_ref<Hos, RR>(host: Hos, rel_ref: RR) -> UriBuf
    where
        Hos: AsRef<str>,
        RR: AsRef<RelRef>,
    {
        let host = host.as_ref();
        let rel_ref = rel_ref
            .as_ref()
            .trim_leading_dot_slashes()
            .trim_leading_slashes();

        // UNWRAP-SAFETY: This is safe because we are fully
        // escaping the host and we already know rel_ref to
        // be well-formed.
        uri_format!("//{}/{}", host.escape_uri().full(), rel_ref).unwrap()
    }

    /// Constructs a `UriBuf` from a scheme, host and an optional port number.
    ///
    /// The host should not be percent encoded. If the given scheme contains invalid
    /// characters, this method will panic.
    pub fn from_scheme_host_port<Sch, Hos>(scheme: Sch, host: Hos, port: Option<u16>) -> UriBuf
    where
        Sch: Into<String>,
        Hos: AsRef<str>,
    {
        let mut ret = scheme.into();

        // Make sure that the scheme is proper. This is likely a string constant,
        // so we go ahead and panic if we detect that it is busted.
        assert_eq!(
            ret.find(|c: char| !(c.is_ascii_alphanumeric() || c == '+' || c == '-' || c == '.')),
            None,
            "Scheme contains invalid characters: {:?}",
            ret
        );

        ret.push_str("://");

        let mut host = host.as_ref();

        // Trim enclosing brackets.
        if host.starts_with('[') && host.ends_with(']') {
            host = &host[1..host.len()];
        }

        if host.find(':').is_some() {
            ret.push('[');
            ret.extend(host.escape_uri());
            ret.push(']');
        }

        if let Some(port) = port {
            write!(ret, ":{}", port).unwrap();
        }

        unsafe { Self::from_string_unchecked(ret) }
    }

    /// Attempts to create a new [`UriBuf`] from a string slice.
    pub fn from_str<S: AsRef<str>>(s: S) -> Result<UriBuf, ParseError> {
        let s = s.as_ref();
        let components = UriRawComponents::from_str(s)?;

        if components.uri_type().can_borrow_as_uri() {
            Ok(unsafe { Self::from_string_unchecked(s.to_string()) })
        } else {
            Err(ParseError::new("Missing scheme or authority", None))
        }
    }

    /// Attempts to create a new [`UriBuf`] from a [`String`].
    pub fn from_string(s: String) -> Result<UriBuf, ParseError> {
        let components = UriRawComponents::from_str(s.as_str())?;

        if components.uri_type().can_borrow_as_uri() {
            Ok(unsafe { Self::from_string_unchecked(s) })
        } else {
            Err(ParseError::new("Missing scheme or authority", None))
        }
    }

    /// Attempts to create a new [`UriBuf`] from a `UriRef` slice.
    pub fn from_uri<S: AsRef<UriRef>>(s: S) -> Option<UriBuf> {
        if s.as_ref().uri_type().can_borrow_as_uri() {
            Some(UriBuf(s.as_ref().to_uri_ref_buf()))
        } else {
            None
        }
    }
}

impl UriBuf {
    /// Borrows a [`Uri`] slice containing this URI.
    #[inline(always)]
    pub fn as_uri(&self) -> &Uri {
        unsafe { Uri::from_str_unchecked(self.as_str()) }
    }
}

/// # Manipulation
impl UriBuf {
    /// Using this URI as the base, performs "relative resolution" to the given instance
    /// implementing [`AnyUriRef`], updating the content of this `UriBuf` with the result.
    pub fn resolve<T: AnyUriRef + ?Sized>(&mut self, dest: &T) -> Result<(), ResolveError> {
        self.0.resolve(dest)
    }

    /// Replaces the path, query, and fragment with that from `rel`.
    pub fn replace_path(&mut self, rel: &RelRef) {
        self.0.replace_path(rel)
    }
}

inherits_uri_ref_buf!(UriBuf);

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

    #[test]
    fn test_from_str() {
        assert_eq!(
            <UriBuf as FromStr>::from_str("https://www.google.com/"),
            UriBuf::from_str("https://www.google.com/")
        );
    }

    #[test]
    fn test_try_from() {
        assert_eq!(
            UriBuf::try_from("https://www.google.com/"),
            UriBuf::from_str("https://www.google.com/")
        );

        assert_eq!(
            UriBuf::try_from("https://www.google.com/".to_string()),
            UriBuf::from_str("https://www.google.com/")
        );

        assert_eq!(
            UriBuf::try_from(&"https://www.google.com/".to_string()),
            UriBuf::try_from("https://www.google.com/"),
        );
    }
}
<|endoftext|>


<commit_before>// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

use super::*;
use std::convert::TryFrom;
use std::fmt::Write;
use std::ops::Deref;
use std::str::FromStr;

/// Sized, heap-allocated string type guaranteed to contain a well-formed [IETF-RFC3986] URI
/// or [network path](enum.UriType.html#variant.NetworkPath).
///
/// The unsized counterpart is [`Uri`](crate::Uri).
///
/// This type implements [`std::ops::Deref<Uri>`], so you can also use all of the
/// methods from [`Uri`] on this type.
///
/// [IETF-RFC3986]: https://tools.ietf.org/html/rfc3986
#[derive(Clone, Eq, Hash)]
pub struct UriBuf(pub(super) UriRefBuf);

impl_uri_buf_traits!(UriBuf, Uri);

impl Deref for UriBuf {
    type Target = Uri;

    fn deref(&self) -> &Self::Target {
        self.as_uri()
    }
}

impl AsRef<Uri> for UriBuf {
    fn as_ref(&self) -> &Uri {
        self.as_uri()
    }
}

impl From<&Uri> for UriBuf {
    fn from(x: &Uri) -> Self {
        x.to_uri_buf()
    }
}

impl FromStr for UriBuf {
    type Err = ParseError;

    fn from_str(input: &str) -> Result<Self, Self::Err> {
        Self::from_str(input)
    }
}

impl TryFrom<&str> for UriBuf {
    type Error = ParseError;

    fn try_from(value: &str) -> Result<Self, Self::Error> {
        Self::from_str(value)
    }
}

impl TryFrom<String> for UriBuf {
    type Error = ParseError;

    fn try_from(value: String) -> Result<Self, Self::Error> {
        Self::from_string(value)
    }
}

impl<'a> TryFrom<&'a String> for UriBuf {
    type Error = <Self as TryFrom<&'a str>>::Error;

    fn try_from(value: &'a String) -> Result<Self, Self::Error> {
        <Self as TryFrom<&'a str>>::try_from(value.as_str())
    }
}

impl std::fmt::Display for UriBuf {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::result::Result<(), std::fmt::Error> {
        self.write_to(f)
    }
}

/// # Unsafe Methods
///
/// `UriBuf` needs some unsafe methods in order to function properly. This section is where
/// they are all located.
impl UriBuf {
    /// Unchecked version of [`UriBuf::from_string`].
    ///
    /// # Safety
    ///
    /// This method is marked as unsafe because it allows you to construct a `UriBuf` with
    /// a value that is not a well-formed URI reference.
    #[inline(always)]
    pub unsafe fn from_string_unchecked(s: String) -> UriBuf {
        UriBuf(UriRefBuf::from_string_unchecked(s))
    }
}

impl UriBuf {
    /// Creates a new `UriBuf` from *unescaped* component values.
    pub fn new<Sch, Hos, Pat, Que, Frg>(
        scheme: Sch,
        host: Hos,
        port: Option<u16>,
        path: Pat,
        query: Option<Que>,
        fragment: Option<Frg>,
    ) -> UriBuf
    where
        Sch: Into<String>,
        Hos: AsRef<str>,
        Pat: AsRef<str>,
        Que: AsRef<str>,
        Frg: AsRef<str>,
    {
        let mut ret: String = Self::from_scheme_host_port(scheme, host, port).into();

        let mut path = path.as_ref();

        if path.starts_with('/') {
            path = &path[1..];
        }

        let path_segment_iter =
            path.split('/')
                .filter_map(|seg| if seg == "." { None } else { Some(seg) });

        for seg in path_segment_iter {
            ret.push('/');
            ret.extend(seg.escape_uri());
        }

        if let Some(query) = query {
            let mut first = true;
            ret.push('?');
            for seg in query.as_ref().split(|c| c == '&' || c == ';') {
                if first {
                    first = false;
                } else {
                    ret.push('&');
                }
                ret.extend(seg.escape_uri().for_query());
            }
        }

        if let Some(fragment) = fragment {
            ret.push('#');
            ret.extend(fragment.as_ref().escape_uri().for_fragment());
        }

        unsafe { Self::from_string_unchecked(ret) }
    }

    /// Constructs a `UriBuf` from a scheme and authority.
    ///
    /// The authority should not be percent encoded. If the given scheme contains invalid
    /// characters, this method will panic.
    ///
    /// # Example
    ///
    /// ```
    /// use async_coap_uri::prelude::*;
    /// let authority = "user@[2001:0db8:85a3::1%en2]:8080";
    ///
    /// let uri_buf = UriBuf::from_scheme_authority("http", authority);
    ///
    /// assert_eq!(uri_buf, uri!("http://user@[2001:0db8:85a3::1%25en2]:8080"));
    ///
    /// ```
    pub fn from_scheme_authority<Sch, Aut>(scheme: Sch, authority: Aut) -> UriBuf
    where
        Sch: Into<String>,
        Aut: AsRef<str>,
    {
        let mut ret = scheme.into();

        // Make sure that the scheme is proper. This is likely a string constant,
        // so we go ahead and panic if we detect that it is busted.
        assert_eq!(
            ret.find(|c: char| !(c.is_ascii_alphanumeric() || c == '+' || c == '-' || c == '.')),
            None,
            "Scheme contains invalid characters: {:?}",
            ret
        );

        ret.push_str("://");

        ret.extend(authority.as_ref().escape_uri().for_authority());

        unsafe { Self::from_string_unchecked(ret) }
    }

    /// Constructs a network path from a host and a relative reference.
    ///
    /// # Example
    ///
    /// ```
    /// use async_coap_uri::prelude::*;
    /// let host = "example.com";
    /// let rel_ref = rel_ref!("./foobar?q");
    ///
    /// let uri_buf = UriBuf::from_host_rel_ref(host, rel_ref);
    ///
    /// assert_eq!(uri_buf, uri!("//example.com/foobar?q"));
    ///
    /// ```
    pub fn from_host_rel_ref<Hos, RR>(host: Hos, rel_ref: RR) -> UriBuf
    where
        Hos: AsRef<str>,
        RR: AsRef<RelRef>,
    {
        let host = host.as_ref();
        let rel_ref = rel_ref
            .as_ref()
            .trim_leading_dot_slashes()
            .trim_leading_slashes();

        // UNWRAP-SAFETY: This is safe because we are fully
        // escaping the host and we already know rel_ref to
        // be well-formed.
        uri_format!("//{}/{}", host.escape_uri().full(), rel_ref).unwrap()
    }

    /// Constructs a `UriBuf` from a scheme, host and an optional port number.
    ///
    /// The host should not be percent encoded. If the given scheme contains invalid
    /// characters, this method will panic.
    pub fn from_scheme_host_port<Sch, Hos>(scheme: Sch, host: Hos, port: Option<u16>) -> UriBuf
    where
        Sch: Into<String>,
        Hos: AsRef<str>,
    {
        let mut ret = scheme.into();

        // Make sure that the scheme is proper. This is likely a string constant,
        // so we go ahead and panic if we detect that it is busted.
        assert_eq!(
            ret.find(|c: char| !(c.is_ascii_alphanumeric() || c == '+' || c == '-' || c == '.')),
            None,
            "Scheme contains invalid characters: {:?}",
            ret
        );

        ret.push_str("://");

        let mut host = host.as_ref();

        // Trim enclosing brackets.
        if host.starts_with('[') && host.ends_with(']') {
            host = &host[1..host.len()];
        }

        if host.find(':').is_some() {
            ret.push('[');
            ret.extend(host.escape_uri());
            ret.push(']');
        }

        if let Some(port) = port {
            write!(ret, ":{}", port).unwrap();
        }

        unsafe { Self::from_string_unchecked(ret) }
    }

    /// Attempts to create a new [`UriBuf`] from a string slice.
    pub fn from_str<S: AsRef<str>>(s: S) -> Result<UriBuf, ParseError> {
        let s = s.as_ref();
        let components = UriRawComponents::from_str(s)?;

        if components.uri_type().can_borrow_as_uri() {
            Ok(unsafe { Self::from_string_unchecked(s.to_string()) })
        } else {
            Err(ParseError::new("Missing scheme or authority", None))
        }
    }

    /// Attempts to create a new [`UriBuf`] from a [`String`].
    pub fn from_string(s: String) -> Result<UriBuf, ParseError> {
        let components = UriRawComponents::from_str(s.as_str())?;

        if components.uri_type().can_borrow_as_uri() {
            Ok(unsafe { Self::from_string_unchecked(s) })
        } else {
            Err(ParseError::new("Missing scheme or authority", None))
        }
    }

    /// Attempts to create a new [`UriBuf`] from a `UriRef` slice.
    pub fn from_uri<S: AsRef<UriRef>>(s: S) -> Option<UriBuf> {
        if s.as_ref().uri_type().can_borrow_as_uri() {
            Some(UriBuf(s.as_ref().to_uri_ref_buf()))
        } else {
            None
        }
    }
}

impl UriBuf {
    /// Borrows a [`Uri`] slice containing this URI.
    #[inline(always)]
    pub fn as_uri(&self) -> &Uri {
        unsafe { Uri::from_str_unchecked(self.as_str()) }
    }
}

/// # Manipulation
impl UriBuf {
    /// Using this URI as the base, performs "relative resolution" to the given instance
    /// implementing [`AnyUriRef`], updating the content of this `UriBuf` with the result.
    pub fn resolve<T: AnyUriRef + ?Sized>(&mut self, dest: &T) -> Result<(), ResolveError> {
        self.0.resolve(dest)
    }

    /// Replaces the path, query, and fragment with that from `rel`.
    pub fn replace_path(&mut self, rel: &RelRef) {
        self.0.replace_path(rel)
    }
}

inherits_uri_ref_buf!(UriBuf);

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

    #[test]
    fn test_from_str() {
        assert_eq!(
            <UriBuf as FromStr>::from_str("https://www.google.com/"),
            UriBuf::from_str("https://www.google.com/")
        );
    }

    #[test]
    fn test_try_from() {
        assert_eq!(
            UriBuf::try_from("https://www.google.com/"),
            UriBuf::from_str("https://www.google.com/")
        );

        assert_eq!(
            UriBuf::try_from("https://www.google.com/".to_string()),
            UriBuf::from_str("https://www.google.com/")
        );

        assert_eq!(
            UriBuf::try_from(&"https://www.google.com/".to_string()),
            UriBuf::try_from("https://www.google.com/"),
        );
    }
}
<commit_msg>async-coap-uri: fix doc link for UriBuf<commit_after>// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

use super::*;
use std::convert::TryFrom;
use std::fmt::Write;
use std::ops::Deref;
use std::str::FromStr;

/// Sized, heap-allocated string type guaranteed to contain a well-formed [IETF-RFC3986] URI
/// or [network path](enum.UriType.html#variant.NetworkPath).
///
/// The unsized counterpart is [`Uri`].
///
/// This type implements [`Deref<Target = Uri>`], so you can also use all of the
/// methods from [`Uri`] on this type.
///
/// [IETF-RFC3986]: https://tools.ietf.org/html/rfc3986
/// [`Deref<Target = Uri>`]: std::ops::Deref
#[derive(Clone, Eq, Hash)]
pub struct UriBuf(pub(super) UriRefBuf);

impl_uri_buf_traits!(UriBuf, Uri);

impl Deref for UriBuf {
    type Target = Uri;

    fn deref(&self) -> &Self::Target {
        self.as_uri()
    }
}

impl AsRef<Uri> for UriBuf {
    fn as_ref(&self) -> &Uri {
        self.as_uri()
    }
}

impl From<&Uri> for UriBuf {
    fn from(x: &Uri) -> Self {
        x.to_uri_buf()
    }
}

impl FromStr for UriBuf {
    type Err = ParseError;

    fn from_str(input: &str) -> Result<Self, Self::Err> {
        Self::from_str(input)
    }
}

impl TryFrom<&str> for UriBuf {
    type Error = ParseError;

    fn try_from(value: &str) -> Result<Self, Self::Error> {
        Self::from_str(value)
    }
}

impl TryFrom<String> for UriBuf {
    type Error = ParseError;

    fn try_from(value: String) -> Result<Self, Self::Error> {
        Self::from_string(value)
    }
}

impl<'a> TryFrom<&'a String> for UriBuf {
    type Error = <Self as TryFrom<&'a str>>::Error;

    fn try_from(value: &'a String) -> Result<Self, Self::Error> {
        <Self as TryFrom<&'a str>>::try_from(value.as_str())
    }
}

impl std::fmt::Display for UriBuf {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::result::Result<(), std::fmt::Error> {
        self.write_to(f)
    }
}

/// # Unsafe Methods
///
/// `UriBuf` needs some unsafe methods in order to function properly. This section is where
/// they are all located.
impl UriBuf {
    /// Unchecked version of [`UriBuf::from_string`].
    ///
    /// # Safety
    ///
    /// This method is marked as unsafe because it allows you to construct a `UriBuf` with
    /// a value that is not a well-formed URI reference.
    #[inline(always)]
    pub unsafe fn from_string_unchecked(s: String) -> UriBuf {
        UriBuf(UriRefBuf::from_string_unchecked(s))
    }
}

impl UriBuf {
    /// Creates a new `UriBuf` from *unescaped* component values.
    pub fn new<Sch, Hos, Pat, Que, Frg>(
        scheme: Sch,
        host: Hos,
        port: Option<u16>,
        path: Pat,
        query: Option<Que>,
        fragment: Option<Frg>,
    ) -> UriBuf
    where
        Sch: Into<String>,
        Hos: AsRef<str>,
        Pat: AsRef<str>,
        Que: AsRef<str>,
        Frg: AsRef<str>,
    {
        let mut ret: String = Self::from_scheme_host_port(scheme, host, port).into();

        let mut path = path.as_ref();

        if path.starts_with('/') {
            path = &path[1..];
        }

        let path_segment_iter =
            path.split('/')
                .filter_map(|seg| if seg == "." { None } else { Some(seg) });

        for seg in path_segment_iter {
            ret.push('/');
            ret.extend(seg.escape_uri());
        }

        if let Some(query) = query {
            let mut first = true;
            ret.push('?');
            for seg in query.as_ref().split(|c| c == '&' || c == ';') {
                if first {
                    first = false;
                } else {
                    ret.push('&');
                }
                ret.extend(seg.escape_uri().for_query());
            }
        }

        if let Some(fragment) = fragment {
            ret.push('#');
            ret.extend(fragment.as_ref().escape_uri().for_fragment());
        }

        unsafe { Self::from_string_unchecked(ret) }
    }

    /// Constructs a `UriBuf` from a scheme and authority.
    ///
    /// The authority should not be percent encoded. If the given scheme contains invalid
    /// characters, this method will panic.
    ///
    /// # Example
    ///
    /// ```
    /// use async_coap_uri::prelude::*;
    /// let authority = "user@[2001:0db8:85a3::1%en2]:8080";
    ///
    /// let uri_buf = UriBuf::from_scheme_authority("http", authority);
    ///
    /// assert_eq!(uri_buf, uri!("http://user@[2001:0db8:85a3::1%25en2]:8080"));
    ///
    /// ```
    pub fn from_scheme_authority<Sch, Aut>(scheme: Sch, authority: Aut) -> UriBuf
    where
        Sch: Into<String>,
        Aut: AsRef<str>,
    {
        let mut ret = scheme.into();

        // Make sure that the scheme is proper. This is likely a string constant,
        // so we go ahead and panic if we detect that it is busted.
        assert_eq!(
            ret.find(|c: char| !(c.is_ascii_alphanumeric() || c == '+' || c == '-' || c == '.')),
            None,
            "Scheme contains invalid characters: {:?}",
            ret
        );

        ret.push_str("://");

        ret.extend(authority.as_ref().escape_uri().for_authority());

        unsafe { Self::from_string_unchecked(ret) }
    }

    /// Constructs a network path from a host and a relative reference.
    ///
    /// # Example
    ///
    /// ```
    /// use async_coap_uri::prelude::*;
    /// let host = "example.com";
    /// let rel_ref = rel_ref!("./foobar?q");
    ///
    /// let uri_buf = UriBuf::from_host_rel_ref(host, rel_ref);
    ///
    /// assert_eq!(uri_buf, uri!("//example.com/foobar?q"));
    ///
    /// ```
    pub fn from_host_rel_ref<Hos, RR>(host: Hos, rel_ref: RR) -> UriBuf
    where
        Hos: AsRef<str>,
        RR: AsRef<RelRef>,
    {
        let host = host.as_ref();
        let rel_ref = rel_ref
            .as_ref()
            .trim_leading_dot_slashes()
            .trim_leading_slashes();

        // UNWRAP-SAFETY: This is safe because we are fully
        // escaping the host and we already know rel_ref to
        // be well-formed.
        uri_format!("//{}/{}", host.escape_uri().full(), rel_ref).unwrap()
    }

    /// Constructs a `UriBuf` from a scheme, host and an optional port number.
    ///
    /// The host should not be percent encoded. If the given scheme contains invalid
    /// characters, this method will panic.
    pub fn from_scheme_host_port<Sch, Hos>(scheme: Sch, host: Hos, port: Option<u16>) -> UriBuf
    where
        Sch: Into<String>,
        Hos: AsRef<str>,
    {
        let mut ret = scheme.into();

        // Make sure that the scheme is proper. This is likely a string constant,
        // so we go ahead and panic if we detect that it is busted.
        assert_eq!(
            ret.find(|c: char| !(c.is_ascii_alphanumeric() || c == '+' || c == '-' || c == '.')),
            None,
            "Scheme contains invalid characters: {:?}",
            ret
        );

        ret.push_str("://");

        let mut host = host.as_ref();

        // Trim enclosing brackets.
        if host.starts_with('[') && host.ends_with(']') {
            host = &host[1..host.len()];
        }

        if host.find(':').is_some() {
            ret.push('[');
            ret.extend(host.escape_uri());
            ret.push(']');
        }

        if let Some(port) = port {
            write!(ret, ":{}", port).unwrap();
        }

        unsafe { Self::from_string_unchecked(ret) }
    }

    /// Attempts to create a new [`UriBuf`] from a string slice.
    pub fn from_str<S: AsRef<str>>(s: S) -> Result<UriBuf, ParseError> {
        let s = s.as_ref();
        let components = UriRawComponents::from_str(s)?;

        if components.uri_type().can_borrow_as_uri() {
            Ok(unsafe { Self::from_string_unchecked(s.to_string()) })
        } else {
            Err(ParseError::new("Missing scheme or authority", None))
        }
    }

    /// Attempts to create a new [`UriBuf`] from a [`String`].
    pub fn from_string(s: String) -> Result<UriBuf, ParseError> {
        let components = UriRawComponents::from_str(s.as_str())?;

        if components.uri_type().can_borrow_as_uri() {
            Ok(unsafe { Self::from_string_unchecked(s) })
        } else {
            Err(ParseError::new("Missing scheme or authority", None))
        }
    }

    /// Attempts to create a new [`UriBuf`] from a `UriRef` slice.
    pub fn from_uri<S: AsRef<UriRef>>(s: S) -> Option<UriBuf> {
        if s.as_ref().uri_type().can_borrow_as_uri() {
            Some(UriBuf(s.as_ref().to_uri_ref_buf()))
        } else {
            None
        }
    }
}

impl UriBuf {
    /// Borrows a [`Uri`] slice containing this URI.
    #[inline(always)]
    pub fn as_uri(&self) -> &Uri {
        unsafe { Uri::from_str_unchecked(self.as_str()) }
    }
}

/// # Manipulation
impl UriBuf {
    /// Using this URI as the base, performs "relative resolution" to the given instance
    /// implementing [`AnyUriRef`], updating the content of this `UriBuf` with the result.
    pub fn resolve<T: AnyUriRef + ?Sized>(&mut self, dest: &T) -> Result<(), ResolveError> {
        self.0.resolve(dest)
    }

    /// Replaces the path, query, and fragment with that from `rel`.
    pub fn replace_path(&mut self, rel: &RelRef) {
        self.0.replace_path(rel)
    }
}

inherits_uri_ref_buf!(UriBuf);

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

    #[test]
    fn test_from_str() {
        assert_eq!(
            <UriBuf as FromStr>::from_str("https://www.google.com/"),
            UriBuf::from_str("https://www.google.com/")
        );
    }

    #[test]
    fn test_try_from() {
        assert_eq!(
            UriBuf::try_from("https://www.google.com/"),
            UriBuf::from_str("https://www.google.com/")
        );

        assert_eq!(
            UriBuf::try_from("https://www.google.com/".to_string()),
            UriBuf::from_str("https://www.google.com/")
        );

        assert_eq!(
            UriBuf::try_from(&"https://www.google.com/".to_string()),
            UriBuf::try_from("https://www.google.com/"),
        );
    }
}
<|endoftext|>


<commit_before>use std::collections::VecDeque;

use search::term::Term;
use search::index::Index;
use search::index::reader::IndexReader;
use search::query::{Query, TermMatcher};

pub enum QuerySetIterator<'a, T: IndexReader<'a>> {
    None,
    All {
        iter: T::AllDocRefIterator,
    },
    Term {
        iter: T::TermDocRefIterator,
    },
    Conjunction {
        iter_a: Box<QuerySetIterator<'a, T>>,
        iter_b: Box<QuerySetIterator<'a, T>>,
        initialised: bool,
        current_doc_a: Option<u64>,
        current_doc_b: Option<u64>,
    },
    Disjunction {
        iter_a: Box<QuerySetIterator<'a, T>>,
        iter_b: Box<QuerySetIterator<'a, T>>,
        initialised: bool,
        current_doc_a: Option<u64>,
        current_doc_b: Option<u64>,
    },
    Exclusion {
        iter_a: Box<QuerySetIterator<'a, T>>,
        iter_b: Box<QuerySetIterator<'a, T>>,
        initialised: bool,
        current_doc_a: Option<u64>,
        current_doc_b: Option<u64>,
    },
}

impl <'a, T: IndexReader<'a>> Iterator for QuerySetIterator<'a, T> {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        match *self {
            QuerySetIterator::None => None,
            QuerySetIterator::All{ref mut iter} => {
                iter.next()
            }
            QuerySetIterator::Term{ref mut iter} => {
                iter.next()
            }
            QuerySetIterator::Conjunction{ref mut iter_a, ref mut iter_b, ref mut initialised, ref mut current_doc_a, ref mut current_doc_b} => {
                if !*initialised {
                    *current_doc_a = iter_a.next();
                    *current_doc_b = iter_b.next();
                    *initialised = true;
                }

                loop {
                    if let Some(doc_id_a) = *current_doc_a {
                        if let Some(doc_id_b) = *current_doc_b {
                            if doc_id_a == doc_id_b {
                                *current_doc_a = iter_a.next();
                                *current_doc_b = iter_b.next();

                                return Some(doc_id_a);
                            } else if doc_id_a < doc_id_b {
                                *current_doc_a = iter_a.next();
                            } else {
                                *current_doc_b = iter_b.next();
                            }
                        } else {
                            return None;
                        }
                    } else {
                        return None;
                    }
                }
            }
            QuerySetIterator::Disjunction{ref mut iter_a, ref mut iter_b, ref mut initialised, ref mut current_doc_a, ref mut current_doc_b} => {
                if !*initialised {
                    *current_doc_a = iter_a.next();
                    *current_doc_b = iter_b.next();
                    *initialised = true;
                }

                if let Some(doc_id_a) = *current_doc_a {
                    if let Some(doc_id_b) = *current_doc_b {
                        if doc_id_a == doc_id_b {
                            *current_doc_a = iter_a.next();
                            *current_doc_b = iter_a.next();
                            Some(doc_id_a)
                        } else if doc_id_a < doc_id_b {
                            *current_doc_a = iter_a.next();
                            Some(doc_id_a)
                        } else {
                            *current_doc_b = iter_b.next();
                            Some(doc_id_b)
                        }
                    } else {
                        *current_doc_a = iter_a.next();
                        Some(doc_id_a)
                    }
                } else {
                    if let Some(doc_id_b) = *current_doc_b {
                        *current_doc_b = iter_b.next();
                        Some(doc_id_b)
                    } else {
                        None
                    }
                }
            }
            QuerySetIterator::Exclusion{ref mut iter_a, ref mut iter_b, ref mut initialised, ref mut current_doc_a, ref mut current_doc_b} => {
                if !*initialised {
                    *current_doc_a = iter_a.next();
                    *current_doc_b = iter_b.next();
                    *initialised = true;
                }

                loop {
                    if let Some(doc_id_a) = *current_doc_a {
                        if let Some(doc_id_b) = *current_doc_b {
                            if doc_id_a == doc_id_b {
                                *current_doc_a = iter_a.next();
                                *current_doc_b = iter_b.next();
                            } else if doc_id_a < doc_id_b {
                                *current_doc_a = iter_a.next();
                                return Some(doc_id_a);
                            } else {
                                *current_doc_b = iter_b.next();
                            }
                        } else {
                            *current_doc_a = iter_a.next();
                            return Some(doc_id_a);
                        }
                    } else {
                        return None;
                    }
                }
            }
        }
    }
}

fn build_conjunction_iterator<'a, T: IndexReader<'a>>(mut iters: VecDeque<QuerySetIterator<'a, T>>) -> QuerySetIterator<'a, T>  {
    if iters.len() == 0 {
        return QuerySetIterator::None;
    }

    // TODO: Order by lowest probability first

    let mut new_iters = VecDeque::with_capacity(iters.len() / 2 + 1);
    while iters.len() >= 2 {
        new_iters.push_back(QuerySetIterator::Conjunction {
            iter_a: Box::new(iters.pop_front().unwrap()),
            iter_b: Box::new(iters.pop_back().unwrap()),
            initialised: false,
            current_doc_a: None,
            current_doc_b: None,
        });
    }

    if iters.len() == 1 {
        new_iters.push_back(iters.pop_front().unwrap());
    }

    if new_iters.len() == 1 {
        // Done!
        new_iters.pop_front().unwrap()
    } else if new_iters.len() > 1 {
        // Still some compaction to do
        build_conjunction_iterator(new_iters)
    } else {
        // Check at top of this function should prevent this
        unreachable!()
    }
}

fn build_disjunction_iterator<'a, T: IndexReader<'a>>(mut iters: VecDeque<QuerySetIterator<'a, T>>) -> QuerySetIterator<'a, T> {
    if iters.len() == 0 {
        return QuerySetIterator::None;
    }

    // TODO: Order by lowest probability first

    let mut new_iters = VecDeque::with_capacity(iters.len() / 2 + 1);
    while iters.len() >= 2 {
        new_iters.push_back(QuerySetIterator::Disjunction {
            iter_a: Box::new(iters.pop_front().unwrap()),
            iter_b: Box::new(iters.pop_back().unwrap()),
            initialised: false,
            current_doc_a: None,
            current_doc_b: None,
        });
    }

    if iters.len() == 1 {
        new_iters.push_back(iters.pop_front().unwrap());
    }

    if new_iters.len() == 1 {
        // Done!
        new_iters.pop_front().unwrap()
    } else if new_iters.len() > 1 {
        // Still some compaction to do
        build_disjunction_iterator(new_iters)
    } else {
        // Check at top of this function should prevent this
        unreachable!()
    }
}

pub fn build_iterator_from_query<'a, T: IndexReader<'a>>(reader: &'a T, query: &Query) -> QuerySetIterator<'a, T> {
    match *query {
        Query::MatchAll => {
            QuerySetIterator::All {
                iter: reader.iter_docids_all(),
            }
        }
        Query::MatchNone => {
            QuerySetIterator::None
        }
        Query::MatchTerm{ref field, ref term, ref matcher} => {
            match *matcher {
                TermMatcher::Exact => {
                    QuerySetIterator::Term {
                        iter: reader.iter_docids_with_term(term.clone(), field.clone()),
                    }
                }
                TermMatcher::Prefix => {
                    // Find all terms in the index that match the prefix
                    let terms = match *term {
                         Term::String(ref term) => {
                             reader.iter_terms().filter_map(|k| {
                                 if let Term::String(ref k) = *k {
                                     if k.starts_with(term) {
                                         return Some(Term::String(k.clone()));
                                     }
                                 }

                                 None
                             }).collect::<Vec<Term>>()
                         }
                         _ => return QuerySetIterator::None,
                    };

                    match terms.len() {
                        0 => QuerySetIterator::None,
                        1 => {
                            QuerySetIterator::Term {
                                iter: reader.iter_docids_with_term(terms[0].clone(), field.clone()),
                            }
                        }
                        _ => {
                            // Produce a disjunction iterator for all the terms
                            let mut iters = VecDeque::new();
                            for term in terms.iter() {
                                iters.push_back(QuerySetIterator::Term {
                                    iter: reader.iter_docids_with_term(term.clone(), field.clone()),
                                });
                            }

                            build_disjunction_iterator(iters)
                        }
                    }
                }
            }
        }
        Query::Conjunction{ref queries} => {
            let mut iters = VecDeque::with_capacity(queries.len());

            for query in queries.iter() {
                iters.push_back(build_iterator_from_query(reader, query));
            }

            build_conjunction_iterator(iters)
        }
        Query::Disjunction{ref queries} => {
            let mut iters = VecDeque::with_capacity(queries.len());

            for query in queries.iter() {
                iters.push_back(build_iterator_from_query(reader, query));
            }

            build_disjunction_iterator(iters)
        }
        Query::NDisjunction{ref queries, minimum_should_match} => {
            // TODO
            QuerySetIterator::None
        }
        Query::DisjunctionMax{ref queries} => {
            let mut iters = VecDeque::with_capacity(queries.len());

            for query in queries.iter() {
                iters.push_back(build_iterator_from_query(reader, query));
            }

            build_disjunction_iterator(iters)
        }
        Query::Exclude{ref query, ref exclude} => {
            QuerySetIterator::Exclusion {
                iter_a: Box::new(build_iterator_from_query(reader, query)),
                iter_b: Box::new(build_iterator_from_query(reader, exclude)),
                initialised: false,
                current_doc_a: None,
                current_doc_b: None,
            }
        }
        Query::Filter{ref query, ref filter} => {
            QuerySetIterator::Conjunction {
                iter_a: Box::new(build_iterator_from_query(reader, query)),
                iter_b: Box::new(build_iterator_from_query(reader, filter)),
                initialised: false,
                current_doc_a: None,
                current_doc_b: None,
            }
        }
        Query::Score{ref query, mul, add} => {
            build_iterator_from_query(reader, query)
        }
    }
}

pub fn next_doc(index: &Index, query: &Query, position: Option<u64>) -> Option<(u64, f64)> {
    match *query {
        Query::MatchAll => {
            match index.store.next_doc_all(position) {
                Some(doc_id) => Some((doc_id, 1.0f64)),
                None => None
            }
        }
        Query::MatchNone => None,
        Query::MatchTerm{ref field, ref term, ref matcher} => {
            match *matcher {
                TermMatcher::Exact => {
                    match index.store.next_doc(term, field, position) {
                        Some((doc_id, term_freq)) => {
                            let score: f64 = term_freq as f64; // TODO: better scoring

                            Some((doc_id, score))
                        }
                        None => None
                    }
                }
                TermMatcher::Prefix => {
                    // Find all terms in the index that match the prefix
                    let terms = match *term {
                         Term::String(ref term) => {
                             index.store.iter_terms().filter_map(|k| {
                                 if let Term::String(ref k) = *k {
                                     if k.starts_with(term) {
                                         return Some(Term::String(k.clone()));
                                     }
                                 }

                                 None
                             }).collect::<Vec<Term>>()
                         }
                         _ => return None,
                    };

                    // Generate a sub query to match those terms
                    let sub_query = match terms.len() {
                        0 => return None,
                        1 => {
                            Query::MatchTerm {
                                field: field.clone(),
                                term: terms[0].clone(),
                                matcher: TermMatcher::Exact,
                            }
                        }
                        _ => {
                            // Produce a disjunction query for all the terms
                            // FIXME: Building a query on every call to next_doc could get very
                            // slow if there are many terms that match the prefix
                            Query::Disjunction {
                                queries: terms.iter().map(|term| {
                                    Query::MatchTerm {
                                        field: field.clone(),
                                        term: term.clone(),
                                        matcher: TermMatcher::Exact,
                                    }
                                }).collect(),
                            }
                        }
                    };

                    // Run the sub query
                    next_doc(index, &sub_query, position)
                }
            }
        }
        Query::Conjunction{ref queries} => {
            // Iterate the docs returned from the first query and check they match the other queries
            // TODO: Instead of just using the first query, we should pick the query that's least
            //       likely to match something so we don't have to shift though as much junk.
            let first_query = match queries.first() {
                Some(query) => query,
                None => return None,
            };

            let mut position = next_doc(index, first_query, position);

            while let Some(pos) = position {
                let mut broke = false;
                let mut total_score = pos.1;

                for query in queries.iter().skip(1) {
                    // Check the next query matches this document
                    match next_doc(index, query, Some(pos.0 - 1)) {
                        Some(query_doc) => {
                            if query_doc.0 != pos.0 {
                                // This query doesn't match the document
                                broke = true;
                                break;
                            }

                            total_score += query_doc.1
                        }
                        None => {
                            // This query doesn't match the document
                            broke = true;
                            break;
                        }
                    }
                }

                if !broke {
                    return Some((pos.0, total_score / queries.len() as f64));
                }

                position = next_doc(index, first_query, Some(pos.0));
            }

            None
        }
        Query::Disjunction{ref queries} => {
            // Get next doc for each query and return lowest one
            // If multiple queries match the same document, average out the scores
            let mut lowest_doc: Option<(u64, f64)> = None;

            for query in queries.iter() {
                if let Some((doc_id, score)) = next_doc(index, query, position) {
                    match lowest_doc {
                        Some((matched_doc_id, total_score)) => {
                            if doc_id < matched_doc_id {
                                // Before current matched doc. Override it
                                lowest_doc = Some((doc_id, score));
                            } else if doc_id == matched_doc_id {
                                // Same doc, merge scores
                                lowest_doc = Some((doc_id, total_score + score))
                            }
                        }
                        None => {
                            // First match
                            lowest_doc = Some((doc_id, score));
                        }
                    }
                }
            }

            match lowest_doc {
                Some((matched_doc_id, total_score)) => Some((matched_doc_id, total_score / queries.len() as f64)),
                None => None,
            }
        }
        Query::NDisjunction{ref queries, minimum_should_match} => {
            None
        }
        Query::DisjunctionMax{ref queries} => {
            // Get next doc for each query and return lowest one
            // If multiple queries match the same document, average out the scores
            let mut lowest_doc: Option<(u64, f64)> = None;

            for query in queries.iter() {
                if let Some((doc_id, score)) = next_doc(index, query, position) {
                    match lowest_doc {
                        Some((matched_doc_id, max_score)) => {
                            if doc_id < matched_doc_id {
                                // Before current matched doc. Override it
                                lowest_doc = Some((doc_id, score));
                            } else if doc_id == matched_doc_id {
                                // Same doc, merge scores
                                if score > max_score {
                                    lowest_doc = Some((doc_id, score))
                                }
                            }
                        }
                        None => {
                            // First match
                            lowest_doc = Some((doc_id, score));
                        }
                    }
                }
            }

            lowest_doc
        }
        Query::Filter{ref query, ref filter} => {
            let mut query_doc = match next_doc(index, query, position) {
                Some((doc_id, score)) => (doc_id, score),
                None => return None,
            };

            loop {
                // See if there is a filter in the same doc
                let filter_doc = match next_doc(index, filter, Some(query_doc.0 - 1)) {
                    Some((doc_id, score)) => doc_id,
                    None => return None,
                };

                if query_doc.0 == filter_doc {
                    return Some(query_doc)
                }

                query_doc = match next_doc(index, query, Some(query_doc.0)) {
                    Some((doc_id, score)) => (doc_id, score),
                    None => return None,
                };
            }
        }
        Query::Exclude{ref query, ref exclude} => {
            let mut query_doc = match next_doc(index, query, position) {
                Some((doc_id, score)) => (doc_id, score),
                None => return None,
            };

            loop {
                // See if there is not a filter in the same doc
                let exclude_doc = match next_doc(index, exclude, Some(query_doc.0 - 1)) {
                    Some((doc_id, score)) => doc_id,
                    None => return Some(query_doc),
                };

                if query_doc.0 == exclude_doc {
                    query_doc = match next_doc(index, query, Some(query_doc.0)) {
                        Some((doc_id, score)) => (doc_id, score),
                        None => return None,
                    };
                    continue
                }

                return Some(query_doc)
            }
        }
        Query::Score{ref query, mul, add} => {
            match next_doc(index, query, position) {
                Some((doc_id, score)) => {
                    Some((doc_id, score.mul_add(mul, add)))
                }
                None => None
            }
        }
    }
}
<commit_msg>Removed next_doc function<commit_after>use std::collections::VecDeque;

use search::term::Term;
use search::index::reader::IndexReader;
use search::query::{Query, TermMatcher};

pub enum QuerySetIterator<'a, T: IndexReader<'a>> {
    None,
    All {
        iter: T::AllDocRefIterator,
    },
    Term {
        iter: T::TermDocRefIterator,
    },
    Conjunction {
        iter_a: Box<QuerySetIterator<'a, T>>,
        iter_b: Box<QuerySetIterator<'a, T>>,
        initialised: bool,
        current_doc_a: Option<u64>,
        current_doc_b: Option<u64>,
    },
    Disjunction {
        iter_a: Box<QuerySetIterator<'a, T>>,
        iter_b: Box<QuerySetIterator<'a, T>>,
        initialised: bool,
        current_doc_a: Option<u64>,
        current_doc_b: Option<u64>,
    },
    Exclusion {
        iter_a: Box<QuerySetIterator<'a, T>>,
        iter_b: Box<QuerySetIterator<'a, T>>,
        initialised: bool,
        current_doc_a: Option<u64>,
        current_doc_b: Option<u64>,
    },
}

impl <'a, T: IndexReader<'a>> Iterator for QuerySetIterator<'a, T> {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        match *self {
            QuerySetIterator::None => None,
            QuerySetIterator::All{ref mut iter} => {
                iter.next()
            }
            QuerySetIterator::Term{ref mut iter} => {
                iter.next()
            }
            QuerySetIterator::Conjunction{ref mut iter_a, ref mut iter_b, ref mut initialised, ref mut current_doc_a, ref mut current_doc_b} => {
                if !*initialised {
                    *current_doc_a = iter_a.next();
                    *current_doc_b = iter_b.next();
                    *initialised = true;
                }

                loop {
                    if let Some(doc_id_a) = *current_doc_a {
                        if let Some(doc_id_b) = *current_doc_b {
                            if doc_id_a == doc_id_b {
                                *current_doc_a = iter_a.next();
                                *current_doc_b = iter_b.next();

                                return Some(doc_id_a);
                            } else if doc_id_a < doc_id_b {
                                *current_doc_a = iter_a.next();
                            } else {
                                *current_doc_b = iter_b.next();
                            }
                        } else {
                            return None;
                        }
                    } else {
                        return None;
                    }
                }
            }
            QuerySetIterator::Disjunction{ref mut iter_a, ref mut iter_b, ref mut initialised, ref mut current_doc_a, ref mut current_doc_b} => {
                if !*initialised {
                    *current_doc_a = iter_a.next();
                    *current_doc_b = iter_b.next();
                    *initialised = true;
                }

                if let Some(doc_id_a) = *current_doc_a {
                    if let Some(doc_id_b) = *current_doc_b {
                        if doc_id_a == doc_id_b {
                            *current_doc_a = iter_a.next();
                            *current_doc_b = iter_a.next();
                            Some(doc_id_a)
                        } else if doc_id_a < doc_id_b {
                            *current_doc_a = iter_a.next();
                            Some(doc_id_a)
                        } else {
                            *current_doc_b = iter_b.next();
                            Some(doc_id_b)
                        }
                    } else {
                        *current_doc_a = iter_a.next();
                        Some(doc_id_a)
                    }
                } else {
                    if let Some(doc_id_b) = *current_doc_b {
                        *current_doc_b = iter_b.next();
                        Some(doc_id_b)
                    } else {
                        None
                    }
                }
            }
            QuerySetIterator::Exclusion{ref mut iter_a, ref mut iter_b, ref mut initialised, ref mut current_doc_a, ref mut current_doc_b} => {
                if !*initialised {
                    *current_doc_a = iter_a.next();
                    *current_doc_b = iter_b.next();
                    *initialised = true;
                }

                loop {
                    if let Some(doc_id_a) = *current_doc_a {
                        if let Some(doc_id_b) = *current_doc_b {
                            if doc_id_a == doc_id_b {
                                *current_doc_a = iter_a.next();
                                *current_doc_b = iter_b.next();
                            } else if doc_id_a < doc_id_b {
                                *current_doc_a = iter_a.next();
                                return Some(doc_id_a);
                            } else {
                                *current_doc_b = iter_b.next();
                            }
                        } else {
                            *current_doc_a = iter_a.next();
                            return Some(doc_id_a);
                        }
                    } else {
                        return None;
                    }
                }
            }
        }
    }
}

fn build_conjunction_iterator<'a, T: IndexReader<'a>>(mut iters: VecDeque<QuerySetIterator<'a, T>>) -> QuerySetIterator<'a, T>  {
    if iters.len() == 0 {
        return QuerySetIterator::None;
    }

    // TODO: Order by lowest probability first

    let mut new_iters = VecDeque::with_capacity(iters.len() / 2 + 1);
    while iters.len() >= 2 {
        new_iters.push_back(QuerySetIterator::Conjunction {
            iter_a: Box::new(iters.pop_front().unwrap()),
            iter_b: Box::new(iters.pop_back().unwrap()),
            initialised: false,
            current_doc_a: None,
            current_doc_b: None,
        });
    }

    if iters.len() == 1 {
        new_iters.push_back(iters.pop_front().unwrap());
    }

    if new_iters.len() == 1 {
        // Done!
        new_iters.pop_front().unwrap()
    } else if new_iters.len() > 1 {
        // Still some compaction to do
        build_conjunction_iterator(new_iters)
    } else {
        // Check at top of this function should prevent this
        unreachable!()
    }
}

fn build_disjunction_iterator<'a, T: IndexReader<'a>>(mut iters: VecDeque<QuerySetIterator<'a, T>>) -> QuerySetIterator<'a, T> {
    if iters.len() == 0 {
        return QuerySetIterator::None;
    }

    // TODO: Order by lowest probability first

    let mut new_iters = VecDeque::with_capacity(iters.len() / 2 + 1);
    while iters.len() >= 2 {
        new_iters.push_back(QuerySetIterator::Disjunction {
            iter_a: Box::new(iters.pop_front().unwrap()),
            iter_b: Box::new(iters.pop_back().unwrap()),
            initialised: false,
            current_doc_a: None,
            current_doc_b: None,
        });
    }

    if iters.len() == 1 {
        new_iters.push_back(iters.pop_front().unwrap());
    }

    if new_iters.len() == 1 {
        // Done!
        new_iters.pop_front().unwrap()
    } else if new_iters.len() > 1 {
        // Still some compaction to do
        build_disjunction_iterator(new_iters)
    } else {
        // Check at top of this function should prevent this
        unreachable!()
    }
}

pub fn build_iterator_from_query<'a, T: IndexReader<'a>>(reader: &'a T, query: &Query) -> QuerySetIterator<'a, T> {
    match *query {
        Query::MatchAll => {
            QuerySetIterator::All {
                iter: reader.iter_docids_all(),
            }
        }
        Query::MatchNone => {
            QuerySetIterator::None
        }
        Query::MatchTerm{ref field, ref term, ref matcher} => {
            match *matcher {
                TermMatcher::Exact => {
                    QuerySetIterator::Term {
                        iter: reader.iter_docids_with_term(term.clone(), field.clone()),
                    }
                }
                TermMatcher::Prefix => {
                    // Find all terms in the index that match the prefix
                    let terms = match *term {
                         Term::String(ref term) => {
                             reader.iter_terms().filter_map(|k| {
                                 if let Term::String(ref k) = *k {
                                     if k.starts_with(term) {
                                         return Some(Term::String(k.clone()));
                                     }
                                 }

                                 None
                             }).collect::<Vec<Term>>()
                         }
                         _ => return QuerySetIterator::None,
                    };

                    match terms.len() {
                        0 => QuerySetIterator::None,
                        1 => {
                            QuerySetIterator::Term {
                                iter: reader.iter_docids_with_term(terms[0].clone(), field.clone()),
                            }
                        }
                        _ => {
                            // Produce a disjunction iterator for all the terms
                            let mut iters = VecDeque::new();
                            for term in terms.iter() {
                                iters.push_back(QuerySetIterator::Term {
                                    iter: reader.iter_docids_with_term(term.clone(), field.clone()),
                                });
                            }

                            build_disjunction_iterator(iters)
                        }
                    }
                }
            }
        }
        Query::Conjunction{ref queries} => {
            let mut iters = VecDeque::with_capacity(queries.len());

            for query in queries.iter() {
                iters.push_back(build_iterator_from_query(reader, query));
            }

            build_conjunction_iterator(iters)
        }
        Query::Disjunction{ref queries} => {
            let mut iters = VecDeque::with_capacity(queries.len());

            for query in queries.iter() {
                iters.push_back(build_iterator_from_query(reader, query));
            }

            build_disjunction_iterator(iters)
        }
        Query::NDisjunction{ref queries, minimum_should_match} => {
            // TODO
            QuerySetIterator::None
        }
        Query::DisjunctionMax{ref queries} => {
            let mut iters = VecDeque::with_capacity(queries.len());

            for query in queries.iter() {
                iters.push_back(build_iterator_from_query(reader, query));
            }

            build_disjunction_iterator(iters)
        }
        Query::Exclude{ref query, ref exclude} => {
            QuerySetIterator::Exclusion {
                iter_a: Box::new(build_iterator_from_query(reader, query)),
                iter_b: Box::new(build_iterator_from_query(reader, exclude)),
                initialised: false,
                current_doc_a: None,
                current_doc_b: None,
            }
        }
        Query::Filter{ref query, ref filter} => {
            QuerySetIterator::Conjunction {
                iter_a: Box::new(build_iterator_from_query(reader, query)),
                iter_b: Box::new(build_iterator_from_query(reader, filter)),
                initialised: false,
                current_doc_a: None,
                current_doc_b: None,
            }
        }
        Query::Score{ref query, mul, add} => {
            build_iterator_from_query(reader, query)
        }
    }
}
<|endoftext|>