<commit_before>use prolog::ast::*;
use prolog::compile::*;
use prolog::heap_print::*;
use prolog::tabled_rc::*;

mod machine_errors;
pub(super) mod machine_state;
#[macro_use]
mod machine_state_impl;
mod system_calls;

use prolog::machine::machine_state::*;

use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::mem::swap;
use std::ops::Index;
use std::rc::Rc;

pub struct MachineCodeIndex<'a> {
    pub code_dir: &'a mut CodeDir,
    pub op_dir: &'a mut OpDir,
//  pub modules: &'a ModuleDir // &'a HashMap<ClauseName, Module>
}

pub struct Machine {
    ms: MachineState,
    call_policy: Box<CallPolicy>,
    cut_policy: Box<CutPolicy>,
    code: Code,
    pub(super) code_dir: CodeDir,
    pub(super) op_dir: OpDir,
    term_dir: TermDir,
    modules: ModuleDir, // HashMap<ClauseName, Module>,
    cached_query: Option
}

impl Index<LocalCodePtr> for Machine {
    type Output = Line;

    fn index(&self, ptr: LocalCodePtr) -> &Self::Output {
        match ptr {
            LocalCodePtr::TopLevel(_, p) => {
                match &self.cached_query {
                    &Some(ref cq) => &cq[p],
                    &None => panic!("Out-of-bounds top level index.")
                }
            },
            LocalCodePtr::DirEntry(p, _) => &self.code[p]
        }
    }
}

impl<'a> SubModuleUser for MachineCodeIndex<'a> {
    fn op_dir(&mut self) -> &mut OpDir {
        self.op_dir
    }

    fn insert_dir_entry(&mut self, name: ClauseName, arity: usize, idx: ModuleCodeIndex) {
        if let Some(ref mut code_idx) = self.code_dir.get_mut(&(name.clone(), arity)) {
            if !code_idx.is_undefined() {
                println!("warning: overwriting {}/{}", &name, arity);
            }

            set_code_index!(code_idx, idx.0, idx.1);

            return;
        }

        self.code_dir.insert((name, arity), CodeIndex::from(idx));
    }
}

static LISTS: &str   = include_str!("../lib/lists.pl");
static CONTROL: &str = include_str!("../lib/control.pl");
static QUEUES: &str  = include_str!("../lib/queues.pl");

impl Machine {
    pub fn new() -> Self {
        let mut wam = Machine {
            ms: MachineState::new(Rc::new(RefCell::new(HashSet::new()))),
            call_policy: Box::new(DefaultCallPolicy {}),
            cut_policy: Box::new(DefaultCutPolicy {}),
            code: Code::new(),
            code_dir: CodeDir::new(),
            term_dir: TermDir::new(),
            op_dir: default_op_dir(),
            modules: HashMap::new(),
            cached_query: None
        };

        compile_listing(&mut wam, BUILTINS);
        wam.use_module_in_toplevel(clause_name!("builtins"));

        compile_listing(&mut wam, LISTS);
        compile_listing(&mut wam, CONTROL);
        compile_listing(&mut wam, QUEUES);

        wam
    }

    fn remove_module(&mut self, module_name: ClauseName) {
        let iter = if let Some(submodule) = self.modules.get(&module_name) {
            submodule.module_decl.exports.iter().cloned()
        } else {
            return;
        };

        for (name, arity) in iter {
            let name = name.defrock_brackets();

            match self.code_dir.get(&(name.clone(), arity)).cloned() {
                Some(CodeIndex (ref code_idx)) => {
                    if &code_idx.borrow().1 != &module_name {
                        continue;
                    }

                    self.code_dir.remove(&(name.clone(), arity));

                    // remove or respecify ops.
                    if arity == 2 {
                        if let Some((_, _, mod_name)) = self.op_dir.get(&(name.clone(), Fixity::In)).cloned()
                        {
                            if mod_name == module_name {
                                self.op_dir.remove(&(name.clone(), Fixity::In));
                            }
                        }
                    } else if arity == 1 {
                        if let Some((_, _, mod_name)) = self.op_dir.get(&(name.clone(), Fixity::Pre)).cloned()
                        {
                            if mod_name == module_name {
                                self.op_dir.remove(&(name.clone(), Fixity::Pre));
                            }
                        }

                        if let Some((_, _, mod_name)) = self.op_dir.get(&(name.clone(), Fixity::Post)).cloned()
                        {
                            if mod_name == module_name {
                                self.op_dir.remove(&(name.clone(), Fixity::Post));
                            }
                        }
                    }
                },
                _ => {}
            };
        }
    }

    pub fn failed(&self) -> bool {
        self.ms.fail
    }

    pub fn atom_tbl(&self) -> TabledData<Atom> {
        self.ms.atom_tbl.clone()
    }

    pub fn use_qualified_module_in_toplevel(&mut self, name: ClauseName, exports: Vec<PredicateKey>)
                                            -> EvalSession
    {
        self.remove_module(name.clone());

        match self.modules.get_mut(&name) {
            Some(ref mut module) => {
                let mut indices = MachineCodeIndex { code_dir: &mut self.code_dir,
                                                     op_dir: &mut self.op_dir };

                indices.use_qualified_module(module, &exports)
            },
            None => EvalSession::from(SessionError::ModuleNotFound)
        }
    }

    pub fn use_module_in_toplevel(&mut self, name: ClauseName) -> EvalSession {
        self.remove_module(name.clone());

        match self.modules.get_mut(&name) {
            Some(ref mut module) => {
                let mut indices = MachineCodeIndex { code_dir: &mut self.code_dir,
                                                     op_dir: &mut self.op_dir };

                indices.use_module(module)
            },
            None => EvalSession::from(SessionError::ModuleNotFound)
        }
    }

    pub fn get_module(&self, name: ClauseName) -> Option<&Module> {
        self.modules.get(&name)
    }

    pub fn add_batched_code(&mut self, mut code: Code, code_dir: CodeDir) {
        self.code.append(&mut code);
        self.code_dir.extend(code_dir.into_iter());
    }

    pub fn add_batched_ops(&mut self, op_dir: OpDir) {
        self.op_dir.extend(op_dir.into_iter());
    }

    pub fn add_module(&mut self, module: Module, code: Code) {
        self.modules.insert(module.module_decl.name.clone(), module);
        self.code.extend(code.into_iter());
    }

    pub fn add_user_code(&mut self, name: ClauseName, arity: usize, code: Code, pred: Predicate)
                         -> EvalSession
    {
        match self.code_dir.get(&(name.clone(), arity)) {
            Some(&CodeIndex (ref idx)) if idx.borrow().1 != clause_name!("user") =>
                    return EvalSession::from(SessionError::ImpermissibleEntry(format!("{}/{}",
                                                                                name,
                                                                                arity))),
            _ => {}
        };

        let offset = self.code.len();

        self.code.extend(code.into_iter());
        self.term_dir.insert((name.clone(), arity), pred);

        let idx = self.code_dir.entry((name, arity))
            .or_insert(CodeIndex::from((offset, clause_name!("user"))));

        set_code_index!(idx, IndexPtr::Index(offset), clause_name!("user"));
        EvalSession::EntrySuccess
    }

    pub fn code_size(&self) -> usize {
        self.code.len()
    }

    fn cached_query_size(&self) -> usize {
        match &self.cached_query {
            &Some(ref query) => query.len(),
            _ => 0
        }
    }

    fn lookup_instr(&self, p: CodePtr) -> Option<Line> {
        match p {
            CodePtr::Local(LocalCodePtr::TopLevel(_, p)) =>
                match &self.cached_query {
                    &Some(ref cq) => Some(cq[p].clone()),
                    &None => None
                },
            CodePtr::Local(LocalCodePtr::DirEntry(p, _)) =>
                Some(self.code[p].clone()),
            CodePtr::BuiltInClause(built_in, _) =>
                Some(call_clause!(ClauseType::BuiltIn(built_in.clone()), built_in.arity(),
                                  0, self.ms.last_call)),
            CodePtr::CallN(arity, _) =>
                Some(call_clause!(ClauseType::CallN, arity, 0, self.ms.last_call))
        }
    }

    fn execute_instr(&mut self)
    {
        let instr = match self.lookup_instr(self.ms.p.clone()) {
            Some(instr) => instr,
            None => return
        };

        match instr {
            Line::Arithmetic(ref arith_instr) =>
                self.ms.execute_arith_instr(arith_instr),
            Line::Choice(ref choice_instr) =>
                self.ms.execute_choice_instr(choice_instr, &mut self.call_policy),
            Line::Cut(ref cut_instr) =>
                self.ms.execute_cut_instr(cut_instr, &mut self.cut_policy),
            Line::Control(ref control_instr) => {
                let code_dirs = CodeDirs::new(&self.code_dir, &self.modules);
                self.ms.execute_ctrl_instr(code_dirs, &mut self.call_policy,
                                           &mut self.cut_policy, control_instr)
            },
            Line::Fact(ref fact) => {
                for fact_instr in fact {
                    if self.failed() {
                        break;
                    }

                    self.ms.execute_fact_instr(&fact_instr);
                }

                self.ms.p += 1;
            },
            Line::Indexing(ref indexing_instr) =>
                self.ms.execute_indexing_instr(&indexing_instr),
            Line::IndexedChoice(ref choice_instr) =>
                self.ms.execute_indexed_choice_instr(choice_instr, &mut self.call_policy),
            Line::Query(ref query) => {
                for query_instr in query {
                    if self.failed() {
                        break;
                    }

                    self.ms.execute_query_instr(&query_instr);
                }

                self.ms.p += 1;
            }
        }
    }

    fn backtrack(&mut self)
    {
        if self.ms.b > 0 {
            let b = self.ms.b - 1;

            self.ms.b0 = self.ms.or_stack[b].b0;
            self.ms.p  = self.ms.or_stack[b].bp.clone();

            if let CodePtr::Local(LocalCodePtr::TopLevel(_, p)) = self.ms.p {
                self.ms.fail = p == 0;
            } else {
                self.ms.fail = false;
            }
        } else {
            self.ms.p = CodePtr::Local(LocalCodePtr::TopLevel(0, 0));
        }
    }

    fn query_stepper<'a>(&mut self)
    {
        loop {
            self.execute_instr();

            if self.failed() {
                self.backtrack();
            }

            match self.ms.p {
                CodePtr::Local(LocalCodePtr::DirEntry(p, _)) if p < self.code.len() => {},
                CodePtr::Local(_) => break,
                _ => {}
            };
        }
    }

    fn record_var_places(&self, chunk_num: usize, alloc_locs: &AllocVarDict,
                         heap_locs: &mut HeapVarDict)
    {
        for (var, var_data) in alloc_locs {
            match var_data {
                &VarData::Perm(p) if p > 0 => {
                    let e = self.ms.e;
                    let r = var_data.as_reg_type().reg_num();
                    let addr = self.ms.and_stack[e][r].clone();

                    heap_locs.insert(var.clone(), addr);
                },
                &VarData::Temp(cn, _, _) if cn == chunk_num => {
                    let r = var_data.as_reg_type();

                    if r.reg_num() != 0 {
                        let addr = self.ms[r].clone();
                        heap_locs.insert(var.clone(), addr);
                    }
                },
                _ => {}
            }
        }
    }

    fn run_query(&mut self, alloc_locs: &AllocVarDict, heap_locs: &mut HeapVarDict)
    {
        let end_ptr = top_level_code_ptr!(0, self.cached_query_size());

        while self.ms.p < end_ptr {
            if let CodePtr::Local(LocalCodePtr::TopLevel(mut cn, p)) = self.ms.p {
                match &self[LocalCodePtr::TopLevel(cn, p)] {
                    &Line::Control(ref ctrl_instr) if ctrl_instr.is_jump_instr() => {
                        self.record_var_places(cn, alloc_locs, heap_locs);
                        cn += 1;
                    },
                    _ => {}
                }

                self.ms.p = top_level_code_ptr!(cn, p);
            }

            self.query_stepper();

            match self.ms.p {
                CodePtr::Local(LocalCodePtr::TopLevel(_, p)) if p > 0 => {},
                _ => {
                    if heap_locs.is_empty() {
                        self.record_var_places(0, alloc_locs, heap_locs);
                    }

                    break;
                }
            };
        }
    }

    fn fail(&mut self, heap_locs: &HeapVarDict) -> EvalSession
    {
        if self.ms.ball.stub.len() > 0 {
            let h = self.ms.heap.h;
            self.ms.copy_and_align_ball_to_heap();

            let error_str = self.ms.print_exception(Addr::HeapCell(h),
                                                    &heap_locs,
                                                    TermFormatter {},
                                                    PrinterOutputter::new())
                                .result();

            EvalSession::from(SessionError::QueryFailureWithException(error_str))
        } else {
            EvalSession::from(SessionError::QueryFailure)
        }
    }

    pub fn submit_query(&mut self, code: Code, alloc_locs: AllocVarDict) -> EvalSession
    {
        let mut heap_locs = HashMap::new();

        self.cached_query = Some(code);
        self.run_query(&alloc_locs, &mut heap_locs);

        if self.failed() {
            self.fail(&heap_locs)
        } else {
            EvalSession::InitialQuerySuccess(alloc_locs, heap_locs)
        }
    }

    pub fn continue_query(&mut self, alloc_l: &AllocVarDict, heap_l: &mut HeapVarDict) -> EvalSession
    {
        if !self.or_stack_is_empty() {
            let b = self.ms.b - 1;
            self.ms.p = self.ms.or_stack[b].bp.clone();

            if let CodePtr::Local(LocalCodePtr::TopLevel(_, 0)) = self.ms.p {
                return EvalSession::from(SessionError::QueryFailure);
            }

            self.run_query(alloc_l, heap_l);

            if self.failed() {
                self.fail(&heap_l)
            } else {
                EvalSession::SubsequentQuerySuccess
            }
        } else {
            EvalSession::from(SessionError::QueryFailure)
        }
    }

    pub fn heap_view<Outputter>(&self, var_dir: &HeapVarDict, mut output: Outputter) -> Outputter
       where Outputter: HCValueOutputter
    {
        for (var, addr) in var_dir {
            let fmt = TermFormatter {};
            output = self.ms.print_var_eq(var.clone(), addr.clone(), var_dir, fmt, output);
        }

        output
    }

    pub fn or_stack_is_empty(&self) -> bool {
        self.ms.b == 0
    }

    pub fn clear(&mut self) {
        let mut machine = Machine::new();
        swap(self, &mut machine);
    }

    pub fn reset(&mut self) {
        self.cut_policy = Box::new(DefaultCutPolicy {});
        self.ms.reset();
    }
}
<commit_msg>get rid of LocalMachineCodeIndex<commit_after>use prolog::ast::*;
use prolog::compile::*;
use prolog::heap_print::*;
use prolog::tabled_rc::*;

mod machine_errors;
pub(super) mod machine_state;
#[macro_use]
mod machine_state_impl;
mod system_calls;

use prolog::machine::machine_state::*;

use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::mem::swap;
use std::ops::Index;
use std::rc::Rc;

pub struct MachineCodeIndex<'a> {
    pub code_dir: &'a mut CodeDir,
    pub op_dir: &'a mut OpDir,
    pub modules: &'a ModuleDir
}

impl ClauseType {
    pub(super)
    fn lookup(indices: &mut MachineCodeIndex, name: ClauseName, arity: usize, fixity: Option<Fixity>) -> Self
    {
        match ClauseType::from(name, arity, fixity) {
            ClauseType::Named(name, _) => {
                let idx = indices.code_dir.entry((name.clone(), arity))
                    .or_insert(CodeIndex::default());

                ClauseType::Named(name, idx.clone())
            },
            ClauseType::Op(name, fixity, _) => {
                let idx = indices.code_dir.entry((name.clone(), arity))
                    .or_insert(CodeIndex::default());

                ClauseType::Op(name, fixity, idx.clone())
            },
            ct => ct
        }
    }
}

pub struct Machine {
    ms: MachineState,
    call_policy: Box<CallPolicy>,
    cut_policy: Box<CutPolicy>,
    code: Code,
    pub(super) code_dir: CodeDir,
    pub(super) op_dir: OpDir,
    term_dir: TermDir,
    pub(super) modules: ModuleDir,
    cached_query: Option
}

impl Index<LocalCodePtr> for Machine {
    type Output = Line;

    fn index(&self, ptr: LocalCodePtr) -> &Self::Output {
        match ptr {
            LocalCodePtr::TopLevel(_, p) => {
                match &self.cached_query {
                    &Some(ref cq) => &cq[p],
                    &None => panic!("Out-of-bounds top level index.")
                }
            },
            LocalCodePtr::DirEntry(p, _) => &self.code[p]
        }
    }
}

impl<'a> SubModuleUser for MachineCodeIndex<'a> {
    fn op_dir(&mut self) -> &mut OpDir {
        self.op_dir
    }

    fn insert_dir_entry(&mut self, name: ClauseName, arity: usize, idx: ModuleCodeIndex) {
        if let Some(ref mut code_idx) = self.code_dir.get_mut(&(name.clone(), arity)) {
            if !code_idx.is_undefined() {
                println!("warning: overwriting {}/{}", &name, arity);
            }

            set_code_index!(code_idx, idx.0, idx.1);

            return;
        }

        self.code_dir.insert((name, arity), CodeIndex::from(idx));
    }
}

static LISTS: &str   = include_str!("../lib/lists.pl");
static CONTROL: &str = include_str!("../lib/control.pl");
static QUEUES: &str  = include_str!("../lib/queues.pl");

impl Machine {
    pub fn new() -> Self {
        let mut wam = Machine {
            ms: MachineState::new(Rc::new(RefCell::new(HashSet::new()))),
            call_policy: Box::new(DefaultCallPolicy {}),
            cut_policy: Box::new(DefaultCutPolicy {}),
            code: Code::new(),
            code_dir: CodeDir::new(),
            term_dir: TermDir::new(),
            op_dir: default_op_dir(),
            modules: HashMap::new(),
            cached_query: None
        };

        compile_listing(&mut wam, BUILTINS);
        wam.use_module_in_toplevel(clause_name!("builtins"));

        compile_listing(&mut wam, LISTS);
        compile_listing(&mut wam, CONTROL);
        compile_listing(&mut wam, QUEUES);

        wam
    }

    fn remove_module(&mut self, module_name: ClauseName) {
        let iter = if let Some(submodule) = self.modules.get(&module_name) {
            submodule.module_decl.exports.iter().cloned()
        } else {
            return;
        };

        for (name, arity) in iter {
            let name = name.defrock_brackets();

            match self.code_dir.get(&(name.clone(), arity)).cloned() {
                Some(CodeIndex (ref code_idx)) => {
                    if &code_idx.borrow().1 != &module_name {
                        continue;
                    }

                    self.code_dir.remove(&(name.clone(), arity));

                    // remove or respecify ops.
                    if arity == 2 {
                        if let Some((_, _, mod_name)) = self.op_dir.get(&(name.clone(), Fixity::In)).cloned()
                        {
                            if mod_name == module_name {
                                self.op_dir.remove(&(name.clone(), Fixity::In));
                            }
                        }
                    } else if arity == 1 {
                        if let Some((_, _, mod_name)) = self.op_dir.get(&(name.clone(), Fixity::Pre)).cloned()
                        {
                            if mod_name == module_name {
                                self.op_dir.remove(&(name.clone(), Fixity::Pre));
                            }
                        }

                        if let Some((_, _, mod_name)) = self.op_dir.get(&(name.clone(), Fixity::Post)).cloned()
                        {
                            if mod_name == module_name {
                                self.op_dir.remove(&(name.clone(), Fixity::Post));
                            }
                        }
                    }
                },
                _ => {}
            };
        }
    }

    pub fn failed(&self) -> bool {
        self.ms.fail
    }

    pub fn atom_tbl(&self) -> TabledData<Atom> {
        self.ms.atom_tbl.clone()
    }

    pub fn use_qualified_module_in_toplevel(&mut self, name: ClauseName, exports: Vec<PredicateKey>) -> EvalSession
    {
        self.remove_module(name.clone());

        if let Some(mut module) = self.modules.remove(&name) {
            let result = {
                let mut indices = machine_code_index!(&mut self.code_dir, &mut self.op_dir,
                                                      &self.modules);
                indices.use_qualified_module(&mut module, &exports)
            };

            self.modules.insert(name, module);
            result
        } else {
            EvalSession::from(SessionError::ModuleNotFound)
        }
    }

    pub fn use_module_in_toplevel(&mut self, name: ClauseName) -> EvalSession
    {
        self.remove_module(name.clone());

        if let Some(mut module) = self.modules.remove(&name) {
            let result = {
                let mut indices = machine_code_index!(&mut self.code_dir, &mut self.op_dir,
                                                      &self.modules);
                indices.use_module(&mut module)
            };

            self.modules.insert(name, module);
            result
        } else {
            EvalSession::from(SessionError::ModuleNotFound)
        }
    }

    pub fn get_module(&self, name: ClauseName) -> Option<&Module> {
        self.modules.get(&name)
    }

    pub fn add_batched_code(&mut self, mut code: Code, code_dir: CodeDir) {
        self.code.append(&mut code);
        self.code_dir.extend(code_dir.into_iter());
    }

    pub fn add_batched_ops(&mut self, op_dir: OpDir) {
        self.op_dir.extend(op_dir.into_iter());
    }

    pub fn add_module(&mut self, module: Module, code: Code) {
        self.modules.insert(module.module_decl.name.clone(), module);
        self.code.extend(code.into_iter());
    }

    pub fn add_user_code(&mut self, name: ClauseName, arity: usize, code: Code, pred: Predicate)
                         -> EvalSession
    {
        match self.code_dir.get(&(name.clone(), arity)) {
            Some(&CodeIndex (ref idx)) if idx.borrow().1 != clause_name!("user") =>
                if !(&CodeIndex(idx.clone())).is_undefined() {
                    return EvalSession::from(SessionError::ImpermissibleEntry(format!("{}/{}",
                                                                                name,
                                                                                arity)))
                },
            _ => {}
        };

        let offset = self.code.len();

        self.code.extend(code.into_iter());
        self.term_dir.insert((name.clone(), arity), pred);

        let idx = self.code_dir.entry((name, arity))
            .or_insert(CodeIndex::from((offset, clause_name!("user"))));

        set_code_index!(idx, IndexPtr::Index(offset), clause_name!("user"));
        EvalSession::EntrySuccess
    }

    pub fn code_size(&self) -> usize {
        self.code.len()
    }

    fn cached_query_size(&self) -> usize {
        match &self.cached_query {
            &Some(ref query) => query.len(),
            _ => 0
        }
    }

    fn lookup_instr(&self, p: CodePtr) -> Option<Line> {
        match p {
            CodePtr::Local(LocalCodePtr::TopLevel(_, p)) =>
                match &self.cached_query {
                    &Some(ref cq) => Some(cq[p].clone()),
                    &None => None
                },
            CodePtr::Local(LocalCodePtr::DirEntry(p, _)) =>
                Some(self.code[p].clone()),
            CodePtr::BuiltInClause(built_in, _) =>
                Some(call_clause!(ClauseType::BuiltIn(built_in.clone()), built_in.arity(),
                                  0, self.ms.last_call)),
            CodePtr::CallN(arity, _) =>
                Some(call_clause!(ClauseType::CallN, arity, 0, self.ms.last_call))
        }
    }

    fn execute_instr(&mut self)
    {
        let instr = match self.lookup_instr(self.ms.p.clone()) {
            Some(instr) => instr,
            None => return
        };

        match instr {
            Line::Arithmetic(ref arith_instr) =>
                self.ms.execute_arith_instr(arith_instr),
            Line::Choice(ref choice_instr) =>
                self.ms.execute_choice_instr(choice_instr, &mut self.call_policy),
            Line::Cut(ref cut_instr) =>
                self.ms.execute_cut_instr(cut_instr, &mut self.cut_policy),
            Line::Control(ref control_instr) => {
                let code_dirs = CodeDirs::new(&self.code_dir, &self.modules);
                self.ms.execute_ctrl_instr(code_dirs, &mut self.call_policy,
                                           &mut self.cut_policy, control_instr)
            },
            Line::Fact(ref fact) => {
                for fact_instr in fact {
                    if self.failed() {
                        break;
                    }

                    self.ms.execute_fact_instr(&fact_instr);
                }

                self.ms.p += 1;
            },
            Line::Indexing(ref indexing_instr) =>
                self.ms.execute_indexing_instr(&indexing_instr),
            Line::IndexedChoice(ref choice_instr) =>
                self.ms.execute_indexed_choice_instr(choice_instr, &mut self.call_policy),
            Line::Query(ref query) => {
                for query_instr in query {
                    if self.failed() {
                        break;
                    }

                    self.ms.execute_query_instr(&query_instr);
                }

                self.ms.p += 1;
            }
        }
    }

    fn backtrack(&mut self)
    {
        if self.ms.b > 0 {
            let b = self.ms.b - 1;

            self.ms.b0 = self.ms.or_stack[b].b0;
            self.ms.p  = self.ms.or_stack[b].bp.clone();

            if let CodePtr::Local(LocalCodePtr::TopLevel(_, p)) = self.ms.p {
                self.ms.fail = p == 0;
            } else {
                self.ms.fail = false;
            }
        } else {
            self.ms.p = CodePtr::Local(LocalCodePtr::TopLevel(0, 0));
        }
    }

    fn query_stepper<'a>(&mut self)
    {
        loop {
            self.execute_instr();

            if self.failed() {
                self.backtrack();
            }

            match self.ms.p {
                CodePtr::Local(LocalCodePtr::DirEntry(p, _)) if p < self.code.len() => {},
                CodePtr::Local(_) => break,
                _ => {}
            };
        }
    }

    fn record_var_places(&self, chunk_num: usize, alloc_locs: &AllocVarDict,
                         heap_locs: &mut HeapVarDict)
    {
        for (var, var_data) in alloc_locs {
            match var_data {
                &VarData::Perm(p) if p > 0 => {
                    let e = self.ms.e;
                    let r = var_data.as_reg_type().reg_num();
                    let addr = self.ms.and_stack[e][r].clone();

                    heap_locs.insert(var.clone(), addr);
                },
                &VarData::Temp(cn, _, _) if cn == chunk_num => {
                    let r = var_data.as_reg_type();

                    if r.reg_num() != 0 {
                        let addr = self.ms[r].clone();
                        heap_locs.insert(var.clone(), addr);
                    }
                },
                _ => {}
            }
        }
    }

    fn run_query(&mut self, alloc_locs: &AllocVarDict, heap_locs: &mut HeapVarDict)
    {
        let end_ptr = top_level_code_ptr!(0, self.cached_query_size());

        while self.ms.p < end_ptr {
            if let CodePtr::Local(LocalCodePtr::TopLevel(mut cn, p)) = self.ms.p {
                match &self[LocalCodePtr::TopLevel(cn, p)] {
                    &Line::Control(ref ctrl_instr) if ctrl_instr.is_jump_instr() => {
                        self.record_var_places(cn, alloc_locs, heap_locs);
                        cn += 1;
                    },
                    _ => {}
                }

                self.ms.p = top_level_code_ptr!(cn, p);
            }

            self.query_stepper();

            match self.ms.p {
                CodePtr::Local(LocalCodePtr::TopLevel(_, p)) if p > 0 => {},
                _ => {
                    if heap_locs.is_empty() {
                        self.record_var_places(0, alloc_locs, heap_locs);
                    }

                    break;
                }
            };
        }
    }

    fn fail(&mut self, heap_locs: &HeapVarDict) -> EvalSession
    {
        if self.ms.ball.stub.len() > 0 {
            let h = self.ms.heap.h;
            self.ms.copy_and_align_ball_to_heap();

            let error_str = self.ms.print_exception(Addr::HeapCell(h),
                                                    &heap_locs,
                                                    TermFormatter {},
                                                    PrinterOutputter::new())
                                .result();

            EvalSession::from(SessionError::QueryFailureWithException(error_str))
        } else {
            EvalSession::from(SessionError::QueryFailure)
        }
    }

    pub fn submit_query(&mut self, code: Code, alloc_locs: AllocVarDict) -> EvalSession
    {
        let mut heap_locs = HashMap::new();

        self.cached_query = Some(code);
        self.run_query(&alloc_locs, &mut heap_locs);

        if self.failed() {
            self.fail(&heap_locs)
        } else {
            EvalSession::InitialQuerySuccess(alloc_locs, heap_locs)
        }
    }

    pub fn continue_query(&mut self, alloc_l: &AllocVarDict, heap_l: &mut HeapVarDict) -> EvalSession
    {
        if !self.or_stack_is_empty() {
            let b = self.ms.b - 1;
            self.ms.p = self.ms.or_stack[b].bp.clone();

            if let CodePtr::Local(LocalCodePtr::TopLevel(_, 0)) = self.ms.p {
                return EvalSession::from(SessionError::QueryFailure);
            }

            self.run_query(alloc_l, heap_l);

            if self.failed() {
                self.fail(&heap_l)
            } else {
                EvalSession::SubsequentQuerySuccess
            }
        } else {
            EvalSession::from(SessionError::QueryFailure)
        }
    }

    pub fn heap_view<Outputter>(&self, var_dir: &HeapVarDict, mut output: Outputter) -> Outputter
       where Outputter: HCValueOutputter
    {
        for (var, addr) in var_dir {
            let fmt = TermFormatter {};
            output = self.ms.print_var_eq(var.clone(), addr.clone(), var_dir, fmt, output);
        }

        output
    }

    pub fn or_stack_is_empty(&self) -> bool {
        self.ms.b == 0
    }

    pub fn clear(&mut self) {
        let mut machine = Machine::new();
        swap(self, &mut machine);
    }

    pub fn reset(&mut self) {
        self.cut_policy = Box::new(DefaultCutPolicy {});
        self.ms.reset();
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// 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.

extern crate habitat_common as common;
extern crate habitat_core as hcore;
#[macro_use]
extern crate habitat_sup as sup;
#[macro_use]
extern crate log;
extern crate env_logger;
extern crate ansi_term;
extern crate libc;
#[macro_use]
extern crate clap;
extern crate url;

use std::net::{SocketAddr, ToSocketAddrs};
use std::path::{Path, PathBuf};
use std::result;
use std::str::FromStr;

use ansi_term::Colour::{Red, Yellow};
use clap::{App, ArgMatches};
use hcore::env as henv;
use hcore::crypto::{default_cache_key_path, SymKey};
use hcore::crypto::init as crypto_init;
use hcore::package::{PackageArchive, PackageIdent};
use hcore::url::{DEFAULT_DEPOT_URL, DEPOT_URL_ENVVAR};
use url::Url;

use sup::config::{GossipListenAddr, GOSSIP_DEFAULT_PORT};
use sup::error::{Error, Result};
use sup::command;
use sup::http_gateway;
use sup::manager::ManagerConfig;
use sup::manager::service::{ServiceBind, ServiceSpec, Topology, UpdateStrategy};

/// Our output key
static LOGKEY: &'static str = "MN";

/// The version number
const VERSION: &'static str = include_str!(concat!(env!("OUT_DIR"), "/VERSION"));

static RING_ENVVAR: &'static str = "HAB_RING";
static RING_KEY_ENVVAR: &'static str = "HAB_RING_KEY";

fn main() {
    env_logger::init().unwrap();
    if let Err(e) = start() {
        println!("{}", e);
        std::process::exit(1);
    }
}

fn start() -> Result<()> {
    crypto_init();
    let app_matches = cli().get_matches();
    match app_matches.subcommand() {
        ("bash", Some(m)) => sub_bash(m),
        ("config", Some(m)) => sub_config(m),
        ("sh", Some(m)) => sub_sh(m),
        ("start", Some(m)) => sub_start(m),
        _ => unreachable!(),
    }
}

fn cli<'a, 'b>() -> App<'a, 'b> {
    clap_app!(("hab-sup") =>
        (about: "The Habitat Supervisor")
        (version: VERSION)
        (author: "\nAuthors: The Habitat Maintainers <humans@habitat.sh>\n")
        (@setting VersionlessSubcommands)
        (@setting SubcommandRequiredElseHelp)
        (@arg VERBOSE: -v +global "Verbose output; shows line numbers")
        (@arg NO_COLOR: --("no-color") +global "Turn ANSI color off")
        (@subcommand bash =>
            (about: "Start an interactive Bash-like shell")
            (aliases: &["b", "ba", "bas"])
        )
        (@subcommand config =>
            (about: "Displays the default configuration options for a service")
            (aliases: &["c", "co", "con", "conf", "confi"])
            (@arg PKG_IDENT: +required +takes_value
                "A package identifier (ex: core/redis, core/busybox-static/1.42.2)")
        )
        (@subcommand sh =>
            (about: "Start an interactive Bourne-like shell")
            (aliases: &[])
        )
        (@subcommand start =>
            (about: "Start a Habitat-supervised service from a package or artifact")
            (aliases: &["st", "sta", "star"])
            (@arg LISTEN_GOSSIP: --("listen-gossip") +takes_value
                "The listen address for the gossip system [default: 0.0.0.0:9638]")
            (@arg LISTEN_HTTP: --("listen-http") +takes_value
                "The listen address for the HTTP gateway [default: 0.0.0.0:9631]")
            (@arg PEER: --peer +takes_value +multiple
                "The listen address of an initial peer (IP[:PORT])")
            (@arg PERMANENT_PEER: --("permanent-peer") -I "If this service is a permanent peer")
            (@arg RING: --ring -r +takes_value "Ring key name")
            (@arg PKG_IDENT_OR_ARTIFACT: +required
                "A Habitat package identifier (ex: acme/redis) or filepath to a Habitat Artifact \
                (ex: /home/acme-redis-3.0.7-21120102031201-x86_64-linux.hart)")
            (@group SVC_ARGS =>
                (@attributes +multiple requires[PKG_IDENT_OR_ARTIFACT])
                (@arg GROUP: --group +takes_value
                    "The service group; shared config and topology [default: default].")
                (@arg ORGANIZATION: --org +takes_value
                    "The organization that a service is part of")
                (@arg DEPOT_URL: --url -u +takes_value {valid_url}
                    "Use a specific Depot URL (ex: http://depot.example.com/v1/depot)")
                (@arg TOPOLOGY: --topology -t +takes_value {valid_topology}
                    "Service topology; [default: none]")
                (@arg STRATEGY: --strategy -s +takes_value {valid_update_strategy}
                    "The update strategy; [default: none] [values: none, at-once, rolling]")
                (@arg BIND: --bind +takes_value +multiple
                    "One or more service groups to bind to a configuration")
                (@arg CONFIG_DIR: --("config-from") +takes_value {dir_exists}
                    "Use package config from this path, rather than the package itself")
            )
        )
    )
}

fn sub_bash(m: &ArgMatches) -> Result<()> {
    if m.is_present("VERBOSE") {
        sup::output::set_verbose(true);
    }
    if m.is_present("NO_COLOR") {
        sup::output::set_no_color(true);
    }

    command::shell::bash()
}

fn sub_config(m: &ArgMatches) -> Result<()> {
    let ident = try!(PackageIdent::from_str(m.value_of("PKG_IDENT").unwrap()));

    try!(common::command::package::config::start(&ident, "/"));
    Ok(())
}

fn sub_sh(m: &ArgMatches) -> Result<()> {
    if m.is_present("VERBOSE") {
        sup::output::set_verbose(true);
    }
    if m.is_present("NO_COLOR") {
        sup::output::set_no_color(true);
    }

    command::shell::sh()
}

fn sub_start(m: &ArgMatches) -> Result<()> {
    if m.is_present("VERBOSE") {
        sup::output::set_verbose(true);
    }
    if m.is_present("NO_COLOR") {
        sup::output::set_no_color(true);
    }

    let mut cfg = ManagerConfig::default();

    if let Some(addr_str) = m.value_of("LISTEN_GOSSIP") {
        cfg.gossip_listen = try!(GossipListenAddr::from_str(addr_str));
    }
    if let Some(addr_str) = m.value_of("LISTEN_HTTP") {
        cfg.http_listen = try!(http_gateway::ListenAddr::from_str(addr_str));
    }
    if m.is_present("PERMANENT_PEER") {
        cfg.gossip_permanent = true;
    }
    // TODO fn: Clean this up--using a for loop doesn't feel good however an iterator was
    // causing a lot of developer/compiler type confusion
    let mut gossip_peers: Vec<SocketAddr> = Vec::new();
    if let Some(peers) = m.values_of("PEER") {
        for peer in peers {
            let peer_addr = if peer.find(':').is_some() {
                peer.to_string()
            } else {
                format!("{}:{}", peer, GOSSIP_DEFAULT_PORT)
            };
            let addrs: Vec<SocketAddr> = match peer_addr.to_socket_addrs() {
                Ok(addrs) => addrs.collect(),
                Err(e) => {
                    outputln!("Failed to resolve peer: {}", peer_addr);
                    return Err(sup_error!(Error::NameLookup(e)));
                }
            };
            let addr: SocketAddr = addrs[0];
            gossip_peers.push(addr);
        }
    }
    cfg.gossip_peers = gossip_peers;
    let ring = match m.value_of("RING") {
        Some(val) => Some(try!(SymKey::get_latest_pair_for(&val, &default_cache_key_path(None)))),
        None => {
            match henv::var(RING_KEY_ENVVAR) {
                Ok(val) => {
                    let (key, _) = try!(SymKey::write_file_from_str(&val,
                                                                    &default_cache_key_path(None)));
                    Some(key)
                }
                Err(_) => {
                    match henv::var(RING_ENVVAR) {
                        Ok(val) => {
                            Some(try!(SymKey::get_latest_pair_for(&val,
                                                                  &default_cache_key_path(None))))
                        }
                        Err(_) => None,
                    }
                }
            }
        }
    };
    if let Some(ring) = ring {
        cfg.ring = Some(ring.name_with_rev());
    }

    let mut maybe_local_artifact: Option<&str> = None;
    let ident_or_artifact = m.value_of("PKG_IDENT_OR_ARTIFACT").unwrap();
    let ident = if Path::new(ident_or_artifact).is_file() {
        maybe_local_artifact = Some(ident_or_artifact);
        try!(PackageArchive::new(Path::new(ident_or_artifact)).ident())
    } else {
        try!(PackageIdent::from_str(ident_or_artifact))
    };
    let spec = try!(spec_from_matches(&ident, m));

    try!(command::start::package(cfg, spec, maybe_local_artifact));
    Ok(())
}

fn spec_from_matches(ident: &PackageIdent, m: &ArgMatches) -> Result<ServiceSpec> {
    let mut spec = ServiceSpec::default_for(ident.clone());

    if let Some(group) = m.value_of("GROUP") {
        spec.group = group.to_string();
    }
    if let Some(org) = m.value_of("ORGANIZATION") {
        spec.organization = Some(org.to_string());
    }
    let env_or_default = henv::var(DEPOT_URL_ENVVAR).unwrap_or(DEFAULT_DEPOT_URL.to_string());
    let url = m.value_of("DEPOT_URL").unwrap_or(&env_or_default);
    spec.depot_url = String::from(url);
    if let Some(topology) = m.value_of("TOPOLOGY") {
        spec.topology = Topology::from_str(topology)?;
    }
    if let Some(ref strategy) = m.value_of("STRATEGY") {
        spec.update_strategy = UpdateStrategy::from_str(strategy)?;
    }
    if let Some(bind_strs) = m.values_of("BIND") {
        let mut binds = Vec::new();
        for bind_str in bind_strs {
            binds.push(ServiceBind::from_str(bind_str)?);
        }
        spec.binds = binds;
    }
    if let Some(ref config_from) = m.value_of("CONFIG_DIR") {
        spec.config_from = Some(PathBuf::from(config_from));
        outputln!("");
        outputln!("{} Setting '{}' should only be used in development, not production!",
                  Red.bold().paint("WARNING:".to_string()),
                  Yellow.bold().paint(format!("--config-from {}", config_from)));
        outputln!("");
    }

    Ok(spec)
}

fn dir_exists(val: String) -> result::Result<(), String> {
    if Path::new(&val).is_dir() {
        Ok(())
    } else {
        Err(format!("Directory: '{}' cannot be found", &val))
    }
}

fn valid_topology(val: String) -> result::Result<(), String> {
    match Topology::from_str(&val) {
        Ok(_) => Ok(()),
        Err(_) => Err(format!("Service topology: '{}' is not valid", &val)),
    }
}

fn valid_update_strategy(val: String) -> result::Result<(), String> {
    match UpdateStrategy::from_str(&val) {
        Ok(_) => Ok(()),
        Err(_) => Err(format!("Update strategy: '{}' is not valid", &val)),
    }
}

fn valid_url(val: String) -> result::Result<(), String> {
    match Url::parse(&val) {
        Ok(_) => Ok(()),
        Err(_) => Err(format!("URL: '{}' is not valid", &val)),
    }
}
<commit_msg>[sup] Remove un-used `#[macro_use]` in `main.rs`.<commit_after>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// 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.

extern crate habitat_common as common;
extern crate habitat_core as hcore;
#[macro_use]
extern crate habitat_sup as sup;
extern crate log;
extern crate env_logger;
extern crate ansi_term;
extern crate libc;
#[macro_use]
extern crate clap;
extern crate url;

use std::net::{SocketAddr, ToSocketAddrs};
use std::path::{Path, PathBuf};
use std::result;
use std::str::FromStr;

use ansi_term::Colour::{Red, Yellow};
use clap::{App, ArgMatches};
use hcore::env as henv;
use hcore::crypto::{default_cache_key_path, SymKey};
use hcore::crypto::init as crypto_init;
use hcore::package::{PackageArchive, PackageIdent};
use hcore::url::{DEFAULT_DEPOT_URL, DEPOT_URL_ENVVAR};
use url::Url;

use sup::config::{GossipListenAddr, GOSSIP_DEFAULT_PORT};
use sup::error::{Error, Result};
use sup::command;
use sup::http_gateway;
use sup::manager::ManagerConfig;
use sup::manager::service::{ServiceBind, ServiceSpec, Topology, UpdateStrategy};

/// Our output key
static LOGKEY: &'static str = "MN";

/// The version number
const VERSION: &'static str = include_str!(concat!(env!("OUT_DIR"), "/VERSION"));

static RING_ENVVAR: &'static str = "HAB_RING";
static RING_KEY_ENVVAR: &'static str = "HAB_RING_KEY";

fn main() {
    env_logger::init().unwrap();
    if let Err(e) = start() {
        println!("{}", e);
        std::process::exit(1);
    }
}

fn start() -> Result<()> {
    crypto_init();
    let app_matches = cli().get_matches();
    match app_matches.subcommand() {
        ("bash", Some(m)) => sub_bash(m),
        ("config", Some(m)) => sub_config(m),
        ("sh", Some(m)) => sub_sh(m),
        ("start", Some(m)) => sub_start(m),
        _ => unreachable!(),
    }
}

fn cli<'a, 'b>() -> App<'a, 'b> {
    clap_app!(("hab-sup") =>
        (about: "The Habitat Supervisor")
        (version: VERSION)
        (author: "\nAuthors: The Habitat Maintainers <humans@habitat.sh>\n")
        (@setting VersionlessSubcommands)
        (@setting SubcommandRequiredElseHelp)
        (@arg VERBOSE: -v +global "Verbose output; shows line numbers")
        (@arg NO_COLOR: --("no-color") +global "Turn ANSI color off")
        (@subcommand bash =>
            (about: "Start an interactive Bash-like shell")
            (aliases: &["b", "ba", "bas"])
        )
        (@subcommand config =>
            (about: "Displays the default configuration options for a service")
            (aliases: &["c", "co", "con", "conf", "confi"])
            (@arg PKG_IDENT: +required +takes_value
                "A package identifier (ex: core/redis, core/busybox-static/1.42.2)")
        )
        (@subcommand sh =>
            (about: "Start an interactive Bourne-like shell")
            (aliases: &[])
        )
        (@subcommand start =>
            (about: "Start a Habitat-supervised service from a package or artifact")
            (aliases: &["st", "sta", "star"])
            (@arg LISTEN_GOSSIP: --("listen-gossip") +takes_value
                "The listen address for the gossip system [default: 0.0.0.0:9638]")
            (@arg LISTEN_HTTP: --("listen-http") +takes_value
                "The listen address for the HTTP gateway [default: 0.0.0.0:9631]")
            (@arg PEER: --peer +takes_value +multiple
                "The listen address of an initial peer (IP[:PORT])")
            (@arg PERMANENT_PEER: --("permanent-peer") -I "If this service is a permanent peer")
            (@arg RING: --ring -r +takes_value "Ring key name")
            (@arg PKG_IDENT_OR_ARTIFACT: +required
                "A Habitat package identifier (ex: acme/redis) or filepath to a Habitat Artifact \
                (ex: /home/acme-redis-3.0.7-21120102031201-x86_64-linux.hart)")
            (@group SVC_ARGS =>
                (@attributes +multiple requires[PKG_IDENT_OR_ARTIFACT])
                (@arg GROUP: --group +takes_value
                    "The service group; shared config and topology [default: default].")
                (@arg ORGANIZATION: --org +takes_value
                    "The organization that a service is part of")
                (@arg DEPOT_URL: --url -u +takes_value {valid_url}
                    "Use a specific Depot URL (ex: http://depot.example.com/v1/depot)")
                (@arg TOPOLOGY: --topology -t +takes_value {valid_topology}
                    "Service topology; [default: none]")
                (@arg STRATEGY: --strategy -s +takes_value {valid_update_strategy}
                    "The update strategy; [default: none] [values: none, at-once, rolling]")
                (@arg BIND: --bind +takes_value +multiple
                    "One or more service groups to bind to a configuration")
                (@arg CONFIG_DIR: --("config-from") +takes_value {dir_exists}
                    "Use package config from this path, rather than the package itself")
            )
        )
    )
}

fn sub_bash(m: &ArgMatches) -> Result<()> {
    if m.is_present("VERBOSE") {
        sup::output::set_verbose(true);
    }
    if m.is_present("NO_COLOR") {
        sup::output::set_no_color(true);
    }

    command::shell::bash()
}

fn sub_config(m: &ArgMatches) -> Result<()> {
    let ident = try!(PackageIdent::from_str(m.value_of("PKG_IDENT").unwrap()));

    try!(common::command::package::config::start(&ident, "/"));
    Ok(())
}

fn sub_sh(m: &ArgMatches) -> Result<()> {
    if m.is_present("VERBOSE") {
        sup::output::set_verbose(true);
    }
    if m.is_present("NO_COLOR") {
        sup::output::set_no_color(true);
    }

    command::shell::sh()
}

fn sub_start(m: &ArgMatches) -> Result<()> {
    if m.is_present("VERBOSE") {
        sup::output::set_verbose(true);
    }
    if m.is_present("NO_COLOR") {
        sup::output::set_no_color(true);
    }

    let mut cfg = ManagerConfig::default();

    if let Some(addr_str) = m.value_of("LISTEN_GOSSIP") {
        cfg.gossip_listen = try!(GossipListenAddr::from_str(addr_str));
    }
    if let Some(addr_str) = m.value_of("LISTEN_HTTP") {
        cfg.http_listen = try!(http_gateway::ListenAddr::from_str(addr_str));
    }
    if m.is_present("PERMANENT_PEER") {
        cfg.gossip_permanent = true;
    }
    // TODO fn: Clean this up--using a for loop doesn't feel good however an iterator was
    // causing a lot of developer/compiler type confusion
    let mut gossip_peers: Vec<SocketAddr> = Vec::new();
    if let Some(peers) = m.values_of("PEER") {
        for peer in peers {
            let peer_addr = if peer.find(':').is_some() {
                peer.to_string()
            } else {
                format!("{}:{}", peer, GOSSIP_DEFAULT_PORT)
            };
            let addrs: Vec<SocketAddr> = match peer_addr.to_socket_addrs() {
                Ok(addrs) => addrs.collect(),
                Err(e) => {
                    outputln!("Failed to resolve peer: {}", peer_addr);
                    return Err(sup_error!(Error::NameLookup(e)));
                }
            };
            let addr: SocketAddr = addrs[0];
            gossip_peers.push(addr);
        }
    }
    cfg.gossip_peers = gossip_peers;
    let ring = match m.value_of("RING") {
        Some(val) => Some(try!(SymKey::get_latest_pair_for(&val, &default_cache_key_path(None)))),
        None => {
            match henv::var(RING_KEY_ENVVAR) {
                Ok(val) => {
                    let (key, _) = try!(SymKey::write_file_from_str(&val,
                                                                    &default_cache_key_path(None)));
                    Some(key)
                }
                Err(_) => {
                    match henv::var(RING_ENVVAR) {
                        Ok(val) => {
                            Some(try!(SymKey::get_latest_pair_for(&val,
                                                                  &default_cache_key_path(None))))
                        }
                        Err(_) => None,
                    }
                }
            }
        }
    };
    if let Some(ring) = ring {
        cfg.ring = Some(ring.name_with_rev());
    }

    let mut maybe_local_artifact: Option<&str> = None;
    let ident_or_artifact = m.value_of("PKG_IDENT_OR_ARTIFACT").unwrap();
    let ident = if Path::new(ident_or_artifact).is_file() {
        maybe_local_artifact = Some(ident_or_artifact);
        try!(PackageArchive::new(Path::new(ident_or_artifact)).ident())
    } else {
        try!(PackageIdent::from_str(ident_or_artifact))
    };
    let spec = try!(spec_from_matches(&ident, m));

    try!(command::start::package(cfg, spec, maybe_local_artifact));
    Ok(())
}

fn spec_from_matches(ident: &PackageIdent, m: &ArgMatches) -> Result<ServiceSpec> {
    let mut spec = ServiceSpec::default_for(ident.clone());

    if let Some(group) = m.value_of("GROUP") {
        spec.group = group.to_string();
    }
    if let Some(org) = m.value_of("ORGANIZATION") {
        spec.organization = Some(org.to_string());
    }
    let env_or_default = henv::var(DEPOT_URL_ENVVAR).unwrap_or(DEFAULT_DEPOT_URL.to_string());
    let url = m.value_of("DEPOT_URL").unwrap_or(&env_or_default);
    spec.depot_url = String::from(url);
    if let Some(topology) = m.value_of("TOPOLOGY") {
        spec.topology = Topology::from_str(topology)?;
    }
    if let Some(ref strategy) = m.value_of("STRATEGY") {
        spec.update_strategy = UpdateStrategy::from_str(strategy)?;
    }
    if let Some(bind_strs) = m.values_of("BIND") {
        let mut binds = Vec::new();
        for bind_str in bind_strs {
            binds.push(ServiceBind::from_str(bind_str)?);
        }
        spec.binds = binds;
    }
    if let Some(ref config_from) = m.value_of("CONFIG_DIR") {
        spec.config_from = Some(PathBuf::from(config_from));
        outputln!("");
        outputln!("{} Setting '{}' should only be used in development, not production!",
                  Red.bold().paint("WARNING:".to_string()),
                  Yellow.bold().paint(format!("--config-from {}", config_from)));
        outputln!("");
    }

    Ok(spec)
}

fn dir_exists(val: String) -> result::Result<(), String> {
    if Path::new(&val).is_dir() {
        Ok(())
    } else {
        Err(format!("Directory: '{}' cannot be found", &val))
    }
}

fn valid_topology(val: String) -> result::Result<(), String> {
    match Topology::from_str(&val) {
        Ok(_) => Ok(()),
        Err(_) => Err(format!("Service topology: '{}' is not valid", &val)),
    }
}

fn valid_update_strategy(val: String) -> result::Result<(), String> {
    match UpdateStrategy::from_str(&val) {
        Ok(_) => Ok(()),
        Err(_) => Err(format!("Update strategy: '{}' is not valid", &val)),
    }
}

fn valid_url(val: String) -> result::Result<(), String> {
    match Url::parse(&val) {
        Ok(_) => Ok(()),
        Err(_) => Err(format!("URL: '{}' is not valid", &val)),
    }
}
<|endoftext|>


<commit_before>use error::Error;
use client::Client;
use params::{List, Metadata, Timestamp};
use resources::{Currency, Discount, Plan};
use serde_qs as query;

/// The set of parameters that can be used when creating or updating an invoice.
///
/// For more details see https://stripe.com/docs/api#create_invoice, https://stripe.com/docs/api#update_invoice.
#[derive(Default, Serialize)]
pub struct InvoiceParams<'a> {
    #[serde(skip_serializing_if = "Option::is_none")] pub application_fee: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")] pub customer: Option<&'a str>,
    #[serde(skip_serializing_if = "Option::is_none")] pub description: Option<&'a str>,
    #[serde(skip_serializing_if = "Option::is_none")] pub statement_descriptor: Option<&'a str>,
    #[serde(skip_serializing_if = "Option::is_none")] pub subscription: Option<&'a str>,
    #[serde(skip_serializing_if = "Option::is_none")] pub tax_percent: Option<f64>,

    #[serde(skip_serializing_if = "Option::is_none")] pub closed: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")] pub forgiven: Option<bool>,
}

/*
#[derive(Serialize)]
pub struct InvoiceListLinesParams {
    #[serde(skip_serializing_if = "Option::is_none")] pub limit: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")] pub ending_before: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")] pub starting_after: Option<String>,

    ..
}
*/

/// Period is a structure representing a start and end dates.
#[derive(Debug, Deserialize)]
pub struct Period {
    pub start: Timestamp,
    pub end: Timestamp,
}

/// The resource representing a Stripe invoice line item.
///
/// For more details see https://stripe.com/docs/api#invoice_line_item_object.
#[derive(Debug, Deserialize)]
pub struct InvoiceLine {
    pub id: String,
    pub amount: u64,
    pub currency: Currency,
    pub description: Option<String>,
    pub discountable: bool,
    pub livemode: bool,
    pub metadata: Metadata,
    pub period: Period,
    pub plan: Option<Plan>,
    pub proration: bool,
    pub quantity: Option<u64>,
    pub subscription: Option<String>,
    pub subscription_item: Option<String>,
    #[serde(rename = "type")] pub item_type: String, // (invoiceitem, subscription)
}

/// The resource representing a Stripe invoice.
///
/// For more details see https://stripe.com/docs/api#invoice_object.
#[derive(Debug, Deserialize)]
pub struct Invoice {
    pub id: String,
    pub amount_due: u64,
    pub application_fee: Option<u64>,
    pub attempt_count: u64,
    pub attempted: bool,
    pub charge: Option<String>,
    pub closed: bool,
    pub currency: Currency,
    pub customer: String,
    pub date: Timestamp,
    pub description: Option<String>,
    pub discount: Option<Discount>,
    pub ending_balance: Option<i64>,
    pub forgiven: bool,
    pub lines: List<InvoiceLine>,
    pub livemode: bool,
    pub metadata: Metadata,
    pub next_payment_attempt: Option<Timestamp>,
    pub paid: bool,
    pub period_end: Timestamp,
    pub period_start: Timestamp,
    pub receipt_number: Option<String>,
    pub starting_balance: i64,
    pub statment_descriptor: Option<String>,
    pub subscription: Option<String>,
    pub subscription_proration_date: Option<Timestamp>,
    pub subtotal: i64,
    pub tax: Option<i64>,
    pub tax_percent: Option<f64>,
    pub total: i64,
    pub webhooks_delivered_at: Option<Timestamp>,
}

#[derive(Default, Serialize)]
pub struct InvoiceListParams<'a> {
    #[serde(skip_serializing_if = "Option::is_none")] pub limit: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")] pub customer: Option<&'a str>,
}

impl Invoice {
    /// Creates a new invoice.
    ///
    /// For more details see https://stripe.com/docs/api#create_invoice.
    pub fn create(client: &Client, params: InvoiceParams) -> Result<Invoice, Error> {
        client.post("/invoices", params)
    }

    /// Retrieves the details of an invoice.
    ///
    /// For more details see https://stripe.com/docs/api#retrieve_invoice.
    pub fn retrieve(client: &Client, invoice_id: &str) -> Result<Invoice, Error> {
        client.get(&format!("/invoices/{}", invoice_id))
    }

    // TODO: Implement InvoiceListLinesParams
    // pub fn get_lines(client: &Client, invoice_id: &str, params: InvoiceListLinesParams) -> Result<List<InvoiceLine>, Error> {
    //     client.get(&format!("/invoices/{}/lines", invoice_id))
    // }

    // TODO: Implement InvoiceUpcomingParams
    // pub fn get_upcoming(client: &Client, params: InvoiceUpcomingParams) -> Result<Invoice, Error> {
    //     client.get(&format!("/invoices/upcoming?customer={}", invoice_id))
    // }

    /// Pays an invoice.
    ///
    /// For more details see https://stripe.com/docs/api#pay_invoice.
    pub fn pay(client: &Client, invoice_id: &str) -> Result<Invoice, Error> {
        client.post_empty(&format!("/invoices/{}/pay", invoice_id))
    }

    // Lists all invoices with optional customer_id and limit params
    // TODO: Implement InvoiceListParams
    pub fn list(client: &Client, params: InvoiceListParams) -> Result<List<Invoice>, Error> {
        client.get(&format!("/invoices?{}", query::to_string(&params)?))
        // client.get(&format!("/invoices?customer={}", customer_id))
    }
}
<commit_msg>Add update invoice function and doc notes<commit_after>use error::Error;
use client::Client;
use params::{List, Metadata, Timestamp};
use resources::{Currency, Discount, Plan};
use serde_qs as query;

/// The set of parameters that can be used when creating or updating an invoice.
///
/// For more details see https://stripe.com/docs/api#create_invoice, https://stripe.com/docs/api#update_invoice.
#[derive(Default, Serialize)]
pub struct InvoiceParams<'a> {
    #[serde(skip_serializing_if = "Option::is_none")] pub application_fee: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")] pub customer: Option<&'a str>,
    #[serde(skip_serializing_if = "Option::is_none")] pub description: Option<&'a str>,
    #[serde(skip_serializing_if = "Option::is_none")] pub statement_descriptor: Option<&'a str>,
    #[serde(skip_serializing_if = "Option::is_none")] pub subscription: Option<&'a str>,
    #[serde(skip_serializing_if = "Option::is_none")] pub tax_percent: Option<f64>,

    #[serde(skip_serializing_if = "Option::is_none")] pub closed: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")] pub forgiven: Option<bool>,
}

/*
#[derive(Serialize)]
pub struct InvoiceListLinesParams {
    #[serde(skip_serializing_if = "Option::is_none")] pub limit: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")] pub ending_before: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")] pub starting_after: Option<String>,

    ..
}
*/

/// Period is a structure representing a start and end dates.
#[derive(Debug, Deserialize)]
pub struct Period {
    pub start: Timestamp,
    pub end: Timestamp,
}

/// The resource representing a Stripe invoice line item.
///
/// For more details see https://stripe.com/docs/api#invoice_line_item_object.
#[derive(Debug, Deserialize)]
pub struct InvoiceLine {
    pub id: String,
    pub amount: u64,
    pub currency: Currency,
    pub description: Option<String>,
    pub discountable: bool,
    pub livemode: bool,
    pub metadata: Metadata,
    pub period: Period,
    pub plan: Option<Plan>,
    pub proration: bool,
    pub quantity: Option<u64>,
    pub subscription: Option<String>,
    pub subscription_item: Option<String>,
    #[serde(rename = "type")] pub item_type: String, // (invoiceitem, subscription)
}

/// The resource representing a Stripe invoice.
///
/// For more details see https://stripe.com/docs/api#invoice_object.
#[derive(Debug, Deserialize)]
pub struct Invoice {
    pub id: String,
    pub amount_due: u64,
    pub application_fee: Option<u64>,
    pub attempt_count: u64,
    pub attempted: bool,
    pub charge: Option<String>,
    pub closed: bool,
    pub currency: Currency,
    pub customer: String,
    pub date: Timestamp,
    pub description: Option<String>,
    pub discount: Option<Discount>,
    pub ending_balance: Option<i64>,
    pub forgiven: bool,
    pub lines: List<InvoiceLine>,
    pub livemode: bool,
    pub metadata: Metadata,
    pub next_payment_attempt: Option<Timestamp>,
    pub paid: bool,
    pub period_end: Timestamp,
    pub period_start: Timestamp,
    pub receipt_number: Option<String>,
    pub starting_balance: i64,
    pub statment_descriptor: Option<String>,
    pub subscription: Option<String>,
    pub subscription_proration_date: Option<Timestamp>,
    pub subtotal: i64,
    pub tax: Option<i64>,
    pub tax_percent: Option<f64>,
    pub total: i64,
    pub webhooks_delivered_at: Option<Timestamp>,
}

#[derive(Default, Serialize)]
pub struct InvoiceListParams<'a> {
    #[serde(skip_serializing_if = "Option::is_none")] pub limit: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")] pub customer: Option<&'a str>,
}

impl Invoice {
    /// Creates a new invoice.
    ///
    /// For more details see https://stripe.com/docs/api#create_invoice.
    pub fn create(client: &Client, params: InvoiceParams) -> Result<Invoice, Error> {
        client.post("/invoices", params)
    }

    /// Retrieves the details of an invoice.
    ///
    /// For more details see https://stripe.com/docs/api#retrieve_invoice.
    pub fn retrieve(client: &Client, invoice_id: &str) -> Result<Invoice, Error> {
        client.get(&format!("/invoices/{}", invoice_id))
    }

    // TODO: Implement InvoiceListLinesParams
    // pub fn get_lines(client: &Client, invoice_id: &str, params: InvoiceListLinesParams) -> Result<List<InvoiceLine>, Error> {
    //     client.get(&format!("/invoices/{}/lines", invoice_id))
    // }

    // TODO: Implement InvoiceUpcomingParams
    // pub fn get_upcoming(client: &Client, params: InvoiceUpcomingParams) -> Result<Invoice, Error> {
    //     client.get(&format!("/invoices/upcoming?customer={}", invoice_id))
    // }

    /// Pays an invoice.
    ///
    /// For more details see https://stripe.com/docs/api#pay_invoice.
    pub fn pay(client: &Client, invoice_id: &str) -> Result<Invoice, Error> {
        client.post_empty(&format!("/invoices/{}/pay", invoice_id))
    }

    /// Updates an invoice.
    ///
    /// For more details see https://stripe.com/docs/api#update_invoice.
    pub fn update(client: &Client, invoice_id: &str, params: InvoiceParams) -> Result<Invoice, Error> {
        client.post(&format!("/invoices/{}", invoice_id), &params)
    }

    /// Lists all invoices with optional customer_id and limit params
    ///
    /// For more details see https://stripe.com/docs/api#list_invoices.
    pub fn list(client: &Client, params: InvoiceListParams) -> Result<List<Invoice>, Error> {
        client.get(&format!("/invoices?{}", query::to_string(&params)?))
    }
}
<|endoftext|>


<commit_before>use std::cmp;
use std::f64;
use std::collections;

use itertools::Itertools;
use num::rational;

use bio::stats::combinatorics::combinations;
use bio::stats::logprobs;
use bio::stats::logprobs::LogProb;

use model;

pub type MeanExpression = rational::Ratio<u32>;

fn likelihood(x: u32, count: u32, count_exact: u32, readout_model: &model::Readout) -> LogProb {
    let x = x as u64;
    let count = count as u64;
    let count_exact = count_exact as u64;

    let x_c = cmp::min(count, x);
    let x_m = count - x_c;

    let imin = if count_exact + x_c > count { count_exact + x_c - count } else { 0 };
    let imax = cmp::min(count_exact, x) + 1;

    let summands = (imin..imax).map(|i| {
        let combs = combinations(count_exact, i).ln() +
                    combinations(count - count_exact, x_c - i).ln();
        let prob = readout_model.prob_call_exact * i as f64 +
                   readout_model.prob_call_mismatch * (x_c - i) as f64 +
                   readout_model.prob_miscall_exact * (count_exact - i) as f64 +
                   readout_model.prob_miscall_mismatch * (x_m + i - count_exact) as f64;
        (combs + prob).min(0.0)
    }).collect_vec();
    // removed (combinations(count, x_c).ln() as f64) since we already sum over all possibilities
    let likelihood = readout_model.prob_missed * (x - x_c) as f64 + logprobs::log_prob_sum(&summands).min(0.0);
    assert!(!likelihood.is_nan());
    likelihood
}

pub struct Expression {
    offset: u32,
    likelihoods: Vec<LogProb>,
    marginal: LogProb
}

impl Expression {
    pub fn new(count: u32, count_exact: u32, readout_model: &model::Readout) -> Self {
        let offset = if count_exact > 50 { count_exact - 50 } else { 0 };
        let likelihoods = (offset..count + 50).map(|x| likelihood(x, count, count_exact, readout_model)).collect_vec();
        // calculate (marginal / flat_prior)
        let marginal = logprobs::log_prob_sum(&likelihoods);
        // TODO trim values such that zero probabilities are not stored
        let mut expression = Expression {
            offset: offset,
            likelihoods: likelihoods,
            marginal: marginal
        };
        expression.refine_interval();
        expression
    }

    fn refine_interval(&mut self) {
        let min_prob = 0.000001f64.ln();
        let map = self.map();
        let mut min_x = self.min_x();
        for x in (self.min_x()..map).rev() {
            if self.posterior_prob(x) < min_prob {
                min_x = x;
                break;
            }
        }
        let mut max_x = self.max_x();
        for x in map..self.max_x() {
            if self.posterior_prob(x) < min_prob {
                max_x = x;
                break;
            }
        }
        self.likelihoods = self.likelihoods[(min_x - self.offset) as usize..(max_x - self.offset) as usize].to_vec();
        self.offset = min_x;
    }

    pub fn pmf(&self) -> Vec<(u32, LogProb)> {
        (self.min_x()..self.max_x()).map(|x| (x, self.posterior_prob(x))).collect_vec()
    }

    pub fn posterior_prob(&self, x: u32) -> LogProb {
        let x = x as usize;
        if x < self.offset as usize || x >= self.offset as usize + self.likelihoods.len() {
            f64::NEG_INFINITY
        }
        else {
            self.likelihoods[x - self.offset as usize] - self.marginal
        }
    }

    /// The maximum a posteriori probability estimate.
    pub fn map(&self) -> u32 {
        let mut max_x = self.min_x();
        let mut max = self.posterior_prob(max_x);
        for x in self.min_x() + 1..self.max_x() {
            let p = self.posterior_prob(x);
            if p >= max {
                max_x = x;
                max = p;
            }
        }
        max_x
    }

    pub fn expected_value(&self) -> f64 {
        (self.min_x()..self.max_x()).map(|x| x as f64 * self.posterior_prob(x).exp()).fold(0.0, |s, e| s + e)
    }

    pub fn variance(&self) -> f64 {
        (self.min_x()..self.max_x()).map(|x| (x as f64 - self.expected_value()).powi(2) * self.posterior_prob(x).exp()).fold(0.0, |s, e| s + e)
    }

    pub fn min_x(&self) -> u32 {
        self.offset
    }

    pub fn max_x(&self) -> u32 {
        self.offset + self.likelihoods.len() as u32
    }
}

pub struct ExpressionSet {
    pmf: collections::HashMap<MeanExpression, LogProb>
}

impl ExpressionSet {
    pub fn new(expressions: &[Expression]) -> Self {
        let mut pmf = collections::HashMap::new();

        fn dfs(
            i: usize, expression_sum: u32, posterior_prob: LogProb,
            pmf: &mut collections::HashMap<MeanExpression, LogProb>,
            expressions: &[Expression]
        ) {
            if i < expressions.len() {
                let ref expr = expressions[i];
                for x in expr.min_x()..expr.max_x() {
                    dfs(i + 1, expression_sum + x, posterior_prob + expr.posterior_prob(x), pmf, expressions);
                }
            }
            else {
                let mean = rational::Ratio::new(expression_sum, expressions.len() as u32);
                if pmf.contains_key(&mean) {
                    let p = pmf.get_mut(&mean).unwrap();
                    *p = logprobs::log_prob_add(*p, posterior_prob);
                }
                else {
                    pmf.insert(mean, posterior_prob);
                }
            }
        }

        dfs(0, 0, 0.0, &mut pmf, expressions);

        ExpressionSet {
            pmf: pmf
        }
    }

    /// Unorderered iteration over probability mass function (PMF).
    pub fn pmf(&self) -> collections::hash_map::Iter<MeanExpression, LogProb> {
        self.pmf.iter()
    }

    /// Conditional expectation of mean expression.
    pub fn expected_value(&self) -> f64 {
        self.pmf.iter().map(|(mean, prob)| {
            *mean.numer() as f64 / *mean.denom() as f64 * prob.exp()
        }).fold(0.0, |s, e| s + e)
    }
}

#[cfg(test)]
mod tests {
    #![allow(non_upper_case_globals)]

    use itertools::Itertools;
    use nalgebra::ApproxEq;
    use bio::stats::logprobs::Prob;

    use super::{likelihood, Expression, ExpressionSet};
    use model;

    const N: u8 = 16;
    const m: u8 = 4;
    const p0: Prob = 0.04;
    const p1: Prob = 0.1;

    fn setup() -> model::Readout {
        model::Readout::new(N, m, p0, p1)
    }

    #[test]
    fn test_likelihood() {
        let readout = setup();
        assert!(likelihood(0, 0, 0, &readout).exp().approx_eq(&1.0));
        assert!(likelihood(1, 1, 1, &readout).exp().approx_eq(&0.4019988717840602));
    }

    #[test]
    fn test_posterior_prob() {
        let readout = setup();
        let expression = Expression::new(50, 10, &readout);
        println!("{:?}", (0..55).map(|x| expression.posterior_prob(x).exp()).collect_vec());
        assert!(false);
        let expression = Expression::new(5, 5, &readout);
        // check if x=5 yields highest probability
        assert_eq!((0..21).sorted_by(|&x, &y| {
            expression.posterior_prob(x).partial_cmp(&expression.posterior_prob(y)).unwrap()
        })[20], 5);
        // check that expression beyond window yields zero probability
        assert!(expression.posterior_prob(100).exp().approx_eq(&0.0));
    }

    #[test]
    fn test_set_expected_value() {
        let readout = setup();
        let expressions = [
            Expression::new(5, 1, &readout),
            Expression::new(5, 1, &readout),
            Expression::new(5, 1, &readout),
            Expression::new(1, 0, &readout)
        ];
        let expression_set = ExpressionSet::new(&expressions);
        // calculate expected value directly
        let expected_value = expressions.iter().map(|e| e.expected_value()).fold(0.0, |s, e| s + e) / expressions.len() as f64;
        assert!(expression_set.expected_value().approx_eq(&expected_value));
        // check if x=5 yields highest probability
        //assert_eq!((0..20).sorted_by(|&x, &y| {
        //    expression_set.posterior_prob(x).partial_cmp(&expression_set.posterior_prob(y)).unwrap()
        //})[19], 5);
    }
}
<commit_msg>Test.<commit_after>use std::cmp;
use std::f64;
use std::collections;

use itertools::Itertools;
use num::rational;

use bio::stats::combinatorics::combinations;
use bio::stats::logprobs;
use bio::stats::logprobs::LogProb;

use model;

pub type MeanExpression = rational::Ratio<u32>;

fn likelihood(x: u32, count: u32, count_exact: u32, readout_model: &model::Readout) -> LogProb {
    let x = x as u64;
    let count = count as u64;
    let count_exact = count_exact as u64;

    let x_c = cmp::min(count, x);
    let x_m = count - x_c;

    let imin = if count_exact + x_c > count { count_exact + x_c - count } else { 0 };
    let imax = cmp::min(count_exact, x) + 1;

    let summands = (imin..imax).map(|i| {
        let combs = combinations(count_exact, i).ln() +
                    combinations(count - count_exact, x_c - i).ln();
        let prob = readout_model.prob_call_exact * i as f64 +
                   readout_model.prob_call_mismatch * (x_c - i) as f64 +
                   readout_model.prob_miscall_exact * (count_exact - i) as f64 +
                   readout_model.prob_miscall_mismatch * (x_m + i - count_exact) as f64;
        (combs + prob).min(0.0)
    }).collect_vec();
    // removed (combinations(count, x_c).ln() as f64) since we already sum over all possibilities
    let likelihood = readout_model.prob_missed * (x - x_c) as f64 + logprobs::log_prob_sum(&summands).min(0.0);
    assert!(!likelihood.is_nan());
    likelihood
}

pub struct Expression {
    offset: u32,
    likelihoods: Vec<LogProb>,
    marginal: LogProb
}

impl Expression {
    pub fn new(count: u32, count_exact: u32, readout_model: &model::Readout) -> Self {
        let offset = if count_exact > 50 { count_exact - 50 } else { 0 };
        let likelihoods = (offset..count + 50).map(|x| likelihood(x, count, count_exact, readout_model)).collect_vec();
        // calculate (marginal / flat_prior)
        let marginal = logprobs::log_prob_sum(&likelihoods);
        // TODO trim values such that zero probabilities are not stored
        let mut expression = Expression {
            offset: offset,
            likelihoods: likelihoods,
            marginal: marginal
        };
        expression.refine_interval();
        expression
    }

    fn refine_interval(&mut self) {
        let min_prob = 0.000001f64.ln();
        let map = self.map();
        let mut min_x = self.min_x();
        for x in (self.min_x()..map).rev() {
            if self.posterior_prob(x) < min_prob {
                min_x = x;
                break;
            }
        }
        let mut max_x = self.max_x();
        for x in map..self.max_x() {
            if self.posterior_prob(x) < min_prob {
                max_x = x;
                break;
            }
        }
        self.likelihoods = self.likelihoods[(min_x - self.offset) as usize..(max_x - self.offset) as usize].to_vec();
        self.offset = min_x;
    }

    pub fn pmf(&self) -> Vec<(u32, LogProb)> {
        (self.min_x()..self.max_x()).map(|x| (x, self.posterior_prob(x))).collect_vec()
    }

    pub fn posterior_prob(&self, x: u32) -> LogProb {
        let x = x as usize;
        if x < self.offset as usize || x >= self.offset as usize + self.likelihoods.len() {
            f64::NEG_INFINITY
        }
        else {
            self.likelihoods[x - self.offset as usize] - self.marginal
        }
    }

    /// The maximum a posteriori probability estimate.
    pub fn map(&self) -> u32 {
        let mut max_x = self.min_x();
        let mut max = self.posterior_prob(max_x);
        for x in self.min_x() + 1..self.max_x() {
            let p = self.posterior_prob(x);
            if p >= max {
                max_x = x;
                max = p;
            }
        }
        max_x
    }

    pub fn expected_value(&self) -> f64 {
        (self.min_x()..self.max_x()).map(|x| x as f64 * self.posterior_prob(x).exp()).fold(0.0, |s, e| s + e)
    }

    pub fn variance(&self) -> f64 {
        (self.min_x()..self.max_x()).map(|x| (x as f64 - self.expected_value()).powi(2) * self.posterior_prob(x).exp()).fold(0.0, |s, e| s + e)
    }

    pub fn min_x(&self) -> u32 {
        self.offset
    }

    pub fn max_x(&self) -> u32 {
        self.offset + self.likelihoods.len() as u32
    }
}

pub struct ExpressionSet {
    pmf: collections::HashMap<MeanExpression, LogProb>
}

impl ExpressionSet {
    pub fn new(expressions: &[Expression]) -> Self {
        let mut pmf = collections::HashMap::new();

        fn dfs(
            i: usize, expression_sum: u32, posterior_prob: LogProb,
            pmf: &mut collections::HashMap<MeanExpression, LogProb>,
            expressions: &[Expression]
        ) {
            if i < expressions.len() {
                let ref expr = expressions[i];
                for x in expr.min_x()..expr.max_x() {
                    dfs(i + 1, expression_sum + x, posterior_prob + expr.posterior_prob(x), pmf, expressions);
                }
            }
            else {
                let mean = rational::Ratio::new(expression_sum, expressions.len() as u32);
                if pmf.contains_key(&mean) {
                    let p = pmf.get_mut(&mean).unwrap();
                    *p = logprobs::log_prob_add(*p, posterior_prob);
                }
                else {
                    pmf.insert(mean, posterior_prob);
                }
            }
        }

        dfs(0, 0, 0.0, &mut pmf, expressions);

        ExpressionSet {
            pmf: pmf
        }
    }

    /// Unorderered iteration over probability mass function (PMF).
    pub fn pmf(&self) -> collections::hash_map::Iter<MeanExpression, LogProb> {
        self.pmf.iter()
    }

    /// Conditional expectation of mean expression.
    pub fn expected_value(&self) -> f64 {
        self.pmf.iter().map(|(mean, prob)| {
            *mean.numer() as f64 / *mean.denom() as f64 * prob.exp()
        }).fold(0.0, |s, e| s + e)
    }
}

#[cfg(test)]
mod tests {
    #![allow(non_upper_case_globals)]

    use itertools::Itertools;
    use nalgebra::ApproxEq;
    use bio::stats::logprobs::Prob;

    use super::{likelihood, Expression, ExpressionSet};
    use model;

    const N: u8 = 16;
    const m: u8 = 4;
    const p0: Prob = 0.04;
    const p1: Prob = 0.1;

    fn setup() -> model::Readout {
        model::Readout::new(N, m, p0, p1)
    }

    #[test]
    fn test_likelihood() {
        let readout = setup();
        assert!(likelihood(0, 0, 0, &readout).exp().approx_eq(&1.0));
        assert!(likelihood(1, 1, 1, &readout).exp().approx_eq(&0.4019988717840602));
    }

    #[test]
    fn test_posterior_prob() {
        let readout = setup();
        let expression = Expression::new(500, 10, &readout);
        println!("{:?}", (490..510).map(|x| expression.posterior_prob(x).exp()).collect_vec());
        assert!(false);
        let expression = Expression::new(5, 5, &readout);
        // check if x=5 yields highest probability
        assert_eq!((0..21).sorted_by(|&x, &y| {
            expression.posterior_prob(x).partial_cmp(&expression.posterior_prob(y)).unwrap()
        })[20], 5);
        // check that expression beyond window yields zero probability
        assert!(expression.posterior_prob(100).exp().approx_eq(&0.0));
    }

    #[test]
    fn test_set_expected_value() {
        let readout = setup();
        let expressions = [
            Expression::new(5, 1, &readout),
            Expression::new(5, 1, &readout),
            Expression::new(5, 1, &readout),
            Expression::new(1, 0, &readout)
        ];
        let expression_set = ExpressionSet::new(&expressions);
        // calculate expected value directly
        let expected_value = expressions.iter().map(|e| e.expected_value()).fold(0.0, |s, e| s + e) / expressions.len() as f64;
        assert!(expression_set.expected_value().approx_eq(&expected_value));
        // check if x=5 yields highest probability
        //assert_eq!((0..20).sorted_by(|&x, &y| {
        //    expression_set.posterior_prob(x).partial_cmp(&expression_set.posterior_prob(y)).unwrap()
        //})[19], 5);
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::f64;
use std::iter::Peekable;
use std::iter::FromIterator;
use std::str::Chars;

use token;

pub struct Lexer<'a> {
    raw_input: &'a str,
    ast: Vec<token::Token>,
    pub errors: Vec<String>
}

impl<'a> Lexer<'a> {
    pub fn new(input: &str) -> Lexer {
        Lexer {
            raw_input: input,
            ast: Vec::new(),
            errors: vec![]
        }
    }

    pub fn lex(&mut self) {
        let tokens = vec![];

        let ast: Vec<token::Token> = self.consume_input(self.raw_input, tokens);
        self.ast = ast;
    }

    fn consume_input(&mut self, raw_input: &str, mut tokens: Vec<token::Token>) -> Vec<token::Token> {
        let mut chars = raw_input.chars().peekable();

        // Iterate over each character in the input
        match chars.clone().peek() {
            Some(c) if c.is_whitespace() => { chars.next(); tokens = self.consume_input(&String::from_iter(chars)[..], tokens); },
            Some(c) if c.is_numeric() => {
                // Grab the number (allowing for possibly decimals)
                let mut number = self.consume_number(&mut chars);
                // Add a numeric token to the list of tokens
                let number = match number.parse() {
                    Ok(val) => {
                        tokens.push(token::Token::DecimalNumber(val));
                    },
                    Err(E) => self.errors.push(format!("FATAL: {}", E))
                };

                tokens = self.consume_input(&String::from_iter(chars)[..], tokens);
            },
            Some(c) if *c == '+' || *c == '-' => {
                // Add the operator and advance the iterator
                self.add_op_and_continue(*c, 2, &mut chars, &mut tokens);
                tokens = self.consume_input(&String::from_iter(chars)[..], tokens);
            },
            Some(c) if *c == '*' || *c == '/' => {
                // Add the operator and advance the iterator
                self.add_op_and_continue(*c, 4, &mut chars, &mut tokens);
                tokens = self.consume_input(&String::from_iter(chars)[..], tokens);
            },
            Some(c) => self.errors.push(format!("Unknown identifier: {}", c)),
            None => ()
        }

        tokens
    }

    // Consumes the iterator until it reaches the end of a number
    fn consume_number(&mut self, it: &mut Peekable<Chars>) -> String {
        let mut chars = vec![];

        // Loop over every character until we reach a non-numeric one
        loop {
            match it.peek() {
                Some(c) if c.is_numeric() || *c == '.' => chars.push(*c),
                Some(c) if !c.is_numeric() => break,
                //Some(c) => println!("Peeking at: {}", c),
                None => break,
                _ => ()
            }
            it.next();
        }

        // Return out number as a String
        chars.into_iter().collect()
    }

    fn add_op_and_continue(&mut self, c: char, precedence: u32, chars: &mut Peekable<Chars>, tokens: &mut Vec<token::Token>) {
        chars.next();
        tokens.push(token::Token::Operator(c, precedence));
    }
}

impl<'a> fmt::Display for Lexer<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {

        let mut result = String::new();
        let ast = self.ast.to_vec();

        for t in ast {
            match t {
                token::Token::Operator(c, p) => result.push(c),
                token::Token::DecimalNumber(n) => result.push_str(&n.to_string()[..]),
                _ => ()
            };
            result.push_str(" "); // Space separated
        }

        write!(f, "{}", result)
    }
}
<commit_msg>Added some comments around the Display trait implementation<commit_after>use std::fmt;
use std::f64;
use std::iter::Peekable;
use std::iter::FromIterator;
use std::str::Chars;

use token;

pub struct Lexer<'a> {
    raw_input: &'a str,
    ast: Vec<token::Token>,
    pub errors: Vec<String>
}

impl<'a> Lexer<'a> {
    pub fn new(input: &str) -> Lexer {
        Lexer {
            raw_input: input,
            ast: Vec::new(),
            errors: vec![]
        }
    }

    pub fn lex(&mut self) {
        let tokens = vec![];

        let ast: Vec<token::Token> = self.consume_input(self.raw_input, tokens);
        self.ast = ast;
    }

    fn consume_input(&mut self, raw_input: &str, mut tokens: Vec<token::Token>) -> Vec<token::Token> {
        let mut chars = raw_input.chars().peekable();

        // Iterate over each character in the input
        match chars.clone().peek() {
            Some(c) if c.is_whitespace() => { chars.next(); tokens = self.consume_input(&String::from_iter(chars)[..], tokens); },
            Some(c) if c.is_numeric() => {
                // Grab the number (allowing for possibly decimals)
                let mut number = self.consume_number(&mut chars);
                // Add a numeric token to the list of tokens
                let number = match number.parse() {
                    Ok(val) => {
                        tokens.push(token::Token::DecimalNumber(val));
                    },
                    Err(E) => self.errors.push(format!("FATAL: {}", E))
                };

                tokens = self.consume_input(&String::from_iter(chars)[..], tokens);
            },
            Some(c) if *c == '+' || *c == '-' => {
                // Add the operator and advance the iterator
                self.add_op_and_continue(*c, 2, &mut chars, &mut tokens);
                tokens = self.consume_input(&String::from_iter(chars)[..], tokens);
            },
            Some(c) if *c == '*' || *c == '/' => {
                // Add the operator and advance the iterator
                self.add_op_and_continue(*c, 4, &mut chars, &mut tokens);
                tokens = self.consume_input(&String::from_iter(chars)[..], tokens);
            },
            Some(c) => self.errors.push(format!("Unknown identifier: {}", c)),
            None => ()
        }

        tokens
    }

    // Consumes the iterator until it reaches the end of a number
    fn consume_number(&mut self, it: &mut Peekable<Chars>) -> String {
        let mut chars = vec![];

        // Loop over every character until we reach a non-numeric one
        loop {
            match it.peek() {
                Some(c) if c.is_numeric() || *c == '.' => chars.push(*c),
                Some(c) if !c.is_numeric() => break,
                //Some(c) => println!("Peeking at: {}", c),
                None => break,
                _ => ()
            }
            it.next();
        }

        // Return out number as a String
        chars.into_iter().collect()
    }

    fn add_op_and_continue(&mut self, c: char, precedence: u32, chars: &mut Peekable<Chars>, tokens: &mut Vec<token::Token>) {
        chars.next();
        tokens.push(token::Token::Operator(c, precedence));
    }
}

// Allows the Lexer AST to be printed as a string
impl<'a> fmt::Display for Lexer<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {

        let mut result = String::new(); // String to output the result to
        let ast = self.ast.to_vec();    // Copy the AST into its own vector so we can consume it

        // Loop over each item in the AST and print a String representation of it
        for t in ast {
            match t {
                token::Token::Operator(c, p) => result.push(c),
                token::Token::DecimalNumber(n) => result.push_str(&n.to_string()[..]),
                _ => ()
            };
            result.push_str(" "); // Space separated
        }

        // Return the result
        write!(f, "{}", result)
    }
}
<|endoftext|>


<commit_before>use std::path;
use std::process;

use clap::{App, Arg, SubCommand};

use config;
use models;
use web;
use utils;

macro_rules! clap_dispatch {
    (
        $matches:expr;
        {
            $( $name:ident ($matches_name:pat $(,$arg_name:ident as $varname:ident),*) => $callback:expr ),*
        }
    ) => {
        match $matches.subcommand_name() {
            $(
                Some(stringify!($name)) => {
                    let matches = $matches.subcommand_matches(stringify!($name)).unwrap();
                    $(
                        let $varname = matches.value_of(stringify!($arg_name)).unwrap();
                    )*
                    let $matches_name = matches;
                    $callback;
                }
            )*
            Some(x) => {
                panic!("Internal error: Command not covered: {}", x);
            },
            None => {
                println!("Subcommand required. See --help for help.");
                process::exit(1);
            }
        }
    }
}

pub fn main() {
    let matches =
        App::new("mysteryshack")
        .version(env!("CARGO_PKG_VERSION"))
        .author("Markus Unterwaditzer & contributors")
        .about("A remoteStorage server.")
        .args_from_usage("-c, --config=[FILE] 'Use specified config file, defaults to ./config'")
        .subcommand(SubCommand::with_name("serve")
                    .about("Start server"))
        .subcommand(SubCommand::with_name("user")
                    .about("User management")
                    .subcommand(SubCommand::with_name("create")
                                .about("Create a new user")
                                .arg(Arg::with_name("USERNAME")
                                     .help("The username")
                                     .required(true)
                                     .index(1)))
                    .subcommand(SubCommand::with_name("delete")
                                .about("Delete a user")
                                .arg(Arg::with_name("USERNAME")
                                     .help("The username")
                                     .required(true)
                                     .index(1))))
        .get_matches();

    let config_path = path::Path::new(matches.value_of("config").unwrap_or("./config"));

    let config = match config::Config::read_file(config_path) {
        Ok(x) => x,
        Err(e) => {
            println!("Failed to parse config: {}", e);
            process::exit(1);
        }
    };

    clap_dispatch!(matches; {
        serve(_) => web::run_server(config),
        user(user_matches) => clap_dispatch!(user_matches; {
            create(_, USERNAME as username) => {
                let password_hash = match models::PasswordHash::from_password(
                    utils::double_prompt("Password for new user: ")) {
                    Ok(x) => x,
                    Err(e) => {
                        println!("Failed to hash password: {}", e);
                        process::exit(1);
                    }
                };

                match models::User::create(&config.data_path, username).map(|user| {
                    user.set_password_hash(password_hash)
                }) {
                    Ok(_) => (),
                    Err(e) => {
                        println!("Failed to create user {}: {}", username, e);
                        process::exit(1);
                    }
                };

                println!("Successfully created user {}", username);
            },
            delete(_, USERNAME as username) => {
                let user = match models::User::get(&config.data_path, username) {
                    Some(x) => x,
                    None => {
                        println!("User does not exist: {}", username);
                        process::exit(1);
                    }
                };

                println!("You are about to delete the user {} and ALL the user's user data. This
                         process is irreversible.", username);
                if !utils::prompt_confirm(format!("Do you want to delete the user {}?", username), false) {
                    println!("Aborted!");
                    process::exit(1);
                }

                match user.delete() {
                    Ok(_) => println!("Successfully deleted user {}.", username),
                    Err(e) => {
                        println!("Failed to delete user: {:?}", e);
                        process::exit(1);
                    }
                };
            }
        })
    });
}
<commit_msg>Fix CLI macro on stable rust<commit_after>use std::path;
use std::process;

use clap::{App, Arg, SubCommand};

use config;
use models;
use web;
use utils;

macro_rules! clap_dispatch {
    (
        $matches:expr;
        // FIXME: $matches_name absolutely requires a trailing comma, couldn't get it to work under
        // stable rust:
        //
        // foo(_) => ()   // Errors
        // foo(_,) => ()  // Works
        {
            $( $name:ident ($matches_name:pat, $($arg_name:ident as $varname:ident),*) => $callback:expr ),*
        }
    ) => {
        match $matches.subcommand_name() {
            $(
                Some(stringify!($name)) => {
                    let matches = $matches.subcommand_matches(stringify!($name)).unwrap();
                    $(
                        let $varname = matches.value_of(stringify!($arg_name)).unwrap();
                    )*
                    let $matches_name = matches;
                    $callback;
                }
            )*
            Some(x) => {
                panic!("Internal error: Command not covered: {}", x);
            },
            None => {
                println!("Subcommand required. See --help for help.");
                process::exit(1);
            }
        }
    }
}

pub fn main() {
    let matches =
        App::new("mysteryshack")
        .version(env!("CARGO_PKG_VERSION"))
        .author("Markus Unterwaditzer & contributors")
        .about("A remoteStorage server.")
        .args_from_usage("-c, --config=[FILE] 'Use specified config file, defaults to ./config'")
        .subcommand(SubCommand::with_name("serve")
                    .about("Start server"))
        .subcommand(SubCommand::with_name("user")
                    .about("User management")
                    .subcommand(SubCommand::with_name("create")
                                .about("Create a new user")
                                .arg(Arg::with_name("USERNAME")
                                     .help("The username")
                                     .required(true)
                                     .index(1)))
                    .subcommand(SubCommand::with_name("delete")
                                .about("Delete a user")
                                .arg(Arg::with_name("USERNAME")
                                     .help("The username")
                                     .required(true)
                                     .index(1))))
        .get_matches();

    let config_path = path::Path::new(matches.value_of("config").unwrap_or("./config"));

    let config = match config::Config::read_file(config_path) {
        Ok(x) => x,
        Err(e) => {
            println!("Failed to parse config: {}", e);
            process::exit(1);
        }
    };

    clap_dispatch!(matches; {
        serve(_,) => web::run_server(config),
        user(user_matches,) => clap_dispatch!(user_matches; {
            create(_, USERNAME as username) => {
                let password_hash = match models::PasswordHash::from_password(
                    utils::double_prompt("Password for new user: ")) {
                    Ok(x) => x,
                    Err(e) => {
                        println!("Failed to hash password: {}", e);
                        process::exit(1);
                    }
                };

                match models::User::create(&config.data_path, username).map(|user| {
                    user.set_password_hash(password_hash)
                }) {
                    Ok(_) => (),
                    Err(e) => {
                        println!("Failed to create user {}: {}", username, e);
                        process::exit(1);
                    }
                };

                println!("Successfully created user {}", username);
            },
            delete(_, USERNAME as username) => {
                let user = match models::User::get(&config.data_path, username) {
                    Some(x) => x,
                    None => {
                        println!("User does not exist: {}", username);
                        process::exit(1);
                    }
                };

                println!("You are about to delete the user {} and ALL the user's user data. This
                         process is irreversible.", username);
                if !utils::prompt_confirm(format!("Do you want to delete the user {}?", username), false) {
                    println!("Aborted!");
                    process::exit(1);
                }

                match user.delete() {
                    Ok(_) => println!("Successfully deleted user {}.", username),
                    Err(e) => {
                        println!("Failed to delete user: {:?}", e);
                        process::exit(1);
                    }
                };
            }
        })
    });
}
<|endoftext|>