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

use crate::{util::get_fragment_filename, ModuleMetadata};
use common::WithLocation;
use fnv::FnvHashMap;
use graphql_ir::{
    Argument, ConstantValue, Directive, FragmentDefinition, OperationDefinition, Program,
    Selection, Transformed, Transformer, Value,
};
use interner::{Intern, StringKey};
use lazy_static::lazy_static;
use schema::{Schema, TypeReference};

lazy_static! {
    pub static ref DATA_DRIVEN_DEPENDENCY_METADATA_KEY: StringKey =
        "__dataDrivenDependencyMetadata".intern();
}

pub fn generate_data_driven_dependency_metadata(program: &Program) -> Program {
    let mut transformer = GenerateDataDrivenDependencyMetadata::new(program);
    transformer
        .transform_program(program)
        .replace_or_else(|| program.clone())
}

struct GenerateDataDrivenDependencyMetadata<'s> {
    pub program: &'s Program,
}

impl<'s> GenerateDataDrivenDependencyMetadata<'s> {
    fn new(program: &'s Program) -> Self {
        GenerateDataDrivenDependencyMetadata { program }
    }

    fn generate_data_driven_dependency_for_selections(
        &mut self,
        type_: TypeReference,
        selections: &[Selection],
    ) -> Option<Directive> {
        let mut processing_queue: Vec<ProcessingItem<'_>> = vec![ProcessingItem {
            plural: false,
            parent_type: type_,
            selections,
        }];

        let mut module_entries: ModuleEntries = Default::default();

        while let Some(processing_item) = processing_queue.pop() {
            for selection in processing_item.selections {
                match selection {
                    Selection::ScalarField(_) | Selection::FragmentSpread(_) => {}
                    Selection::LinkedField(linked_filed) => {
                        let field_type = &self
                            .program
                            .schema
                            .field(linked_filed.definition.item)
                            .type_;
                        processing_queue.push(ProcessingItem {
                            plural: processing_item.plural || field_type.is_list(),
                            parent_type: field_type.clone(),
                            selections: &linked_filed.selections,
                        });
                    }
                    Selection::InlineFragment(inline_fragment) => {
                        let parent_type = match inline_fragment.type_condition {
                            Some(type_) => TypeReference::Named(type_),
                            None => processing_item.parent_type.clone(),
                        };
                        if let Some(module_metadata) =
                            ModuleMetadata::find(&inline_fragment.directives)
                        {
                            let id = module_metadata.module_id;
                            let component = module_metadata.module_name;

                            let fragment_spread = inline_fragment
                                .selections
                                .iter()
                                .find(|item| matches!(item, Selection::FragmentSpread(_)));
                            // This is expected to be a fragment spread
                            let fragment_name = match fragment_spread {
                                Some(Selection::FragmentSpread(spread)) => spread.fragment.item,
                                _ => panic!("Expected to have a fragment spread"),
                            };

                            let type_name = self
                                .program
                                .schema
                                .get_type_string(&processing_item.parent_type);
                            module_entries
                                .entry(id)
                                .and_modify(|module_entry| {
                                    module_entry.branches.insert(
                                        type_name.clone(),
                                        Branch {
                                            component,
                                            fragment: get_fragment_filename(fragment_name),
                                        },
                                    );
                                })
                                .or_insert(ModuleEntry {
                                    id,
                                    branches: {
                                        let mut map: FnvHashMap<String, Branch> =
                                            Default::default();
                                        map.insert(
                                            type_name.clone(),
                                            Branch {
                                                component,
                                                fragment: get_fragment_filename(fragment_name),
                                            },
                                        );
                                        map
                                    },
                                    plural: processing_item.plural,
                                });
                        }
                        processing_queue.push(ProcessingItem {
                            plural: processing_item.plural,
                            parent_type,
                            selections: &inline_fragment.selections,
                        });
                    }
                    Selection::Condition(condition) => {
                        processing_queue.push(ProcessingItem {
                            plural: processing_item.plural,
                            parent_type: processing_item.parent_type.clone(),
                            selections: &condition.selections,
                        });
                    }
                }
            }
        }

        if !module_entries.is_empty() {
            Some(create_metadata_directive(module_entries))
        } else {
            None
        }
    }
}

#[derive(Debug, Copy, Clone)]
struct Branch {
    component: StringKey,
    fragment: StringKey,
}

type ModuleEntries = FnvHashMap<StringKey, ModuleEntry>;

#[derive(Debug)]
struct ModuleEntry {
    id: StringKey,
    branches: FnvHashMap<String, Branch>,
    plural: bool,
}

#[derive(Debug)]
struct ProcessingItem<'a> {
    plural: bool,
    parent_type: TypeReference,
    selections: &'a [Selection],
}

impl<'s> Transformer for GenerateDataDrivenDependencyMetadata<'s> {
    const NAME: &'static str = "GenerateDataDrivenDependencyMetadata";
    const VISIT_ARGUMENTS: bool = false;
    const VISIT_DIRECTIVES: bool = false;

    fn transform_operation(
        &mut self,
        operation: &OperationDefinition,
    ) -> Transformed<OperationDefinition> {
        if let Some(query_type) = self.program.schema.query_type() {
            let generated_directive = self.generate_data_driven_dependency_for_selections(
                TypeReference::Named(query_type),
                &operation.selections,
            );
            if let Some(generated_directive) = generated_directive {
                let mut next_directives: Vec<Directive> =
                    Vec::with_capacity(operation.directives.len() + 1);
                for directive in operation.directives.iter() {
                    next_directives.push(directive.clone());
                }

                next_directives.push(generated_directive);

                Transformed::Replace(OperationDefinition {
                    directives: next_directives,
                    ..operation.clone()
                })
            } else {
                Transformed::Keep
            }
        } else {
            Transformed::Keep
        }
    }

    fn transform_fragment(
        &mut self,
        fragment: &FragmentDefinition,
    ) -> Transformed<FragmentDefinition> {
        let generated_directive = self.generate_data_driven_dependency_for_selections(
            TypeReference::Named(fragment.type_condition),
            &fragment.selections,
        );
        if let Some(generated_directive) = generated_directive {
            let mut next_directives: Vec<Directive> =
                Vec::with_capacity(fragment.directives.len() + 1);
            for directive in fragment.directives.iter() {
                next_directives.push(directive.clone());
            }
            next_directives.push(generated_directive);

            Transformed::Replace(FragmentDefinition {
                directives: next_directives,
                ..fragment.clone()
            })
        } else {
            Transformed::Keep
        }
    }
}

fn create_metadata_directive(module_entries: FnvHashMap<StringKey, ModuleEntry>) -> Directive {
    let mut arguments: Vec<Argument> = Vec::with_capacity(module_entries.len());
    for (key, module_entry) in module_entries {
        arguments.push(Argument {
            name: WithLocation::generated(key),
            value: WithLocation::generated(Value::Constant(ConstantValue::String(From::from(
                module_entry,
            )))),
        })
    }
    arguments.sort_unstable_by(|a, b| a.name.item.cmp(&b.name.item));

    Directive {
        name: WithLocation::generated(*DATA_DRIVEN_DEPENDENCY_METADATA_KEY),
        arguments,
        data: None,
    }
}

impl From<ModuleEntry> for StringKey {
    fn from(module_entry: ModuleEntry) -> Self {
        let mut serialized_branches: Vec<(String, String)> =
            Vec::with_capacity(module_entry.branches.len());
        for (id, branch) in module_entry.branches.iter() {
            serialized_branches.push((
                id.clone(),
                format!(
                    "\"{}\":{{\"component\":\"{}\",\"fragment\":\"{}\"}}",
                    id, branch.component, branch.fragment
                ),
            ));
        }

        serialized_branches.sort_unstable_by(|a, b| a.0.cmp(&b.0));

        format!(
            "{{\"branches\":{{{}}},\"plural\":{}}}",
            serialized_branches
                .into_iter()
                .map(|(_, value)| value)
                .collect::<Vec<String>>()
                .join(","),
            module_entry.plural
        )
        .intern()
    }
}
<commit_msg>lint fixes for nightly builds<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::{util::get_fragment_filename, ModuleMetadata};
use common::WithLocation;
use fnv::FnvHashMap;
use graphql_ir::{
    Argument, ConstantValue, Directive, FragmentDefinition, OperationDefinition, Program,
    Selection, Transformed, Transformer, Value,
};
use interner::{Intern, StringKey};
use lazy_static::lazy_static;
use schema::{Schema, TypeReference};

lazy_static! {
    pub static ref DATA_DRIVEN_DEPENDENCY_METADATA_KEY: StringKey =
        "__dataDrivenDependencyMetadata".intern();
}

pub fn generate_data_driven_dependency_metadata(program: &Program) -> Program {
    let mut transformer = GenerateDataDrivenDependencyMetadata::new(program);
    transformer
        .transform_program(program)
        .replace_or_else(|| program.clone())
}

struct GenerateDataDrivenDependencyMetadata<'s> {
    pub program: &'s Program,
}

impl<'s> GenerateDataDrivenDependencyMetadata<'s> {
    fn new(program: &'s Program) -> Self {
        GenerateDataDrivenDependencyMetadata { program }
    }

    fn generate_data_driven_dependency_for_selections(
        &mut self,
        type_: TypeReference,
        selections: &[Selection],
    ) -> Option<Directive> {
        let mut processing_queue: Vec<ProcessingItem<'_>> = vec![ProcessingItem {
            plural: false,
            parent_type: type_,
            selections,
        }];

        let mut module_entries: ModuleEntries = Default::default();

        while let Some(processing_item) = processing_queue.pop() {
            for selection in processing_item.selections {
                match selection {
                    Selection::ScalarField(_) | Selection::FragmentSpread(_) => {}
                    Selection::LinkedField(linked_filed) => {
                        let field_type = &self
                            .program
                            .schema
                            .field(linked_filed.definition.item)
                            .type_;
                        processing_queue.push(ProcessingItem {
                            plural: processing_item.plural || field_type.is_list(),
                            parent_type: field_type.clone(),
                            selections: &linked_filed.selections,
                        });
                    }
                    Selection::InlineFragment(inline_fragment) => {
                        let parent_type = match inline_fragment.type_condition {
                            Some(type_) => TypeReference::Named(type_),
                            None => processing_item.parent_type.clone(),
                        };
                        if let Some(module_metadata) =
                            ModuleMetadata::find(&inline_fragment.directives)
                        {
                            let id = module_metadata.module_id;
                            let component = module_metadata.module_name;

                            let fragment_spread = inline_fragment
                                .selections
                                .iter()
                                .find(|item| matches!(item, Selection::FragmentSpread(_)));
                            // This is expected to be a fragment spread
                            let fragment_name = match fragment_spread {
                                Some(Selection::FragmentSpread(spread)) => spread.fragment.item,
                                _ => panic!("Expected to have a fragment spread"),
                            };

                            let type_name = self
                                .program
                                .schema
                                .get_type_string(&processing_item.parent_type);
                            module_entries
                                .entry(id)
                                .and_modify(|module_entry| {
                                    module_entry.branches.insert(
                                        type_name.clone(),
                                        Branch {
                                            component,
                                            fragment: get_fragment_filename(fragment_name),
                                        },
                                    );
                                })
                                .or_insert(ModuleEntry {
                                    branches: {
                                        let mut map: FnvHashMap<String, Branch> =
                                            Default::default();
                                        map.insert(
                                            type_name.clone(),
                                            Branch {
                                                component,
                                                fragment: get_fragment_filename(fragment_name),
                                            },
                                        );
                                        map
                                    },
                                    plural: processing_item.plural,
                                });
                        }
                        processing_queue.push(ProcessingItem {
                            plural: processing_item.plural,
                            parent_type,
                            selections: &inline_fragment.selections,
                        });
                    }
                    Selection::Condition(condition) => {
                        processing_queue.push(ProcessingItem {
                            plural: processing_item.plural,
                            parent_type: processing_item.parent_type.clone(),
                            selections: &condition.selections,
                        });
                    }
                }
            }
        }

        if !module_entries.is_empty() {
            Some(create_metadata_directive(module_entries))
        } else {
            None
        }
    }
}

#[derive(Debug, Copy, Clone)]
struct Branch {
    component: StringKey,
    fragment: StringKey,
}

type ModuleEntries = FnvHashMap<StringKey, ModuleEntry>;

#[derive(Debug)]
struct ModuleEntry {
    branches: FnvHashMap<String, Branch>,
    plural: bool,
}

#[derive(Debug)]
struct ProcessingItem<'a> {
    plural: bool,
    parent_type: TypeReference,
    selections: &'a [Selection],
}

impl<'s> Transformer for GenerateDataDrivenDependencyMetadata<'s> {
    const NAME: &'static str = "GenerateDataDrivenDependencyMetadata";
    const VISIT_ARGUMENTS: bool = false;
    const VISIT_DIRECTIVES: bool = false;

    fn transform_operation(
        &mut self,
        operation: &OperationDefinition,
    ) -> Transformed<OperationDefinition> {
        if let Some(query_type) = self.program.schema.query_type() {
            let generated_directive = self.generate_data_driven_dependency_for_selections(
                TypeReference::Named(query_type),
                &operation.selections,
            );
            if let Some(generated_directive) = generated_directive {
                let mut next_directives: Vec<Directive> =
                    Vec::with_capacity(operation.directives.len() + 1);
                for directive in operation.directives.iter() {
                    next_directives.push(directive.clone());
                }

                next_directives.push(generated_directive);

                Transformed::Replace(OperationDefinition {
                    directives: next_directives,
                    ..operation.clone()
                })
            } else {
                Transformed::Keep
            }
        } else {
            Transformed::Keep
        }
    }

    fn transform_fragment(
        &mut self,
        fragment: &FragmentDefinition,
    ) -> Transformed<FragmentDefinition> {
        let generated_directive = self.generate_data_driven_dependency_for_selections(
            TypeReference::Named(fragment.type_condition),
            &fragment.selections,
        );
        if let Some(generated_directive) = generated_directive {
            let mut next_directives: Vec<Directive> =
                Vec::with_capacity(fragment.directives.len() + 1);
            for directive in fragment.directives.iter() {
                next_directives.push(directive.clone());
            }
            next_directives.push(generated_directive);

            Transformed::Replace(FragmentDefinition {
                directives: next_directives,
                ..fragment.clone()
            })
        } else {
            Transformed::Keep
        }
    }
}

fn create_metadata_directive(module_entries: FnvHashMap<StringKey, ModuleEntry>) -> Directive {
    let mut arguments: Vec<Argument> = Vec::with_capacity(module_entries.len());
    for (key, module_entry) in module_entries {
        arguments.push(Argument {
            name: WithLocation::generated(key),
            value: WithLocation::generated(Value::Constant(ConstantValue::String(From::from(
                module_entry,
            )))),
        })
    }
    arguments.sort_unstable_by(|a, b| a.name.item.cmp(&b.name.item));

    Directive {
        name: WithLocation::generated(*DATA_DRIVEN_DEPENDENCY_METADATA_KEY),
        arguments,
        data: None,
    }
}

impl From<ModuleEntry> for StringKey {
    fn from(module_entry: ModuleEntry) -> Self {
        let mut serialized_branches: Vec<(String, String)> =
            Vec::with_capacity(module_entry.branches.len());
        for (id, branch) in module_entry.branches.iter() {
            serialized_branches.push((
                id.clone(),
                format!(
                    "\"{}\":{{\"component\":\"{}\",\"fragment\":\"{}\"}}",
                    id, branch.component, branch.fragment
                ),
            ));
        }

        serialized_branches.sort_unstable_by(|a, b| a.0.cmp(&b.0));

        format!(
            "{{\"branches\":{{{}}},\"plural\":{}}}",
            serialized_branches
                .into_iter()
                .map(|(_, value)| value)
                .collect::<Vec<String>>()
                .join(","),
            module_entry.plural
        )
        .intern()
    }
}
<|endoftext|>


<commit_before>fn main() {
    println!("Hello, world!");
}
<commit_msg>Requiring externals<commit_after>extern crate piston;
extern crate graphics;
extern crate glutin_window;
extern crate opengl_graphics;

use piston::window::WindowSettings;
use piston::event_loop::*;
use piston::input::*;
use glutin_window::GlutinWindow as Window;
use opengl_graphics::{ GlGraphics, OpenGL };

fn main() {
    println!("Hello, world!");
}
<|endoftext|>


<commit_before>use core::cmp::{Eq, Ord};
use iter::BaseIter;

pub trait PersistentMap<K: Copy Eq Ord, V: Copy> {
  pure fn get(k: K) -> Option<V>;
  pure fn put(k: K, v: V) -> self;
  pure fn delete(k: K) -> self;
}

enum RBColor {
  Red,
  Black
}

pub enum RBMap<K: Copy Eq Ord, V: Copy> {
  //   color    left         k  value       right
  Tree(RBColor, @RBMap<K,V>, K, Option<V>, @RBMap<K,V>),
  Leaf
}

pure fn RBMap<K: Copy Eq Ord, V: Copy>(key: K, value: V) -> @RBMap<K, V> {
  @Tree(Red, @Leaf, key, Some(value), @Leaf)
}

priv impl<K: Copy Eq Ord, V: Copy> @RBMap<K, V> {
  pure fn modify(k: K, new_value: Option<V>) -> @RBMap<K, V> {
    match self {
      @Leaf => @Tree(Red, self, k, copy new_value, self),
      @Tree(color, left, key, original_value, right) => {
        if (k.lt(&key)) {
          self.balanceLeft(color, left.modify(k, new_value), key, original_value, right)
        } else if (k == key) {
          @Tree(color, left, k, copy new_value, right)
        } else {
          self.balanceRight(color, left, key, original_value, right.modify(k, new_value))
        }
      }
    }
  }

  pure fn balanceLeft(c: RBColor, l: @RBMap<K,V>, k: K, v: Option<V>, r: @RBMap<K,V>) -> @RBMap<K,V> {
    match (c, l, k, v, r) {
      (Black,@Tree(Red,@Tree(Red,a,xK,xV,b),yK,yV,c),zK,zV,d) =>
        @Tree(Red,@Tree(Black,a,xK,xV,b),yK,yV,@Tree(Black,c,zK,zV,d)),
      (c,a,xK,xV,b) =>
        @Tree(c,a,xK,xV,b)
    }
  }

  pure fn balanceRight(c: RBColor, l: @RBMap<K,V>, k: K, v: Option<V>, r: @RBMap<K,V>) -> @RBMap<K,V> {
    match (c, l, k, v, r) {
      (Black, @Tree(Red, t1, k1, v1, t2), k0, v0, @Tree(Red, t3, k2, v2, t4)) =>
        @Tree(Red, @Tree(Black, t1, k1, v1, t2), k0, v0, @Tree(Black, t3, k2, v2, t4)),
      (c, l, k, v, @Tree(Red, t1, k1, v1, t2)) =>
        @Tree(c, @Tree(Red, l, k, v, t1), k1, v1, t2),
      (c,a,xK,xV,b) =>
        @Tree(c,a,xK,xV,b)
    }
  }
}

/**
 * A purely functional Left-Leaning Red-Black Tree.
 */
impl<K: Copy Eq Ord, V: Copy> @RBMap<K, V> : PersistentMap<K, V> {
  pure fn get(k: K) -> Option<V> {
    match self {
      @Leaf => None,
      @Tree(_, left, key, maybe_value, right) => {
        if (k.lt(&key)) {
          left.get(k)
        } else if (k.gt(&key)) {
          right.get(k)
        } else {
          copy maybe_value
        }
      }
    }
  }

  pure fn put(k: K, new_value: V) -> @RBMap<K, V> {
    self.modify(k, Some(new_value))
  }

  pure fn delete(k: K) -> @RBMap<K, V> {
    self.modify(k, None)
  }
}

trait BareIter<A> {
  pure fn each(&self, blk: fn(v: A) -> bool);
}

impl<K: Copy Eq Ord, V: Copy> RBMap<K, V>: BareIter<(K, V)> {
  pure fn each(&self, f: fn((K, V)) -> bool) {
    match *self {
      Leaf => (),
      Tree(_, left, key, maybe_value, right) => {
        left.each(f);
        match maybe_value {
          Some(value) => f((key, value)),
          None => false,
        };
        right.each(f);
      }
    }
  }
}

#[test]
fn test_rb_tree() {
  let v1 = RBMap("stevej", 150);
  let v2 = v1.put("thatstacy", 187);

  assert(v2.get("stevej") == Some(150));
  assert(v2.get("thatstacy") == Some(187));

  let v3 = v2.delete("stevej");
  assert(v3.get("stevej") == None);
  assert(v3.get("thatstacy") == Some(187));

  let v4 = v2.put("jeremy", 16);

  let v5 = v4.put("ev", 20);
  let v6 = v5.put("zhana", 51573);

  assert(v6.get("jeremy") == Some(16));
  assert(v6.get("stevej") == Some(150));
  assert(v6.get("thatstacy") == Some(187));
  assert(v6.get("ev") == Some(20));
  assert(v6.get("zhana") == Some(51573));
}

#[test]
fn test_base_iter_each() {
  let v1 = RBMap(1, 0);
  let v2 = v1.put(4, 0);
  let v3 = v2.put(3, 0);
  let v4 = v3.put(5, 0);
  let v5 = v4.put(2, 0);

  let n = @mut 1;

  fn t(n: @mut int, kv: &(&int, &int)) -> bool{
    match kv {
      &(k, _) => {
        assert (*n == *k);
        *n += 1;
      }
    }
    false
  }

  v5.each(|z| t(n, z));
}
<commit_msg>fixing tests<commit_after>use core::cmp::{Eq, Ord};
use iter::BaseIter;

pub trait PersistentMap<K: Copy Eq Ord, V: Copy> {
  pure fn get(k: K) -> Option<V>;
  pure fn put(k: K, v: V) -> self;
  pure fn delete(k: K) -> self;
}

enum RBColor {
  Red,
  Black
}

pub enum RBMap<K: Copy Eq Ord, V: Copy> {
  //   color    left         k  value       right
  Tree(RBColor, @RBMap<K,V>, K, Option<V>, @RBMap<K,V>),
  Leaf
}

pure fn RBMap<K: Copy Eq Ord, V: Copy>(key: K, value: V) -> @RBMap<K, V> {
  @Tree(Red, @Leaf, key, Some(value), @Leaf)
}

priv impl<K: Copy Eq Ord, V: Copy> @RBMap<K, V> {
  pure fn modify(k: K, new_value: Option<V>) -> @RBMap<K, V> {
    match self {
      @Leaf => @Tree(Red, self, k, copy new_value, self),
      @Tree(color, left, key, original_value, right) => {
        if (k.lt(&key)) {
          self.balanceLeft(color, left.modify(k, new_value), key, original_value, right)
        } else if (k == key) {
          @Tree(color, left, k, copy new_value, right)
        } else {
          self.balanceRight(color, left, key, original_value, right.modify(k, new_value))
        }
      }
    }
  }

  pure fn balanceLeft(c: RBColor, l: @RBMap<K,V>, k: K, v: Option<V>, r: @RBMap<K,V>) -> @RBMap<K,V> {
    match (c, l, k, v, r) {
      (Black,@Tree(Red,@Tree(Red,a,xK,xV,b),yK,yV,c),zK,zV,d) =>
        @Tree(Red,@Tree(Black,a,xK,xV,b),yK,yV,@Tree(Black,c,zK,zV,d)),
      (c,a,xK,xV,b) =>
        @Tree(c,a,xK,xV,b)
    }
  }

  pure fn balanceRight(c: RBColor, l: @RBMap<K,V>, k: K, v: Option<V>, r: @RBMap<K,V>) -> @RBMap<K,V> {
    match (c, l, k, v, r) {
      (Black, @Tree(Red, t1, k1, v1, t2), k0, v0, @Tree(Red, t3, k2, v2, t4)) =>
        @Tree(Red, @Tree(Black, t1, k1, v1, t2), k0, v0, @Tree(Black, t3, k2, v2, t4)),
      (c, l, k, v, @Tree(Red, t1, k1, v1, t2)) =>
        @Tree(c, @Tree(Red, l, k, v, t1), k1, v1, t2),
      (c,a,xK,xV,b) =>
        @Tree(c,a,xK,xV,b)
    }
  }
}

/**
 * A purely functional Left-Leaning Red-Black Tree.
 */
impl<K: Copy Eq Ord, V: Copy> @RBMap<K, V> : PersistentMap<K, V> {
  pure fn get(k: K) -> Option<V> {
    match self {
      @Leaf => None,
      @Tree(_, left, key, maybe_value, right) => {
        if (k.lt(&key)) {
          left.get(k)
        } else if (k.gt(&key)) {
          right.get(k)
        } else {
          copy maybe_value
        }
      }
    }
  }

  pure fn put(k: K, new_value: V) -> @RBMap<K, V> {
    self.modify(k, Some(new_value))
  }

  pure fn delete(k: K) -> @RBMap<K, V> {
    self.modify(k, None)
  }
}

trait BareIter<A> {
  pure fn each(&self, blk: fn(v: A) -> bool);
}

impl<K: Copy Eq Ord, V: Copy> RBMap<K, V>: BareIter<(K, V)> {
  pure fn each(&self, f: fn((K, V)) -> bool) {
    match *self {
      Leaf => (),
      Tree(_, left, key, maybe_value, right) => {
        left.each(f);
        match maybe_value {
          Some(value) => f((key, value)),
          None => false,
        };
        right.each(f);
      }
    }
  }
}

#[test]
fn test_rb_tree() {
  let v1 = RBMap("stevej", 150);
  let v2 = v1.put("thatstacy", 187);

  assert(v2.get("stevej") == Some(150));
  assert(v2.get("thatstacy") == Some(187));

  let v3 = v2.delete("stevej");
  assert(v3.get("stevej") == None);
  assert(v3.get("thatstacy") == Some(187));

  let v4 = v2.put("jeremy", 16);

  let v5 = v4.put("ev", 20);
  let v6 = v5.put("zhana", 51573);

  assert(v6.get("jeremy") == Some(16));
  assert(v6.get("stevej") == Some(150));
  assert(v6.get("thatstacy") == Some(187));
  assert(v6.get("ev") == Some(20));
  assert(v6.get("zhana") == Some(51573));
}

#[test]
fn test_base_iter_each() {
  let v1 = RBMap(1, 0);
  let v2 = v1.put(4, 0);
  let v3 = v2.put(3, 0);
  let v4 = v3.put(5, 0);
  let v5 = v4.put(2, 0);

  let n = @mut 1;

  fn t(n: @mut int, kv: (int, int)) -> bool{
    match kv {
      (k, _) => {
        assert (*n == k);
        *n += 1;
      }
    }
    false
  }

  v5.each(|z| t(n, z));
}
<|endoftext|>


<commit_before>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

//! Configuration for a resolver
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::time::Duration;

use trust_dns_proto::rr::Name;

/// Configuration for the upstream nameservers to use for resolution
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ResolverConfig {
    // base search domain
    domain: Option<Name>,
    // search domains
    search: Vec<Name>,
    // nameservers to use for resolution.
    name_servers: Vec<NameServerConfig>,
}

impl ResolverConfig {
    /// Creates a new empty configuration
    pub fn new() -> Self {
        ResolverConfig {
            // TODO: this should get the hostname and use the basename as the default
            domain: None,
            search: vec![],
            name_servers: vec![],
        }
    }

    /// Create a ResolverConfig with all parts specified
    ///
    /// # Arguments
    ///
    /// * `domain` - domain of the entity querying results. If the `Name` being looked up is not an FQDN, then this is the first part appended to attempt a lookup. `ndots` in the `ResolverOption` does take precedence over this.
    /// * `search` - additional search domains that are attempted if the `Name` is not found in `domain`, defaults to `vec![]`
    /// * `name_servers` - set of name servers to use for lookups, defaults are Google: `8.8.8.8`, `8.8.4.4` and `2001:4860:4860::8888`, `2001:4860:4860::8844`
    pub fn from_parts(
        domain: Option<Name>,
        search: Vec<Name>,
        name_servers: Vec<NameServerConfig>,
    ) -> Self {
        ResolverConfig {
            domain,
            search,
            name_servers,
        }
    }

    /// Returns the local domain
    ///
    /// By default any names will be appended to all non-fully-qualified-domain names, and searched for after any ndots rules
    pub fn domain(&self) -> Option<&Name> {
        self.domain.as_ref()
    }

    /// Set the domain of the entity querying results.
    pub fn set_domain(&mut self, domain: Name) {
        self.domain = Some(domain.clone());
        self.search = vec![domain];
    }

    /// Returns the search domains
    ///
    /// These will be queried after any local domain and then in the order of the set of search domains
    pub fn search(&self) -> &[Name] {
        &self.search
    }

    /// Add a search domain
    pub fn add_search(&mut self, search: Name) {
        self.search.push(search)
    }

    // TODO: consider allowing options per NameServer... like different timeouts?
    /// Add the configuration for a name server
    pub fn add_name_server(&mut self, name_server: NameServerConfig) {
        self.name_servers.push(name_server);
    }

    /// Returns a reference to the name servers
    pub fn name_servers(&self) -> &[NameServerConfig] {
        &self.name_servers
    }
}

impl Default for ResolverConfig {
    /// Creates a default configuration, using `8.8.8.8`, `8.8.4.4` and `2001:4860:4860::8888`, `2001:4860:4860::8844` (thank you, Google).
    ///
    /// Please see Google's [privacy statement](https://developers.google.com/speed/public-dns/privacy) for important information about what they track, many ISP's track similar information in DNS. To use the the system configuration see: `Resolver::from_system_conf` and `ResolverFuture::from_system_conf`
    fn default() -> Self {
        let domain = None;
        let google_ns1 = NameServerConfig {
            socket_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 53),
            protocol: Protocol::Udp,
        };

        let google_ns2 = NameServerConfig {
            socket_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 4, 4)), 53),
            protocol: Protocol::Udp,
        };

        let google_v6_ns1 = NameServerConfig {
            socket_addr: SocketAddr::new(
                IpAddr::V6(Ipv6Addr::new(
                    0x20,
                    0x01,
                    0x48,
                    0x60,
                    0x48,
                    0x60,
                    0x88,
                    0x88,
                )),
                53,
            ),
            protocol: Protocol::Udp,
        };

        let google_v6_ns2 = NameServerConfig {
            socket_addr: SocketAddr::new(
                IpAddr::V6(Ipv6Addr::new(
                    0x20,
                    0x01,
                    0x48,
                    0x60,
                    0x48,
                    0x60,
                    0x88,
                    0x44,
                )),
                53,
            ),
            protocol: Protocol::Udp,
        };

        ResolverConfig {
            domain,
            search: vec![],
            name_servers: vec![google_ns1, google_ns2, google_v6_ns1, google_v6_ns2],
        }
    }
}

/// The protocol on which a NameServer should be communicated with
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Protocol {
    /// UDP is the traditional DNS port, this is generally the correct choice
    Udp,
    /// TCP can be used for large queries, but not all NameServers support it
    Tcp,
    // TODO: add client certificate for mTLS?
    // Tls,
}

impl Protocol {
    /// Returns true if this is a datagram oriented protocol, e.g. UDP
    pub fn is_datagram(&self) -> bool {
        match *self {
            Protocol::Udp => true,
            Protocol::Tcp => false,
        }
    }

    /// Returns true if this is a stream oriented protocol, e.g. TCP
    pub fn is_stream(&self) -> bool {
        !self.is_datagram()
    }
}

/// Configuration for the NameServer
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct NameServerConfig {
    /// The address which the DNS NameServer is registered at.
    pub socket_addr: SocketAddr,
    /// The protocol to use when communicating with the NameServer.
    pub protocol: Protocol,
}

/// The lookup ip strategy
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LookupIpStrategy {
    /// Only query for A (Ipv4) records
    Ipv4Only,
    /// Only query for AAAA (Ipv6) records
    Ipv6Only,
    /// Query for A and AAAA in parallel
    Ipv4AndIpv6,
    /// Query for Ipv6 if that fails, query for Ipv4
    Ipv6thenIpv4,
    /// Query for Ipv4 if that fails, query for Ipv6 (default)
    Ipv4thenIpv6,
}

impl Default for LookupIpStrategy {
    /// Returns Ipv4AndIpv6 as the default.
    fn default() -> Self {
        LookupIpStrategy::Ipv4thenIpv6
    }
}

/// Configuration for the Resolver
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
#[allow(dead_code)] // TODO: remove after all params are supported
pub struct ResolverOpts {
    /// Sets the number of dots that must appear (unless it's a final dot representing the root)
    ///  that must appear before a query is assumted to include the TLD. The default is one, which
    ///  means that `www` would never be assumed to be a TLD, and would always be appended to either
    ///  the search
    pub ndots: usize,
    /// Specify the timeout for a request. Defaults to 5 seconds
    pub timeout: Duration,
    /// Number of attempts before giving up. Defaults to 2
    pub attempts: usize,
    /// Rotate through the resource records in the response (if there is more than one for a given name)
    pub(crate) rotate: bool,
    /// Validate the names in the response, not implemented don't really see the point unless you need to support
    ///  badly configured DNS
    pub(crate) check_names: bool,
    /// Enable edns, for larger records
    pub(crate) edns0: bool,
    /// Use DNSSec to validate the request
    pub validate: bool,
    /// The ip_strategy for the Resolver to use when lookup Ipv4 or Ipv6 addresses
    pub ip_strategy: LookupIpStrategy,
    /// Cache size is in number of records (some records can be large)
    pub cache_size: usize,
    /// Check /ect/hosts file before dns requery (only works for unix like OS)
    pub use_hosts_file: bool,
}

impl Default for ResolverOpts {
    /// Default values for the Reolver configuration.
    ///
    /// This follows the resolv.conf defaults as defined in the [Linux man pages](http://man7.org/linux/man-pages/man5/resolv.conf.5.html)
    fn default() -> Self {
        ResolverOpts {
            ndots: 1,
            timeout: Duration::from_secs(5),
            attempts: 2,
            rotate: false,
            check_names: true,
            edns0: false,
            validate: false,
            ip_strategy: LookupIpStrategy::default(),
            cache_size: 32,
            use_hosts_file: true,
        }
    }
}
<commit_msg>fixes the ipv6 addr construction for google NS<commit_after>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

//! Configuration for a resolver
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::time::Duration;

use trust_dns_proto::rr::Name;

/// Configuration for the upstream nameservers to use for resolution
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ResolverConfig {
    // base search domain
    domain: Option<Name>,
    // search domains
    search: Vec<Name>,
    // nameservers to use for resolution.
    name_servers: Vec<NameServerConfig>,
}

impl ResolverConfig {
    /// Creates a new empty configuration
    pub fn new() -> Self {
        ResolverConfig {
            // TODO: this should get the hostname and use the basename as the default
            domain: None,
            search: vec![],
            name_servers: vec![],
        }
    }

    /// Create a ResolverConfig with all parts specified
    ///
    /// # Arguments
    ///
    /// * `domain` - domain of the entity querying results. If the `Name` being looked up is not an FQDN, then this is the first part appended to attempt a lookup. `ndots` in the `ResolverOption` does take precedence over this.
    /// * `search` - additional search domains that are attempted if the `Name` is not found in `domain`, defaults to `vec![]`
    /// * `name_servers` - set of name servers to use for lookups, defaults are Google: `8.8.8.8`, `8.8.4.4` and `2001:4860:4860::8888`, `2001:4860:4860::8844`
    pub fn from_parts(
        domain: Option<Name>,
        search: Vec<Name>,
        name_servers: Vec<NameServerConfig>,
    ) -> Self {
        ResolverConfig {
            domain,
            search,
            name_servers,
        }
    }

    /// Returns the local domain
    ///
    /// By default any names will be appended to all non-fully-qualified-domain names, and searched for after any ndots rules
    pub fn domain(&self) -> Option<&Name> {
        self.domain.as_ref()
    }

    /// Set the domain of the entity querying results.
    pub fn set_domain(&mut self, domain: Name) {
        self.domain = Some(domain.clone());
        self.search = vec![domain];
    }

    /// Returns the search domains
    ///
    /// These will be queried after any local domain and then in the order of the set of search domains
    pub fn search(&self) -> &[Name] {
        &self.search
    }

    /// Add a search domain
    pub fn add_search(&mut self, search: Name) {
        self.search.push(search)
    }

    // TODO: consider allowing options per NameServer... like different timeouts?
    /// Add the configuration for a name server
    pub fn add_name_server(&mut self, name_server: NameServerConfig) {
        self.name_servers.push(name_server);
    }

    /// Returns a reference to the name servers
    pub fn name_servers(&self) -> &[NameServerConfig] {
        &self.name_servers
    }
}

impl Default for ResolverConfig {
    /// Creates a default configuration, using `8.8.8.8`, `8.8.4.4` and `2001:4860:4860::8888`, `2001:4860:4860::8844` (thank you, Google).
    ///
    /// Please see Google's [privacy statement](https://developers.google.com/speed/public-dns/privacy) for important information about what they track, many ISP's track similar information in DNS. To use the the system configuration see: `Resolver::from_system_conf` and `ResolverFuture::from_system_conf`
    fn default() -> Self {
        let domain = None;
        let google_ns1 = NameServerConfig {
            socket_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 53),
            protocol: Protocol::Udp,
        };

        let google_ns2 = NameServerConfig {
            socket_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 4, 4)), 53),
            protocol: Protocol::Udp,
        };

        let google_v6_ns1 = NameServerConfig {
            socket_addr: SocketAddr::new(
                IpAddr::V6(Ipv6Addr::new(0x2001, 0x4860, 0x4860, 0, 0, 0, 0, 0x8888)),
                53,
            ),
            protocol: Protocol::Udp,
        };

        let google_v6_ns2 = NameServerConfig {
            socket_addr: SocketAddr::new(
                IpAddr::V6(Ipv6Addr::new(0x2001, 0x4860, 0x4860, 0, 0, 0, 0, 0x8844)),
                53,
            ),
            protocol: Protocol::Udp,
        };

        ResolverConfig {
            domain,
            search: vec![],
            name_servers: vec![google_ns1, google_ns2, google_v6_ns1, google_v6_ns2],
        }
    }
}

/// The protocol on which a NameServer should be communicated with
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Protocol {
    /// UDP is the traditional DNS port, this is generally the correct choice
    Udp,
    /// TCP can be used for large queries, but not all NameServers support it
    Tcp,
    // TODO: add client certificate for mTLS?
    // Tls,
}

impl Protocol {
    /// Returns true if this is a datagram oriented protocol, e.g. UDP
    pub fn is_datagram(&self) -> bool {
        match *self {
            Protocol::Udp => true,
            Protocol::Tcp => false,
        }
    }

    /// Returns true if this is a stream oriented protocol, e.g. TCP
    pub fn is_stream(&self) -> bool {
        !self.is_datagram()
    }
}

/// Configuration for the NameServer
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct NameServerConfig {
    /// The address which the DNS NameServer is registered at.
    pub socket_addr: SocketAddr,
    /// The protocol to use when communicating with the NameServer.
    pub protocol: Protocol,
}

/// The lookup ip strategy
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LookupIpStrategy {
    /// Only query for A (Ipv4) records
    Ipv4Only,
    /// Only query for AAAA (Ipv6) records
    Ipv6Only,
    /// Query for A and AAAA in parallel
    Ipv4AndIpv6,
    /// Query for Ipv6 if that fails, query for Ipv4
    Ipv6thenIpv4,
    /// Query for Ipv4 if that fails, query for Ipv6 (default)
    Ipv4thenIpv6,
}

impl Default for LookupIpStrategy {
    /// Returns Ipv4AndIpv6 as the default.
    fn default() -> Self {
        LookupIpStrategy::Ipv4thenIpv6
    }
}

/// Configuration for the Resolver
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
#[allow(dead_code)] // TODO: remove after all params are supported
pub struct ResolverOpts {
    /// Sets the number of dots that must appear (unless it's a final dot representing the root)
    ///  that must appear before a query is assumted to include the TLD. The default is one, which
    ///  means that `www` would never be assumed to be a TLD, and would always be appended to either
    ///  the search
    pub ndots: usize,
    /// Specify the timeout for a request. Defaults to 5 seconds
    pub timeout: Duration,
    /// Number of attempts before giving up. Defaults to 2
    pub attempts: usize,
    /// Rotate through the resource records in the response (if there is more than one for a given name)
    pub(crate) rotate: bool,
    /// Validate the names in the response, not implemented don't really see the point unless you need to support
    ///  badly configured DNS
    pub(crate) check_names: bool,
    /// Enable edns, for larger records
    pub(crate) edns0: bool,
    /// Use DNSSec to validate the request
    pub validate: bool,
    /// The ip_strategy for the Resolver to use when lookup Ipv4 or Ipv6 addresses
    pub ip_strategy: LookupIpStrategy,
    /// Cache size is in number of records (some records can be large)
    pub cache_size: usize,
    /// Check /ect/hosts file before dns requery (only works for unix like OS)
    pub use_hosts_file: bool,
}

impl Default for ResolverOpts {
    /// Default values for the Reolver configuration.
    ///
    /// This follows the resolv.conf defaults as defined in the [Linux man pages](http://man7.org/linux/man-pages/man5/resolv.conf.5.html)
    fn default() -> Self {
        ResolverOpts {
            ndots: 1,
            timeout: Duration::from_secs(5),
            attempts: 2,
            rotate: false,
            check_names: true,
            edns0: false,
            validate: false,
            ip_strategy: LookupIpStrategy::default(),
            cache_size: 32,
            use_hosts_file: true,
        }
    }
}
<|endoftext|>


<commit_before>//! A channel for sending a single message between asynchronous tasks.

use std::marker::Unpin;
use std::mem::PinMut;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering::SeqCst;
use std::error::Error;
use std::fmt;

use futures_core::{Future, Poll};
use futures_core::task::{self, Waker};

use crate::lock::Lock;

/// A future for a value that will be provided by another asynchronous task.
///
/// This is created by the [`channel`](channel) function.
#[must_use = "futures do nothing unless polled"]
#[derive(Debug)]
pub struct Receiver<T> {
    inner: Arc<Inner<T>>,
}

/// A means of transmitting a single value to another task.
///
/// This is created by the [`channel`](channel) function.
#[derive(Debug)]
pub struct Sender<T> {
    inner: Arc<Inner<T>>,
}

// The channels do not ever project PinMut to the inner T
impl<T> Unpin for Receiver<T> {}
impl<T> Unpin for Sender<T> {}

/// Internal state of the `Receiver`/`Sender` pair above. This is all used as
/// the internal synchronization between the two for send/recv operations.
#[derive(Debug)]
struct Inner<T> {
    /// Indicates whether this oneshot is complete yet. This is filled in both
    /// by `Sender::drop` and by `Receiver::drop`, and both sides interpret it
    /// appropriately.
    ///
    /// For `Receiver`, if this is `true`, then it's guaranteed that `data` is
    /// unlocked and ready to be inspected.
    ///
    /// For `Sender` if this is `true` then the oneshot has gone away and it
    /// can return ready from `poll_cancel`.
    complete: AtomicBool,

    /// The actual data being transferred as part of this `Receiver`. This is
    /// filled in by `Sender::complete` and read by `Receiver::poll`.
    ///
    /// Note that this is protected by `Lock`, but it is in theory safe to
    /// replace with an `UnsafeCell` as it's actually protected by `complete`
    /// above. I wouldn't recommend doing this, however, unless someone is
    /// supremely confident in the various atomic orderings here and there.
    data: Lock<Option<T>>,

    /// Field to store the task which is blocked in `Receiver::poll`.
    ///
    /// This is filled in when a oneshot is polled but not ready yet. Note that
    /// the `Lock` here, unlike in `data` above, is important to resolve races.
    /// Both the `Receiver` and the `Sender` halves understand that if they
    /// can't acquire the lock then some important interference is happening.
    rx_task: Lock<Option<Waker>>,

    /// Like `rx_task` above, except for the task blocked in
    /// `Sender::poll_cancel`. Additionally, `Lock` cannot be `UnsafeCell`.
    tx_task: Lock<Option<Waker>>,
}

/// Creates a new one-shot channel for sending values across asynchronous tasks.
///
/// This function is similar to Rust's channel constructor found in the standard library.
/// Two halves are returned, the first of which is a `Sender` handle, used to
/// signal the end of a computation and provide its value. The second half is a
/// `Receiver` which implements the `Future` trait, resolving to the value that
/// was given to the `Sender` handle.
///
/// Each half can be separately owned and sent across tasks.
///
/// # Examples
///
/// ```
/// extern crate futures;
///
/// use std::thread;
///
/// use futures::channel::oneshot;
/// use futures::*;
///
/// fn main() {
///     let (tx, rx) = oneshot::channel::<i32>();
///
/// # let t =
///     thread::spawn(|| {
///         let future = rx.map(|i| {
///             println!("got: {:?}", i);
///         });
///         // ...
/// # return future;
///     });
///
///     tx.send(3).unwrap();
/// # futures::executor::block_on(t.join().unwrap());
/// }
/// ```
pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
    let inner = Arc::new(Inner::new());
    let receiver = Receiver {
        inner: inner.clone(),
    };
    let sender = Sender {
        inner,
    };
    (sender, receiver)
}

impl<T> Inner<T> {
    fn new() -> Inner<T> {
        Inner {
            complete: AtomicBool::new(false),
            data: Lock::new(None),
            rx_task: Lock::new(None),
            tx_task: Lock::new(None),
        }
    }

    fn send(&self, t: T) -> Result<(), T> {
        if self.complete.load(SeqCst) {
            return Err(t)
        }

        // Note that this lock acquisition may fail if the receiver
        // is closed and sets the `complete` flag to true, whereupon
        // the receiver may call `poll()`.
        if let Some(mut slot) = self.data.try_lock() {
            assert!(slot.is_none());
            *slot = Some(t);
            drop(slot);

            // If the receiver called `close()` between the check at the
            // start of the function, and the lock being released, then
            // the receiver may not be around to receive it, so try to
            // pull it back out.
            if self.complete.load(SeqCst) {
                // If lock acquisition fails, then receiver is actually
                // receiving it, so we're good.
                if let Some(mut slot) = self.data.try_lock() {
                    if let Some(t) = slot.take() {
                        return Err(t);
                    }
                }
            }
            Ok(())
        } else {
            // Must have been closed
            Err(t)
        }
    }

    fn poll_cancel(&self, cx: &mut task::Context) -> Poll<()> {
        // Fast path up first, just read the flag and see if our other half is
        // gone. This flag is set both in our destructor and the oneshot
        // destructor, but our destructor hasn't run yet so if it's set then the
        // oneshot is gone.
        if self.complete.load(SeqCst) {
            return Poll::Ready(())
        }

        // If our other half is not gone then we need to park our current task
        // and move it into the `notify_cancel` slot to get notified when it's
        // actually gone.
        //
        // If `try_lock` fails, then the `Receiver` is in the process of using
        // it, so we can deduce that it's now in the process of going away and
        // hence we're canceled. If it succeeds then we just store our handle.
        //
        // Crucially we then check `oneshot_gone` *again* before we return.
        // While we were storing our handle inside `notify_cancel` the `Receiver`
        // may have been dropped. The first thing it does is set the flag, and
        // if it fails to acquire the lock it assumes that we'll see the flag
        // later on. So... we then try to see the flag later on!
        let handle = cx.waker().clone();
        match self.tx_task.try_lock() {
            Some(mut p) => *p = Some(handle),
            None => return Poll::Ready(()),
        }
        if self.complete.load(SeqCst) {
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }

    fn is_canceled(&self) -> bool {
        self.complete.load(SeqCst)
    }

    fn drop_tx(&self) {
        // Flag that we're a completed `Sender` and try to wake up a receiver.
        // Whether or not we actually stored any data will get picked up and
        // translated to either an item or cancellation.
        //
        // Note that if we fail to acquire the `rx_task` lock then that means
        // we're in one of two situations:
        //
        // 1. The receiver is trying to block in `poll`
        // 2. The receiver is being dropped
        //
        // In the first case it'll check the `complete` flag after it's done
        // blocking to see if it succeeded. In the latter case we don't need to
        // wake up anyone anyway. So in both cases it's ok to ignore the `None`
        // case of `try_lock` and bail out.
        //
        // The first case crucially depends on `Lock` using `SeqCst` ordering
        // under the hood. If it instead used `Release` / `Acquire` ordering,
        // then it would not necessarily synchronize with `inner.complete`
        // and deadlock might be possible, as was observed in
        // https://github.com/rust-lang-nursery/futures-rs/pull/219.
        self.complete.store(true, SeqCst);
        if let Some(mut slot) = self.rx_task.try_lock() {
            if let Some(task) = slot.take() {
                drop(slot);
                task.wake();
            }
        }
    }

    fn close_rx(&self) {
        // Flag our completion and then attempt to wake up the sender if it's
        // blocked. See comments in `drop` below for more info
        self.complete.store(true, SeqCst);
        if let Some(mut handle) = self.tx_task.try_lock() {
            if let Some(task) = handle.take() {
                drop(handle);
                task.wake()
            }
        }
    }

    fn try_recv(&self) -> Result<Option<T>, Canceled> {
        // If we're complete, either `::close_rx` or `::drop_tx` was called.
        // We can assume a successful send if data is present.
        if self.complete.load(SeqCst) {
            if let Some(mut slot) = self.data.try_lock() {
                if let Some(data) = slot.take() {
                    return Ok(Some(data));
                }
            }
            Err(Canceled)
        } else {
            Ok(None)
        }
    }

    fn recv(&self, cx: &mut task::Context) -> Poll<Result<T, Canceled>> {
        // Check to see if some data has arrived. If it hasn't then we need to
        // block our task.
        //
        // Note that the acquisition of the `rx_task` lock might fail below, but
        // the only situation where this can happen is during `Sender::drop`
        // when we are indeed completed already. If that's happening then we
        // know we're completed so keep going.
        let done = if self.complete.load(SeqCst) {
            true
        } else {
            let task = cx.waker().clone();
            match self.rx_task.try_lock() {
                Some(mut slot) => { *slot = Some(task); false },
                None => true,
            }
        };

        // If we're `done` via one of the paths above, then look at the data and
        // figure out what the answer is. If, however, we stored `rx_task`
        // successfully above we need to check again if we're completed in case
        // a message was sent while `rx_task` was locked and couldn't notify us
        // otherwise.
        //
        // If we're not done, and we're not complete, though, then we've
        // successfully blocked our task and we return `Pending`.
        if done || self.complete.load(SeqCst) {
            // If taking the lock fails, the sender will realise that the we're
            // `done` when it checks the `complete` flag on the way out, and will
            // treat the send as a failure.
            if let Some(mut slot) = self.data.try_lock() {
                if let Some(data) = slot.take() {
                    return Poll::Ready(Ok(data));
                }
            }
            Poll::Ready(Err(Canceled))
        } else {
            Poll::Pending
        }
    }

    fn drop_rx(&self) {
        // Indicate to the `Sender` that we're done, so any future calls to
        // `poll_cancel` are weeded out.
        self.complete.store(true, SeqCst);

        // If we've blocked a task then there's no need for it to stick around,
        // so we need to drop it. If this lock acquisition fails, though, then
        // it's just because our `Sender` is trying to take the task, so we
        // let them take care of that.
        if let Some(mut slot) = self.rx_task.try_lock() {
            let task = slot.take();
            drop(slot);
            drop(task);
        }

        // Finally, if our `Sender` wants to get notified of us going away, it
        // would have stored something in `tx_task`. Here we try to peel that
        // out and unpark it.
        //
        // Note that the `try_lock` here may fail, but only if the `Sender` is
        // in the process of filling in the task. If that happens then we
        // already flagged `complete` and they'll pick that up above.
        if let Some(mut handle) = self.tx_task.try_lock() {
            if let Some(task) = handle.take() {
                drop(handle);
                task.wake()
            }
        }
    }
}

impl<T> Sender<T> {
    /// Completes this oneshot with a successful result.
    ///
    /// This function will consume `self` and indicate to the other end, the
    /// [`Receiver`](Receiver), that the value provided is the result of the
    /// computation this represents.
    ///
    /// If the value is successfully enqueued for the remote end to receive,
    /// then `Ok(())` is returned. If the receiving end was dropped before
    /// this function was called, however, then `Err` is returned with the value
    /// provided.
    pub fn send(self, t: T) -> Result<(), T> {
        self.inner.send(t)
    }

    /// Polls this `Sender` half to detect whether its associated
    /// [`Receiver`](Receiver) with has been dropped.
    ///
    /// # Return values
    ///
    /// If `Ok(Ready)` is returned then the associated `Receiver` has been
    /// dropped, which means any work required for sending should be canceled.
    ///
    /// If `Ok(Pending)` is returned then the associated `Receiver` is still
    /// alive and may be able to receive a message if sent. The current task,
    /// however, is scheduled to receive a notification if the corresponding
    /// `Receiver` goes away.
    pub fn poll_cancel(&mut self, cx: &mut task::Context) -> Poll<()> {
        self.inner.poll_cancel(cx)
    }

    /// Tests to see whether this `Sender`'s corresponding `Receiver`
    /// has been dropped.
    ///
    /// Unlike [`poll_cancel`](Sender::poll_cancel), this function does not
    /// enqueue a task for wakeup upon cancellation, but merely reports the
    /// current state, which may be subject to concurrent modification.
    pub fn is_canceled(&self) -> bool {
        self.inner.is_canceled()
    }
}

impl<T> Drop for Sender<T> {
    fn drop(&mut self) {
        self.inner.drop_tx()
    }
}

/// Error returned from a [`Receiver`](Receiver) when the corresponding
/// [`Sender`](Sender) is dropped.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct Canceled;

impl fmt::Display for Canceled {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "oneshot canceled")
    }
}

impl Error for Canceled {
    fn description(&self) -> &str {
        "oneshot canceled"
    }
}

impl<T> Receiver<T> {
    /// Gracefully close this receiver, preventing any subsequent attempts to
    /// send to it.
    ///
    /// Any `send` operation which happens after this method returns is
    /// guaranteed to fail. After calling this method, you can use
    /// [`Receiver::poll`](Receiver::poll) to determine whether a message had
    /// previously been sent.
    pub fn close(&mut self) {
        self.inner.close_rx()
    }

    /// Attempts to receive a message outside of the context of a task.
    ///
    /// Useful when a [`Context`](Context) is not available such as within a
    /// `Drop` impl.
    ///
    /// Does not schedule a task wakeup or have any other side effects.
    ///
    /// A return value of `None` must be considered immediately stale (out of
    /// date) unless [`::close`](Receiver::close) has been called first.
    ///
    /// Returns an error if the sender was dropped.
    pub fn try_recv(&mut self) -> Result<Option<T>, Canceled> {
        self.inner.try_recv()
    }
}

impl<T> Future for Receiver<T> {
    type Output = Result<T, Canceled>;

    fn poll(self: PinMut<Self>, cx: &mut task::Context) -> Poll<Result<T, Canceled>> {
        self.inner.recv(cx)
    }
}

impl<T> Drop for Receiver<T> {
    fn drop(&mut self) {
        self.inner.drop_rx()
    }
}
<commit_msg>Align naming with mpsc::channel() docs<commit_after>//! A channel for sending a single message between asynchronous tasks.

use std::marker::Unpin;
use std::mem::PinMut;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering::SeqCst;
use std::error::Error;
use std::fmt;

use futures_core::{Future, Poll};
use futures_core::task::{self, Waker};

use crate::lock::Lock;

/// A future for a value that will be provided by another asynchronous task.
///
/// This is created by the [`channel`](channel) function.
#[must_use = "futures do nothing unless polled"]
#[derive(Debug)]
pub struct Receiver<T> {
    inner: Arc<Inner<T>>,
}

/// A means of transmitting a single value to another task.
///
/// This is created by the [`channel`](channel) function.
#[derive(Debug)]
pub struct Sender<T> {
    inner: Arc<Inner<T>>,
}

// The channels do not ever project PinMut to the inner T
impl<T> Unpin for Receiver<T> {}
impl<T> Unpin for Sender<T> {}

/// Internal state of the `Receiver`/`Sender` pair above. This is all used as
/// the internal synchronization between the two for send/recv operations.
#[derive(Debug)]
struct Inner<T> {
    /// Indicates whether this oneshot is complete yet. This is filled in both
    /// by `Sender::drop` and by `Receiver::drop`, and both sides interpret it
    /// appropriately.
    ///
    /// For `Receiver`, if this is `true`, then it's guaranteed that `data` is
    /// unlocked and ready to be inspected.
    ///
    /// For `Sender` if this is `true` then the oneshot has gone away and it
    /// can return ready from `poll_cancel`.
    complete: AtomicBool,

    /// The actual data being transferred as part of this `Receiver`. This is
    /// filled in by `Sender::complete` and read by `Receiver::poll`.
    ///
    /// Note that this is protected by `Lock`, but it is in theory safe to
    /// replace with an `UnsafeCell` as it's actually protected by `complete`
    /// above. I wouldn't recommend doing this, however, unless someone is
    /// supremely confident in the various atomic orderings here and there.
    data: Lock<Option<T>>,

    /// Field to store the task which is blocked in `Receiver::poll`.
    ///
    /// This is filled in when a oneshot is polled but not ready yet. Note that
    /// the `Lock` here, unlike in `data` above, is important to resolve races.
    /// Both the `Receiver` and the `Sender` halves understand that if they
    /// can't acquire the lock then some important interference is happening.
    rx_task: Lock<Option<Waker>>,

    /// Like `rx_task` above, except for the task blocked in
    /// `Sender::poll_cancel`. Additionally, `Lock` cannot be `UnsafeCell`.
    tx_task: Lock<Option<Waker>>,
}

/// Creates a new one-shot channel for sending values across asynchronous tasks.
///
/// This function is similar to Rust's channel constructor found in the standard library.
/// Two halves are returned, the first of which is a `Sender` handle, used to
/// signal the end of a computation and provide its value. The second half is a
/// `Receiver` which implements the `Future` trait, resolving to the value that
/// was given to the `Sender` handle.
///
/// Each half can be separately owned and sent across tasks.
///
/// # Examples
///
/// ```
/// extern crate futures;
///
/// use std::thread;
///
/// use futures::channel::oneshot;
/// use futures::*;
///
/// fn main() {
///     let (sender, receiver) = oneshot::channel::<i32>();
///
/// # let t =
///     thread::spawn(|| {
///         let future = receiver.map(|i| {
///             println!("got: {:?}", i);
///         });
///         // ...
/// # return future;
///     });
///
///     sender.send(3).unwrap();
/// # futures::executor::block_on(t.join().unwrap());
/// }
/// ```
pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
    let inner = Arc::new(Inner::new());
    let receiver = Receiver {
        inner: inner.clone(),
    };
    let sender = Sender {
        inner,
    };
    (sender, receiver)
}

impl<T> Inner<T> {
    fn new() -> Inner<T> {
        Inner {
            complete: AtomicBool::new(false),
            data: Lock::new(None),
            rx_task: Lock::new(None),
            tx_task: Lock::new(None),
        }
    }

    fn send(&self, t: T) -> Result<(), T> {
        if self.complete.load(SeqCst) {
            return Err(t)
        }

        // Note that this lock acquisition may fail if the receiver
        // is closed and sets the `complete` flag to true, whereupon
        // the receiver may call `poll()`.
        if let Some(mut slot) = self.data.try_lock() {
            assert!(slot.is_none());
            *slot = Some(t);
            drop(slot);

            // If the receiver called `close()` between the check at the
            // start of the function, and the lock being released, then
            // the receiver may not be around to receive it, so try to
            // pull it back out.
            if self.complete.load(SeqCst) {
                // If lock acquisition fails, then receiver is actually
                // receiving it, so we're good.
                if let Some(mut slot) = self.data.try_lock() {
                    if let Some(t) = slot.take() {
                        return Err(t);
                    }
                }
            }
            Ok(())
        } else {
            // Must have been closed
            Err(t)
        }
    }

    fn poll_cancel(&self, cx: &mut task::Context) -> Poll<()> {
        // Fast path up first, just read the flag and see if our other half is
        // gone. This flag is set both in our destructor and the oneshot
        // destructor, but our destructor hasn't run yet so if it's set then the
        // oneshot is gone.
        if self.complete.load(SeqCst) {
            return Poll::Ready(())
        }

        // If our other half is not gone then we need to park our current task
        // and move it into the `notify_cancel` slot to get notified when it's
        // actually gone.
        //
        // If `try_lock` fails, then the `Receiver` is in the process of using
        // it, so we can deduce that it's now in the process of going away and
        // hence we're canceled. If it succeeds then we just store our handle.
        //
        // Crucially we then check `oneshot_gone` *again* before we return.
        // While we were storing our handle inside `notify_cancel` the `Receiver`
        // may have been dropped. The first thing it does is set the flag, and
        // if it fails to acquire the lock it assumes that we'll see the flag
        // later on. So... we then try to see the flag later on!
        let handle = cx.waker().clone();
        match self.tx_task.try_lock() {
            Some(mut p) => *p = Some(handle),
            None => return Poll::Ready(()),
        }
        if self.complete.load(SeqCst) {
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }

    fn is_canceled(&self) -> bool {
        self.complete.load(SeqCst)
    }

    fn drop_tx(&self) {
        // Flag that we're a completed `Sender` and try to wake up a receiver.
        // Whether or not we actually stored any data will get picked up and
        // translated to either an item or cancellation.
        //
        // Note that if we fail to acquire the `rx_task` lock then that means
        // we're in one of two situations:
        //
        // 1. The receiver is trying to block in `poll`
        // 2. The receiver is being dropped
        //
        // In the first case it'll check the `complete` flag after it's done
        // blocking to see if it succeeded. In the latter case we don't need to
        // wake up anyone anyway. So in both cases it's ok to ignore the `None`
        // case of `try_lock` and bail out.
        //
        // The first case crucially depends on `Lock` using `SeqCst` ordering
        // under the hood. If it instead used `Release` / `Acquire` ordering,
        // then it would not necessarily synchronize with `inner.complete`
        // and deadlock might be possible, as was observed in
        // https://github.com/rust-lang-nursery/futures-rs/pull/219.
        self.complete.store(true, SeqCst);
        if let Some(mut slot) = self.rx_task.try_lock() {
            if let Some(task) = slot.take() {
                drop(slot);
                task.wake();
            }
        }
    }

    fn close_rx(&self) {
        // Flag our completion and then attempt to wake up the sender if it's
        // blocked. See comments in `drop` below for more info
        self.complete.store(true, SeqCst);
        if let Some(mut handle) = self.tx_task.try_lock() {
            if let Some(task) = handle.take() {
                drop(handle);
                task.wake()
            }
        }
    }

    fn try_recv(&self) -> Result<Option<T>, Canceled> {
        // If we're complete, either `::close_rx` or `::drop_tx` was called.
        // We can assume a successful send if data is present.
        if self.complete.load(SeqCst) {
            if let Some(mut slot) = self.data.try_lock() {
                if let Some(data) = slot.take() {
                    return Ok(Some(data));
                }
            }
            Err(Canceled)
        } else {
            Ok(None)
        }
    }

    fn recv(&self, cx: &mut task::Context) -> Poll<Result<T, Canceled>> {
        // Check to see if some data has arrived. If it hasn't then we need to
        // block our task.
        //
        // Note that the acquisition of the `rx_task` lock might fail below, but
        // the only situation where this can happen is during `Sender::drop`
        // when we are indeed completed already. If that's happening then we
        // know we're completed so keep going.
        let done = if self.complete.load(SeqCst) {
            true
        } else {
            let task = cx.waker().clone();
            match self.rx_task.try_lock() {
                Some(mut slot) => { *slot = Some(task); false },
                None => true,
            }
        };

        // If we're `done` via one of the paths above, then look at the data and
        // figure out what the answer is. If, however, we stored `rx_task`
        // successfully above we need to check again if we're completed in case
        // a message was sent while `rx_task` was locked and couldn't notify us
        // otherwise.
        //
        // If we're not done, and we're not complete, though, then we've
        // successfully blocked our task and we return `Pending`.
        if done || self.complete.load(SeqCst) {
            // If taking the lock fails, the sender will realise that the we're
            // `done` when it checks the `complete` flag on the way out, and will
            // treat the send as a failure.
            if let Some(mut slot) = self.data.try_lock() {
                if let Some(data) = slot.take() {
                    return Poll::Ready(Ok(data));
                }
            }
            Poll::Ready(Err(Canceled))
        } else {
            Poll::Pending
        }
    }

    fn drop_rx(&self) {
        // Indicate to the `Sender` that we're done, so any future calls to
        // `poll_cancel` are weeded out.
        self.complete.store(true, SeqCst);

        // If we've blocked a task then there's no need for it to stick around,
        // so we need to drop it. If this lock acquisition fails, though, then
        // it's just because our `Sender` is trying to take the task, so we
        // let them take care of that.
        if let Some(mut slot) = self.rx_task.try_lock() {
            let task = slot.take();
            drop(slot);
            drop(task);
        }

        // Finally, if our `Sender` wants to get notified of us going away, it
        // would have stored something in `tx_task`. Here we try to peel that
        // out and unpark it.
        //
        // Note that the `try_lock` here may fail, but only if the `Sender` is
        // in the process of filling in the task. If that happens then we
        // already flagged `complete` and they'll pick that up above.
        if let Some(mut handle) = self.tx_task.try_lock() {
            if let Some(task) = handle.take() {
                drop(handle);
                task.wake()
            }
        }
    }
}

impl<T> Sender<T> {
    /// Completes this oneshot with a successful result.
    ///
    /// This function will consume `self` and indicate to the other end, the
    /// [`Receiver`](Receiver), that the value provided is the result of the
    /// computation this represents.
    ///
    /// If the value is successfully enqueued for the remote end to receive,
    /// then `Ok(())` is returned. If the receiving end was dropped before
    /// this function was called, however, then `Err` is returned with the value
    /// provided.
    pub fn send(self, t: T) -> Result<(), T> {
        self.inner.send(t)
    }

    /// Polls this `Sender` half to detect whether its associated
    /// [`Receiver`](Receiver) with has been dropped.
    ///
    /// # Return values
    ///
    /// If `Ok(Ready)` is returned then the associated `Receiver` has been
    /// dropped, which means any work required for sending should be canceled.
    ///
    /// If `Ok(Pending)` is returned then the associated `Receiver` is still
    /// alive and may be able to receive a message if sent. The current task,
    /// however, is scheduled to receive a notification if the corresponding
    /// `Receiver` goes away.
    pub fn poll_cancel(&mut self, cx: &mut task::Context) -> Poll<()> {
        self.inner.poll_cancel(cx)
    }

    /// Tests to see whether this `Sender`'s corresponding `Receiver`
    /// has been dropped.
    ///
    /// Unlike [`poll_cancel`](Sender::poll_cancel), this function does not
    /// enqueue a task for wakeup upon cancellation, but merely reports the
    /// current state, which may be subject to concurrent modification.
    pub fn is_canceled(&self) -> bool {
        self.inner.is_canceled()
    }
}

impl<T> Drop for Sender<T> {
    fn drop(&mut self) {
        self.inner.drop_tx()
    }
}

/// Error returned from a [`Receiver`](Receiver) when the corresponding
/// [`Sender`](Sender) is dropped.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct Canceled;

impl fmt::Display for Canceled {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "oneshot canceled")
    }
}

impl Error for Canceled {
    fn description(&self) -> &str {
        "oneshot canceled"
    }
}

impl<T> Receiver<T> {
    /// Gracefully close this receiver, preventing any subsequent attempts to
    /// send to it.
    ///
    /// Any `send` operation which happens after this method returns is
    /// guaranteed to fail. After calling this method, you can use
    /// [`Receiver::poll`](Receiver::poll) to determine whether a message had
    /// previously been sent.
    pub fn close(&mut self) {
        self.inner.close_rx()
    }

    /// Attempts to receive a message outside of the context of a task.
    ///
    /// Useful when a [`Context`](Context) is not available such as within a
    /// `Drop` impl.
    ///
    /// Does not schedule a task wakeup or have any other side effects.
    ///
    /// A return value of `None` must be considered immediately stale (out of
    /// date) unless [`::close`](Receiver::close) has been called first.
    ///
    /// Returns an error if the sender was dropped.
    pub fn try_recv(&mut self) -> Result<Option<T>, Canceled> {
        self.inner.try_recv()
    }
}

impl<T> Future for Receiver<T> {
    type Output = Result<T, Canceled>;

    fn poll(self: PinMut<Self>, cx: &mut task::Context) -> Poll<Result<T, Canceled>> {
        self.inner.recv(cx)
    }
}

impl<T> Drop for Receiver<T> {
    fn drop(&mut self) {
        self.inner.drop_rx()
    }
}
<|endoftext|>