<commit_before>//! Rust simulator code generation.

use crate::code_writer;
use crate::module;

use std::collections::HashMap;
use std::io::Write;

struct RegNames {
    value_name: String,
    next_name: String,
}

struct Context<'a> {
    reg_names: HashMap<*const module::Signal<'a>, RegNames>,
}

enum ValueType {
    Bool,
    U32,
    U64,
    U128,
}

impl ValueType {
    fn from_bit_width(bit_width: u32) -> ValueType {
        if bit_width == 1 {
            ValueType::Bool
        } else if bit_width <= 32 {
            ValueType::U32
        } else if bit_width <= 64 {
            ValueType::U64
        } else if bit_width <= 128 {
            ValueType::U128
        } else {
            unreachable!()
        }
    }

    fn name(&self) -> &'static str {
        match self {
            ValueType::Bool => "bool",
            ValueType::U32 => "u32",
            ValueType::U64 => "u64",
            ValueType::U128 => "u128",
        }
    }

    fn bit_width(&self) -> u32 {
        match self {
            ValueType::Bool => 1,
            ValueType::U32 => 32,
            ValueType::U64 => 64,
            ValueType::U128 => 128,
        }
    }
}

pub fn generate<W: Write>(m: &module::Module, w: &mut W) -> Result<(), code_writer::Error> {
    let mut c = Context {
        reg_names: HashMap::new(),
    };

    for (_, output) in m.outputs().iter() {
        gather_regs(&output.source, &mut c);
    }

    let mut w = code_writer::CodeWriter::new(w);

    w.append_line("#[allow(non_camel_case_types)]")?;
    w.append_line("#[derive(Default)]")?;
    w.append_line(&format!("pub struct {} {{", m.name))?;
    w.indent();

    let inputs = m.inputs();
    if inputs.len() > 0 {
        w.append_line("// Inputs")?;
        for (name, input) in inputs.iter() {
            w.append_line(&format!(
                "pub {}: {}, // {} bit(s)",
                name,
                ValueType::from_bit_width(input.bit_width()).name(),
                input.bit_width()
            ))?;
        }
    }

    let outputs = m.outputs();
    if outputs.len() > 0 {
        w.append_line("// Outputs")?;
        for (name, output) in outputs.iter() {
            w.append_line(&format!(
                "pub {}: {}, // {} bit(s)",
                name,
                ValueType::from_bit_width(output.source.bit_width()).name(),
                output.source.bit_width()
            ))?;
        }
    }

    if c.reg_names.len() > 0 {
        w.append_line("// Regs")?;
        for (reg, names) in c.reg_names.iter() {
            let reg = unsafe { &**reg as &module::Signal };
            let type_name = ValueType::from_bit_width(reg.bit_width()).name();
            w.append_line(&format!(
                "{}: {}, // {} bit(s)",
                names.value_name,
                type_name,
                reg.bit_width()
            ))?;
            w.append_line(&format!("{}: {},", names.next_name, type_name))?;
        }
    }

    w.unindent()?;
    w.append_line("}")?;
    w.append_newline()?;

    w.append_line(&format!("impl {} {{", m.name))?;
    w.indent();

    if c.reg_names.len() > 0 {
        w.append_line("pub fn reset(&mut self) {")?;
        w.indent();

        for (reg, names) in c.reg_names.iter() {
            let reg = unsafe { &**reg as &module::Signal };
            match reg.data {
                module::SignalData::Reg {
                    ref initial_value,
                    bit_width,
                    ..
                } => {
                    w.append_indent()?;
                    w.append(&format!("self.{} = ", names.value_name))?;
                    gen_value(initial_value, bit_width, &mut w)?;
                    w.append(";")?;
                    w.append_newline()?;
                }
                _ => unreachable!(),
            }
        }

        w.unindent()?;
        w.append_line("}")?;
        w.append_newline()?;

        w.append_line("pub fn posedge_clk(&mut self) {")?;
        w.indent();

        for names in c.reg_names.values() {
            w.append_line(&format!(
                "self.{} = self.{};",
                names.value_name, names.next_name
            ))?;
        }

        w.unindent()?;
        w.append_line("}")?;
        w.append_newline()?;
    }

    w.append_line("#[allow(unused_parens)]")?;
    w.append_line("pub fn prop(&mut self) {")?;
    w.indent();

    for (name, output) in m.outputs().iter() {
        w.append_indent()?;
        w.append(&format!("self.{} = ", name))?;
        gen_expr(&output.source, &c, &mut w)?;
        w.append(";")?;
        w.append_newline()?;
    }

    for (reg, names) in c.reg_names.iter() {
        let reg = unsafe { &**reg as &module::Signal };
        match reg.data {
            module::SignalData::Reg { ref next, .. } => {
                w.append_indent()?;
                w.append(&format!("self.{} = ", names.next_name))?;
                gen_expr(next.borrow().unwrap(), &c, &mut w)?;
                w.append(";")?;
                w.append_newline()?;
            }
            _ => unreachable!(),
        }
    }

    w.unindent()?;
    w.append_line("}")?;

    w.unindent()?;
    w.append_line("}")?;
    w.append_newline()?;

    Ok(())
}

fn gather_regs<'a>(signal: &'a module::Signal<'a>, c: &mut Context<'a>) {
    match signal.data {
        module::SignalData::Lit { .. } | module::SignalData::Input { .. } => (),

        module::SignalData::Reg { ref next, .. } => {
            if c.reg_names.contains_key(&(signal as *const _)) {
                return;
            }
            let value_name = format!("__reg{}", c.reg_names.len());
            let next_name = format!("{}_next", value_name);
            c.reg_names.insert(
                signal,
                RegNames {
                    value_name,
                    next_name,
                },
            );
            gather_regs(next.borrow().expect("Discovered undriven register(s)"), c);
            // TODO: Proper error and test(s)
        }

        module::SignalData::UnOp { source, .. } => {
            gather_regs(source, c);
        }
        module::SignalData::BinOp { lhs, rhs, .. } => {
            gather_regs(lhs, c);
            gather_regs(rhs, c);
        }

        module::SignalData::Bit { source, .. } => {
            gather_regs(source, c);
        }
        module::SignalData::Bits { source, .. } => {
            gather_regs(source, c);
        }

        module::SignalData::Mux { a, b, sel } => {
            gather_regs(sel, c);
            gather_regs(b, c);
            gather_regs(a, c);
        }
    }
}

fn gen_expr<'a, W: Write>(
    signal: &'a module::Signal<'a>,
    c: &Context<'a>,
    w: &mut code_writer::CodeWriter<W>,
) -> Result<(), code_writer::Error> {
    match signal.data {
        module::SignalData::Lit {
            ref value,
            bit_width,
        } => {
            gen_value(value, bit_width, w)?;
        }

        module::SignalData::Input {
            ref name,
            bit_width,
        } => {
            let target_type = ValueType::from_bit_width(bit_width);
            if target_type.bit_width() > 1 && target_type.bit_width() != bit_width {
                w.append("(")?;
            }
            w.append(&format!("self.{}", name))?;
            if target_type.bit_width() > 1 && target_type.bit_width() != bit_width {
                w.append(&format!(" & 0x{:x}", (1u128 << bit_width) - 1))?;
                w.append(")")?;
            }
        }

        module::SignalData::Reg { .. } => {
            w.append(&format!(
                "self.{}",
                c.reg_names[&(signal as *const _)].value_name
            ))?;
        }

        module::SignalData::UnOp { source, op } => {
            let bit_width = source.bit_width();
            let target_type = ValueType::from_bit_width(bit_width);
            if target_type.bit_width() > 1 && target_type.bit_width() != bit_width {
                w.append("(")?;
            }
            w.append(match op {
                module::UnOp::Not => "!",
            })?;
            gen_expr(source, c, w)?;
            if target_type.bit_width() > 1 && target_type.bit_width() != bit_width {
                w.append(&format!(" & 0x{:x}", (1u128 << bit_width) - 1))?;
                w.append(")")?;
            }
        }
        module::SignalData::BinOp { lhs, rhs, op } => {
            w.append("(")?;
            gen_expr(lhs, c, w)?;
            w.append(&format!(
                " {} ",
                match op {
                    module::BinOp::BitAnd => '&',
                    module::BinOp::BitOr => '|',
                }
            ))?;
            gen_expr(rhs, c, w)?;
            w.append(")")?;
        }

        module::SignalData::Bit { source, index } => {
            let bit_width = source.bit_width();
            let source_type = ValueType::from_bit_width(bit_width);
            if source_type.bit_width() > 1 {
                w.append("((")?;
            }
            gen_expr(source, c, w)?;
            if source_type.bit_width() > 1 {
                w.append(&format!(" >> {}) & 1) == 1", index))?;
            }
        }
        module::SignalData::Bits {
            source, range_low, ..
        } => {
            let source_bit_width = source.bit_width();
            let source_type = ValueType::from_bit_width(source_bit_width);
            let target_bit_width = signal.bit_width();
            let target_type = ValueType::from_bit_width(target_bit_width);
            if source_type.bit_width() > target_type.bit_width() {
                w.append("(")?;
            }
            if target_bit_width < source_bit_width
                && (target_type.bit_width() == 1 || target_type.bit_width() != target_bit_width)
            {
                w.append("(")?;
            }
            gen_expr(source, c, w)?;
            if range_low > 0 {
                w.append(&format!(" >> {}", range_low))?;
            }
            if target_bit_width < source_bit_width
                && (target_type.bit_width() == 1 || target_type.bit_width() != target_bit_width)
            {
                // TODO: It's probably better to do this masking before the type cast for more than one bit
                w.append(&format!(") & {}", (1u128 << target_bit_width) - 1))?;
            }
            if source_type.bit_width() > target_type.bit_width() {
                w.append(")")?;
                if target_type.bit_width() == 1 {
                    w.append(" == 1")?;
                } else {
                    w.append(&format!(" as {}", target_type.name()))?;
                }
            }
        }

        module::SignalData::Mux { a, b, sel } => {
            w.append("if ")?;
            gen_expr(sel, c, w)?;
            w.append(" { ")?;
            gen_expr(b, c, w)?;
            w.append(" } else { ")?;
            gen_expr(a, c, w)?;
            w.append(" }")?;
        }
    }

    Ok(())
}

fn gen_value<W: Write>(
    value: &module::Value,
    bit_width: u32,
    w: &mut code_writer::CodeWriter<W>,
) -> Result<(), code_writer::Error> {
    let type_name = ValueType::from_bit_width(bit_width).name();
    w.append(&match value {
        module::Value::Bool(value) => {
            if bit_width == 1 {
                format!("{}", value)
            } else {
                format!("0x{:x}{}", if *value { 1 } else { 0 }, type_name)
            }
        }
        module::Value::U8(value) => format!("0x{:x}{}", value, type_name),
        module::Value::U16(value) => format!("0x{:x}{}", value, type_name),
        module::Value::U32(value) => format!("0x{:x}{}", value, type_name),
        module::Value::U64(value) => format!("0x{:x}{}", value, type_name),
        module::Value::U128(value) => format!("0x{:x}{}", value, type_name),
    })
}
<commit_msg>Always specify bitmasks in hex for sim gen<commit_after>//! Rust simulator code generation.

use crate::code_writer;
use crate::module;

use std::collections::HashMap;
use std::io::Write;

struct RegNames {
    value_name: String,
    next_name: String,
}

struct Context<'a> {
    reg_names: HashMap<*const module::Signal<'a>, RegNames>,
}

enum ValueType {
    Bool,
    U32,
    U64,
    U128,
}

impl ValueType {
    fn from_bit_width(bit_width: u32) -> ValueType {
        if bit_width == 1 {
            ValueType::Bool
        } else if bit_width <= 32 {
            ValueType::U32
        } else if bit_width <= 64 {
            ValueType::U64
        } else if bit_width <= 128 {
            ValueType::U128
        } else {
            unreachable!()
        }
    }

    fn name(&self) -> &'static str {
        match self {
            ValueType::Bool => "bool",
            ValueType::U32 => "u32",
            ValueType::U64 => "u64",
            ValueType::U128 => "u128",
        }
    }

    fn bit_width(&self) -> u32 {
        match self {
            ValueType::Bool => 1,
            ValueType::U32 => 32,
            ValueType::U64 => 64,
            ValueType::U128 => 128,
        }
    }
}

pub fn generate<W: Write>(m: &module::Module, w: &mut W) -> Result<(), code_writer::Error> {
    let mut c = Context {
        reg_names: HashMap::new(),
    };

    for (_, output) in m.outputs().iter() {
        gather_regs(&output.source, &mut c);
    }

    let mut w = code_writer::CodeWriter::new(w);

    w.append_line("#[allow(non_camel_case_types)]")?;
    w.append_line("#[derive(Default)]")?;
    w.append_line(&format!("pub struct {} {{", m.name))?;
    w.indent();

    let inputs = m.inputs();
    if inputs.len() > 0 {
        w.append_line("// Inputs")?;
        for (name, input) in inputs.iter() {
            w.append_line(&format!(
                "pub {}: {}, // {} bit(s)",
                name,
                ValueType::from_bit_width(input.bit_width()).name(),
                input.bit_width()
            ))?;
        }
    }

    let outputs = m.outputs();
    if outputs.len() > 0 {
        w.append_line("// Outputs")?;
        for (name, output) in outputs.iter() {
            w.append_line(&format!(
                "pub {}: {}, // {} bit(s)",
                name,
                ValueType::from_bit_width(output.source.bit_width()).name(),
                output.source.bit_width()
            ))?;
        }
    }

    if c.reg_names.len() > 0 {
        w.append_line("// Regs")?;
        for (reg, names) in c.reg_names.iter() {
            let reg = unsafe { &**reg as &module::Signal };
            let type_name = ValueType::from_bit_width(reg.bit_width()).name();
            w.append_line(&format!(
                "{}: {}, // {} bit(s)",
                names.value_name,
                type_name,
                reg.bit_width()
            ))?;
            w.append_line(&format!("{}: {},", names.next_name, type_name))?;
        }
    }

    w.unindent()?;
    w.append_line("}")?;
    w.append_newline()?;

    w.append_line(&format!("impl {} {{", m.name))?;
    w.indent();

    if c.reg_names.len() > 0 {
        w.append_line("pub fn reset(&mut self) {")?;
        w.indent();

        for (reg, names) in c.reg_names.iter() {
            let reg = unsafe { &**reg as &module::Signal };
            match reg.data {
                module::SignalData::Reg {
                    ref initial_value,
                    bit_width,
                    ..
                } => {
                    w.append_indent()?;
                    w.append(&format!("self.{} = ", names.value_name))?;
                    gen_value(initial_value, bit_width, &mut w)?;
                    w.append(";")?;
                    w.append_newline()?;
                }
                _ => unreachable!(),
            }
        }

        w.unindent()?;
        w.append_line("}")?;
        w.append_newline()?;

        w.append_line("pub fn posedge_clk(&mut self) {")?;
        w.indent();

        for names in c.reg_names.values() {
            w.append_line(&format!(
                "self.{} = self.{};",
                names.value_name, names.next_name
            ))?;
        }

        w.unindent()?;
        w.append_line("}")?;
        w.append_newline()?;
    }

    w.append_line("#[allow(unused_parens)]")?;
    w.append_line("pub fn prop(&mut self) {")?;
    w.indent();

    for (name, output) in m.outputs().iter() {
        w.append_indent()?;
        w.append(&format!("self.{} = ", name))?;
        gen_expr(&output.source, &c, &mut w)?;
        w.append(";")?;
        w.append_newline()?;
    }

    for (reg, names) in c.reg_names.iter() {
        let reg = unsafe { &**reg as &module::Signal };
        match reg.data {
            module::SignalData::Reg { ref next, .. } => {
                w.append_indent()?;
                w.append(&format!("self.{} = ", names.next_name))?;
                gen_expr(next.borrow().unwrap(), &c, &mut w)?;
                w.append(";")?;
                w.append_newline()?;
            }
            _ => unreachable!(),
        }
    }

    w.unindent()?;
    w.append_line("}")?;

    w.unindent()?;
    w.append_line("}")?;
    w.append_newline()?;

    Ok(())
}

fn gather_regs<'a>(signal: &'a module::Signal<'a>, c: &mut Context<'a>) {
    match signal.data {
        module::SignalData::Lit { .. } | module::SignalData::Input { .. } => (),

        module::SignalData::Reg { ref next, .. } => {
            if c.reg_names.contains_key(&(signal as *const _)) {
                return;
            }
            let value_name = format!("__reg{}", c.reg_names.len());
            let next_name = format!("{}_next", value_name);
            c.reg_names.insert(
                signal,
                RegNames {
                    value_name,
                    next_name,
                },
            );
            gather_regs(next.borrow().expect("Discovered undriven register(s)"), c);
            // TODO: Proper error and test(s)
        }

        module::SignalData::UnOp { source, .. } => {
            gather_regs(source, c);
        }
        module::SignalData::BinOp { lhs, rhs, .. } => {
            gather_regs(lhs, c);
            gather_regs(rhs, c);
        }

        module::SignalData::Bit { source, .. } => {
            gather_regs(source, c);
        }
        module::SignalData::Bits { source, .. } => {
            gather_regs(source, c);
        }

        module::SignalData::Mux { a, b, sel } => {
            gather_regs(sel, c);
            gather_regs(b, c);
            gather_regs(a, c);
        }
    }
}

fn gen_expr<'a, W: Write>(
    signal: &'a module::Signal<'a>,
    c: &Context<'a>,
    w: &mut code_writer::CodeWriter<W>,
) -> Result<(), code_writer::Error> {
    match signal.data {
        module::SignalData::Lit {
            ref value,
            bit_width,
        } => {
            gen_value(value, bit_width, w)?;
        }

        module::SignalData::Input {
            ref name,
            bit_width,
        } => {
            let target_type = ValueType::from_bit_width(bit_width);
            if target_type.bit_width() > 1 && target_type.bit_width() != bit_width {
                w.append("(")?;
            }
            w.append(&format!("self.{}", name))?;
            if target_type.bit_width() > 1 && target_type.bit_width() != bit_width {
                w.append(&format!(" & 0x{:x}", (1u128 << bit_width) - 1))?;
                w.append(")")?;
            }
        }

        module::SignalData::Reg { .. } => {
            w.append(&format!(
                "self.{}",
                c.reg_names[&(signal as *const _)].value_name
            ))?;
        }

        module::SignalData::UnOp { source, op } => {
            let bit_width = source.bit_width();
            let target_type = ValueType::from_bit_width(bit_width);
            if target_type.bit_width() > 1 && target_type.bit_width() != bit_width {
                w.append("(")?;
            }
            w.append(match op {
                module::UnOp::Not => "!",
            })?;
            gen_expr(source, c, w)?;
            if target_type.bit_width() > 1 && target_type.bit_width() != bit_width {
                w.append(&format!(" & 0x{:x}", (1u128 << bit_width) - 1))?;
                w.append(")")?;
            }
        }
        module::SignalData::BinOp { lhs, rhs, op } => {
            w.append("(")?;
            gen_expr(lhs, c, w)?;
            w.append(&format!(
                " {} ",
                match op {
                    module::BinOp::BitAnd => '&',
                    module::BinOp::BitOr => '|',
                }
            ))?;
            gen_expr(rhs, c, w)?;
            w.append(")")?;
        }

        module::SignalData::Bit { source, index } => {
            let bit_width = source.bit_width();
            let source_type = ValueType::from_bit_width(bit_width);
            if source_type.bit_width() > 1 {
                w.append("((")?;
            }
            gen_expr(source, c, w)?;
            if source_type.bit_width() > 1 {
                w.append(&format!(" >> {}) & 0x1) == 1", index))?;
            }
        }
        module::SignalData::Bits {
            source, range_low, ..
        } => {
            let source_bit_width = source.bit_width();
            let source_type = ValueType::from_bit_width(source_bit_width);
            let target_bit_width = signal.bit_width();
            let target_type = ValueType::from_bit_width(target_bit_width);
            if source_type.bit_width() > target_type.bit_width() {
                w.append("(")?;
            }
            if target_bit_width < source_bit_width
                && (target_type.bit_width() == 1 || target_type.bit_width() != target_bit_width)
            {
                w.append("(")?;
            }
            gen_expr(source, c, w)?;
            if range_low > 0 {
                w.append(&format!(" >> {}", range_low))?;
            }
            if target_bit_width < source_bit_width
                && (target_type.bit_width() == 1 || target_type.bit_width() != target_bit_width)
            {
                // TODO: It's probably better to do this masking before the type cast for more than one bit
                w.append(&format!(") & 0x{:x}", (1u128 << target_bit_width) - 1))?;
            }
            if source_type.bit_width() > target_type.bit_width() {
                w.append(")")?;
                if target_type.bit_width() == 1 {
                    w.append(" == 1")?;
                } else {
                    w.append(&format!(" as {}", target_type.name()))?;
                }
            }
        }

        module::SignalData::Mux { a, b, sel } => {
            w.append("if ")?;
            gen_expr(sel, c, w)?;
            w.append(" { ")?;
            gen_expr(b, c, w)?;
            w.append(" } else { ")?;
            gen_expr(a, c, w)?;
            w.append(" }")?;
        }
    }

    Ok(())
}

fn gen_value<W: Write>(
    value: &module::Value,
    bit_width: u32,
    w: &mut code_writer::CodeWriter<W>,
) -> Result<(), code_writer::Error> {
    let type_name = ValueType::from_bit_width(bit_width).name();
    w.append(&match value {
        module::Value::Bool(value) => {
            if bit_width == 1 {
                format!("{}", value)
            } else {
                format!("0x{:x}{}", if *value { 1 } else { 0 }, type_name)
            }
        }
        module::Value::U8(value) => format!("0x{:x}{}", value, type_name),
        module::Value::U16(value) => format!("0x{:x}{}", value, type_name),
        module::Value::U32(value) => format!("0x{:x}{}", value, type_name),
        module::Value::U64(value) => format!("0x{:x}{}", value, type_name),
        module::Value::U128(value) => format!("0x{:x}{}", value, type_name),
    })
}
<|endoftext|>


<commit_before>use eyre::Result;
use graph_triples::{Relation, Resource};
use node_address::{Address};
use std::{
    collections::HashMap,
    path::{Path, PathBuf},
};
use stencila_schema::*;

// Re-exports
pub use kernels::{KernelSelector, KernelSpace, TaskResult};

mod executable;
pub use executable::*;

type Addresses = HashMap<String, Address>;
type Relations = HashMap<Resource, Vec<(Relation, Resource)>>;

/// Compile a node
///
/// Compiling a document involves walking over the node tree and compiling each
/// individual node so that it is ready to be built & executed. This includes
/// (but is not limited to):
///
/// - for those node types needing to be accesses directly (e.g. executable nodes) ensuring
///   they have an `id` and recording their address
/// - for executable nodes (e.g. `CodeChunk`) performing semantic analysis of the code
/// - determining dependencies within and between documents and other resources
#[tracing::instrument(skip(node))]
pub fn compile(node: &mut Node, path: &Path, project: &Path) -> Result<(Addresses, Relations)> {
    let mut address = Address::default();
    let mut context = CompileContext {
        path: PathBuf::from(path),
        project: PathBuf::from(project),
        ..Default::default()
    };
    node.compile(&mut address, &mut context)?;

    let addresses = context.addresses;
    let relations = context.relations.into_iter().collect();
    Ok((addresses, relations))
}

/// Execute a node
#[tracing::instrument(skip(node, kernels))]
pub async fn execute<Type>(node: &mut Type, kernels: &mut KernelSpace) -> Result<()>
where
    Type: Executable + Send,
{
    node.execute(kernels).await
}
<commit_msg>refactor(Node execution): Use exported type<commit_after>use eyre::Result;
use graph_triples::Relations;
use node_address::Address;
use std::{
    collections::HashMap,
    path::{Path, PathBuf},
};
use stencila_schema::*;

// Re-exports
pub use kernels::{KernelSelector, KernelSpace, TaskResult};

mod executable;
pub use executable::*;

/// A map of node ids to their address
///
/// Used to enable faster access to a node based on it's id.
type Addresses = HashMap<String, Address>;

/// Compile a node
///
/// Compiling a document involves walking over the node tree and compiling each
/// individual node so that it is ready to be built & executed. This includes
/// (but is not limited to):
///
/// - for those node types needing to be accesses directly (e.g. executable nodes) ensuring
///   they have an `id` and recording their address
/// - for executable nodes (e.g. `CodeChunk`) performing semantic analysis of the code
/// - determining dependencies within and between documents and other resources
#[tracing::instrument(skip(node))]
pub fn compile(node: &mut Node, path: &Path, project: &Path) -> Result<(Addresses, Relations)> {
    let mut address = Address::default();
    let mut context = CompileContext {
        path: PathBuf::from(path),
        project: PathBuf::from(project),
        ..Default::default()
    };
    node.compile(&mut address, &mut context)?;

    let addresses = context.addresses;
    let relations = context.relations.into_iter().collect();
    Ok((addresses, relations))
}

/// Execute a node
#[tracing::instrument(skip(node, kernels))]
pub async fn execute<Type>(node: &mut Type, kernels: &mut KernelSpace) -> Result<()>
where
    Type: Executable + Send,
{
    node.execute(kernels).await
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fmt::{self, Formatter, Display};
use std::sync::Arc;
use std::sync::mpsc::SyncSender;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Instant;
use std::str;

use rocksdb::{DB, Writable, WriteBatch};
use kvproto::raft_serverpb::{RaftApplyState, RegionLocalState, PeerState};
use kvproto::eraftpb::Snapshot as RaftSnapshot;

use util::worker::Runnable;
use util::{escape, rocksdb};
use raftstore::store::engine::{Mutable, Snapshot, Iterable, IterOption};
use raftstore::store::peer_storage::{JOB_STATUS_FINISHED, JOB_STATUS_CANCELLED, JOB_STATUS_FAILED,
                                     JOB_STATUS_CANCELLING, JOB_STATUS_PENDING, JOB_STATUS_RUNNING};
use raftstore::store::{self, check_abort, SnapManager, SnapKey, SnapEntry, ApplyOptions, keys,
                       Peekable};
use raftstore::store::snap::{Error, Result};
use storage::CF_RAFT;

use super::metrics::*;

/// region related task.
pub enum Task {
    Gen {
        region_id: u64,
        notifier: SyncSender<RaftSnapshot>,
    },
    Apply {
        region_id: u64,
        status: Arc<AtomicUsize>,
    },
    /// Destroy data between [start_key, end_key).
    ///
    /// The deletion may and may not succeed.
    Destroy {
        region_id: u64,
        start_key: Vec<u8>,
        end_key: Vec<u8>,
    },
}

impl Task {
    pub fn destroy(region_id: u64, start_key: Vec<u8>, end_key: Vec<u8>) -> Task {
        Task::Destroy {
            region_id: region_id,
            start_key: start_key,
            end_key: end_key,
        }
    }
}

impl Display for Task {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        match *self {
            Task::Gen { region_id, .. } => write!(f, "Snap gen for {}", region_id),
            Task::Apply { region_id, .. } => write!(f, "Snap apply for {}", region_id),
            Task::Destroy { region_id, ref start_key, ref end_key } => {
                write!(f,
                       "Destroy {} [{}, {})",
                       region_id,
                       escape(start_key),
                       escape(end_key))
            }
        }
    }
}

// TODO: use threadpool to do task concurrently
pub struct Runner {
    db: Arc<DB>,
    batch_size: usize,
    mgr: SnapManager,
}

impl Runner {
    pub fn new(db: Arc<DB>, mgr: SnapManager, batch_size: usize) -> Runner {
        Runner {
            db: db,
            mgr: mgr,
            batch_size: batch_size,
        }
    }

    fn generate_snap(&self, region_id: u64, notifier: SyncSender<RaftSnapshot>) -> Result<()> {
        // do we need to check leader here?
        let raw_snap = Snapshot::new(self.db.clone());

        let snap = box_try!(store::do_snapshot(self.mgr.clone(), &raw_snap, region_id));
        if let Err(e) = notifier.try_send(snap) {
            info!("[region {}] failed to notify snap result, maybe leadership has changed, \
                   ignore: {:?}",
                  region_id,
                  e);
        }
        Ok(())
    }

    fn handle_gen(&self, region_id: u64, notifier: SyncSender<RaftSnapshot>) {
        SNAP_COUNTER_VEC.with_label_values(&["generate", "all"]).inc();
        let gen_histogram = SNAP_HISTOGRAM.with_label_values(&["generate"]);
        let timer = gen_histogram.start_timer();

        if let Err(e) = self.generate_snap(region_id, notifier) {
            error!("[region {}] failed to generate snap: {:?}!!!", region_id, e);
            return;
        }

        SNAP_COUNTER_VEC.with_label_values(&["generate", "success"]).inc();
        timer.observe_duration();
    }

    fn delete_all_in_range(&self,
                           start_key: &[u8],
                           end_key: &[u8],
                           abort: &AtomicUsize)
                           -> Result<()> {
        let mut wb = WriteBatch::new();
        let mut size_cnt = 0;
        for cf in self.db.cf_names() {
            try!(check_abort(abort));
            let handle = box_try!(rocksdb::get_cf_handle(&self.db, cf));

            let iter_opt = IterOption::new(Some(end_key.to_vec()), false);
            let mut it = box_try!(self.db.new_iterator_cf(cf, iter_opt));

            try!(check_abort(abort));
            it.seek(start_key.into());
            while it.valid() {
                {
                    let key = it.key();
                    if key >= end_key {
                        break;
                    }

                    box_try!(wb.delete_cf(handle, key));
                    size_cnt += key.len();
                    if size_cnt >= self.batch_size {
                        // Can't use write_without_wal here.
                        // Otherwise it may cause dirty data when applying snapshot.
                        box_try!(self.db.write(wb));
                        wb = WriteBatch::new();
                        size_cnt = 0;
                    }
                };
                try!(check_abort(abort));
                if !it.next() {
                    break;
                }
            }
        }

        if wb.count() > 0 {
            box_try!(self.db.write(wb));
        }
        Ok(())
    }

    fn apply_snap(&self, region_id: u64, abort: Arc<AtomicUsize>) -> Result<()> {
        info!("[region {}] begin apply snap data", region_id);
        try!(check_abort(&abort));
        let region_key = keys::region_state_key(region_id);
        let mut region_state: RegionLocalState = match box_try!(self.db.get_msg(&region_key)) {
            Some(state) => state,
            None => return Err(box_err!("failed to get region_state from {}", escape(&region_key))),
        };

        // clear up origin data.
        let region = region_state.get_region().clone();
        let start_key = keys::enc_start_key(&region);
        let end_key = keys::enc_end_key(&region);
        box_try!(self.delete_all_in_range(&start_key, &end_key, &abort));

        let state_key = keys::apply_state_key(region_id);
        let apply_state: RaftApplyState = match box_try!(self.db.get_msg_cf(CF_RAFT, &state_key)) {
            Some(state) => state,
            None => return Err(box_err!("failed to get raftstate from {}", escape(&state_key))),
        };
        let term = apply_state.get_truncated_state().get_term();
        let idx = apply_state.get_truncated_state().get_index();
        let snap_key = SnapKey::new(region_id, term, idx);
        self.mgr.register(snap_key.clone(), SnapEntry::Applying);
        defer!({
            self.mgr.deregister(&snap_key, &SnapEntry::Applying);
        });
        let mut s = box_try!(self.mgr.get_snapshot_for_applying(&snap_key));
        if !s.exists() {
            return Err(box_err!("missing snapshot file {}", s.path()));
        }
        try!(check_abort(&abort));
        let timer = Instant::now();
        let options = ApplyOptions {
            db: self.db.clone(),
            region: region.clone(),
            abort: abort.clone(),
            write_batch_size: self.batch_size,
        };
        try!(s.apply(options));
        region_state.set_state(PeerState::Normal);
        box_try!(self.db.put_msg(&region_key, &region_state));
        info!("[region {}] apply new data takes {:?}",
              region_id,
              timer.elapsed());
        Ok(())
    }

    fn handle_apply(&self, region_id: u64, status: Arc<AtomicUsize>) {
        status.compare_and_swap(JOB_STATUS_PENDING, JOB_STATUS_RUNNING, Ordering::SeqCst);
        SNAP_COUNTER_VEC.with_label_values(&["apply", "all"]).inc();
        let apply_histogram = SNAP_HISTOGRAM.with_label_values(&["apply"]);
        let timer = apply_histogram.start_timer();

        match self.apply_snap(region_id, status.clone()) {
            Ok(()) => {
                status.swap(JOB_STATUS_FINISHED, Ordering::SeqCst);
                SNAP_COUNTER_VEC.with_label_values(&["apply", "success"]).inc();
            }
            Err(Error::Abort) => {
                warn!("applying snapshot for region {} is aborted.", region_id);
                assert_eq!(status.swap(JOB_STATUS_CANCELLED, Ordering::SeqCst),
                           JOB_STATUS_CANCELLING);
                SNAP_COUNTER_VEC.with_label_values(&["apply", "abort"]).inc();
            }
            Err(e) => {
                error!("failed to apply snap: {:?}!!!", e);
                status.swap(JOB_STATUS_FAILED, Ordering::SeqCst);
                SNAP_COUNTER_VEC.with_label_values(&["apply", "fail"]).inc();
            }
        }

        timer.observe_duration();
    }

    fn handle_destroy(&mut self, region_id: u64, start_key: Vec<u8>, end_key: Vec<u8>) {
        info!("[region {}] deleting data in [{}, {})",
              region_id,
              escape(&start_key),
              escape(&end_key));
        let status = AtomicUsize::new(JOB_STATUS_PENDING);
        if let Err(e) = self.delete_all_in_range(&start_key, &end_key, &status) {
            error!("failed to delete data in [{}, {}): {:?}",
                   escape(&start_key),
                   escape(&end_key),
                   e);
        }
    }
}

impl Runnable<Task> for Runner {
    fn run(&mut self, task: Task) {
        match task {
            Task::Gen { region_id, notifier } => self.handle_gen(region_id, notifier),
            Task::Apply { region_id, status } => self.handle_apply(region_id, status),
            Task::Destroy { region_id, start_key, end_key } => {
                self.handle_destroy(region_id, start_key, end_key)
            }
        }
    }
}
<commit_msg>raftstore: generate snap concurrently (#1963)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fmt::{self, Formatter, Display};
use std::sync::Arc;
use std::sync::mpsc::SyncSender;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Instant;
use std::str;

use rocksdb::{DB, Writable, WriteBatch};
use kvproto::raft_serverpb::{RaftApplyState, RegionLocalState, PeerState};
use kvproto::eraftpb::Snapshot as RaftSnapshot;
use threadpool::ThreadPool;

use util::worker::Runnable;
use util::{escape, rocksdb};
use raftstore::store::engine::{Mutable, Snapshot, Iterable, IterOption};
use raftstore::store::peer_storage::{JOB_STATUS_FINISHED, JOB_STATUS_CANCELLED, JOB_STATUS_FAILED,
                                     JOB_STATUS_CANCELLING, JOB_STATUS_PENDING, JOB_STATUS_RUNNING};
use raftstore::store::{self, check_abort, SnapManager, SnapKey, SnapEntry, ApplyOptions, keys,
                       Peekable};
use raftstore::store::snap::{Error, Result};
use storage::CF_RAFT;

use super::metrics::*;

const GENERATE_POOL_SIZE: usize = 2;

/// region related task.
pub enum Task {
    Gen {
        region_id: u64,
        notifier: SyncSender<RaftSnapshot>,
    },
    Apply {
        region_id: u64,
        status: Arc<AtomicUsize>,
    },
    /// Destroy data between [start_key, end_key).
    ///
    /// The deletion may and may not succeed.
    Destroy {
        region_id: u64,
        start_key: Vec<u8>,
        end_key: Vec<u8>,
    },
}

impl Task {
    pub fn destroy(region_id: u64, start_key: Vec<u8>, end_key: Vec<u8>) -> Task {
        Task::Destroy {
            region_id: region_id,
            start_key: start_key,
            end_key: end_key,
        }
    }
}

impl Display for Task {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        match *self {
            Task::Gen { region_id, .. } => write!(f, "Snap gen for {}", region_id),
            Task::Apply { region_id, .. } => write!(f, "Snap apply for {}", region_id),
            Task::Destroy { region_id, ref start_key, ref end_key } => {
                write!(f,
                       "Destroy {} [{}, {})",
                       region_id,
                       escape(start_key),
                       escape(end_key))
            }
        }
    }
}

#[derive(Clone)]
struct SnapContext {
    db: Arc<DB>,
    batch_size: usize,
    mgr: SnapManager,
}

impl SnapContext {
    fn generate_snap(&self, region_id: u64, notifier: SyncSender<RaftSnapshot>) -> Result<()> {
        // do we need to check leader here?
        let raw_snap = Snapshot::new(self.db.clone());

        let snap = box_try!(store::do_snapshot(self.mgr.clone(), &raw_snap, region_id));
        if let Err(e) = notifier.try_send(snap) {
            info!("[region {}] failed to notify snap result, maybe leadership has changed, \
                   ignore: {:?}",
                  region_id,
                  e);
        }
        Ok(())
    }

    fn handle_gen(&self, region_id: u64, notifier: SyncSender<RaftSnapshot>) {
        SNAP_COUNTER_VEC.with_label_values(&["generate", "all"]).inc();
        let gen_histogram = SNAP_HISTOGRAM.with_label_values(&["generate"]);
        let timer = gen_histogram.start_timer();

        if let Err(e) = self.generate_snap(region_id, notifier) {
            error!("[region {}] failed to generate snap: {:?}!!!", region_id, e);
            return;
        }

        SNAP_COUNTER_VEC.with_label_values(&["generate", "success"]).inc();
        timer.observe_duration();
    }

    fn delete_all_in_range(&self,
                           start_key: &[u8],
                           end_key: &[u8],
                           abort: &AtomicUsize)
                           -> Result<()> {
        let mut wb = WriteBatch::new();
        let mut size_cnt = 0;
        for cf in self.db.cf_names() {
            try!(check_abort(abort));
            let handle = box_try!(rocksdb::get_cf_handle(&self.db, cf));

            let iter_opt = IterOption::new(Some(end_key.to_vec()), false);
            let mut it = box_try!(self.db.new_iterator_cf(cf, iter_opt));

            try!(check_abort(abort));
            it.seek(start_key.into());
            while it.valid() {
                {
                    let key = it.key();
                    if key >= end_key {
                        break;
                    }

                    box_try!(wb.delete_cf(handle, key));
                    size_cnt += key.len();
                    if size_cnt >= self.batch_size {
                        // Can't use write_without_wal here.
                        // Otherwise it may cause dirty data when applying snapshot.
                        box_try!(self.db.write(wb));
                        wb = WriteBatch::new();
                        size_cnt = 0;
                    }
                };
                try!(check_abort(abort));
                if !it.next() {
                    break;
                }
            }
        }

        if wb.count() > 0 {
            box_try!(self.db.write(wb));
        }
        Ok(())
    }

    fn apply_snap(&self, region_id: u64, abort: Arc<AtomicUsize>) -> Result<()> {
        info!("[region {}] begin apply snap data", region_id);
        try!(check_abort(&abort));
        let region_key = keys::region_state_key(region_id);
        let mut region_state: RegionLocalState = match box_try!(self.db.get_msg(&region_key)) {
            Some(state) => state,
            None => return Err(box_err!("failed to get region_state from {}", escape(&region_key))),
        };

        // clear up origin data.
        let region = region_state.get_region().clone();
        let start_key = keys::enc_start_key(&region);
        let end_key = keys::enc_end_key(&region);
        box_try!(self.delete_all_in_range(&start_key, &end_key, &abort));

        let state_key = keys::apply_state_key(region_id);
        let apply_state: RaftApplyState = match box_try!(self.db.get_msg_cf(CF_RAFT, &state_key)) {
            Some(state) => state,
            None => return Err(box_err!("failed to get raftstate from {}", escape(&state_key))),
        };
        let term = apply_state.get_truncated_state().get_term();
        let idx = apply_state.get_truncated_state().get_index();
        let snap_key = SnapKey::new(region_id, term, idx);
        self.mgr.register(snap_key.clone(), SnapEntry::Applying);
        defer!({
            self.mgr.deregister(&snap_key, &SnapEntry::Applying);
        });
        let mut s = box_try!(self.mgr.get_snapshot_for_applying(&snap_key));
        if !s.exists() {
            return Err(box_err!("missing snapshot file {}", s.path()));
        }
        try!(check_abort(&abort));
        let timer = Instant::now();
        let options = ApplyOptions {
            db: self.db.clone(),
            region: region.clone(),
            abort: abort.clone(),
            write_batch_size: self.batch_size,
        };
        try!(s.apply(options));
        region_state.set_state(PeerState::Normal);
        box_try!(self.db.put_msg(&region_key, &region_state));
        info!("[region {}] apply new data takes {:?}",
              region_id,
              timer.elapsed());
        Ok(())
    }

    fn handle_apply(&self, region_id: u64, status: Arc<AtomicUsize>) {
        status.compare_and_swap(JOB_STATUS_PENDING, JOB_STATUS_RUNNING, Ordering::SeqCst);
        SNAP_COUNTER_VEC.with_label_values(&["apply", "all"]).inc();
        let apply_histogram = SNAP_HISTOGRAM.with_label_values(&["apply"]);
        let timer = apply_histogram.start_timer();

        match self.apply_snap(region_id, status.clone()) {
            Ok(()) => {
                status.swap(JOB_STATUS_FINISHED, Ordering::SeqCst);
                SNAP_COUNTER_VEC.with_label_values(&["apply", "success"]).inc();
            }
            Err(Error::Abort) => {
                warn!("applying snapshot for region {} is aborted.", region_id);
                assert_eq!(status.swap(JOB_STATUS_CANCELLED, Ordering::SeqCst),
                           JOB_STATUS_CANCELLING);
                SNAP_COUNTER_VEC.with_label_values(&["apply", "abort"]).inc();
            }
            Err(e) => {
                error!("failed to apply snap: {:?}!!!", e);
                status.swap(JOB_STATUS_FAILED, Ordering::SeqCst);
                SNAP_COUNTER_VEC.with_label_values(&["apply", "fail"]).inc();
            }
        }

        timer.observe_duration();
    }

    fn handle_destroy(&mut self, region_id: u64, start_key: Vec<u8>, end_key: Vec<u8>) {
        info!("[region {}] deleting data in [{}, {})",
              region_id,
              escape(&start_key),
              escape(&end_key));
        let status = AtomicUsize::new(JOB_STATUS_PENDING);
        if let Err(e) = self.delete_all_in_range(&start_key, &end_key, &status) {
            error!("failed to delete data in [{}, {}): {:?}",
                   escape(&start_key),
                   escape(&end_key),
                   e);
        }
    }
}

pub struct Runner {
    pool: ThreadPool,
    ctx: SnapContext,
}

impl Runner {
    pub fn new(db: Arc<DB>, mgr: SnapManager, batch_size: usize) -> Runner {
        Runner {
            pool: ThreadPool::new_with_name(thd_name!("snap generator"), GENERATE_POOL_SIZE),
            ctx: SnapContext {
                db: db,
                mgr: mgr,
                batch_size: batch_size,
            },
        }
    }
}

impl Runnable<Task> for Runner {
    fn run(&mut self, task: Task) {
        match task {
            Task::Gen { region_id, notifier } => {
                // It safe for now to handle generating and applying snapshot concurrently,
                // but it may not when merge is implemented.
                let ctx = self.ctx.clone();
                self.pool.execute(move || ctx.handle_gen(region_id, notifier))
            }
            Task::Apply { region_id, status } => self.ctx.handle_apply(region_id, status),
            Task::Destroy { region_id, start_key, end_key } => {
                self.ctx.handle_destroy(region_id, start_key, end_key)
            }
        }
    }
}
<|endoftext|>


<commit_before>use parser::ast::*;
use parser::ast::Elem::*;
use parser::ast::Expr::*;
use parser::ast::Stmt::*;

use error::ParseError;

pub trait Visitor<'v> : Sized {
    fn visit_ast(&mut self, a: &'v Ast) {
        walk_ast(self, a)
    }

    fn visit_fct(&mut self, a: &'v Function) {
        walk_stmt(self, &a.block)
    }

    fn visit_param(&mut self, p: &'v Param) {
        walk_param(self, p);
    }

    fn visit_stmt(&mut self, s: &'v Stmt) {
        walk_stmt(self, s)
    }

    fn visit_expr(&mut self, e: &'v Expr) {
        walk_expr(self, e)
    }
}

pub fn walk_ast<'v, V: Visitor<'v>>(v: &mut V, a: &'v Ast) {
    for e in &a.elements {
        match *e {
            ElemFunction(ref f) => v.visit_fct(f),
            ElemUnknown => unreachable!()
        }
    }
}

pub fn walk_fct<'v, V: Visitor<'v>>(v: &mut V, f: &'v Function) {
    v.visit_stmt(&f.block)
}

pub fn walk_param<'v, V: Visitor<'v>>(v: &mut V, f: &'v Param) {

}

pub fn walk_stmt<'v, V: Visitor<'v>>(v: &mut V, s: &'v Stmt) {
    match *s {
        StmtVar(ref value) => {
            if let Some(ref e) = value.expr {
                v.visit_expr(e);
            }
        }

        StmtWhile(ref value) => {
            v.visit_expr(&value.cond);
            v.visit_stmt(&value.block);
        }

        StmtLoop(ref value) => {
            v.visit_stmt(&value.block);
        }

        StmtIf(ref value) => {
            v.visit_expr(&value.cond);
            v.visit_stmt(&value.then_block);

            if let Some(ref b) = value.else_block {
                v.visit_stmt(b);
            }
        }

        StmtExpr(ref value) => {
            v.visit_expr(&value.expr);
        }

        StmtBlock(ref value) => {
            for stmt in &value.stmts {
                v.visit_stmt(stmt);
            }
        }

        StmtReturn(ref value) => {
            if let Some(ref e) = value.expr {
                v.visit_expr(e);
            }
        }

        StmtBreak(_) => { }
        StmtContinue(_) => { }
    }
}

pub fn walk_expr<'v, V: Visitor<'v>>(v: &mut V, e: &'v Expr) {
    match *e {
        ExprUn(ref value) => {
            v.visit_expr(&value.opnd);
        }

        ExprBin(ref value) => {
            v.visit_expr(&value.lhs);
            v.visit_expr(&value.rhs);
        }

        ExprAssign(ref value) => {
            v.visit_expr(&value.lhs);
            v.visit_expr(&value.rhs);
        }

        ExprLitInt(_) => {}
        ExprLitStr(_) => {}
        ExprLitBool(_) => {}
        ExprIdent(_) => {}
    }
}
<commit_msg>ast: added walk-functions for param and type<commit_after>use parser::ast::*;
use parser::ast::Elem::*;
use parser::ast::Expr::*;
use parser::ast::Stmt::*;
use parser::ast::Type::*;

use error::ParseError;

pub trait Visitor<'v> : Sized {
    fn visit_ast(&mut self, a: &'v Ast) {
        walk_ast(self, a);
    }

    fn visit_fct(&mut self, a: &'v Function) {
        walk_stmt(self, &a.block);
    }

    fn visit_param(&mut self, p: &'v Param) {
        walk_param(self, p);
    }

    fn visit_type(&mut self, t: &'v Type) {
        walk_type(self, t);
    }

    fn visit_stmt(&mut self, s: &'v Stmt) {
        walk_stmt(self, s);
    }

    fn visit_expr(&mut self, e: &'v Expr) {
        walk_expr(self, e);
    }
}

pub fn walk_ast<'v, V: Visitor<'v>>(v: &mut V, a: &'v Ast) {
    for e in &a.elements {
        match *e {
            ElemFunction(ref f) => v.visit_fct(f),
            ElemUnknown => unreachable!()
        }
    }
}

pub fn walk_fct<'v, V: Visitor<'v>>(v: &mut V, f: &'v Function) {
    v.visit_stmt(&f.block)
}

pub fn walk_param<'v, V: Visitor<'v>>(v: &mut V, p: &'v Param) {
    v.visit_type(&p.data_type);
}

pub fn walk_type<'v, V: Visitor<'v>>(v: &mut V, t: &'v Type) {
    match *t {
        TypeBasic(_) => { },
        TypeTuple(ref tuple) => {
            for ty in &tuple.subtypes {
                v.visit_type(ty);
            }
        }

        TypePtr(ref ptr) => {
            v.visit_type(&ptr.subtype);
        }

        TypeArray(ref array) => {
            v.visit_type(&array.subtype);
        }
    }
}

pub fn walk_stmt<'v, V: Visitor<'v>>(v: &mut V, s: &'v Stmt) {
    match *s {
        StmtVar(ref value) => {
            if let Some(ref e) = value.expr {
                v.visit_expr(e);
            }
        }

        StmtWhile(ref value) => {
            v.visit_expr(&value.cond);
            v.visit_stmt(&value.block);
        }

        StmtLoop(ref value) => {
            v.visit_stmt(&value.block);
        }

        StmtIf(ref value) => {
            v.visit_expr(&value.cond);
            v.visit_stmt(&value.then_block);

            if let Some(ref b) = value.else_block {
                v.visit_stmt(b);
            }
        }

        StmtExpr(ref value) => {
            v.visit_expr(&value.expr);
        }

        StmtBlock(ref value) => {
            for stmt in &value.stmts {
                v.visit_stmt(stmt);
            }
        }

        StmtReturn(ref value) => {
            if let Some(ref e) = value.expr {
                v.visit_expr(e);
            }
        }

        StmtBreak(_) => { }
        StmtContinue(_) => { }
    }
}

pub fn walk_expr<'v, V: Visitor<'v>>(v: &mut V, e: &'v Expr) {
    match *e {
        ExprUn(ref value) => {
            v.visit_expr(&value.opnd);
        }

        ExprBin(ref value) => {
            v.visit_expr(&value.lhs);
            v.visit_expr(&value.rhs);
        }

        ExprAssign(ref value) => {
            v.visit_expr(&value.lhs);
            v.visit_expr(&value.rhs);
        }

        ExprLitInt(_) => {}
        ExprLitStr(_) => {}
        ExprLitBool(_) => {}
        ExprIdent(_) => {}
    }
}
<|endoftext|>


<commit_before>use rustc_errors::{Applicability, DiagnosticBuilder};

use rustc_ast::ptr::P;
use rustc_ast::token;
use rustc_ast::tokenstream::{DelimSpan, TokenStream};
use rustc_ast::{self as ast, *};
use rustc_ast_pretty::pprust;
use rustc_expand::base::*;
use rustc_parse::parser::Parser;
use rustc_span::symbol::{sym, Ident, Symbol};
use rustc_span::{Span, DUMMY_SP};

pub fn expand_assert<'cx>(
    cx: &'cx mut ExtCtxt<'_>,
    sp: Span,
    tts: TokenStream,
) -> Box<dyn MacResult + 'cx> {
    let Assert { cond_expr, custom_message } = match parse_assert(cx, sp, tts) {
        Ok(assert) => assert,
        Err(mut err) => {
            err.emit();
            return DummyResult::any(sp);
        }
    };

    // `core::panic` and `std::panic` are different macros, so we use call-site
    // context to pick up whichever is currently in scope.
    let sp = cx.with_call_site_ctxt(sp);

    let panic_call = if let Some(tokens) = custom_message {
        // Pass the custom message to panic!().
        cx.expr(
            sp,
            ExprKind::MacCall(MacCall {
                path: Path::from_ident(Ident::new(sym::panic, sp)),
                args: P(MacArgs::Delimited(
                    DelimSpan::from_single(sp),
                    MacDelimiter::Parenthesis,
                    tokens,
                )),
                prior_type_ascription: None,
            }),
        )
    } else {
        // Pass our own message directly to $crate::panicking::panic(),
        // because it might contain `{` and `}` that should always be
        // passed literally.
        cx.expr_call_global(
            sp,
            cx.std_path(&[sym::panicking, sym::panic]),
            vec![cx.expr_str(
                DUMMY_SP,
                Symbol::intern(&format!(
                    "assertion failed: {}",
                    pprust::expr_to_string(&cond_expr).escape_debug()
                )),
            )],
        )
    };
    let if_expr =
        cx.expr_if(sp, cx.expr(sp, ExprKind::Unary(UnOp::Not, cond_expr)), panic_call, None);
    MacEager::expr(if_expr)
}

struct Assert {
    cond_expr: P<ast::Expr>,
    custom_message: Option<TokenStream>,
}

fn parse_assert<'a>(
    cx: &mut ExtCtxt<'a>,
    sp: Span,
    stream: TokenStream,
) -> Result<Assert, DiagnosticBuilder<'a>> {
    let mut parser = cx.new_parser_from_tts(stream);

    if parser.token == token::Eof {
        let mut err = cx.struct_span_err(sp, "macro requires a boolean expression as an argument");
        err.span_label(sp, "boolean expression required");
        return Err(err);
    }

    let cond_expr = parser.parse_expr()?;

    // Some crates use the `assert!` macro in the following form (note extra semicolon):
    //
    // assert!(
    //     my_function();
    // );
    //
    // Emit an error about semicolon and suggest removing it.
    if parser.token == token::Semi {
        let mut err = cx.struct_span_err(sp, "macro requires an expression as an argument");
        err.span_suggestion(
            parser.token.span,
            "try removing semicolon",
            String::new(),
            Applicability::MaybeIncorrect,
        );
        err.emit();

        parser.bump();
    }

    // Some crates use the `assert!` macro in the following form (note missing comma before
    // message):
    //
    // assert!(true "error message");
    //
    // Emit an error and suggest inserting a comma.
    let custom_message =
        if let token::Literal(token::Lit { kind: token::Str, .. }) = parser.token.kind {
            let mut err = cx.struct_span_err(parser.token.span, "unexpected string literal");
            let comma_span = parser.prev_token.span.shrink_to_hi();
            err.span_suggestion_short(
                comma_span,
                "try adding a comma",
                ", ".to_string(),
                Applicability::MaybeIncorrect,
            );
            err.emit();

            parse_custom_message(&mut parser)
        } else if parser.eat(&token::Comma) {
            parse_custom_message(&mut parser)
        } else {
            None
        };

    if parser.token != token::Eof {
        return parser.unexpected();
    }

    Ok(Assert { cond_expr, custom_message })
}

fn parse_custom_message(parser: &mut Parser<'_>) -> Option<TokenStream> {
    let ts = parser.parse_tokens();
    if !ts.is_empty() { Some(ts) } else { None }
}
<commit_msg>Expand assert!(expr, args..) to include $crate for hygiene on 2021.<commit_after>use rustc_errors::{Applicability, DiagnosticBuilder};

use rustc_ast::ptr::P;
use rustc_ast::token;
use rustc_ast::tokenstream::{DelimSpan, TokenStream};
use rustc_ast::{self as ast, *};
use rustc_ast_pretty::pprust;
use rustc_expand::base::*;
use rustc_parse::parser::Parser;
use rustc_span::symbol::{sym, Ident, Symbol};
use rustc_span::{Span, DUMMY_SP};

pub fn expand_assert<'cx>(
    cx: &'cx mut ExtCtxt<'_>,
    sp: Span,
    tts: TokenStream,
) -> Box<dyn MacResult + 'cx> {
    let Assert { cond_expr, custom_message } = match parse_assert(cx, sp, tts) {
        Ok(assert) => assert,
        Err(mut err) => {
            err.emit();
            return DummyResult::any(sp);
        }
    };

    let is_2021 = sp.rust_2021();

    // `core::panic` and `std::panic` are different macros, so we use call-site
    // context to pick up whichever is currently in scope.
    let sp = cx.with_call_site_ctxt(sp);

    let panic_call = if let Some(tokens) = custom_message {
        let path = if is_2021 {
            // On edition 2021, we always call `$crate::panic!()`.
            Path {
                span: sp,
                segments: cx
                    .std_path(&[sym::panic])
                    .into_iter()
                    .map(|ident| PathSegment::from_ident(ident))
                    .collect(),
                tokens: None,
            }
        } else {
            // Before edition 2021, we call `panic!()` unqualified,
            // such that it calls either `std::panic!()` or `core::panic!()`.
            Path::from_ident(Ident::new(sym::panic, sp))
        };
        // Pass the custom message to panic!().
        cx.expr(
            sp,
            ExprKind::MacCall(MacCall {
                path,
                args: P(MacArgs::Delimited(
                    DelimSpan::from_single(sp),
                    MacDelimiter::Parenthesis,
                    tokens,
                )),
                prior_type_ascription: None,
            }),
        )
    } else {
        // Pass our own message directly to $crate::panicking::panic(),
        // because it might contain `{` and `}` that should always be
        // passed literally.
        cx.expr_call_global(
            sp,
            cx.std_path(&[sym::panicking, sym::panic]),
            vec![cx.expr_str(
                DUMMY_SP,
                Symbol::intern(&format!(
                    "assertion failed: {}",
                    pprust::expr_to_string(&cond_expr).escape_debug()
                )),
            )],
        )
    };
    let if_expr =
        cx.expr_if(sp, cx.expr(sp, ExprKind::Unary(UnOp::Not, cond_expr)), panic_call, None);
    MacEager::expr(if_expr)
}

struct Assert {
    cond_expr: P<ast::Expr>,
    custom_message: Option<TokenStream>,
}

fn parse_assert<'a>(
    cx: &mut ExtCtxt<'a>,
    sp: Span,
    stream: TokenStream,
) -> Result<Assert, DiagnosticBuilder<'a>> {
    let mut parser = cx.new_parser_from_tts(stream);

    if parser.token == token::Eof {
        let mut err = cx.struct_span_err(sp, "macro requires a boolean expression as an argument");
        err.span_label(sp, "boolean expression required");
        return Err(err);
    }

    let cond_expr = parser.parse_expr()?;

    // Some crates use the `assert!` macro in the following form (note extra semicolon):
    //
    // assert!(
    //     my_function();
    // );
    //
    // Emit an error about semicolon and suggest removing it.
    if parser.token == token::Semi {
        let mut err = cx.struct_span_err(sp, "macro requires an expression as an argument");
        err.span_suggestion(
            parser.token.span,
            "try removing semicolon",
            String::new(),
            Applicability::MaybeIncorrect,
        );
        err.emit();

        parser.bump();
    }

    // Some crates use the `assert!` macro in the following form (note missing comma before
    // message):
    //
    // assert!(true "error message");
    //
    // Emit an error and suggest inserting a comma.
    let custom_message =
        if let token::Literal(token::Lit { kind: token::Str, .. }) = parser.token.kind {
            let mut err = cx.struct_span_err(parser.token.span, "unexpected string literal");
            let comma_span = parser.prev_token.span.shrink_to_hi();
            err.span_suggestion_short(
                comma_span,
                "try adding a comma",
                ", ".to_string(),
                Applicability::MaybeIncorrect,
            );
            err.emit();

            parse_custom_message(&mut parser)
        } else if parser.eat(&token::Comma) {
            parse_custom_message(&mut parser)
        } else {
            None
        };

    if parser.token != token::Eof {
        return parser.unexpected();
    }

    Ok(Assert { cond_expr, custom_message })
}

fn parse_custom_message(parser: &mut Parser<'_>) -> Option<TokenStream> {
    let ts = parser.parse_tokens();
    if !ts.is_empty() { Some(ts) } else { None }
}
<|endoftext|>


<commit_before>use std::rc::*;
use super::treenode::*;

///
/// Trait implemented by types that can work as a tree node index
///
pub trait TreeNodeIndex {
    ///
    /// Finds the tree node corresponding to the specified index in the tree
    ///
    fn lookup_index<'a>(&self, parent_node: &'a TreeNode) -> Option<&'a Rc<TreeNode>>;
}

impl TreeNodeIndex for usize {
    ///
    /// Finds the tree node corresponding to the specified index in the tree
    ///
    fn lookup_index<'a>(&self, parent_node: &'a TreeNode) -> Option<&'a Rc<TreeNode>> {
        let mut pos = *self;
        let mut current_child = parent_node.get_child_ref().to_owned();

        loop {
            match current_child {
                None        => return None,
                Some(child) => {
                    if pos == 0 { return current_child; }

                    pos = pos-1;
                    current_child = child.get_sibling_ref().to_owned();
                }
            }
        }
    }
}

impl<'b> TreeNodeIndex for &'b str {
    ///
    /// Finds the tree node corresponding to the specified index in the tree
    ///
    /// When searching by tag, we match only the first item that we find.
    ///
    fn lookup_index<'a>(&self, parent_node: &'a TreeNode) -> Option<&'a Rc<TreeNode>> {
        let mut current_child = parent_node.get_child_ref().to_owned();

        loop {
            match current_child {
                None        => return None,
                Some(child) => {
                    if (*child).get_tag() == *self {
                        return current_child;
                    }

                    current_child = child.get_sibling_ref().to_owned();
                }
            }
        }
    }
}

///
/// Provides the ability to reference the children of a tree node by looking up a particular index
///
pub trait TreeNodeLookup : TreeNode {
    ///
    /// Looks up a child node at a particular index (panics if the child does not exist)
    ///
    fn get_child_at<'a, TIndex: TreeNodeIndex>(&'a self, index: TIndex) -> &'a TreeNode;

    ///
    /// Looks up a child node at a particular index
    ///
    fn get_child_ref_at<'a, TIndex: TreeNodeIndex>(&'a self, index: TIndex) -> Option<&'a Rc<TreeNode>>;
}

impl<T: TreeNode> TreeNodeLookup for T {
    ///
    /// Looks up a child node at a particular index (panics if the child does not exist)
    ///
    fn get_child_at<'a, TIndex: TreeNodeIndex>(&'a self, index: TIndex) -> &'a TreeNode {
        let opt_node = index.lookup_index(self);
        let node_ref = opt_node.unwrap();

        &**node_ref
    }

    ///
    /// Looks up a child node at a particular index
    ///
    fn get_child_ref_at<'a, TIndex: TreeNodeIndex>(&'a self, index: TIndex) -> Option<&'a Rc<TreeNode>> {
        index.lookup_index(self)
    }
}

#[cfg(test)]
mod treenode_index_tests {
    use super::super::treenode::*;
    use super::super::basictree::*;
    use std::rc::*;

    #[test]
    fn lookup_usize() {
        let mut tree = BasicTree::new("test", ());
        let first_child = Rc::new(BasicTree::new("first_child", ()));

        tree.set_child_ref(first_child);

        let lookup = 0.lookup_index(&tree);
        assert!(lookup.is_some());
        assert!(lookup.unwrap().get_tag() == "first_child");
        assert!(tree.get_sibling_ref().is_none());
    }

    #[test]
    fn can_get_first_child() {
        let mut tree = BasicTree::new("test", ());
        let first_child = Rc::new(BasicTree::new("first_child", ()));

        tree.set_child_ref(first_child);

        assert!((tree.get_child_at(0).get_tag()) == "first_child");
        assert!(tree.get_sibling_ref().is_none());
    }

    #[test]
    fn can_get_first_child_by_string() {
        let mut tree = BasicTree::new("test", ());
        let first_child = Rc::new(BasicTree::new("first_child", ()));

        tree.set_child_ref(first_child);

        assert!((tree.get_child_at("first_child").get_tag()) == "first_child");
        assert!(tree.get_sibling_ref().is_none());
    }
}
<commit_msg>Annoyingly T: TreeNode does not match TreeNode itself so we have to duplicate code to apply the exact same thing to the trait as well as objects that implement the trait.<commit_after>use std::rc::*;
use super::treenode::*;

///
/// Trait implemented by types that can work as a tree node index
///
pub trait TreeNodeIndex {
    ///
    /// Finds the tree node corresponding to the specified index in the tree
    ///
    fn lookup_index<'a>(&self, parent_node: &'a TreeNode) -> Option<&'a Rc<TreeNode>>;
}

impl TreeNodeIndex for usize {
    ///
    /// Finds the tree node corresponding to the specified index in the tree
    ///
    fn lookup_index<'a>(&self, parent_node: &'a TreeNode) -> Option<&'a Rc<TreeNode>> {
        let mut pos = *self;
        let mut current_child = parent_node.get_child_ref().to_owned();

        loop {
            match current_child {
                None        => return None,
                Some(child) => {
                    if pos == 0 { return current_child; }

                    pos = pos-1;
                    current_child = child.get_sibling_ref().to_owned();
                }
            }
        }
    }
}

impl<'b> TreeNodeIndex for &'b str {
    ///
    /// Finds the tree node corresponding to the specified index in the tree
    ///
    /// When searching by tag, we match only the first item that we find.
    ///
    fn lookup_index<'a>(&self, parent_node: &'a TreeNode) -> Option<&'a Rc<TreeNode>> {
        let mut current_child = parent_node.get_child_ref().to_owned();

        loop {
            match current_child {
                None        => return None,
                Some(child) => {
                    if (*child).get_tag() == *self {
                        return current_child;
                    }

                    current_child = child.get_sibling_ref().to_owned();
                }
            }
        }
    }
}

///
/// Provides the ability to reference the children of a tree node by looking up a particular index
///
pub trait TreeNodeLookup {
    ///
    /// Looks up a child node at a particular index (panics if the child does not exist)
    ///
    fn get_child_at<'a, TIndex: TreeNodeIndex>(&'a self, index: TIndex) -> &'a TreeNode;

    ///
    /// Looks up a child node at a particular index
    ///
    fn get_child_ref_at<'a, TIndex: TreeNodeIndex>(&'a self, index: TIndex) -> Option<&'a Rc<TreeNode>>;
}

impl<T: TreeNode> TreeNodeLookup for T {
    ///
    /// Looks up a child node at a particular index (panics if the child does not exist)
    ///
    fn get_child_at<'a, TIndex: TreeNodeIndex>(&'a self, index: TIndex) -> &'a TreeNode {
        let opt_node = index.lookup_index(self);
        let node_ref = opt_node.unwrap();

        &**node_ref
    }

    ///
    /// Looks up a child node at a particular index
    ///
    fn get_child_ref_at<'a, TIndex: TreeNodeIndex>(&'a self, index: TIndex) -> Option<&'a Rc<TreeNode>> {
        index.lookup_index(self)
    }
}

// Annoying that the above definition doesn't cover the 'raw treenode' case
impl TreeNodeLookup for TreeNode {
    ///
    /// Looks up a child node at a particular index (panics if the child does not exist)
    ///
    fn get_child_at<'a, TIndex: TreeNodeIndex>(&'a self, index: TIndex) -> &'a TreeNode {
        let opt_node = index.lookup_index(self);
        let node_ref = opt_node.unwrap();

        &**node_ref
    }

    ///
    /// Looks up a child node at a particular index
    ///
    fn get_child_ref_at<'a, TIndex: TreeNodeIndex>(&'a self, index: TIndex) -> Option<&'a Rc<TreeNode>> {
        index.lookup_index(self)
    }
}

#[cfg(test)]
mod treenode_index_tests {
    use super::super::treenode::*;
    use super::super::basictree::*;
    use std::rc::*;

    #[test]
    fn lookup_usize() {
        let mut tree = BasicTree::new("test", ());
        let first_child = Rc::new(BasicTree::new("first_child", ()));

        tree.set_child_ref(first_child);

        let lookup = 0.lookup_index(&tree);
        assert!(lookup.is_some());
        assert!(lookup.unwrap().get_tag() == "first_child");
        assert!(tree.get_sibling_ref().is_none());
    }

    #[test]
    fn can_get_first_child() {
        let mut tree = BasicTree::new("test", ());
        let first_child = Rc::new(BasicTree::new("first_child", ()));

        tree.set_child_ref(first_child);

        assert!((tree.get_child_at(0).get_tag()) == "first_child");
        assert!(tree.get_sibling_ref().is_none());
    }

    #[test]
    fn can_get_first_child_by_string() {
        let mut tree = BasicTree::new("test", ());
        let first_child = Rc::new(BasicTree::new("first_child", ()));

        tree.set_child_ref(first_child);

        assert!((tree.get_child_at("first_child").get_tag()) == "first_child");
        assert!(tree.get_sibling_ref().is_none());
    }
}
<|endoftext|>


<commit_before>use bindings::*;
use std::slice::bytes::copy_memory;

static KEY_BYTES: uint = 32;
static NONCE_BYTES: uint = 24;
static ZERO_BYTES: uint = 32;

pub struct SecretKey ([u8, ..KEY_BYTES]);

impl SecretKey {
    pub fn from_str(str: &str) -> SecretKey {
        SecretKey::from_slice(str.as_bytes())
    }

    pub fn from_slice(slice: &[u8]) -> SecretKey {
        assert!(slice.len() <= KEY_BYTES);

        let mut sized = [0u8, ..KEY_BYTES];
        copy_memory(sized, slice);

        SecretKey(sized)
    }
}

pub struct SecretBox {
    sk: SecretKey
}

impl SecretBox {
    pub fn new(sk: SecretKey) -> SecretBox {
        SecretBox { sk: sk }
    }

    pub fn encrypt(&self, msg: &[u8]) -> (Vec<u8>, Vec<u8>) {
        unsafe {
            let mut stretched: Vec<u8> = Vec::from_elem(ZERO_BYTES, 0u8);
            stretched.push_all(msg);

            let mut nonce = Vec::from_elem(NONCE_BYTES, 0u8);
            randombytes(nonce.as_mut_ptr(), NONCE_BYTES as u64);

            let SecretKey(sk) = self.sk;

            let mut cipher = Vec::from_elem(stretched.len(), 0u8);

            match crypto_secretbox(cipher.as_mut_ptr(),
                                   stretched.as_ptr(),
                                   stretched.len() as u64,
                                   nonce.as_ptr(),
                                   sk.as_ptr()) {
                0 => (cipher, nonce),
                _ => fail!("crypto_secretbox failed")
            }
        }
    }

    pub fn decrypt(&self, cipher: &[u8], nonce: &[u8]) -> Vec<u8> {
        unsafe {
            let mut msg = Vec::from_elem(cipher.len(), 0u8);
            let SecretKey(sk) = self.sk;

            match crypto_secretbox_open(msg.as_mut_ptr(),
                                        cipher.as_ptr(),
                                        cipher.len() as u64,
                                        nonce.as_ptr(),
                                        sk.as_ptr()) {
                0 => Vec::from_slice(msg.slice(ZERO_BYTES, msg.len())),
                _ => fail!("crypto_secretbox_open failed")
            }
        }
    }
}

#[test]
fn test_secretbox_sanity() {
    for i in range(0 as uint, 16) {
        let msg = Vec::from_elem(i * 4, i as u8);

        let sb = SecretBox::new(SecretKey::from_str("some passkey"));
        let (encr, nonce) = sb.encrypt(msg.as_slice());

        println!("enc:\t{}\nnonce:\t{}", encr, nonce);

        let decr = sb.decrypt(encr.as_slice(), nonce.as_slice());

        println!("dec:\t{}", decr);

        assert!(decr == msg);
    }
}
<commit_msg>Encapsulate cipher text and nonce in SecretMsg struct<commit_after>use bindings::*;
use std::slice::bytes::copy_memory;

static KEY_BYTES: uint = 32;
static NONCE_BYTES: uint = 24;
static ZERO_BYTES: uint = 32;

pub struct SecretKey ([u8, ..KEY_BYTES]);

impl SecretKey {
    pub fn from_str(str: &str) -> SecretKey {
        SecretKey::from_slice(str.as_bytes())
    }

    pub fn from_slice(slice: &[u8]) -> SecretKey {
        assert!(slice.len() <= KEY_BYTES);

        let mut sized = [0u8, ..KEY_BYTES];
        copy_memory(sized, slice);

        SecretKey(sized)
    }
}

pub struct SecretMsg {
    pub nonce: [u8, ..NONCE_BYTES],
    pub cipher: Vec<u8>
}

pub struct SecretBox {
    sk: SecretKey
}

impl SecretBox {
    pub fn new(sk: SecretKey) -> SecretBox {
        SecretBox { sk: sk }
    }

    pub fn encrypt(&self, msg: &[u8]) -> SecretMsg {
        unsafe {
            let mut stretched: Vec<u8> = Vec::from_elem(ZERO_BYTES, 0u8);
            stretched.push_all(msg);

            let mut nonce = [0u8, ..NONCE_BYTES];
            randombytes(nonce.as_mut_ptr(), NONCE_BYTES as u64);

            let SecretKey(sk) = self.sk;

            let mut cipher = Vec::from_elem(stretched.len(), 0u8);

            match crypto_secretbox(cipher.as_mut_ptr(),
                                   stretched.as_ptr(),
                                   stretched.len() as u64,
                                   nonce.as_ptr(),
                                   sk.as_ptr()) {
                0 => SecretMsg { nonce: nonce, cipher: cipher },
                _ => fail!("crypto_secretbox failed")
            }
        }
    }

    pub fn decrypt(&self, msg: &SecretMsg) -> Vec<u8> {
        let &SecretMsg { ref nonce, ref cipher } = msg;

        unsafe {
            let mut msg = Vec::from_elem(cipher.len(), 0u8);
            let SecretKey(sk) = self.sk;

            match crypto_secretbox_open(msg.as_mut_ptr(),
                                        cipher.as_ptr(),
                                        cipher.len() as u64,
                                        nonce.as_ptr(),
                                        sk.as_ptr()) {
                0 => Vec::from_slice(msg.slice(ZERO_BYTES, msg.len())),
                _ => fail!("crypto_secretbox_open failed")
            }
        }
    }
}

#[test]
fn test_secretbox_sanity() {
    for i in range(0 as uint, 16) {
        let msg = Vec::from_elem(i * 4, i as u8);

        let sb = SecretBox::new(SecretKey::from_str("some secret key"));
        let SecretMsg { nonce, cipher } = sb.encrypt(msg.as_slice());

        println!("enc:\t{}\nnonce:\t{}", cipher, Vec::from_slice(nonce));

        let decr = sb.decrypt(&SecretMsg { nonce: nonce, cipher: cipher });

        println!("dec:\t{}", decr);

        assert!(decr == msg);
    }
}

#[test]
fn test_secretbox_uniqueness() {
    let msg = Vec::from_elem(128, 0x53u8);

    let box1 = SecretBox::new(SecretKey::from_str("1"));
    let box2 = SecretBox::new(SecretKey::from_str(""));

    let SecretMsg { nonce: n1, cipher: c1 } = box1.encrypt(msg.as_slice());
    let SecretMsg { nonce: n2, cipher: c2 } = box2.encrypt(msg.as_slice());

    assert!(n1 != n2);
    assert!(c1 != c2);
    assert!(c1 != msg);
    assert!(c2 != msg);
}
<|endoftext|>