<commit_before>use hash::KEY_SIZE;
use hash::Hash160;
use std::net;
use std::collections::VecDeque;
use std::mem;

const K : usize = 20;
const BUCKET_DEPTH : usize = K;

/// Kademlia routing table, with 160 buckets of `BUCKET_DEPTH` (k) node
/// identifiers each.
///
/// The structure employs least-recently seen eviction. Conflicts generated
/// by evicting a node by inserting a newer one remain tracked, so they can
/// can be resolved later.
pub struct RoutingTable {
   parent_key : Hash160,
   buckets    : Vec<Bucket>,
   conflicts  : Vec<EvictionConflict>,
}

/// Bucket size is estimated to be small enough not to warrant
/// the downsides of using a linked list.
#[derive(Debug, Clone)]
struct Bucket {
   entries  : VecDeque<NodeInfo>,
}

/// Represents a conflict derived from attempting to insert a node in a full
/// bucket.
#[derive(Debug,Clone)]
struct EvictionConflict {
   evicted  : NodeInfo,
   inserted : NodeInfo
}

#[derive(Debug, Clone, PartialEq)]
pub struct NodeInfo {
   pub key  : Hash160,
   pub ip   : net::IpAddr,
   pub port : u16,
}

#[derive(Debug, PartialEq)]
pub enum LookupResult {
   Found(NodeInfo),
   ClosestNodes(Vec<NodeInfo>),
   Myself,
}

impl RoutingTable {
   pub fn new(parent_key : Hash160) -> RoutingTable {
      let mut buckets = Vec::<Bucket>::with_capacity(KEY_SIZE);
      for _ in 0..KEY_SIZE {
         buckets.push(Bucket::new());
      }

      RoutingTable {
         parent_key : parent_key,
         buckets : buckets,
         conflicts  : Vec::new()
      }
   }

   /// Inserts a node in the routing table. Employs least-recently-seen eviction
   /// by kicking out the oldest node in case the bucket is full, and registering
   /// an eviction conflict that can be revised later.
   pub fn insert_node(&mut self, info : NodeInfo) {
      if let Some(index) = self.bucket_for_node(&info.key) {
         let bucket = &mut self.buckets[index];
         bucket.entries.retain(|ref stored_info| info.key != stored_info.key);
         if bucket.entries.len() == BUCKET_DEPTH {
            let conflict = EvictionConflict {
               evicted  : bucket.entries.pop_front().unwrap(),
               inserted : info.clone()
            };
            self.conflicts.push(conflict);
         }
         bucket.entries.push_back(info);
      }
   }

   /// Performs a node lookup on the routing table. The lookup result may
   /// contain the specific node, a list of up to the N closest nodes, or
   /// report that the parent node itself was requested.
   pub fn lookup(&self, key : &Hash160, n : usize) -> LookupResult {
      if *key == self.parent_key {
         return LookupResult::Myself;
      }

      match self.specific_node(key) {
         Some(info) => LookupResult::Found(info),
         None => LookupResult::ClosestNodes(self.closest_n_nodes_to(key, n)),
      }
   }

   /// Returns a table entry for the specific node with a given hash.
   fn specific_node(&self, key : &Hash160) -> Option<NodeInfo> {
      if let Some(index) = self.bucket_for_node(key){
         let bucket = &self.buckets[index];
         return bucket.entries.iter().find(|ref info| *key == info.key).cloned();
      }
      None
   }

   fn closest_n_nodes_to(&self, key : &Hash160, n : usize) -> Vec<NodeInfo>{
      let mut closest = Vec::with_capacity(n);
      // Unwrap because we shouldn't be calling this private method
      // with our own key as argument.
      let ideal_bucket_index = self.bucket_for_node(key).unwrap();

      let ascent_to_ideal = 1..ideal_bucket_index;
      let ascent_from_above_ideal = (ideal_bucket_index+1)..KEY_SIZE;
      let lookup_order = (ideal_bucket_index..ideal_bucket_index)
                         .chain(ascent_to_ideal)
                         .chain(ascent_from_above_ideal);

      for bucket_index in lookup_order {
         let mut nodes_from_bucket = self.buckets[bucket_index].clone().entries.into_iter().collect::<Vec<NodeInfo>>();
         nodes_from_bucket.sort_by_key(|ref info| &info.key ^ key);
         let space_left = closest.capacity() - closest.len();
         nodes_from_bucket.truncate(space_left);
         closest.append(&mut nodes_from_bucket);
         if closest.capacity() == closest.len() {
            break;
         }
      }
      closest
   }

   /// Returns the appropriate position for a node, by computing
   /// the last bit of their distance set to 1. None if we are
   /// attempting to add the parent key.
   fn bucket_for_node(&self, key : &Hash160) -> Option<usize> {
       (&self.parent_key ^ key).height()
   }

   fn revert_conflict(&mut self, conflict : EvictionConflict) {
      if let Some(index) = self.bucket_for_node(&conflict.inserted.key) {
         let bucket  = &mut self.buckets[index];
         let evictor = &mut bucket.entries.iter_mut().find(|ref info| conflict.inserted.key == info.key).unwrap();
         mem::replace::<NodeInfo>(evictor, conflict.evicted);
      }
   }
}

impl Bucket {
   pub fn new() -> Bucket {
      Bucket { entries : VecDeque::<NodeInfo>::with_capacity(BUCKET_DEPTH) }
   }

}

#[cfg(test)]
mod tests {
    use super::*;
    use hash::Hash160;
    use std::net;
    use std::str::FromStr;

    #[test]
    fn new_routing_table_is_empty() {
       let parent_key = Hash160::random();
       let table = RoutingTable::new(parent_key);
       let test_hash = Hash160::random();

       assert!(table.specific_node(&test_hash).is_none());
    }

    #[test]
    fn inserting_and_retrieving_specific_node() {
       let node_info = NodeInfo {
          key  : Hash160::random(),
          ip   : net::IpAddr::from_str("0.0.0.0").unwrap(),
          port : 50000
       };

       let mut table = RoutingTable::new(Hash160::random());
       table.insert_node(node_info.clone());
       assert_eq!(table.specific_node(&node_info.key), Some(node_info));
    }

    impl RoutingTable {
       pub fn fill_bucket(&mut self, bucket_index : usize, fill_quantity : u8) {
          // Otherwise this helper function becomes quite complex.
          assert!(bucket_index > 7);
          let any_ip = net::IpAddr::from_str("0.0.0.0").unwrap();
          for i in 0..fill_quantity {
             let mut key = self.parent_key.clone();
             key.flip_bit(bucket_index);

             key.raw[0] = i as u8;
             let info = NodeInfo {
                key  : key,
                ip   : any_ip.clone(),
                port : 0
             };
             self.insert_node(info);
          }
       }
    }

    #[test]
    fn inserting_in_a_full_bucket_causes_eviction_conflict() {
       let mut parent_key = Hash160::blank();
       parent_key.raw[1] = 1; // This will guarantee all nodes will fall on the same bucket.
       let any_ip = net::IpAddr::from_str("0.0.0.0").unwrap();

       let mut table = RoutingTable::new(parent_key);

       table.fill_bucket(8, super::BUCKET_DEPTH as u8);
       assert!(table.conflicts.is_empty());

       // When we add another node to the same bucket, we cause a conflict.
       let mut key = Hash160::blank();
       key.raw[0] = 0xFF;
       let info = NodeInfo {
          key  : key,
          ip   : any_ip.clone(),
          port : 0
       };
       table.insert_node(info);
       assert_eq!(table.conflicts.len(), 1);

       let conflict = table.conflicts.first().unwrap();
       // We evicted the oldest node, which has a blank key.
       assert!(table.specific_node(&conflict.inserted.key).is_some());
    }

    #[test]
    fn reverting_an_eviction_conflict_reinserts_the_evicted_node_in_place_of_evictor() {
       let mut parent_key = Hash160::blank();
       parent_key.raw[1] = 1; // This will guarantee all nodes will fall on the same bucket.
       let any_ip = net::IpAddr::from_str("0.0.0.0").unwrap();

       let mut table = RoutingTable::new(parent_key);

       table.fill_bucket(8, super::BUCKET_DEPTH as u8);
       assert!(table.conflicts.is_empty());

       // When we add another node to the same bucket, we cause a conflict.
       let mut key = Hash160::blank();
       key.raw[0] = 0xFF;
       let info = NodeInfo {
          key  : key,
          ip   : any_ip.clone(),
          port : 0
       };
       table.insert_node(info);
       assert_eq!(table.conflicts.len(), 1);
       let conflict = table.conflicts.pop().unwrap();

       table.revert_conflict(conflict.clone());
       // The evictor has been removed.
       assert!(table.specific_node(&conflict.inserted.key).is_none());
       // And the evicted has been reinserted.
       assert!(table.specific_node(&conflict.evicted.key).is_some());
    }

    #[test]
    fn lookup_for_a_stored_node() {
       let mut table = RoutingTable::new(Hash160::random());
       let node = NodeInfo {
          key  : Hash160::random(),
          ip   : net::IpAddr::from_str("0.0.0.0").unwrap(),
          port : 0,
       };
       table.insert_node(node.clone());

       assert_eq!(table.lookup(&node.key, 20), LookupResult::Found(node));
    }

    #[test]
    fn lookup_for_a_not_stored_node() {
       let mut table = RoutingTable::new(Hash160::random());
       table.fill_bucket(11, 5);
       table.fill_bucket(12, 5);
       table.fill_bucket(13, 5);

       // We construct a node key that would fall in bucket 12, but isn't stored.
       let mut node_key = table.parent_key.clone();
       node_key.flip_bit(12);
       node_key.raw[0] = 0xFF;

       if let LookupResult::ClosestNodes(nodes) = table.lookup(&node_key,20) {
          assert_eq!(nodes.len(), 15);

          // Ensure they are ordered by ascending distance
          for (current_node, next_node) in nodes.iter().zip(nodes.iter().skip(1)) {
             assert!((&current_node.key ^ &node_key) <= (&next_node.key ^ &node_key));
          }
       }
       else {
          panic!("We shouldn't have found the node!");
       }
    }

    #[test]
    fn ascending_lookup_on_a_sparse_table() {
       let mut table = RoutingTable::new(Hash160::random());
       for i in (10..50).filter(|x| x%2 == 0) {
         table.fill_bucket(i, 2);
       }
       let mut node_key = table.parent_key.clone();
       node_key.flip_bit(8); // Bucket 8
       node_key.raw[0] = 0xFF;
       if let LookupResult::ClosestNodes(nodes) = table.lookup(&node_key,5) {
          assert_eq!(nodes.len(), 5);

          // Ensure they are ordered by ascending distance
          for (current_node, next_node) in nodes.iter().zip(nodes.iter().skip(1)) {
             assert!((&current_node.key ^ &node_key) <= (&next_node.key ^ &node_key));
          }
       }
       else {
          panic!("We shouldn't have found the node!");
       }
    }

    #[test]
    fn descending_lookup_on_a_sparse_table() {
       let mut table = RoutingTable::new(Hash160::random());
       for i in (10..50).filter(|x| x%2 == 0) {
         table.fill_bucket(i, 2);
       }
       let mut node_key = table.parent_key.clone();
       node_key.flip_bit(51); // Bucket 51
       node_key.raw[0] = 0xFF;
       if let LookupResult::ClosestNodes(nodes) = table.lookup(&node_key,5) {
          assert_eq!(nodes.len(), 5);

          // Ensure they are ordered by ascending distance
          for (current_node, next_node) in nodes.iter().zip(nodes.iter().skip(1)) {
             assert!((&current_node.key ^ &node_key) <= (&next_node.key ^ &node_key));
          }
       }
       else {
          panic!("We shouldn't have found the node!");
       }
    }

    #[test]
    fn lookup_on_a_sparse_table() {
       let mut table = RoutingTable::new(Hash160::random());
       for i in (10..50).filter(|x| x%2 == 0) {
         table.fill_bucket(i, 2);
       }
       let mut node_key = table.parent_key.clone();
       node_key.flip_bit(25); // Bucket 25
       node_key.raw[0] = 0xFF;
       if let LookupResult::ClosestNodes(nodes) = table.lookup(&node_key,5) {
          assert_eq!(nodes.len(), 5);

          // Ensure they are ordered by ascending distance
          for (current_node, next_node) in nodes.iter().zip(nodes.iter().skip(1)) {
             assert!((&current_node.key ^ &node_key) <= (&next_node.key ^ &node_key));
          }
       }
       else {
          panic!("We shouldn't have found the node!");
       }
    }

}
<commit_msg>Master key algorithm<commit_after>use hash::KEY_SIZE;
use hash::Hash160;
use std::net;
use std::collections::VecDeque;
use std::mem;

const K : usize = 20;
const BUCKET_DEPTH : usize = K;

/// Kademlia routing table, with 160 buckets of `BUCKET_DEPTH` (k) node
/// identifiers each.
///
/// The structure employs least-recently seen eviction. Conflicts generated
/// by evicting a node by inserting a newer one remain tracked, so they can
/// can be resolved later.
pub struct RoutingTable {
   parent_key : Hash160,
   buckets    : Vec<Bucket>,
   conflicts  : Vec<EvictionConflict>,
}

/// Bucket size is estimated to be small enough not to warrant
/// the downsides of using a linked list.
#[derive(Debug, Clone)]
struct Bucket {
   entries  : VecDeque<NodeInfo>,
}

/// Represents a conflict derived from attempting to insert a node in a full
/// bucket.
#[derive(Debug,Clone)]
struct EvictionConflict {
   evicted  : NodeInfo,
   inserted : NodeInfo
}

#[derive(Debug, Clone, PartialEq)]
pub struct NodeInfo {
   pub key  : Hash160,
   pub ip   : net::IpAddr,
   pub port : u16,
}

#[derive(Debug, PartialEq)]
pub enum LookupResult {
   Found(NodeInfo),
   ClosestNodes(Vec<NodeInfo>),
   Myself,
}

impl RoutingTable {
   pub fn new(parent_key : Hash160) -> RoutingTable {
      let mut buckets = Vec::<Bucket>::with_capacity(KEY_SIZE);
      for _ in 0..KEY_SIZE {
         buckets.push(Bucket::new());
      }

      RoutingTable {
         parent_key : parent_key,
         buckets : buckets,
         conflicts  : Vec::new()
      }
   }

   /// Inserts a node in the routing table. Employs least-recently-seen eviction
   /// by kicking out the oldest node in case the bucket is full, and registering
   /// an eviction conflict that can be revised later.
   pub fn insert_node(&mut self, info : NodeInfo) {
      if let Some(index) = self.bucket_for_node(&info.key) {
         let bucket = &mut self.buckets[index];
         bucket.entries.retain(|ref stored_info| info.key != stored_info.key);
         if bucket.entries.len() == BUCKET_DEPTH {
            let conflict = EvictionConflict {
               evicted  : bucket.entries.pop_front().unwrap(),
               inserted : info.clone()
            };
            self.conflicts.push(conflict);
         }
         bucket.entries.push_back(info);
      }
   }

   /// Performs a node lookup on the routing table. The lookup result may
   /// contain the specific node, a list of up to the N closest nodes, or
   /// report that the parent node itself was requested.
   pub fn lookup(&self, key : &Hash160, n : usize) -> LookupResult {
      if *key == self.parent_key {
         return LookupResult::Myself;
      }

      match self.specific_node(key) {
         Some(info) => LookupResult::Found(info),
         None => LookupResult::ClosestNodes(self.closest_n_nodes_to(key, n)),
      }
   }

   /// Returns a table entry for the specific node with a given hash.
   fn specific_node(&self, key : &Hash160) -> Option<NodeInfo> {
      if let Some(index) = self.bucket_for_node(key){
         let bucket = &self.buckets[index];
         return bucket.entries.iter().find(|ref info| *key == info.key).cloned();
      }
      None
   }

   fn closest_n_nodes_to(&self, key : &Hash160, n : usize) -> Vec<NodeInfo>{
      let mut closest = Vec::with_capacity(n);
      let mut master_key = key ^ &self.parent_key;
      let ideal_bucket = master_key.height().unwrap();
      let mut floor = ideal_bucket;
      let mut ceiling = floor + 1;

      while ceiling > 1 {
         println!("Trying buckets {} to {}", floor, ceiling);
         for bucket_index in floor..ceiling {
            print!(" doot");
            let mut nodes_from_bucket = self.buckets[bucket_index].clone().entries.into_iter().collect::<Vec<NodeInfo>>();
            nodes_from_bucket.sort_by_key(|ref info| &info.key ^ key);
            let space_left = closest.capacity() - closest.len();
            nodes_from_bucket.truncate(space_left);
            closest.append(&mut nodes_from_bucket);
            if closest.capacity() == closest.len() {
               break;
            }
         }
         ceiling = floor;
         master_key.flip_bit(floor);
         floor = match master_key.height() {
            Some(height) => height,
            None => 1
         }
      }

      // If the master key algorithm wasn't enough, just pour from the higher buckets
      for bucket_index in (ideal_bucket+1)..KEY_SIZE {
         if closest.capacity() == closest.len() {
            break;
         }
         let mut nodes_from_bucket = self.buckets[bucket_index].clone().entries.into_iter().collect::<Vec<NodeInfo>>();
         nodes_from_bucket.sort_by_key(|ref info| &info.key ^ key);
         let space_left = closest.capacity() - closest.len();
         nodes_from_bucket.truncate(space_left);
         closest.append(&mut nodes_from_bucket);
      }
      closest
   }

   /// Returns the appropriate position for a node, by computing
   /// the last bit of their distance set to 1. None if we are
   /// attempting to add the parent key.
   fn bucket_for_node(&self, key : &Hash160) -> Option<usize> {
       (&self.parent_key ^ key).height()
   }

   fn revert_conflict(&mut self, conflict : EvictionConflict) {
      if let Some(index) = self.bucket_for_node(&conflict.inserted.key) {
         let bucket  = &mut self.buckets[index];
         let evictor = &mut bucket.entries.iter_mut().find(|ref info| conflict.inserted.key == info.key).unwrap();
         mem::replace::<NodeInfo>(evictor, conflict.evicted);
      }
   }
}

impl Bucket {
   pub fn new() -> Bucket {
      Bucket { entries : VecDeque::<NodeInfo>::with_capacity(BUCKET_DEPTH) }
   }

}

#[cfg(test)]
mod tests {
    use super::*;
    use hash::Hash160;
    use std::net;
    use std::str::FromStr;

    #[test]
    fn new_routing_table_is_empty() {
       let parent_key = Hash160::random();
       let table = RoutingTable::new(parent_key);
       let test_hash = Hash160::random();

       assert!(table.specific_node(&test_hash).is_none());
    }

    #[test]
    fn inserting_and_retrieving_specific_node() {
       let node_info = NodeInfo {
          key  : Hash160::random(),
          ip   : net::IpAddr::from_str("0.0.0.0").unwrap(),
          port : 50000
       };

       let mut table = RoutingTable::new(Hash160::random());
       table.insert_node(node_info.clone());
       assert_eq!(table.specific_node(&node_info.key), Some(node_info));
    }

    impl RoutingTable {
       pub fn fill_bucket(&mut self, bucket_index : usize, fill_quantity : u8) {
          // Otherwise this helper function becomes quite complex.
          assert!(bucket_index > 7);
          let any_ip = net::IpAddr::from_str("0.0.0.0").unwrap();
          for i in 0..fill_quantity {
             let mut key = self.parent_key.clone();
             key.flip_bit(bucket_index);

             key.raw[0] = i as u8;
             let info = NodeInfo {
                key  : key,
                ip   : any_ip.clone(),
                port : 0
             };
             self.insert_node(info);
          }
       }
    }

    #[test]
    fn inserting_in_a_full_bucket_causes_eviction_conflict() {
       let mut parent_key = Hash160::blank();
       parent_key.raw[1] = 1; // This will guarantee all nodes will fall on the same bucket.
       let any_ip = net::IpAddr::from_str("0.0.0.0").unwrap();

       let mut table = RoutingTable::new(parent_key);

       table.fill_bucket(8, super::BUCKET_DEPTH as u8);
       assert!(table.conflicts.is_empty());

       // When we add another node to the same bucket, we cause a conflict.
       let mut key = Hash160::blank();
       key.raw[0] = 0xFF;
       let info = NodeInfo {
          key  : key,
          ip   : any_ip.clone(),
          port : 0
       };
       table.insert_node(info);
       assert_eq!(table.conflicts.len(), 1);

       let conflict = table.conflicts.first().unwrap();
       // We evicted the oldest node, which has a blank key.
       assert!(table.specific_node(&conflict.inserted.key).is_some());
    }

    #[test]
    fn reverting_an_eviction_conflict_reinserts_the_evicted_node_in_place_of_evictor() {
       let mut parent_key = Hash160::blank();
       parent_key.raw[1] = 1; // This will guarantee all nodes will fall on the same bucket.
       let any_ip = net::IpAddr::from_str("0.0.0.0").unwrap();

       let mut table = RoutingTable::new(parent_key);

       table.fill_bucket(8, super::BUCKET_DEPTH as u8);
       assert!(table.conflicts.is_empty());

       // When we add another node to the same bucket, we cause a conflict.
       let mut key = Hash160::blank();
       key.raw[0] = 0xFF;
       let info = NodeInfo {
          key  : key,
          ip   : any_ip.clone(),
          port : 0
       };
       table.insert_node(info);
       assert_eq!(table.conflicts.len(), 1);
       let conflict = table.conflicts.pop().unwrap();

       table.revert_conflict(conflict.clone());
       // The evictor has been removed.
       assert!(table.specific_node(&conflict.inserted.key).is_none());
       // And the evicted has been reinserted.
       assert!(table.specific_node(&conflict.evicted.key).is_some());
    }

    #[test]
    fn lookup_for_a_stored_node() {
       let mut table = RoutingTable::new(Hash160::random());
       let node = NodeInfo {
          key  : Hash160::random(),
          ip   : net::IpAddr::from_str("0.0.0.0").unwrap(),
          port : 0,
       };
       table.insert_node(node.clone());

       assert_eq!(table.lookup(&node.key, 20), LookupResult::Found(node));
    }

    #[test]
    fn lookup_for_a_not_stored_node() {
       let mut table = RoutingTable::new(Hash160::random());
       table.fill_bucket(11, 5);
       table.fill_bucket(12, 5);
       table.fill_bucket(13, 5);

       // We construct a node key that would fall in bucket 12, but isn't stored.
       let mut node_key = table.parent_key.clone();
       node_key.flip_bit(12);
       node_key.raw[0] = 0xFF;

       if let LookupResult::ClosestNodes(nodes) = table.lookup(&node_key,20) {
          assert_eq!(nodes.len(), 15);

          // Ensure they are ordered by ascending distance
          for (current_node, next_node) in nodes.iter().zip(nodes.iter().skip(1)) {
             assert!((&current_node.key ^ &node_key) <= (&next_node.key ^ &node_key));
          }
       }
       else {
          panic!("We shouldn't have found the node!");
       }
    }

    #[test]
    fn ascending_lookup_on_a_sparse_table() {
       let mut table = RoutingTable::new(Hash160::random());
       for i in (10..50).filter(|x| x%2 == 0) {
         table.fill_bucket(i, 2);
       }
       let mut node_key = table.parent_key.clone();
       node_key.flip_bit(8); // Bucket 8
       if let LookupResult::ClosestNodes(nodes) = table.lookup(&node_key,5) {
          assert_eq!(nodes.len(), 5);

          // Ensure they are ordered by ascending distance
          for (current_node, next_node) in nodes.iter().zip(nodes.iter().skip(1)) {
             assert!((&current_node.key ^ &node_key) <= (&next_node.key ^ &node_key));
          }
       }
       else {
          panic!("We shouldn't have found the node!");
       }
    }

    #[test]
    fn descending_lookup_on_a_sparse_table() {
       let mut table = RoutingTable::new(Hash160::random());
       for i in (10..50).filter(|x| x%2 == 0) {
         table.fill_bucket(i, 2);
       }
       let mut node_key = table.parent_key.clone();
       node_key.flip_bit(51); // Bucket 51
       node_key.raw[0] = 0xFF;
       if let LookupResult::ClosestNodes(nodes) = table.lookup(&node_key,5) {
          assert_eq!(nodes.len(), 5);

          // Ensure they are ordered by ascending distance
          for (current_node, next_node) in nodes.iter().zip(nodes.iter().skip(1)) {
             assert!((&current_node.key ^ &node_key) <= (&next_node.key ^ &node_key));
          }
       }
       else {
          panic!("We shouldn't have found the node!");
       }
    }

    #[test]
    fn lookup_on_a_sparse_table() {
       let mut table = RoutingTable::new(Hash160::random());
       for i in (10..50).filter(|x| x%2 == 0) {
         table.fill_bucket(i, 2);
       }
       let mut node_key = table.parent_key.clone();
       node_key.flip_bit(25); // Bucket 25
       node_key.raw[0] = 0xFF;
       if let LookupResult::ClosestNodes(nodes) = table.lookup(&node_key,5) {
          assert_eq!(nodes.len(), 5);

          // Ensure they are ordered by ascending distance
          for (current_node, next_node) in nodes.iter().zip(nodes.iter().skip(1)) {
             assert!((&current_node.key ^ &node_key) <= (&next_node.key ^ &node_key));
          }
       }
       else {
          panic!("We shouldn't have found the node!");
       }
    }

}
<|endoftext|>


<commit_before>use glob;

use fs::DotFilter;
use fs::filter::{FileFilter, SortField, SortCase, IgnorePatterns};

use options::{flags, Misfire};
use options::parser::Matches;

impl FileFilter {

    /// Determines the set of file filter options to use, based on the user’s
    /// command-line arguments.
    pub fn deduce(matches: &Matches) -> Result<FileFilter, Misfire> {
        Ok(FileFilter {
            list_dirs_first: matches.has(&flags::DIRS_FIRST),
            reverse:         matches.has(&flags::REVERSE),
            sort_field:      SortField::deduce(matches)?,
            dot_filter:      DotFilter::deduce(matches)?,
            ignore_patterns: IgnorePatterns::deduce(matches)?,
        })
    }
}

impl Default for SortField {
    fn default() -> SortField {
        SortField::Name(SortCase::Sensitive)
    }
}

const SORTS: &[&str] = &[ "name", "Name", "size", "extension",
                          "Extension", "modified", "accessed",
                          "created", "inode", "type", "none" ];

impl SortField {

    /// Determine the sort field to use, based on the presence of a “sort”
    /// argument. This will return `Err` if the option is there, but does not
    /// correspond to a valid field.
    fn deduce(matches: &Matches) -> Result<SortField, Misfire> {
        let word = match matches.get(&flags::SORT) {
            Some(w)  => w,
            None     => return Ok(SortField::default()),
        };

        if word == "name" || word == "filename" {
            Ok(SortField::Name(SortCase::Sensitive))
        }
        else if word == "Name" || word == "Filename" {
            Ok(SortField::Name(SortCase::Insensitive))
        }
        else if word == "size" || word == "filesize" {
            Ok(SortField::Size)
        }
        else if word == "ext" || word == "extension" {
            Ok(SortField::Extension(SortCase::Sensitive))
        }
        else if word == "Ext" || word == "Extension" {
            Ok(SortField::Extension(SortCase::Insensitive))
        }
        else if word == "mod" || word == "modified" {
            Ok(SortField::ModifiedDate)
        }
        else if word == "acc" || word == "accessed" {
            Ok(SortField::AccessedDate)
        }
        else if word == "cr" || word == "created" {
            Ok(SortField::CreatedDate)
        }
        else if word == "inode" {
            Ok(SortField::FileInode)
        }
        else if word == "type" {
            Ok(SortField::FileType)
        }
        else if word == "none" {
            Ok(SortField::Unsorted)
        }
        else {
            Err(Misfire::bad_argument(&flags::SORT, word, SORTS))
        }
    }
}

impl DotFilter {
    pub fn deduce(matches: &Matches) -> Result<DotFilter, Misfire> {
        let dots = match matches.count(&flags::ALL) {
            0 => return Ok(DotFilter::JustFiles),
            1 => DotFilter::Dotfiles,
            _ => DotFilter::DotfilesAndDots,
        };

        if matches.has(&flags::TREE) {
            Err(Misfire::TreeAllAll)
        }
        else {
            Ok(dots)
        }
    }
}

impl IgnorePatterns {

    /// Determines the set of file filter options to use, based on the user’s
    /// command-line arguments.
    pub fn deduce(matches: &Matches) -> Result<IgnorePatterns, Misfire> {
        let patterns = match matches.get(&flags::IGNORE_GLOB) {
            None => Ok(Vec::new()),
            Some(is) => is.to_string_lossy().split('|').map(|a| glob::Pattern::new(a)).collect(),
        }?;

        // TODO: is to_string_lossy really the best way to handle
        // invalid UTF-8 there?

        Ok(IgnorePatterns { patterns })
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::ffi::OsString;
    use options::flags;

    pub fn os(input: &'static str) -> OsString {
        let mut os = OsString::new();
        os.push(input);
        os
    }

    macro_rules! test {
        ($name:ident: $type:ident <- $inputs:expr => $result:expr) => {
            #[test]
            fn $name() {
                use options::parser::{parse, Args, Arg};
                use std::ffi::OsString;

                static TEST_ARGS: &[&Arg] = &[ &flags::SORT ];

                let bits = $inputs.as_ref().into_iter().map(|&o| os(o)).collect::<Vec<OsString>>();
                let results = parse(&Args(TEST_ARGS), bits.iter());
                assert_eq!($type::deduce(results.as_ref().unwrap()), $result);
            }
        };
    }

    mod sort_fields {
        use super::*;

        // Default behaviour
        test!(empty:         SortField <- []                  => Ok(SortField::default()));

        // Sort field arguments
        test!(one_arg:       SortField <- ["--sort=cr"]       => Ok(SortField::CreatedDate));
        test!(one_long:      SortField <- ["--sort=size"]     => Ok(SortField::Size));
        test!(one_short:     SortField <- ["-saccessed"]      => Ok(SortField::AccessedDate));
        test!(lowercase:     SortField <- ["--sort", "name"]  => Ok(SortField::Name(SortCase::Sensitive)));
        test!(uppercase:     SortField <- ["--sort", "Name"]  => Ok(SortField::Name(SortCase::Insensitive)));

        // Errors
        test!(error:         SortField <- ["--sort=colour"]   => Err(Misfire::bad_argument(&flags::SORT, &os("colour"), super::SORTS)));

        // Overriding
        test!(overridden:    SortField <- ["--sort=cr",       "--sort", "mod"]     => Ok(SortField::ModifiedDate));
        test!(overridden_2:  SortField <- ["--sort", "none",  "--sort=Extension"]  => Ok(SortField::Extension(SortCase::Insensitive)));
    }
}
<commit_msg>Add tests for dot filters<commit_after>use glob;

use fs::DotFilter;
use fs::filter::{FileFilter, SortField, SortCase, IgnorePatterns};

use options::{flags, Misfire};
use options::parser::Matches;

impl FileFilter {

    /// Determines the set of file filter options to use, based on the user’s
    /// command-line arguments.
    pub fn deduce(matches: &Matches) -> Result<FileFilter, Misfire> {
        Ok(FileFilter {
            list_dirs_first: matches.has(&flags::DIRS_FIRST),
            reverse:         matches.has(&flags::REVERSE),
            sort_field:      SortField::deduce(matches)?,
            dot_filter:      DotFilter::deduce(matches)?,
            ignore_patterns: IgnorePatterns::deduce(matches)?,
        })
    }
}

impl Default for SortField {
    fn default() -> SortField {
        SortField::Name(SortCase::Sensitive)
    }
}

const SORTS: &[&str] = &[ "name", "Name", "size", "extension",
                          "Extension", "modified", "accessed",
                          "created", "inode", "type", "none" ];

impl SortField {

    /// Determine the sort field to use, based on the presence of a “sort”
    /// argument. This will return `Err` if the option is there, but does not
    /// correspond to a valid field.
    fn deduce(matches: &Matches) -> Result<SortField, Misfire> {
        let word = match matches.get(&flags::SORT) {
            Some(w)  => w,
            None     => return Ok(SortField::default()),
        };

        if word == "name" || word == "filename" {
            Ok(SortField::Name(SortCase::Sensitive))
        }
        else if word == "Name" || word == "Filename" {
            Ok(SortField::Name(SortCase::Insensitive))
        }
        else if word == "size" || word == "filesize" {
            Ok(SortField::Size)
        }
        else if word == "ext" || word == "extension" {
            Ok(SortField::Extension(SortCase::Sensitive))
        }
        else if word == "Ext" || word == "Extension" {
            Ok(SortField::Extension(SortCase::Insensitive))
        }
        else if word == "mod" || word == "modified" {
            Ok(SortField::ModifiedDate)
        }
        else if word == "acc" || word == "accessed" {
            Ok(SortField::AccessedDate)
        }
        else if word == "cr" || word == "created" {
            Ok(SortField::CreatedDate)
        }
        else if word == "inode" {
            Ok(SortField::FileInode)
        }
        else if word == "type" {
            Ok(SortField::FileType)
        }
        else if word == "none" {
            Ok(SortField::Unsorted)
        }
        else {
            Err(Misfire::bad_argument(&flags::SORT, word, SORTS))
        }
    }
}

impl DotFilter {
    pub fn deduce(matches: &Matches) -> Result<DotFilter, Misfire> {
        let dots = match matches.count(&flags::ALL) {
            0 => return Ok(DotFilter::JustFiles),
            1 => DotFilter::Dotfiles,
            _ => DotFilter::DotfilesAndDots,
        };

        if matches.has(&flags::TREE) {
            Err(Misfire::TreeAllAll)
        }
        else {
            Ok(dots)
        }
    }
}

impl IgnorePatterns {

    /// Determines the set of file filter options to use, based on the user’s
    /// command-line arguments.
    pub fn deduce(matches: &Matches) -> Result<IgnorePatterns, Misfire> {
        let patterns = match matches.get(&flags::IGNORE_GLOB) {
            None => Ok(Vec::new()),
            Some(is) => is.to_string_lossy().split('|').map(|a| glob::Pattern::new(a)).collect(),
        }?;

        // TODO: is to_string_lossy really the best way to handle
        // invalid UTF-8 there?

        Ok(IgnorePatterns { patterns })
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::ffi::OsString;
    use options::flags;

    pub fn os(input: &'static str) -> OsString {
        let mut os = OsString::new();
        os.push(input);
        os
    }

    macro_rules! test {
        ($name:ident: $type:ident <- $inputs:expr => $result:expr) => {
            #[test]
            fn $name() {
                use options::parser::{parse, Args, Arg};
                use std::ffi::OsString;

                static TEST_ARGS: &[&Arg] = &[ &flags::SORT, &flags::ALL, &flags::TREE ];

                let bits = $inputs.as_ref().into_iter().map(|&o| os(o)).collect::<Vec<OsString>>();
                let results = parse(&Args(TEST_ARGS), bits.iter());
                assert_eq!($type::deduce(results.as_ref().unwrap()), $result);
            }
        };
    }

    mod sort_fields {
        use super::*;

        // Default behaviour
        test!(empty:         SortField <- []                  => Ok(SortField::default()));

        // Sort field arguments
        test!(one_arg:       SortField <- ["--sort=cr"]       => Ok(SortField::CreatedDate));
        test!(one_long:      SortField <- ["--sort=size"]     => Ok(SortField::Size));
        test!(one_short:     SortField <- ["-saccessed"]      => Ok(SortField::AccessedDate));
        test!(lowercase:     SortField <- ["--sort", "name"]  => Ok(SortField::Name(SortCase::Sensitive)));
        test!(uppercase:     SortField <- ["--sort", "Name"]  => Ok(SortField::Name(SortCase::Insensitive)));

        // Errors
        test!(error:         SortField <- ["--sort=colour"]   => Err(Misfire::bad_argument(&flags::SORT, &os("colour"), super::SORTS)));

        // Overriding
        test!(overridden:    SortField <- ["--sort=cr",       "--sort", "mod"]     => Ok(SortField::ModifiedDate));
        test!(overridden_2:  SortField <- ["--sort", "none",  "--sort=Extension"]  => Ok(SortField::Extension(SortCase::Insensitive)));
    }

    mod dot_filters {
        use super::*;

        // Default behaviour
        test!(empty:      DotFilter <- []               => Ok(DotFilter::JustFiles));

        // --all
        test!(all:        DotFilter <- ["--all"]        => Ok(DotFilter::Dotfiles));
        test!(all_all:    DotFilter <- ["--all", "-a"]  => Ok(DotFilter::DotfilesAndDots));
        test!(all_all_2:  DotFilter <- ["-aa"]          => Ok(DotFilter::DotfilesAndDots));

        // --all and --tree
        test!(tree:       DotFilter <- ["-Taa"]         => Err(Misfire::TreeAllAll));
    }
}
<|endoftext|>


<commit_before>// 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.

//! # Node Builder
//!
//! ```
//! use commands::parser::builder::*;
//!
//! let mut tree = CommandTree::new();
//! tree.command(Command::new("again")
//!                  .hidden(false)
//!                  .parameter(Parameter::new("test")
//!                                 .required(false)
//!                                 .help("This is just a test parameter.")
//!                                 .finalize())
//!                  .finalize());
//! ```

use std::rc::Rc;
use super::nodes::*;

/// Indicate the type of parameter, so that the correct class and node
/// structures are created.
#[derive(Clone, Copy, PartialEq)]
pub enum ParameterKind {
    /// This parameter is a `FlagParameter`.
    Flag,
    /// This parameter is a `NamedParameter`.
    Named,
    /// This parameter is a `SimpleParameter`.
    Simple,
}

/// Store a command tree while populating it. This can be used
/// to construct a `RootNode` to be used with the `Parser`.
pub struct CommandTree {
    commands: Vec<Command>,
}

impl CommandTree {
    /// Create a new `CommandTree`.
    pub fn new() -> Self {
        CommandTree { commands: vec![] }
    }

    /// Add a `Command` to the `CommandTree`.
    pub fn command(&mut self, command: Command) {
        self.commands.push(command);
    }

    /// Construct the `CommandTree` and produce a `RootNode`.
    pub fn finalize(&self) -> Rc<RootNode> {
        let mut successors: Vec<Rc<Node>> = vec![];
        for c in &self.commands {
            successors.push(self.build_command(c));
        }
        RootNode::new(successors)
    }

    fn build_command(&self, command: &Command) -> Rc<Node> {
        let mut parameters: Vec<Rc<ParameterNode>> = vec![];
        let mut successors: Vec<Rc<Node>> = vec![];
        for parameter in &command.parameters {
            match parameter.parameter_kind {
                ParameterKind::Flag => {
                    self.build_flag_parameter(parameter, &mut parameters, &mut successors);
                }
                ParameterKind::Named => {
                    self.build_named_parameter(parameter, &mut parameters, &mut successors);
                }
                ParameterKind::Simple => {
                    self.build_simple_parameter(parameter, &mut parameters, &mut successors);
                }
            };
        }
        let c = CommandNode::new(&*command.name,
                                 command.help_text.clone(),
                                 command.hidden,
                                 command.priority,
                                 successors,
                                 None,
                                 parameters);
        Rc::new(c)
    }

    fn build_flag_parameter(&self,
                            parameter: &Parameter,
                            parameters: &mut Vec<Rc<ParameterNode>>,
                            successors: &mut Vec<Rc<Node>>) {
        let p = FlagParameterNode::new(&*parameter.name,
                                       parameter.help_text.clone(),
                                       parameter.hidden,
                                       parameter.priority.unwrap_or(PRIORITY_DEFAULT),
                                       vec![],
                                       parameter.repeatable,
                                       None,
                                       parameter.required);
        let fp = Rc::new(p);
        parameters.push(fp.clone());
        successors.push(fp);
    }

    fn build_named_parameter(&self,
                             parameter: &Parameter,
                             parameters: &mut Vec<Rc<ParameterNode>>,
                             successors: &mut Vec<Rc<Node>>) {
        let p = Rc::new(NamedParameterNode::new(&*parameter.name,
                                                parameter.help_text.clone(),
                                                parameter.hidden,
                                                parameter.priority.unwrap_or(PRIORITY_PARAMETER),
                                                vec![],
                                                parameter.repeatable,
                                                None,
                                                parameter.required));
        parameters.push(p.clone());
        let n = Rc::new(ParameterNameNode::new(&*parameter.name,
                                               parameter.hidden,
                                               PRIORITY_DEFAULT,
                                               vec![p.clone()],
                                               parameter.repeatable,
                                               Some(p.clone()),
                                               p.clone()));
        successors.push(n);
        for alias in &parameter.aliases {
            let a = Rc::new(ParameterNameNode::new(&*alias,
                                                   parameter.hidden,
                                                   PRIORITY_DEFAULT,
                                                   vec![p.clone()],
                                                   parameter.repeatable,
                                                   Some(p.clone()),
                                                   p.clone()));
            successors.push(a);
        }
    }

    fn build_simple_parameter(&self,
                              parameter: &Parameter,
                              parameters: &mut Vec<Rc<ParameterNode>>,
                              successors: &mut Vec<Rc<Node>>) {
        let p = SimpleParameterNode::new(&*parameter.name,
                                         parameter.help_text.clone(),
                                         parameter.hidden,
                                         parameter.priority.unwrap_or(PRIORITY_PARAMETER),
                                         vec![],
                                         parameter.repeatable,
                                         None,
                                         parameter.required);
        let sp = Rc::new(p);
        parameters.push(sp.clone());
        successors.push(sp);
    }
}

/// Description of a command to be added to the `CommandTree`.
#[derive(Clone)]
pub struct Command {
    hidden: bool,
    priority: i32,
    name: String,
    help_text: Option<String>,
    parameters: Vec<Parameter>,
    wrapped_root: Option<String>,
}

impl Command {
    /// Construct a default (blank) command with the given `name`.
    pub fn new(name: &str) -> Self {
        Command {
            hidden: false,
            priority: PRIORITY_DEFAULT,
            name: name.to_string(),
            help_text: None,
            parameters: vec![],
            wrapped_root: None,
        }
    }

    /// Mark the command as hidden. Hidden commands will match
    /// within the parser, but are not listed during completion.
    pub fn hidden(&mut self, hidden: bool) -> &mut Self {
        self.hidden = hidden;
        self
    }

    /// Give the command a priority. This is used when sorting
    /// out conflicts during matching and completion.
    pub fn priority(&mut self, priority: i32) -> &mut Self {
        self.priority = priority;
        self
    }

    /// Supply help text for the command.
    pub fn help(&mut self, help_text: &str) -> &mut Self {
        self.help_text = Some(help_text.to_string());
        self
    }

    /// Add a `Parameter` to the command.
    pub fn parameter(&mut self, parameter: Parameter) -> &mut Self {
        self.parameters.push(parameter);
        self
    }

    /// Create a `WrapperNode` instead of a `CommandNode`. The
    /// `wrapped_root` signifies the path to the command that should
    /// be wrapped by this command.
    pub fn wraps(&mut self, wrapped_root: String) -> &mut Self {
        self.wrapped_root = Some(wrapped_root);
        self
    }

    /// Return an instance of `Command` that can be passed to the
    /// `CommandTree`. This is used to terminate the series of construction
    /// methods used to initialize and configure the command.
    pub fn finalize(&self) -> Self {
        self.clone()
    }
}

/// Description of a parameter to be added to the `Command`.
#[derive(Clone)]
pub struct Parameter {
    hidden: bool,
    priority: Option<i32>,
    name: String,
    repeatable: bool,
    aliases: Vec<String>,
    help_text: Option<String>,
    required: bool,
    parameter_kind: ParameterKind,
}

impl Parameter {
    /// Construct a default (blank) parameter with the given `name`.
    pub fn new(name: &str) -> Self {
        Parameter {
            hidden: false,
            priority: None,
            name: name.to_string(),
            repeatable: false,
            aliases: vec![],
            help_text: None,
            required: false,
            parameter_kind: ParameterKind::Simple,
        }
    }

    /// Mark the parameter as hidden. Hidden parameters will match
    /// within the parser, but are not listed during completion.
    pub fn hidden(&mut self, hidden: bool) -> &mut Self {
        self.hidden = hidden;
        self
    }

    /// Give the parameter a priority. This is used when sorting
    /// out conflicts during matching and completion.
    ///
    /// The `priority` of a `Parameter` defaults to `PRIORITY_PARAMETER`
    /// except for when the `kind` is `ParameterKind::Flag` in which
    /// case, the default will be `PRIORITY_DEFAULT`.
    pub fn priority(&mut self, priority: i32) -> &mut Self {
        self.priority = Some(priority);
        self
    }

    /// Establish whether or not this parameter is repeatable.
    /// Repeated parameters produce a vector of values and can
    /// be given multiple times within a single command invocation.
    pub fn repeatable(&mut self, repeatable: bool) -> &mut Self {
        self.repeatable = repeatable;
        self
    }

    /// Add an alias that this parameter can use.
    pub fn alias(&mut self, alias: &str) -> &mut Self {
        self.aliases.push(alias.to_string());
        self
    }

    /// Supply the help text for the parameter.
    pub fn help(&mut self, help_text: &str) -> &mut Self {
        self.help_text = Some(help_text.to_string());
        self
    }

    /// Establish whether or not this parameter is required.
    pub fn required(&mut self, required: bool) -> &mut Self {
        self.required = required;
        self
    }

    /// Set which type of `ParameterNode` is supposed to be created
    /// to represent this parameter.
    pub fn kind(&mut self, kind: ParameterKind) -> &mut Self {
        self.parameter_kind = kind;
        self
    }

    /// Return an instance of `Parameter` that can be passed to the
    /// `Command`. This is used to terminate the series of construction
    /// methods used to initialize and configure the parameter.
    pub fn finalize(&self) -> Self {
        self.clone()
    }
}
<commit_msg>Note that aliases are only valid for named parameters.<commit_after>// 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.

//! # Node Builder
//!
//! ```
//! use commands::parser::builder::*;
//!
//! let mut tree = CommandTree::new();
//! tree.command(Command::new("again")
//!                  .hidden(false)
//!                  .parameter(Parameter::new("test")
//!                                 .required(false)
//!                                 .help("This is just a test parameter.")
//!                                 .finalize())
//!                  .finalize());
//! ```

use std::rc::Rc;
use super::nodes::*;

/// Indicate the type of parameter, so that the correct class and node
/// structures are created.
#[derive(Clone, Copy, PartialEq)]
pub enum ParameterKind {
    /// This parameter is a `FlagParameter`.
    Flag,
    /// This parameter is a `NamedParameter`.
    Named,
    /// This parameter is a `SimpleParameter`.
    Simple,
}

/// Store a command tree while populating it. This can be used
/// to construct a `RootNode` to be used with the `Parser`.
pub struct CommandTree {
    commands: Vec<Command>,
}

impl CommandTree {
    /// Create a new `CommandTree`.
    pub fn new() -> Self {
        CommandTree { commands: vec![] }
    }

    /// Add a `Command` to the `CommandTree`.
    pub fn command(&mut self, command: Command) {
        self.commands.push(command);
    }

    /// Construct the `CommandTree` and produce a `RootNode`.
    pub fn finalize(&self) -> Rc<RootNode> {
        let mut successors: Vec<Rc<Node>> = vec![];
        for c in &self.commands {
            successors.push(self.build_command(c));
        }
        RootNode::new(successors)
    }

    fn build_command(&self, command: &Command) -> Rc<Node> {
        let mut parameters: Vec<Rc<ParameterNode>> = vec![];
        let mut successors: Vec<Rc<Node>> = vec![];
        for parameter in &command.parameters {
            match parameter.parameter_kind {
                ParameterKind::Flag => {
                    self.build_flag_parameter(parameter, &mut parameters, &mut successors);
                }
                ParameterKind::Named => {
                    self.build_named_parameter(parameter, &mut parameters, &mut successors);
                }
                ParameterKind::Simple => {
                    self.build_simple_parameter(parameter, &mut parameters, &mut successors);
                }
            };
        }
        let c = CommandNode::new(&*command.name,
                                 command.help_text.clone(),
                                 command.hidden,
                                 command.priority,
                                 successors,
                                 None,
                                 parameters);
        Rc::new(c)
    }

    fn build_flag_parameter(&self,
                            parameter: &Parameter,
                            parameters: &mut Vec<Rc<ParameterNode>>,
                            successors: &mut Vec<Rc<Node>>) {
        let p = FlagParameterNode::new(&*parameter.name,
                                       parameter.help_text.clone(),
                                       parameter.hidden,
                                       parameter.priority.unwrap_or(PRIORITY_DEFAULT),
                                       vec![],
                                       parameter.repeatable,
                                       None,
                                       parameter.required);
        let fp = Rc::new(p);
        parameters.push(fp.clone());
        successors.push(fp);
    }

    fn build_named_parameter(&self,
                             parameter: &Parameter,
                             parameters: &mut Vec<Rc<ParameterNode>>,
                             successors: &mut Vec<Rc<Node>>) {
        let p = Rc::new(NamedParameterNode::new(&*parameter.name,
                                                parameter.help_text.clone(),
                                                parameter.hidden,
                                                parameter.priority.unwrap_or(PRIORITY_PARAMETER),
                                                vec![],
                                                parameter.repeatable,
                                                None,
                                                parameter.required));
        parameters.push(p.clone());
        let n = Rc::new(ParameterNameNode::new(&*parameter.name,
                                               parameter.hidden,
                                               PRIORITY_DEFAULT,
                                               vec![p.clone()],
                                               parameter.repeatable,
                                               Some(p.clone()),
                                               p.clone()));
        successors.push(n);
        for alias in &parameter.aliases {
            let a = Rc::new(ParameterNameNode::new(&*alias,
                                                   parameter.hidden,
                                                   PRIORITY_DEFAULT,
                                                   vec![p.clone()],
                                                   parameter.repeatable,
                                                   Some(p.clone()),
                                                   p.clone()));
            successors.push(a);
        }
    }

    fn build_simple_parameter(&self,
                              parameter: &Parameter,
                              parameters: &mut Vec<Rc<ParameterNode>>,
                              successors: &mut Vec<Rc<Node>>) {
        let p = SimpleParameterNode::new(&*parameter.name,
                                         parameter.help_text.clone(),
                                         parameter.hidden,
                                         parameter.priority.unwrap_or(PRIORITY_PARAMETER),
                                         vec![],
                                         parameter.repeatable,
                                         None,
                                         parameter.required);
        let sp = Rc::new(p);
        parameters.push(sp.clone());
        successors.push(sp);
    }
}

/// Description of a command to be added to the `CommandTree`.
#[derive(Clone)]
pub struct Command {
    hidden: bool,
    priority: i32,
    name: String,
    help_text: Option<String>,
    parameters: Vec<Parameter>,
    wrapped_root: Option<String>,
}

impl Command {
    /// Construct a default (blank) command with the given `name`.
    pub fn new(name: &str) -> Self {
        Command {
            hidden: false,
            priority: PRIORITY_DEFAULT,
            name: name.to_string(),
            help_text: None,
            parameters: vec![],
            wrapped_root: None,
        }
    }

    /// Mark the command as hidden. Hidden commands will match
    /// within the parser, but are not listed during completion.
    pub fn hidden(&mut self, hidden: bool) -> &mut Self {
        self.hidden = hidden;
        self
    }

    /// Give the command a priority. This is used when sorting
    /// out conflicts during matching and completion.
    pub fn priority(&mut self, priority: i32) -> &mut Self {
        self.priority = priority;
        self
    }

    /// Supply help text for the command.
    pub fn help(&mut self, help_text: &str) -> &mut Self {
        self.help_text = Some(help_text.to_string());
        self
    }

    /// Add a `Parameter` to the command.
    pub fn parameter(&mut self, parameter: Parameter) -> &mut Self {
        self.parameters.push(parameter);
        self
    }

    /// Create a `WrapperNode` instead of a `CommandNode`. The
    /// `wrapped_root` signifies the path to the command that should
    /// be wrapped by this command.
    pub fn wraps(&mut self, wrapped_root: String) -> &mut Self {
        self.wrapped_root = Some(wrapped_root);
        self
    }

    /// Return an instance of `Command` that can be passed to the
    /// `CommandTree`. This is used to terminate the series of construction
    /// methods used to initialize and configure the command.
    pub fn finalize(&self) -> Self {
        self.clone()
    }
}

/// Description of a parameter to be added to the `Command`.
#[derive(Clone)]
pub struct Parameter {
    hidden: bool,
    priority: Option<i32>,
    name: String,
    repeatable: bool,
    aliases: Vec<String>,
    help_text: Option<String>,
    required: bool,
    parameter_kind: ParameterKind,
}

impl Parameter {
    /// Construct a default (blank) parameter with the given `name`.
    pub fn new(name: &str) -> Self {
        Parameter {
            hidden: false,
            priority: None,
            name: name.to_string(),
            repeatable: false,
            aliases: vec![],
            help_text: None,
            required: false,
            parameter_kind: ParameterKind::Simple,
        }
    }

    /// Mark the parameter as hidden. Hidden parameters will match
    /// within the parser, but are not listed during completion.
    pub fn hidden(&mut self, hidden: bool) -> &mut Self {
        self.hidden = hidden;
        self
    }

    /// Give the parameter a priority. This is used when sorting
    /// out conflicts during matching and completion.
    ///
    /// The `priority` of a `Parameter` defaults to `PRIORITY_PARAMETER`
    /// except for when the `kind` is `ParameterKind::Flag` in which
    /// case, the default will be `PRIORITY_DEFAULT`.
    pub fn priority(&mut self, priority: i32) -> &mut Self {
        self.priority = Some(priority);
        self
    }

    /// Establish whether or not this parameter is repeatable.
    /// Repeated parameters produce a vector of values and can
    /// be given multiple times within a single command invocation.
    pub fn repeatable(&mut self, repeatable: bool) -> &mut Self {
        self.repeatable = repeatable;
        self
    }

    /// Add an alias that this parameter can use.
    ///
    /// Aliases are currently only valid for parameters of `kind`
    /// `ParameterKind::Named`.
    pub fn alias(&mut self, alias: &str) -> &mut Self {
        self.aliases.push(alias.to_string());
        self
    }

    /// Supply the help text for the parameter.
    pub fn help(&mut self, help_text: &str) -> &mut Self {
        self.help_text = Some(help_text.to_string());
        self
    }

    /// Establish whether or not this parameter is required.
    pub fn required(&mut self, required: bool) -> &mut Self {
        self.required = required;
        self
    }

    /// Set which type of `ParameterNode` is supposed to be created
    /// to represent this parameter.
    pub fn kind(&mut self, kind: ParameterKind) -> &mut Self {
        self.parameter_kind = kind;
        self
    }

    /// Return an instance of `Parameter` that can be passed to the
    /// `Command`. This is used to terminate the series of construction
    /// methods used to initialize and configure the parameter.
    pub fn finalize(&self) -> Self {
        self.clone()
    }
}
<|endoftext|>


<commit_before>#![feature(plugin_registrar, quote, rustc_private, box_patterns)]

extern crate rustc_plugin;
#[macro_use] pub extern crate syntax;
#[macro_use] extern crate syntax_pos;
pub extern crate rustc_errors as errors;

use syntax::ast;
use syntax::codemap;
use syntax::ext::base::{ExtCtxt, MacResult, MacEager, DummyResult};
use syntax::parse;
use syntax::parse::token;
use syntax::fold::Folder;
use syntax::util::small_vector::SmallVector;
use rustc_plugin::Registry;
use std::io::Read;
use std::fs::File;
use std::path::Path;

use rustast::{AstBuilder, DUMMY_SP, TokenTree};

mod translate;
mod grammar;
mod rustast;

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(
            token::intern("peg"),
            syntax::ext::base::IdentTT(Box::new(expand_peg_str), None, false));

    reg.register_syntax_extension(
            token::intern("peg_file"),
            syntax::ext::base::IdentTT(Box::new(expand_peg_file), None, false));
}

fn expand_peg_str<'s>(cx: &'s mut ExtCtxt, sp: codemap::Span, ident: ast::Ident, tts: Vec<TokenTree>) -> Box<MacResult + 's> {
    let source = match parse_arg(cx, &tts) {
        Some(source) => source,
        None => return DummyResult::any(sp),
    };

    expand_peg(cx, sp, ident, &source)
}

fn expand_peg_file<'s>(cx: &'s mut ExtCtxt, sp: codemap::Span, ident: ast::Ident, tts: Vec<TokenTree>) -> Box<MacResult + 's> {
    let fname = match parse_arg(cx, &tts) {
        Some(fname) => fname,
        None => return DummyResult::any(sp),
    };

    let path = Path::new(&cx.codemap().span_to_filename(sp)).parent().unwrap().join(&fname);

    let mut source = String::new();
    if let Err(e) = File::open(&path).map(|mut f| f.read_to_string(&mut source)) {
        cx.span_err(sp, &e.to_string());
        return DummyResult::any(sp);
    }

    cx.codemap().new_filemap(format!("{}", path.display()), None, "".to_string());

    expand_peg(cx, sp, ident, &source)
}

fn expand_peg(cx: &mut ExtCtxt, sp: codemap::Span, ident: ast::Ident, source: &str) -> Box<MacResult + 'static> {
    let grammar_def = grammar::grammar(source);

    let grammar_def = match grammar_def {
      Ok(grammar_def) => grammar_def,
      Err(msg) => {
        cx.span_err(sp, &format!("{}", msg));
        return DummyResult::any(sp)
      }
    };

    let ast = translate::compile_grammar(cx, &grammar_def);

    // #![allow(non_snake_case, unused)]
    let allow = cx.attribute(DUMMY_SP, cx.meta_list(DUMMY_SP, token::InternedString::new("allow"), vec![
        cx.meta_list_item_word(DUMMY_SP, token::InternedString::new("non_snake_case")),
        cx.meta_list_item_word(DUMMY_SP, token::InternedString::new("unused")),
    ]));

    MacEager::items(SmallVector::one(cx.item_mod(sp, sp, ident, vec![allow], ast.items.clone())))
}

fn parse_arg(cx: &mut ExtCtxt, tts: &[TokenTree]) -> Option<String> {
    use syntax::print::pprust;

    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg(),
                                                tts.to_vec());
    // The `expand_expr` method is called so that any macro calls in the
    // parsed expression are expanded.
    let arg = cx.expander().fold_expr(panictry!(parser.parse_expr()));
    match arg.node {
        ast::ExprKind::Lit(ref spanned) => {
            match spanned.node {
                ast::LitKind::Str(ref n, _) => {
                    if !parser.eat(&token::Eof) {
                        cx.span_err(parser.span,
                                    "expected only one string literal");
                        return None
                    }
                    return Some(n.to_string())
                }
                _ => {}
            }
        }
        _ => {}
    }

    let err = format!("expected string literal but got `{}`",
                      pprust::expr_to_string(&*arg));
    cx.span_err(parser.span, &err);
    None
}
<commit_msg>rust: Fix for CrateConfig change<commit_after>#![feature(plugin_registrar, quote, rustc_private, box_patterns)]

extern crate rustc_plugin;
#[macro_use] pub extern crate syntax;
#[macro_use] extern crate syntax_pos;
pub extern crate rustc_errors as errors;

use syntax::ast;
use syntax::codemap;
use syntax::ext::base::{ExtCtxt, MacResult, MacEager, DummyResult};
use syntax::parse;
use syntax::parse::token;
use syntax::fold::Folder;
use syntax::util::small_vector::SmallVector;
use rustc_plugin::Registry;
use std::io::Read;
use std::fs::File;
use std::path::Path;

use rustast::{AstBuilder, DUMMY_SP, TokenTree};

mod translate;
mod grammar;
mod rustast;

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(
            token::intern("peg"),
            syntax::ext::base::IdentTT(Box::new(expand_peg_str), None, false));

    reg.register_syntax_extension(
            token::intern("peg_file"),
            syntax::ext::base::IdentTT(Box::new(expand_peg_file), None, false));
}

fn expand_peg_str<'s>(cx: &'s mut ExtCtxt, sp: codemap::Span, ident: ast::Ident, tts: Vec<TokenTree>) -> Box<MacResult + 's> {
    let source = match parse_arg(cx, &tts) {
        Some(source) => source,
        None => return DummyResult::any(sp),
    };

    expand_peg(cx, sp, ident, &source)
}

fn expand_peg_file<'s>(cx: &'s mut ExtCtxt, sp: codemap::Span, ident: ast::Ident, tts: Vec<TokenTree>) -> Box<MacResult + 's> {
    let fname = match parse_arg(cx, &tts) {
        Some(fname) => fname,
        None => return DummyResult::any(sp),
    };

    let path = Path::new(&cx.codemap().span_to_filename(sp)).parent().unwrap().join(&fname);

    let mut source = String::new();
    if let Err(e) = File::open(&path).map(|mut f| f.read_to_string(&mut source)) {
        cx.span_err(sp, &e.to_string());
        return DummyResult::any(sp);
    }

    cx.codemap().new_filemap(format!("{}", path.display()), None, "".to_string());

    expand_peg(cx, sp, ident, &source)
}

fn expand_peg(cx: &mut ExtCtxt, sp: codemap::Span, ident: ast::Ident, source: &str) -> Box<MacResult + 'static> {
    let grammar_def = grammar::grammar(source);

    let grammar_def = match grammar_def {
      Ok(grammar_def) => grammar_def,
      Err(msg) => {
        cx.span_err(sp, &format!("{}", msg));
        return DummyResult::any(sp)
      }
    };

    let ast = translate::compile_grammar(cx, &grammar_def);

    // #![allow(non_snake_case, unused)]
    let allow = cx.attribute(DUMMY_SP, cx.meta_list(DUMMY_SP, token::InternedString::new("allow"), vec![
        cx.meta_list_item_word(DUMMY_SP, token::InternedString::new("non_snake_case")),
        cx.meta_list_item_word(DUMMY_SP, token::InternedString::new("unused")),
    ]));

    MacEager::items(SmallVector::one(cx.item_mod(sp, sp, ident, vec![allow], ast.items.clone())))
}

fn parse_arg(cx: &mut ExtCtxt, tts: &[TokenTree]) -> Option<String> {
    use syntax::print::pprust;

    let mut parser = parse::new_parser_from_tts(cx.parse_sess(), cx.cfg().clone(),
                                                tts.to_vec());
    // The `expand_expr` method is called so that any macro calls in the
    // parsed expression are expanded.
    let arg = cx.expander().fold_expr(panictry!(parser.parse_expr()));
    match arg.node {
        ast::ExprKind::Lit(ref spanned) => {
            match spanned.node {
                ast::LitKind::Str(ref n, _) => {
                    if !parser.eat(&token::Eof) {
                        cx.span_err(parser.span,
                                    "expected only one string literal");
                        return None
                    }
                    return Some(n.to_string())
                }
                _ => {}
            }
        }
        _ => {}
    }

    let err = format!("expected string literal but got `{}`",
                      pprust::expr_to_string(&*arg));
    cx.span_err(parser.span, &err);
    None
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Anatoly Ikorsky
//
// 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. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

use futures_util::future::Either;

use crate::{
    connection_like::{
        streamless::Streamless, ConnectionLike, ConnectionLikeWrapper, StmtCacheResult,
    },
    error::*,
    io,
    prelude::FromRow,
    queryable::{query_result::QueryResult, BinaryProtocol},
    Column, Params, Row,
    Value::{self},
};
use mysql_common::{
    constants::MAX_PAYLOAD_LEN,
    packets::{parse_stmt_packet, ComStmtExecuteRequestBuilder, ComStmtSendLongData},
};

/// Inner statement representation.
#[derive(Eq, PartialEq, Clone, Debug)]
pub struct InnerStmt {
    /// Positions and names of named parameters
    pub named_params: Option<Vec<String>>,
    pub params: Option<Vec<Column>>,
    pub columns: Option<Vec<Column>>,
    pub statement_id: u32,
    pub num_columns: u16,
    pub num_params: u16,
    pub warning_count: u16,
}

impl InnerStmt {
    // TODO: Consume payload?
    pub fn new(payload: &[u8], named_params: Option<Vec<String>>) -> Result<InnerStmt> {
        let packet = parse_stmt_packet(payload)?;

        Ok(InnerStmt {
            named_params,
            statement_id: packet.statement_id(),
            num_columns: packet.num_columns(),
            num_params: packet.num_params(),
            warning_count: packet.warning_count(),
            params: None,
            columns: None,
        })
    }
}

/// Prepared statement.
#[derive(Debug)]
pub struct Stmt<T> {
    conn_like: Option<Either<T, Streamless<T>>>,
    inner: InnerStmt,
    /// None => In use elsewhere
    /// Some(Cached) => Should not be closed
    /// Some(NotCached(_)) => Should be closed
    cached: Option<StmtCacheResult>,
}

pub fn new<T>(conn_like: T, inner: InnerStmt, cached: StmtCacheResult) -> Stmt<T>
where
    T: ConnectionLike + Sized + 'static,
{
    Stmt::new(conn_like, inner, cached)
}

impl<T> Stmt<T>
where
    T: ConnectionLike + Sized + 'static,
{
    fn new(conn_like: T, inner: InnerStmt, cached: StmtCacheResult) -> Stmt<T> {
        Stmt {
            conn_like: Some(Either::Left(conn_like)),
            inner,
            cached: Some(cached),
        }
    }

    /// Returns statement identifier.
    pub fn id(&self) -> u32 {
        self.inner.statement_id
    }

    /// Returns statement columns.
    ///
    /// ```rust
    /// # use mysql_async::test_misc::get_opts;
    /// use mysql_async::{consts::{ColumnFlags, ColumnType}, Pool};
    /// use mysql_async::prelude::*;
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<(), mysql_async::error::Error> {
    ///     let pool = Pool::new(get_opts());
    ///     let conn = pool.get_conn().await?;
    ///
    ///     let stmt = conn.prepare("SELECT 'foo', CAST(42 AS UNSIGNED)").await?;
    ///
    ///     let columns = stmt.columns();
    ///
    ///     assert_eq!(columns.len(), 2);
    ///
    ///     assert_eq!(columns[0].column_type(), ColumnType::MYSQL_TYPE_VAR_STRING);
    ///
    ///     assert_eq!(columns[1].column_type(), ColumnType::MYSQL_TYPE_LONGLONG);
    ///     assert!(columns[1].flags().contains(ColumnFlags::UNSIGNED_FLAG));
    ///
    ///     Ok(())
    /// }
    /// ```
    pub fn columns(&self) -> &[Column] {
        self.inner
            .columns
            .as_ref()
            .map(|c| &c[..])
            .unwrap_or_default()
    }

    /// Returns statement parameters.
    ///
    /// ```rust
    /// # use mysql_async::test_misc::get_opts;
    /// use mysql_async::{consts::{ColumnFlags, ColumnType}, Pool};
    /// use mysql_async::prelude::*;
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<(), mysql_async::error::Error> {
    ///     let pool = Pool::new(get_opts());
    ///     let conn = pool.get_conn().await?;
    ///
    ///     let stmt = conn.prepare("SELECT ?, ?").await?;
    ///
    ///     assert_eq!(stmt.params().len(), 2);
    ///
    ///     Ok(())
    /// }
    /// ```
    pub fn params(&self) -> &[Column] {
        self.inner
            .params
            .as_ref()
            .map(|c| &c[..])
            .unwrap_or_default()
    }

    async fn send_long_data(self, params: Vec<Value>) -> Result<Self> {
        let mut this = self;

        for (i, value) in params.into_iter().enumerate() {
            if let Value::Bytes(bytes) = value {
                let chunks = bytes.chunks(MAX_PAYLOAD_LEN - 6);
                let chunks = chunks.chain(if bytes.is_empty() {
                    Some(&[][..])
                } else {
                    None
                });
                for chunk in chunks {
                    let com = ComStmtSendLongData::new(this.inner.statement_id, i, chunk);
                    this = this.write_command_raw(com.into()).await?;
                }
            }
        }

        Ok(this)
    }

    async fn execute_positional<U>(self, params: U) -> Result<QueryResult<Self, BinaryProtocol>>
    where
        U: ::std::ops::Deref<Target = [Value]>,
        U: IntoIterator<Item = Value>,
        U: Send + 'static,
    {
        if self.inner.num_params as usize != params.len() {
            Err(DriverError::StmtParamsMismatch {
                required: self.inner.num_params,
                supplied: params.len() as u16,
            })?
        }

        let params = params.into_iter().collect::<Vec<_>>();

        let (body, as_long_data) =
            ComStmtExecuteRequestBuilder::new(self.inner.statement_id).build(&*params);

        let this = if as_long_data {
            self.send_long_data(params).await?
        } else {
            self
        };

        this.write_command_raw(body)
            .await?
            .read_result_set(None)
            .await
    }

    async fn execute_named(self, params: Params) -> Result<QueryResult<Self, BinaryProtocol>> {
        if self.inner.named_params.is_none() {
            let error = DriverError::NamedParamsForPositionalQuery.into();
            return Err(error);
        }

        let positional_params =
            match params.into_positional(self.inner.named_params.as_ref().unwrap()) {
                Ok(positional_params) => positional_params,
                Err(error) => return Err(error.into()),
            };

        match positional_params {
            Params::Positional(params) => self.execute_positional(params).await,
            _ => unreachable!(),
        }
    }

    async fn execute_empty(self) -> Result<QueryResult<Self, BinaryProtocol>> {
        if self.inner.num_params > 0 {
            let error = DriverError::StmtParamsMismatch {
                required: self.inner.num_params,
                supplied: 0,
            }
            .into();
            return Err(error);
        }

        let (body, _) = ComStmtExecuteRequestBuilder::new(self.inner.statement_id).build(&[]);
        let this = self.write_command_raw(body).await?;
        this.read_result_set(None).await
    }

    /// See `Queryable::execute`
    pub async fn execute<P>(self, params: P) -> Result<QueryResult<Self, BinaryProtocol>>
    where
        P: Into<Params>,
    {
        let params = params.into();
        match params {
            Params::Positional(params) => self.execute_positional(params).await,
            Params::Named(_) => self.execute_named(params).await,
            Params::Empty => self.execute_empty().await,
        }
    }

    /// See `Queryable::first`
    pub async fn first<P, R>(self, params: P) -> Result<(Self, Option<R>)>
    where
        P: Into<Params> + 'static,
        R: FromRow,
    {
        let result = self.execute(params).await?;
        let (this, mut rows) = result.collect_and_drop::<Row>().await?;
        if rows.len() > 1 {
            Ok((this, Some(FromRow::from_row(rows.swap_remove(0)))))
        } else {
            Ok((this, rows.pop().map(FromRow::from_row)))
        }
    }

    /// See `Queryable::batch`
    pub async fn batch<I, P>(self, params_iter: I) -> Result<Self>
    where
        I: IntoIterator<Item = P>,
        I::IntoIter: Send + 'static,
        Params: From<P>,
        P: 'static,
    {
        let mut params_iter = params_iter.into_iter().map(Params::from);
        let mut this = self;
        loop {
            match params_iter.next() {
                Some(params) => {
                    this = this.execute(params).await?.drop_result().await?;
                }
                None => break Ok(this),
            }
        }
    }

    /// This will close statement (if it's not in the cache) and resolve to a wrapped queryable.
    pub async fn close(mut self) -> Result<T> {
        let cached = self.cached.take();
        match self.conn_like {
            Some(Either::Left(conn_like)) => {
                if let Some(StmtCacheResult::NotCached(stmt_id)) = cached {
                    conn_like.close_stmt(stmt_id).await
                } else {
                    Ok(conn_like)
                }
            }
            _ => unreachable!(),
        }
    }

    pub(crate) fn unwrap(mut self) -> (T, Option<StmtCacheResult>) {
        match self.conn_like {
            Some(Either::Left(conn_like)) => (conn_like, self.cached.take()),
            _ => unreachable!(),
        }
    }
}

impl<T: ConnectionLike + 'static> ConnectionLikeWrapper for Stmt<T> {
    type ConnLike = T;

    fn take_stream(self) -> (Streamless<Self>, io::Stream)
    where
        Self: Sized,
    {
        let Stmt {
            conn_like,
            inner,
            cached,
        } = self;
        match conn_like {
            Some(Either::Left(conn_like)) => {
                let (streamless, stream) = conn_like.take_stream();
                let this = Stmt {
                    conn_like: Some(Either::Right(streamless)),
                    inner,
                    cached,
                };
                (Streamless::new(this), stream)
            }
            _ => unreachable!(),
        }
    }

    fn return_stream(&mut self, stream: io::Stream) {
        let conn_like = self.conn_like.take().unwrap();
        match conn_like {
            Either::Right(streamless) => {
                self.conn_like = Some(Either::Left(streamless.return_stream(stream)));
            }
            _ => unreachable!(),
        }
    }

    fn conn_like_ref(&self) -> &Self::ConnLike {
        match self.conn_like {
            Some(Either::Left(ref conn_like)) => conn_like,
            _ => unreachable!(),
        }
    }

    fn conn_like_mut(&mut self) -> &mut Self::ConnLike {
        match self.conn_like {
            Some(Either::Left(ref mut conn_like)) => conn_like,
            _ => unreachable!(),
        }
    }
}
<commit_msg>Doctest fix for mariadb<commit_after>// Copyright (c) 2017 Anatoly Ikorsky
//
// 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. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

use futures_util::future::Either;

use crate::{
    connection_like::{
        streamless::Streamless, ConnectionLike, ConnectionLikeWrapper, StmtCacheResult,
    },
    error::*,
    io,
    prelude::FromRow,
    queryable::{query_result::QueryResult, BinaryProtocol},
    Column, Params, Row,
    Value::{self},
};
use mysql_common::{
    constants::MAX_PAYLOAD_LEN,
    packets::{parse_stmt_packet, ComStmtExecuteRequestBuilder, ComStmtSendLongData},
};

/// Inner statement representation.
#[derive(Eq, PartialEq, Clone, Debug)]
pub struct InnerStmt {
    /// Positions and names of named parameters
    pub named_params: Option<Vec<String>>,
    pub params: Option<Vec<Column>>,
    pub columns: Option<Vec<Column>>,
    pub statement_id: u32,
    pub num_columns: u16,
    pub num_params: u16,
    pub warning_count: u16,
}

impl InnerStmt {
    // TODO: Consume payload?
    pub fn new(payload: &[u8], named_params: Option<Vec<String>>) -> Result<InnerStmt> {
        let packet = parse_stmt_packet(payload)?;

        Ok(InnerStmt {
            named_params,
            statement_id: packet.statement_id(),
            num_columns: packet.num_columns(),
            num_params: packet.num_params(),
            warning_count: packet.warning_count(),
            params: None,
            columns: None,
        })
    }
}

/// Prepared statement.
#[derive(Debug)]
pub struct Stmt<T> {
    conn_like: Option<Either<T, Streamless<T>>>,
    inner: InnerStmt,
    /// None => In use elsewhere
    /// Some(Cached) => Should not be closed
    /// Some(NotCached(_)) => Should be closed
    cached: Option<StmtCacheResult>,
}

pub fn new<T>(conn_like: T, inner: InnerStmt, cached: StmtCacheResult) -> Stmt<T>
where
    T: ConnectionLike + Sized + 'static,
{
    Stmt::new(conn_like, inner, cached)
}

impl<T> Stmt<T>
where
    T: ConnectionLike + Sized + 'static,
{
    fn new(conn_like: T, inner: InnerStmt, cached: StmtCacheResult) -> Stmt<T> {
        Stmt {
            conn_like: Some(Either::Left(conn_like)),
            inner,
            cached: Some(cached),
        }
    }

    /// Returns statement identifier.
    pub fn id(&self) -> u32 {
        self.inner.statement_id
    }

    /// Returns statement columns.
    ///
    /// ```rust
    /// # use mysql_async::test_misc::get_opts;
    /// use mysql_async::{consts::{ColumnFlags, ColumnType}, Pool};
    /// use mysql_async::prelude::*;
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<(), mysql_async::error::Error> {
    ///     let pool = Pool::new(get_opts());
    ///     let conn = pool.get_conn().await?;
    ///
    ///     let stmt = conn.prepare("SELECT 'foo', CAST(42 AS UNSIGNED)").await?;
    ///
    ///     let columns = stmt.columns();
    ///
    ///     assert_eq!(columns.len(), 2);
    ///     assert_eq!(columns[0].column_type(), ColumnType::MYSQL_TYPE_VAR_STRING);
    ///     assert!(columns[1].flags().contains(ColumnFlags::UNSIGNED_FLAG));
    ///
    ///     Ok(())
    /// }
    /// ```
    pub fn columns(&self) -> &[Column] {
        self.inner
            .columns
            .as_ref()
            .map(|c| &c[..])
            .unwrap_or_default()
    }

    /// Returns statement parameters.
    ///
    /// ```rust
    /// # use mysql_async::test_misc::get_opts;
    /// use mysql_async::{consts::{ColumnFlags, ColumnType}, Pool};
    /// use mysql_async::prelude::*;
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<(), mysql_async::error::Error> {
    ///     let pool = Pool::new(get_opts());
    ///     let conn = pool.get_conn().await?;
    ///
    ///     let stmt = conn.prepare("SELECT ?, ?").await?;
    ///
    ///     assert_eq!(stmt.params().len(), 2);
    ///
    ///     Ok(())
    /// }
    /// ```
    pub fn params(&self) -> &[Column] {
        self.inner
            .params
            .as_ref()
            .map(|c| &c[..])
            .unwrap_or_default()
    }

    async fn send_long_data(self, params: Vec<Value>) -> Result<Self> {
        let mut this = self;

        for (i, value) in params.into_iter().enumerate() {
            if let Value::Bytes(bytes) = value {
                let chunks = bytes.chunks(MAX_PAYLOAD_LEN - 6);
                let chunks = chunks.chain(if bytes.is_empty() {
                    Some(&[][..])
                } else {
                    None
                });
                for chunk in chunks {
                    let com = ComStmtSendLongData::new(this.inner.statement_id, i, chunk);
                    this = this.write_command_raw(com.into()).await?;
                }
            }
        }

        Ok(this)
    }

    async fn execute_positional<U>(self, params: U) -> Result<QueryResult<Self, BinaryProtocol>>
    where
        U: ::std::ops::Deref<Target = [Value]>,
        U: IntoIterator<Item = Value>,
        U: Send + 'static,
    {
        if self.inner.num_params as usize != params.len() {
            Err(DriverError::StmtParamsMismatch {
                required: self.inner.num_params,
                supplied: params.len() as u16,
            })?
        }

        let params = params.into_iter().collect::<Vec<_>>();

        let (body, as_long_data) =
            ComStmtExecuteRequestBuilder::new(self.inner.statement_id).build(&*params);

        let this = if as_long_data {
            self.send_long_data(params).await?
        } else {
            self
        };

        this.write_command_raw(body)
            .await?
            .read_result_set(None)
            .await
    }

    async fn execute_named(self, params: Params) -> Result<QueryResult<Self, BinaryProtocol>> {
        if self.inner.named_params.is_none() {
            let error = DriverError::NamedParamsForPositionalQuery.into();
            return Err(error);
        }

        let positional_params =
            match params.into_positional(self.inner.named_params.as_ref().unwrap()) {
                Ok(positional_params) => positional_params,
                Err(error) => return Err(error.into()),
            };

        match positional_params {
            Params::Positional(params) => self.execute_positional(params).await,
            _ => unreachable!(),
        }
    }

    async fn execute_empty(self) -> Result<QueryResult<Self, BinaryProtocol>> {
        if self.inner.num_params > 0 {
            let error = DriverError::StmtParamsMismatch {
                required: self.inner.num_params,
                supplied: 0,
            }
            .into();
            return Err(error);
        }

        let (body, _) = ComStmtExecuteRequestBuilder::new(self.inner.statement_id).build(&[]);
        let this = self.write_command_raw(body).await?;
        this.read_result_set(None).await
    }

    /// See `Queryable::execute`
    pub async fn execute<P>(self, params: P) -> Result<QueryResult<Self, BinaryProtocol>>
    where
        P: Into<Params>,
    {
        let params = params.into();
        match params {
            Params::Positional(params) => self.execute_positional(params).await,
            Params::Named(_) => self.execute_named(params).await,
            Params::Empty => self.execute_empty().await,
        }
    }

    /// See `Queryable::first`
    pub async fn first<P, R>(self, params: P) -> Result<(Self, Option<R>)>
    where
        P: Into<Params> + 'static,
        R: FromRow,
    {
        let result = self.execute(params).await?;
        let (this, mut rows) = result.collect_and_drop::<Row>().await?;
        if rows.len() > 1 {
            Ok((this, Some(FromRow::from_row(rows.swap_remove(0)))))
        } else {
            Ok((this, rows.pop().map(FromRow::from_row)))
        }
    }

    /// See `Queryable::batch`
    pub async fn batch<I, P>(self, params_iter: I) -> Result<Self>
    where
        I: IntoIterator<Item = P>,
        I::IntoIter: Send + 'static,
        Params: From<P>,
        P: 'static,
    {
        let mut params_iter = params_iter.into_iter().map(Params::from);
        let mut this = self;
        loop {
            match params_iter.next() {
                Some(params) => {
                    this = this.execute(params).await?.drop_result().await?;
                }
                None => break Ok(this),
            }
        }
    }

    /// This will close statement (if it's not in the cache) and resolve to a wrapped queryable.
    pub async fn close(mut self) -> Result<T> {
        let cached = self.cached.take();
        match self.conn_like {
            Some(Either::Left(conn_like)) => {
                if let Some(StmtCacheResult::NotCached(stmt_id)) = cached {
                    conn_like.close_stmt(stmt_id).await
                } else {
                    Ok(conn_like)
                }
            }
            _ => unreachable!(),
        }
    }

    pub(crate) fn unwrap(mut self) -> (T, Option<StmtCacheResult>) {
        match self.conn_like {
            Some(Either::Left(conn_like)) => (conn_like, self.cached.take()),
            _ => unreachable!(),
        }
    }
}

impl<T: ConnectionLike + 'static> ConnectionLikeWrapper for Stmt<T> {
    type ConnLike = T;

    fn take_stream(self) -> (Streamless<Self>, io::Stream)
    where
        Self: Sized,
    {
        let Stmt {
            conn_like,
            inner,
            cached,
        } = self;
        match conn_like {
            Some(Either::Left(conn_like)) => {
                let (streamless, stream) = conn_like.take_stream();
                let this = Stmt {
                    conn_like: Some(Either::Right(streamless)),
                    inner,
                    cached,
                };
                (Streamless::new(this), stream)
            }
            _ => unreachable!(),
        }
    }

    fn return_stream(&mut self, stream: io::Stream) {
        let conn_like = self.conn_like.take().unwrap();
        match conn_like {
            Either::Right(streamless) => {
                self.conn_like = Some(Either::Left(streamless.return_stream(stream)));
            }
            _ => unreachable!(),
        }
    }

    fn conn_like_ref(&self) -> &Self::ConnLike {
        match self.conn_like {
            Some(Either::Left(ref conn_like)) => conn_like,
            _ => unreachable!(),
        }
    }

    fn conn_like_mut(&mut self) -> &mut Self::ConnLike {
        match self.conn_like {
            Some(Either::Left(ref mut conn_like)) => conn_like,
            _ => unreachable!(),
        }
    }
}
<|endoftext|>