<commit_before>// Copyright 2022 Google LLC
//
// 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
//
//    https://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.

//! Module which understands C++ constructor synthesis rules.

#[cfg_attr(test, derive(Eq, PartialEq))]
pub(super) struct ImplicitConstructorsNeeded {
    pub(super) default_constructor: bool,
    pub(super) copy_constructor_taking_t: bool,
    pub(super) copy_constructor_taking_const_t: bool,
    pub(super) move_constructor: bool,
}

#[derive(Default)]
pub(super) struct ExplicitItemsFound {
    pub(super) move_constructor: bool,
    pub(super) copy_constructor: bool,
    pub(super) any_other_constructor: bool,
    pub(super) any_bases_or_fields_lack_const_copy_constructors: bool,
    pub(super) any_bases_or_fields_have_deleted_or_inaccessible_copy_constructors: bool,
    pub(super) destructor: bool,
    pub(super) any_bases_have_deleted_or_inaccessible_destructors: bool,
    pub(super) copy_assignment_operator: bool,
    pub(super) move_assignment_operator: bool,
    pub(super) has_rvalue_reference_fields: bool,
}

pub(super) fn determine_implicit_constructors(
    explicits: ExplicitItemsFound,
) -> ImplicitConstructorsNeeded {
    let any_constructor =
        explicits.copy_constructor || explicits.move_constructor || explicits.any_other_constructor;
    // If no user-declared constructors of any kind are provided for a class type (struct, class, or union),
    // the compiler will always declare a default constructor as an inline public member of its class.
    let default_constructor = !any_constructor;

    // If no user-defined copy constructors are provided for a class type (struct, class, or union),
    // the compiler will always declare a copy constructor as a non-explicit inline public member of its class.
    // This implicitly-declared copy constructor has the form T::T(const T&) if all of the following are true:
    //  each direct and virtual base B of T has a copy constructor whose parameters are const B& or const volatile B&;
    //  each non-static data member M of T of class type or array of class type has a copy constructor whose parameters are const M& or const volatile M&.

    // The implicitly-declared or defaulted copy constructor for class T is defined as deleted if any of the following conditions are true:
    // T is a union-like class and has a variant member with non-trivial copy constructor; // we don't support unions anyway
    // T has a user-defined move constructor or move assignment operator (this condition only causes the implicitly-declared, not the defaulted, copy constructor to be deleted).
    // T has non-static data members that cannot be copied (have deleted, inaccessible, or ambiguous copy constructors);
    // T has direct or virtual base class that cannot be copied (has deleted, inaccessible, or ambiguous copy constructors);
    // T has direct or virtual base class with a deleted or inaccessible destructor;
    // T has a data member of rvalue reference type;
    let copy_constructor_is_deleted = explicits.move_constructor
        || explicits.move_assignment_operator
        || explicits.any_bases_or_fields_have_deleted_or_inaccessible_copy_constructors
        || explicits.any_bases_have_deleted_or_inaccessible_destructors
        || explicits.has_rvalue_reference_fields;

    let (copy_constructor_taking_const_t, copy_constructor_taking_t) =
        if explicits.copy_constructor || copy_constructor_is_deleted {
            (false, false)
        } else if explicits.any_bases_or_fields_lack_const_copy_constructors {
            (false, true)
        } else {
            (true, false)
        };

    // If no user-defined move constructors are provided for a class type (struct, class, or union), and all of the following is true:
    // there are no user-declared copy constructors;
    // there are no user-declared copy assignment operators;
    // there are no user-declared move assignment operators;
    // there is no user-declared destructor.
    // then the compiler will declare a move constructor
    let move_constructor = !(explicits.move_constructor
        || explicits.copy_constructor
        || explicits.destructor
        || explicits.copy_assignment_operator
        || explicits.move_assignment_operator);

    ImplicitConstructorsNeeded {
        default_constructor,
        copy_constructor_taking_t,
        copy_constructor_taking_const_t,
        move_constructor,
    }
}

#[cfg(test)]
mod tests {
    use super::determine_implicit_constructors;

    use super::ExplicitItemsFound;

    #[test]
    fn test_simple() {
        let inputs = ExplicitItemsFound::default();
        let outputs = determine_implicit_constructors(inputs);
        assert!(outputs.default_constructor);
        assert!(outputs.copy_constructor_taking_const_t);
        assert!(!outputs.copy_constructor_taking_t);
        assert!(outputs.move_constructor);
    }

    #[test]
    fn test_with_destructor() {
        let inputs = ExplicitItemsFound {
            destructor: true,
            ..Default::default()
        };
        let outputs = determine_implicit_constructors(inputs);
        assert!(outputs.default_constructor);
        assert!(outputs.copy_constructor_taking_const_t);
        assert!(!outputs.copy_constructor_taking_t);
        assert!(!outputs.move_constructor);
    }

    #[test]
    fn test_with_pesky_base() {
        let inputs = ExplicitItemsFound {
            any_bases_or_fields_lack_const_copy_constructors: true,
            ..Default::default()
        };
        let outputs = determine_implicit_constructors(inputs);
        assert!(outputs.default_constructor);
        assert!(!outputs.copy_constructor_taking_const_t);
        assert!(!outputs.copy_constructor_taking_t);
        assert!(!outputs.move_constructor);
    }

    #[test]
    fn test_with_user_defined_move_constructor() {
        let inputs = ExplicitItemsFound {
            move_constructor: true,
            ..Default::default()
        };
        let outputs = determine_implicit_constructors(inputs);
        assert!(!outputs.default_constructor);
        assert!(!outputs.copy_constructor_taking_const_t);
        assert!(!outputs.copy_constructor_taking_t);
        assert!(!outputs.move_constructor);
    }

    #[test]
    fn test_with_user_defined_misc_constructor() {
        let inputs = ExplicitItemsFound {
            any_other_constructor: true,
            ..Default::default()
        };
        let outputs = determine_implicit_constructors(inputs);
        assert!(!outputs.default_constructor);
        assert!(outputs.copy_constructor_taking_const_t);
        assert!(!outputs.copy_constructor_taking_t);
        assert!(outputs.move_constructor);
    }
}
<commit_msg>correct dumb booboo<commit_after>// Copyright 2022 Google LLC
//
// 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
//
//    https://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.

//! Module which understands C++ constructor synthesis rules.

#[cfg_attr(test, derive(Eq, PartialEq))]
pub(super) struct ImplicitConstructorsNeeded {
    pub(super) default_constructor: bool,
    pub(super) copy_constructor_taking_t: bool,
    pub(super) copy_constructor_taking_const_t: bool,
    pub(super) move_constructor: bool,
}

#[derive(Default)]
pub(super) struct ExplicitItemsFound {
    pub(super) move_constructor: bool,
    pub(super) copy_constructor: bool,
    pub(super) any_other_constructor: bool,
    pub(super) any_bases_or_fields_lack_const_copy_constructors: bool,
    pub(super) any_bases_or_fields_have_deleted_or_inaccessible_copy_constructors: bool,
    pub(super) destructor: bool,
    pub(super) any_bases_have_deleted_or_inaccessible_destructors: bool,
    pub(super) copy_assignment_operator: bool,
    pub(super) move_assignment_operator: bool,
    pub(super) has_rvalue_reference_fields: bool,
}

pub(super) fn determine_implicit_constructors(
    explicits: ExplicitItemsFound,
) -> ImplicitConstructorsNeeded {
    let any_constructor =
        explicits.copy_constructor || explicits.move_constructor || explicits.any_other_constructor;
    // If no user-declared constructors of any kind are provided for a class type (struct, class, or union),
    // the compiler will always declare a default constructor as an inline public member of its class.
    let default_constructor = !any_constructor;

    // If no user-defined copy constructors are provided for a class type (struct, class, or union),
    // the compiler will always declare a copy constructor as a non-explicit inline public member of its class.
    // This implicitly-declared copy constructor has the form T::T(const T&) if all of the following are true:
    //  each direct and virtual base B of T has a copy constructor whose parameters are const B& or const volatile B&;
    //  each non-static data member M of T of class type or array of class type has a copy constructor whose parameters are const M& or const volatile M&.

    // The implicitly-declared or defaulted copy constructor for class T is defined as deleted if any of the following conditions are true:
    // T is a union-like class and has a variant member with non-trivial copy constructor; // we don't support unions anyway
    // T has a user-defined move constructor or move assignment operator (this condition only causes the implicitly-declared, not the defaulted, copy constructor to be deleted).
    // T has non-static data members that cannot be copied (have deleted, inaccessible, or ambiguous copy constructors);
    // T has direct or virtual base class that cannot be copied (has deleted, inaccessible, or ambiguous copy constructors);
    // T has direct or virtual base class with a deleted or inaccessible destructor;
    // T has a data member of rvalue reference type;
    let copy_constructor_is_deleted = explicits.move_constructor
        || explicits.move_assignment_operator
        || explicits.any_bases_or_fields_have_deleted_or_inaccessible_copy_constructors
        || explicits.any_bases_have_deleted_or_inaccessible_destructors
        || explicits.has_rvalue_reference_fields;

    let (copy_constructor_taking_const_t, copy_constructor_taking_t) =
        if explicits.copy_constructor || copy_constructor_is_deleted {
            (false, false)
        } else if explicits.any_bases_or_fields_lack_const_copy_constructors {
            (false, true)
        } else {
            (true, false)
        };

    // If no user-defined move constructors are provided for a class type (struct, class, or union), and all of the following is true:
    // there are no user-declared copy constructors;
    // there are no user-declared copy assignment operators;
    // there are no user-declared move assignment operators;
    // there is no user-declared destructor.
    // then the compiler will declare a move constructor
    let move_constructor = !(explicits.move_constructor
        || explicits.copy_constructor
        || explicits.destructor
        || explicits.copy_assignment_operator
        || explicits.move_assignment_operator);

    ImplicitConstructorsNeeded {
        default_constructor,
        copy_constructor_taking_t,
        copy_constructor_taking_const_t,
        move_constructor,
    }
}

#[cfg(test)]
mod tests {
    use super::determine_implicit_constructors;

    use super::ExplicitItemsFound;

    #[test]
    fn test_simple() {
        let inputs = ExplicitItemsFound::default();
        let outputs = determine_implicit_constructors(inputs);
        assert!(outputs.default_constructor);
        assert!(outputs.copy_constructor_taking_const_t);
        assert!(!outputs.copy_constructor_taking_t);
        assert!(outputs.move_constructor);
    }

    #[test]
    fn test_with_destructor() {
        let inputs = ExplicitItemsFound {
            destructor: true,
            ..Default::default()
        };
        let outputs = determine_implicit_constructors(inputs);
        assert!(outputs.default_constructor);
        assert!(outputs.copy_constructor_taking_const_t);
        assert!(!outputs.copy_constructor_taking_t);
        assert!(!outputs.move_constructor);
    }

    #[test]
    fn test_with_pesky_base() {
        let inputs = ExplicitItemsFound {
            any_bases_or_fields_lack_const_copy_constructors: true,
            ..Default::default()
        };
        let outputs = determine_implicit_constructors(inputs);
        assert!(outputs.default_constructor);
        assert!(!outputs.copy_constructor_taking_const_t);
        assert!(outputs.copy_constructor_taking_t);
        assert!(!outputs.move_constructor);
    }

    #[test]
    fn test_with_user_defined_move_constructor() {
        let inputs = ExplicitItemsFound {
            move_constructor: true,
            ..Default::default()
        };
        let outputs = determine_implicit_constructors(inputs);
        assert!(!outputs.default_constructor);
        assert!(!outputs.copy_constructor_taking_const_t);
        assert!(!outputs.copy_constructor_taking_t);
        assert!(!outputs.move_constructor);
    }

    #[test]
    fn test_with_user_defined_misc_constructor() {
        let inputs = ExplicitItemsFound {
            any_other_constructor: true,
            ..Default::default()
        };
        let outputs = determine_implicit_constructors(inputs);
        assert!(!outputs.default_constructor);
        assert!(outputs.copy_constructor_taking_const_t);
        assert!(!outputs.copy_constructor_taking_t);
        assert!(outputs.move_constructor);
    }
}
<|endoftext|>


<commit_before>//! The `ast` module contains a lot of useful functionality
//! to create and walk through the ast (abstract syntaxtree)

use frontend::lexer::Token;
use backend::zcode::zfile;
use backend::zcode::zfile::{FormattingState};
use std::collections::HashMap;

//==============================
// ast

pub struct AST {
    passages: Vec<ASTNode>
}

 /// add zcode based on tokens
fn gen_zcode<'a>(node: &'a ASTNode, state: FormattingState, mut out: &mut zfile::Zfile, mut manager: &mut CodeGenManager<'a>) {
    let mut state_copy = state.clone();
    let mut set_formatting = false;

    match node {
        &ASTNode::Passage(ref node) => {
            match &node.category {
                &Token::TokPassageName(ref name) => {
                    out.routine(name, 0);
                },
                _ => {
                    debug!("no match 1");
                }
            };

            for child in &node.childs {
                gen_zcode(child, state_copy, out, manager);
            }

            out.op_newline();
            out.op_call_1n("system_check_links");

        },
        &ASTNode::Default(ref t) => {
            match &t.category {
                &Token::TokText(ref s) => {
                    out.gen_print_ops(s);
                },
                &Token::TokNewLine => {
                    out.op_newline();
                },
                &Token::TokFormatBoldStart => {
                    state_copy.bold = true;
                    out.op_set_text_style(state_copy.bold, state_copy.inverted, state_copy.mono, state_copy.italic);
                    set_formatting = true;
                },
                &Token::TokFormatItalicStart => {
                    state_copy.italic = true;
                    out.op_set_text_style(state_copy.bold, state_copy.inverted, state_copy.mono, state_copy.italic);
                    set_formatting = true;
                },
                &Token::TokPassageLink (ref name, ref link) => {
                    out.op_call_2n_with_address("system_add_link", link);

                    out.op_set_text_style(state_copy.bold, true, state_copy.mono, state_copy.italic);
                    let link_text = format!("{}[", name);
                    out.op_print(&link_text);
                    out.op_print_num_var(16);
                    out.op_print("]");
                    out.op_set_text_style(state_copy.bold, state_copy.inverted, state_copy.mono, state_copy.italic);
                    set_formatting = true;
                },
                &Token::TokAssign(ref var, ref operator) => {
                    if operator == "=" || operator == "to" {
                        if !manager.symbol_table.is_known_symbol(var) {
                            manager.symbol_table.insert_new_symbol(&var);
                        }
                        let symbol_id = manager.symbol_table.get_symbol_id(var);
                        if t.childs.len() == 1 {
                            match t.childs[0].as_default().category {
                                Token::TokInt(value) => {
                                    out.op_store_u16(symbol_id, value as u16);
                                },
                                Token::TokBoolean(ref bool_val) => {
                                    out.op_store_u8(symbol_id, boolstr_to_u8(&*bool_val));
                                }
                                _ => { }
                            }
                        } else {
                            debug!("Assign Expression currently not supported.");
                        }
                    }
                },
                &Token::TokIf => {
                    if t.childs.len() < 2 {
                        panic!("Unsupported if-expression!");
                    }

                    let mut compare: u8 = 1;

                    // check if the first node is a pseudonode
                    let pseudo_node = match t.childs[0].as_default().category {
                        Token::TokPseudo => t.childs[0].as_default(),
                        _ =>  panic!("Unsupported if-expression!")
                    };

                    // Check if first token is variable
                    let var_name = match pseudo_node.childs[0].as_default().category {
                        Token::TokVariable(ref var) => var,
                        _ =>  panic!("Unsupported if-expression!")
                    };

                    if pseudo_node.childs.len() > 1 {
                        // Check if second token is compare operator
                        match pseudo_node.childs[1].as_default().category {
                            Token::TokCompOp(ref op) => {
                                match &*(*op) {
                                    "==" | "is" => {} ,
                                    _ => panic!("Unsupported Compare Operator!")
                                }
                            }, _ =>  panic!("Unsupported if-expression!")
                        }

                        // Check if third token is number
                        compare = match pseudo_node.childs[2].as_default().category {
                            Token::TokInt(ref value) => {
                                *value as u8
                            },
                            Token::TokBoolean(ref bool_val) => {
                                boolstr_to_u8(&*bool_val)
                            }, _ => panic!("Unsupported assign value!")
                        };
                    }

                    let symbol_id = manager.symbol_table.get_symbol_id(&*var_name);
                    let if_id = manager.ids_if.start_next();

                    let if_label = &format!("if_{}", if_id);
                    let after_if_label = &format!("after_if_{}", if_id);
                    let after_else_label = &format!("after_else_{}", if_id);
                    out.op_je(symbol_id, compare, if_label);
                    out.op_jump(after_if_label);
                    out.label(if_label);

                    for i in 1..t.childs.len() {
                        gen_zcode(&t.childs[i], state_copy, out, manager)
                    }

                    out.op_jump(after_else_label);
                    out.label(after_if_label);
                },
                &Token::TokElse => {
                    for child in &t.childs {
                        gen_zcode(child, state_copy, out, manager)
                    }
                },
                &Token::TokEndIf => {
                    let after_else_label = &format!("after_else_{}", manager.ids_if.pop_id());
                    out.label(after_else_label);
                },
                _ => {
                    debug!("no match 2");
                }
            };
            if set_formatting {
                for child in &t.childs {
                    gen_zcode(child, state_copy, out, manager);
                }
                out.op_set_text_style(false, false, false, false);
                out.op_set_text_style(state.bold, state.inverted, state.mono, state.italic);
            }
        }
    };

}

fn boolstr_to_u8(string: &str) -> u8 {
    match string {
        "true" => 1 as u8,
        _ => 0 as u8
    }
}

impl AST {
    /// convert ast to zcode
    pub fn to_zcode(& self, out: &mut zfile::Zfile) {
        let mut manager = CodeGenManager::new();
        let state = FormattingState {bold: false, italic: false, mono: false, inverted: false};
        for child in &self.passages {
            gen_zcode(child, state, out, &mut manager);
        }
    }

    pub fn new() -> AST {
        AST {
            passages: Vec::new()
        }
    }

    /// prints the tree
    pub fn print(&self) {
        debug!("Abstract Syntax Tree: ");
        for child in &self.passages {
            child.print(0);
        }
        debug!("");
    }

    /// adds a passage to the path in the ast
    pub fn add_passage(&mut self, token: Token) {
        let node = ASTNode::Passage(NodePassage { category: token, childs: Vec::new() });
        self.passages.push(node);
    }

    /// adds a child to the path in the ast
    pub fn add_child(&mut self, path: &Vec<usize>, token: Token) {
        if let Some(index) = path.first() {
            let mut new_path: Vec<usize> = path.to_vec();
            new_path.remove(0);

            self.passages[*index].add_child(new_path, token)
        } else {
            self.passages.push(ASTNode::Default(NodeDefault { category: token, childs: Vec::new() }));
        }
    }

    /// counts the childs of the path in the asts
    pub fn count_childs(&self, path: Vec<usize>) -> usize {
        if let Some(index) = path.first() {
            let mut new_path: Vec<usize> = path.to_vec();
            new_path.remove(0);

            self.passages[*index].count_childs(new_path)
        } else {
            self.passages.len()
        }
    }
}

// ================================
// node types
enum ASTNode {
    Default (NodeDefault),
    Passage (NodePassage)
}

struct NodePassage {
    category: Token,
    pub childs: Vec<ASTNode>,
    /*tags: Vec<ASTNode>*/
}

struct NodeDefault {
    category: Token,
    childs: Vec<ASTNode>
}

struct CodeGenManager<'a> {
    ids_if: IdentifierProvider,
    symbol_table: SymbolTable<'a>
}

struct IdentifierProvider {
    current_id: u32,
    id_stack: Vec<u32>
}

struct SymbolTable<'a> {
    current_id: u8,
    symbol_map: HashMap<&'a str, u8>
}

impl <'a> CodeGenManager<'a> {
    pub fn new() -> CodeGenManager<'a> {
        CodeGenManager {
            ids_if: IdentifierProvider::new(),
            symbol_table: SymbolTable::new()
        }
    }
}

impl IdentifierProvider {
    pub fn new() -> IdentifierProvider {
        IdentifierProvider {
            current_id: 0,
            id_stack: Vec::new()
        }
    }

    // Returns a new id and pushes it onto the stack
    pub fn start_next(&mut self) -> u32 {
        let id = self.current_id;
        self.current_id += 1;
        self.id_stack.push(id);
        id
    }

    // Pops the last id from the stack
    pub fn pop_id(&mut self) -> u32 {
        self.id_stack.pop().unwrap()
    }
}

impl <'a> SymbolTable<'a> {
    pub fn new() -> SymbolTable<'a> {
        SymbolTable {
            current_id: 25,
            symbol_map: HashMap::<&str, u8>::new()
        }
    }

    // Inserts a symbol into the table, assigning a new id
    pub fn insert_new_symbol(&mut self, symbol: &'a str) {
        debug!("Assigned id {} to variable {}", self.current_id, symbol);
        self.symbol_map.insert(symbol, self.current_id);
        self.current_id += 1;
    }

    // Checks if the symbol is already existent in the table
    pub fn is_known_symbol(&self, symbol: &str) -> bool {
        self.symbol_map.contains_key(symbol)
    }

    // Returns the id for a given symbol
    // (check if is_known_symbol, otherwise panics)
    pub fn get_symbol_id(&self, symbol: &str) -> u8 {
        *self.symbol_map.get(symbol).unwrap()
    }
}

impl ASTNode {
    /// adds an child to the path in the ast
    pub fn add_child(&mut self, path: Vec<usize>, token: Token) {
        if let Some(index) = path.first() {
            let mut new_path: Vec<usize> = path.to_vec();
            new_path.remove(0);

            match self {
                &mut ASTNode::Default(ref mut node) => node.childs[*index].add_child(new_path, token),
                &mut ASTNode::Passage(ref mut node) => node.childs[*index].add_child(new_path, token),
            }
        } else {
            match self {
                &mut ASTNode::Default(ref mut node) => node.childs.push(ASTNode::Default(NodeDefault { category: token, childs: Vec::new() } )),
                &mut ASTNode::Passage(ref mut node) => node.childs.push(ASTNode::Default(NodeDefault { category: token, childs: Vec::new() } )),
            }
        }
    }

    /// counts the childs of the current path in the ast
    pub fn count_childs(&self, path: Vec<usize>) -> usize {
        if let Some(index) = path.first() {
            let mut new_path: Vec<usize> = path.to_vec();
            new_path.remove(0);

            match self {
                &ASTNode::Default(ref node) => node.childs[*index].count_childs(new_path),
                &ASTNode::Passage(ref node) => node.childs[*index].count_childs(new_path),
            }
        } else {
            match self {
                &ASTNode::Default(ref node) => node.childs.len(),
                &ASTNode::Passage(ref node) => node.childs.len(),
            }
        }
    }

    /// prints an node of an ast
    pub fn print(&self, indent: usize) {
        let mut spaces = "".to_string();
        for _ in 0..indent {
            spaces.push_str(" ");
        }

        match self {
            &ASTNode::Passage(ref t) => {
                debug!("{}|- : {:?}", spaces, t.category);
                for child in &t.childs {
                    child.print(indent+2);
                }
            },
            &ASTNode::Default(ref t) => {
                debug!("{}|- : {:?}", spaces, t.category);
                for child in &t.childs {
                    child.print(indent+2);
                }
            }
        }
    }

    pub fn as_default(&self) -> &NodeDefault {
        match self {
            &ASTNode::Default(ref def) => def,
            _ => panic!("Node cannot be unwrapped as NodeDefault!")
        }
    }
}
<commit_msg>Moved formatting state into CodeGenManager<commit_after>//! The `ast` module contains a lot of useful functionality
//! to create and walk through the ast (abstract syntaxtree)

use frontend::lexer::Token;
use backend::zcode::zfile;
use backend::zcode::zfile::{FormattingState};
use std::collections::HashMap;

//==============================
// ast

pub struct AST {
    passages: Vec<ASTNode>
}

 /// add zcode based on tokens
fn gen_zcode<'a>(node: &'a ASTNode, mut out: &mut zfile::Zfile, mut manager: &mut CodeGenManager<'a>) {
    let mut state_copy = manager.format_state.clone();
    let mut set_formatting = false;

    match node {
        &ASTNode::Passage(ref node) => {
            match &node.category {
                &Token::TokPassageName(ref name) => {
                    out.routine(name, 0);
                },
                _ => {
                    debug!("no match 1");
                }
            };

            for child in &node.childs {
                gen_zcode(child, out, manager)
            }

            out.op_newline();
            out.op_call_1n("system_check_links");

        },
        &ASTNode::Default(ref t) => {
            match &t.category {
                &Token::TokText(ref s) => {
                    out.gen_print_ops(s);
                },
                &Token::TokNewLine => {
                    out.op_newline();
                },
                &Token::TokFormatBoldStart => {
                    state_copy.bold = true;
                    out.op_set_text_style(state_copy.bold, state_copy.inverted, state_copy.mono, state_copy.italic);
                    set_formatting = true;
                },
                &Token::TokFormatItalicStart => {
                    state_copy.italic = true;
                    out.op_set_text_style(state_copy.bold, state_copy.inverted, state_copy.mono, state_copy.italic);
                    set_formatting = true;
                },
                &Token::TokPassageLink (ref name, ref link) => {
                    out.op_call_2n_with_address("system_add_link", link);

                    out.op_set_text_style(state_copy.bold, true, state_copy.mono, state_copy.italic);
                    let link_text = format!("{}[", name);
                    out.op_print(&link_text);
                    out.op_print_num_var(16);
                    out.op_print("]");
                    out.op_set_text_style(state_copy.bold, state_copy.inverted, state_copy.mono, state_copy.italic);
                    set_formatting = true;
                },
                &Token::TokAssign(ref var, ref operator) => {
                    if operator == "=" || operator == "to" {
                        if !manager.symbol_table.is_known_symbol(var) {
                            manager.symbol_table.insert_new_symbol(&var);
                        }
                        let symbol_id = manager.symbol_table.get_symbol_id(var);
                        if t.childs.len() == 1 {
                            match t.childs[0].as_default().category {
                                Token::TokInt(value) => {
                                    out.op_store_u16(symbol_id, value as u16);
                                },
                                Token::TokBoolean(ref bool_val) => {
                                    out.op_store_u8(symbol_id, boolstr_to_u8(&*bool_val));
                                }
                                _ => { }
                            }
                        } else {
                            debug!("Assign Expression currently not supported.");
                        }
                    }
                },
                &Token::TokIf => {
                    if t.childs.len() < 2 {
                        panic!("Unsupported if-expression!");
                    }

                    let mut compare: u8 = 1;

                    // check if the first node is a pseudonode
                    let pseudo_node = match t.childs[0].as_default().category {
                        Token::TokPseudo => t.childs[0].as_default(),
                        _ =>  panic!("Unsupported if-expression!")
                    };

                    // Check if first token is variable
                    let var_name = match pseudo_node.childs[0].as_default().category {
                        Token::TokVariable(ref var) => var,
                        _ =>  panic!("Unsupported if-expression!")
                    };

                    if pseudo_node.childs.len() > 1 {
                        // Check if second token is compare operator
                        match pseudo_node.childs[1].as_default().category {
                            Token::TokCompOp(ref op) => {
                                match &*(*op) {
                                    "==" | "is" => {} ,
                                    _ => panic!("Unsupported Compare Operator!")
                                }
                            }, _ =>  panic!("Unsupported if-expression!")
                        }

                        // Check if third token is number
                        compare = match pseudo_node.childs[2].as_default().category {
                            Token::TokInt(ref value) => {
                                *value as u8
                            },
                            Token::TokBoolean(ref bool_val) => {
                                boolstr_to_u8(&*bool_val)
                            }, _ => panic!("Unsupported assign value!")
                        };
                    }

                    let symbol_id = manager.symbol_table.get_symbol_id(&*var_name);
                    let if_id = manager.ids_if.start_next();

                    let if_label = &format!("if_{}", if_id);
                    let after_if_label = &format!("after_if_{}", if_id);
                    let after_else_label = &format!("after_else_{}", if_id);
                    out.op_je(symbol_id, compare, if_label);
                    out.op_jump(after_if_label);
                    out.label(if_label);

                    for i in 1..t.childs.len() {
                        gen_zcode(&t.childs[i], out, manager)
                    }

                    out.op_jump(after_else_label);
                    out.label(after_if_label);
                },
                &Token::TokElse => {
                    for child in &t.childs {
                        gen_zcode(child, out, manager)
                    }
                },
                &Token::TokEndIf => {
                    let after_else_label = &format!("after_else_{}", manager.ids_if.pop_id());
                    out.label(after_else_label);
                },
                _ => {
                    debug!("no match 2");
                }
            };
            if set_formatting {
                for child in &t.childs {
                    gen_zcode(child, out, manager)
                }
                out.op_set_text_style(false, false, false, false);
                let state = manager.format_state;
                out.op_set_text_style(state.bold, state.inverted, state.mono, state.italic);
            }
        }
    };

}

fn boolstr_to_u8(string: &str) -> u8 {
    match string {
        "true" => 1 as u8,
        _ => 0 as u8
    }
}

impl AST {
    /// convert ast to zcode
    pub fn to_zcode(& self, out: &mut zfile::Zfile) {
        let mut manager = CodeGenManager::new();
        for child in &self.passages {
            gen_zcode(child, out, &mut manager);
        }
    }

    pub fn new() -> AST {
        AST {
            passages: Vec::new()
        }
    }

    /// prints the tree
    pub fn print(&self) {
        debug!("Abstract Syntax Tree: ");
        for child in &self.passages {
            child.print(0);
        }
        debug!("");
    }

    /// adds a passage to the path in the ast
    pub fn add_passage(&mut self, token: Token) {
        let node = ASTNode::Passage(NodePassage { category: token, childs: Vec::new() });
        self.passages.push(node);
    }

    /// adds a child to the path in the ast
    pub fn add_child(&mut self, path: &Vec<usize>, token: Token) {
        if let Some(index) = path.first() {
            let mut new_path: Vec<usize> = path.to_vec();
            new_path.remove(0);

            self.passages[*index].add_child(new_path, token)
        } else {
            self.passages.push(ASTNode::Default(NodeDefault { category: token, childs: Vec::new() }));
        }
    }

    /// counts the childs of the path in the asts
    pub fn count_childs(&self, path: Vec<usize>) -> usize {
        if let Some(index) = path.first() {
            let mut new_path: Vec<usize> = path.to_vec();
            new_path.remove(0);

            self.passages[*index].count_childs(new_path)
        } else {
            self.passages.len()
        }
    }
}

// ================================
// node types
enum ASTNode {
    Default (NodeDefault),
    Passage (NodePassage)
}

struct NodePassage {
    category: Token,
    pub childs: Vec<ASTNode>,
    /*tags: Vec<ASTNode>*/
}

struct NodeDefault {
    category: Token,
    childs: Vec<ASTNode>
}

struct CodeGenManager<'a> {
    ids_if: IdentifierProvider,
    symbol_table: SymbolTable<'a>,
    format_state: FormattingState
}

struct IdentifierProvider {
    current_id: u32,
    id_stack: Vec<u32>
}

struct SymbolTable<'a> {
    current_id: u8,
    symbol_map: HashMap<&'a str, u8>
}

impl <'a> CodeGenManager<'a> {
    pub fn new() -> CodeGenManager<'a> {
        CodeGenManager {
            ids_if: IdentifierProvider::new(),
            symbol_table: SymbolTable::new(),
            format_state: FormattingState {bold: false, italic: false, mono: false, inverted: false}
        }
    }
}

impl IdentifierProvider {
    pub fn new() -> IdentifierProvider {
        IdentifierProvider {
            current_id: 0,
            id_stack: Vec::new()
        }
    }

    // Returns a new id and pushes it onto the stack
    pub fn start_next(&mut self) -> u32 {
        let id = self.current_id;
        self.current_id += 1;
        self.id_stack.push(id);
        id
    }

    // Pops the last id from the stack
    pub fn pop_id(&mut self) -> u32 {
        self.id_stack.pop().unwrap()
    }
}

impl <'a> SymbolTable<'a> {
    pub fn new() -> SymbolTable<'a> {
        SymbolTable {
            current_id: 25,
            symbol_map: HashMap::<&str, u8>::new()
        }
    }

    // Inserts a symbol into the table, assigning a new id
    pub fn insert_new_symbol(&mut self, symbol: &'a str) {
        debug!("Assigned id {} to variable {}", self.current_id, symbol);
        self.symbol_map.insert(symbol, self.current_id);
        self.current_id += 1;
    }

    // Checks if the symbol is already existent in the table
    pub fn is_known_symbol(&self, symbol: &str) -> bool {
        self.symbol_map.contains_key(symbol)
    }

    // Returns the id for a given symbol
    // (check if is_known_symbol, otherwise panics)
    pub fn get_symbol_id(&self, symbol: &str) -> u8 {
        *self.symbol_map.get(symbol).unwrap()
    }
}

impl ASTNode {
    /// adds an child to the path in the ast
    pub fn add_child(&mut self, path: Vec<usize>, token: Token) {
        if let Some(index) = path.first() {
            let mut new_path: Vec<usize> = path.to_vec();
            new_path.remove(0);

            match self {
                &mut ASTNode::Default(ref mut node) => node.childs[*index].add_child(new_path, token),
                &mut ASTNode::Passage(ref mut node) => node.childs[*index].add_child(new_path, token),
            }
        } else {
            match self {
                &mut ASTNode::Default(ref mut node) => node.childs.push(ASTNode::Default(NodeDefault { category: token, childs: Vec::new() } )),
                &mut ASTNode::Passage(ref mut node) => node.childs.push(ASTNode::Default(NodeDefault { category: token, childs: Vec::new() } )),
            }
        }
    }

    /// counts the childs of the current path in the ast
    pub fn count_childs(&self, path: Vec<usize>) -> usize {
        if let Some(index) = path.first() {
            let mut new_path: Vec<usize> = path.to_vec();
            new_path.remove(0);

            match self {
                &ASTNode::Default(ref node) => node.childs[*index].count_childs(new_path),
                &ASTNode::Passage(ref node) => node.childs[*index].count_childs(new_path),
            }
        } else {
            match self {
                &ASTNode::Default(ref node) => node.childs.len(),
                &ASTNode::Passage(ref node) => node.childs.len(),
            }
        }
    }

    /// prints an node of an ast
    pub fn print(&self, indent: usize) {
        let mut spaces = "".to_string();
        for _ in 0..indent {
            spaces.push_str(" ");
        }

        match self {
            &ASTNode::Passage(ref t) => {
                debug!("{}|- : {:?}", spaces, t.category);
                for child in &t.childs {
                    child.print(indent+2);
                }
            },
            &ASTNode::Default(ref t) => {
                debug!("{}|- : {:?}", spaces, t.category);
                for child in &t.childs {
                    child.print(indent+2);
                }
            }
        }
    }

    pub fn as_default(&self) -> &NodeDefault {
        match self {
            &ASTNode::Default(ref def) => def,
            _ => panic!("Node cannot be unwrapped as NodeDefault!")
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2017-2019 int08h LLC

// 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.

// for value_t_or_exit!()
#[macro_use]
extern crate clap;

use ring::rand;
use ring::rand::SecureRandom;

use byteorder::{LittleEndian, ReadBytesExt};

use chrono::offset::Utc;
use chrono::TimeZone;

use std::collections::HashMap;
use std::fs::File;
use std::io::Write;
use std::iter::Iterator;
use std::net::{SocketAddr, ToSocketAddrs, UdpSocket};

use clap::{App, Arg};
use roughenough::merkle::root_from_paths;
use roughenough::sign::Verifier;
use roughenough::{
    roughenough_version, RtMessage, Tag, CERTIFICATE_CONTEXT, SIGNED_RESPONSE_CONTEXT,
};

fn create_nonce() -> [u8; 64] {
    let rng = rand::SystemRandom::new();
    let mut nonce = [0u8; 64];
    rng.fill(&mut nonce).unwrap();

    nonce
}

fn make_request(nonce: &[u8]) -> Vec<u8> {
    let mut msg = RtMessage::new(1);
    msg.add_field(Tag::NONC, nonce).unwrap();
    msg.pad_to_kilobyte();

    msg.encode().unwrap()
}

fn receive_response(sock: &mut UdpSocket) -> RtMessage {
    let mut buf = [0; 744];
    let resp_len = sock.recv_from(&mut buf).unwrap().0;

    RtMessage::from_bytes(&buf[0..resp_len]).unwrap()
}

fn stress_test_forever(addr: &SocketAddr) -> ! {
    if !addr.ip().is_loopback() {
        panic!("Cannot use non-loopback address {} for stress testing", addr.ip());
    }

    println!("Stress testing!");

    let nonce = create_nonce();
    let socket = UdpSocket::bind("0.0.0.0:0").expect("Couldn't open UDP socket");
    let request = make_request(&nonce);
    loop {
        socket.send_to(&request, addr).unwrap();
    }
}

struct ResponseHandler {
    pub_key: Option<Vec<u8>>,
    msg: HashMap<Tag, Vec<u8>>,
    srep: HashMap<Tag, Vec<u8>>,
    cert: HashMap<Tag, Vec<u8>>,
    dele: HashMap<Tag, Vec<u8>>,
    nonce: [u8; 64],
}

struct ParsedResponse {
    verified: bool,
    midpoint: u64,
    radius: u32,
}

impl ResponseHandler {
    pub fn new(pub_key: Option<Vec<u8>>, response: RtMessage, nonce: [u8; 64]) -> ResponseHandler {
        let msg = response.into_hash_map();
        let srep = RtMessage::from_bytes(&msg[&Tag::SREP])
            .unwrap()
            .into_hash_map();
        let cert = RtMessage::from_bytes(&msg[&Tag::CERT])
            .unwrap()
            .into_hash_map();
        let dele = RtMessage::from_bytes(&cert[&Tag::DELE])
            .unwrap()
            .into_hash_map();

        ResponseHandler {
            pub_key,
            msg,
            srep,
            cert,
            dele,
            nonce,
        }
    }

    pub fn extract_time(&self) -> ParsedResponse {
        let midpoint = self.srep[&Tag::MIDP]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();
        let radius = self.srep[&Tag::RADI]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();

        let verified = if self.pub_key.is_some() {
            self.validate_dele();
            self.validate_srep();
            self.validate_merkle();
            self.validate_midpoint(midpoint);
            true
        } else {
            false
        };

        ParsedResponse {
            verified,
            midpoint,
            radius,
        }
    }

    fn validate_dele(&self) {
        let mut full_cert = Vec::from(CERTIFICATE_CONTEXT.as_bytes());
        full_cert.extend(&self.cert[&Tag::DELE]);

        assert!(
            self.validate_sig(
                self.pub_key.as_ref().unwrap(),
                &self.cert[&Tag::SIG],
                &full_cert
            ),
            "Invalid signature on DELE tag, response may not be authentic"
        );
    }

    fn validate_srep(&self) {
        let mut full_srep = Vec::from(SIGNED_RESPONSE_CONTEXT.as_bytes());
        full_srep.extend(&self.msg[&Tag::SREP]);

        assert!(
            self.validate_sig(&self.dele[&Tag::PUBK], &self.msg[&Tag::SIG], &full_srep),
            "Invalid signature on SREP tag, response may not be authentic"
        );
    }

    fn validate_merkle(&self) {
        let srep = RtMessage::from_bytes(&self.msg[&Tag::SREP])
            .unwrap()
            .into_hash_map();
        let index = self.msg[&Tag::INDX]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();
        let paths = &self.msg[&Tag::PATH];

        let hash = root_from_paths(index as usize, &self.nonce, paths);

        assert_eq!(
            hash, srep[&Tag::ROOT],
            "Nonce is not present in the response's merkle tree"
        );
    }

    fn validate_midpoint(&self, midpoint: u64) {
        let mint = self.dele[&Tag::MINT]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();
        let maxt = self.dele[&Tag::MAXT]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();

        assert!(
            midpoint >= mint,
            "Response midpoint {} lies *before* delegation span ({}, {})",
            midpoint, mint, maxt
        );
        assert!(
            midpoint <= maxt,
            "Response midpoint {} lies *after* delegation span ({}, {})",
            midpoint, mint, maxt
        );
    }

    fn validate_sig(&self, public_key: &[u8], sig: &[u8], data: &[u8]) -> bool {
        let mut verifier = Verifier::new(public_key);
        verifier.update(data);
        verifier.verify(sig)
    }
}

fn main() {
    let matches = App::new("roughenough client")
    .version(roughenough_version().as_ref())
    .arg(Arg::with_name("host")
      .required(true)
      .help("The Roughtime server to connect to")
      .takes_value(true))
    .arg(Arg::with_name("port")
      .required(true)
      .help("The Roughtime server port to connect to")
      .takes_value(true))
    .arg(Arg::with_name("verbose")
      .short("v")
      .long("verbose")
      .help("Print more output"))
    .arg(Arg::with_name("json")
      .short("j")
      .long("json")
      .help("Print output in JSON"))
    .arg(Arg::with_name("public-key")
      .short("p")
      .long("public-key")
      .takes_value(true)
      .help("The server public key used to validate responses. If unset, no validation will be performed"))
    .arg(Arg::with_name("time-format")
      .short("f")
      .long("time-format")
      .takes_value(true)
      .help("The strftime format string used to print the time recieved from the server")
      .default_value("%b %d %Y %H:%M:%S")
    )
    .arg(Arg::with_name("num-requests")
      .short("n")
      .long("num-requests")
      .takes_value(true)
      .help("The number of requests to make to the server (each from a different source port). This is mainly useful for testing batch response handling")
      .default_value("1")
    )
    .arg(Arg::with_name("stress")
      .short("s")
      .long("stress")
      .help("Stress-tests the server by sending the same request as fast as possible. Please only use this on your own server")
    )
    .arg(Arg::with_name("output")
      .short("o")
      .long("output")
      .takes_value(true)
      .help("Writes all requsts to the specified file, in addition to sending them to the server. Useful for generating fuzer inputs")
    )
    .get_matches();

    let host = matches.value_of("host").unwrap();
    let port = value_t_or_exit!(matches.value_of("port"), u16);
    let verbose = matches.is_present("verbose");
    let json = matches.is_present("json");
    let num_requests = value_t_or_exit!(matches.value_of("num-requests"), u16) as usize;
    let time_format = matches.value_of("time-format").unwrap();
    let stress = matches.is_present("stress");
    let pub_key = matches
        .value_of("public-key")
        .map(|pkey| hex::decode(pkey).expect("Error parsing public key!"));
    let out = matches.value_of("output");

    if verbose {
        println!("Requesting time from: {:?}:{:?}", host, port);
    }

    let addr = (host, port).to_socket_addrs().unwrap().next().unwrap();

    if stress {
        stress_test_forever(&addr)
    }

    let mut requests = Vec::with_capacity(num_requests);
    let mut file = out.map(|o| File::create(o).expect("Failed to create file!"));

    for _ in 0..num_requests {
        let nonce = create_nonce();
        let socket = UdpSocket::bind("0.0.0.0:0").expect("Couldn't open UDP socket");
        let request = make_request(&nonce);

        if let Some(f) = file.as_mut() {
            f.write_all(&request).expect("Failed to write to file!")
        }

        requests.push((nonce, request, socket));
    }

    for &mut (_, ref request, ref mut socket) in &mut requests {
        socket.send_to(request, addr).unwrap();
    }

    for (nonce, _, mut socket) in requests {
        let resp = receive_response(&mut socket);

        let ParsedResponse {
            verified,
            midpoint,
            radius,
        } = ResponseHandler::new(pub_key.clone(), resp.clone(), nonce).extract_time();

        let map = resp.into_hash_map();
        let index = map[&Tag::INDX]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();

        let seconds = midpoint / 10_u64.pow(6);
        let nsecs = (midpoint - (seconds * 10_u64.pow(6))) * 10_u64.pow(3);
        let spec = Utc.timestamp(seconds as i64, nsecs as u32);
        let out = spec.format(time_format).to_string();
        let verify_str = if verified { "Yes" } else { "No" };

        if verbose {
            println!(
                "Received time from server: midpoint={:?}, radius={:?}, verified={} (merkle_index={})",
                out, radius, verify_str, index
            );
        }

        if json {
            println!(
                r#"{{ "midpoint": {:?}, "radius": {:?}, "verified": {}, "merkle_index": {} }}"#,
                out, radius, verified, index
            );
        } else {
            println!("{}", out);
        }
    }
}

<commit_msg>Print Non-Data Output to Standard Error<commit_after>// Copyright 2017-2019 int08h LLC

// 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.

// for value_t_or_exit!()
#[macro_use]
extern crate clap;

use ring::rand;
use ring::rand::SecureRandom;

use byteorder::{LittleEndian, ReadBytesExt};

use chrono::offset::Utc;
use chrono::TimeZone;

use std::collections::HashMap;
use std::fs::File;
use std::io::Write;
use std::iter::Iterator;
use std::net::{SocketAddr, ToSocketAddrs, UdpSocket};

use clap::{App, Arg};
use roughenough::merkle::root_from_paths;
use roughenough::sign::Verifier;
use roughenough::{
    roughenough_version, RtMessage, Tag, CERTIFICATE_CONTEXT, SIGNED_RESPONSE_CONTEXT,
};

fn create_nonce() -> [u8; 64] {
    let rng = rand::SystemRandom::new();
    let mut nonce = [0u8; 64];
    rng.fill(&mut nonce).unwrap();

    nonce
}

fn make_request(nonce: &[u8]) -> Vec<u8> {
    let mut msg = RtMessage::new(1);
    msg.add_field(Tag::NONC, nonce).unwrap();
    msg.pad_to_kilobyte();

    msg.encode().unwrap()
}

fn receive_response(sock: &mut UdpSocket) -> RtMessage {
    let mut buf = [0; 744];
    let resp_len = sock.recv_from(&mut buf).unwrap().0;

    RtMessage::from_bytes(&buf[0..resp_len]).unwrap()
}

fn stress_test_forever(addr: &SocketAddr) -> ! {
    if !addr.ip().is_loopback() {
        panic!("Cannot use non-loopback address {} for stress testing", addr.ip());
    }

    println!("Stress testing!");

    let nonce = create_nonce();
    let socket = UdpSocket::bind("0.0.0.0:0").expect("Couldn't open UDP socket");
    let request = make_request(&nonce);
    loop {
        socket.send_to(&request, addr).unwrap();
    }
}

struct ResponseHandler {
    pub_key: Option<Vec<u8>>,
    msg: HashMap<Tag, Vec<u8>>,
    srep: HashMap<Tag, Vec<u8>>,
    cert: HashMap<Tag, Vec<u8>>,
    dele: HashMap<Tag, Vec<u8>>,
    nonce: [u8; 64],
}

struct ParsedResponse {
    verified: bool,
    midpoint: u64,
    radius: u32,
}

impl ResponseHandler {
    pub fn new(pub_key: Option<Vec<u8>>, response: RtMessage, nonce: [u8; 64]) -> ResponseHandler {
        let msg = response.into_hash_map();
        let srep = RtMessage::from_bytes(&msg[&Tag::SREP])
            .unwrap()
            .into_hash_map();
        let cert = RtMessage::from_bytes(&msg[&Tag::CERT])
            .unwrap()
            .into_hash_map();
        let dele = RtMessage::from_bytes(&cert[&Tag::DELE])
            .unwrap()
            .into_hash_map();

        ResponseHandler {
            pub_key,
            msg,
            srep,
            cert,
            dele,
            nonce,
        }
    }

    pub fn extract_time(&self) -> ParsedResponse {
        let midpoint = self.srep[&Tag::MIDP]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();
        let radius = self.srep[&Tag::RADI]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();

        let verified = if self.pub_key.is_some() {
            self.validate_dele();
            self.validate_srep();
            self.validate_merkle();
            self.validate_midpoint(midpoint);
            true
        } else {
            false
        };

        ParsedResponse {
            verified,
            midpoint,
            radius,
        }
    }

    fn validate_dele(&self) {
        let mut full_cert = Vec::from(CERTIFICATE_CONTEXT.as_bytes());
        full_cert.extend(&self.cert[&Tag::DELE]);

        assert!(
            self.validate_sig(
                self.pub_key.as_ref().unwrap(),
                &self.cert[&Tag::SIG],
                &full_cert
            ),
            "Invalid signature on DELE tag, response may not be authentic"
        );
    }

    fn validate_srep(&self) {
        let mut full_srep = Vec::from(SIGNED_RESPONSE_CONTEXT.as_bytes());
        full_srep.extend(&self.msg[&Tag::SREP]);

        assert!(
            self.validate_sig(&self.dele[&Tag::PUBK], &self.msg[&Tag::SIG], &full_srep),
            "Invalid signature on SREP tag, response may not be authentic"
        );
    }

    fn validate_merkle(&self) {
        let srep = RtMessage::from_bytes(&self.msg[&Tag::SREP])
            .unwrap()
            .into_hash_map();
        let index = self.msg[&Tag::INDX]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();
        let paths = &self.msg[&Tag::PATH];

        let hash = root_from_paths(index as usize, &self.nonce, paths);

        assert_eq!(
            hash, srep[&Tag::ROOT],
            "Nonce is not present in the response's merkle tree"
        );
    }

    fn validate_midpoint(&self, midpoint: u64) {
        let mint = self.dele[&Tag::MINT]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();
        let maxt = self.dele[&Tag::MAXT]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();

        assert!(
            midpoint >= mint,
            "Response midpoint {} lies *before* delegation span ({}, {})",
            midpoint, mint, maxt
        );
        assert!(
            midpoint <= maxt,
            "Response midpoint {} lies *after* delegation span ({}, {})",
            midpoint, mint, maxt
        );
    }

    fn validate_sig(&self, public_key: &[u8], sig: &[u8], data: &[u8]) -> bool {
        let mut verifier = Verifier::new(public_key);
        verifier.update(data);
        verifier.verify(sig)
    }
}

fn main() {
    let matches = App::new("roughenough client")
    .version(roughenough_version().as_ref())
    .arg(Arg::with_name("host")
      .required(true)
      .help("The Roughtime server to connect to")
      .takes_value(true))
    .arg(Arg::with_name("port")
      .required(true)
      .help("The Roughtime server port to connect to")
      .takes_value(true))
    .arg(Arg::with_name("verbose")
      .short("v")
      .long("verbose")
      .help("Print more output"))
    .arg(Arg::with_name("json")
      .short("j")
      .long("json")
      .help("Print output in JSON"))
    .arg(Arg::with_name("public-key")
      .short("p")
      .long("public-key")
      .takes_value(true)
      .help("The server public key used to validate responses. If unset, no validation will be performed"))
    .arg(Arg::with_name("time-format")
      .short("f")
      .long("time-format")
      .takes_value(true)
      .help("The strftime format string used to print the time recieved from the server")
      .default_value("%b %d %Y %H:%M:%S")
    )
    .arg(Arg::with_name("num-requests")
      .short("n")
      .long("num-requests")
      .takes_value(true)
      .help("The number of requests to make to the server (each from a different source port). This is mainly useful for testing batch response handling")
      .default_value("1")
    )
    .arg(Arg::with_name("stress")
      .short("s")
      .long("stress")
      .help("Stress-tests the server by sending the same request as fast as possible. Please only use this on your own server")
    )
    .arg(Arg::with_name("output")
      .short("o")
      .long("output")
      .takes_value(true)
      .help("Writes all requsts to the specified file, in addition to sending them to the server. Useful for generating fuzer inputs")
    )
    .get_matches();

    let host = matches.value_of("host").unwrap();
    let port = value_t_or_exit!(matches.value_of("port"), u16);
    let verbose = matches.is_present("verbose");
    let json = matches.is_present("json");
    let num_requests = value_t_or_exit!(matches.value_of("num-requests"), u16) as usize;
    let time_format = matches.value_of("time-format").unwrap();
    let stress = matches.is_present("stress");
    let pub_key = matches
        .value_of("public-key")
        .map(|pkey| hex::decode(pkey).expect("Error parsing public key!"));
    let out = matches.value_of("output");

    if verbose {
        eprintln!("Requesting time from: {:?}:{:?}", host, port);
    }

    let addr = (host, port).to_socket_addrs().unwrap().next().unwrap();

    if stress {
        stress_test_forever(&addr)
    }

    let mut requests = Vec::with_capacity(num_requests);
    let mut file = out.map(|o| File::create(o).expect("Failed to create file!"));

    for _ in 0..num_requests {
        let nonce = create_nonce();
        let socket = UdpSocket::bind("0.0.0.0:0").expect("Couldn't open UDP socket");
        let request = make_request(&nonce);

        if let Some(f) = file.as_mut() {
            f.write_all(&request).expect("Failed to write to file!")
        }

        requests.push((nonce, request, socket));
    }

    for &mut (_, ref request, ref mut socket) in &mut requests {
        socket.send_to(request, addr).unwrap();
    }

    for (nonce, _, mut socket) in requests {
        let resp = receive_response(&mut socket);

        let ParsedResponse {
            verified,
            midpoint,
            radius,
        } = ResponseHandler::new(pub_key.clone(), resp.clone(), nonce).extract_time();

        let map = resp.into_hash_map();
        let index = map[&Tag::INDX]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();

        let seconds = midpoint / 10_u64.pow(6);
        let nsecs = (midpoint - (seconds * 10_u64.pow(6))) * 10_u64.pow(3);
        let spec = Utc.timestamp(seconds as i64, nsecs as u32);
        let out = spec.format(time_format).to_string();
        let verify_str = if verified { "Yes" } else { "No" };

        if verbose {
            eprintln!(
                "Received time from server: midpoint={:?}, radius={:?}, verified={} (merkle_index={})",
                out, radius, verify_str, index
            );
        }

        if json {
            println!(
                r#"{{ "midpoint": {:?}, "radius": {:?}, "verified": {}, "merkle_index": {} }}"#,
                out, radius, verified, index
            );
        } else {
            println!("{}", out);
        }
    }
}

<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{prelude::*, LintContext};
use guppy::{
    graph::{PackageGraph, PackageMetadata},
    PackageId,
};
use std::path::Path;

/// Represents a linter that runs once per package.
pub trait PackageLinter: Linter {
    fn run<'l>(
        &self,
        ctx: &PackageContext<'l>,
        out: &mut LintFormatter<'l, '_>,
    ) -> Result<RunStatus<'l>>;
}

/// Lint context for an individual package.
#[derive(Copy, Clone, Debug)]
pub struct PackageContext<'l> {
    project_ctx: ProjectContext<'l>,
    workspace_path: &'l Path,
    metadata: &'l PackageMetadata,
}

impl<'l> PackageContext<'l> {
    pub fn new(
        project_ctx: ProjectContext<'l>,
        package_graph: &'l PackageGraph,
        workspace_path: &'l Path,
        id: &PackageId,
    ) -> Self {
        Self {
            project_ctx,
            workspace_path,
            metadata: package_graph.metadata(id).expect("package id is valid"),
        }
    }

    /// Returns the relative path for this package in the workspace.
    pub fn workspace_path(&self) -> &'l Path {
        self.workspace_path
    }

    /// Returns the metadata for this package.
    pub fn metadata(&self) -> &'l PackageMetadata {
        self.metadata
    }
}

impl<'l> LintContext<'l> for PackageContext<'l> {
    fn kind(&self) -> LintKind<'l> {
        LintKind::Package {
            name: self.metadata.name(),
            workspace_path: self.workspace_path,
        }
    }
}
<commit_msg>[x-lint] Add project context accessor<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{prelude::*, LintContext};
use guppy::{
    graph::{PackageGraph, PackageMetadata},
    PackageId,
};
use std::path::Path;

/// Represents a linter that runs once per package.
pub trait PackageLinter: Linter {
    fn run<'l>(
        &self,
        ctx: &PackageContext<'l>,
        out: &mut LintFormatter<'l, '_>,
    ) -> Result<RunStatus<'l>>;
}

/// Lint context for an individual package.
#[derive(Copy, Clone, Debug)]
pub struct PackageContext<'l> {
    project_ctx: ProjectContext<'l>,
    workspace_path: &'l Path,
    metadata: &'l PackageMetadata,
}

impl<'l> PackageContext<'l> {
    pub fn new(
        project_ctx: ProjectContext<'l>,
        package_graph: &'l PackageGraph,
        workspace_path: &'l Path,
        id: &PackageId,
    ) -> Self {
        Self {
            project_ctx,
            workspace_path,
            metadata: package_graph.metadata(id).expect("package id is valid"),
        }
    }

    /// Returns the project context
    pub fn project_ctx(&self) -> &ProjectContext<'l> {
        &self.project_ctx
    }

    /// Returns the relative path for this package in the workspace.
    pub fn workspace_path(&self) -> &'l Path {
        self.workspace_path
    }

    /// Returns the metadata for this package.
    pub fn metadata(&self) -> &'l PackageMetadata {
        self.metadata
    }
}

impl<'l> LintContext<'l> for PackageContext<'l> {
    fn kind(&self) -> LintKind<'l> {
        LintKind::Package {
            name: self.metadata.name(),
            workspace_path: self.workspace_path,
        }
    }
}
<|endoftext|>


<commit_before>extern crate assert_cli;

#[cfg(test)]
mod integration {
    use assert_cli;

    #[test]
    fn mhost_wo_args() {
        mhost()
            .fails()
            .and()
            .stderr().contains("error: The following required arguments were not provided:")
            .unwrap()
    }

    #[test]
    fn mhost_simple_lookup() {
        mhost()
            .with_args(&["www.example.com"])
            .succeeds()
            .and()
            .stdout().contains("* IPv4: 93.184.216.34")
            .and()
            .stdout().contains("* IPv6: 2606:2800:220:1:248:1893:25c8:1946")
            .unwrap()
    }

    #[test]
    fn mhost_simple_ipv4_reverse_lookup() {
        mhost()
            .with_args(&["8.8.8.8"])
            .succeeds()
            .and()
            .stdout().contains("* PTR: google-public-dns-a.google.com.")
            .unwrap()
    }

    #[test]
    fn mhost_simple_ipv6_reverse_lookup() {
        mhost()
            .with_args(&["2001:4860:4860::8888"])
            .succeeds()
            .and()
            .stdout().contains("* PTR: google-public-dns-a.google.com.")
            .unwrap()
    }

    fn mhost() -> assert_cli::Assert {
        ::std::env::set_var("RUST_BACKTRACE", "1");
        assert_cli::Assert::command(&["./target/debug/mhost"])
    }
}
<commit_msg>Distinguish cargo release type for integration tests<commit_after>extern crate assert_cli;

#[cfg(test)]
mod integration {
    use assert_cli;

    #[test]
    fn mhost_wo_args() {
        mhost()
            .fails()
            .and()
            .stderr().contains("error: The following required arguments were not provided:")
            .unwrap()
    }

    #[test]
    fn mhost_simple_lookup() {
        mhost()
            .with_args(&["www.example.com"])
            .succeeds()
            .and()
            .stdout().contains("* IPv4: 93.184.216.34")
            .and()
            .stdout().contains("* IPv6: 2606:2800:220:1:248:1893:25c8:1946")
            .unwrap()
    }

    #[test]
    fn mhost_simple_ipv4_reverse_lookup() {
        mhost()
            .with_args(&["8.8.8.8"])
            .succeeds()
            .and()
            .stdout().contains("* PTR: google-public-dns-a.google.com.")
            .unwrap()
    }

    #[test]
    fn mhost_simple_ipv6_reverse_lookup() {
        mhost()
            .with_args(&["2001:4860:4860::8888"])
            .succeeds()
            .and()
            .stdout().contains("* PTR: google-public-dns-a.google.com.")
            .unwrap()
    }

    // Since local development uses `cargo test` and Travis CI uses `cargo test --release` we need to distinguish,
    // which binary to call -- `#[cfg(debug_assertions)]` to the rescue.

    #[cfg(debug_assertions)]
    fn mhost() -> assert_cli::Assert {
        ::std::env::set_var("RUST_BACKTRACE", "1");
        assert_cli::Assert::command(&["./target/debug/mhost"])
    }

    #[cfg(not(debug_assertions))]
    fn mhost() -> assert_cli::Assert {
        ::std::env::set_var("RUST_BACKTRACE", "1");
        assert_cli::Assert::command(&["./target/release/mhost"])
    }
}
<|endoftext|>


<commit_before>// We add this `extern crate` here to ensure that bindgen is up-to-date and
// rebuilt, even though we aren't using any of its types or functions here, only
// indirectly calling the executable.
#[allow(dead_code)]
extern crate bindgen;

use std::env;
use std::fs;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::process;

const TEST_BATCH_DEFAULT_SIZE: usize = 16;

fn spawn_run_bindgen<P, Q, R>(run_bindgen: P, bindgen: Q, header: R) -> process::Child
    where P: AsRef<Path>,
          Q: AsRef<Path>,
          R: AsRef<Path>
{
    let run_bindgen = run_bindgen.as_ref();
    let bindgen = bindgen.as_ref();
    let header = header.as_ref();

    // Convert from "tests/headers/foo.hpp" to "tests/expectations/foo.rs" by
    // saving the filename, popping off "headers/foo.hpp", pushing
    // "expectations", pushing the saved filename, and finally modifying the
    // extension.

    let mut expected = PathBuf::from(header);
    let file_name = expected.file_name()
        .expect("Should have filename")
        .to_os_string();
    expected.pop();
    expected.pop();
    expected.push("expectations");
    expected.push(file_name);
    expected.set_extension("rs");

    let mut cmd = process::Command::new(run_bindgen);
    cmd.stdout(process::Stdio::piped())
        .stderr(process::Stdio::piped())
        .arg(bindgen)
        .arg(header)
        .arg(expected);

    if cfg!(feature = "llvm_stable") {
        cmd.arg("--feature")
            .arg("llvm_stable");
    }

    cmd.spawn()
        .expect("Should be able to spawn run-bindgen.py child process")
}

#[test]
fn run_bindgen_tests() {
    let crate_root = env::var("CARGO_MANIFEST_DIR")
        .expect("should have CARGO_MANIFEST_DIR environment variable");

    let mut run_bindgen = PathBuf::from(&crate_root);
    run_bindgen.push("tests");
    run_bindgen.push("tools");
    run_bindgen.push("run-bindgen.py");

    let mut bindgen = PathBuf::from(&crate_root);
    bindgen.push("target");
    bindgen.push("debug");
    bindgen.push("bindgen");
    if !bindgen.is_file() {
        panic!("{} is not a file! Build bindgen before running tests.",
               bindgen.display());
    }

    let mut headers_dir = PathBuf::from(&crate_root);
    headers_dir.push("tests");
    headers_dir.push("headers");

    let entries = fs::read_dir(&headers_dir)
        .expect("Should read directory")
        .map(|result| result.expect("Should read directory entry"));

    let tests = entries.filter(|entry| {
        match entry.path().extension().map(|s| s.to_str()) {
            Some(Some("h")) |
            Some(Some("hpp")) => true,
            _ => false,
        }
    }).collect::<Vec<_>>();

    let batch_size = env::var("BINDGEN_TEST_BATCH_SIZE")
        .ok()
        .and_then(|x| x.parse::<usize>().ok())
        .unwrap_or(TEST_BATCH_DEFAULT_SIZE);

    // Spawn batch_size child to run in parallel
    // and wait on all of them before processing the next batch

    let children = tests.chunks(batch_size).map(|x| {
        x.iter().map(|entry| {
            let child = spawn_run_bindgen(run_bindgen.clone(), bindgen.clone(), entry.path());
            (entry.path(), child)
        }).collect::<Vec<_>>()
    });

    let failures: Vec<_> = children.flat_map(|x| {
        x.into_iter().filter_map(|(path, mut child)| {
            let passed = child.wait()
                .expect("Should wait on child process")
                .success();

            if passed { None } else { Some((path, child)) }
        })
    })
    .collect();

    let num_failures = failures.len();

    for (path, child) in failures {
        println!("FAIL: {}", path.display());

        let mut buf = String::new();

        child.stdout
            .expect("should have stdout piped")
            .read_to_string(&mut buf)
            .expect("should read child's stdout");
        for line in buf.lines() {
            println!("child stdout> {}", line);
        }

        child.stderr
            .expect("should have stderr piped")
            .read_to_string(&mut buf)
            .expect("should read child's stderr");
        for line in buf.lines() {
            println!("child stderr> {}", line);
        }
    }

    if num_failures > 0 {
        panic!("{} test failures!", num_failures);
    }
}
<commit_msg>Auto merge of #174 - fitzgen:small-typos, r=emilio<commit_after>// We add this `extern crate` here to ensure that bindgen is up-to-date and
// rebuilt, even though we aren't using any of its types or functions here, only
// indirectly calling the executable.
#[allow(dead_code)]
extern crate bindgen;

use std::env;
use std::fs;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::process;

const TEST_BATCH_DEFAULT_SIZE: usize = 16;

fn spawn_run_bindgen<P, Q, R>(run_bindgen: P, bindgen: Q, header: R) -> process::Child
    where P: AsRef<Path>,
          Q: AsRef<Path>,
          R: AsRef<Path>
{
    let run_bindgen = run_bindgen.as_ref();
    let bindgen = bindgen.as_ref();
    let header = header.as_ref();

    // Convert from "tests/headers/foo.hpp" to "tests/expectations/foo.rs" by
    // saving the filename, popping off "headers/foo.hpp", pushing
    // "expectations", pushing the saved filename, and finally modifying the
    // extension.

    let mut expected = PathBuf::from(header);
    let file_name = expected.file_name()
        .expect("Should have filename")
        .to_os_string();
    expected.pop();
    expected.pop();
    expected.push("expectations");
    expected.push(file_name);
    expected.set_extension("rs");

    let mut cmd = process::Command::new(run_bindgen);
    cmd.stdout(process::Stdio::piped())
        .stderr(process::Stdio::piped())
        .arg(bindgen)
        .arg(header)
        .arg(expected);

    if cfg!(feature = "llvm_stable") {
        cmd.arg("--feature")
            .arg("llvm_stable");
    }

    cmd.spawn()
        .expect("Should be able to spawn run-bindgen.py child process")
}

#[test]
fn run_bindgen_tests() {
    let crate_root = env::var("CARGO_MANIFEST_DIR")
        .expect("should have CARGO_MANIFEST_DIR environment variable");

    let mut run_bindgen = PathBuf::from(&crate_root);
    run_bindgen.push("tests");
    run_bindgen.push("tools");
    run_bindgen.push("run-bindgen.py");

    let mut bindgen = PathBuf::from(&crate_root);
    bindgen.push("target");
    bindgen.push("debug");
    bindgen.push("bindgen");
    if !bindgen.is_file() {
        panic!("{} is not a file! Build bindgen before running tests.",
               bindgen.display());
    }

    let mut headers_dir = PathBuf::from(&crate_root);
    headers_dir.push("tests");
    headers_dir.push("headers");

    let entries = fs::read_dir(&headers_dir)
        .expect("Should read directory")
        .map(|result| result.expect("Should read directory entry"));

    let tests = entries.filter(|entry| {
        match entry.path().extension().map(|s| s.to_str()) {
            Some(Some("h")) |
            Some(Some("hpp")) => true,
            _ => false,
        }
    }).collect::<Vec<_>>();

    let batch_size = env::var("BINDGEN_TEST_BATCH_SIZE")
        .ok()
        .and_then(|x| x.parse::<usize>().ok())
        .unwrap_or(TEST_BATCH_DEFAULT_SIZE);

    // Spawn `batch_size` children to run in parallel and wait on all of them
    // before processing the next batch. This puts a limit on the resources
    // consumed when testing, so that we don't overload the system.

    let children = tests.chunks(batch_size).map(|x| {
        x.iter().map(|entry| {
            let child = spawn_run_bindgen(run_bindgen.clone(), bindgen.clone(), entry.path());
            (entry.path(), child)
        }).collect::<Vec<_>>()
    });

    let failures: Vec<_> = children.flat_map(|x| {
        x.into_iter().filter_map(|(path, mut child)| {
            let passed = child.wait()
                .expect("Should wait on child process")
                .success();

            if passed { None } else { Some((path, child)) }
        })
    })
    .collect();

    let num_failures = failures.len();

    for (path, child) in failures {
        println!("FAIL: {}", path.display());

        let mut buf = String::new();

        child.stdout
            .expect("should have stdout piped")
            .read_to_string(&mut buf)
            .expect("should read child's stdout");
        for line in buf.lines() {
            println!("child stdout> {}", line);
        }

        child.stderr
            .expect("should have stderr piped")
            .read_to_string(&mut buf)
            .expect("should read child's stderr");
        for line in buf.lines() {
            println!("child stderr> {}", line);
        }
    }

    if num_failures > 0 {
        panic!("{} test failures!", num_failures);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use anyhow::Result;

use bytecode::function_target::FunctionTarget;
use move_model::ast::Spec;

use crate::workflow::WorkflowOptions;

pub(crate) fn inline_all_exp_in_spec(
    _options: &WorkflowOptions,
    _target: FunctionTarget,
    spec: Spec,
) -> Result<Spec> {
    Ok(spec)
}
<commit_msg>[move-prover][spec-flatten] unroll function call, lambda, and block exp<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use anyhow::Result;
use std::collections::BTreeMap;

use bytecode::function_target::FunctionTarget;
use move_model::{
    ast::{Condition, ConditionKind, Exp, ExpData, Operation, Spec, TempIndex},
    model::GlobalEnv,
    symbol::Symbol,
};

use crate::workflow::WorkflowOptions;

pub(crate) fn inline_all_exp_in_spec(
    _options: &WorkflowOptions,
    target: FunctionTarget,
    spec: Spec,
) -> Result<Spec> {
    let env = target.global_env();
    let inliner = ExpInliner { env };

    let Spec {
        loc,
        conditions,
        properties,
        on_impl,
    } = spec;

    let mut local_vars = BTreeMap::new();
    let mut new_conditions = vec![];
    for cond in conditions {
        let Condition {
            loc,
            kind,
            properties,
            exp,
            additional_exps,
        } = cond;

        match &kind {
            ConditionKind::LetPre(sym) | ConditionKind::LetPost(sym) => {
                let var_exp = inliner.inline_exp(&exp, None, Some(&local_vars));
                local_vars.insert(*sym, var_exp);
            }
            _ => {
                let new_exp = inliner.inline_exp(&exp, None, Some(&local_vars));
                let new_additional_exps = additional_exps
                    .into_iter()
                    .map(|e| inliner.inline_exp(&e, None, Some(&local_vars)))
                    .collect();
                let new_cond = Condition {
                    loc,
                    kind,
                    properties,
                    exp: new_exp,
                    additional_exps: new_additional_exps,
                };
                new_conditions.push(new_cond);
            }
        }
    }

    let new_spec = Spec {
        loc,
        conditions: new_conditions,
        properties,
        on_impl,
    };
    Ok(new_spec)
}

struct ExpInliner<'env> {
    env: &'env GlobalEnv,
}

impl ExpInliner<'_> {
    fn inline_exp(
        &self,
        exp: &Exp,
        temp_var_repl: Option<&BTreeMap<TempIndex, Exp>>,
        local_var_repl: Option<&BTreeMap<Symbol, Exp>>,
    ) -> Exp {
        use Operation::*;

        let mut rewriter = |e: Exp| match e.as_ref() {
            ExpData::LocalVar(_, sym) => match local_var_repl {
                None => Err(e),
                Some(var_map) => Ok(var_map.get(sym).unwrap().clone()),
            },
            ExpData::Temporary(_, idx) => match temp_var_repl {
                None => Err(e),
                Some(var_map) => Ok(var_map.get(idx).unwrap().clone()),
            },
            ExpData::Call(node_id, Function(mid, fid, _), args) => {
                let callee_menv = self.env.get_module(*mid);
                let callee_decl = callee_menv.get_spec_fun(*fid);
                debug_assert_eq!(args.len(), callee_decl.params.len());
                if callee_decl.is_native || callee_decl.uninterpreted || callee_decl.body.is_none()
                {
                    Err(e)
                } else {
                    let mut callee_local_vars =
                        local_var_repl.cloned().unwrap_or_else(BTreeMap::new);
                    for (arg_exp, (sym, _)) in args
                        .iter()
                        .map(|e| self.inline_exp(e, temp_var_repl, local_var_repl))
                        .zip(callee_decl.params.iter())
                    {
                        callee_local_vars.insert(*sym, arg_exp);
                    }

                    let callee_targs = self.env.get_node_instantiation(*node_id);
                    let callee_body = ExpData::rewrite_node_id(
                        callee_decl.body.as_ref().unwrap().clone(),
                        &mut |id| ExpData::instantiate_node(self.env, id, &callee_targs),
                    );
                    Ok(self.inline_exp(&callee_body, temp_var_repl, Some(&callee_local_vars)))
                }
            }
            ExpData::Invoke(_, lambda, args) => match lambda.as_ref() {
                ExpData::Lambda(_, locals, body) => {
                    debug_assert_eq!(args.len(), locals.len());
                    let mut lambda_local_vars =
                        local_var_repl.cloned().unwrap_or_else(BTreeMap::new);
                    for (arg_exp, decl) in args
                        .iter()
                        .map(|e| self.inline_exp(e, temp_var_repl, local_var_repl))
                        .zip(locals)
                    {
                        lambda_local_vars.insert(decl.name, arg_exp);
                    }
                    Ok(self.inline_exp(body, temp_var_repl, Some(&lambda_local_vars)))
                }
                _ => Err(e),
            },
            ExpData::Block(_, var_decls, body) => {
                let mut block_local_vars = local_var_repl.cloned().unwrap_or_else(BTreeMap::new);
                for var_decl in var_decls {
                    let var_exp = self.inline_exp(
                        var_decl.binding.as_ref().unwrap(),
                        temp_var_repl,
                        Some(&block_local_vars),
                    );
                    block_local_vars.insert(var_decl.name, var_exp);
                }
                Ok(self.inline_exp(body, temp_var_repl, Some(&block_local_vars)))
            }
            _ => Err(e),
        };
        ExpData::rewrite(exp.clone(), &mut rewriter)
    }
}
<|endoftext|>