<commit_before>use node_types::{escape_name, node_type_name, Entry, Field, Storage, TypeName};
use std::collections::BTreeMap as Map;
use std::collections::BTreeSet as Set;
use std::fmt;
use std::path::Path;
use tracing::{error, info, span, Level};
use tree_sitter::{Language, Node, Parser, Tree};

pub struct Extractor {
    pub parser: Parser,
    pub schema: Vec<Entry>,
}

pub fn create(language: Language, schema: Vec<Entry>) -> Extractor {
    let mut parser = Parser::new();
    parser.set_language(language).unwrap();

    Extractor { parser, schema }
}
impl Extractor {
    pub fn extract<'a>(&'a mut self, path: &Path) -> std::io::Result<Program> {
        let span = span!(
            Level::TRACE,
            "extract",
            file = %path.display()
        );

        let _enter = span.enter();

        info!("extracting: {}", path.display());

        let source = std::fs::read(&path)?;
        let tree = &self
            .parser
            .parse(&source, None)
            .expect("Failed to parse file");
        let mut visitor = Visitor {
            source: &source,
            trap_output: vec![TrapEntry::Comment(format!(
                "Auto-generated TRAP file for {}",
                path.display()
            ))],
            counter: -1,
            // TODO: should we handle path strings that are not valid UTF8 better?
            path: format!("{}", path.display()),
            stack: Vec::new(),
            tables: build_schema_lookup(&self.schema),
            union_types: build_union_type_lookup(&self.schema),
        };
        traverse(&tree, &mut visitor);

        &self.parser.reset();
        Ok(Program(visitor.trap_output))
    }
}

fn build_schema_lookup<'a>(schema: &'a Vec<Entry>) -> Map<&'a TypeName, &'a Entry> {
    let mut map = std::collections::BTreeMap::new();
    for entry in schema {
        if let Entry::Table { type_name, .. } = entry {
            map.insert(type_name, entry);
        }
    }
    map
}

fn build_union_type_lookup<'a>(schema: &'a Vec<Entry>) -> Map<&'a TypeName, &'a Set<TypeName>> {
    let mut union_types = std::collections::BTreeMap::new();
    for entry in schema {
        if let Entry::Union { type_name, members } = entry {
            union_types.insert(type_name, members);
        }
    }
    union_types
}

struct Visitor<'a> {
    /// The file path of the source code (as string)
    path: String,
    /// The source code as a UTF-8 byte array
    source: &'a Vec<u8>,
    /// The accumulated trap entries
    trap_output: Vec<TrapEntry>,
    /// A counter for generating fresh labels
    counter: i32,
    /// A lookup table from type name to dbscheme table entries
    tables: Map<&'a TypeName, &'a Entry>,
    /// A lookup table for union types mapping a type name to its direct members
    union_types: Map<&'a TypeName, &'a Set<TypeName>>,
    /// A stack for gathering information from hild nodes. Whenever a node is entered
    /// an empty list is pushed. All children append their data (field name, label, type) to
    /// the the list. When the visitor leaves a node the list containing the child data is popped
    /// from the stack and matched against the dbscheme for the node. If the expectations are met
    /// the corresponding row definitions are added to the trap_output.
    stack: Vec<Vec<(Option<&'static str>, Label, TypeName)>>,
}

impl Visitor<'_> {
    fn enter_node(&mut self, node: Node) -> bool {
        if node.is_error() {
            error!("{}:{}: parse error", &self.path, node.start_position().row);
            return false;
        }
        if node.is_missing() {
            error!(
                "{}:{}: parse expecting '{}'",
                &self.path,
                node.start_position().row,
                node.kind()
            );
            return false;
        }

        if node.is_extra() {
            return false;
        }

        self.stack.push(Vec::new());
        return true;
    }

    fn leave_node(&mut self, field_name: Option<&'static str>, node: Node) {
        if node.is_extra() || node.is_error() || node.is_missing() {
            return;
        }
        let child_nodes = self.stack.pop().expect("Vistor: empty stack");
        let table = self.tables.get(&TypeName {
            kind: node.kind().to_owned(),
            named: node.is_named(),
        });
        if let Some(Entry::Table { fields, .. }) = table {
            self.counter += 1;
            let id = Label::Normal(self.counter);
            let loc = Label::Location(self.counter);
            self.trap_output.push(TrapEntry::New(id));
            self.trap_output.push(TrapEntry::New(loc));
            self.trap_output.push(location_for(&self.path, loc, node));
            let table_name = node_type_name(node.kind(), node.is_named());
            let args: Option<Vec<Arg>>;
            if fields.is_empty() {
                args = Some(vec![sliced_source_arg(self.source, node)]);
            } else {
                args = self.complex_node(&node, fields, child_nodes, id);
            }
            if let Some(args) = args {
                self.trap_output
                    .push(TrapEntry::Definition(table_name, id, args, loc));
            }
            if let Some(parent) = self.stack.last_mut() {
                parent.push((
                    field_name,
                    id,
                    TypeName {
                        kind: node.kind().to_owned(),
                        named: node.is_named(),
                    },
                ))
            };
        } else {
            error!(
                "{}:{}: unknown table type: '{}'",
                &self.path,
                node.start_position().row,
                node.kind()
            );
        }
    }
    fn complex_node(
        &mut self,
        node: &Node,
        fields: &Vec<Field>,
        child_nodes: Vec<(Option<&str>, Label, TypeName)>,
        parent_id: Label,
    ) -> Option<Vec<Arg>> {
        let mut map: Map<&Option<String>, (&Field, Vec<Label>)> = std::collections::BTreeMap::new();
        for field in fields {
            map.insert(&field.name, (field, Vec::new()));
        }
        for (child_field, child_id, child_type) in child_nodes {
            if let Some((field, values)) = map.get_mut(&child_field.map(|x| x.to_owned())) {
                //TODO: handle error and missing nodes
                if self.type_matches(&child_type, &field.types) {
                    values.push(child_id);
                } else if field.name.is_some() {
                    error!(
                        "{}:{}: type mismatch for field {}::{} with type {:?} != {:?}",
                        &self.path,
                        node.start_position().row,
                        node.kind(),
                        child_field.unwrap_or("child"),
                        child_type,
                        field.types
                    )
                }
            } else {
                if child_field.is_some() || child_type.named {
                    error!(
                        "{}:{}: value for unknown field: {}::{} and type {:?}",
                        &self.path,
                        node.start_position().row,
                        node.kind(),
                        &child_field.unwrap_or("child"),
                        &child_type
                    );
                }
            }
        }
        let mut args = Vec::new();
        let mut is_valid = true;
        for field in fields {
            let child_ids = &map.get(&field.name).unwrap().1;
            match &field.storage {
                Storage::Column => {
                    if child_ids.len() == 1 {
                        args.push(Arg::Label(*child_ids.first().unwrap()));
                    } else {
                        is_valid = false;
                        error!(
                            "{}:{}: {} for field: {}::{}",
                            &self.path,
                            node.start_position().row,
                            if child_ids.is_empty() {
                                "missing value"
                            } else {
                                "too many values"
                            },
                            node.kind(),
                            &field.get_name()
                        )
                    }
                }
                Storage::Table => {
                    for (index, child_id) in child_ids.iter().enumerate() {
                        self.trap_output.push(TrapEntry::ChildOf(
                            node_type_name(&field.parent.kind, field.parent.named),
                            parent_id,
                            field.get_name(),
                            Index(index),
                            *child_id,
                        ));
                    }
                }
            }
        }
        if is_valid {
            Some(args)
        } else {
            None
        }
    }
    fn type_matches(&self, tp: &TypeName, types: &Set<TypeName>) -> bool {
        if types.contains(tp) {
            return true;
        }
        for other in types.iter() {
            if let Some(x) = self.union_types.get(other) {
                if self.type_matches(tp, x) {
                    return true;
                }
            }
        }
        return false;
    }
}

// Emit a slice of a source file as an Arg.
fn sliced_source_arg(source: &Vec<u8>, n: Node) -> Arg {
    let range = n.byte_range();
    Arg::String(String::from(
        std::str::from_utf8(&source[range.start..range.end]).expect("Failed to decode string"),
    ))
}

// Emit a 'Located' TrapEntry for the provided node, appropriately calibrated.
fn location_for<'a>(fp: &String, label: Label, n: Node) -> TrapEntry {
    let start_line = n.start_position().row;
    let start_col = n.start_position().column;
    let end_line = n.end_position().row;
    let end_col = n.end_position().column;
    TrapEntry::Located(vec![
        Arg::Label(label),
        Arg::String(fp.to_owned()),
        Arg::Int(start_line),
        Arg::Int(start_col),
        Arg::Int(end_line),
        Arg::Int(end_col),
    ])
}

fn traverse(tree: &Tree, visitor: &mut Visitor) {
    let cursor = &mut tree.walk();
    visitor.enter_node(cursor.node());
    let mut recurse = true;
    loop {
        if recurse && cursor.goto_first_child() {
            recurse = visitor.enter_node(cursor.node());
        } else {
            visitor.leave_node(cursor.field_name(), cursor.node());

            if cursor.goto_next_sibling() {
                recurse = visitor.enter_node(cursor.node());
            } else if cursor.goto_parent() {
                recurse = false;
            } else {
                break;
            }
        }
    }
}

pub struct Program(Vec<TrapEntry>);

impl fmt::Display for Program {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut text = String::new();
        for trap_entry in &self.0 {
            text.push_str(&format!("{}\n", trap_entry));
        }
        write!(f, "{}", text)
    }
}

enum TrapEntry {
    // @id = *@
    New(Label),
    // @node_def(self, arg?, location)@
    Definition(String, Label, Vec<Arg>, Label),
    // @node_child(self, index, parent)@
    ChildOf(String, Label, String, Index, Label),
    // @location(loc, path, r1, c1, r2, c2)
    Located(Vec<Arg>),
    Comment(String),
}
impl fmt::Display for TrapEntry {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            TrapEntry::New(id) => write!(f, "{} = *", id),
            TrapEntry::Definition(n, id, args, loc) => {
                let mut args_str = String::new();
                for arg in args {
                    args_str.push_str(&format!("{}, ", arg));
                }
                write!(
                    f,
                    "{}({}, {}{})",
                    escape_name(&format!("{}_def", &n)),
                    id,
                    args_str,
                    loc
                )
            }
            TrapEntry::ChildOf(pname, id, fname, idx, p) => write!(
                f,
                "{}({}, {}, {})",
                escape_name(&format!("{}_{}", &pname, &fname)),
                id,
                idx,
                p
            ),
            TrapEntry::Located(args) => write!(
                f,
                "location({}, {}, {}, {}, {}, {})",
                args.get(0).unwrap(),
                args.get(1).unwrap(),
                args.get(2).unwrap(),
                args.get(3).unwrap(),
                args.get(4).unwrap(),
                args.get(5).unwrap(),
            ),
            TrapEntry::Comment(line) => write!(f, "// {}", line),
        }
    }
}

#[derive(Debug, Copy, Clone)]
enum Label {
    // Identifiers of the form #0, #1...
    Normal(i32), // #0, #1, etc.
    // Location identifiers of the form #0_loc, #1_loc...
    Location(i32), // #0_loc, #1_loc, etc.
}

impl fmt::Display for Label {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Label::Normal(x) => write!(f, "#{}", x),
            Label::Location(x) => write!(f, "#{}_loc", x),
        }
    }
}

// Numeric indices.
#[derive(Debug, Copy, Clone)]
struct Index(usize);

impl fmt::Display for Index {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

// Some untyped argument to a TrapEntry.
#[derive(Debug)]
enum Arg {
    Label(Label),
    Int(usize),
    String(String),
}

impl fmt::Display for Arg {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Arg::Label(x) => write!(f, "{}", x),
            Arg::Int(x) => write!(f, "{}", x),
            Arg::String(x) => write!(f, "\"{}\"", x.replace("\"", "\"\"")),
        }
    }
}
<commit_msg>Adjust source locations<commit_after>use node_types::{escape_name, node_type_name, Entry, Field, Storage, TypeName};
use std::collections::BTreeMap as Map;
use std::collections::BTreeSet as Set;
use std::fmt;
use std::path::Path;
use tracing::{error, info, span, Level};
use tree_sitter::{Language, Node, Parser, Tree};

pub struct Extractor {
    pub parser: Parser,
    pub schema: Vec<Entry>,
}

pub fn create(language: Language, schema: Vec<Entry>) -> Extractor {
    let mut parser = Parser::new();
    parser.set_language(language).unwrap();

    Extractor { parser, schema }
}
impl Extractor {
    pub fn extract<'a>(&'a mut self, path: &Path) -> std::io::Result<Program> {
        let span = span!(
            Level::TRACE,
            "extract",
            file = %path.display()
        );

        let _enter = span.enter();

        info!("extracting: {}", path.display());

        let source = std::fs::read(&path)?;
        let tree = &self
            .parser
            .parse(&source, None)
            .expect("Failed to parse file");
        let mut visitor = Visitor {
            source: &source,
            trap_output: vec![TrapEntry::Comment(format!(
                "Auto-generated TRAP file for {}",
                path.display()
            ))],
            counter: -1,
            // TODO: should we handle path strings that are not valid UTF8 better?
            path: format!("{}", path.display()),
            stack: Vec::new(),
            tables: build_schema_lookup(&self.schema),
            union_types: build_union_type_lookup(&self.schema),
        };
        traverse(&tree, &mut visitor);

        &self.parser.reset();
        Ok(Program(visitor.trap_output))
    }
}

fn build_schema_lookup<'a>(schema: &'a Vec<Entry>) -> Map<&'a TypeName, &'a Entry> {
    let mut map = std::collections::BTreeMap::new();
    for entry in schema {
        if let Entry::Table { type_name, .. } = entry {
            map.insert(type_name, entry);
        }
    }
    map
}

fn build_union_type_lookup<'a>(schema: &'a Vec<Entry>) -> Map<&'a TypeName, &'a Set<TypeName>> {
    let mut union_types = std::collections::BTreeMap::new();
    for entry in schema {
        if let Entry::Union { type_name, members } = entry {
            union_types.insert(type_name, members);
        }
    }
    union_types
}

struct Visitor<'a> {
    /// The file path of the source code (as string)
    path: String,
    /// The source code as a UTF-8 byte array
    source: &'a Vec<u8>,
    /// The accumulated trap entries
    trap_output: Vec<TrapEntry>,
    /// A counter for generating fresh labels
    counter: i32,
    /// A lookup table from type name to dbscheme table entries
    tables: Map<&'a TypeName, &'a Entry>,
    /// A lookup table for union types mapping a type name to its direct members
    union_types: Map<&'a TypeName, &'a Set<TypeName>>,
    /// A stack for gathering information from hild nodes. Whenever a node is entered
    /// an empty list is pushed. All children append their data (field name, label, type) to
    /// the the list. When the visitor leaves a node the list containing the child data is popped
    /// from the stack and matched against the dbscheme for the node. If the expectations are met
    /// the corresponding row definitions are added to the trap_output.
    stack: Vec<Vec<(Option<&'static str>, Label, TypeName)>>,
}

impl Visitor<'_> {
    fn enter_node(&mut self, node: Node) -> bool {
        if node.is_error() {
            error!("{}:{}: parse error", &self.path, node.start_position().row);
            return false;
        }
        if node.is_missing() {
            error!(
                "{}:{}: parse error: expecting '{}'",
                &self.path,
                node.start_position().row,
                node.kind()
            );
            return false;
        }

        if node.is_extra() {
            return false;
        }

        self.stack.push(Vec::new());
        return true;
    }

    fn leave_node(&mut self, field_name: Option<&'static str>, node: Node) {
        if node.is_extra() || node.is_error() || node.is_missing() {
            return;
        }
        let child_nodes = self.stack.pop().expect("Vistor: empty stack");
        let table = self.tables.get(&TypeName {
            kind: node.kind().to_owned(),
            named: node.is_named(),
        });
        if let Some(Entry::Table { fields, .. }) = table {
            self.counter += 1;
            let id = Label::Normal(self.counter);
            let loc = Label::Location(self.counter);
            self.trap_output.push(TrapEntry::New(id));
            self.trap_output.push(TrapEntry::New(loc));
            self.trap_output
                .push(location_for(&self.source, &self.path, loc, node));
            let table_name = node_type_name(node.kind(), node.is_named());
            let args: Option<Vec<Arg>>;
            if fields.is_empty() {
                args = Some(vec![sliced_source_arg(self.source, node)]);
            } else {
                args = self.complex_node(&node, fields, child_nodes, id);
            }
            if let Some(args) = args {
                self.trap_output
                    .push(TrapEntry::Definition(table_name, id, args, loc));
            }
            if let Some(parent) = self.stack.last_mut() {
                parent.push((
                    field_name,
                    id,
                    TypeName {
                        kind: node.kind().to_owned(),
                        named: node.is_named(),
                    },
                ))
            };
        } else {
            error!(
                "{}:{}: unknown table type: '{}'",
                &self.path,
                node.start_position().row,
                node.kind()
            );
        }
    }
    fn complex_node(
        &mut self,
        node: &Node,
        fields: &Vec<Field>,
        child_nodes: Vec<(Option<&str>, Label, TypeName)>,
        parent_id: Label,
    ) -> Option<Vec<Arg>> {
        let mut map: Map<&Option<String>, (&Field, Vec<Label>)> = std::collections::BTreeMap::new();
        for field in fields {
            map.insert(&field.name, (field, Vec::new()));
        }
        for (child_field, child_id, child_type) in child_nodes {
            if let Some((field, values)) = map.get_mut(&child_field.map(|x| x.to_owned())) {
                //TODO: handle error and missing nodes
                if self.type_matches(&child_type, &field.types) {
                    values.push(child_id);
                } else if field.name.is_some() {
                    error!(
                        "{}:{}: type mismatch for field {}::{} with type {:?} != {:?}",
                        &self.path,
                        node.start_position().row,
                        node.kind(),
                        child_field.unwrap_or("child"),
                        child_type,
                        field.types
                    )
                }
            } else {
                if child_field.is_some() || child_type.named {
                    error!(
                        "{}:{}: value for unknown field: {}::{} and type {:?}",
                        &self.path,
                        node.start_position().row,
                        node.kind(),
                        &child_field.unwrap_or("child"),
                        &child_type
                    );
                }
            }
        }
        let mut args = Vec::new();
        let mut is_valid = true;
        for field in fields {
            let child_ids = &map.get(&field.name).unwrap().1;
            match &field.storage {
                Storage::Column => {
                    if child_ids.len() == 1 {
                        args.push(Arg::Label(*child_ids.first().unwrap()));
                    } else {
                        is_valid = false;
                        error!(
                            "{}:{}: {} for field: {}::{}",
                            &self.path,
                            node.start_position().row,
                            if child_ids.is_empty() {
                                "missing value"
                            } else {
                                "too many values"
                            },
                            node.kind(),
                            &field.get_name()
                        )
                    }
                }
                Storage::Table => {
                    for (index, child_id) in child_ids.iter().enumerate() {
                        self.trap_output.push(TrapEntry::ChildOf(
                            node_type_name(&field.parent.kind, field.parent.named),
                            parent_id,
                            field.get_name(),
                            Index(index),
                            *child_id,
                        ));
                    }
                }
            }
        }
        if is_valid {
            Some(args)
        } else {
            None
        }
    }
    fn type_matches(&self, tp: &TypeName, types: &Set<TypeName>) -> bool {
        if types.contains(tp) {
            return true;
        }
        for other in types.iter() {
            if let Some(x) = self.union_types.get(other) {
                if self.type_matches(tp, x) {
                    return true;
                }
            }
        }
        return false;
    }
}

// Emit a slice of a source file as an Arg.
fn sliced_source_arg(source: &Vec<u8>, n: Node) -> Arg {
    let range = n.byte_range();
    Arg::String(String::from(
        std::str::from_utf8(&source[range.start..range.end]).expect("Failed to decode string"),
    ))
}

// Emit a 'Located' TrapEntry for the provided node, appropriately calibrated.
fn location_for<'a>(source: &Vec<u8>, fp: &String, label: Label, n: Node) -> TrapEntry {
    let start_line = n.start_position().row + 1;
    let start_col = n.start_position().column + 1;
    let mut end_line = n.end_position().row + 1;
    let mut end_col = n.end_position().column;
    if end_col == 0 {
        if start_line >= end_line {
            // the range is empty, clip it to sensible values
            end_line = start_line;
            end_col = start_col - 1;
        } else {
            let mut index = n.end_byte();
            // end_col = 0 means that we are the start of a line
            // unfortunately 0 is invalid as column number, therefore
            // we should update the end location to be the end of the
            // previous line
            if index > 0 && index <= source.len() {
                index -= 1;
                if source[index] != b'\n' {
                    error!("expecting a line break symbol, but none found while correcting end column value");
                }
                end_line -= 1;
                end_col = 1;
                while index > 0 && source[index - 1] != b'\n' {
                    index -= 1;
                    end_col += 1;
                }
            } else {
                error!(
                    "cannot correct end column value: end_byte index {} is not in range [0,{}]",
                    index,
                    source.len()
                );
            }
        }
    }
    TrapEntry::Located(vec![
        Arg::Label(label),
        Arg::String(fp.to_owned()),
        Arg::Int(start_line),
        Arg::Int(start_col),
        Arg::Int(end_line),
        Arg::Int(end_col),
    ])
}

fn traverse(tree: &Tree, visitor: &mut Visitor) {
    let cursor = &mut tree.walk();
    visitor.enter_node(cursor.node());
    let mut recurse = true;
    loop {
        if recurse && cursor.goto_first_child() {
            recurse = visitor.enter_node(cursor.node());
        } else {
            visitor.leave_node(cursor.field_name(), cursor.node());

            if cursor.goto_next_sibling() {
                recurse = visitor.enter_node(cursor.node());
            } else if cursor.goto_parent() {
                recurse = false;
            } else {
                break;
            }
        }
    }
}

pub struct Program(Vec<TrapEntry>);

impl fmt::Display for Program {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut text = String::new();
        for trap_entry in &self.0 {
            text.push_str(&format!("{}\n", trap_entry));
        }
        write!(f, "{}", text)
    }
}

enum TrapEntry {
    // @id = *@
    New(Label),
    // @node_def(self, arg?, location)@
    Definition(String, Label, Vec<Arg>, Label),
    // @node_child(self, index, parent)@
    ChildOf(String, Label, String, Index, Label),
    // @location(loc, path, r1, c1, r2, c2)
    Located(Vec<Arg>),
    Comment(String),
}
impl fmt::Display for TrapEntry {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            TrapEntry::New(id) => write!(f, "{} = *", id),
            TrapEntry::Definition(n, id, args, loc) => {
                let mut args_str = String::new();
                for arg in args {
                    args_str.push_str(&format!("{}, ", arg));
                }
                write!(
                    f,
                    "{}({}, {}{})",
                    escape_name(&format!("{}_def", &n)),
                    id,
                    args_str,
                    loc
                )
            }
            TrapEntry::ChildOf(pname, id, fname, idx, p) => write!(
                f,
                "{}({}, {}, {})",
                escape_name(&format!("{}_{}", &pname, &fname)),
                id,
                idx,
                p
            ),
            TrapEntry::Located(args) => write!(
                f,
                "location({}, {}, {}, {}, {}, {})",
                args.get(0).unwrap(),
                args.get(1).unwrap(),
                args.get(2).unwrap(),
                args.get(3).unwrap(),
                args.get(4).unwrap(),
                args.get(5).unwrap(),
            ),
            TrapEntry::Comment(line) => write!(f, "// {}", line),
        }
    }
}

#[derive(Debug, Copy, Clone)]
enum Label {
    // Identifiers of the form #0, #1...
    Normal(i32), // #0, #1, etc.
    // Location identifiers of the form #0_loc, #1_loc...
    Location(i32), // #0_loc, #1_loc, etc.
}

impl fmt::Display for Label {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Label::Normal(x) => write!(f, "#{}", x),
            Label::Location(x) => write!(f, "#{}_loc", x),
        }
    }
}

// Numeric indices.
#[derive(Debug, Copy, Clone)]
struct Index(usize);

impl fmt::Display for Index {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

// Some untyped argument to a TrapEntry.
#[derive(Debug)]
enum Arg {
    Label(Label),
    Int(usize),
    String(String),
}

impl fmt::Display for Arg {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Arg::Label(x) => write!(f, "{}", x),
            Arg::Int(x) => write!(f, "{}", x),
            Arg::String(x) => write!(f, "\"{}\"", x.replace("\"", "\"\"")),
        }
    }
}
<|endoftext|>


<commit_before>use std::mem::swap;
use std::ops::Deref;

use unicode_width::UnicodeWidthStr;

use strcursor::StrCursor;

pub struct Buffer {
    front_buf: String,
    back_buf: String,
    pos: usize
}

impl Buffer {

    pub fn new() -> Buffer {
        Buffer {
            front_buf: String::new(),
            back_buf: String::new(),
            pos: 0
        }
    }

    pub fn swap(&mut self) {
        swap(&mut self.front_buf, &mut self.back_buf);
        self.pos = self.front_buf.len();
    }

    pub fn reset(&mut self) {
        self.front_buf.clear();
        self.pos = 0;
    }

    pub fn replace(&mut self, s: &str) {
        self.front_buf.clear();
        self.front_buf.extend(s.chars());
        self.pos = s.len();
    }

    pub fn insert_char_at_cursor(&mut self, c: char) {
        self.front_buf.insert(self.pos, c);
        self.pos += c.len_utf8();
    }

    pub fn insert_chars_at_cursor(&mut self, s: String) {
        for c in s.chars() {
            self.insert_char_at_cursor(c);
        }
    }

    pub fn delete_char_left_of_cursor(&mut self) {
        if self.move_left() {
            self.front_buf.remove(self.pos);
        }
    }

    pub fn delete_char_right_of_cursor(&mut self) -> bool {
        if self.pos < self.front_buf.len() {
            self.front_buf.remove(self.pos);
            return true;
        } else {
            return false;
        }
    }

    fn cursor(&self) -> StrCursor {
        StrCursor::new_at_left_of_byte_pos(self.front_buf.deref(), self.pos)
    }

    fn prev_pos(&self) -> Option<usize> {
        self.cursor().at_prev().map(|c| c.byte_pos())
    }

    pub fn move_left(&mut self) -> bool {
        match self.prev_pos() {
            Some(pos) => {
                self.pos = pos;
                true
            },
            None => false
        }
    }

    fn next_pos(&self) -> Option<usize> {
        self.cursor().at_next().map(|c| c.byte_pos())
    }

    pub fn move_right(&mut self) -> bool {
        match self.next_pos() {
            Some(pos) => {
                self.pos = pos;
                true
            },
            None => false
        }
    }

    pub fn move_start(&mut self) {
        self.pos = 0;
    }

    pub fn move_end(&mut self) {
        self.pos = self.front_buf.len();
    }

    fn char_pos(&self) -> usize {
         UnicodeWidthStr::width(self.cursor().slice_before())
    }

    pub fn get_line(&self, prompt: &str) -> Vec<u8> {
        let mut seq = Vec::new();
        seq.extend("\r".as_bytes());
        seq.extend(prompt.as_bytes());
        seq.extend(self.front_buf.as_bytes());
        seq.extend("\x1b[0K".as_bytes());
        seq.extend(&format!("\r\x1b[{}C", prompt.len() + self.char_pos()).into_bytes());
        seq
    }

    pub fn to_string(self) -> String {
        self.front_buf
    }

}

#[test]
fn move_and_insert_ascii() {
    let mut buf = Buffer::new();
    buf.insert_char_at_cursor('a');
    buf.move_left();
    buf.insert_char_at_cursor('x');
    buf.move_left();
    buf.move_right();
    buf.move_right();
    buf.insert_char_at_cursor('b');
    buf.move_start();
    buf.insert_char_at_cursor('w');
    buf.move_end();
    buf.insert_char_at_cursor('c');
    assert_eq!(buf.to_string(), "wxabc".to_string());
}

#[test]
fn move_and_insert_cyrillic() {
    let mut buf = Buffer::new();
    buf.insert_char_at_cursor('Й');
    buf.move_left();
    buf.insert_char_at_cursor('ч');
    buf.move_left();
    buf.move_right();
    buf.move_right();
    buf.insert_char_at_cursor('Њ');
    buf.move_start();
    buf.insert_char_at_cursor('Ѿ');
    buf.move_end();
    buf.insert_char_at_cursor('Җ');
    assert_eq!(buf.to_string(), "ѾчЙЊҖ".to_string());
}
<commit_msg>add cjk buffer test<commit_after>use std::mem::swap;
use std::ops::Deref;

use unicode_width::UnicodeWidthStr;

use strcursor::StrCursor;

pub struct Buffer {
    front_buf: String,
    back_buf: String,
    pos: usize
}

impl Buffer {

    pub fn new() -> Buffer {
        Buffer {
            front_buf: String::new(),
            back_buf: String::new(),
            pos: 0
        }
    }

    pub fn swap(&mut self) {
        swap(&mut self.front_buf, &mut self.back_buf);
        self.pos = self.front_buf.len();
    }

    pub fn reset(&mut self) {
        self.front_buf.clear();
        self.pos = 0;
    }

    pub fn replace(&mut self, s: &str) {
        self.front_buf.clear();
        self.front_buf.extend(s.chars());
        self.pos = s.len();
    }

    pub fn insert_char_at_cursor(&mut self, c: char) {
        self.front_buf.insert(self.pos, c);
        self.pos += c.len_utf8();
    }

    pub fn insert_chars_at_cursor(&mut self, s: String) {
        for c in s.chars() {
            self.insert_char_at_cursor(c);
        }
    }

    pub fn delete_char_left_of_cursor(&mut self) {
        if self.move_left() {
            self.front_buf.remove(self.pos);
        }
    }

    pub fn delete_char_right_of_cursor(&mut self) -> bool {
        if self.pos < self.front_buf.len() {
            self.front_buf.remove(self.pos);
            return true;
        } else {
            return false;
        }
    }

    fn cursor(&self) -> StrCursor {
        StrCursor::new_at_left_of_byte_pos(self.front_buf.deref(), self.pos)
    }

    fn prev_pos(&self) -> Option<usize> {
        self.cursor().at_prev().map(|c| c.byte_pos())
    }

    pub fn move_left(&mut self) -> bool {
        match self.prev_pos() {
            Some(pos) => {
                self.pos = pos;
                true
            },
            None => false
        }
    }

    fn next_pos(&self) -> Option<usize> {
        self.cursor().at_next().map(|c| c.byte_pos())
    }

    pub fn move_right(&mut self) -> bool {
        match self.next_pos() {
            Some(pos) => {
                self.pos = pos;
                true
            },
            None => false
        }
    }

    pub fn move_start(&mut self) {
        self.pos = 0;
    }

    pub fn move_end(&mut self) {
        self.pos = self.front_buf.len();
    }

    fn char_pos(&self) -> usize {
         UnicodeWidthStr::width(self.cursor().slice_before())
    }

    pub fn get_line(&self, prompt: &str) -> Vec<u8> {
        let mut seq = Vec::new();
        seq.extend("\r".as_bytes());
        seq.extend(prompt.as_bytes());
        seq.extend(self.front_buf.as_bytes());
        seq.extend("\x1b[0K".as_bytes());
        seq.extend(&format!("\r\x1b[{}C", prompt.len() + self.char_pos()).into_bytes());
        seq
    }

    pub fn to_string(self) -> String {
        self.front_buf
    }

}

#[test]
fn move_and_insert_ascii() {
    let mut buf = Buffer::new();
    buf.insert_char_at_cursor('a');
    buf.move_left();
    buf.insert_char_at_cursor('x');
    buf.move_left();
    buf.move_right();
    buf.move_right();
    buf.insert_char_at_cursor('b');
    buf.move_start();
    buf.insert_char_at_cursor('w');
    buf.move_end();
    buf.insert_char_at_cursor('c');
    assert_eq!(buf.to_string(), "wxabc".to_string());
}

#[test]
fn move_and_insert_cyrillic() {
    let mut buf = Buffer::new();
    buf.insert_char_at_cursor('Й');
    buf.move_left();
    buf.insert_char_at_cursor('ч');
    buf.move_left();
    buf.move_right();
    buf.move_right();
    buf.insert_char_at_cursor('Њ');
    buf.move_start();
    buf.insert_char_at_cursor('Ѿ');
    buf.move_end();
    buf.insert_char_at_cursor('Җ');
    assert_eq!(buf.to_string(), "ѾчЙЊҖ".to_string());
}

#[test]
fn move_and_insert_cjk() {
    let mut buf = Buffer::new();
    buf.insert_char_at_cursor('䩖');
    buf.move_left();
    buf.insert_char_at_cursor('䨻');
    buf.move_left();
    buf.move_right();
    buf.move_right();
    buf.insert_char_at_cursor('䦴');
    buf.move_start();
    buf.insert_char_at_cursor('乫');
    buf.move_end();
    buf.insert_char_at_cursor('憛');
    assert_eq!(buf.to_string(), "乫䨻䩖䦴憛".to_string());
}
<|endoftext|>


<commit_before>use std::fmt;
use std::net::{
    IpAddr,
    Ipv4Addr,
    SocketAddr,
};
use std::str;
use std::sync::mpsc::sync_channel;
use std::time::{Duration, Instant};
use std::thread;

use kudu_pb::master::{
    DeleteTableRequestPB,
    IsCreateTableDoneRequestPB,
    TableIdentifierPB,
};

use backoff::Backoff;
use Error;
use master::MasterProxy;
use Result;
use rpc::Messenger;
use table::TableBuilder;
use TableId;

#[derive(Clone)]
pub struct Client {
    messenger: Messenger,
    master: MasterProxy,
    config: ClientConfig,
}

impl Client {

    pub fn new(config: ClientConfig) -> Client {
        let messenger = Messenger::new().unwrap();
        let master = MasterProxy::new(config.master_addresses(), messenger.clone());
        Client {
            master: master,
            messenger: messenger,
            config: config,
        }
    }

    /// Creates a new Kudu table with the schema and options specified by `builder`. Returns the
    /// new table's ID, or an error on failure.
    pub fn create_table(&self, builder: TableBuilder, deadline: Instant) -> Result<TableId> {
        let (send, recv) = sync_channel(0);
        self.master.create_table(deadline,
                                 try!(builder.into_pb()),
                                 move |resp| send.send(resp).unwrap());
        recv.recv().unwrap().map(|resp| {
            str::from_utf8(resp.get_table_id())
                .map_err(|error| Error::Serialization(format!("{}", error)))
                .and_then(TableId::parse)
        }).unwrap()
    }

    /// Returns `true` if the table is fully created.
    pub fn is_create_table_done<S>(&self, table: S, deadline: Instant) -> Result<bool>
    where S: Into<String> {
        let mut identifier = TableIdentifierPB::new();
        identifier.set_table_name(table.into());
        self.do_is_create_table_done(identifier, deadline)
    }

    /// Returns `true` if the table is fully created.
    pub fn is_create_table_done_by_id(&self, id: &TableId, deadline: Instant) -> Result<bool> {
        let mut identifier = TableIdentifierPB::new();
        identifier.set_table_id(id.to_string().into_bytes());
        self.do_is_create_table_done(identifier, deadline)
    }

    fn do_is_create_table_done(&self, table: TableIdentifierPB, deadline: Instant) -> Result<bool> {
        let mut request = IsCreateTableDoneRequestPB::new();
        request.set_table(table);

        let (send, recv) = sync_channel(0);
        self.master.is_create_table_done(deadline, request, move |resp| send.send(resp).unwrap());
        recv.recv().unwrap().map(|resp| resp.get_done())
    }

    /// Synchronously waits until the table is created. If an error is returned,
    /// the table may not be created yet.
    pub fn wait_for_table_creation<S>(&self, table: S, deadline: Instant) -> Result<()>
    where S: Into<String> {
        let mut identifier = TableIdentifierPB::new();
        identifier.set_table_name(table.into());
        self.do_wait_for_table_creation(identifier, deadline)
    }

    /// Synchronously waits until the table is created. If an error is returned,
    /// the table may not be created yet.
    pub fn wait_for_table_creation_by_id(&self, id: &TableId, deadline: Instant) -> Result<()> {
        let mut identifier = TableIdentifierPB::new();
        identifier.set_table_id(id.to_string().into_bytes());
        self.do_wait_for_table_creation(identifier, deadline)
    }

    fn do_wait_for_table_creation(&self, table: TableIdentifierPB, deadline: Instant) -> Result<()> {
        let mut backoff = Backoff::with_duration_range(16, 4096);

        let mut is_create_table_done = self.do_is_create_table_done(table.clone(), deadline);
        while let Ok(false) = is_create_table_done {
            let sleep_ms = backoff.next_backoff_ms();
            debug!("create table not yet complete, waiting {}ms", sleep_ms);
            // TODO: do delayed send instead of sleep
            thread::sleep(Duration::from_millis(sleep_ms));
            is_create_table_done = self.do_is_create_table_done(table.clone(), deadline);
        }
        is_create_table_done.map(|_| ())
    }

    /// Deletes the Kudu table.
    pub fn delete_table<S>(&self, table: S, deadline: Instant) -> Result<()>
    where S: Into<String> {
        let mut identifier = TableIdentifierPB::new();
        identifier.set_table_name(table.into());
        self.do_delete_table(identifier, deadline)
    }

    /// Deletes the Kudu table.
    pub fn delete_table_by_id(&self, id: &TableId, deadline: Instant) -> Result<()> {
        let mut identifier = TableIdentifierPB::new();
        identifier.set_table_id(id.to_string().into_bytes());
        self.do_delete_table(identifier, deadline)
    }

    pub fn do_delete_table(&self, table: TableIdentifierPB, deadline: Instant) -> Result<()> {
        let mut request = DeleteTableRequestPB::new();
        request.set_table(table);

        let (send, recv) = sync_channel(0);
        self.master.delete_table(deadline, request, move |resp| send.send(resp).unwrap());
        recv.recv().unwrap().map(|_| ())
    }
}

impl fmt::Debug for Client {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Client")
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct ClientConfig {

    /// A seed set of master addresses. Must contain at least one active master in the cluster.
    master_addresses: Vec<SocketAddr>,

    /// The default timeout for admin operations.
    default_admin_operation_timeout: Duration,
}

impl ClientConfig {
    pub fn new(master_addresses: Vec<SocketAddr>) -> ClientConfig {
        ClientConfig {
            master_addresses: master_addresses,
            ..Default::default()
        }
    }

    pub fn master_addresses(&self) -> &[SocketAddr] {
        &self.master_addresses
    }

    pub fn set_master_addresses(&mut self, master_addresses: Vec<SocketAddr>) -> &mut ClientConfig {
        self.master_addresses = master_addresses;
        self
    }

    pub fn default_admin_operation_timeout(&self) -> Duration {
        self.default_admin_operation_timeout
    }

    pub fn set_default_admin_operation_timeout(&mut self, timeout: Duration) -> &mut ClientConfig {
        self.default_admin_operation_timeout = timeout;
        self
    }
}

impl Default for ClientConfig {
    fn default() -> ClientConfig {
        ClientConfig {
            master_addresses: vec![SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 7051)],
            default_admin_operation_timeout: Duration::from_secs(30),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::{Duration, Instant};

    use mini_cluster::MiniCluster;
    use SchemaBuilder;
    use super::*;
    use TableBuilder;
    use DataType;

    use env_logger;

    fn deadline() -> Instant {
        Instant::now() + Duration::from_secs(5)
    }

    #[test]
    fn create_and_delete_table() {
        let _ = env_logger::init();
        let cluster = MiniCluster::default();

        let client = Client::new(ClientConfig::new(cluster.master_addrs().to_owned()));

        let mut schema_builder = SchemaBuilder::new();
        schema_builder.add_column("key", DataType::Int32).set_not_null();
        schema_builder.add_column("val", DataType::Int32);
        schema_builder.set_primary_key(vec!["key".to_owned()]);

        let mut table_builder = TableBuilder::new("create_and_delete_table",
                                                  schema_builder.build().unwrap());
        table_builder.add_hash_partitions(vec!["key".to_owned()], 4);
        table_builder.set_num_replicas(1);

        // The tablet server is real slow coming up.
        // TODO: add MiniCluster::wait_for_startup() or equivalent.
        ::std::thread::sleep_ms(2000);

        let table_id = client.create_table(table_builder, deadline()).unwrap();
        client.wait_for_table_creation_by_id(&table_id, deadline()).unwrap();

        client.delete_table_by_id(&table_id, deadline()).unwrap();
    }
}
<commit_msg>do not panic on table create failure<commit_after>use std::fmt;
use std::net::{
    IpAddr,
    Ipv4Addr,
    SocketAddr,
};
use std::str;
use std::sync::mpsc::sync_channel;
use std::time::{Duration, Instant};
use std::thread;

use kudu_pb::master::{
    DeleteTableRequestPB,
    IsCreateTableDoneRequestPB,
    TableIdentifierPB,
};

use backoff::Backoff;
use Error;
use master::MasterProxy;
use Result;
use rpc::Messenger;
use table::TableBuilder;
use TableId;

#[derive(Clone)]
pub struct Client {
    messenger: Messenger,
    master: MasterProxy,
    config: ClientConfig,
}

impl Client {

    pub fn new(config: ClientConfig) -> Client {
        let messenger = Messenger::new().unwrap();
        let master = MasterProxy::new(config.master_addresses(), messenger.clone());
        Client {
            master: master,
            messenger: messenger,
            config: config,
        }
    }

    /// Creates a new Kudu table with the schema and options specified by `builder`. Returns the
    /// new table's ID, or an error on failure.
    pub fn create_table(&self, builder: TableBuilder, deadline: Instant) -> Result<TableId> {
        let (send, recv) = sync_channel(0);
        self.master.create_table(deadline,
                                 try!(builder.into_pb()),
                                 move |resp| send.send(resp).unwrap());
        recv.recv().unwrap().and_then(|resp| {
            str::from_utf8(resp.get_table_id())
                .map_err(|error| Error::Serialization(format!("{}", error)))
                .and_then(TableId::parse)
        })
    }

    /// Returns `true` if the table is fully created.
    pub fn is_create_table_done<S>(&self, table: S, deadline: Instant) -> Result<bool>
    where S: Into<String> {
        let mut identifier = TableIdentifierPB::new();
        identifier.set_table_name(table.into());
        self.do_is_create_table_done(identifier, deadline)
    }

    /// Returns `true` if the table is fully created.
    pub fn is_create_table_done_by_id(&self, id: &TableId, deadline: Instant) -> Result<bool> {
        let mut identifier = TableIdentifierPB::new();
        identifier.set_table_id(id.to_string().into_bytes());
        self.do_is_create_table_done(identifier, deadline)
    }

    fn do_is_create_table_done(&self, table: TableIdentifierPB, deadline: Instant) -> Result<bool> {
        let mut request = IsCreateTableDoneRequestPB::new();
        request.set_table(table);

        let (send, recv) = sync_channel(0);
        self.master.is_create_table_done(deadline, request, move |resp| send.send(resp).unwrap());
        recv.recv().unwrap().map(|resp| resp.get_done())
    }

    /// Synchronously waits until the table is created. If an error is returned,
    /// the table may not be created yet.
    pub fn wait_for_table_creation<S>(&self, table: S, deadline: Instant) -> Result<()>
    where S: Into<String> {
        let mut identifier = TableIdentifierPB::new();
        identifier.set_table_name(table.into());
        self.do_wait_for_table_creation(identifier, deadline)
    }

    /// Synchronously waits until the table is created. If an error is returned,
    /// the table may not be created yet.
    pub fn wait_for_table_creation_by_id(&self, id: &TableId, deadline: Instant) -> Result<()> {
        let mut identifier = TableIdentifierPB::new();
        identifier.set_table_id(id.to_string().into_bytes());
        self.do_wait_for_table_creation(identifier, deadline)
    }

    fn do_wait_for_table_creation(&self, table: TableIdentifierPB, deadline: Instant) -> Result<()> {
        let mut backoff = Backoff::with_duration_range(16, 4096);

        let mut is_create_table_done = self.do_is_create_table_done(table.clone(), deadline);
        while let Ok(false) = is_create_table_done {
            let sleep_ms = backoff.next_backoff_ms();
            debug!("create table not yet complete, waiting {}ms", sleep_ms);
            // TODO: do delayed send instead of sleep
            thread::sleep(Duration::from_millis(sleep_ms));
            is_create_table_done = self.do_is_create_table_done(table.clone(), deadline);
        }
        is_create_table_done.map(|_| ())
    }

    /// Deletes the Kudu table.
    pub fn delete_table<S>(&self, table: S, deadline: Instant) -> Result<()>
    where S: Into<String> {
        let mut identifier = TableIdentifierPB::new();
        identifier.set_table_name(table.into());
        self.do_delete_table(identifier, deadline)
    }

    /// Deletes the Kudu table.
    pub fn delete_table_by_id(&self, id: &TableId, deadline: Instant) -> Result<()> {
        let mut identifier = TableIdentifierPB::new();
        identifier.set_table_id(id.to_string().into_bytes());
        self.do_delete_table(identifier, deadline)
    }

    pub fn do_delete_table(&self, table: TableIdentifierPB, deadline: Instant) -> Result<()> {
        let mut request = DeleteTableRequestPB::new();
        request.set_table(table);

        let (send, recv) = sync_channel(0);
        self.master.delete_table(deadline, request, move |resp| send.send(resp).unwrap());
        recv.recv().unwrap().map(|_| ())
    }
}

impl fmt::Debug for Client {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Client")
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct ClientConfig {

    /// A seed set of master addresses. Must contain at least one active master in the cluster.
    master_addresses: Vec<SocketAddr>,

    /// The default timeout for admin operations.
    default_admin_operation_timeout: Duration,
}

impl ClientConfig {
    pub fn new(master_addresses: Vec<SocketAddr>) -> ClientConfig {
        ClientConfig {
            master_addresses: master_addresses,
            ..Default::default()
        }
    }

    pub fn master_addresses(&self) -> &[SocketAddr] {
        &self.master_addresses
    }

    pub fn set_master_addresses(&mut self, master_addresses: Vec<SocketAddr>) -> &mut ClientConfig {
        self.master_addresses = master_addresses;
        self
    }

    pub fn default_admin_operation_timeout(&self) -> Duration {
        self.default_admin_operation_timeout
    }

    pub fn set_default_admin_operation_timeout(&mut self, timeout: Duration) -> &mut ClientConfig {
        self.default_admin_operation_timeout = timeout;
        self
    }
}

impl Default for ClientConfig {
    fn default() -> ClientConfig {
        ClientConfig {
            master_addresses: vec![SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 7051)],
            default_admin_operation_timeout: Duration::from_secs(30),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::{Duration, Instant};

    use mini_cluster::MiniCluster;
    use SchemaBuilder;
    use super::*;
    use TableBuilder;
    use DataType;

    use env_logger;

    fn deadline() -> Instant {
        Instant::now() + Duration::from_secs(5)
    }

    #[test]
    fn create_and_delete_table() {
        let _ = env_logger::init();
        let cluster = MiniCluster::default();

        let client = Client::new(ClientConfig::new(cluster.master_addrs().to_owned()));

        let mut schema_builder = SchemaBuilder::new();
        schema_builder.add_column("key", DataType::Int32).set_not_null();
        schema_builder.add_column("val", DataType::Int32);
        schema_builder.set_primary_key(vec!["key".to_owned()]);

        let mut table_builder = TableBuilder::new("create_and_delete_table",
                                                  schema_builder.build().unwrap());
        table_builder.add_hash_partitions(vec!["key".to_owned()], 4);
        table_builder.set_num_replicas(1);

        // The tablet server is real slow coming up.
        // TODO: add MiniCluster::wait_for_startup() or equivalent.
        ::std::thread::sleep_ms(2000);

        let table_id = client.create_table(table_builder, deadline()).unwrap();
        client.wait_for_table_creation_by_id(&table_id, deadline()).unwrap();

        client.delete_table_by_id(&table_id, deadline()).unwrap();
    }
}
<|endoftext|>


<commit_before>use std::io::net::udp::UdpSocket;
use std::io::net::ip::{Ipv4Addr, SocketAddr};
use std::io::{IoResult, IoError, IoErrorKind, OtherIoError};
use std::io::Timer;
use std::comm::{TryRecvError, Disconnected, Empty};
use std::time::duration::Duration;
use std::comm::Select;
use packet::{Packet, PacketConnect, PacketAccept, PacketReject, Command, Disconnect};

fn reader_process(mut reader: UdpSocket, reader_sub_out: Sender<Packet>, reader_sub_in: Receiver<Command>, target_addr: SocketAddr, protocol_id: u32) {
    let mut buf = [0, ..255];
    reader.set_timeout(Some(1000));
    loop {
        match reader.recv_from(buf) {
            Ok((amt, src)) => {
                if src == target_addr {
                    match Packet::deserialize(buf.slice_to(amt)) {
                        Ok(packet) => {
                            if packet.protocol_id == protocol_id {
                                reader_sub_out.send(packet);
                            }
                        },
                        Err(_) => ()
                    }
                }
            }
            Err(e) => {
                match e.kind {
                    TimedOut => {
                        match reader_sub_in.try_recv() {
                            Ok(Disconnect) => {
                                break;
                            },
                            Err(Disconnected) => {
                                break;
                            },
                            Err(Empty) => {

                            }
                        }
                    }
                }
            }
        }
    }
}

fn writer_process(mut writer: UdpSocket, writer_sub_out: Sender<Command>, writer_sub_in: Receiver<Packet>, target_addr: SocketAddr) {
    for msg in writer_sub_in.iter() {
        writer.send_to(msg.serialize().as_slice(), target_addr);
    }
}

pub enum ConnectionState {
    CommsDisconnected,
    CommsConnecting,
    CommsConnected
}

/**
 * Clientside implementation of UDP networking
 */
pub struct Client {
    pub addr: SocketAddr,
    pub target_addr: SocketAddr,

    protocol_id: u32,
    connection_state: ConnectionState,

    reader_send: Sender<Command>,
    reader_receive: Receiver<Packet>,
    writer_send: Sender<Packet>,
    writer_receive: Receiver<Command>
}

impl Client {
    /**
     * Connect our Client to a target Server
     */
    pub fn connect(addr: SocketAddr, target_addr: SocketAddr, protocol_id: u32) -> IoResult<Client> {
         match UdpSocket::bind(addr) {
            Ok(reader) => {
                let writer = reader.clone();

                let (reader_out, reader_sub_in) = channel();
                let (reader_sub_out, reader_in) = channel();

                spawn(proc() {
                    reader_process(reader, reader_sub_out, reader_sub_in, target_addr, protocol_id);
                });

                let (writer_out, writer_sub_in) = channel();
                let (writer_sub_out, writer_in) = channel();

                spawn(proc() {
                    writer_process(writer, writer_sub_out, writer_sub_in, target_addr);
                });

                let mut client = Client {
                    addr: addr,
                    target_addr: target_addr,
                    reader_send: reader_out,
                    reader_receive: reader_in,
                    writer_send: writer_out,
                    writer_receive: writer_in,
                    protocol_id: protocol_id,
                    connection_state: CommsDisconnected
                };

                client.connection_dance();
                match client.connection_state {
                    CommsConnected => {
                        Ok(client)
                    },
                    _ => {
                        Err(IoError {
                            kind: OtherIoError,
                            desc: "Failed to connect",
                            detail: None
                        })
                    }
                }
            }
            Err(e) => Err(e)
        }
    }

    /**
     * A blocking connection request
     */
    fn connection_dance(&mut self) {
        self.connection_state = CommsConnecting;
        let mut timer = Timer::new().unwrap();
        let mut attempts = 0u;

        while attempts < 3 && match self.connection_state { CommsConnecting => true, _ => false } {

            self.writer_send.send(Packet {
                protocol_id: self.protocol_id,
                packet_type: PacketConnect,
                packet_content: vec![]
            });

            let timeout = timer.oneshot(Duration::seconds(5));
            //FIXME: Replace with the select! macro when it starts working
            let sel = Select::new();
            let mut reader = sel.handle(&self.reader_receive);
            let mut timeout = sel.handle(&timeout);
            unsafe { reader.add(); timeout.add(); }
            let ret = sel.wait();
            if ret == reader.id() {
                let packet = self.reader_receive.recv();
                {
                    match packet.packet_type {
                        PacketAccept => {
                            self.connection_state = CommsConnected;
                        }
                        PacketReject => {
                            self.connection_state = CommsDisconnected;
                        }
                        _ => (),
                    }
                }
            } else if ret == timeout.id() {
                let () = timeout.recv();
                { attempts += 1; }
            } else {
                fail!("What");
            }
        }

        //We didn't manage to connect
        match self.connection_state {
            CommsConnecting => {
                self.connection_state = CommsDisconnected
            },
            _ => ()
        }
    }

    /**
     * Pop the last event off of our comms queue, if any
     */
    pub fn poll(&mut self) -> Option<Packet> {
        match self.reader_receive.try_recv() {
            Ok(value) => Some(value),
            _ => None
        }
    }
}

impl Drop for Client {

    fn drop(&mut self) {
        self.reader_send.send(Disconnect);
    }
}
<commit_msg>Tidyup<commit_after>use std::io::net::udp::UdpSocket;
use std::io::net::ip::SocketAddr;
use std::io::{IoResult, IoError, OtherIoError, TimedOut};
use std::io::Timer;
use std::comm::{Disconnected, Empty, Select};
use std::time::duration::Duration;
use packet::{Packet, PacketConnect, PacketAccept, PacketReject, Command, Disconnect};

fn reader_process(mut reader: UdpSocket, reader_sub_out: Sender<Packet>, reader_sub_in: Receiver<Command>, target_addr: SocketAddr, protocol_id: u32) {
    let mut buf = [0, ..255];
    reader.set_timeout(Some(1000));
    loop {
        match reader.recv_from(buf) {
            Ok((amt, src)) => {
                if src == target_addr {
                    match Packet::deserialize(buf.slice_to(amt)) {
                        Ok(packet) => {
                            if packet.protocol_id == protocol_id {
                                reader_sub_out.send(packet);
                            }
                        },
                        Err(_) => ()
                    }
                }
            }
            Err(e) => {
                match e.kind {
                    TimedOut => {
                        match reader_sub_in.try_recv() {
                            Ok(Disconnect) => {
                                break;
                            },
                            Err(Disconnected) => {
                                break;
                            },
                            Err(Empty) => {
                                //Keep going
                            }
                        }
                    },
                    _ => ()
                }
            }
        }
    }
}

fn writer_process(mut writer: UdpSocket, writer_sub_out: Sender<Command>, writer_sub_in: Receiver<Packet>, target_addr: SocketAddr) {
    for msg in writer_sub_in.iter() {
        match writer.send_to(msg.serialize().as_slice(), target_addr) {
            Ok(()) => (),
            Err(e) => println!("Error sending data - {}", e)
        }
    }
}

pub enum ConnectionState {
    CommsDisconnected,
    CommsConnecting,
    CommsConnected
}

/**
 * Clientside implementation of UDP networking
 */
pub struct Client {
    pub addr: SocketAddr,
    pub target_addr: SocketAddr,

    protocol_id: u32,
    connection_state: ConnectionState,

    reader_send: Sender<Command>,
    reader_receive: Receiver<Packet>,
    writer_send: Sender<Packet>,
    writer_receive: Receiver<Command>
}

impl Client {
    /**
     * Connect our Client to a target Server
     */
    pub fn connect(addr: SocketAddr, target_addr: SocketAddr, protocol_id: u32) -> IoResult<Client> {
         match UdpSocket::bind(addr) {
            Ok(reader) => {
                let writer = reader.clone();

                let (reader_out, reader_sub_in) = channel();
                let (reader_sub_out, reader_in) = channel();

                spawn(proc() {
                    reader_process(reader, reader_sub_out, reader_sub_in, target_addr, protocol_id);
                });

                let (writer_out, writer_sub_in) = channel();
                let (writer_sub_out, writer_in) = channel();

                spawn(proc() {
                    writer_process(writer, writer_sub_out, writer_sub_in, target_addr);
                });

                let mut client = Client {
                    addr: addr,
                    target_addr: target_addr,
                    reader_send: reader_out,
                    reader_receive: reader_in,
                    writer_send: writer_out,
                    writer_receive: writer_in,
                    protocol_id: protocol_id,
                    connection_state: CommsDisconnected
                };

                if client.connection_dance() {
                    Ok(client)
                } else {
                    Err(IoError {
                        kind: OtherIoError,
                        desc: "Failed to connect",
                        detail: None
                    })
                }
            }
            Err(e) => Err(e)
        }
    }

    /**
     * A blocking connection request
     */
    fn connection_dance(&mut self) -> bool {
        self.connection_state = CommsConnecting;
        let mut timer = Timer::new().unwrap();
        let mut attempts = 0u;

        while attempts < 3 && match self.connection_state { CommsConnecting => true, _ => false } {

            self.writer_send.send(Packet {
                protocol_id: self.protocol_id,
                packet_type: PacketConnect,
                packet_content: vec![]
            });

            let timeout = timer.oneshot(Duration::seconds(5));

            //FIXME: Replace with the select! macro when it starts working
            let sel = Select::new();
            let mut reader = sel.handle(&self.reader_receive);
            let mut timeout = sel.handle(&timeout);
            unsafe { reader.add(); timeout.add(); }
            let ret = sel.wait();
            if ret == reader.id() {
                let packet = self.reader_receive.recv();
                {
                    match packet.packet_type {
                        PacketAccept => {
                            self.connection_state = CommsConnected;
                        }
                        PacketReject => {
                            self.connection_state = CommsDisconnected;
                        }
                        _ => (),
                    }
                }
            } else if ret == timeout.id() {
                let () = timeout.recv();
                { attempts += 1; }
            } else {
                fail!("What");
            }
        }

        //We didn't manage to connect
        match self.connection_state {
            CommsConnecting => {
                self.connection_state = CommsDisconnected;
                false
            },
            CommsDisconnected => false,
            CommsConnected => {
                true
            }
        }
    }

    /**
     * Pop the last event off of our comms queue, if any
     */
    pub fn poll(&mut self) -> Option<Packet> {
        match self.reader_receive.try_recv() {
            Ok(value) => Some(value),
            _ => None
        }
    }
}

impl Drop for Client {

    fn drop(&mut self) {
        self.reader_send.send(Disconnect);
    }
}
<|endoftext|>


<commit_before>use std::path::{Path, PathBuf};

use xdg;
use rustc_version;

fn xdg_dir() -> xdg::BaseDirectories {
    let prefix = Path::new("afl.rs").join(rustc_version()).join(
        pkg_version(),
    );
    xdg::BaseDirectories::with_prefix(prefix).unwrap()
}

pub fn rustc_version() -> String {
    let version_meta = rustc_version::version_meta().unwrap();
    let mut ret = String::from("rustc-");
    ret.push_str(&version_meta.semver.to_string());
    if let Some(commit_hash) = version_meta.commit_hash {
        ret.push('-');
        ret.push_str(&commit_hash[..7]);
    }
    ret
}

fn pkg_version() -> String {
    let mut ret = String::from("afl.rs-");

    let version = env!("CARGO_PKG_VERSION");
    assert!(!version.is_empty());

    ret.push_str(version);
    ret
}

pub fn afl_dir() -> PathBuf {
    xdg_dir().create_data_directory("afl").unwrap()
}

pub fn afl_llvm_rt_dir() -> PathBuf {
    xdg_dir().create_data_directory("afl-llvm-rt").unwrap()
}

#[allow(dead_code)]
pub fn object_file_path() -> PathBuf {
    afl_llvm_rt_dir().join("libafl-llvm-rt.o")
}

pub fn archive_file_path() -> PathBuf {
    afl_llvm_rt_dir().join("libafl-llvm-rt.a")
}
<commit_msg>Utilize constant.<commit_after>use std::path::{Path, PathBuf};

use xdg;
use rustc_version;

fn xdg_dir() -> xdg::BaseDirectories {
    let prefix = Path::new("afl.rs").join(rustc_version()).join(
        pkg_version(),
    );
    xdg::BaseDirectories::with_prefix(prefix).unwrap()
}

const SHORT_COMMIT_HASH_LEN: usize = 7;

pub fn rustc_version() -> String {
    let version_meta = rustc_version::version_meta().unwrap();
    let mut ret = String::from("rustc-");
    ret.push_str(&version_meta.semver.to_string());
    if let Some(commit_hash) = version_meta.commit_hash {
        ret.push('-');
        ret.push_str(&commit_hash[..SHORT_COMMIT_HASH_LEN]);
    }
    ret
}

fn pkg_version() -> String {
    let mut ret = String::from("afl.rs-");

    let version = env!("CARGO_PKG_VERSION");
    assert!(!version.is_empty());

    ret.push_str(version);
    ret
}

pub fn afl_dir() -> PathBuf {
    xdg_dir().create_data_directory("afl").unwrap()
}

pub fn afl_llvm_rt_dir() -> PathBuf {
    xdg_dir().create_data_directory("afl-llvm-rt").unwrap()
}

#[allow(dead_code)]
pub fn object_file_path() -> PathBuf {
    afl_llvm_rt_dir().join("libafl-llvm-rt.o")
}

pub fn archive_file_path() -> PathBuf {
    afl_llvm_rt_dir().join("libafl-llvm-rt.a")
}
<|endoftext|>


<commit_before>//! The configuration module represents the various options that an NSQ consumer
//! can contain.
#![deny(missing_docs)]

/// NsqConfig contains values supported by `nsqd` for consumers.
#[derive(Debug)]
pub struct NsqConfig {
    max_in_flight: usize,
}

impl NsqConfig {
    /// Return a default NsqConfig with sensible defaults.
    pub fn new() -> NsqConfig {
        NsqConfig {
            max_in_flight: 1,
        }
    }

    /// Set the maximum number of messages a consumer can handle before `nsqd`
    /// expects a response.
    pub fn max_in_fligh<'a>(&'a mut self, mif: usize) -> &'a mut NsqConfig {
        self.max_in_flight = mif;
        self
    }
}
<commit_msg>move to the Default trait instead of new<commit_after>//! The configuration module represents the various options that an NSQ consumer
//! can contain.
#![deny(missing_docs)]

/// NsqConfig contains values supported by `nsqd` for consumers.
#[derive(Debug)]
pub struct NsqConfig {
    max_in_flight: usize,
}

impl Default for NsqConfig {
    fn default() -> NsqConfig {
        NsqConfig {
            max_in_flight: 1,
        }
    }
}

impl NsqConfig {
    /// Set the maximum number of messages a consumer can handle before `nsqd`
    /// expects a response.
    pub fn max_in_fligh<'a>(&'a mut self, mif: usize) -> &'a mut NsqConfig {
        self.max_in_flight = mif;
        self
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

extern crate toml;

use lists::{SeparatorTactic, ListTactic};

macro_rules! configuration_option_enum{
    ($e:ident: $( $x:ident ),+ $(,)*) => {
        #[derive(Copy, Clone, Eq, PartialEq, Debug)]
        pub enum $e {
            $( $x ),+
        }

        impl_enum_decodable!($e, $( $x ),+);
    }
}

configuration_option_enum! { NewlineStyle:
    Windows, // \r\n
    Unix, // \n
}

configuration_option_enum! { BraceStyle:
    AlwaysNextLine,
    PreferSameLine,
    // Prefer same line except where there is a where clause, in which case force
    // the brace to the next line.
    SameLineWhere,
}

// How to indent a function's return type.
configuration_option_enum! { ReturnIndent:
    // Aligned with the arguments
    WithArgs,
    // Aligned with the where clause
    WithWhereClause,
}

// How to stle a struct literal.
configuration_option_enum! { StructLitStyle:
    // First line on the same line as the opening brace, all lines aligned with
    // the first line.
    Visual,
    // First line is on a new line and all lines align with block indent.
    Block,
    // FIXME Maybe we should also have an option to align types.
}

configuration_option_enum! { BlockIndentStyle:
    // Same level as parent.
    Inherit,
    // One level deeper than parent.
    Tabbed,
    // Aligned with block open.
    Visual,
}

configuration_option_enum! { Density:
    // Fit as much on one line as possible.
    Compressed,
    // Use more lines.
    Tall,
    // Try to compress if the body is empty.
    CompressedIfEmpty,
}

impl Density {
    pub fn to_list_tactic(self) -> ListTactic {
        match self {
            Density::Compressed => ListTactic::Mixed,
            Density::Tall | Density::CompressedIfEmpty => ListTactic::HorizontalVertical,
        }
    }
}

configuration_option_enum! { LicensePolicy:
    // Do not place license text at top of files
    NoLicense,
    // Use the text in "license" field as the license
    TextLicense,
    // Use a text file as the license text
    FileLicense,
}

configuration_option_enum! { MultilineStyle:
    // Use horizontal layout if it fits in one line, fall back to vertical
    PreferSingle,
    // Use vertical layout
    ForceMulti,
}

impl MultilineStyle {
    pub fn to_list_tactic(self) -> ListTactic {
        match self {
            MultilineStyle::PreferSingle => ListTactic::HorizontalVertical,
            MultilineStyle::ForceMulti => ListTactic::Vertical,
        }
    }
}

configuration_option_enum! { ReportTactic:
    Always,
    Unnumbered,
    Never,
}

// This trait and the following impl blocks are there so that we an use
// UCFS inside the get_docs() function on types for configs.
pub trait ConfigType {
    fn get_variant_names() -> String;
}

impl ConfigType for bool {
    fn get_variant_names() -> String {
        String::from("<boolean>")
    }
}

impl ConfigType for usize {
    fn get_variant_names() -> String {
        String::from("<unsigned integer>")
    }
}

impl ConfigType for String {
    fn get_variant_names() -> String {
        String::from("<string>")
    }
}

pub struct ConfigHelpItem {
    option_name: &'static str,
    doc_string: &'static str,
    variant_names: String,
    default: &'static str,
}

impl ConfigHelpItem {
    pub fn option_name(&self) -> &'static str {
        self.option_name
    }

    pub fn doc_string(&self) -> &'static str {
        self.doc_string
    }

    pub fn variant_names(&self) -> &String {
        &self.variant_names
    }

    pub fn default(&self) -> &'static str {
        self.default
    }
}

macro_rules! create_config {
    ($($i:ident: $ty:ty, $def:expr, $( $dstring:expr ),+ );+ $(;)*) => (
        #[derive(RustcDecodable, Clone)]
        pub struct Config {
            $(pub $i: $ty),+
        }

        // Just like the Config struct but with each property wrapped
        // as Option<T>. This is used to parse a rustfmt.toml that doesn't
        // specity all properties of `Config`.
        // We first parse into `ParsedConfig`, then create a default `Config`
        // and overwrite the properties with corresponding values from `ParsedConfig`
        #[derive(RustcDecodable, Clone)]
        pub struct ParsedConfig {
            $(pub $i: Option<$ty>),+
        }

        impl Config {

            fn fill_from_parsed_config(mut self, parsed: ParsedConfig) -> Config {
            $(
                if let Some(val) = parsed.$i {
                    self.$i = val;
                }
            )+
                self
            }

            pub fn from_toml(toml: &str) -> Config {
                let parsed = toml.parse().unwrap();
                let parsed_config:ParsedConfig = match toml::decode(parsed) {
                    Some(decoded) => decoded,
                    None => {
                        println!("Decoding config file failed. Config:\n{}", toml);
                        let parsed: toml::Value = toml.parse().unwrap();
                        println!("\n\nParsed:\n{:?}", parsed);
                        panic!();
                    }
                };
                Config::default().fill_from_parsed_config(parsed_config)
            }

            pub fn override_value(&mut self, key: &str, val: &str) {
                match key {
                    $(
                        stringify!($i) => {
                            self.$i = val.parse::<$ty>().unwrap();
                        }
                    )+
                    _ => panic!("Bad config key!")
                }
            }

            pub fn print_docs() {
                use std::cmp;
                let max = 0;
                $( let max = cmp::max(max, stringify!($i).len()+1); )+
                let mut space_str = String::with_capacity(max);
                for _ in 0..max {
                    space_str.push(' ');
                }
                println!("\nConfiguration Options:");
                $(
                    let name_raw = stringify!($i);
                    let mut name_out = String::with_capacity(max);
                    for _ in name_raw.len()..max-1 {
                        name_out.push(' ')
                    }
                    name_out.push_str(name_raw);
                    name_out.push(' ');
                    println!("{}{} Default: {:?}",
                             name_out,
                             <$ty>::get_variant_names(),
                             $def);
                    $(
                        println!("{}{}", space_str, $dstring);
                    )+
                    println!("");
                )+
            }
        }

        // Template for the default configuration
        impl Default for Config {
            fn default() -> Config {
                Config {
                    $(
                        $i: $def,
                    )+
                }
            }
        }
    )
}

create_config! {
    max_width: usize, 100, "Maximum width of each line";
    ideal_width: usize, 80, "Ideal width of each line";
    tab_spaces: usize, 4, "Number of spaces per tab";
    fn_call_width: usize, 60,
        "Maximum width of the args of a function call before faling back to vertical formatting";
    struct_lit_width: usize, 16,
        "Maximum width in the body of a struct lit before faling back to vertical formatting";
    newline_style: NewlineStyle, NewlineStyle::Unix, "Unix or Windows line endings";
    fn_brace_style: BraceStyle, BraceStyle::SameLineWhere, "Brace style for functions";
    fn_return_indent: ReturnIndent, ReturnIndent::WithArgs,
        "Location of return type in function declaration";
    fn_args_paren_newline: bool, true, "If function argument parenthases goes on a newline";
    fn_args_density: Density, Density::Tall, "Argument density in functions";
    fn_args_layout: StructLitStyle, StructLitStyle::Visual, "Layout of function arguments";
    fn_arg_indent: BlockIndentStyle, BlockIndentStyle::Visual, "Indent on function arguments";
    // Should we at least try to put the where clause on the same line as the rest of the
    // function decl?
    where_density: Density, Density::CompressedIfEmpty, "Density of a where clause";
    // Visual will be treated like Tabbed
    where_indent: BlockIndentStyle, BlockIndentStyle::Tabbed, "Indentation of a where clause";
    where_layout: ListTactic, ListTactic::Vertical, "Element layout inside a where clause";
    where_pred_indent: BlockIndentStyle, BlockIndentStyle::Visual,
        "Indentation style of a where predicate";
    generics_indent: BlockIndentStyle, BlockIndentStyle::Visual, "Indentation of generics";
    struct_trailing_comma: SeparatorTactic, SeparatorTactic::Vertical,
        "If there is a trailing comma on structs";
    struct_lit_trailing_comma: SeparatorTactic, SeparatorTactic::Vertical,
        "If there is a trailing comma on literal structs";
    struct_lit_style: StructLitStyle, StructLitStyle::Block, "Style of struct definition";
    struct_lit_multiline_style: MultilineStyle, MultilineStyle::PreferSingle,
        "Multilline style on literal structs";
    enum_trailing_comma: bool, true, "Put a trailing comma on enum declarations";
    report_todo: ReportTactic, ReportTactic::Always,
        "Report all occurences of TODO in source file comments";
    report_fixme: ReportTactic, ReportTactic::Never,
        "Report all occurences of FIXME in source file comments";
    // Alphabetically, case sensitive.
    reorder_imports: bool, false, "Reorder import statements alphabetically";
    single_line_if_else: bool, false, "Put else on same line as closing brace for if statements";
    format_strings: bool, true, "Format string literals, or leave as is";
    chains_overflow_last: bool, true, "Allow last call in method chain to break the line";
    take_source_hints: bool, true, "Retain some formatting characteristics from the source code";
    hard_tabs: bool, false, "Use tab characters for indentation, spaces for alignment";
}
<commit_msg>Fix typos in config descriptions<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

extern crate toml;

use lists::{SeparatorTactic, ListTactic};

macro_rules! configuration_option_enum{
    ($e:ident: $( $x:ident ),+ $(,)*) => {
        #[derive(Copy, Clone, Eq, PartialEq, Debug)]
        pub enum $e {
            $( $x ),+
        }

        impl_enum_decodable!($e, $( $x ),+);
    }
}

configuration_option_enum! { NewlineStyle:
    Windows, // \r\n
    Unix, // \n
}

configuration_option_enum! { BraceStyle:
    AlwaysNextLine,
    PreferSameLine,
    // Prefer same line except where there is a where clause, in which case force
    // the brace to the next line.
    SameLineWhere,
}

// How to indent a function's return type.
configuration_option_enum! { ReturnIndent:
    // Aligned with the arguments
    WithArgs,
    // Aligned with the where clause
    WithWhereClause,
}

// How to stle a struct literal.
configuration_option_enum! { StructLitStyle:
    // First line on the same line as the opening brace, all lines aligned with
    // the first line.
    Visual,
    // First line is on a new line and all lines align with block indent.
    Block,
    // FIXME Maybe we should also have an option to align types.
}

configuration_option_enum! { BlockIndentStyle:
    // Same level as parent.
    Inherit,
    // One level deeper than parent.
    Tabbed,
    // Aligned with block open.
    Visual,
}

configuration_option_enum! { Density:
    // Fit as much on one line as possible.
    Compressed,
    // Use more lines.
    Tall,
    // Try to compress if the body is empty.
    CompressedIfEmpty,
}

impl Density {
    pub fn to_list_tactic(self) -> ListTactic {
        match self {
            Density::Compressed => ListTactic::Mixed,
            Density::Tall | Density::CompressedIfEmpty => ListTactic::HorizontalVertical,
        }
    }
}

configuration_option_enum! { LicensePolicy:
    // Do not place license text at top of files
    NoLicense,
    // Use the text in "license" field as the license
    TextLicense,
    // Use a text file as the license text
    FileLicense,
}

configuration_option_enum! { MultilineStyle:
    // Use horizontal layout if it fits in one line, fall back to vertical
    PreferSingle,
    // Use vertical layout
    ForceMulti,
}

impl MultilineStyle {
    pub fn to_list_tactic(self) -> ListTactic {
        match self {
            MultilineStyle::PreferSingle => ListTactic::HorizontalVertical,
            MultilineStyle::ForceMulti => ListTactic::Vertical,
        }
    }
}

configuration_option_enum! { ReportTactic:
    Always,
    Unnumbered,
    Never,
}

// This trait and the following impl blocks are there so that we an use
// UCFS inside the get_docs() function on types for configs.
pub trait ConfigType {
    fn get_variant_names() -> String;
}

impl ConfigType for bool {
    fn get_variant_names() -> String {
        String::from("<boolean>")
    }
}

impl ConfigType for usize {
    fn get_variant_names() -> String {
        String::from("<unsigned integer>")
    }
}

impl ConfigType for String {
    fn get_variant_names() -> String {
        String::from("<string>")
    }
}

pub struct ConfigHelpItem {
    option_name: &'static str,
    doc_string: &'static str,
    variant_names: String,
    default: &'static str,
}

impl ConfigHelpItem {
    pub fn option_name(&self) -> &'static str {
        self.option_name
    }

    pub fn doc_string(&self) -> &'static str {
        self.doc_string
    }

    pub fn variant_names(&self) -> &String {
        &self.variant_names
    }

    pub fn default(&self) -> &'static str {
        self.default
    }
}

macro_rules! create_config {
    ($($i:ident: $ty:ty, $def:expr, $( $dstring:expr ),+ );+ $(;)*) => (
        #[derive(RustcDecodable, Clone)]
        pub struct Config {
            $(pub $i: $ty),+
        }

        // Just like the Config struct but with each property wrapped
        // as Option<T>. This is used to parse a rustfmt.toml that doesn't
        // specity all properties of `Config`.
        // We first parse into `ParsedConfig`, then create a default `Config`
        // and overwrite the properties with corresponding values from `ParsedConfig`
        #[derive(RustcDecodable, Clone)]
        pub struct ParsedConfig {
            $(pub $i: Option<$ty>),+
        }

        impl Config {

            fn fill_from_parsed_config(mut self, parsed: ParsedConfig) -> Config {
            $(
                if let Some(val) = parsed.$i {
                    self.$i = val;
                }
            )+
                self
            }

            pub fn from_toml(toml: &str) -> Config {
                let parsed = toml.parse().unwrap();
                let parsed_config:ParsedConfig = match toml::decode(parsed) {
                    Some(decoded) => decoded,
                    None => {
                        println!("Decoding config file failed. Config:\n{}", toml);
                        let parsed: toml::Value = toml.parse().unwrap();
                        println!("\n\nParsed:\n{:?}", parsed);
                        panic!();
                    }
                };
                Config::default().fill_from_parsed_config(parsed_config)
            }

            pub fn override_value(&mut self, key: &str, val: &str) {
                match key {
                    $(
                        stringify!($i) => {
                            self.$i = val.parse::<$ty>().unwrap();
                        }
                    )+
                    _ => panic!("Bad config key!")
                }
            }

            pub fn print_docs() {
                use std::cmp;
                let max = 0;
                $( let max = cmp::max(max, stringify!($i).len()+1); )+
                let mut space_str = String::with_capacity(max);
                for _ in 0..max {
                    space_str.push(' ');
                }
                println!("\nConfiguration Options:");
                $(
                    let name_raw = stringify!($i);
                    let mut name_out = String::with_capacity(max);
                    for _ in name_raw.len()..max-1 {
                        name_out.push(' ')
                    }
                    name_out.push_str(name_raw);
                    name_out.push(' ');
                    println!("{}{} Default: {:?}",
                             name_out,
                             <$ty>::get_variant_names(),
                             $def);
                    $(
                        println!("{}{}", space_str, $dstring);
                    )+
                    println!("");
                )+
            }
        }

        // Template for the default configuration
        impl Default for Config {
            fn default() -> Config {
                Config {
                    $(
                        $i: $def,
                    )+
                }
            }
        }
    )
}

create_config! {
    max_width: usize, 100, "Maximum width of each line";
    ideal_width: usize, 80, "Ideal width of each line";
    tab_spaces: usize, 4, "Number of spaces per tab";
    fn_call_width: usize, 60,
        "Maximum width of the args of a function call before falling back to vertical formatting";
    struct_lit_width: usize, 16,
        "Maximum width in the body of a struct lit before falling back to vertical formatting";
    newline_style: NewlineStyle, NewlineStyle::Unix, "Unix or Windows line endings";
    fn_brace_style: BraceStyle, BraceStyle::SameLineWhere, "Brace style for functions";
    fn_return_indent: ReturnIndent, ReturnIndent::WithArgs,
        "Location of return type in function declaration";
    fn_args_paren_newline: bool, true, "If function argument parenthesis goes on a newline";
    fn_args_density: Density, Density::Tall, "Argument density in functions";
    fn_args_layout: StructLitStyle, StructLitStyle::Visual, "Layout of function arguments";
    fn_arg_indent: BlockIndentStyle, BlockIndentStyle::Visual, "Indent on function arguments";
    // Should we at least try to put the where clause on the same line as the rest of the
    // function decl?
    where_density: Density, Density::CompressedIfEmpty, "Density of a where clause";
    // Visual will be treated like Tabbed
    where_indent: BlockIndentStyle, BlockIndentStyle::Tabbed, "Indentation of a where clause";
    where_layout: ListTactic, ListTactic::Vertical, "Element layout inside a where clause";
    where_pred_indent: BlockIndentStyle, BlockIndentStyle::Visual,
        "Indentation style of a where predicate";
    generics_indent: BlockIndentStyle, BlockIndentStyle::Visual, "Indentation of generics";
    struct_trailing_comma: SeparatorTactic, SeparatorTactic::Vertical,
        "If there is a trailing comma on structs";
    struct_lit_trailing_comma: SeparatorTactic, SeparatorTactic::Vertical,
        "If there is a trailing comma on literal structs";
    struct_lit_style: StructLitStyle, StructLitStyle::Block, "Style of struct definition";
    struct_lit_multiline_style: MultilineStyle, MultilineStyle::PreferSingle,
        "Multiline style on literal structs";
    enum_trailing_comma: bool, true, "Put a trailing comma on enum declarations";
    report_todo: ReportTactic, ReportTactic::Always,
        "Report all occurrences of TODO in source file comments";
    report_fixme: ReportTactic, ReportTactic::Never,
        "Report all occurrences of FIXME in source file comments";
    // Alphabetically, case sensitive.
    reorder_imports: bool, false, "Reorder import statements alphabetically";
    single_line_if_else: bool, false, "Put else on same line as closing brace for if statements";
    format_strings: bool, true, "Format string literals, or leave as is";
    chains_overflow_last: bool, true, "Allow last call in method chain to break the line";
    take_source_hints: bool, true, "Retain some formatting characteristics from the source code";
    hard_tabs: bool, false, "Use tab characters for indentation, spaces for alignment";
}
<|endoftext|>


<commit_before>// Copyright Istio Authors
//
// 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.

use crate::identity;
use std::net::{IpAddr, Ipv6Addr, SocketAddr};
use std::path::PathBuf;
use std::time::Duration;
use tokio::time;

#[derive(Clone, Debug)]
pub struct Config {
    pub window_size: u32,
    pub connection_window_size: u32,
    pub frame_size: u32,

    pub socks5_addr: SocketAddr,
    pub admin_addr: SocketAddr,
    pub inbound_addr: SocketAddr,
    pub inbound_plaintext_addr: SocketAddr,
    pub outbound_addr: SocketAddr,

    /// The name of the node this ztunnel is running as.
    pub local_node: Option<String>,

    /// XDS address to use. If unset, XDS will not be used.
    pub xds_address: Option<String>,
    /// Filepath to a local xds file for workloads, as YAML.
    pub local_xds_path: Option<String>,
    /// If true, on-demand XDS will be used
    pub xds_on_demand: bool,

    pub auth: identity::AuthSource,

    pub termination_grace_period: time::Duration,

    /// Specify the number of worker threads the Tokio Runtime will use.
    pub num_worker_threads: usize,
}

const DEFAULT_WORKER_THREADS: usize = 2;

impl Default for Config {
    fn default() -> Config {
        // TODO: copy JWT auth logic from CA client and use TLS here (port 15012)
        let xds_address = Some(if std::env::var("KUBERNETES_SERVICE_HOST").is_ok() {
            "https://istiod.istio-system:15012".to_string()
        } else {
            "https://localhost:15012".to_string()
        });
        Config {
            window_size: 4 * 1024 * 1024,
            connection_window_size: 4 * 1024 * 1024,
            frame_size: 1024 * 1024,

            termination_grace_period: Duration::from_secs(5),

            admin_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15021),
            socks5_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15080),
            inbound_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15008),
            inbound_plaintext_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15006),
            outbound_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15001),

            local_node: std::env::var("NODE_NAME").ok(),

            xds_address,
            local_xds_path: std::env::var("LOCAL_XDS_PATH").ok(),
            xds_on_demand: std::env::var("XDS_ON_DEMAND").ok().as_deref() == Some("on"),

            auth: identity::AuthSource::Token(PathBuf::from(
                r"./var/run/secrets/tokens/istio-token",
            )),

            num_worker_threads: std::env::var("ZTUNNEL_WOKRER_THREADS")
                .ok()
                .map(|v| {
                    v.parse::<usize>()
                        .ok()
                        .filter(|n| *n > 0)
                        .unwrap_or(DEFAULT_WORKER_THREADS)
                })
                .unwrap_or(DEFAULT_WORKER_THREADS),
        }
    }
}
<commit_msg>Fix the worker threads environment name typo (#189)<commit_after>// Copyright Istio Authors
//
// 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.

use crate::identity;
use std::net::{IpAddr, Ipv6Addr, SocketAddr};
use std::path::PathBuf;
use std::time::Duration;
use tokio::time;

#[derive(Clone, Debug)]
pub struct Config {
    pub window_size: u32,
    pub connection_window_size: u32,
    pub frame_size: u32,

    pub socks5_addr: SocketAddr,
    pub admin_addr: SocketAddr,
    pub inbound_addr: SocketAddr,
    pub inbound_plaintext_addr: SocketAddr,
    pub outbound_addr: SocketAddr,

    /// The name of the node this ztunnel is running as.
    pub local_node: Option<String>,

    /// XDS address to use. If unset, XDS will not be used.
    pub xds_address: Option<String>,
    /// Filepath to a local xds file for workloads, as YAML.
    pub local_xds_path: Option<String>,
    /// If true, on-demand XDS will be used
    pub xds_on_demand: bool,

    pub auth: identity::AuthSource,

    pub termination_grace_period: time::Duration,

    /// Specify the number of worker threads the Tokio Runtime will use.
    pub num_worker_threads: usize,
}

const DEFAULT_WORKER_THREADS: usize = 2;

impl Default for Config {
    fn default() -> Config {
        // TODO: copy JWT auth logic from CA client and use TLS here (port 15012)
        let xds_address = Some(if std::env::var("KUBERNETES_SERVICE_HOST").is_ok() {
            "https://istiod.istio-system:15012".to_string()
        } else {
            "https://localhost:15012".to_string()
        });
        Config {
            window_size: 4 * 1024 * 1024,
            connection_window_size: 4 * 1024 * 1024,
            frame_size: 1024 * 1024,

            termination_grace_period: Duration::from_secs(5),

            admin_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15021),
            socks5_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15080),
            inbound_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15008),
            inbound_plaintext_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15006),
            outbound_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15001),

            local_node: std::env::var("NODE_NAME").ok(),

            xds_address,
            local_xds_path: std::env::var("LOCAL_XDS_PATH").ok(),
            xds_on_demand: std::env::var("XDS_ON_DEMAND").ok().as_deref() == Some("on"),

            auth: identity::AuthSource::Token(PathBuf::from(
                r"./var/run/secrets/tokens/istio-token",
            )),

            num_worker_threads: std::env::var("ZTUNNEL_WORKER_THREADS")
                .ok()
                .map(|v| {
                    v.parse::<usize>()
                        .ok()
                        .filter(|n| *n > 0)
                        .unwrap_or(DEFAULT_WORKER_THREADS)
                })
                .unwrap_or(DEFAULT_WORKER_THREADS),
        }
    }
}
<|endoftext|>


<commit_before>extern crate std;

pub trait ConfigValue: Send + Clone + std::fmt::Show {
}

pub struct Config {
	params: std::collections::HashMap<String, Box<ConfigValue>>
}

impl ConfigValue for i32 {}
impl ConfigValue for String {}
impl ConfigValue for bool {}
<commit_msg>Minor refactor of config<commit_after>use std::collections::HashMap;
use std::fmt::Show;

pub trait ConfigValue: Send + Clone + Show {
}

pub struct Config {
	params: HashMap<String, Box<ConfigValue>>
}

impl ConfigValue for int {}
impl ConfigValue for String {}
impl ConfigValue for bool {}
<|endoftext|>