<commit_before>// Copyright 2015 The noise-rs developers. For a full listing of the authors,
// refer to the AUTHORS file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![allow(unstable)]

use std::rand::{Rand, Rng, SeedableRng, XorShiftRng};
use std::num::SignedInt;

use math;

const TABLE_SIZE: usize = 256;

#[allow(missing_copy_implementations)]
pub struct Seed {
    values: [usize; TABLE_SIZE * 2],
}

impl Rand for Seed {
    fn rand<R: Rng>(rng: &mut R) -> Seed {
        let mut seq: Vec<usize> = (0..TABLE_SIZE).collect();
        rng.shuffle(&mut *seq);

        // It's unfortunate that this double-initializes the array, but Rust doesn't currently provide a
        // clean way to do this in one pass. Hopefully won't matter, as Seed creation will usually be a
        // one-time event.
        let mut seed = Seed { values: [0; TABLE_SIZE * 2] };
        let seq_it = seq.iter().cycle();
        for (x, y) in seed.values.iter_mut().zip(seq_it) { *x = *y }
        seed
    }
}

impl Seed {
    pub fn new(seed: u32) -> Seed {
        let mut rng: XorShiftRng = SeedableRng::from_seed([1, seed, seed, seed]);
        rng.gen()
    }

    #[inline(always)]
    pub fn get1<T: SignedInt>(&self, x: T) -> usize {
        self.values[math::cast::<T, usize>(x & math::cast(TABLE_SIZE - 1))]
    }

    #[inline(always)]
    pub fn get2<T: SignedInt>(&self, pos: math::Point2<T>) -> usize {
        self.values[self.get1(pos[0]) + math::cast(pos[1] & math::cast(TABLE_SIZE - 1))]
    }

    #[inline(always)]
    pub fn get3<T: SignedInt>(&self, pos: math::Point3<T>) -> usize {
        self.values[self.get2([pos[0], pos[1]]) + math::cast(pos[2] & math::cast(TABLE_SIZE - 1))]
    }

    #[inline(always)]
    pub fn get4<T: SignedInt>(&self, pos: math::Point4<T>) -> usize {
        self.values[self.get3([pos[0], pos[1], pos[2]]) + math::cast(pos[3] & math::cast(TABLE_SIZE - 1))]
    }
}

#[cfg(test)]
mod tests {
    use std::rand::random;
    use perlin::perlin3;
    use super::Seed;

    #[test]
    fn test_random_seed() {
        let _ = perlin3::<f32>(&random(), &[1.0, 2.0, 3.0]);
    }

    #[test]
    fn test_negative_params() {
        let _ = perlin3::<f32>(&Seed::new(0), &[-1.0, 2.0, 3.0]);
    }
}
<commit_msg>Use separate permutation tables for each coordinate<commit_after>// Copyright 2015 The noise-rs developers. For a full listing of the authors,
// refer to the AUTHORS file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![allow(unstable)]

use std::rand::{Rand, Rng, SeedableRng, XorShiftRng};
use std::num::SignedInt;

use math;

const TABLE_SIZE: usize = 256;

#[allow(missing_copy_implementations)]
pub struct Seed {
    x_values: [u8; TABLE_SIZE],
    y_values: [u8; TABLE_SIZE],
    z_values: [u8; TABLE_SIZE],
    w_values: [u8; TABLE_SIZE],
}

impl Rand for Seed {
    fn rand<R: Rng>(rng: &mut R) -> Seed {
        let mut x_values: Vec<u8> = ::std::iter::range_inclusive(0, (TABLE_SIZE - 1) as u8).collect();
        let mut y_values: Vec<u8> = x_values.clone();
        let mut z_values: Vec<u8> = x_values.clone();
        let mut w_values: Vec<u8> = x_values.clone();

        rng.shuffle(&mut *x_values);
        rng.shuffle(&mut *y_values);
        rng.shuffle(&mut *z_values);
        rng.shuffle(&mut *w_values);

        // It's unfortunate that this double-initializes the array, but Rust doesn't currently provide a
        // clean way to do this in one pass. Hopefully won't matter, as Seed creation will usually be a
        // one-time event.
        let mut seed = Seed {
            x_values: [0; TABLE_SIZE],
            y_values: [0; TABLE_SIZE],
            z_values: [0; TABLE_SIZE],
            w_values: [0; TABLE_SIZE],
        };

        let x_iter = x_values.iter().cycle();
        for (x, y) in seed.x_values.iter_mut().zip(x_iter) { *x = *y }
        let y_iter = y_values.iter().cycle();
        for (x, y) in seed.y_values.iter_mut().zip(y_iter) { *x = *y }
        let z_iter = z_values.iter().cycle();
        for (x, y) in seed.z_values.iter_mut().zip(z_iter) { *x = *y }
        let w_iter = w_values.iter().cycle();
        for (x, y) in seed.w_values.iter_mut().zip(w_iter) { *x = *y }
        seed
    }
}

impl Seed {
    pub fn new(seed: u32) -> Seed {
        let mut rng: XorShiftRng = SeedableRng::from_seed([1, seed, seed, seed]);
        rng.gen()
    }

    #[inline(always)]
    pub fn get1<T: SignedInt>(&self, x: T) -> usize {
        self.x_values[math::cast::<T, usize>(x & math::cast(TABLE_SIZE - 1))] as usize
    }

    #[inline(always)]
    pub fn get2<T: SignedInt>(&self, pos: math::Point2<T>) -> usize {
        self.get1(pos[0]) ^ self.y_values[math::cast::<T, usize>(pos[1] & math::cast(TABLE_SIZE - 1))] as usize
    }

    #[inline(always)]
    pub fn get3<T: SignedInt>(&self, pos: math::Point3<T>) -> usize {
        self.get2([pos[0], pos[1]]) ^ self.z_values[math::cast::<T, usize>(pos[2] & math::cast(TABLE_SIZE - 1))] as usize
    }

    #[inline(always)]
    pub fn get4<T: SignedInt>(&self, pos: math::Point4<T>) -> usize {
        self.get3([pos[0], pos[1], pos[2]]) ^ self.w_values[math::cast::<T, usize>(pos[3] & math::cast(TABLE_SIZE - 1))] as usize
    }
}

#[cfg(test)]
mod tests {
    use std::rand::random;
    use perlin::perlin3;
    use super::Seed;

    #[test]
    fn test_random_seed() {
        let _ = perlin3::<f32>(&random(), &[1.0, 2.0, 3.0]);
    }

    #[test]
    fn test_negative_params() {
        let _ = perlin3::<f32>(&Seed::new(0), &[-1.0, 2.0, 3.0]);
    }
}
<|endoftext|>


<commit_before>// Semantics analyzer. This pass plays a few important roles as shown
// below:
//
// - Add types to nodes. For example, a tree that represents "1+2" is
//   typed as INT because the result type of an addition of two
//   integers is integer.
//
// - Resolve variable names based on the C scope rules.
//   Local variables are resolved to offsets from the base pointer.
//   Global variables are resolved to their names.
//
// - Insert nodes to make array-to-pointer conversion explicit.
//   Recall that, in C, "array of T" is automatically converted to
//   "pointer to T" in most contexts.
//
// - Reject bad assignments, such as `1=2+3`.

use crate::parse::*;
use crate::util::*;
use std::collections::HashMap;
use std::mem;
use std::sync::Mutex;

lazy_static! {
    static ref ENV: Mutex<Env> = Mutex::new(new_env(None));
    static ref STR_LABEL: Mutex<i32> = Mutex::new(0);
    static ref GLOBALS: Mutex<Vec<Var>> = Mutex::new(Vec::new());
    static ref STACKSIZE: Mutex<i32> = Mutex::new(0);
}

#[derive(Clone, Debug)]
struct Env {
    vars: HashMap<String, Var>,
    next: Option<Box<Env>>,
}

fn new_env(next: Option<Env>) -> Env {
    Env {
        vars: HashMap::new(),
        next: if let Some(e) = next {
            Some(Box::new(e))
        } else {
            None
        },
    }
}

fn env_var_put(name: String, var: Var) {
    match ENV.lock() {
        Ok(ref mut env) => {
            env.vars.insert(name, var);
        }
        Err(_) => {
            panic!();
        }
    }
}

fn env_push() {
    let mut env = ENV.lock().unwrap();
    *env = new_env(Some(env.clone()));
}

fn env_pop() {
    let mut env = ENV.lock().unwrap();
    *env = *env.next.clone().unwrap();
}

fn new_int(val: i32) -> Node {
    let mut node = alloc_node();
    node.op = NodeType::NUM;
    node.ty = int_ty();
    node.val = val;
    return node;
}

//impl<'a> Env<'a> {
impl Env {
    fn find(&self, name: & String) -> Option<&Var> {
        if let Some(v) = self.vars.get(name) {
            return Some(v);
        } else if let Some(ref next) = self.next {
            return next.find(name);
        } else {
            return None;
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct Var {
    pub ty: Type,
    pub is_local: bool,

    // local
    pub offset: i32,

    // global
    pub name: String,
    pub is_extern: bool,
    pub data: String,
    pub len: usize,
}

fn alloc_var(ty: Type) -> Var {
    Var {
        ty: ty,
        is_local: false,
        offset: 0,
        name: String::new(),
        is_extern: false,
        data: String::new(),
        len: 0,
    }
}

fn new_global(ty: Type, name: String, data: String, len: usize) -> Var {
    let mut var = alloc_var(ty);
    var.is_local = false;
    //var.name = format!(".L.str{}", bump_str_label()).to_string();
    var.name = name;
    var.data = data;
    var.len = len;
    return var;
}

fn init_globals() {
    *GLOBALS.lock().unwrap() = Vec::new();
}

fn globals() -> Vec<Var> {
    GLOBALS.lock().unwrap().clone()
}

fn globals_push(var: Var) {
    let mut globals = GLOBALS.lock().unwrap();
    globals.push(var);
}

fn stacksize() -> i32 {
    *STACKSIZE.lock().unwrap()
}

fn init_stacksize() {
    let mut stacksize = STACKSIZE.lock().unwrap();
    *stacksize = 0;
}

fn set_stacksize(i: i32) {
    let mut stacksize = STACKSIZE.lock().unwrap();
    *stacksize = i;
}

fn add_stacksize(i: i32) {
    let mut stacksize = STACKSIZE.lock().unwrap();
    *stacksize += i;
}

fn bump_str_label() -> i32 {
    let mut str_label = STR_LABEL.lock().unwrap();
    let ret = *str_label;
    *str_label += 1;
    return ret;
}

fn find_var(name: String) -> Option<Var> {
    return match ENV.lock() {
        Ok(ref env) => {
            match env.find(&name) {
                Some(v) => Some(v.clone()),
                None => None,
            }
        }
        Err(_) => {
            panic!();
        }
    }
}

fn swap(p: &mut Node, q: &mut Node) {
    mem::swap(p, q);
}

fn maybe_decay(base: & mut Node, decay: bool) -> & mut Node {
    if !decay || base.ty.ty != CType::ARY {
        return base;
    }

    let tmp = base.clone();
    *base = alloc_node();
    base.op = NodeType::ADDR;
    base.ty = ptr_to(*tmp.clone().ty.ary_of.unwrap());

    base.expr = Some(Box::new(tmp));
    return base;
}

fn check_lval(node: & Node) {
    match node.op {
        NodeType::LVAR | NodeType::GVAR | NodeType::DEREF | NodeType::DOT => {
            return;
        }
        _ => {
            panic!("not an lvalue: {:?} ({})", node.op, node.name);
        }
    }
}

fn walk<'a>(node: &'a mut Node, decay: bool) -> &'a Node {
    match node.op {
        NodeType::NUM => {
            return node;
        }
        NodeType::STR => {
            let node_ty = node.ty.clone();
            let node_data = node.data.clone();

            let name = format!(".L.str{}", bump_str_label()).to_string();
            let var = new_global(node_ty.clone(), name, node_data, node.len);
            globals_push(var.clone());

            *node = alloc_node();
            node.op = NodeType::GVAR;
            node.ty = node_ty;
            node.name = var.name;
            return maybe_decay(node, decay);
        }
        NodeType::IDENT => {
            let var = match find_var(node.name.clone()) {
                None => panic!("undefined variable: {}", node.name),
                Some(var) => var,
            };

            if var.is_local {
                *node = alloc_node();
                node.op = NodeType::LVAR;
                node.ty = var.ty.clone();
                node.offset = var.offset;
                return maybe_decay(node, decay);
            }

            *node = alloc_node();
            node.op = NodeType::GVAR;
            node.ty = var.ty.clone();
            node.name = var.name.clone();
            return maybe_decay(node, decay);
        }
        NodeType::VARDEF => {
            set_stacksize(roundup(stacksize(), node.ty.align));
            add_stacksize(node.ty.size);
            node.offset = stacksize();

            let mut var = alloc_var(node.ty.clone());
            var.is_local = true;
            var.offset = stacksize();
            env_var_put(node.name.clone(), var.clone());

            match node.init {
                Some(_) => {
                    node.init = Some(Box::new(walk(&mut *node.init.clone().unwrap(), true).clone()));
                }
                None => { }
            }
            return node;
        }
        NodeType::IF => {
            node.cond = Some(Box::new(walk(&mut *node.cond.clone().unwrap(), true).clone()));
            node.then = Some(Box::new(walk(&mut *node.then.clone().unwrap(), true).clone()));

            if node.els.is_some() {
                node.els = Some(Box::new(walk(&mut *node.els.clone().unwrap(), true).clone()));
            }
            return node;
        }
        NodeType::FOR => {
            node.init = Some(Box::new(walk(&mut *node.init.clone().unwrap(), true).clone()));
            node.cond = Some(Box::new(walk(&mut *node.cond.clone().unwrap(), true).clone()));
            node.inc = Some(Box::new(walk(&mut *node.inc.clone().unwrap(), true).clone()));
            node.body = Some(Box::new(walk(&mut *node.body.clone().unwrap(), true).clone()));
            return node;
        }
        NodeType::DO_WHILE => {
            node.cond = Some(Box::new(walk(&mut *node.cond.clone().unwrap(), true).clone()));
            node.body = Some(Box::new(walk(&mut *node.body.clone().unwrap(), true).clone()));
            return node;
        }
        NodeType::ADD | NodeType::SUB => {
            node.lhs = Some(Box::new(walk(&mut *node.lhs.clone().unwrap(), true).clone()));
            node.rhs = Some(Box::new(walk(&mut *node.rhs.clone().unwrap(), true).clone()));

            match (&mut node.lhs, &mut node.rhs) {
                (Some(ref mut lhs), Some(ref mut rhs)) => {
                    if rhs.ty.ty == CType::PTR {
                        swap(lhs, rhs);
                    }
                }
                _ => {}
            }
            match node.rhs {
                Some(ref rhs) => {
                    if rhs.ty.ty == CType::PTR {
                        panic!("'pointer {:?} pointer' is not defined", node.op);
                    }
                }
                None => {
                }
            }

            match node.lhs {
                Some(ref lhs) => {
                    node.ty = lhs.ty.clone();
                }
                None => { }
            }

            return node;
        }
        NodeType::EQL => {
            node.lhs = Some(Box::new(walk(&mut *node.lhs.clone().unwrap(), false).clone()));
            check_lval(& *node.lhs.clone().unwrap());

            node.rhs = Some(Box::new(walk(&mut *node.rhs.clone().unwrap(), true).clone()));
            node.ty = node.lhs.clone().unwrap().ty;
            return node;
        }
        NodeType::DOT => {
            node.expr = Some(Box::new(walk(&mut *node.expr.clone().unwrap(), true).clone()));
            if node.expr.clone().unwrap().ty.ty != CType::STRUCT {
                panic!("struct expected before '.'");
            }

            let ty = node.expr.clone().unwrap().ty;
            if ty.members == None {
                panic!("incomplete type: {:?}", node.expr);
            }

            for m in ty.members.unwrap().iter() {
                if m.name != node.name {
                    continue;
                }

                node.ty = m.ty.clone();
                node.offset = m.ty.offset;
                return maybe_decay(node, decay);
            }
            panic!("member missing: {}", node.name);
        }
        NodeType::QUEST => {
            node.cond = Some(Box::new(walk(&mut *node.cond.clone().unwrap(), true).clone()));
            node.then = Some(Box::new(walk(&mut *node.then.clone().unwrap(), true).clone()));
            node.els = Some(Box::new(walk(&mut *node.els.clone().unwrap(), true).clone()));
            node.ty = node.then.clone().unwrap().ty;
            return node;
        }
        NodeType::MUL |
            NodeType::DIV |
            NodeType::MOD |
            NodeType::LT |
            NodeType::OR |
            NodeType::XOR |
            NodeType::AND |
            NodeType::EQ |
            NodeType::NE |
            NodeType::LE |
            NodeType::SHL |
            NodeType::SHR |
            NodeType::LOGAND |
            NodeType::LOGOR => {
                node.lhs = Some(Box::new(walk(&mut *node.lhs.clone().unwrap(), true).clone()));
                node.rhs = Some(Box::new(walk(&mut *node.rhs.clone().unwrap(), true).clone()));

                match node.lhs {
                    Some(ref lhs) => {
                        node.ty = lhs.ty.clone();
                    }
                    None => { }
                }
                return node;
            }
        NodeType::COMMA => {
            node.lhs = Some(Box::new(walk(&mut *node.lhs.clone().unwrap(), true).clone()));
            node.rhs = Some(Box::new(walk(&mut *node.rhs.clone().unwrap(), true).clone()));
            node.ty = node.rhs.clone().unwrap().ty;
            return node;
        }
        NodeType::PRE_INC |
            NodeType::PRE_DEC |
            NodeType::POST_INC |
            NodeType::POST_DEC |
            NodeType::NEG |
            NodeType::EXCLAM => {
            node.expr = Some(Box::new(walk(&mut *node.expr.clone().unwrap(), true).clone()));
            node.ty = node.expr.clone().unwrap().ty;
            return node;
        }
        NodeType::ADDR => {
            node.expr = Some(Box::new(walk(&mut *node.expr.clone().unwrap(), true).clone()));
            check_lval(& *node.expr.clone().unwrap());
            node.ty = ptr_to(node.expr.clone().unwrap().ty);
            return node;
        }
        NodeType::DEREF => {
            node.expr = Some(Box::new(walk(&mut *node.expr.clone().unwrap(), true).clone()));

            match &mut node.expr {
                Some(ref expr) => {
                    if expr.ty.ty != CType::PTR {
                        panic!("operand must be a pointer: {:?}", expr);
                    }

                    if expr.ty.ptr_to.clone().unwrap().ty == CType::VOID {
                        panic!("operand dereference void pointer");
                    }
                }
                None => {}
            }

            node.ty = *node.expr.clone().unwrap().ty.ptr_to.unwrap();
            return node;
        }
        NodeType::RETURN |
            NodeType::EXPR_STMT => {
            node.expr = Some(Box::new(walk(&mut *node.expr.clone().unwrap(), true).clone()));
            return node;
        }
        NodeType::SIZEOF => {
            let mut nexpr = *node.expr.clone().unwrap();
            let expr = walk(&mut nexpr, false);
            let val = expr.ty.size;

            *node = new_int(val);
            return node;
        }
        NodeType::ALIGNOF => {
            let mut e = node.expr.clone().unwrap();
            let expr = walk(&mut e, false);
            *node = new_int(expr.ty.align);
            return node;
        }
        NodeType::CALL => {
            for i in 0..node.args.len() {
                node.args[i] = walk(&mut node.args[i], true).clone();
            }
            node.ty = int_ty();
            return node;
        }
        NodeType::FUNC => {
            for i in 0..node.args.len() {
                node.args[i] = walk(&mut node.args[i], true).clone();
            }
            node.body = Some(Box::new(walk(&mut *node.body.clone().unwrap(), true).clone()));
            return node;
        }
        NodeType::COMP_STMT => {
            env_push();
            for i in 0..node.stmts.len() {
                node.stmts[i] = walk(&mut node.stmts[i], true).clone();
            }
            env_pop();
            return node;
        }
        NodeType::STMT_EXPR => {
            node.body = Some(Box::new(walk(&mut *node.body.clone().unwrap(), true).clone()));
            node.ty = int_ty();
            return node;
        }
        NodeType::NULL => {
            return node;
        }
        _ => {
            panic!("unknown node type");
        }
    }
}

pub fn sema(nodes: &mut Vec<Node>) -> Vec<Var> {
    init_globals();
    //let mut topenv = new_env(None);

    for node in nodes.iter_mut() {
        if node.op == NodeType::VARDEF {
            let mut var = new_global(node.ty.clone(), node.name.clone(), node.data.clone(), node.len);
            var.is_extern = node.is_extern;
            globals_push(var.clone());
            env_var_put(node.name.clone(), var);
            continue;
        }

        assert!(node.op == NodeType::FUNC);

        init_stacksize();
        walk(node, true);
        node.stacksize = stacksize();
    }
    return globals();
}
<commit_msg>Simplify.<commit_after>// Semantics analyzer. This pass plays a few important roles as shown
// below:
//
// - Add types to nodes. For example, a tree that represents "1+2" is
//   typed as INT because the result type of an addition of two
//   integers is integer.
//
// - Resolve variable names based on the C scope rules.
//   Local variables are resolved to offsets from the base pointer.
//   Global variables are resolved to their names.
//
// - Insert nodes to make array-to-pointer conversion explicit.
//   Recall that, in C, "array of T" is automatically converted to
//   "pointer to T" in most contexts.
//
// - Reject bad assignments, such as `1=2+3`.

use crate::parse::*;
use crate::util::*;
use std::collections::HashMap;
use std::mem;
use std::sync::Mutex;

lazy_static! {
    static ref ENV: Mutex<Env> = Mutex::new(new_env(None));
    static ref STR_LABEL: Mutex<i32> = Mutex::new(0);
    static ref GLOBALS: Mutex<Vec<Var>> = Mutex::new(Vec::new());
    static ref STACKSIZE: Mutex<i32> = Mutex::new(0);
}

#[derive(Clone, Debug)]
struct Env {
    vars: HashMap<String, Var>,
    next: Option<Box<Env>>,
}

fn new_env(next: Option<Env>) -> Env {
    Env {
        vars: HashMap::new(),
        next: if let Some(e) = next {
            Some(Box::new(e))
        } else {
            None
        },
    }
}

fn env_var_put(name: String, var: Var) {
    match ENV.lock() {
        Ok(ref mut env) => {
            env.vars.insert(name, var);
        }
        Err(_) => {
            panic!();
        }
    }
}

fn env_push() {
    let mut env = ENV.lock().unwrap();
    *env = new_env(Some(env.clone()));
}

fn env_pop() {
    let mut env = ENV.lock().unwrap();
    *env = *env.next.clone().unwrap();
}

fn new_int(val: i32) -> Node {
    let mut node = alloc_node();
    node.op = NodeType::NUM;
    node.ty = int_ty();
    node.val = val;
    return node;
}

//impl<'a> Env<'a> {
impl Env {
    fn find(&self, name: & String) -> Option<&Var> {
        if let Some(v) = self.vars.get(name) {
            return Some(v);
        } else if let Some(ref next) = self.next {
            return next.find(name);
        } else {
            return None;
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct Var {
    pub ty: Type,
    pub is_local: bool,

    // local
    pub offset: i32,

    // global
    pub name: String,
    pub is_extern: bool,
    pub data: String,
    pub len: usize,
}

fn alloc_var(ty: Type) -> Var {
    Var {
        ty: ty,
        is_local: false,
        offset: 0,
        name: String::new(),
        is_extern: false,
        data: String::new(),
        len: 0,
    }
}

fn new_global(ty: Type, name: String, data: String, len: usize) -> Var {
    let mut var = alloc_var(ty);
    var.is_local = false;
    //var.name = format!(".L.str{}", bump_str_label()).to_string();
    var.name = name;
    var.data = data;
    var.len = len;
    return var;
}

fn init_globals() {
    *GLOBALS.lock().unwrap() = Vec::new();
}

fn globals() -> Vec<Var> {
    GLOBALS.lock().unwrap().clone()
}

fn globals_push(var: Var) {
    let mut globals = GLOBALS.lock().unwrap();
    globals.push(var);
}

fn stacksize() -> i32 {
    *STACKSIZE.lock().unwrap()
}

fn init_stacksize() {
    let mut stacksize = STACKSIZE.lock().unwrap();
    *stacksize = 0;
}

fn set_stacksize(i: i32) {
    let mut stacksize = STACKSIZE.lock().unwrap();
    *stacksize = i;
}

fn add_stacksize(i: i32) {
    let mut stacksize = STACKSIZE.lock().unwrap();
    *stacksize += i;
}

fn bump_str_label() -> i32 {
    let mut str_label = STR_LABEL.lock().unwrap();
    let ret = *str_label;
    *str_label += 1;
    return ret;
}

fn find_var(name: String) -> Option<Var> {
    return match ENV.lock() {
        Ok(ref env) => {
            match env.find(&name) {
                Some(v) => Some(v.clone()),
                None => None,
            }
        }
        Err(_) => {
            panic!();
        }
    }
}

fn swap(p: &mut Node, q: &mut Node) {
    mem::swap(p, q);
}

fn maybe_decay(base: & mut Node, decay: bool) -> & mut Node {
    if !decay || base.ty.ty != CType::ARY {
        return base;
    }

    let tmp = base.clone();
    *base = alloc_node();
    base.op = NodeType::ADDR;
    base.ty = ptr_to(*tmp.clone().ty.ary_of.unwrap());

    base.expr = Some(Box::new(tmp));
    return base;
}

fn check_lval(node: & Node) {
    match node.op {
        NodeType::LVAR | NodeType::GVAR | NodeType::DEREF | NodeType::DOT => { }
        _ => {
            panic!("not an lvalue: {:?} ({})", node.op, node.name);
        }
    }
}

fn walk<'a>(node: &'a mut Node, decay: bool) -> &'a Node {
    match node.op {
        NodeType::NUM |
            NodeType::NULL => {
            return node;
        }
        NodeType::STR => {
            let node_ty = node.ty.clone();
            let node_data = node.data.clone();

            let name = format!(".L.str{}", bump_str_label()).to_string();
            let var = new_global(node_ty.clone(), name, node_data, node.len);
            globals_push(var.clone());

            *node = alloc_node();
            node.op = NodeType::GVAR;
            node.ty = node_ty;
            node.name = var.name;
            return maybe_decay(node, decay);
        }
        NodeType::IDENT => {
            let var = match find_var(node.name.clone()) {
                None => panic!("undefined variable: {}", node.name),
                Some(var) => var,
            };

            if var.is_local {
                *node = alloc_node();
                node.op = NodeType::LVAR;
                node.ty = var.ty.clone();
                node.offset = var.offset;
                return maybe_decay(node, decay);
            }

            *node = alloc_node();
            node.op = NodeType::GVAR;
            node.ty = var.ty.clone();
            node.name = var.name.clone();
            return maybe_decay(node, decay);
        }
        NodeType::VARDEF => {
            set_stacksize(roundup(stacksize(), node.ty.align));
            add_stacksize(node.ty.size);
            node.offset = stacksize();

            let mut var = alloc_var(node.ty.clone());
            var.is_local = true;
            var.offset = stacksize();
            env_var_put(node.name.clone(), var.clone());

            match node.init {
                Some(_) => {
                    node.init = Some(Box::new(walk(&mut *node.init.clone().unwrap(), true).clone()));
                }
                None => { }
            }
            return node;
        }
        NodeType::IF => {
            node.cond = Some(Box::new(walk(&mut *node.cond.clone().unwrap(), true).clone()));
            node.then = Some(Box::new(walk(&mut *node.then.clone().unwrap(), true).clone()));

            if node.els.is_some() {
                node.els = Some(Box::new(walk(&mut *node.els.clone().unwrap(), true).clone()));
            }
            return node;
        }
        NodeType::FOR => {
            node.init = Some(Box::new(walk(&mut *node.init.clone().unwrap(), true).clone()));
            node.cond = Some(Box::new(walk(&mut *node.cond.clone().unwrap(), true).clone()));
            node.inc = Some(Box::new(walk(&mut *node.inc.clone().unwrap(), true).clone()));
            node.body = Some(Box::new(walk(&mut *node.body.clone().unwrap(), true).clone()));
            return node;
        }
        NodeType::DO_WHILE => {
            node.cond = Some(Box::new(walk(&mut *node.cond.clone().unwrap(), true).clone()));
            node.body = Some(Box::new(walk(&mut *node.body.clone().unwrap(), true).clone()));
            return node;
        }
        NodeType::ADD | NodeType::SUB => {
            node.lhs = Some(Box::new(walk(&mut *node.lhs.clone().unwrap(), true).clone()));
            node.rhs = Some(Box::new(walk(&mut *node.rhs.clone().unwrap(), true).clone()));

            match (&mut node.lhs, &mut node.rhs) {
                (Some(ref mut lhs), Some(ref mut rhs)) => {
                    if rhs.ty.ty == CType::PTR {
                        swap(lhs, rhs);
                    }
                }
                _ => {}
            }
            match node.rhs {
                Some(ref rhs) => {
                    if rhs.ty.ty == CType::PTR {
                        panic!("'pointer {:?} pointer' is not defined", node.op);
                    }
                }
                None => {
                }
            }

            match node.lhs {
                Some(ref lhs) => {
                    node.ty = lhs.ty.clone();
                }
                None => { }
            }

            return node;
        }
        NodeType::EQL => {
            node.lhs = Some(Box::new(walk(&mut *node.lhs.clone().unwrap(), false).clone()));
            check_lval(& *node.lhs.clone().unwrap());

            node.rhs = Some(Box::new(walk(&mut *node.rhs.clone().unwrap(), true).clone()));
            node.ty = node.lhs.clone().unwrap().ty;
            return node;
        }
        NodeType::DOT => {
            node.expr = Some(Box::new(walk(&mut *node.expr.clone().unwrap(), true).clone()));
            if node.expr.clone().unwrap().ty.ty != CType::STRUCT {
                panic!("struct expected before '.'");
            }

            let ty = node.expr.clone().unwrap().ty;
            if ty.members == None {
                panic!("incomplete type: {:?}", node.expr);
            }

            for m in ty.members.unwrap().iter() {
                if m.name != node.name {
                    continue;
                }

                node.ty = m.ty.clone();
                node.offset = m.ty.offset;
                return maybe_decay(node, decay);
            }
            panic!("member missing: {}", node.name);
        }
        NodeType::QUEST => {
            node.cond = Some(Box::new(walk(&mut *node.cond.clone().unwrap(), true).clone()));
            node.then = Some(Box::new(walk(&mut *node.then.clone().unwrap(), true).clone()));
            node.els = Some(Box::new(walk(&mut *node.els.clone().unwrap(), true).clone()));
            node.ty = node.then.clone().unwrap().ty;
            return node;
        }
        NodeType::MUL |
            NodeType::DIV |
            NodeType::MOD |
            NodeType::LT |
            NodeType::OR |
            NodeType::XOR |
            NodeType::AND |
            NodeType::EQ |
            NodeType::NE |
            NodeType::LE |
            NodeType::SHL |
            NodeType::SHR |
            NodeType::LOGAND |
            NodeType::LOGOR => {
                node.lhs = Some(Box::new(walk(&mut *node.lhs.clone().unwrap(), true).clone()));
                node.rhs = Some(Box::new(walk(&mut *node.rhs.clone().unwrap(), true).clone()));

                match node.lhs {
                    Some(ref lhs) => {
                        node.ty = lhs.ty.clone();
                    }
                    None => { }
                }
                return node;
            }
        NodeType::COMMA => {
            node.lhs = Some(Box::new(walk(&mut *node.lhs.clone().unwrap(), true).clone()));
            node.rhs = Some(Box::new(walk(&mut *node.rhs.clone().unwrap(), true).clone()));
            node.ty = node.rhs.clone().unwrap().ty;
            return node;
        }
        NodeType::PRE_INC |
            NodeType::PRE_DEC |
            NodeType::POST_INC |
            NodeType::POST_DEC |
            NodeType::NEG |
            NodeType::EXCLAM => {
            node.expr = Some(Box::new(walk(&mut *node.expr.clone().unwrap(), true).clone()));
            node.ty = node.expr.clone().unwrap().ty;
            return node;
        }
        NodeType::ADDR => {
            node.expr = Some(Box::new(walk(&mut *node.expr.clone().unwrap(), true).clone()));
            check_lval(& *node.expr.clone().unwrap());
            node.ty = ptr_to(node.expr.clone().unwrap().ty);
            return node;
        }
        NodeType::DEREF => {
            node.expr = Some(Box::new(walk(&mut *node.expr.clone().unwrap(), true).clone()));

            match &mut node.expr {
                Some(ref expr) => {
                    if expr.ty.ty != CType::PTR {
                        panic!("operand must be a pointer: {:?}", expr);
                    }

                    if expr.ty.ptr_to.clone().unwrap().ty == CType::VOID {
                        panic!("operand dereference void pointer");
                    }
                }
                None => {}
            }

            node.ty = *node.expr.clone().unwrap().ty.ptr_to.unwrap();
            return node;
        }
        NodeType::RETURN |
            NodeType::EXPR_STMT => {
            node.expr = Some(Box::new(walk(&mut *node.expr.clone().unwrap(), true).clone()));
            return node;
        }
        NodeType::SIZEOF => {
            let mut nexpr = *node.expr.clone().unwrap();
            let expr = walk(&mut nexpr, false);
            let val = expr.ty.size;

            *node = new_int(val);
            return node;
        }
        NodeType::ALIGNOF => {
            let mut e = node.expr.clone().unwrap();
            let expr = walk(&mut e, false);
            *node = new_int(expr.ty.align);
            return node;
        }
        NodeType::CALL => {
            for i in 0..node.args.len() {
                node.args[i] = walk(&mut node.args[i], true).clone();
            }
            node.ty = int_ty();
            return node;
        }
        NodeType::FUNC => {
            for i in 0..node.args.len() {
                node.args[i] = walk(&mut node.args[i], true).clone();
            }
            node.body = Some(Box::new(walk(&mut *node.body.clone().unwrap(), true).clone()));
            return node;
        }
        NodeType::COMP_STMT => {
            env_push();
            for i in 0..node.stmts.len() {
                node.stmts[i] = walk(&mut node.stmts[i], true).clone();
            }
            env_pop();
            return node;
        }
        NodeType::STMT_EXPR => {
            node.body = Some(Box::new(walk(&mut *node.body.clone().unwrap(), true).clone()));
            node.ty = int_ty();
            return node;
        }
        _ => {
            panic!("unknown node type");
        }
    }
}

pub fn sema(nodes: &mut Vec<Node>) -> Vec<Var> {
    init_globals();
    //let mut topenv = new_env(None);

    for node in nodes.iter_mut() {
        if node.op == NodeType::VARDEF {
            let mut var = new_global(node.ty.clone(), node.name.clone(), node.data.clone(), node.len);
            var.is_extern = node.is_extern;
            globals_push(var.clone());
            env_var_put(node.name.clone(), var);
            continue;
        }

        assert!(node.op == NodeType::FUNC);

        init_stacksize();
        walk(node, true);
        node.stacksize = stacksize();
    }
    return globals();
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 Robert Clipsham <robert@octarineparrot.com>
//
// 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.

extern crate libc;

use std::clone::Clone;
use std::result::{Ok, Err};
use std::option::{Some};
use std::task::try_future;
use std::io::net::ip::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::iter::Iterator;

use datalink::{datalink_channel, DataLinkChannelType};
use packet::Packet;
use packet::ethernet::{EtherTypes, EthernetHeader, MutableEthernetHeader, EthernetPacket};
use packet::ip::{IpNextHeaderProtocols, IpNextHeaderProtocol};
use packet::ipv4::{Ipv4Header, MutableIpv4Header, Ipv4Packet};
use packet::ipv6::{MutableIpv6Header, Ipv6Packet};
use packet::udp::{UdpHeader, MutableUdpHeader, UdpPacket};
use transport::{udp_header_iter, ipv4_header_iter, transport_channel, TransportProtocol,
                TransportChannelType};
use transport::TransportProtocol::{Ipv4, Ipv6};
use util::NetworkInterface;

const MIN_PACKET_SIZE: uint = 64;
const ETHERNET_HEADER_LEN: uint = 14;
const IPV4_HEADER_LEN: uint = 20;
const IPV6_HEADER_LEN: uint = 40;
const UDP_HEADER_LEN: uint = 8;
const TEST_DATA_LEN: uint = 4;

const IPV4_SOURCE: IpAddr = Ipv4Addr(127, 0, 0, 1);
const IPV4_DESTINATION: IpAddr = Ipv4Addr(127, 0, 0, 1);
const IPV6_SOURCE: IpAddr = Ipv6Addr(0, 0, 0, 0, 0, 0, 0, 1);
const IPV6_DESTINATION: IpAddr = Ipv6Addr(0, 0, 0, 0, 0, 0, 0, 1);

// Use a protocol which is unlikely to have other packets on
const TEST_PROTO: IpNextHeaderProtocol = IpNextHeaderProtocols::Test1;

fn build_ipv4_header(packet: &mut [u8], offset: uint) {
    let mut ip_header = MutableIpv4Header::new(packet.slice_from_mut(offset));

    let total_len = (IPV4_HEADER_LEN + UDP_HEADER_LEN + TEST_DATA_LEN) as u16;

    ip_header.set_version(4);
    ip_header.set_header_length(5);
    ip_header.set_total_length(total_len);
    ip_header.set_ttl(4);
    ip_header.set_next_level_protocol(TEST_PROTO);
    ip_header.set_source(IPV4_SOURCE);
    ip_header.set_destination(IPV4_DESTINATION);
    ip_header.checksum();
}

fn build_ipv6_header(packet: &mut [u8], offset: uint) {
    let mut ip_header = MutableIpv6Header::new(packet.slice_from_mut(offset));

    ip_header.set_version(6);
    ip_header.set_payload_length((UDP_HEADER_LEN + TEST_DATA_LEN) as u16);
    ip_header.set_next_header(TEST_PROTO);
    ip_header.set_hop_limit(4);
    ip_header.set_source(IPV6_SOURCE);
    ip_header.set_destination(IPV6_DESTINATION);
}

fn build_udp_header(packet: &mut [u8], offset: uint) {
    let mut udp_header = MutableUdpHeader::new(packet.slice_from_mut(offset));

    udp_header.set_source(1234); // Arbitary port number
    udp_header.set_destination(1234);
    udp_header.set_length((UDP_HEADER_LEN + TEST_DATA_LEN) as u16);
}

fn build_udp4_packet(packet: &mut [u8], start: uint, msg: &str) {
    build_ipv4_header(packet, start);
    build_udp_header(packet, start + IPV4_HEADER_LEN as uint);

    let data_start = start + IPV4_HEADER_LEN + UDP_HEADER_LEN;
    packet[data_start + 0] = msg.char_at(0) as u8;
    packet[data_start + 1] = msg.char_at(1) as u8;
    packet[data_start + 2] = msg.char_at(2) as u8;
    packet[data_start + 3] = msg.char_at(3) as u8;

    let slice = packet.slice_from_mut(start + IPV4_HEADER_LEN as uint);
    MutableUdpHeader::new(slice).checksum(IPV4_SOURCE, IPV4_DESTINATION, TEST_PROTO);
}

fn build_udp6_packet(packet: &mut [u8], start: uint, msg: &str) {
    build_ipv6_header(packet, start);
    build_udp_header(packet, start + IPV6_HEADER_LEN as uint);

    let data_start = start + IPV6_HEADER_LEN + UDP_HEADER_LEN;
    packet[data_start + 0] = msg.char_at(0) as u8;
    packet[data_start + 1] = msg.char_at(1) as u8;
    packet[data_start + 2] = msg.char_at(2) as u8;
    packet[data_start + 3] = msg.char_at(3) as u8;

    let slice = packet.slice_from_mut(start + IPV6_HEADER_LEN as uint);
    MutableUdpHeader::new(slice).checksum(IPV6_SOURCE, IPV6_DESTINATION, TEST_PROTO);
}

fn get_test_interface() -> NetworkInterface {
    use std::os::getenv;
    use util;

    (*util::get_network_interfaces()
        .as_slice().iter()
        .filter(|x| {
            match getenv("PNET_TEST_IFACE") {
                Some(name) => x.name == name,
                None => x.is_loopback()
            }
        })
        .next()
        .unwrap())
        .clone()
}

// OSes have a nasty habit of tweaking IP fields, so we only check
// the less volatile fields (identification, checksum)
fn check_ipv4_header(packet: &[u8], header: Ipv4Header) {
    let ipv4_header = Ipv4Header::new(packet);

    assert_eq!(header.get_version(), ipv4_header.get_version());
    assert_eq!(header.get_header_length(), ipv4_header.get_header_length());
    assert_eq!(header.get_dscp(), ipv4_header.get_dscp());
    assert_eq!(header.get_ecn(), ipv4_header.get_ecn());
    assert_eq!(header.get_total_length(), ipv4_header.get_total_length());
    assert_eq!(header.get_flags(), ipv4_header.get_flags());
    assert_eq!(header.get_fragment_offset(), ipv4_header.get_fragment_offset());
    assert_eq!(header.get_ttl(), ipv4_header.get_ttl());
    assert_eq!(header.get_next_level_protocol(), ipv4_header.get_next_level_protocol());
    assert_eq!(header.get_source(), ipv4_header.get_source());
    assert_eq!(header.get_destination(), ipv4_header.get_destination());
}

fn layer4(ip: IpAddr, header_len: uint) {
    let mut packet = [0u8, ..IPV6_HEADER_LEN + UDP_HEADER_LEN + TEST_DATA_LEN];
    let packet_len = header_len + UDP_HEADER_LEN + TEST_DATA_LEN;

    match ip {
        Ipv4Addr(..) => {
            build_udp4_packet(packet.as_mut_slice(), 0, "l4i4")
        },
        Ipv6Addr(..) => {
            build_udp6_packet(packet.as_mut_slice(), 0, "l4i6")
        }
    };

    let udp = UdpHeader::new(packet.slice(header_len, packet_len));

    let (tx, rx) = channel();

    let tc = transport_channel(128, TransportChannelType::Layer4(get_proto(ip)));
    let (mut ttx, mut trx) = match tc {
        Ok((tx, rx)) => (tx, rx),
        Err(e) => panic!("layer4: unable to create channel: {}", e),
    };

    let res = try_future( proc() {
        tx.send(());
        pfor!((header, addr) in udp_header_iter(&mut trx) {
            assert_eq!(addr, ip);
            assert_eq!(header, UdpHeader::new(packet.slice(header_len, packet_len)));
            break;
        } on Err(e) {
            panic!("Receive failed for layer4_test(): {}", e);
        })
    });

    rx.recv();
    match ttx.send_to(udp, ip) {
        Ok(res) => assert_eq!(res as uint, UDP_HEADER_LEN + TEST_DATA_LEN),
        Err(e) => panic!("layer4_test failed: {}", e)
    }

    fn get_proto(ip: IpAddr) -> TransportProtocol {
        match ip {
            Ipv4Addr(..) => Ipv4(TEST_PROTO),
            Ipv6Addr(..) => Ipv6(TEST_PROTO)
        }
    }

    match res.unwrap() {
        Err(e) => panic!(e),
        _ => ()
    }
}

#[test]
fn layer4_ipv4() {
    layer4(Ipv4Addr(127, 0, 0, 1), IPV4_HEADER_LEN as uint);
}

#[test]
fn layer4_ipv6() {
    layer4(Ipv6Addr(0, 0, 0, 0, 0, 0, 0, 1), IPV6_HEADER_LEN);
}

#[test]
fn layer3_ipv4() {
    let send_addr = Ipv4Addr(127, 0, 0, 1);
    let mut packet = [0u8, ..IPV4_HEADER_LEN + UDP_HEADER_LEN + TEST_DATA_LEN];

    build_udp4_packet(packet.as_mut_slice(), 0, "l3i4");

    let (tx, rx) = channel();

    let tc =  transport_channel(IPV4_HEADER_LEN + UDP_HEADER_LEN + TEST_DATA_LEN,
                                    TransportChannelType::Layer3(TEST_PROTO));
    let (mut ttx, mut trx) = match tc {
        Ok((tx, rx)) => (tx, rx),
        Err(e) => panic!("layer3: unable to create channel: {}", e),
    };

    let res = try_future( proc() {
        tx.send(());
        pfor!((header, addr) in ipv4_header_iter(&mut trx) {
            assert_eq!(addr, send_addr);
            check_ipv4_header(packet.as_slice(), header);
            let udp_header = UdpHeader::new(header.packet().slice_from(
                                           header.get_header_length() as uint * 4u));
            assert_eq!(udp_header, UdpHeader::new(packet.slice_from(IPV4_HEADER_LEN)));

            assert_eq!(udp_header.packet().slice_from(UDP_HEADER_LEN),
                       packet.slice_from(IPV4_HEADER_LEN + UDP_HEADER_LEN))
            break;
        } on Err(e) {
            panic!("receive failed for layer3_ipv4_test(): {}", e);
        })
    });

    rx.recv();
    match ttx.send_to(Ipv4Header::new(packet.as_slice()), send_addr) {
        Ok(res) => assert_eq!(res as uint, packet.len()),
        Err(e) => panic!("layer3_ipv4_test failed: {}", e)
    }

    match res.unwrap() {
        Err(e) => panic!(e),
        _ => ()
    }

}

#[test]
fn layer2() {
    let interface = get_test_interface();

    let mut packet = [0u8, ..ETHERNET_HEADER_LEN +
                             IPV4_HEADER_LEN +
                             UDP_HEADER_LEN +
                             TEST_DATA_LEN];

    {
        let mut ethernet_header = MutableEthernetHeader::new(packet.as_mut_slice());
        ethernet_header.set_source(interface.mac_address());
        ethernet_header.set_destination(interface.mac_address());
        ethernet_header.set_ethertype(EtherTypes::Ipv4);
    }

    build_udp4_packet(packet.as_mut_slice(), ETHERNET_HEADER_LEN as uint, "l2tt");

    let (tx, rx) = channel();

    let dlc = datalink_channel(&interface,
                               MIN_PACKET_SIZE*2,
                               MIN_PACKET_SIZE*2,
                               DataLinkChannelType::Layer2);
    let (mut dltx, mut dlrx) = match dlc {
        Ok((tx, rx)) => (tx, rx),
        Err(e) => panic!("layer2: unable to create channel: {}", e)
    };

    let res = try_future( proc() {
        tx.send(());
        let mut i = 0u;
        pfor!(eh in dlrx.iter() {
            if i == 10_000 {
                panic!("layer2: did not find matching packet after 10_000 iterations");
            }
            if EthernetHeader::new(packet.as_slice()) == eh {
                return;
            }
            i += 1;
        } on Err(e) {
            panic!("layer2 failed: {}", e);
        })
    });

    rx.recv();
    match dltx.send_to(EthernetHeader::new(packet.as_slice()), None) {
        Some(Ok(())) => (),
        Some(Err(e)) => panic!("layer2_test failed: {}", e),
        None => panic!("Provided buffer too small")
    }

    match res.unwrap() {
        Err(e) => panic!(e),
        _ => ()
    }
}

#[test]
fn check_test_environment() {
    use std::os;
    let tasks = os::getenv("RUST_TEST_TASKS");
    if !tasks.is_some() || tasks.unwrap().as_slice() != "1" {
        panic!("Tests must be run with environment variable RUST_TEST_TASKS=1");
    }

    test_iface();

    #[cfg(not(target_os = "linux"))]
    fn test_iface() {
        let iface = os::getenv("PNET_TEST_IFACE");
        if !iface.is_some() {
            panic!("The environment variable PNET_TEST_IFACE must be set.");
        }
    }

    #[cfg(target_os = "linux")]
    fn test_iface() {}
}
<commit_msg>Fix unwrap deprecation warnings.<commit_after>// Copyright (c) 2014 Robert Clipsham <robert@octarineparrot.com>
//
// 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.

extern crate libc;

use std::clone::Clone;
use std::result::{Ok, Err};
use std::option::{Some};
use std::task::try_future;
use std::io::net::ip::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::iter::Iterator;

use datalink::{datalink_channel, DataLinkChannelType};
use packet::Packet;
use packet::ethernet::{EtherTypes, EthernetHeader, MutableEthernetHeader, EthernetPacket};
use packet::ip::{IpNextHeaderProtocols, IpNextHeaderProtocol};
use packet::ipv4::{Ipv4Header, MutableIpv4Header, Ipv4Packet};
use packet::ipv6::{MutableIpv6Header, Ipv6Packet};
use packet::udp::{UdpHeader, MutableUdpHeader, UdpPacket};
use transport::{udp_header_iter, ipv4_header_iter, transport_channel, TransportProtocol,
                TransportChannelType};
use transport::TransportProtocol::{Ipv4, Ipv6};
use util::NetworkInterface;

const MIN_PACKET_SIZE: uint = 64;
const ETHERNET_HEADER_LEN: uint = 14;
const IPV4_HEADER_LEN: uint = 20;
const IPV6_HEADER_LEN: uint = 40;
const UDP_HEADER_LEN: uint = 8;
const TEST_DATA_LEN: uint = 4;

const IPV4_SOURCE: IpAddr = Ipv4Addr(127, 0, 0, 1);
const IPV4_DESTINATION: IpAddr = Ipv4Addr(127, 0, 0, 1);
const IPV6_SOURCE: IpAddr = Ipv6Addr(0, 0, 0, 0, 0, 0, 0, 1);
const IPV6_DESTINATION: IpAddr = Ipv6Addr(0, 0, 0, 0, 0, 0, 0, 1);

// Use a protocol which is unlikely to have other packets on
const TEST_PROTO: IpNextHeaderProtocol = IpNextHeaderProtocols::Test1;

fn build_ipv4_header(packet: &mut [u8], offset: uint) {
    let mut ip_header = MutableIpv4Header::new(packet.slice_from_mut(offset));

    let total_len = (IPV4_HEADER_LEN + UDP_HEADER_LEN + TEST_DATA_LEN) as u16;

    ip_header.set_version(4);
    ip_header.set_header_length(5);
    ip_header.set_total_length(total_len);
    ip_header.set_ttl(4);
    ip_header.set_next_level_protocol(TEST_PROTO);
    ip_header.set_source(IPV4_SOURCE);
    ip_header.set_destination(IPV4_DESTINATION);
    ip_header.checksum();
}

fn build_ipv6_header(packet: &mut [u8], offset: uint) {
    let mut ip_header = MutableIpv6Header::new(packet.slice_from_mut(offset));

    ip_header.set_version(6);
    ip_header.set_payload_length((UDP_HEADER_LEN + TEST_DATA_LEN) as u16);
    ip_header.set_next_header(TEST_PROTO);
    ip_header.set_hop_limit(4);
    ip_header.set_source(IPV6_SOURCE);
    ip_header.set_destination(IPV6_DESTINATION);
}

fn build_udp_header(packet: &mut [u8], offset: uint) {
    let mut udp_header = MutableUdpHeader::new(packet.slice_from_mut(offset));

    udp_header.set_source(1234); // Arbitary port number
    udp_header.set_destination(1234);
    udp_header.set_length((UDP_HEADER_LEN + TEST_DATA_LEN) as u16);
}

fn build_udp4_packet(packet: &mut [u8], start: uint, msg: &str) {
    build_ipv4_header(packet, start);
    build_udp_header(packet, start + IPV4_HEADER_LEN as uint);

    let data_start = start + IPV4_HEADER_LEN + UDP_HEADER_LEN;
    packet[data_start + 0] = msg.char_at(0) as u8;
    packet[data_start + 1] = msg.char_at(1) as u8;
    packet[data_start + 2] = msg.char_at(2) as u8;
    packet[data_start + 3] = msg.char_at(3) as u8;

    let slice = packet.slice_from_mut(start + IPV4_HEADER_LEN as uint);
    MutableUdpHeader::new(slice).checksum(IPV4_SOURCE, IPV4_DESTINATION, TEST_PROTO);
}

fn build_udp6_packet(packet: &mut [u8], start: uint, msg: &str) {
    build_ipv6_header(packet, start);
    build_udp_header(packet, start + IPV6_HEADER_LEN as uint);

    let data_start = start + IPV6_HEADER_LEN + UDP_HEADER_LEN;
    packet[data_start + 0] = msg.char_at(0) as u8;
    packet[data_start + 1] = msg.char_at(1) as u8;
    packet[data_start + 2] = msg.char_at(2) as u8;
    packet[data_start + 3] = msg.char_at(3) as u8;

    let slice = packet.slice_from_mut(start + IPV6_HEADER_LEN as uint);
    MutableUdpHeader::new(slice).checksum(IPV6_SOURCE, IPV6_DESTINATION, TEST_PROTO);
}

fn get_test_interface() -> NetworkInterface {
    use std::os::getenv;
    use util;

    (*util::get_network_interfaces()
        .as_slice().iter()
        .filter(|x| {
            match getenv("PNET_TEST_IFACE") {
                Some(name) => x.name == name,
                None => x.is_loopback()
            }
        })
        .next()
        .unwrap())
        .clone()
}

// OSes have a nasty habit of tweaking IP fields, so we only check
// the less volatile fields (identification, checksum)
fn check_ipv4_header(packet: &[u8], header: Ipv4Header) {
    let ipv4_header = Ipv4Header::new(packet);

    assert_eq!(header.get_version(), ipv4_header.get_version());
    assert_eq!(header.get_header_length(), ipv4_header.get_header_length());
    assert_eq!(header.get_dscp(), ipv4_header.get_dscp());
    assert_eq!(header.get_ecn(), ipv4_header.get_ecn());
    assert_eq!(header.get_total_length(), ipv4_header.get_total_length());
    assert_eq!(header.get_flags(), ipv4_header.get_flags());
    assert_eq!(header.get_fragment_offset(), ipv4_header.get_fragment_offset());
    assert_eq!(header.get_ttl(), ipv4_header.get_ttl());
    assert_eq!(header.get_next_level_protocol(), ipv4_header.get_next_level_protocol());
    assert_eq!(header.get_source(), ipv4_header.get_source());
    assert_eq!(header.get_destination(), ipv4_header.get_destination());
}

fn layer4(ip: IpAddr, header_len: uint) {
    let mut packet = [0u8, ..IPV6_HEADER_LEN + UDP_HEADER_LEN + TEST_DATA_LEN];
    let packet_len = header_len + UDP_HEADER_LEN + TEST_DATA_LEN;

    match ip {
        Ipv4Addr(..) => {
            build_udp4_packet(packet.as_mut_slice(), 0, "l4i4")
        },
        Ipv6Addr(..) => {
            build_udp6_packet(packet.as_mut_slice(), 0, "l4i6")
        }
    };

    let udp = UdpHeader::new(packet.slice(header_len, packet_len));

    let (tx, rx) = channel();

    let tc = transport_channel(128, TransportChannelType::Layer4(get_proto(ip)));
    let (mut ttx, mut trx) = match tc {
        Ok((tx, rx)) => (tx, rx),
        Err(e) => panic!("layer4: unable to create channel: {}", e),
    };

    let res = try_future( proc() {
        tx.send(());
        pfor!((header, addr) in udp_header_iter(&mut trx) {
            assert_eq!(addr, ip);
            assert_eq!(header, UdpHeader::new(packet.slice(header_len, packet_len)));
            break;
        } on Err(e) {
            panic!("Receive failed for layer4_test(): {}", e);
        })
    });

    rx.recv();
    match ttx.send_to(udp, ip) {
        Ok(res) => assert_eq!(res as uint, UDP_HEADER_LEN + TEST_DATA_LEN),
        Err(e) => panic!("layer4_test failed: {}", e)
    }

    fn get_proto(ip: IpAddr) -> TransportProtocol {
        match ip {
            Ipv4Addr(..) => Ipv4(TEST_PROTO),
            Ipv6Addr(..) => Ipv6(TEST_PROTO)
        }
    }

    match res.into_inner() {
        Err(e) => panic!(e),
        _ => ()
    }
}

#[test]
fn layer4_ipv4() {
    layer4(Ipv4Addr(127, 0, 0, 1), IPV4_HEADER_LEN as uint);
}

#[test]
fn layer4_ipv6() {
    layer4(Ipv6Addr(0, 0, 0, 0, 0, 0, 0, 1), IPV6_HEADER_LEN);
}

#[test]
fn layer3_ipv4() {
    let send_addr = Ipv4Addr(127, 0, 0, 1);
    let mut packet = [0u8, ..IPV4_HEADER_LEN + UDP_HEADER_LEN + TEST_DATA_LEN];

    build_udp4_packet(packet.as_mut_slice(), 0, "l3i4");

    let (tx, rx) = channel();

    let tc =  transport_channel(IPV4_HEADER_LEN + UDP_HEADER_LEN + TEST_DATA_LEN,
                                    TransportChannelType::Layer3(TEST_PROTO));
    let (mut ttx, mut trx) = match tc {
        Ok((tx, rx)) => (tx, rx),
        Err(e) => panic!("layer3: unable to create channel: {}", e),
    };

    let res = try_future( proc() {
        tx.send(());
        pfor!((header, addr) in ipv4_header_iter(&mut trx) {
            assert_eq!(addr, send_addr);
            check_ipv4_header(packet.as_slice(), header);
            let udp_header = UdpHeader::new(header.packet().slice_from(
                                           header.get_header_length() as uint * 4u));
            assert_eq!(udp_header, UdpHeader::new(packet.slice_from(IPV4_HEADER_LEN)));

            assert_eq!(udp_header.packet().slice_from(UDP_HEADER_LEN),
                       packet.slice_from(IPV4_HEADER_LEN + UDP_HEADER_LEN))
            break;
        } on Err(e) {
            panic!("receive failed for layer3_ipv4_test(): {}", e);
        })
    });

    rx.recv();
    match ttx.send_to(Ipv4Header::new(packet.as_slice()), send_addr) {
        Ok(res) => assert_eq!(res as uint, packet.len()),
        Err(e) => panic!("layer3_ipv4_test failed: {}", e)
    }

    match res.into_inner() {
        Err(e) => panic!(e),
        _ => ()
    }

}

#[test]
fn layer2() {
    let interface = get_test_interface();

    let mut packet = [0u8, ..ETHERNET_HEADER_LEN +
                             IPV4_HEADER_LEN +
                             UDP_HEADER_LEN +
                             TEST_DATA_LEN];

    {
        let mut ethernet_header = MutableEthernetHeader::new(packet.as_mut_slice());
        ethernet_header.set_source(interface.mac_address());
        ethernet_header.set_destination(interface.mac_address());
        ethernet_header.set_ethertype(EtherTypes::Ipv4);
    }

    build_udp4_packet(packet.as_mut_slice(), ETHERNET_HEADER_LEN as uint, "l2tt");

    let (tx, rx) = channel();

    let dlc = datalink_channel(&interface,
                               MIN_PACKET_SIZE*2,
                               MIN_PACKET_SIZE*2,
                               DataLinkChannelType::Layer2);
    let (mut dltx, mut dlrx) = match dlc {
        Ok((tx, rx)) => (tx, rx),
        Err(e) => panic!("layer2: unable to create channel: {}", e)
    };

    let res = try_future( proc() {
        tx.send(());
        let mut i = 0u;
        pfor!(eh in dlrx.iter() {
            if i == 10_000 {
                panic!("layer2: did not find matching packet after 10_000 iterations");
            }
            if EthernetHeader::new(packet.as_slice()) == eh {
                return;
            }
            i += 1;
        } on Err(e) {
            panic!("layer2 failed: {}", e);
        })
    });

    rx.recv();
    match dltx.send_to(EthernetHeader::new(packet.as_slice()), None) {
        Some(Ok(())) => (),
        Some(Err(e)) => panic!("layer2_test failed: {}", e),
        None => panic!("Provided buffer too small")
    }

    match res.into_inner() {
        Err(e) => panic!(e),
        _ => ()
    }
}

#[test]
fn check_test_environment() {
    use std::os;
    let tasks = os::getenv("RUST_TEST_TASKS");
    if !tasks.is_some() || tasks.unwrap().as_slice() != "1" {
        panic!("Tests must be run with environment variable RUST_TEST_TASKS=1");
    }

    test_iface();

    #[cfg(not(target_os = "linux"))]
    fn test_iface() {
        let iface = os::getenv("PNET_TEST_IFACE");
        if !iface.is_some() {
            panic!("The environment variable PNET_TEST_IFACE must be set.");
        }
    }

    #[cfg(target_os = "linux")]
    fn test_iface() {}
}
<|endoftext|>


<commit_before>use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, AtomicUsize, ATOMIC_USIZE_INIT, AtomicIsize,
                        Ordering};
use std::sync::mpsc::{self, SyncSender, Receiver};
use std::sync::{Mutex, Arc};
use std::time::Duration;
use std::thread;

use {ManageConnection, CustomizeConnection, Pool, Config};

#[derive(Debug, PartialEq)]
struct FakeConnection(bool);

struct OkManager;

impl ManageConnection for OkManager {
    type Connection = FakeConnection;
    type Error = ();

    fn connect(&self) -> Result<FakeConnection, ()> {
        Ok(FakeConnection(true))
    }

    fn is_valid(&self, _: &mut FakeConnection) -> Result<(), ()> {
        Ok(())
    }

    fn has_broken(&self, _: &mut FakeConnection) -> bool {
        false
    }
}

struct NthConnectFailManager {
    n: Mutex<u32>,
}

impl ManageConnection for NthConnectFailManager {
    type Connection = FakeConnection;
    type Error = ();

    fn connect(&self) -> Result<FakeConnection, ()> {
        let mut n = self.n.lock().unwrap();
        if *n > 0 {
            *n -= 1;
            Ok(FakeConnection(true))
        } else {
            Err(())
        }
    }

    fn is_valid(&self, _: &mut FakeConnection) -> Result<(), ()> {
        Ok(())
    }

    fn has_broken(&self, _: &mut FakeConnection) -> bool {
        false
    }
}

#[test]
fn test_pool_size_ok() {
    let config = Config::builder().pool_size(5).build();
    let manager = NthConnectFailManager { n: Mutex::new(5) };
    let pool = Pool::new(config, manager).unwrap();
    let mut conns = vec![];
    for _ in 0..5 {
        conns.push(pool.get().ok().unwrap());
    }
}

#[test]
fn test_acquire_release() {
    let config = Config::builder().pool_size(2).build();
    let pool = Pool::new(config, OkManager).unwrap();

    let conn1 = pool.get().ok().unwrap();
    let conn2 = pool.get().ok().unwrap();
    drop(conn1);
    let conn3 = pool.get().ok().unwrap();
    drop(conn2);
    drop(conn3);
}

#[test]
fn test_is_send_sync() {
    fn is_send_sync<T: Send + Sync>() {}
    is_send_sync::<Pool<OkManager>>();
}

#[test]
fn test_issue_2_unlocked_during_is_valid() {
    struct BlockingChecker {
        first: AtomicBool,
        s: Mutex<SyncSender<()>>,
        r: Mutex<Receiver<()>>,
    }

    impl ManageConnection for BlockingChecker {
        type Connection = FakeConnection;
        type Error = ();

        fn connect(&self) -> Result<FakeConnection, ()> {
            Ok(FakeConnection(true))
        }

        fn is_valid(&self, _: &mut FakeConnection) -> Result<(), ()> {
            if self.first.compare_and_swap(true, false, Ordering::SeqCst) {
                self.s.lock().unwrap().send(()).unwrap();
                self.r.lock().unwrap().recv().unwrap();
            }
            Ok(())
        }

        fn has_broken(&self, _: &mut FakeConnection) -> bool {
            false
        }
    }

    let (s1, r1) = mpsc::sync_channel(0);
    let (s2, r2) = mpsc::sync_channel(0);

    let config = Config::builder()
                     .test_on_check_out(true)
                     .pool_size(2)
                     .build();
    let manager = BlockingChecker {
        first: AtomicBool::new(true),
        s: Mutex::new(s1),
        r: Mutex::new(r2),
    };
    let pool = Arc::new(Pool::new(config, manager).unwrap());

    let p2 = pool.clone();
    let t = thread::spawn(move || {
        p2.get().ok().unwrap();
    });

    r1.recv().unwrap();
    // get call by other task has triggered the health check
    pool.get().ok().unwrap();
    s2.send(()).ok().unwrap();

    t.join().ok().unwrap();
}

#[test]
fn test_drop_on_broken() {
    static DROPPED: AtomicBool = ATOMIC_BOOL_INIT;
    DROPPED.store(false, Ordering::SeqCst);

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.store(true, Ordering::SeqCst);
        }
    }

    struct Handler;

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = ();

        fn connect(&self) -> Result<Connection, ()> {
            Ok(Connection)
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), ()> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            true
        }
    }

    let pool = Pool::new(Default::default(), Handler).unwrap();

    drop(pool.get().ok().unwrap());

    assert!(DROPPED.load(Ordering::SeqCst));
}

#[test]
fn test_initialization_failure() {
    let config = Config::builder()
                     .connection_timeout_ms(1000)
                     .build();
    let manager = NthConnectFailManager { n: Mutex::new(0) };
    Pool::new(config, manager).err().unwrap();
}

#[test]
fn test_get_timeout() {
    let config = Config::builder()
                     .pool_size(1)
                     .connection_timeout(Duration::from_secs(1))
                     .build();
    let pool = Pool::new(config, OkManager).unwrap();
    let _c = pool.get().unwrap();
    pool.get().err().unwrap();
}

#[test]
fn test_connection_customizer() {
    static DROPPED: AtomicBool = ATOMIC_BOOL_INIT;
    DROPPED.store(false, Ordering::SeqCst);

    struct Connection(i32);

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.store(true, Ordering::SeqCst);
        }
    }

    struct Handler;

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = ();

        fn connect(&self) -> Result<Connection, ()> {
            Ok(Connection(0))
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), ()> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            true
        }
    }

    struct Customizer;

    impl CustomizeConnection<Connection, ()> for Customizer {
        fn on_acquire(&self, conn: &mut Connection) -> Result<(), ()> {
            if !DROPPED.load(Ordering::SeqCst) {
                Err(())
            } else {
                conn.0 = 1;
                Ok(())
            }
        }
    }

    let config = Config::builder()
                     .connection_customizer(Box::new(Customizer))
                     .build();
    let pool = Pool::new(config, Handler).unwrap();

    let conn = pool.get().unwrap();
    assert_eq!(1, conn.0);
    assert!(DROPPED.load(Ordering::SeqCst));
}

#[test]
fn test_idle_timeout() {
    static DROPPED: AtomicUsize = ATOMIC_USIZE_INIT;

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.fetch_add(1, Ordering::SeqCst);
        }
    }

    struct Handler(AtomicIsize);

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = ();

        fn connect(&self) -> Result<Connection, ()> {
            if self.0.fetch_sub(1, Ordering::SeqCst) > 0 {
                Ok(Connection)
            } else {
                Err(())
            }
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), ()> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            false
        }
    }

    let config = Config::builder()
                     .pool_size(5)
                     .idle_timeout(Some(Duration::from_secs(1)))
                     .build();
    let pool = Pool::new_inner(config, Handler(AtomicIsize::new(5)), 1).unwrap();
    let conn = pool.get().unwrap();
    thread::sleep(Duration::from_secs(2));
    assert_eq!(4, DROPPED.load(Ordering::SeqCst));
    drop(conn);
    assert_eq!(4, DROPPED.load(Ordering::SeqCst));
}

#[test]
fn test_max_lifetime() {
    static DROPPED: AtomicUsize = ATOMIC_USIZE_INIT;

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.fetch_add(1, Ordering::SeqCst);
        }
    }

    struct Handler(AtomicIsize);

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = ();

        fn connect(&self) -> Result<Connection, ()> {
            if self.0.fetch_sub(1, Ordering::SeqCst) > 0 {
                Ok(Connection)
            } else {
                Err(())
            }
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), ()> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            false
        }
    }

    let config = Config::builder()
                     .pool_size(5)
                     .max_lifetime(Some(Duration::from_secs(1)))
                     .connection_timeout(Duration::from_secs(1))
                     .build();
    let pool = Pool::new_inner(config, Handler(AtomicIsize::new(5)), 1).unwrap();
    let conn = pool.get().unwrap();
    thread::sleep(Duration::from_secs(2));
    assert_eq!(4, DROPPED.load(Ordering::SeqCst));
    drop(conn);
    thread::sleep(Duration::from_secs(2));
    assert_eq!(5, DROPPED.load(Ordering::SeqCst));
    assert!(pool.get().is_err());
}

#[test]
fn min_idle() {
    static CREATED: AtomicUsize = ATOMIC_USIZE_INIT;

    struct Connection;

    struct Handler;

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = ();

        fn connect(&self) -> Result<Connection, ()> {
            CREATED.fetch_add(1, Ordering::SeqCst);
            Ok(Connection)
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), ()> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            false
        }
    }

    let config = Config::builder()
                     .pool_size(5)
                     .min_idle(Some(2))
                     .build();
    let pool = Pool::new(config, Handler).unwrap();
    assert_eq!(2, CREATED.load(Ordering::SeqCst));
    let _conns = (0..5).map(|_| pool.get().unwrap()).collect::<Vec<_>>();
    assert_eq!(5, CREATED.load(Ordering::SeqCst));
}

#[test]
fn conns_drop_on_pool_drop() {
    static DROPPED: AtomicUsize = ATOMIC_USIZE_INIT;

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.fetch_add(1, Ordering::SeqCst);
        }
    }

    struct Handler;

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = ();

        fn connect(&self) -> Result<Connection, ()> {
            Ok(Connection)
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), ()> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            false
        }
    }

    let config = Config::builder()
                    .max_lifetime(Some(Duration::from_secs(10)))
                    .pool_size(10)
                    .build();
    let pool = Pool::new(config, Handler).unwrap();
    drop(pool);
    assert_eq!(DROPPED.load(Ordering::SeqCst), 10);
}
<commit_msg>Make conns_drop_on_pool_drop less flaky<commit_after>use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, AtomicUsize, ATOMIC_USIZE_INIT, AtomicIsize,
                        Ordering};
use std::sync::mpsc::{self, SyncSender, Receiver};
use std::sync::{Mutex, Arc};
use std::time::Duration;
use std::thread;

use {ManageConnection, CustomizeConnection, Pool, Config};

#[derive(Debug, PartialEq)]
struct FakeConnection(bool);

struct OkManager;

impl ManageConnection for OkManager {
    type Connection = FakeConnection;
    type Error = ();

    fn connect(&self) -> Result<FakeConnection, ()> {
        Ok(FakeConnection(true))
    }

    fn is_valid(&self, _: &mut FakeConnection) -> Result<(), ()> {
        Ok(())
    }

    fn has_broken(&self, _: &mut FakeConnection) -> bool {
        false
    }
}

struct NthConnectFailManager {
    n: Mutex<u32>,
}

impl ManageConnection for NthConnectFailManager {
    type Connection = FakeConnection;
    type Error = ();

    fn connect(&self) -> Result<FakeConnection, ()> {
        let mut n = self.n.lock().unwrap();
        if *n > 0 {
            *n -= 1;
            Ok(FakeConnection(true))
        } else {
            Err(())
        }
    }

    fn is_valid(&self, _: &mut FakeConnection) -> Result<(), ()> {
        Ok(())
    }

    fn has_broken(&self, _: &mut FakeConnection) -> bool {
        false
    }
}

#[test]
fn test_pool_size_ok() {
    let config = Config::builder().pool_size(5).build();
    let manager = NthConnectFailManager { n: Mutex::new(5) };
    let pool = Pool::new(config, manager).unwrap();
    let mut conns = vec![];
    for _ in 0..5 {
        conns.push(pool.get().ok().unwrap());
    }
}

#[test]
fn test_acquire_release() {
    let config = Config::builder().pool_size(2).build();
    let pool = Pool::new(config, OkManager).unwrap();

    let conn1 = pool.get().ok().unwrap();
    let conn2 = pool.get().ok().unwrap();
    drop(conn1);
    let conn3 = pool.get().ok().unwrap();
    drop(conn2);
    drop(conn3);
}

#[test]
fn test_is_send_sync() {
    fn is_send_sync<T: Send + Sync>() {}
    is_send_sync::<Pool<OkManager>>();
}

#[test]
fn test_issue_2_unlocked_during_is_valid() {
    struct BlockingChecker {
        first: AtomicBool,
        s: Mutex<SyncSender<()>>,
        r: Mutex<Receiver<()>>,
    }

    impl ManageConnection for BlockingChecker {
        type Connection = FakeConnection;
        type Error = ();

        fn connect(&self) -> Result<FakeConnection, ()> {
            Ok(FakeConnection(true))
        }

        fn is_valid(&self, _: &mut FakeConnection) -> Result<(), ()> {
            if self.first.compare_and_swap(true, false, Ordering::SeqCst) {
                self.s.lock().unwrap().send(()).unwrap();
                self.r.lock().unwrap().recv().unwrap();
            }
            Ok(())
        }

        fn has_broken(&self, _: &mut FakeConnection) -> bool {
            false
        }
    }

    let (s1, r1) = mpsc::sync_channel(0);
    let (s2, r2) = mpsc::sync_channel(0);

    let config = Config::builder()
                     .test_on_check_out(true)
                     .pool_size(2)
                     .build();
    let manager = BlockingChecker {
        first: AtomicBool::new(true),
        s: Mutex::new(s1),
        r: Mutex::new(r2),
    };
    let pool = Arc::new(Pool::new(config, manager).unwrap());

    let p2 = pool.clone();
    let t = thread::spawn(move || {
        p2.get().ok().unwrap();
    });

    r1.recv().unwrap();
    // get call by other task has triggered the health check
    pool.get().ok().unwrap();
    s2.send(()).ok().unwrap();

    t.join().ok().unwrap();
}

#[test]
fn test_drop_on_broken() {
    static DROPPED: AtomicBool = ATOMIC_BOOL_INIT;
    DROPPED.store(false, Ordering::SeqCst);

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.store(true, Ordering::SeqCst);
        }
    }

    struct Handler;

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = ();

        fn connect(&self) -> Result<Connection, ()> {
            Ok(Connection)
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), ()> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            true
        }
    }

    let pool = Pool::new(Default::default(), Handler).unwrap();

    drop(pool.get().ok().unwrap());

    assert!(DROPPED.load(Ordering::SeqCst));
}

#[test]
fn test_initialization_failure() {
    let config = Config::builder()
                     .connection_timeout_ms(1000)
                     .build();
    let manager = NthConnectFailManager { n: Mutex::new(0) };
    Pool::new(config, manager).err().unwrap();
}

#[test]
fn test_get_timeout() {
    let config = Config::builder()
                     .pool_size(1)
                     .connection_timeout(Duration::from_secs(1))
                     .build();
    let pool = Pool::new(config, OkManager).unwrap();
    let _c = pool.get().unwrap();
    pool.get().err().unwrap();
}

#[test]
fn test_connection_customizer() {
    static DROPPED: AtomicBool = ATOMIC_BOOL_INIT;
    DROPPED.store(false, Ordering::SeqCst);

    struct Connection(i32);

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.store(true, Ordering::SeqCst);
        }
    }

    struct Handler;

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = ();

        fn connect(&self) -> Result<Connection, ()> {
            Ok(Connection(0))
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), ()> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            true
        }
    }

    struct Customizer;

    impl CustomizeConnection<Connection, ()> for Customizer {
        fn on_acquire(&self, conn: &mut Connection) -> Result<(), ()> {
            if !DROPPED.load(Ordering::SeqCst) {
                Err(())
            } else {
                conn.0 = 1;
                Ok(())
            }
        }
    }

    let config = Config::builder()
                     .connection_customizer(Box::new(Customizer))
                     .build();
    let pool = Pool::new(config, Handler).unwrap();

    let conn = pool.get().unwrap();
    assert_eq!(1, conn.0);
    assert!(DROPPED.load(Ordering::SeqCst));
}

#[test]
fn test_idle_timeout() {
    static DROPPED: AtomicUsize = ATOMIC_USIZE_INIT;

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.fetch_add(1, Ordering::SeqCst);
        }
    }

    struct Handler(AtomicIsize);

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = ();

        fn connect(&self) -> Result<Connection, ()> {
            if self.0.fetch_sub(1, Ordering::SeqCst) > 0 {
                Ok(Connection)
            } else {
                Err(())
            }
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), ()> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            false
        }
    }

    let config = Config::builder()
                     .pool_size(5)
                     .idle_timeout(Some(Duration::from_secs(1)))
                     .build();
    let pool = Pool::new_inner(config, Handler(AtomicIsize::new(5)), 1).unwrap();
    let conn = pool.get().unwrap();
    thread::sleep(Duration::from_secs(2));
    assert_eq!(4, DROPPED.load(Ordering::SeqCst));
    drop(conn);
    assert_eq!(4, DROPPED.load(Ordering::SeqCst));
}

#[test]
fn test_max_lifetime() {
    static DROPPED: AtomicUsize = ATOMIC_USIZE_INIT;

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.fetch_add(1, Ordering::SeqCst);
        }
    }

    struct Handler(AtomicIsize);

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = ();

        fn connect(&self) -> Result<Connection, ()> {
            if self.0.fetch_sub(1, Ordering::SeqCst) > 0 {
                Ok(Connection)
            } else {
                Err(())
            }
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), ()> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            false
        }
    }

    let config = Config::builder()
                     .pool_size(5)
                     .max_lifetime(Some(Duration::from_secs(1)))
                     .connection_timeout(Duration::from_secs(1))
                     .build();
    let pool = Pool::new_inner(config, Handler(AtomicIsize::new(5)), 1).unwrap();
    let conn = pool.get().unwrap();
    thread::sleep(Duration::from_secs(2));
    assert_eq!(4, DROPPED.load(Ordering::SeqCst));
    drop(conn);
    thread::sleep(Duration::from_secs(2));
    assert_eq!(5, DROPPED.load(Ordering::SeqCst));
    assert!(pool.get().is_err());
}

#[test]
fn min_idle() {
    static CREATED: AtomicUsize = ATOMIC_USIZE_INIT;

    struct Connection;

    struct Handler;

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = ();

        fn connect(&self) -> Result<Connection, ()> {
            CREATED.fetch_add(1, Ordering::SeqCst);
            Ok(Connection)
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), ()> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            false
        }
    }

    let config = Config::builder()
                     .pool_size(5)
                     .min_idle(Some(2))
                     .build();
    let pool = Pool::new(config, Handler).unwrap();
    assert_eq!(2, CREATED.load(Ordering::SeqCst));
    let _conns = (0..5).map(|_| pool.get().unwrap()).collect::<Vec<_>>();
    assert_eq!(5, CREATED.load(Ordering::SeqCst));
}

#[test]
fn conns_drop_on_pool_drop() {
    static DROPPED: AtomicUsize = ATOMIC_USIZE_INIT;

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.fetch_add(1, Ordering::SeqCst);
        }
    }

    struct Handler;

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = ();

        fn connect(&self) -> Result<Connection, ()> {
            Ok(Connection)
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), ()> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            false
        }
    }

    let config = Config::builder()
                    .max_lifetime(Some(Duration::from_secs(10)))
                    .pool_size(10)
                    .build();
    let pool = Pool::new(config, Handler).unwrap();
    drop(pool);
    for _ in 0..10 {
        if DROPPED.load(Ordering::SeqCst) == 10 {
            return;
        }
        thread::sleep(Duration::from_secs(1));
    }
    panic!("timed out waiting for connections to drop");
}
<|endoftext|>


<commit_before>use countminsketch::CountMinSketch;
use std::collections::{BTreeSet, HashMap};
use std::collections::hash_map::Entry;
use std::cmp::Ordering;
use std::fmt;
use std::hash::Hash;
use std::rc::Rc;

struct TreeEntry<T>
where
    T: Eq + Ord,
{
    obj: Rc<T>,
    n: usize,
}

impl<T> Clone for TreeEntry<T>
where
    T: Eq + Ord,
{
    fn clone(&self) -> Self {
        TreeEntry {
            obj: Rc::clone(&self.obj),
            n: self.n,
        }
    }
}

impl<T> PartialEq for TreeEntry<T>
where
    T: Eq + Ord,
{
    fn eq(&self, other: &TreeEntry<T>) -> bool {
        self.obj == other.obj
    }
}

impl<T> Eq for TreeEntry<T>
where
    T: Eq + Ord,
{
}

impl<T> PartialOrd for TreeEntry<T>
where
    T: Eq + Ord,
{
    fn partial_cmp(&self, other: &TreeEntry<T>) -> Option<Ordering> {
        match self.n.cmp(&other.n) {
            Ordering::Greater => Some(Ordering::Greater),
            Ordering::Less => Some(Ordering::Less),
            Ordering::Equal => Some(self.obj.cmp(&other.obj)),
        }
    }
}

impl<T> Ord for TreeEntry<T>
where
    T: Eq + Ord,
{
    fn cmp(&self, other: &TreeEntry<T>) -> Ordering {
        self.partial_cmp(other).unwrap()
    }
}

/// Top-K implementation.
///
/// This data structure keeps the `k` most frequent data points of a stream.
#[derive(Clone)]
pub struct TopK<T>
where
    T: Clone + Eq + Hash + Ord,
{
    cms: CountMinSketch,
    obj2count: HashMap<Rc<T>, usize>,
    tree: BTreeSet<TreeEntry<T>>,
    k: usize,
}

impl<T> TopK<T>
where
    T: Clone + Eq + Hash + Ord,
{
    /// Create new Top-K data structure.
    ///
    /// - `k`: number of elements to remember
    /// - `cms`: CountMinSketch that is used for guessing the frequency of elements not currently
    ///   hold. Refer to its documentation about parameter selection.
    ///
    /// Panics if `k == 0`.
    pub fn new(k: usize, cms: CountMinSketch) -> TopK<T> {
        assert!(k > 0);
        TopK {
            cms: cms,
            obj2count: HashMap::new(),
            tree: BTreeSet::new(),
            k: k,
        }
    }

    /// Number of data points to remember.
    pub fn k(&self) -> usize {
        return self.k;
    }

    /// Observe a data point.
    pub fn add(&mut self, obj: T) {
        // create Rc so we can use obj in multiple data structs
        let rc = Rc::new(obj);

        // always increase CountMinSketch counts
        self.cms.add(&rc);

        // check if entry is in top K
        let size = self.obj2count.len();
        let mut update_obj2count: Option<(usize, Rc<T>)> = None;
        match self.obj2count.entry(Rc::clone(&rc)) {
            Entry::Occupied(mut o) => {
                // it is => increase exact counter
                let mut n = o.get_mut();
                *n += 1;

                // update tree data
                let mut entry = TreeEntry {
                    obj: Rc::clone(&rc),
                    n: *n - 1,
                };
                self.tree.remove(&entry);
                entry.n += 1;
                self.tree.insert(entry);
            }
            Entry::Vacant(v) => {
                // it's not => check capicity
                if size < self.k {
                    // space left => add to top k
                    v.insert(1);
                    self.tree.insert(TreeEntry {
                        obj: Rc::clone(&rc),
                        n: 1,
                    });
                } else {
                    // not enough space => check if it would be a top k element
                    let count = self.cms.query_point(&rc);
                    let min: TreeEntry<T> = (*self.tree.iter().next().unwrap()).clone();
                    if count > min.n {
                        // => kick out minimal element of top k
                        self.tree.remove(&min);
                        self.tree.insert(TreeEntry {
                            obj: Rc::clone(&rc),
                            n: count,
                        });

                        // delay obj2count changes, because it is currently borrowed
                        update_obj2count = Some((count, min.obj));
                    }
                }
            }
        }

        // trick the borrow checker
        if let Some((count, obj)) = update_obj2count {
            self.obj2count.insert(rc, count);
            self.obj2count.remove(&obj);
        }
    }

    /// Returns top-k values.
    ///
    /// The result may contain less than `k` values if less than `k` unique data points where
    /// observed.
    pub fn values(&self) -> Vec<T> {
        self.tree.iter().map(|x| (*x.obj).clone()).collect()
    }

    /// Check whether the sampler is empty (i.e. no data points observer so far)
    pub fn is_empty(&self) -> bool {
        self.obj2count.is_empty()
    }

    /// Clear sampler, so that it behaves like it never observed any data points.
    pub fn clear(&mut self) {
        self.cms.clear();
        self.tree.clear();
        self.obj2count.clear();
    }
}

impl<T> fmt::Debug for TopK<T>
where
    T: Clone + Eq + Hash + Ord,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "TopK {{ k: {} }}", self.k)
    }
}

impl<T> Extend<T> for TopK<T>
where
    T: Clone + Eq + Hash + Ord,
{
    fn extend<S: IntoIterator<Item = T>>(&mut self, iter: S) {
        for elem in iter {
            self.add(elem);
        }
    }
}

#[cfg(test)]
mod tests {
    use countminsketch::CountMinSketch;
    use super::TopK;

    #[test]
    fn getter() {
        let cms = CountMinSketch::with_params(10, 20);
        let tk: TopK<usize> = TopK::new(2, cms);

        assert_eq!(tk.k(), 2);
    }

    #[test]
    fn add_1() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);

        tk.add(1);
        assert_eq!(tk.values(), vec![1]);
    }

    #[test]
    fn add_2_same() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);

        tk.add(1);
        tk.add(1);
        assert_eq!(tk.values(), vec![1]);
    }

    #[test]
    fn add_2_different() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);

        tk.add(1);
        tk.add(2);
        assert_eq!(tk.values(), vec![1, 2]);
    }

    #[test]
    fn add_n() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);

        for i in 0..5 {
            tk.add(i);
        }
        for _ in 0..100 {
            tk.add(99);
        }
        for _ in 0..100 {
            tk.add(100);
        }
        for i in 0..5 {
            tk.add(i);
        }
        assert_eq!(tk.values(), vec![99, 100]);
    }

    #[test]
    fn is_empty() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);
        assert_eq!(tk.is_empty(), true);

        tk.add(0);
        assert_eq!(tk.is_empty(), false);
    }

    #[test]
    fn clear() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);
        tk.add(0);

        tk.clear();
        assert_eq!(tk.is_empty(), true);

        tk.add(1);
        assert_eq!(tk.values(), vec![1]);
    }

    #[test]
    fn clone() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk1 = TopK::new(2, cms);
        tk1.add(0);

        let mut tk2 = tk1.clone();
        tk2.add(1);

        assert_eq!(tk1.values(), vec![0]);
        assert_eq!(tk2.values(), vec![0, 1]);
    }

    #[test]
    fn debug() {
        let cms = CountMinSketch::with_params(10, 20);
        let tk: TopK<usize> = TopK::new(2, cms);
        assert_eq!(format!("{:?}", tk), "TopK { k: 2 }");
    }

    #[test]
    fn extend() {
        let cms = CountMinSketch::with_params(10, 10);
        let mut tk = TopK::new(2, cms);

        tk.extend(vec![0, 1]);
        assert_eq!(tk.values(), vec![0, 1]);
    }
}
<commit_msg>topk: .values() -> .iter()<commit_after>use countminsketch::CountMinSketch;
use std::collections::{BTreeSet, HashMap};
use std::collections::hash_map::Entry;
use std::cmp::Ordering;
use std::fmt;
use std::hash::Hash;
use std::rc::Rc;

struct TreeEntry<T>
where
    T: Eq + Ord,
{
    obj: Rc<T>,
    n: usize,
}

impl<T> Clone for TreeEntry<T>
where
    T: Eq + Ord,
{
    fn clone(&self) -> Self {
        TreeEntry {
            obj: Rc::clone(&self.obj),
            n: self.n,
        }
    }
}

impl<T> PartialEq for TreeEntry<T>
where
    T: Eq + Ord,
{
    fn eq(&self, other: &TreeEntry<T>) -> bool {
        self.obj == other.obj
    }
}

impl<T> Eq for TreeEntry<T>
where
    T: Eq + Ord,
{
}

impl<T> PartialOrd for TreeEntry<T>
where
    T: Eq + Ord,
{
    fn partial_cmp(&self, other: &TreeEntry<T>) -> Option<Ordering> {
        match self.n.cmp(&other.n) {
            Ordering::Greater => Some(Ordering::Greater),
            Ordering::Less => Some(Ordering::Less),
            Ordering::Equal => Some(self.obj.cmp(&other.obj)),
        }
    }
}

impl<T> Ord for TreeEntry<T>
where
    T: Eq + Ord,
{
    fn cmp(&self, other: &TreeEntry<T>) -> Ordering {
        self.partial_cmp(other).unwrap()
    }
}

/// Top-K implementation.
///
/// This data structure keeps the `k` most frequent data points of a stream.
#[derive(Clone)]
pub struct TopK<T>
where
    T: Clone + Eq + Hash + Ord,
{
    cms: CountMinSketch,
    obj2count: HashMap<Rc<T>, usize>,
    tree: BTreeSet<TreeEntry<T>>,
    k: usize,
}

impl<T> TopK<T>
where
    T: Clone + Eq + Hash + Ord,
{
    /// Create new Top-K data structure.
    ///
    /// - `k`: number of elements to remember
    /// - `cms`: CountMinSketch that is used for guessing the frequency of elements not currently
    ///   hold. Refer to its documentation about parameter selection.
    ///
    /// Panics if `k == 0`.
    pub fn new(k: usize, cms: CountMinSketch) -> TopK<T> {
        assert!(k > 0);
        TopK {
            cms: cms,
            obj2count: HashMap::new(),
            tree: BTreeSet::new(),
            k: k,
        }
    }

    /// Number of data points to remember.
    pub fn k(&self) -> usize {
        return self.k;
    }

    /// Observe a data point.
    pub fn add(&mut self, obj: T) {
        // create Rc so we can use obj in multiple data structs
        let rc = Rc::new(obj);

        // always increase CountMinSketch counts
        self.cms.add(&rc);

        // check if entry is in top K
        let size = self.obj2count.len();
        let mut update_obj2count: Option<(usize, Rc<T>)> = None;
        match self.obj2count.entry(Rc::clone(&rc)) {
            Entry::Occupied(mut o) => {
                // it is => increase exact counter
                let mut n = o.get_mut();
                *n += 1;

                // update tree data
                let mut entry = TreeEntry {
                    obj: Rc::clone(&rc),
                    n: *n - 1,
                };
                self.tree.remove(&entry);
                entry.n += 1;
                self.tree.insert(entry);
            }
            Entry::Vacant(v) => {
                // it's not => check capicity
                if size < self.k {
                    // space left => add to top k
                    v.insert(1);
                    self.tree.insert(TreeEntry {
                        obj: Rc::clone(&rc),
                        n: 1,
                    });
                } else {
                    // not enough space => check if it would be a top k element
                    let count = self.cms.query_point(&rc);
                    let min: TreeEntry<T> = (*self.tree.iter().next().unwrap()).clone();
                    if count > min.n {
                        // => kick out minimal element of top k
                        self.tree.remove(&min);
                        self.tree.insert(TreeEntry {
                            obj: Rc::clone(&rc),
                            n: count,
                        });

                        // delay obj2count changes, because it is currently borrowed
                        update_obj2count = Some((count, min.obj));
                    }
                }
            }
        }

        // trick the borrow checker
        if let Some((count, obj)) = update_obj2count {
            self.obj2count.insert(rc, count);
            self.obj2count.remove(&obj);
        }
    }

    /// Iterates over collected top-k values.
    ///
    /// The result may contain less than `k` values if less than `k` unique data points where
    /// observed.
    pub fn iter<'a>(&'a self) -> impl 'a + Iterator<Item=T> {
        self.tree.iter().map(|x| (*x.obj).clone())
    }

    /// Check whether the sampler is empty (i.e. no data points observer so far)
    pub fn is_empty(&self) -> bool {
        self.obj2count.is_empty()
    }

    /// Clear sampler, so that it behaves like it never observed any data points.
    pub fn clear(&mut self) {
        self.cms.clear();
        self.tree.clear();
        self.obj2count.clear();
    }
}

impl<T> fmt::Debug for TopK<T>
where
    T: Clone + Eq + Hash + Ord,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "TopK {{ k: {} }}", self.k)
    }
}

impl<T> Extend<T> for TopK<T>
where
    T: Clone + Eq + Hash + Ord,
{
    fn extend<S: IntoIterator<Item = T>>(&mut self, iter: S) {
        for elem in iter {
            self.add(elem);
        }
    }
}

#[cfg(test)]
mod tests {
    use countminsketch::CountMinSketch;
    use super::TopK;

    #[test]
    fn getter() {
        let cms = CountMinSketch::with_params(10, 20);
        let tk: TopK<usize> = TopK::new(2, cms);

        assert_eq!(tk.k(), 2);
    }

    #[test]
    fn add_1() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);

        tk.add(1);
        assert_eq!(tk.iter().collect::<Vec<u32>>(), vec![1]);
    }

    #[test]
    fn add_2_same() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);

        tk.add(1);
        tk.add(1);
        assert_eq!(tk.iter().collect::<Vec<u32>>(), vec![1]);
    }

    #[test]
    fn add_2_different() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);

        tk.add(1);
        tk.add(2);
        assert_eq!(tk.iter().collect::<Vec<u32>>(), vec![1, 2]);
    }

    #[test]
    fn add_n() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);

        for i in 0..5 {
            tk.add(i);
        }
        for _ in 0..100 {
            tk.add(99);
        }
        for _ in 0..100 {
            tk.add(100);
        }
        for i in 0..5 {
            tk.add(i);
        }
        assert_eq!(tk.iter().collect::<Vec<u32>>(), vec![99, 100]);
    }

    #[test]
    fn is_empty() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);
        assert_eq!(tk.is_empty(), true);

        tk.add(0);
        assert_eq!(tk.is_empty(), false);
    }

    #[test]
    fn clear() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk = TopK::new(2, cms);
        tk.add(0);

        tk.clear();
        assert_eq!(tk.is_empty(), true);

        tk.add(1);
        assert_eq!(tk.iter().collect::<Vec<u32>>(), vec![1]);
    }

    #[test]
    fn clone() {
        let cms = CountMinSketch::with_params(10, 20);
        let mut tk1 = TopK::new(2, cms);
        tk1.add(0);

        let mut tk2 = tk1.clone();
        tk2.add(1);

        assert_eq!(tk1.iter().collect::<Vec<u32>>(), vec![0]);
        assert_eq!(tk2.iter().collect::<Vec<u32>>(), vec![0, 1]);
    }

    #[test]
    fn debug() {
        let cms = CountMinSketch::with_params(10, 20);
        let tk: TopK<usize> = TopK::new(2, cms);
        assert_eq!(format!("{:?}", tk), "TopK { k: 2 }");
    }

    #[test]
    fn extend() {
        let cms = CountMinSketch::with_params(10, 10);
        let mut tk = TopK::new(2, cms);

        tk.extend(vec![0, 1]);
        assert_eq!(tk.iter().collect::<Vec<u32>>(), vec![0, 1]);
    }
}
<|endoftext|>