<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 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("\"", "\"\"")),
        }
    }
}
<commit_msg>Add comment<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 {
    // Tree-sitter row, column values are 0-based while CodeQL starts
    // counting at 1. In addition Tree-sitter's row and column for the
    // end position are exclusive while CodeQL's end positions are inclusive.
    // This means that all values should be incremented by 1 and in addition the
    // end position needs to be shift 1 to the left. In most cases this means
    // simply incrementing all values except the end column except in cases where
    // the end column is 0 (start of a line). In such cases the end position must be
    // set to the end of the previous line.
    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::collections::HashMap;
use hash_ring::{NodeInfo, HashRing};

use super::{MemcacheResult, Commands, };
use super::connection::Connection;

// Actual memcached client. Holds all connections and perform appropriate job using
// Consistent Hash Ring to balanced between servers and virtual replicas.
pub struct Client {
    ring : HashRing<NodeInfo>,
    connections : HashMap<String, Connection>
}

impl Client {
    // `replicas` are copies of nodes that point to real servers.
    fn new(nodes: Vec<NodeInfo>, replicas: isize) -> MemcacheResult<Client> {
        let mut new_client = Client {
           ring : HashRing::new(nodes.clone(), replicas),
           connections : HashMap::with_capacity(nodes.len())
        };

        // since BufStream can't be clonned. The connections reside in Client. So node lookups are then mapped
        // and commands executed using the right connection regardles if a read or virtual node.
        for n in nodes {
            let conn = try!{Connection::connect(n.host, n.port)};
            new_client.connections.insert(n.to_string(), conn);
        }

        Ok(new_client)
    }
}

impl Commands for Client {

    fn set(&mut self, key: &str, value: &[u8], exptime: isize, flags: u16) -> MemcacheResult<bool> {
        let node = self.ring.get_node(key.to_string());
        let conn = self.connections.get_mut(&node.to_string()).expect("Inexistent Connection");
        conn.set(key, value, exptime, flags)
    }

    fn get(&mut self, key: &str) -> MemcacheResult<Option<(Vec<u8>, u16)>> {
        let node = self.ring.get_node(key.to_string());
        let conn = self.connections.get_mut(&node.to_string()).expect("Inexistent Connection");
        conn.get(key)
    }

    fn flush(&mut self) -> MemcacheResult<()> {
        for (_, conn) in self.connections.iter_mut() {
            match conn.flush() {
                Ok(_) => continue,
                e => return e
            }
        }
        Ok(())
    }

    fn delete(&mut self, key: &str) -> MemcacheResult<bool> {
        let node = self.ring.get_node(key.to_string());
        let conn = self.connections.get_mut(&node.to_string()).expect("Inexistent Connection");
        conn.delete(key)
    }

    fn incr(&mut self, key: &str, value: u64) -> MemcacheResult<Option<(isize)>> {
        let node = self.ring.get_node(key.to_string());
        let conn = self.connections.get_mut(&node.to_string()).expect("Inexistent Connection");
        conn.incr(key, value)
    }

    fn decr(&mut self, key: &str, value: u64) -> MemcacheResult<Option<(isize)>> {
        let node = self.ring.get_node(key.to_string());
        let conn = self.connections.get_mut(&node.to_string()).expect("Inexistent Connection");
        conn.decr(key, value)
    }

}

#[test]
fn test_client() {
    let mut nodes: Vec<NodeInfo> = Vec::new();
    nodes.push(NodeInfo{host: "localhost", port: 2333});
    nodes.push(NodeInfo{host: "localhost", port: 2334});

    let client = Client::new(nodes, 2);
    assert! { client.is_ok() };
}

#[test]
fn test_flush() {
    let mut nodes: Vec<NodeInfo> = Vec::new();
    nodes.push(NodeInfo{host: "localhost", port: 2333});
    nodes.push(NodeInfo{host: "localhost", port: 2334});

    let mut client = Client::new(nodes, 2).ok().unwrap();
    assert!{ client.flush().is_ok() };
}

#[test]
fn test_set() {
    let mut nodes: Vec<NodeInfo> = Vec::new();
    nodes.push(NodeInfo{host: "localhost", port: 2333});
    nodes.push(NodeInfo{host: "localhost", port: 2334});

    let mut client = Client::new(nodes, 2).ok().unwrap();
    assert!{ client.flush().is_ok() };
    assert!{ client.set("fooc", b"bar", 10, 0).ok().unwrap() == true };
}

#[test]
fn test_get() {
    let mut nodes: Vec<NodeInfo> = Vec::new();
    nodes.push(NodeInfo{host: "localhost", port: 2333});
    nodes.push(NodeInfo{host: "localhost", port: 2334});

    let mut client = Client::new(nodes, 2).ok().unwrap();

    assert!{ client.flush().is_ok() };
    assert!{ client.get("fooc").ok().unwrap() == None };

    assert!{ client.set("fooc", b"bar", 0, 10).ok().unwrap() == true };
    let result = client.get("fooc");
    let result_tuple = result.ok().unwrap().unwrap();
    assert!{ result_tuple.0 == b"bar" };
    assert!{ result_tuple.1 == 10 };
}

#[test]
fn test_delete() {
    let mut nodes: Vec<NodeInfo> = Vec::new();
    nodes.push(NodeInfo{host: "localhost", port: 2333});
    nodes.push(NodeInfo{host: "localhost", port: 2334});

    let mut client = Client::new(nodes, 2).ok().unwrap();

    assert!{ client.flush().is_ok() };
    assert!{ client.delete("fooc").ok().unwrap() == false };
}

#[test]
fn test_incr() {
    let mut nodes: Vec<NodeInfo> = Vec::new();
    nodes.push(NodeInfo{host: "localhost", port: 2333});
    nodes.push(NodeInfo{host: "localhost", port: 2334});

    let mut client = Client::new(nodes, 2).ok().unwrap();
    assert!{ client.flush().is_ok() };
    let mut result = client.incr("liec", 42);
    assert!{ result.ok().unwrap() == None };

    assert!{ client.flush().is_ok() };
    client.set("truthc", b"42", 0, 0).ok().unwrap();
    result = client.incr("truthc", 1);
    assert!{ result.ok().unwrap().unwrap() == 43 };
}

#[test]
fn test_decr() {
    let mut nodes: Vec<NodeInfo> = Vec::new();
    nodes.push(NodeInfo{host: "localhost", port: 2333});
    nodes.push(NodeInfo{host: "localhost", port: 2334});
    let mut client = Client::new(nodes, 2).ok().unwrap();
    assert!{ client.flush().is_ok() };

    let mut result = client.decr("lie", 42);
    assert!{ result.ok().unwrap() == None };

    assert!{ client.flush().is_ok() };
    client.set("truthc", b"42", 0, 0).ok().unwrap();
    result = client.decr("truthc", 1);
    assert!{ result.ok().unwrap().unwrap() == 41 };
}
<commit_msg>add Client#get_connection_by_key<commit_after>use std::collections::HashMap;
use hash_ring::{NodeInfo, HashRing};

use super::{MemcacheResult, Commands, };
use super::connection::Connection;

// Actual memcached client. Holds all connections and perform appropriate job using
// Consistent Hash Ring to balanced between servers and virtual replicas.
pub struct Client {
    ring : HashRing<NodeInfo>,
    connections : HashMap<String, Connection>
}

impl Client {
    // `replicas` are copies of nodes that point to real servers.
    fn new(nodes: Vec<NodeInfo>, replicas: isize) -> MemcacheResult<Client> {
        let mut new_client = Client {
           ring : HashRing::new(nodes.clone(), replicas),
           connections : HashMap::with_capacity(nodes.len())
        };

        // since BufStream can't be clonned. The connections reside in Client. So node lookups are then mapped
        // and commands executed using the right connection regardles if a read or virtual node.
        for n in nodes {
            let conn = try!{Connection::connect(n.host, n.port)};
            new_client.connections.insert(n.to_string(), conn);
        }

        Ok(new_client)
    }
}

impl Client {

    fn get_connection_by_key(&mut self, key :&str) -> &mut Connection {
        let node = self.ring.get_node(key.to_string());
        let conn = self.connections.get_mut(&node.to_string()).expect("Inexistent Connection");
        return conn;
    }

}

impl Commands for Client {

    fn set(&mut self, key: &str, value: &[u8], exptime: isize, flags: u16) -> MemcacheResult<bool> {
        let conn = self.get_connection_by_key(key);
        conn.set(key, value, exptime, flags)
    }

    fn get(&mut self, key: &str) -> MemcacheResult<Option<(Vec<u8>, u16)>> {
        let conn = self.get_connection_by_key(key);
        conn.get(key)
    }

    fn flush(&mut self) -> MemcacheResult<()> {
        for (_, conn) in self.connections.iter_mut() {
            match conn.flush() {
                Ok(_) => continue,
                e => return e
            }
        }
        Ok(())
    }

    fn delete(&mut self, key: &str) -> MemcacheResult<bool> {
        let conn = self.get_connection_by_key(key);
        conn.delete(key)
    }

    fn incr(&mut self, key: &str, value: u64) -> MemcacheResult<Option<(isize)>> {
        let conn = self.get_connection_by_key(key);
        conn.incr(key, value)
    }

    fn decr(&mut self, key: &str, value: u64) -> MemcacheResult<Option<(isize)>> {
        let conn = self.get_connection_by_key(key);
        conn.decr(key, value)
    }

}

#[test]
fn test_client() {
    let mut nodes: Vec<NodeInfo> = Vec::new();
    nodes.push(NodeInfo{host: "localhost", port: 2333});
    nodes.push(NodeInfo{host: "localhost", port: 2334});

    let client = Client::new(nodes, 2);
    assert! { client.is_ok() };
}

#[test]
fn test_flush() {
    let mut nodes: Vec<NodeInfo> = Vec::new();
    nodes.push(NodeInfo{host: "localhost", port: 2333});
    nodes.push(NodeInfo{host: "localhost", port: 2334});

    let mut client = Client::new(nodes, 2).ok().unwrap();
    assert!{ client.flush().is_ok() };
}

#[test]
fn test_set() {
    let mut nodes: Vec<NodeInfo> = Vec::new();
    nodes.push(NodeInfo{host: "localhost", port: 2333});
    nodes.push(NodeInfo{host: "localhost", port: 2334});

    let mut client = Client::new(nodes, 2).ok().unwrap();
    assert!{ client.flush().is_ok() };
    assert!{ client.set("fooc", b"bar", 10, 0).ok().unwrap() == true };
}

#[test]
fn test_get() {
    let mut nodes: Vec<NodeInfo> = Vec::new();
    nodes.push(NodeInfo{host: "localhost", port: 2333});
    nodes.push(NodeInfo{host: "localhost", port: 2334});

    let mut client = Client::new(nodes, 2).ok().unwrap();

    assert!{ client.flush().is_ok() };
    assert!{ client.get("fooc").ok().unwrap() == None };

    assert!{ client.set("fooc", b"bar", 0, 10).ok().unwrap() == true };
    let result = client.get("fooc");
    let result_tuple = result.ok().unwrap().unwrap();
    assert!{ result_tuple.0 == b"bar" };
    assert!{ result_tuple.1 == 10 };
}

#[test]
fn test_delete() {
    let mut nodes: Vec<NodeInfo> = Vec::new();
    nodes.push(NodeInfo{host: "localhost", port: 2333});
    nodes.push(NodeInfo{host: "localhost", port: 2334});

    let mut client = Client::new(nodes, 2).ok().unwrap();

    assert!{ client.flush().is_ok() };
    assert!{ client.delete("fooc").ok().unwrap() == false };
}

#[test]
fn test_incr() {
    let mut nodes: Vec<NodeInfo> = Vec::new();
    nodes.push(NodeInfo{host: "localhost", port: 2333});
    nodes.push(NodeInfo{host: "localhost", port: 2334});

    let mut client = Client::new(nodes, 2).ok().unwrap();
    assert!{ client.flush().is_ok() };
    let mut result = client.incr("liec", 42);
    assert!{ result.ok().unwrap() == None };

    assert!{ client.flush().is_ok() };
    client.set("truthc", b"42", 0, 0).ok().unwrap();
    result = client.incr("truthc", 1);
    assert!{ result.ok().unwrap().unwrap() == 43 };
}

#[test]
fn test_decr() {
    let mut nodes: Vec<NodeInfo> = Vec::new();
    nodes.push(NodeInfo{host: "localhost", port: 2333});
    nodes.push(NodeInfo{host: "localhost", port: 2334});
    let mut client = Client::new(nodes, 2).ok().unwrap();
    assert!{ client.flush().is_ok() };

    let mut result = client.decr("lie", 42);
    assert!{ result.ok().unwrap() == None };

    assert!{ client.flush().is_ok() };
    client.set("truthc", b"42", 0, 0).ok().unwrap();
    result = client.decr("truthc", 1);
    assert!{ result.ok().unwrap().unwrap() == 41 };
}
<|endoftext|>


<commit_before>use std::cell::{RefCell, RefMut, UnsafeCell};
use std::collections::VecDeque;
use std::fmt::Write;
use std::rc::Rc;
use std::time::{Duration, Instant};
use std::{env, fmt, net};

use bytes::BytesMut;
use futures::{future, Future};
use futures_cpupool::CpuPool;
use http::StatusCode;
use lazycell::LazyCell;
use parking_lot::Mutex;
use time;
use tokio_current_thread::spawn;
use tokio_timer::{sleep, Delay};

use super::channel::Node;
use super::message::{Request, RequestPool};
use super::KeepAlive;
use body::Body;
use httpresponse::{HttpResponse, HttpResponseBuilder, HttpResponsePool};

/// Env variable for default cpu pool size
const ENV_CPU_POOL_VAR: &str = "ACTIX_CPU_POOL";

lazy_static! {
    pub(crate) static ref DEFAULT_CPUPOOL: Mutex<CpuPool> = {
        let default = match env::var(ENV_CPU_POOL_VAR) {
            Ok(val) => {
                if let Ok(val) = val.parse() {
                    val
                } else {
                    error!("Can not parse ACTIX_CPU_POOL value");
                    20
                }
            }
            Err(_) => 20,
        };
        Mutex::new(CpuPool::new(default))
    };
}

/// Various server settings
pub struct ServerSettings {
    addr: net::SocketAddr,
    secure: bool,
    host: String,
    cpu_pool: LazyCell<CpuPool>,
    responses: &'static HttpResponsePool,
}

impl Clone for ServerSettings {
    fn clone(&self) -> Self {
        ServerSettings {
            addr: self.addr,
            secure: self.secure,
            host: self.host.clone(),
            cpu_pool: LazyCell::new(),
            responses: HttpResponsePool::get_pool(),
        }
    }
}

impl Default for ServerSettings {
    fn default() -> Self {
        ServerSettings {
            addr: "127.0.0.1:8080".parse().unwrap(),
            secure: false,
            host: "localhost:8080".to_owned(),
            responses: HttpResponsePool::get_pool(),
            cpu_pool: LazyCell::new(),
        }
    }
}

impl ServerSettings {
    /// Crate server settings instance
    pub(crate) fn new(
        addr: net::SocketAddr, host: &str, secure: bool,
    ) -> ServerSettings {
        let host = host.to_owned();
        let cpu_pool = LazyCell::new();
        let responses = HttpResponsePool::get_pool();
        ServerSettings {
            addr,
            secure,
            host,
            cpu_pool,
            responses,
        }
    }

    /// Returns the socket address of the local half of this TCP connection
    pub fn local_addr(&self) -> net::SocketAddr {
        self.addr
    }

    /// Returns true if connection is secure(https)
    pub fn secure(&self) -> bool {
        self.secure
    }

    /// Returns host header value
    pub fn host(&self) -> &str {
        &self.host
    }

    /// Returns default `CpuPool` for server
    pub fn cpu_pool(&self) -> &CpuPool {
        self.cpu_pool.borrow_with(|| DEFAULT_CPUPOOL.lock().clone())
    }

    #[inline]
    pub(crate) fn get_response(&self, status: StatusCode, body: Body) -> HttpResponse {
        HttpResponsePool::get_response(&self.responses, status, body)
    }

    #[inline]
    pub(crate) fn get_response_builder(
        &self, status: StatusCode,
    ) -> HttpResponseBuilder {
        HttpResponsePool::get_builder(&self.responses, status)
    }
}

// "Sun, 06 Nov 1994 08:49:37 GMT".len()
const DATE_VALUE_LENGTH: usize = 29;

/// Http service configuration
pub struct ServiceConfig<H>(Rc<Inner<H>>);

struct Inner<H> {
    handler: H,
    keep_alive: Option<Duration>,
    client_timeout: u64,
    client_shutdown: u64,
    ka_enabled: bool,
    bytes: Rc<SharedBytesPool>,
    messages: &'static RequestPool,
    node: RefCell<Node<()>>,
    date: UnsafeCell<(bool, Date)>,
}

impl<H> Clone for ServiceConfig<H> {
    fn clone(&self) -> Self {
        ServiceConfig(self.0.clone())
    }
}

impl<H> ServiceConfig<H> {
    /// Create instance of `ServiceConfig`
    pub(crate) fn new(
        handler: H, keep_alive: KeepAlive, client_timeout: u64, client_shutdown: u64,
        settings: ServerSettings,
    ) -> ServiceConfig<H> {
        let (keep_alive, ka_enabled) = match keep_alive {
            KeepAlive::Timeout(val) => (val as u64, true),
            KeepAlive::Os | KeepAlive::Tcp(_) => (0, true),
            KeepAlive::Disabled => (0, false),
        };
        let keep_alive = if ka_enabled && keep_alive > 0 {
            Some(Duration::from_secs(keep_alive))
        } else {
            None
        };

        ServiceConfig(Rc::new(Inner {
            handler,
            keep_alive,
            ka_enabled,
            client_timeout,
            client_shutdown,
            bytes: Rc::new(SharedBytesPool::new()),
            messages: RequestPool::pool(settings),
            node: RefCell::new(Node::head()),
            date: UnsafeCell::new((false, Date::new())),
        }))
    }

    /// Create worker settings builder.
    pub fn build(handler: H) -> ServiceConfigBuilder<H> {
        ServiceConfigBuilder::new(handler)
    }

    pub(crate) fn head(&self) -> RefMut<Node<()>> {
        self.0.node.borrow_mut()
    }

    pub(crate) fn handler(&self) -> &H {
        &self.0.handler
    }

    #[inline]
    /// Keep alive duration if configured.
    pub fn keep_alive(&self) -> Option<Duration> {
        self.0.keep_alive
    }

    #[inline]
    /// Return state of connection keep-alive funcitonality
    pub fn keep_alive_enabled(&self) -> bool {
        self.0.ka_enabled
    }

    pub(crate) fn get_bytes(&self) -> BytesMut {
        self.0.bytes.get_bytes()
    }

    pub(crate) fn release_bytes(&self, bytes: BytesMut) {
        self.0.bytes.release_bytes(bytes)
    }

    pub(crate) fn get_request(&self) -> Request {
        RequestPool::get(self.0.messages)
    }

    fn update_date(&self) {
        // Unsafe: WorkerSetting is !Sync and !Send
        unsafe { (*self.0.date.get()).0 = false };
    }
}

impl<H: 'static> ServiceConfig<H> {
    #[inline]
    /// Client timeout for first request.
    pub fn client_timer(&self) -> Option<Delay> {
        let delay = self.0.client_timeout;
        if delay != 0 {
            Some(Delay::new(self.now() + Duration::from_millis(delay)))
        } else {
            None
        }
    }

    /// Client timeout for first request.
    pub fn client_timer_expire(&self) -> Option<Instant> {
        let delay = self.0.client_timeout;
        if delay != 0 {
            Some(self.now() + Duration::from_millis(delay))
        } else {
            None
        }
    }

    /// Client shutdown timer
    pub fn client_shutdown_timer(&self) -> Option<Instant> {
        let delay = self.0.client_shutdown;
        if delay != 0 {
            Some(self.now() + Duration::from_millis(delay))
        } else {
            None
        }
    }

    #[inline]
    /// Return keep-alive timer delay is configured.
    pub fn keep_alive_timer(&self) -> Option<Delay> {
        if let Some(ka) = self.0.keep_alive {
            Some(Delay::new(self.now() + ka))
        } else {
            None
        }
    }

    /// Keep-alive expire time
    pub fn keep_alive_expire(&self) -> Option<Instant> {
        if let Some(ka) = self.0.keep_alive {
            Some(self.now() + ka)
        } else {
            None
        }
    }

    pub(crate) fn set_date(&self, dst: &mut BytesMut, full: bool) {
        // Unsafe: WorkerSetting is !Sync and !Send
        let date_bytes = unsafe {
            let date = &mut (*self.0.date.get());
            if !date.0 {
                date.1.update();
                date.0 = true;

                // periodic date update
                let s = self.clone();
                spawn(sleep(Duration::from_millis(500)).then(move |_| {
                    s.update_date();
                    future::ok(())
                }));
            }
            &date.1.bytes
        };
        if full {
            let mut buf: [u8; 39] = [0; 39];
            buf[..6].copy_from_slice(b"date: ");
            buf[6..35].copy_from_slice(date_bytes);
            buf[35..].copy_from_slice(b"\r\n\r\n");
            dst.extend_from_slice(&buf);
        } else {
            dst.extend_from_slice(date_bytes);
        }
    }

    #[inline]
    pub(crate) fn now(&self) -> Instant {
        unsafe {
            let date = &mut (*self.0.date.get());
            if !date.0 {
                date.1.update();
                date.0 = true;

                // periodic date update
                let s = self.clone();
                spawn(sleep(Duration::from_millis(500)).then(move |_| {
                    s.update_date();
                    future::ok(())
                }));
            }
            date.1.current
        }
    }
}

/// A service config builder
///
/// This type can be used to construct an instance of `ServiceConfig` through a
/// builder-like pattern.
pub struct ServiceConfigBuilder<H> {
    handler: H,
    keep_alive: KeepAlive,
    client_timeout: u64,
    client_shutdown: u64,
    host: String,
    addr: net::SocketAddr,
    secure: bool,
}

impl<H> ServiceConfigBuilder<H> {
    /// Create instance of `ServiceConfigBuilder`
    pub fn new(handler: H) -> ServiceConfigBuilder<H> {
        ServiceConfigBuilder {
            handler,
            keep_alive: KeepAlive::Timeout(5),
            client_timeout: 5000,
            client_shutdown: 5000,
            secure: false,
            host: "localhost".to_owned(),
            addr: "127.0.0.1:8080".parse().unwrap(),
        }
    }

    /// Enable secure flag for current server.
    ///
    /// By default this flag is set to false.
    pub fn secure(mut self) -> Self {
        self.secure = true;
        self
    }

    /// Set server keep-alive setting.
    ///
    /// By default keep alive is set to a 5 seconds.
    pub fn keep_alive<T: Into<KeepAlive>>(mut self, val: T) -> Self {
        self.keep_alive = val.into();
        self
    }

    /// Set server client timeout in milliseconds for first request.
    ///
    /// Defines a timeout for reading client request header. If a client does not transmit
    /// the entire set headers within this time, the request is terminated with
    /// the 408 (Request Time-out) error.
    ///
    /// To disable timeout set value to 0.
    ///
    /// By default client timeout is set to 5000 milliseconds.
    pub fn client_timeout(mut self, val: u64) -> Self {
        self.client_timeout = val;
        self
    }

    /// Set server connection shutdown timeout in milliseconds.
    ///
    /// Defines a timeout for shutdown connection. If a shutdown procedure does not complete
    /// within this time, the request is dropped. This timeout affects only secure connections.
    ///
    /// To disable timeout set value to 0.
    ///
    /// By default client timeout is set to 5000 milliseconds.
    pub fn client_shutdown(mut self, val: u64) -> Self {
        self.client_shutdown = val;
        self
    }

    /// Set server host name.
    ///
    /// Host name is used by application router aa a hostname for url
    /// generation. Check [ConnectionInfo](./dev/struct.ConnectionInfo.
    /// html#method.host) documentation for more information.
    ///
    /// By default host name is set to a "localhost" value.
    pub fn server_hostname(mut self, val: &str) -> Self {
        self.host = val.to_owned();
        self
    }

    /// Set server ip address.
    ///
    /// Host name is used by application router aa a hostname for url
    /// generation. Check [ConnectionInfo](./dev/struct.ConnectionInfo.
    /// html#method.host) documentation for more information.
    ///
    /// By default server address is set to a "127.0.0.1:8080"
    pub fn server_address<S: net::ToSocketAddrs>(mut self, addr: S) -> Self {
        match addr.to_socket_addrs() {
            Err(err) => error!("Can not convert to SocketAddr: {}", err),
            Ok(mut addrs) => if let Some(addr) = addrs.next() {
                self.addr = addr;
            },
        }
        self
    }

    /// Finish service configuration and create `ServiceConfig` object.
    pub fn finish(self) -> ServiceConfig<H> {
        let settings = ServerSettings::new(self.addr, &self.host, self.secure);
        let client_shutdown = if self.secure { self.client_shutdown } else { 0 };

        ServiceConfig::new(
            self.handler,
            self.keep_alive,
            self.client_timeout,
            client_shutdown,
            settings,
        )
    }
}

struct Date {
    current: Instant,
    bytes: [u8; DATE_VALUE_LENGTH],
    pos: usize,
}

impl Date {
    fn new() -> Date {
        let mut date = Date {
            current: Instant::now(),
            bytes: [0; DATE_VALUE_LENGTH],
            pos: 0,
        };
        date.update();
        date
    }
    fn update(&mut self) {
        self.pos = 0;
        self.current = Instant::now();
        write!(self, "{}", time::at_utc(time::get_time()).rfc822()).unwrap();
    }
}

impl fmt::Write for Date {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        let len = s.len();
        self.bytes[self.pos..self.pos + len].copy_from_slice(s.as_bytes());
        self.pos += len;
        Ok(())
    }
}

#[derive(Debug)]
pub(crate) struct SharedBytesPool(RefCell<VecDeque<BytesMut>>);

impl SharedBytesPool {
    pub fn new() -> SharedBytesPool {
        SharedBytesPool(RefCell::new(VecDeque::with_capacity(128)))
    }

    pub fn get_bytes(&self) -> BytesMut {
        if let Some(bytes) = self.0.borrow_mut().pop_front() {
            bytes
        } else {
            BytesMut::new()
        }
    }

    pub fn release_bytes(&self, mut bytes: BytesMut) {
        let v = &mut self.0.borrow_mut();
        if v.len() < 128 {
            bytes.clear();
            v.push_front(bytes);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures::future;
    use tokio::runtime::current_thread;

    #[test]
    fn test_date_len() {
        assert_eq!(DATE_VALUE_LENGTH, "Sun, 06 Nov 1994 08:49:37 GMT".len());
    }

    #[test]
    fn test_date() {
        let mut rt = current_thread::Runtime::new().unwrap();

        let _ = rt.block_on(future::lazy(|| {
            let settings = ServiceConfig::<()>::new(
                (),
                KeepAlive::Os,
                0,
                0,
                ServerSettings::default(),
            );
            let mut buf1 = BytesMut::with_capacity(DATE_VALUE_LENGTH + 10);
            settings.set_date(&mut buf1, true);
            let mut buf2 = BytesMut::with_capacity(DATE_VALUE_LENGTH + 10);
            settings.set_date(&mut buf2, true);
            assert_eq!(buf1, buf2);
            future::ok::<_, ()>(())
        }));
    }
}
<commit_msg>refactor date rendering<commit_after>use std::cell::{Cell, RefCell, RefMut};
use std::collections::VecDeque;
use std::fmt::Write;
use std::rc::Rc;
use std::time::{Duration, Instant};
use std::{env, fmt, net};

use bytes::BytesMut;
use futures::{future, Future};
use futures_cpupool::CpuPool;
use http::StatusCode;
use lazycell::LazyCell;
use parking_lot::Mutex;
use time;
use tokio_current_thread::spawn;
use tokio_timer::{sleep, Delay};

use super::channel::Node;
use super::message::{Request, RequestPool};
use super::KeepAlive;
use body::Body;
use httpresponse::{HttpResponse, HttpResponseBuilder, HttpResponsePool};

/// Env variable for default cpu pool size
const ENV_CPU_POOL_VAR: &str = "ACTIX_CPU_POOL";

lazy_static! {
    pub(crate) static ref DEFAULT_CPUPOOL: Mutex<CpuPool> = {
        let default = match env::var(ENV_CPU_POOL_VAR) {
            Ok(val) => {
                if let Ok(val) = val.parse() {
                    val
                } else {
                    error!("Can not parse ACTIX_CPU_POOL value");
                    20
                }
            }
            Err(_) => 20,
        };
        Mutex::new(CpuPool::new(default))
    };
}

/// Various server settings
pub struct ServerSettings {
    addr: net::SocketAddr,
    secure: bool,
    host: String,
    cpu_pool: LazyCell<CpuPool>,
    responses: &'static HttpResponsePool,
}

impl Clone for ServerSettings {
    fn clone(&self) -> Self {
        ServerSettings {
            addr: self.addr,
            secure: self.secure,
            host: self.host.clone(),
            cpu_pool: LazyCell::new(),
            responses: HttpResponsePool::get_pool(),
        }
    }
}

impl Default for ServerSettings {
    fn default() -> Self {
        ServerSettings {
            addr: "127.0.0.1:8080".parse().unwrap(),
            secure: false,
            host: "localhost:8080".to_owned(),
            responses: HttpResponsePool::get_pool(),
            cpu_pool: LazyCell::new(),
        }
    }
}

impl ServerSettings {
    /// Crate server settings instance
    pub(crate) fn new(
        addr: net::SocketAddr, host: &str, secure: bool,
    ) -> ServerSettings {
        let host = host.to_owned();
        let cpu_pool = LazyCell::new();
        let responses = HttpResponsePool::get_pool();
        ServerSettings {
            addr,
            secure,
            host,
            cpu_pool,
            responses,
        }
    }

    /// Returns the socket address of the local half of this TCP connection
    pub fn local_addr(&self) -> net::SocketAddr {
        self.addr
    }

    /// Returns true if connection is secure(https)
    pub fn secure(&self) -> bool {
        self.secure
    }

    /// Returns host header value
    pub fn host(&self) -> &str {
        &self.host
    }

    /// Returns default `CpuPool` for server
    pub fn cpu_pool(&self) -> &CpuPool {
        self.cpu_pool.borrow_with(|| DEFAULT_CPUPOOL.lock().clone())
    }

    #[inline]
    pub(crate) fn get_response(&self, status: StatusCode, body: Body) -> HttpResponse {
        HttpResponsePool::get_response(&self.responses, status, body)
    }

    #[inline]
    pub(crate) fn get_response_builder(
        &self, status: StatusCode,
    ) -> HttpResponseBuilder {
        HttpResponsePool::get_builder(&self.responses, status)
    }
}

// "Sun, 06 Nov 1994 08:49:37 GMT".len()
const DATE_VALUE_LENGTH: usize = 29;

/// Http service configuration
pub struct ServiceConfig<H>(Rc<Inner<H>>);

struct Inner<H> {
    handler: H,
    keep_alive: Option<Duration>,
    client_timeout: u64,
    client_shutdown: u64,
    ka_enabled: bool,
    bytes: Rc<SharedBytesPool>,
    messages: &'static RequestPool,
    node: RefCell<Node<()>>,
    date: Cell<Option<Date>>,
}

impl<H> Clone for ServiceConfig<H> {
    fn clone(&self) -> Self {
        ServiceConfig(self.0.clone())
    }
}

impl<H> ServiceConfig<H> {
    /// Create instance of `ServiceConfig`
    pub(crate) fn new(
        handler: H, keep_alive: KeepAlive, client_timeout: u64, client_shutdown: u64,
        settings: ServerSettings,
    ) -> ServiceConfig<H> {
        let (keep_alive, ka_enabled) = match keep_alive {
            KeepAlive::Timeout(val) => (val as u64, true),
            KeepAlive::Os | KeepAlive::Tcp(_) => (0, true),
            KeepAlive::Disabled => (0, false),
        };
        let keep_alive = if ka_enabled && keep_alive > 0 {
            Some(Duration::from_secs(keep_alive))
        } else {
            None
        };

        ServiceConfig(Rc::new(Inner {
            handler,
            keep_alive,
            ka_enabled,
            client_timeout,
            client_shutdown,
            bytes: Rc::new(SharedBytesPool::new()),
            messages: RequestPool::pool(settings),
            node: RefCell::new(Node::head()),
            date: Cell::new(None),
        }))
    }

    /// Create worker settings builder.
    pub fn build(handler: H) -> ServiceConfigBuilder<H> {
        ServiceConfigBuilder::new(handler)
    }

    pub(crate) fn head(&self) -> RefMut<Node<()>> {
        self.0.node.borrow_mut()
    }

    pub(crate) fn handler(&self) -> &H {
        &self.0.handler
    }

    #[inline]
    /// Keep alive duration if configured.
    pub fn keep_alive(&self) -> Option<Duration> {
        self.0.keep_alive
    }

    #[inline]
    /// Return state of connection keep-alive funcitonality
    pub fn keep_alive_enabled(&self) -> bool {
        self.0.ka_enabled
    }

    pub(crate) fn get_bytes(&self) -> BytesMut {
        self.0.bytes.get_bytes()
    }

    pub(crate) fn release_bytes(&self, bytes: BytesMut) {
        self.0.bytes.release_bytes(bytes)
    }

    pub(crate) fn get_request(&self) -> Request {
        RequestPool::get(self.0.messages)
    }
}

impl<H: 'static> ServiceConfig<H> {
    #[inline]
    /// Client timeout for first request.
    pub fn client_timer(&self) -> Option<Delay> {
        let delay = self.0.client_timeout;
        if delay != 0 {
            Some(Delay::new(self.now() + Duration::from_millis(delay)))
        } else {
            None
        }
    }

    /// Client timeout for first request.
    pub fn client_timer_expire(&self) -> Option<Instant> {
        let delay = self.0.client_timeout;
        if delay != 0 {
            Some(self.now() + Duration::from_millis(delay))
        } else {
            None
        }
    }

    /// Client shutdown timer
    pub fn client_shutdown_timer(&self) -> Option<Instant> {
        let delay = self.0.client_shutdown;
        if delay != 0 {
            Some(self.now() + Duration::from_millis(delay))
        } else {
            None
        }
    }

    #[inline]
    /// Return keep-alive timer delay is configured.
    pub fn keep_alive_timer(&self) -> Option<Delay> {
        if let Some(ka) = self.0.keep_alive {
            Some(Delay::new(self.now() + ka))
        } else {
            None
        }
    }

    /// Keep-alive expire time
    pub fn keep_alive_expire(&self) -> Option<Instant> {
        if let Some(ka) = self.0.keep_alive {
            Some(self.now() + ka)
        } else {
            None
        }
    }

    fn check_date(&self) {
        if unsafe { &*self.0.date.as_ptr() }.is_none() {
            self.0.date.set(Some(Date::new()));

            // periodic date update
            let s = self.clone();
            spawn(sleep(Duration::from_millis(500)).then(move |_| {
                s.0.date.set(None);
                future::ok(())
            }));
        }
    }

    pub(crate) fn set_date(&self, dst: &mut BytesMut, full: bool) {
        self.check_date();

        let date = &unsafe { &*self.0.date.as_ptr() }.as_ref().unwrap().bytes;

        if full {
            let mut buf: [u8; 39] = [0; 39];
            buf[..6].copy_from_slice(b"date: ");
            buf[6..35].copy_from_slice(date);
            buf[35..].copy_from_slice(b"\r\n\r\n");
            dst.extend_from_slice(&buf);
        } else {
            dst.extend_from_slice(date);
        }
    }

    #[inline]
    pub(crate) fn now(&self) -> Instant {
        self.check_date();
        unsafe { &*self.0.date.as_ptr() }.as_ref().unwrap().current
    }
}

/// A service config builder
///
/// This type can be used to construct an instance of `ServiceConfig` through a
/// builder-like pattern.
pub struct ServiceConfigBuilder<H> {
    handler: H,
    keep_alive: KeepAlive,
    client_timeout: u64,
    client_shutdown: u64,
    host: String,
    addr: net::SocketAddr,
    secure: bool,
}

impl<H> ServiceConfigBuilder<H> {
    /// Create instance of `ServiceConfigBuilder`
    pub fn new(handler: H) -> ServiceConfigBuilder<H> {
        ServiceConfigBuilder {
            handler,
            keep_alive: KeepAlive::Timeout(5),
            client_timeout: 5000,
            client_shutdown: 5000,
            secure: false,
            host: "localhost".to_owned(),
            addr: "127.0.0.1:8080".parse().unwrap(),
        }
    }

    /// Enable secure flag for current server.
    ///
    /// By default this flag is set to false.
    pub fn secure(mut self) -> Self {
        self.secure = true;
        self
    }

    /// Set server keep-alive setting.
    ///
    /// By default keep alive is set to a 5 seconds.
    pub fn keep_alive<T: Into<KeepAlive>>(mut self, val: T) -> Self {
        self.keep_alive = val.into();
        self
    }

    /// Set server client timeout in milliseconds for first request.
    ///
    /// Defines a timeout for reading client request header. If a client does not transmit
    /// the entire set headers within this time, the request is terminated with
    /// the 408 (Request Time-out) error.
    ///
    /// To disable timeout set value to 0.
    ///
    /// By default client timeout is set to 5000 milliseconds.
    pub fn client_timeout(mut self, val: u64) -> Self {
        self.client_timeout = val;
        self
    }

    /// Set server connection shutdown timeout in milliseconds.
    ///
    /// Defines a timeout for shutdown connection. If a shutdown procedure does not complete
    /// within this time, the request is dropped. This timeout affects only secure connections.
    ///
    /// To disable timeout set value to 0.
    ///
    /// By default client timeout is set to 5000 milliseconds.
    pub fn client_shutdown(mut self, val: u64) -> Self {
        self.client_shutdown = val;
        self
    }

    /// Set server host name.
    ///
    /// Host name is used by application router aa a hostname for url
    /// generation. Check [ConnectionInfo](./dev/struct.ConnectionInfo.
    /// html#method.host) documentation for more information.
    ///
    /// By default host name is set to a "localhost" value.
    pub fn server_hostname(mut self, val: &str) -> Self {
        self.host = val.to_owned();
        self
    }

    /// Set server ip address.
    ///
    /// Host name is used by application router aa a hostname for url
    /// generation. Check [ConnectionInfo](./dev/struct.ConnectionInfo.
    /// html#method.host) documentation for more information.
    ///
    /// By default server address is set to a "127.0.0.1:8080"
    pub fn server_address<S: net::ToSocketAddrs>(mut self, addr: S) -> Self {
        match addr.to_socket_addrs() {
            Err(err) => error!("Can not convert to SocketAddr: {}", err),
            Ok(mut addrs) => if let Some(addr) = addrs.next() {
                self.addr = addr;
            },
        }
        self
    }

    /// Finish service configuration and create `ServiceConfig` object.
    pub fn finish(self) -> ServiceConfig<H> {
        let settings = ServerSettings::new(self.addr, &self.host, self.secure);
        let client_shutdown = if self.secure { self.client_shutdown } else { 0 };

        ServiceConfig::new(
            self.handler,
            self.keep_alive,
            self.client_timeout,
            client_shutdown,
            settings,
        )
    }
}

#[derive(Copy, Clone)]
struct Date {
    current: Instant,
    bytes: [u8; DATE_VALUE_LENGTH],
    pos: usize,
}

impl Date {
    fn new() -> Date {
        let mut date = Date {
            current: Instant::now(),
            bytes: [0; DATE_VALUE_LENGTH],
            pos: 0,
        };
        date.update();
        date
    }
    fn update(&mut self) {
        self.pos = 0;
        self.current = Instant::now();
        write!(self, "{}", time::at_utc(time::get_time()).rfc822()).unwrap();
    }
}

impl fmt::Write for Date {
    fn write_str(&mut self, s: &str) -> fmt::Result {
        let len = s.len();
        self.bytes[self.pos..self.pos + len].copy_from_slice(s.as_bytes());
        self.pos += len;
        Ok(())
    }
}

#[derive(Debug)]
pub(crate) struct SharedBytesPool(RefCell<VecDeque<BytesMut>>);

impl SharedBytesPool {
    pub fn new() -> SharedBytesPool {
        SharedBytesPool(RefCell::new(VecDeque::with_capacity(128)))
    }

    pub fn get_bytes(&self) -> BytesMut {
        if let Some(bytes) = self.0.borrow_mut().pop_front() {
            bytes
        } else {
            BytesMut::new()
        }
    }

    pub fn release_bytes(&self, mut bytes: BytesMut) {
        let v = &mut self.0.borrow_mut();
        if v.len() < 128 {
            bytes.clear();
            v.push_front(bytes);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures::future;
    use tokio::runtime::current_thread;

    #[test]
    fn test_date_len() {
        assert_eq!(DATE_VALUE_LENGTH, "Sun, 06 Nov 1994 08:49:37 GMT".len());
    }

    #[test]
    fn test_date() {
        let mut rt = current_thread::Runtime::new().unwrap();

        let _ = rt.block_on(future::lazy(|| {
            let settings = ServiceConfig::<()>::new(
                (),
                KeepAlive::Os,
                0,
                0,
                ServerSettings::default(),
            );
            let mut buf1 = BytesMut::with_capacity(DATE_VALUE_LENGTH + 10);
            settings.set_date(&mut buf1, true);
            let mut buf2 = BytesMut::with_capacity(DATE_VALUE_LENGTH + 10);
            settings.set_date(&mut buf2, true);
            assert_eq!(buf1, buf2);
            future::ok::<_, ()>(())
        }));
    }
}
<|endoftext|>


<commit_before>//! To see how color numbers map to actual colors in your terminal run
//! `cargo run --example colors`. Use tab to swap fg/bg colors.

use yaml_rust::Yaml;
use yaml_rust::YamlLoader;

use std::env::home_dir;
use std::fs::File;
use std::io::Read;
use std::path::PathBuf;

#[derive(Debug, Clone)]
pub struct Server {
    /// Address of the server
    pub server_addr: String,

    /// Port of the server
    pub server_port: u16,

    /// Hostname to be used in connection registration
    pub hostname: String,

    /// Real name to be used in connection registration
    pub real_name: String,

    /// Nicks to try when connecting to this server. tiny tries these sequentially, and starts
    /// adding trailing underscores to the last one if none of the nicks are available.
    pub nicks: Vec<String>,

    /// Commands to automatically run after joining to the server. Useful for e.g. registering the
    /// nick via nickserv or joining channels. Uses tiny command syntax.
    pub auto_cmds: Vec<String>,
}

/// Similar to `Server`, but used when connecting via the `/connect` command.
#[derive(Debug, Clone)]
pub struct Defaults {
    pub nicks: Vec<String>,
    pub hostname: String,
    pub realname: String,
    pub auto_cmds: Vec<String>
}

pub fn get_defaults() -> Defaults {
    Defaults {
        nicks: vec!["tiny_user".to_owned()],
        hostname: "tiny".to_owned(),
        realname: "anonymous tiny user".to_owned(),
        auto_cmds: vec![],
    }
}

fn get_config_path() -> PathBuf {
    let mut config_path = home_dir().unwrap();
    config_path.push(".tinyrc.yml");
    config_path
}

pub fn read_config() -> Option<(Vec<Server>, Defaults)> {

    // sigh ... what a mess

    let config_str = {
        let config_path = get_config_path();
        if !config_path.exists() { return None; }
        let mut config_file = File::open(get_config_path()).unwrap();
        let mut config_str = String::new();
        config_file.read_to_string(&mut config_str).unwrap();
        config_str
    };

    let config_yaml = YamlLoader::load_from_str(&config_str).unwrap();

    let servers_yaml: &Yaml =
        config_yaml[0].as_hash().unwrap()
            .get(&Yaml::String("servers".to_owned())).unwrap();
                // duh, allocation for lookup

    let defaults_yaml: &Yaml =
        config_yaml[0].as_hash().unwrap()
            .get(&Yaml::String("defaults".to_owned())).unwrap();

    let mut servers = vec![];

    let yaml_addr_key = Yaml::String("addr".to_owned());
    let yaml_port_key = Yaml::String("port".to_owned());
    let yaml_hostname_key = Yaml::String("hostname".to_owned());
    let yaml_realname_key = Yaml::String("realname".to_owned());
    let yaml_nicks_key = Yaml::String("nicks".to_owned());
    let yaml_auto_cmds_key = Yaml::String("auto_cmds".to_owned());

    for server in servers_yaml.as_vec().unwrap().into_iter() {
        let server_hash = server.as_hash().unwrap();
        let server_addr =
            server_hash.get(&yaml_addr_key).unwrap().as_str().unwrap().to_owned();
        let server_port =
            server_hash.get(&yaml_port_key).unwrap().as_i64().unwrap().to_owned() as u16;
        let hostname =
            server_hash.get(&yaml_hostname_key).unwrap().as_str().unwrap().to_owned();
        let real_name =
            server_hash.get(&yaml_realname_key).unwrap().as_str().unwrap().to_owned();
        let nicks: Vec<String> =
            server_hash.get(&yaml_nicks_key).unwrap().as_vec().unwrap()
                .into_iter().map(|s| s.as_str().unwrap().to_owned()).collect();
        let auto_cmds: Vec<String> =
            server_hash.get(&yaml_auto_cmds_key).unwrap().as_vec().unwrap()
                .into_iter().map(|s| s.as_str().unwrap().to_owned()).collect();

        servers.push(Server {
            server_addr, server_port, hostname, real_name, nicks, auto_cmds
        });
    }

    let defaults = defaults_yaml.as_hash().unwrap();
    let defaults = Defaults {
        nicks: defaults.get(&yaml_nicks_key).unwrap().as_vec().unwrap()
                   .into_iter().map(|s| s.as_str().unwrap().to_owned()).collect(),
        hostname: defaults.get(&yaml_hostname_key).unwrap().as_str().unwrap().to_owned(),
        realname: defaults.get(&yaml_realname_key).unwrap().as_str().unwrap().to_owned(),
        auto_cmds: defaults.get(&yaml_auto_cmds_key).unwrap().as_vec().unwrap()
                       .into_iter().map(|s| s.as_str().unwrap().to_owned()).collect(),
    };

    Some((servers, defaults))
}

////////////////////////////////////////////////////////////////////////////////////////////////////
// Colors
////////////////////////////////////////////////////////////////////////////////////////////////////

/// Colors used to highlight nicks
pub static NICK_COLORS: [u8; 15] =
    [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15 ];

pub use termbox_simple::*;

#[derive(Debug, Clone, Copy)]
pub struct Style {
    /// Termbox fg.
    pub fg: u16,

    /// Termbox bg.
    pub bg: u16,
}

pub const CLEAR: Style =
    Style {
        fg: TB_DEFAULT,
        bg: TB_DEFAULT,
    };

pub const USER_MSG: Style =
    Style {
        fg: 15,
        bg: TB_DEFAULT,
    };

pub const ERR_MSG: Style =
    Style {
        fg: 15 | TB_BOLD,
        bg: 1,
    };

pub const TOPIC: Style =
    Style {
        fg: 14 | TB_BOLD,
        bg: TB_DEFAULT,
    };

pub const CURSOR: Style =
    USER_MSG;

pub const JOIN: Style =
    Style {
        fg: 242,
        bg: TB_DEFAULT,
    };

pub const PART: Style =
    Style {
        fg: 242,
        bg: TB_DEFAULT,
    };

pub const NICK: Style =
    Style {
        fg: 242,
        bg: TB_DEFAULT,
    };

pub const EXIT_DIALOGUE: Style =
    Style {
        fg: TB_DEFAULT,
        bg: 4,
    };

pub const HIGHLIGHT: Style =
    Style {
        fg: 9 | TB_BOLD,
        bg: TB_DEFAULT,
    };

// Currently unused
// pub const MENTION: Style =
//     Style {
//         fg: 220,
//         bg: TB_DEFAULT,
//     };

pub const COMPLETION: Style =
    Style {
        fg: 84,
        bg: TB_DEFAULT,
    };

pub const TIMESTAMP: Style =
    Style {
        fg: 15 | TB_BOLD,
        bg: TB_DEFAULT,
    };

pub const TAB_ACTIVE: Style =
    Style {
        fg: 15 | TB_BOLD,
        bg: 0,
    };

pub const TAB_NORMAL: Style =
    Style {
        fg: 8,
        bg: 0,
    };

pub const TAB_IMPORTANT: Style =
    Style {
        fg: 9 | TB_BOLD,
        bg: 0,
    };

pub const TAB_HIGHLIGHT: Style =
    Style {
        fg: 5,
        bg: 0,
    };
<commit_msg>Tweak colors<commit_after>//! To see how color numbers map to actual colors in your terminal run
//! `cargo run --example colors`. Use tab to swap fg/bg colors.

use yaml_rust::Yaml;
use yaml_rust::YamlLoader;

use std::env::home_dir;
use std::fs::File;
use std::io::Read;
use std::path::PathBuf;

#[derive(Debug, Clone)]
pub struct Server {
    /// Address of the server
    pub server_addr: String,

    /// Port of the server
    pub server_port: u16,

    /// Hostname to be used in connection registration
    pub hostname: String,

    /// Real name to be used in connection registration
    pub real_name: String,

    /// Nicks to try when connecting to this server. tiny tries these sequentially, and starts
    /// adding trailing underscores to the last one if none of the nicks are available.
    pub nicks: Vec<String>,

    /// Commands to automatically run after joining to the server. Useful for e.g. registering the
    /// nick via nickserv or joining channels. Uses tiny command syntax.
    pub auto_cmds: Vec<String>,
}

/// Similar to `Server`, but used when connecting via the `/connect` command.
#[derive(Debug, Clone)]
pub struct Defaults {
    pub nicks: Vec<String>,
    pub hostname: String,
    pub realname: String,
    pub auto_cmds: Vec<String>
}

pub fn get_defaults() -> Defaults {
    Defaults {
        nicks: vec!["tiny_user".to_owned()],
        hostname: "tiny".to_owned(),
        realname: "anonymous tiny user".to_owned(),
        auto_cmds: vec![],
    }
}

fn get_config_path() -> PathBuf {
    let mut config_path = home_dir().unwrap();
    config_path.push(".tinyrc.yml");
    config_path
}

pub fn read_config() -> Option<(Vec<Server>, Defaults)> {

    // sigh ... what a mess

    let config_str = {
        let config_path = get_config_path();
        if !config_path.exists() { return None; }
        let mut config_file = File::open(get_config_path()).unwrap();
        let mut config_str = String::new();
        config_file.read_to_string(&mut config_str).unwrap();
        config_str
    };

    let config_yaml = YamlLoader::load_from_str(&config_str).unwrap();

    let servers_yaml: &Yaml =
        config_yaml[0].as_hash().unwrap()
            .get(&Yaml::String("servers".to_owned())).unwrap();
                // duh, allocation for lookup

    let defaults_yaml: &Yaml =
        config_yaml[0].as_hash().unwrap()
            .get(&Yaml::String("defaults".to_owned())).unwrap();

    let mut servers = vec![];

    let yaml_addr_key = Yaml::String("addr".to_owned());
    let yaml_port_key = Yaml::String("port".to_owned());
    let yaml_hostname_key = Yaml::String("hostname".to_owned());
    let yaml_realname_key = Yaml::String("realname".to_owned());
    let yaml_nicks_key = Yaml::String("nicks".to_owned());
    let yaml_auto_cmds_key = Yaml::String("auto_cmds".to_owned());

    for server in servers_yaml.as_vec().unwrap().into_iter() {
        let server_hash = server.as_hash().unwrap();
        let server_addr =
            server_hash.get(&yaml_addr_key).unwrap().as_str().unwrap().to_owned();
        let server_port =
            server_hash.get(&yaml_port_key).unwrap().as_i64().unwrap().to_owned() as u16;
        let hostname =
            server_hash.get(&yaml_hostname_key).unwrap().as_str().unwrap().to_owned();
        let real_name =
            server_hash.get(&yaml_realname_key).unwrap().as_str().unwrap().to_owned();
        let nicks: Vec<String> =
            server_hash.get(&yaml_nicks_key).unwrap().as_vec().unwrap()
                .into_iter().map(|s| s.as_str().unwrap().to_owned()).collect();
        let auto_cmds: Vec<String> =
            server_hash.get(&yaml_auto_cmds_key).unwrap().as_vec().unwrap()
                .into_iter().map(|s| s.as_str().unwrap().to_owned()).collect();

        servers.push(Server {
            server_addr, server_port, hostname, real_name, nicks, auto_cmds
        });
    }

    let defaults = defaults_yaml.as_hash().unwrap();
    let defaults = Defaults {
        nicks: defaults.get(&yaml_nicks_key).unwrap().as_vec().unwrap()
                   .into_iter().map(|s| s.as_str().unwrap().to_owned()).collect(),
        hostname: defaults.get(&yaml_hostname_key).unwrap().as_str().unwrap().to_owned(),
        realname: defaults.get(&yaml_realname_key).unwrap().as_str().unwrap().to_owned(),
        auto_cmds: defaults.get(&yaml_auto_cmds_key).unwrap().as_vec().unwrap()
                       .into_iter().map(|s| s.as_str().unwrap().to_owned()).collect(),
    };

    Some((servers, defaults))
}

////////////////////////////////////////////////////////////////////////////////////////////////////
// Colors
////////////////////////////////////////////////////////////////////////////////////////////////////

/// Colors used to highlight nicks
pub static NICK_COLORS: [u8; 13] =
    [ 1, 2, 3, 4, 5, 6, 7, 9, 10, 11, 12, 13, 14 ];

pub use termbox_simple::*;

#[derive(Debug, Clone, Copy)]
pub struct Style {
    /// Termbox fg.
    pub fg: u16,

    /// Termbox bg.
    pub bg: u16,
}

pub const CLEAR: Style =
    Style {
        fg: TB_DEFAULT,
        bg: TB_DEFAULT,
    };

pub const USER_MSG: Style =
    Style {
        fg: 0,
        bg: TB_DEFAULT,
    };

pub const ERR_MSG: Style =
    Style {
        fg: 0 | TB_BOLD,
        bg: 1,
    };

pub const TOPIC: Style =
    Style {
        fg: 14 | TB_BOLD,
        bg: TB_DEFAULT,
    };

pub const CURSOR: Style =
    USER_MSG;

pub const JOIN: Style =
    Style {
        fg: 242,
        bg: TB_DEFAULT,
    };

pub const PART: Style =
    Style {
        fg: 242,
        bg: TB_DEFAULT,
    };

pub const NICK: Style =
    Style {
        fg: 242,
        bg: TB_DEFAULT,
    };

pub const EXIT_DIALOGUE: Style =
    Style {
        fg: TB_DEFAULT,
        bg: 4,
    };

pub const HIGHLIGHT: Style =
    Style {
        fg: 9 | TB_BOLD,
        bg: TB_DEFAULT,
    };

// Currently unused
// pub const MENTION: Style =
//     Style {
//         fg: 220,
//         bg: TB_DEFAULT,
//     };

pub const COMPLETION: Style =
    Style {
        fg: 84,
        bg: TB_DEFAULT,
    };

pub const TIMESTAMP: Style =
    Style {
        fg: 0 | TB_BOLD,
        bg: TB_DEFAULT,
    };

pub const TAB_ACTIVE: Style =
    Style {
        fg: 0 | TB_BOLD,
        bg: 0,
    };

pub const TAB_NORMAL: Style =
    Style {
        fg: 8,
        bg: 0,
    };

pub const TAB_IMPORTANT: Style =
    Style {
        fg: 9 | TB_BOLD,
        bg: 0,
    };

pub const TAB_HIGHLIGHT: Style =
    Style {
        fg: 5,
        bg: 0,
    };
<|endoftext|>