<commit_before>// Copyright (c) 2016 The vulkano developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use std::mem;

use parse;
use enums;

/// Translates all the structs that are contained in the SPIR-V document as Rust structs.
pub fn write_structs(doc: &parse::Spirv) -> String {
    let mut result = String::new();

    for instruction in &doc.instructions {
        match *instruction {
            parse::Instruction::TypeStruct { result_id, ref member_types } => {
                result.push_str(&write_struct(doc, result_id, member_types));
                result.push_str("\n");
            },
            _ => ()
        }
    }

    result
}

/// Writes a single struct.
fn write_struct(doc: &parse::Spirv, struct_id: u32, members: &[u32]) -> String {
    let name = ::name_from_id(doc, struct_id);

    let mut members_defs = Vec::with_capacity(members.len());

    // Contains the offset of the next field.
    // Equals to `None` if there's a runtime-sized field in there.
    let mut current_rust_offset = Some(0);

    for (num, &member) in members.iter().enumerate() {
        let (ty, rust_size, rust_align) = type_from_id(doc, member);
        let member_name = ::member_name_from_id(doc, struct_id, num as u32);

        // Ignore the whole struct is a member is built in, which includes
        // `gl_Position` for example.
        if is_builtin_member(doc, struct_id, num as u32) {
            return String::new();
        }

        let spirv_offset = doc.instructions.iter().filter_map(|i| {
            match *i {
                parse::Instruction::MemberDecorate { target_id, member,
                                                   decoration: enums::Decoration::DecorationOffset,
                                                   ref params } if target_id == struct_id &&
                                                                   member as usize == num =>
                {
                    return Some(params[0]);
                },
                _ => ()
            };

            None
        }).next().expect(&format!("Struct member `{}` of `{}` is missing an `Offset` decoration",
                                  member_name, name)) as usize;

        // We need to add a dummy field if necessary.
        {
            let current_rust_offset = current_rust_offset.as_mut().expect("Found runtime-sized member in non-final position");

            // Updating current_rust_offset to take the alignment of the next field into account
            *current_rust_offset = if *current_rust_offset == 0 {
                0
            } else {
                (1 + (*current_rust_offset - 1) / rust_align) * rust_align
            };

            if spirv_offset != *current_rust_offset {
                let diff = spirv_offset.checked_sub(*current_rust_offset).unwrap();
                members_defs.push(format!("_dummy: [u8; {}]", diff));       // FIXME: fix name if there are multiple dummies
                *current_rust_offset += diff;
            }
        }

        // Updating `current_rust_offset`.
        if let Some(s) = rust_size {
            *current_rust_offset.as_mut().unwrap() += s;
        } else {
            current_rust_offset = None;
        }

        members_defs.push(format!("pub {name}: {ty}", name = member_name, ty = ty));
    }

    // We can only derive common traits if there's no unsized member in the struct.
    let derive = if current_rust_offset.is_some() {
        "#[derive(Copy, Clone, Debug, Default)]\n"
    } else {
        ""
    };

    format!("#[repr(C)]\n{derive}\
             pub struct {name} {{\n\t{members}\n}}\n",
            derive = derive, name = name, members = members_defs.join(",\n\t"))
}

/// Returns true if a `BuiltIn` decorator is applied on a struct member.
fn is_builtin_member(doc: &parse::Spirv, id: u32, member_id: u32) -> bool {
    for instruction in &doc.instructions {
        match *instruction {
            parse::Instruction::MemberDecorate { target_id, member,
                                                 decoration: enums::Decoration::DecorationBuiltIn,
                                                 .. } if target_id == id && member == member_id =>
            {
                return true;
            },
            _ => ()
        }
    }

    false
}

/// Returns the type name to put in the Rust struct, and its size and alignment.
///
/// The size can be `None` if it's only known at runtime.
fn type_from_id(doc: &parse::Spirv, searched: u32) -> (String, Option<usize>, usize) {
    for instruction in doc.instructions.iter() {
        match instruction {
            &parse::Instruction::TypeBool { result_id } if result_id == searched => {
                return ("bool".to_owned(), Some(mem::size_of::<bool>()), mem::align_of::<bool>())
            },
            &parse::Instruction::TypeInt { result_id, width, signedness } if result_id == searched => {
                // FIXME: width
                return ("i32".to_owned(), Some(mem::size_of::<i32>()), mem::align_of::<i32>())
            },
            &parse::Instruction::TypeFloat { result_id, width } if result_id == searched => {
                // FIXME: width
                return ("f32".to_owned(), Some(mem::size_of::<f32>()), mem::align_of::<f32>())
            },
            &parse::Instruction::TypeVector { result_id, component_id, count } if result_id == searched => {
                debug_assert_eq!(mem::align_of::<[u32; 3]>(), mem::align_of::<u32>());
                let (t, t_size, t_align) = type_from_id(doc, component_id);
                return (format!("[{}; {}]", t, count), t_size.map(|s| s * count as usize), t_align);
            },
            &parse::Instruction::TypeMatrix { result_id, column_type_id, column_count } if result_id == searched => {
                // FIXME: row-major or column-major
                debug_assert_eq!(mem::align_of::<[u32; 3]>(), mem::align_of::<u32>());
                let (t, t_size, t_align) = type_from_id(doc, column_type_id);
                return (format!("[{}; {}]", t, column_count), t_size.map(|s| s * column_count as usize), t_align);
            },
            &parse::Instruction::TypeArray { result_id, type_id, length_id } if result_id == searched => {
                debug_assert_eq!(mem::align_of::<[u32; 3]>(), mem::align_of::<u32>());
                let (t, t_size, t_align) = type_from_id(doc, type_id);
                let len = doc.instructions.iter().filter_map(|e| {
                    match e { &parse::Instruction::Constant { result_id, ref data, .. } if result_id == length_id => Some(data.clone()), _ => None }
                }).next().expect("failed to find array length");
                let len = len.iter().rev().fold(0u64, |a, &b| (a << 32) | b as u64);
                return (format!("[{}; {}]", t, len), t_size.map(|s| s * len as usize), t_align);       // FIXME:
            },
            &parse::Instruction::TypeRuntimeArray { result_id, type_id } if result_id == searched => {
                debug_assert_eq!(mem::align_of::<[u32; 3]>(), mem::align_of::<u32>());
                let (t, _, t_align) = type_from_id(doc, type_id);
                return (format!("[{}]", t), None, t_align);
            },
            &parse::Instruction::TypeStruct { result_id, ref member_types } if result_id == searched => {
                let name = ::name_from_id(doc, result_id);
                let size = member_types.iter().filter_map(|&t| type_from_id(doc, t).1).fold(0, |a,b|a+b);
                let align = member_types.iter().map(|&t| type_from_id(doc, t).2).max().unwrap_or(1);
                return (name, Some(size), align);
            },
            _ => ()
        }
    }

    panic!("Type #{} not found", searched)
}
<commit_msg>Ignore local structs that don't have Offset decorations<commit_after>// Copyright (c) 2016 The vulkano developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use std::mem;

use parse;
use enums;

/// Translates all the structs that are contained in the SPIR-V document as Rust structs.
pub fn write_structs(doc: &parse::Spirv) -> String {
    let mut result = String::new();

    for instruction in &doc.instructions {
        match *instruction {
            parse::Instruction::TypeStruct { result_id, ref member_types } => {
                result.push_str(&write_struct(doc, result_id, member_types));
                result.push_str("\n");
            },
            _ => ()
        }
    }

    result
}

/// Writes a single struct.
fn write_struct(doc: &parse::Spirv, struct_id: u32, members: &[u32]) -> String {
    let name = ::name_from_id(doc, struct_id);

    let mut members_defs = Vec::with_capacity(members.len());

    // Contains the offset of the next field.
    // Equals to `None` if there's a runtime-sized field in there.
    let mut current_rust_offset = Some(0);

    for (num, &member) in members.iter().enumerate() {
        let (ty, rust_size, rust_align) = type_from_id(doc, member);
        let member_name = ::member_name_from_id(doc, struct_id, num as u32);

        // Ignore the whole struct is a member is built in, which includes
        // `gl_Position` for example.
        if is_builtin_member(doc, struct_id, num as u32) {
            return String::new();
        }

        let spirv_offset = doc.instructions.iter().filter_map(|i| {
            match *i {
                parse::Instruction::MemberDecorate { target_id, member,
                                                   decoration: enums::Decoration::DecorationOffset,
                                                   ref params } if target_id == struct_id &&
                                                                   member as usize == num =>
                {
                    return Some(params[0]);
                },
                _ => ()
            };

            None
        }).next();

        // Some structs don't have `Offset` decorations, in the case they are used as local
        // variables only. Ignoring these.
        let spirv_offset = match spirv_offset {
            Some(o) => o as usize,
            None => return String::new()
        };

        // We need to add a dummy field if necessary.
        {
            let current_rust_offset = current_rust_offset.as_mut().expect("Found runtime-sized member in non-final position");

            // Updating current_rust_offset to take the alignment of the next field into account
            *current_rust_offset = if *current_rust_offset == 0 {
                0
            } else {
                (1 + (*current_rust_offset - 1) / rust_align) * rust_align
            };

            if spirv_offset != *current_rust_offset {
                let diff = spirv_offset.checked_sub(*current_rust_offset).unwrap();
                members_defs.push(format!("_dummy: [u8; {}]", diff));       // FIXME: fix name if there are multiple dummies
                *current_rust_offset += diff;
            }
        }

        // Updating `current_rust_offset`.
        if let Some(s) = rust_size {
            *current_rust_offset.as_mut().unwrap() += s;
        } else {
            current_rust_offset = None;
        }

        members_defs.push(format!("pub {name}: {ty}", name = member_name, ty = ty));
    }

    // We can only derive common traits if there's no unsized member in the struct.
    let derive = if current_rust_offset.is_some() {
        "#[derive(Copy, Clone, Debug, Default)]\n"
    } else {
        ""
    };

    format!("#[repr(C)]\n{derive}\
             pub struct {name} {{\n\t{members}\n}}\n",
            derive = derive, name = name, members = members_defs.join(",\n\t"))
}

/// Returns true if a `BuiltIn` decorator is applied on a struct member.
fn is_builtin_member(doc: &parse::Spirv, id: u32, member_id: u32) -> bool {
    for instruction in &doc.instructions {
        match *instruction {
            parse::Instruction::MemberDecorate { target_id, member,
                                                 decoration: enums::Decoration::DecorationBuiltIn,
                                                 .. } if target_id == id && member == member_id =>
            {
                return true;
            },
            _ => ()
        }
    }

    false
}

/// Returns the type name to put in the Rust struct, and its size and alignment.
///
/// The size can be `None` if it's only known at runtime.
fn type_from_id(doc: &parse::Spirv, searched: u32) -> (String, Option<usize>, usize) {
    for instruction in doc.instructions.iter() {
        match instruction {
            &parse::Instruction::TypeBool { result_id } if result_id == searched => {
                return ("bool".to_owned(), Some(mem::size_of::<bool>()), mem::align_of::<bool>())
            },
            &parse::Instruction::TypeInt { result_id, width, signedness } if result_id == searched => {
                // FIXME: width
                return ("i32".to_owned(), Some(mem::size_of::<i32>()), mem::align_of::<i32>())
            },
            &parse::Instruction::TypeFloat { result_id, width } if result_id == searched => {
                // FIXME: width
                return ("f32".to_owned(), Some(mem::size_of::<f32>()), mem::align_of::<f32>())
            },
            &parse::Instruction::TypeVector { result_id, component_id, count } if result_id == searched => {
                debug_assert_eq!(mem::align_of::<[u32; 3]>(), mem::align_of::<u32>());
                let (t, t_size, t_align) = type_from_id(doc, component_id);
                return (format!("[{}; {}]", t, count), t_size.map(|s| s * count as usize), t_align);
            },
            &parse::Instruction::TypeMatrix { result_id, column_type_id, column_count } if result_id == searched => {
                // FIXME: row-major or column-major
                debug_assert_eq!(mem::align_of::<[u32; 3]>(), mem::align_of::<u32>());
                let (t, t_size, t_align) = type_from_id(doc, column_type_id);
                return (format!("[{}; {}]", t, column_count), t_size.map(|s| s * column_count as usize), t_align);
            },
            &parse::Instruction::TypeArray { result_id, type_id, length_id } if result_id == searched => {
                debug_assert_eq!(mem::align_of::<[u32; 3]>(), mem::align_of::<u32>());
                let (t, t_size, t_align) = type_from_id(doc, type_id);
                let len = doc.instructions.iter().filter_map(|e| {
                    match e { &parse::Instruction::Constant { result_id, ref data, .. } if result_id == length_id => Some(data.clone()), _ => None }
                }).next().expect("failed to find array length");
                let len = len.iter().rev().fold(0u64, |a, &b| (a << 32) | b as u64);
                return (format!("[{}; {}]", t, len), t_size.map(|s| s * len as usize), t_align);       // FIXME:
            },
            &parse::Instruction::TypeRuntimeArray { result_id, type_id } if result_id == searched => {
                debug_assert_eq!(mem::align_of::<[u32; 3]>(), mem::align_of::<u32>());
                let (t, _, t_align) = type_from_id(doc, type_id);
                return (format!("[{}]", t), None, t_align);
            },
            &parse::Instruction::TypeStruct { result_id, ref member_types } if result_id == searched => {
                let name = ::name_from_id(doc, result_id);
                let size = member_types.iter().filter_map(|&t| type_from_id(doc, t).1).fold(0, |a,b|a+b);
                let align = member_types.iter().map(|&t| type_from_id(doc, t).2).max().unwrap_or(1);
                return (name, Some(size), align);
            },
            _ => ()
        }
    }

    panic!("Type #{} not found", searched)
}
<|endoftext|>


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

use crate::{
    block::Block,
    common::{Payload, Round},
    quorum_cert::QuorumCert,
    vote_proposal::{MaybeSignedVoteProposal, VoteProposal},
};
use executor_types::StateComputeResult;
use libra_crypto::hash::HashValue;
use libra_types::block_info::BlockInfo;
use std::fmt::{Display, Formatter};

/// ExecutedBlocks are managed in a speculative tree, the committed blocks form a chain. Besides
/// block data, each executed block also has other derived meta data which could be regenerated from
/// blocks.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ExecutedBlock {
    /// Block data that cannot be regenerated.
    block: Block,
    /// The state_compute_result is calculated for all the pending blocks prior to insertion to
    /// the tree. The execution results are not persisted: they're recalculated again for the
    /// pending blocks upon restart.
    state_compute_result: StateComputeResult,
}

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

impl ExecutedBlock {
    pub fn new(block: Block, state_compute_result: StateComputeResult) -> Self {
        Self {
            block,
            state_compute_result,
        }
    }

    pub fn block(&self) -> &Block {
        &self.block
    }

    pub fn id(&self) -> HashValue {
        self.block().id()
    }

    pub fn payload(&self) -> Option<&Payload> {
        self.block().payload()
    }

    pub fn parent_id(&self) -> HashValue {
        self.quorum_cert().certified_block().id()
    }

    pub fn quorum_cert(&self) -> &QuorumCert {
        self.block().quorum_cert()
    }

    pub fn round(&self) -> Round {
        self.block().round()
    }

    pub fn timestamp_usecs(&self) -> u64 {
        self.block().timestamp_usecs()
    }

    pub fn compute_result(&self) -> &StateComputeResult {
        &self.state_compute_result
    }

    pub fn block_info(&self) -> BlockInfo {
        self.block().gen_block_info(
            self.compute_result().root_hash(),
            self.compute_result().version(),
            self.compute_result().epoch_state().clone(),
        )
    }

    pub fn maybe_signed_vote_proposal(&self) -> MaybeSignedVoteProposal {
        MaybeSignedVoteProposal {
            vote_proposal: VoteProposal::new(
                self.compute_result().extension_proof(),
                self.block.clone(),
                self.compute_result().epoch_state().clone(),
            ),
            signature: self.compute_result().signature().clone(),
        }
    }
}
<commit_msg>[logging] Removed debug output from ExecutedBlock<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    block::Block,
    common::{Payload, Round},
    quorum_cert::QuorumCert,
    vote_proposal::{MaybeSignedVoteProposal, VoteProposal},
};
use executor_types::StateComputeResult;
use libra_crypto::hash::HashValue;
use libra_types::block_info::BlockInfo;
use std::fmt::{Debug, Display, Formatter};

/// ExecutedBlocks are managed in a speculative tree, the committed blocks form a chain. Besides
/// block data, each executed block also has other derived meta data which could be regenerated from
/// blocks.
#[derive(Clone, Eq, PartialEq)]
pub struct ExecutedBlock {
    /// Block data that cannot be regenerated.
    block: Block,
    /// The state_compute_result is calculated for all the pending blocks prior to insertion to
    /// the tree. The execution results are not persisted: they're recalculated again for the
    /// pending blocks upon restart.
    state_compute_result: StateComputeResult,
}

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

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

impl ExecutedBlock {
    pub fn new(block: Block, state_compute_result: StateComputeResult) -> Self {
        Self {
            block,
            state_compute_result,
        }
    }

    pub fn block(&self) -> &Block {
        &self.block
    }

    pub fn id(&self) -> HashValue {
        self.block().id()
    }

    pub fn payload(&self) -> Option<&Payload> {
        self.block().payload()
    }

    pub fn parent_id(&self) -> HashValue {
        self.quorum_cert().certified_block().id()
    }

    pub fn quorum_cert(&self) -> &QuorumCert {
        self.block().quorum_cert()
    }

    pub fn round(&self) -> Round {
        self.block().round()
    }

    pub fn timestamp_usecs(&self) -> u64 {
        self.block().timestamp_usecs()
    }

    pub fn compute_result(&self) -> &StateComputeResult {
        &self.state_compute_result
    }

    pub fn block_info(&self) -> BlockInfo {
        self.block().gen_block_info(
            self.compute_result().root_hash(),
            self.compute_result().version(),
            self.compute_result().epoch_state().clone(),
        )
    }

    pub fn maybe_signed_vote_proposal(&self) -> MaybeSignedVoteProposal {
        MaybeSignedVoteProposal {
            vote_proposal: VoteProposal::new(
                self.compute_result().extension_proof(),
                self.block.clone(),
                self.compute_result().epoch_state().clone(),
            ),
            signature: self.compute_result().signature().clone(),
        }
    }
}
<|endoftext|>


<commit_before>use crate::{
    script::{EvalErr, EvalErrType},
    serializer::*,
};
use std::io::{self, Cursor};

pub type SkipFlags = u8;

#[allow(clippy::identity_op)]
pub const SKIP_NONE: u8 = 1 << 0;
pub const SKIP_REWARD_TX: u8 = 1 << 1;

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum BlockErr {
    InvalidBlockHeight,
    InvalidMerkleRoot,
    InvalidSignature,
    InvalidPrevHash,
    InvalidHash,
    Tx(TxErr),
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum TxErr {
    ScriptEval(EvalErr),
    ScriptHashMismatch,
    ScriptRetFalse,
    Arithmetic,
    InsufficientBalance,
    InvalidFeeAmount,
    TooManySignatures,
    TxTooLarge,
    TxProhibited,
    TxExpired,
    TxDupe,
}

impl TxErr {
    pub fn serialize(self, buf: &mut Vec<u8>) {
        match self {
            TxErr::ScriptEval(err) => {
                buf.push(0);
                buf.push_u32(err.pos);
                buf.push(err.err as u8);
            }
            TxErr::ScriptHashMismatch => buf.push(1),
            TxErr::ScriptRetFalse => buf.push(2),
            TxErr::Arithmetic => buf.push(3),
            TxErr::InsufficientBalance => buf.push(4),
            TxErr::InvalidFeeAmount => buf.push(5),
            TxErr::TooManySignatures => buf.push(6),
            TxErr::TxTooLarge => buf.push(7),
            TxErr::TxProhibited => buf.push(8),
            TxErr::TxExpired => buf.push(9),
            TxErr::TxDupe => buf.push(10),
        }
    }

    pub fn deserialize(cursor: &mut Cursor<&[u8]>) -> io::Result<Self> {
        let tag = cursor.take_u8()?;
        Ok(match tag {
            0 => {
                let pos = cursor.take_u32()?;
                let kind = match cursor.take_u8()? {
                    t if t == EvalErrType::UnexpectedEOF as u8 => EvalErrType::UnexpectedEOF,
                    t if t == EvalErrType::UnknownOp as u8 => EvalErrType::UnknownOp,
                    t if t == EvalErrType::InvalidItemOnStack as u8 => {
                        EvalErrType::InvalidItemOnStack
                    }
                    t if t == EvalErrType::StackOverflow as u8 => EvalErrType::StackOverflow,
                    t if t == EvalErrType::StackUnderflow as u8 => EvalErrType::StackUnderflow,
                    _ => {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidData,
                            "failed to deserialize EvalErrType",
                        ))
                    }
                };
                TxErr::ScriptEval(EvalErr::new(pos, kind))
            }
            1 => TxErr::ScriptHashMismatch,
            2 => TxErr::ScriptRetFalse,
            3 => TxErr::Arithmetic,
            4 => TxErr::InsufficientBalance,
            5 => TxErr::InvalidFeeAmount,
            6 => TxErr::TooManySignatures,
            7 => TxErr::TxTooLarge,
            8 => TxErr::TxProhibited,
            9 => TxErr::TxExpired,
            10 => TxErr::TxDupe,
            _ => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "failed to deserialize TxErr",
                ))
            }
        })
    }
}
<commit_msg>Use hex for TxErr serialization<commit_after>use crate::{
    script::{EvalErr, EvalErrType},
    serializer::*,
};
use std::io::{self, Cursor};

pub type SkipFlags = u8;

#[allow(clippy::identity_op)]
pub const SKIP_NONE: u8 = 1 << 0;
pub const SKIP_REWARD_TX: u8 = 1 << 1;

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum BlockErr {
    InvalidBlockHeight,
    InvalidMerkleRoot,
    InvalidSignature,
    InvalidPrevHash,
    InvalidHash,
    Tx(TxErr),
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum TxErr {
    ScriptEval(EvalErr),
    ScriptHashMismatch,
    ScriptRetFalse,
    Arithmetic,
    InsufficientBalance,
    InvalidFeeAmount,
    TooManySignatures,
    TxTooLarge,
    TxProhibited,
    TxExpired,
    TxDupe,
}

impl TxErr {
    pub fn serialize(self, buf: &mut Vec<u8>) {
        match self {
            TxErr::ScriptEval(err) => {
                buf.push(0x00);
                buf.push_u32(err.pos);
                buf.push(err.err as u8);
            }
            TxErr::ScriptHashMismatch => buf.push(0x01),
            TxErr::ScriptRetFalse => buf.push(0x02),
            TxErr::Arithmetic => buf.push(0x03),
            TxErr::InsufficientBalance => buf.push(0x04),
            TxErr::InvalidFeeAmount => buf.push(0x05),
            TxErr::TooManySignatures => buf.push(0x06),
            TxErr::TxTooLarge => buf.push(0x07),
            TxErr::TxProhibited => buf.push(0x08),
            TxErr::TxExpired => buf.push(0x09),
            TxErr::TxDupe => buf.push(0x0A),
        }
    }

    pub fn deserialize(cursor: &mut Cursor<&[u8]>) -> io::Result<Self> {
        let tag = cursor.take_u8()?;
        Ok(match tag {
            0x00 => {
                let pos = cursor.take_u32()?;
                let kind = match cursor.take_u8()? {
                    t if t == EvalErrType::UnexpectedEOF as u8 => EvalErrType::UnexpectedEOF,
                    t if t == EvalErrType::UnknownOp as u8 => EvalErrType::UnknownOp,
                    t if t == EvalErrType::InvalidItemOnStack as u8 => {
                        EvalErrType::InvalidItemOnStack
                    }
                    t if t == EvalErrType::StackOverflow as u8 => EvalErrType::StackOverflow,
                    t if t == EvalErrType::StackUnderflow as u8 => EvalErrType::StackUnderflow,
                    _ => {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidData,
                            "failed to deserialize EvalErrType",
                        ))
                    }
                };
                TxErr::ScriptEval(EvalErr::new(pos, kind))
            }
            0x01 => TxErr::ScriptHashMismatch,
            0x02 => TxErr::ScriptRetFalse,
            0x03 => TxErr::Arithmetic,
            0x04 => TxErr::InsufficientBalance,
            0x05 => TxErr::InvalidFeeAmount,
            0x06 => TxErr::TooManySignatures,
            0x07 => TxErr::TxTooLarge,
            0x08 => TxErr::TxProhibited,
            0x09 => TxErr::TxExpired,
            0x0A => TxErr::TxDupe,
            _ => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "failed to deserialize TxErr",
                ))
            }
        })
    }
}
<|endoftext|>


<commit_before>/**
 *   Mles-utils to be used with Mles client or server.
 *   Copyright (C) 2017  Mles developers
 *
 *   This program is free software: you can redistribute it and/or modify
 *   it under the terms of the GNU General Public License as published by
 *   the Free Software Foundation, either version 3 of the License, or
 *   (at your option) any later version.
 *
 *   This program is distributed in the hope that it will be useful,
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *   GNU General Public License for more details.
 *
 *   You should have received a copy of the GNU General Public License
 *   along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
extern crate tokio_core;
extern crate futures;
extern crate mles_utils;

use std::{thread, process, env};
use std::sync::mpsc::{Sender, Receiver};
use std::sync::mpsc::channel;
use std::net::TcpStream;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::net::TcpListener;
use std::collections::HashMap;
use std::io::{Read, Error};
use std::io::Write;
use std::time::Duration;
use std::option::Option;
use std::str;
use mles_utils::*;

const HDRL: u64 = 4;
const KEYL: u64 = 8;
const DELAY: u64 = 50;

fn main_old() {
    let mut peer = "".to_string();
    let mut argcnt = 0;
    for arg in env::args() {
        argcnt += 1;
        if 2 == argcnt {
            peer = arg;
            peer += ":8077";
            match peer.parse::<SocketAddr>() {
                Ok(_) => {},
                    Err(err) => {
                        println!("Error: {}\nUsage: mles [peer-address]", err);
                        process::exit(1);
                    },
            }
        }
        if argcnt > 2 {
            println!("Usage: mles [peer-address]");
            process::exit(1);
        }
    }
    let peer = match peer.parse::<SocketAddr>() {
        Ok(addr) => addr,
        Err(_) => {
            SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0)
        },
    };

    let address = "0.0.0.0:8077";
    let listener = match TcpListener::bind(&address) {
        Ok(listener) => listener,
        Err(err) => {
            println!("Error: {}", err);
            process::exit(1);
        },
    };

    let keyval =match env::var("MLES_KEY") {
        Ok(val) => val,
        Err(_) => "".to_string(),
    };

    let mut spawned = HashMap::new();
    let option: Option<Duration> = Some(Duration::from_millis(DELAY));
    let addr = listener.local_addr().unwrap();
    println!("Listening for connections on {}", addr);
    let (tx, rx) = channel();
    let (removedtx, removedrx) = channel();

    for socket in listener.incoming() {
        /* Check first has anybody removed channels */
        let mut remrx = true;
        while remrx {
            match removedrx.try_recv() {
                Ok(val) => {
                    let removed_channel: String = val;
                    println!("Removing unused channel {}", removed_channel.as_str());
                    spawned.remove(&removed_channel);
                    remrx = true;
                },
                Err(_) => { remrx = false; }
            }
        }
        /* 1. Read incoming msg channel
         * 2. If it does not exist, spawn new thread
         * 3. Send socket to thread
         */
        let socket = socket.unwrap();
        let stream = socket.try_clone().unwrap();
        let tuple = read_n(&stream, HDRL);
        let status = tuple.0;
        match status {
            Ok(0) => {
               continue;
            },
            Ok(_) => {},
            _ => {
               continue;
            },
        }
        let buf = tuple.1;
        if 0 == buf.len() {
            continue;
        }
        if read_hdr_type(buf.as_slice()) != 'M' as u32 {
            println!("Incorrect payload type");
            continue;
        }
        // read key
        let tuple = read_n(&stream, KEYL);
        let status = tuple.0;
        match status {
            Ok(0) => {
               continue;
            },
            Ok(_) => {},
            _ => {
               continue;
            },
        }
        // verify key
        let hkey;
        let key = read_key(tuple.1);
        if 0 == keyval.len() {
            let paddr = match stream.peer_addr() {
                Ok(paddr) => paddr,
                Err(_) => {
                    let addr = "0.0.0.0:0";
                    let addr = addr.parse::<SocketAddr>().unwrap();
                    addr
                }
            };
            hkey = do_hash(&paddr);
        }
        else {
            hkey = do_hash(&keyval);
        }
        if hkey != key {
            println!("Incorrect remote key");
            continue;
        }
        let tuple = read_n(&stream, read_hdr_len(buf.as_slice()) as u64);
        let status = tuple.0;
        match status {
            Ok(0) => {
               continue;
            },
            Ok(_) => {},
            _ => {
               continue;
            },
        }
        let buf = tuple.1;
        let decoded_msg = message_decode(buf.as_slice());
        if 0 == decoded_msg.channel.len() {
            continue;
        }
        let _val = socket.set_nodelay(true);
        let _val = socket.set_read_timeout(option);
        if !spawned.contains_key(decoded_msg.channel.as_str()) {
            let tx = tx.clone();
            let removedtx = removedtx.clone();
            let msg = decoded_msg.clone();
            thread::spawn(move|| process_channel(peer, key, tx, removedtx, msg));

            let thr_feed = rx.recv().unwrap();
            spawned.insert(decoded_msg.channel.clone(), thr_feed);
        }
        let thr_socket = spawned.get_mut(&decoded_msg.channel).unwrap();
        thr_socket.send(socket).unwrap();
    }
}

fn process_channel(peer: SocketAddr, key: u64, tx: Sender<Sender<TcpStream>>, removedtx: Sender<String>, msg: Msg ) {
    let mut cnt = 0;
    let (thr_tx, thr_rx): (Sender<TcpStream>, Receiver<TcpStream>) = channel();
    println!("Spawned: New channel {} created!", msg.channel);
    let mut users = HashMap::new();
    let mut messages: Vec<Vec<u8>> = Vec::new();
    tx.send(thr_tx.clone()).unwrap();

    // just try once during channel creation
    if 0 != peer.port() {
        match TcpStream::connect(peer) {
            Ok(mut sock) => {
                let option: Option<Duration> = Some(Duration::from_millis(DELAY));
                let _val = sock.set_nodelay(true);
                let _val = sock.set_read_timeout(option);
                let encoded_msg = message_encode(&msg);
                let keyv = write_key(key);
                let mut msgv = write_hdr(encoded_msg.len());
                msgv.extend(keyv);
                msgv.extend(encoded_msg);
                sock.write(msgv.as_slice()).unwrap();

                cnt += 1;
                println!("Adding peer {}", cnt);
                users.insert(cnt, sock);
            },
            Err(_) => {
                println!("Could not connect to peer {}", peer);
            },
        }
    }

    loop {
        let mut removals = Vec::new();
        let mut newuser = true;
        while newuser {
            match thr_rx.try_recv() {
                Ok(val) => {
                    let mut thr: TcpStream = val;
                    cnt += 1;
                    println!("Adding user {}", cnt);
                    users.insert(cnt, thr.try_clone().unwrap());

                    /* If a new user, all push messages to her */
                    for buf in &messages {
                        thr.write(buf.as_slice()).unwrap();
                    }
                    newuser = true;
                },
                Err(_) => { newuser = false; }
            }
        }
        for (user, thr_socket) in &users {
            let stream = thr_socket.try_clone().unwrap();
            let tuple = read_n(&stream, HDRL);
            let status = tuple.0;
            match status {
                Ok(0) => {
                    removals.push(user.clone());
                },
                    _ => {}
            }
            let mut buf = tuple.1;
            if 0 == buf.len() {
                continue;
            }
            if read_hdr_type(buf.as_slice()) != 'M' as u32 {
                continue;
            }
            let hdr_len = read_hdr_len(buf.as_slice()) as u64;
            if 0 == hdr_len {
                continue;
            }
            // read key
            let tuple = read_n(&stream, KEYL);
            let status = tuple.0;
            match status {
                Ok(0) => {
                    continue;
                },
                    Ok(_) => {},
                    _ => {
                        continue;
                    },
            }
            let key = tuple.1;
            //ignore key
            buf.extend(key);
            let tuple = read_n(&stream, hdr_len);
            let status = tuple.0;
            match status {
                Ok(0) => {
                    removals.push(user.clone());
                },
                    _ => {}
            }
            let payload = tuple.1;
            if payload.len() != (hdr_len as usize) {
                continue;
            }
            buf.extend(payload);
            for (another_user, mut thr_sock) in &users {
                if user != another_user {
                    thr_sock.write(buf.as_slice()).unwrap();
                }
            }
            /* Add to local db */
            messages.push(buf);
        }
        for removal in &removals {
            println!("Removing user {}", removal);
            users.remove(removal);
        }
        if cnt > 0 && users.is_empty() {
            removedtx.send(msg.channel).unwrap();
            break;
        }
    }
}

fn read_n<R>(reader: R, bytes_to_read: u64) -> (Result<usize, Error>, Vec<u8>)
where R: Read,
{
    let mut buf = vec![];
    let mut chunk = reader.take(bytes_to_read);
    let status = chunk.read_to_end(&mut buf);
    (status, buf)
}

/* Tokio version */
use std::collections::HashMap;
use std::rc::Rc;
use std::cell::RefCell;
use std::iter;
use std::env;
use std::io::{Error, ErrorKind};

use tokio_core::net::TcpListener;
use tokio_core::reactor::Core;
use tokio_core::io::{self, Io};

use futures::stream::{self, Stream};
use futures::Future;
use mles_utils::*;

//static mut BUF: Vec<u8> = vec![0,4];

fn main() {
    let addr = env::args().nth(1).unwrap_or("0.0.0.0:8077".to_string());
    let addr = addr.parse().unwrap();
    let mut cnt = 0;

    // Create the event loop and TCP listener we'll accept connections on.
    let mut core = Core::new().unwrap();
    let handle = core.handle();
    let socket = TcpListener::bind(&addr, &handle).unwrap();
    println!("Listening on: {}", addr);

    // This is a single-threaded server, so we can just use Rc and RefCell to
    // store the map of all connections we know about.
    let connections = Rc::new(RefCell::new(HashMap::new()));

    let srv = socket.incoming().for_each(move |(stream, addr)| {
        println!("New Connection: {}", addr);
        let (reader, writer) = stream.split();

        // Create a channel for our stream, which other sockets will use to
        // send us messages. Then register our address with the stream to send
        // data to us.
        let (tx, rx) = futures::sync::mpsc::unbounded();
        cnt += 1;
        connections.borrow_mut().insert(cnt, tx);

        // Define here what we do for the actual I/O. That is, read a bunch of
        // frames from the socket and dispatch them while we also write any frames
        // from other sockets.
        let connections_inner = connections.clone();

        // Model the read portion of this socket by mapping an infinite
        // iterator to each frame off the socket. This "loop" is then
        // terminated with an error once we hit EOF on the socket.
        let iter = stream::iter(iter::repeat(()).map(Ok::<(), Error>));
        let socket_reader = iter.fold(reader, move |reader, _| {
            // Read a off header the socket, failing if we're at EOF
            let frame = io::read_exact(reader, vec![0;4]);
            let frame = frame.and_then(|(reader, payload)| {
                if payload.len() == 0 {
                    Err(Error::new(ErrorKind::BrokenPipe, "broken pipe"))
                } else {
                    if read_hdr_type(payload.as_slice()) != 'M' as u32 {
                        return Err(Error::new(ErrorKind::BrokenPipe, "incorrect header"));
                    }
                    let hdr_len = read_hdr_len(payload.as_slice());
                    if 0 == hdr_len {
                        return Err(Error::new(ErrorKind::BrokenPipe, "incorrect header len"));
                    }
                    Ok((reader, payload))
                }
            });
            //let pframe = frame.and_then(|reader, payload)| {
            //        let pframe = io::read_exact(reader, vec![0;hdr_len]);
            //        let pframe = pframe.and_then(|(reader, buf)| {
            //            if buf.len() == 0 {
            //                return Err(Error::new(ErrorKind::BrokenPipe, "broken pipe"));
           //             }
           //             Ok((reader, buf))
            //        });

            // Send frame to all other connected clients
            let connections = connections_inner.clone();
            frame.map(move |(reader, message)| {
                println!("{}: {:?}", cnt, message);
                let mut conns = connections.borrow_mut();
                let iter = conns.iter_mut()
                .filter(|&(&k, _)| k != cnt)
                .map(|(_, v)| v);
                for tx in iter {
                    println!("Sending {}: {:?}", cnt, message);
                    tx.send(message.clone()).unwrap();
                }
                reader
            })
        });

        // Whenever we receive a string on the Receiver, we write it to
        // `WriteHalf<TcpStream>`.
        let socket_writer = rx.fold(writer, |writer, msg| {
            let amt = io::write_all(writer, msg);
            let amt = amt.map(|(writer, _)| writer);
            amt.map_err(|_| ())
        });

        // Now that we've got futures representing each half of the socket, we
        // use the `select` combinator to wait for either half to be done to
        // tear down the other. Then we spawn off the result.
        let connections = connections.clone();
        let socket_reader = socket_reader.map_err(|_| ());
        let connection = socket_reader.map(|_| ()).select(socket_writer.map(|_| ()));
        handle.spawn(connection.then(move |_| {
            connections.borrow_mut().remove(&cnt);
            println!("Connection {} closed.", addr);
            Ok(())
        }));

        Ok(())
    });

    // execute server
    core.run(srv).unwrap();
}
<commit_msg>Read header with Tokio.<commit_after>/**
 *   Mles-utils to be used with Mles client or server.
 *   Copyright (C) 2017  Mles developers
 *
 *   This program is free software: you can redistribute it and/or modify
 *   it under the terms of the GNU General Public License as published by
 *   the Free Software Foundation, either version 3 of the License, or
 *   (at your option) any later version.
 *
 *   This program is distributed in the hope that it will be useful,
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *   GNU General Public License for more details.
 *
 *   You should have received a copy of the GNU General Public License
 *   along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
extern crate tokio_core;
extern crate futures;
extern crate mles_utils;

use std::{thread, process, env};
use std::sync::mpsc::{Sender, Receiver};
use std::sync::mpsc::channel;
use std::net::TcpStream;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::net::TcpListener;
use std::collections::HashMap;
use std::io::{Read, Error};
use std::io::Write;
use std::time::Duration;
use std::option::Option;
use std::str;
use mles_utils::*;

const HDRL: u64 = 4;
const KEYL: u64 = 8;
const DELAY: u64 = 50;

fn main_old() {
    let mut peer = "".to_string();
    let mut argcnt = 0;
    for arg in env::args() {
        argcnt += 1;
        if 2 == argcnt {
            peer = arg;
            peer += ":8077";
            match peer.parse::<SocketAddr>() {
                Ok(_) => {},
                    Err(err) => {
                        println!("Error: {}\nUsage: mles [peer-address]", err);
                        process::exit(1);
                    },
            }
        }
        if argcnt > 2 {
            println!("Usage: mles [peer-address]");
            process::exit(1);
        }
    }
    let peer = match peer.parse::<SocketAddr>() {
        Ok(addr) => addr,
        Err(_) => {
            SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0)
        },
    };

    let address = "0.0.0.0:8077";
    let listener = match TcpListener::bind(&address) {
        Ok(listener) => listener,
        Err(err) => {
            println!("Error: {}", err);
            process::exit(1);
        },
    };

    let keyval =match env::var("MLES_KEY") {
        Ok(val) => val,
        Err(_) => "".to_string(),
    };

    let mut spawned = HashMap::new();
    let option: Option<Duration> = Some(Duration::from_millis(DELAY));
    let addr = listener.local_addr().unwrap();
    println!("Listening for connections on {}", addr);
    let (tx, rx) = channel();
    let (removedtx, removedrx) = channel();

    for socket in listener.incoming() {
        /* Check first has anybody removed channels */
        let mut remrx = true;
        while remrx {
            match removedrx.try_recv() {
                Ok(val) => {
                    let removed_channel: String = val;
                    println!("Removing unused channel {}", removed_channel.as_str());
                    spawned.remove(&removed_channel);
                    remrx = true;
                },
                Err(_) => { remrx = false; }
            }
        }
        /* 1. Read incoming msg channel
         * 2. If it does not exist, spawn new thread
         * 3. Send socket to thread
         */
        let socket = socket.unwrap();
        let stream = socket.try_clone().unwrap();
        let tuple = read_n(&stream, HDRL);
        let status = tuple.0;
        match status {
            Ok(0) => {
               continue;
            },
            Ok(_) => {},
            _ => {
               continue;
            },
        }
        let buf = tuple.1;
        if 0 == buf.len() {
            continue;
        }
        if read_hdr_type(buf.as_slice()) != 'M' as u32 {
            println!("Incorrect payload type");
            continue;
        }
        // read key
        let tuple = read_n(&stream, KEYL);
        let status = tuple.0;
        match status {
            Ok(0) => {
               continue;
            },
            Ok(_) => {},
            _ => {
               continue;
            },
        }
        // verify key
        let hkey;
        let key = read_key(tuple.1);
        if 0 == keyval.len() {
            let paddr = match stream.peer_addr() {
                Ok(paddr) => paddr,
                Err(_) => {
                    let addr = "0.0.0.0:0";
                    let addr = addr.parse::<SocketAddr>().unwrap();
                    addr
                }
            };
            hkey = do_hash(&paddr);
        }
        else {
            hkey = do_hash(&keyval);
        }
        if hkey != key {
            println!("Incorrect remote key");
            continue;
        }
        let tuple = read_n(&stream, read_hdr_len(buf.as_slice()) as u64);
        let status = tuple.0;
        match status {
            Ok(0) => {
               continue;
            },
            Ok(_) => {},
            _ => {
               continue;
            },
        }
        let buf = tuple.1;
        let decoded_msg = message_decode(buf.as_slice());
        if 0 == decoded_msg.channel.len() {
            continue;
        }
        let _val = socket.set_nodelay(true);
        let _val = socket.set_read_timeout(option);
        if !spawned.contains_key(decoded_msg.channel.as_str()) {
            let tx = tx.clone();
            let removedtx = removedtx.clone();
            let msg = decoded_msg.clone();
            thread::spawn(move|| process_channel(peer, key, tx, removedtx, msg));

            let thr_feed = rx.recv().unwrap();
            spawned.insert(decoded_msg.channel.clone(), thr_feed);
        }
        let thr_socket = spawned.get_mut(&decoded_msg.channel).unwrap();
        thr_socket.send(socket).unwrap();
    }
}

fn process_channel(peer: SocketAddr, key: u64, tx: Sender<Sender<TcpStream>>, removedtx: Sender<String>, msg: Msg ) {
    let mut cnt = 0;
    let (thr_tx, thr_rx): (Sender<TcpStream>, Receiver<TcpStream>) = channel();
    println!("Spawned: New channel {} created!", msg.channel);
    let mut users = HashMap::new();
    let mut messages: Vec<Vec<u8>> = Vec::new();
    tx.send(thr_tx.clone()).unwrap();

    // just try once during channel creation
    if 0 != peer.port() {
        match TcpStream::connect(peer) {
            Ok(mut sock) => {
                let option: Option<Duration> = Some(Duration::from_millis(DELAY));
                let _val = sock.set_nodelay(true);
                let _val = sock.set_read_timeout(option);
                let encoded_msg = message_encode(&msg);
                let keyv = write_key(key);
                let mut msgv = write_hdr(encoded_msg.len());
                msgv.extend(keyv);
                msgv.extend(encoded_msg);
                sock.write(msgv.as_slice()).unwrap();

                cnt += 1;
                println!("Adding peer {}", cnt);
                users.insert(cnt, sock);
            },
            Err(_) => {
                println!("Could not connect to peer {}", peer);
            },
        }
    }

    loop {
        let mut removals = Vec::new();
        let mut newuser = true;
        while newuser {
            match thr_rx.try_recv() {
                Ok(val) => {
                    let mut thr: TcpStream = val;
                    cnt += 1;
                    println!("Adding user {}", cnt);
                    users.insert(cnt, thr.try_clone().unwrap());

                    /* If a new user, all push messages to her */
                    for buf in &messages {
                        thr.write(buf.as_slice()).unwrap();
                    }
                    newuser = true;
                },
                Err(_) => { newuser = false; }
            }
        }
        for (user, thr_socket) in &users {
            let stream = thr_socket.try_clone().unwrap();
            let tuple = read_n(&stream, HDRL);
            let status = tuple.0;
            match status {
                Ok(0) => {
                    removals.push(user.clone());
                },
                    _ => {}
            }
            let mut buf = tuple.1;
            if 0 == buf.len() {
                continue;
            }
            if read_hdr_type(buf.as_slice()) != 'M' as u32 {
                continue;
            }
            let hdr_len = read_hdr_len(buf.as_slice()) as u64;
            if 0 == hdr_len {
                continue;
            }
            // read key
            let tuple = read_n(&stream, KEYL);
            let status = tuple.0;
            match status {
                Ok(0) => {
                    continue;
                },
                    Ok(_) => {},
                    _ => {
                        continue;
                    },
            }
            let key = tuple.1;
            //ignore key
            buf.extend(key);
            let tuple = read_n(&stream, hdr_len);
            let status = tuple.0;
            match status {
                Ok(0) => {
                    removals.push(user.clone());
                },
                    _ => {}
            }
            let payload = tuple.1;
            if payload.len() != (hdr_len as usize) {
                continue;
            }
            buf.extend(payload);
            for (another_user, mut thr_sock) in &users {
                if user != another_user {
                    thr_sock.write(buf.as_slice()).unwrap();
                }
            }
            /* Add to local db */
            messages.push(buf);
        }
        for removal in &removals {
            println!("Removing user {}", removal);
            users.remove(removal);
        }
        if cnt > 0 && users.is_empty() {
            removedtx.send(msg.channel).unwrap();
            break;
        }
    }
}

fn read_n<R>(reader: R, bytes_to_read: u64) -> (Result<usize, Error>, Vec<u8>)
where R: Read,
{
    let mut buf = vec![];
    let mut chunk = reader.take(bytes_to_read);
    let status = chunk.read_to_end(&mut buf);
    (status, buf)
}

/* Tokio version */
use std::collections::HashMap;
use std::rc::Rc;
use std::cell::RefCell;
use std::iter;
use std::env;
use std::io::{Error, ErrorKind};

use tokio_core::net::TcpListener;
use tokio_core::reactor::Core;
use tokio_core::io::{self, Io};

use futures::stream::{self, Stream};
use futures::Future;
use mles_utils::*;

//static mut BUF: Vec<u8> = vec![0,4];

fn main() {
    let addr = env::args().nth(1).unwrap_or("0.0.0.0:8077".to_string());
    let addr = addr.parse().unwrap();
    let mut cnt = 0;

    // Create the event loop and TCP listener we'll accept connections on.
    let mut core = Core::new().unwrap();
    let handle = core.handle();
    let socket = TcpListener::bind(&addr, &handle).unwrap();
    println!("Listening on: {}", addr);

    // This is a single-threaded server, so we can just use Rc and RefCell to
    // store the map of all connections we know about.
    let connections = Rc::new(RefCell::new(HashMap::new()));

    let srv = socket.incoming().for_each(move |(stream, addr)| {
        println!("New Connection: {}", addr);
        let (reader, writer) = stream.split();

        // Create a channel for our stream, which other sockets will use to
        // send us messages. Then register our address with the stream to send
        // data to us.
        let (tx, rx) = futures::sync::mpsc::unbounded();
        cnt += 1;
        connections.borrow_mut().insert(cnt, tx);

        // Define here what we do for the actual I/O. That is, read a bunch of
        // frames from the socket and dispatch them while we also write any frames
        // from other sockets.
        let connections_inner = connections.clone();

        // Model the read portion of this socket by mapping an infinite
        // iterator to each frame off the socket. This "loop" is then
        // terminated with an error once we hit EOF on the socket.
        let iter = stream::iter(iter::repeat(()).map(Ok::<(), Error>));
        let socket_reader = iter.fold(reader, move |reader, _| {
            // Read a off header the socket, failing if we're at EOF
            let frame = io::read_exact(reader, vec![0;4]);
            let frame = frame.and_then(move |(reader, payload)| {
                if payload.len() == 0 {
                    Err(Error::new(ErrorKind::BrokenPipe, "broken pipe"))
                } else {
                    if read_hdr_type(payload.as_slice()) != 'M' as u32 {
                        return Err(Error::new(ErrorKind::BrokenPipe, "incorrect header"));
                    }
                    let hdr_len = read_hdr_len(payload.as_slice());
                    if 0 == hdr_len {
                        return Err(Error::new(ErrorKind::BrokenPipe, "incorrect header len"));
                    }
                    Ok((reader, hdr_len))
                }
            });

            let frame = frame.and_then(move |(reader, hdr_len)| {
                println!("Hdrlen {}", hdr_len);
                let tframe = io::read_exact(reader, vec![0;hdr_len]);
                let tframe = tframe.and_then(move |(reader, payload)| {
                    let decoded_message = message_decode(payload.as_slice());
                    println!("Message {:?}", decoded_message);
                    Ok((reader, message_encode(&decoded_message)))
                });
                tframe
            });

            println!("Next to send ");
            // Send frame to all other connected clients
            let connections = connections_inner.clone();
            frame.map(move |(reader, message)| {
                println!("Message {}: {:?}", cnt, message);
                let mut conns = connections.borrow_mut();
                let iter = conns.iter_mut()
                .filter(|&(&k, _)| k != cnt)
                .map(|(_, v)| v);
                for tx in iter {
                    println!("Sending {}: {:?}", cnt, message);
                    tx.send(message.clone()).unwrap();
                }
                reader
            })
        });

        // Whenever we receive a string on the Receiver, we write it to
        // `WriteHalf<TcpStream>`.
        let socket_writer = rx.fold(writer, |writer, msg| {
            let amt = io::write_all(writer, msg);
            let amt = amt.map(|(writer, _)| writer);
            amt.map_err(|_| ())
        });

        // Now that we've got futures representing each half of the socket, we
        // use the `select` combinator to wait for either half to be done to
        // tear down the other. Then we spawn off the result.
        let connections = connections.clone();
        let socket_reader = socket_reader.map_err(|_| ());
        let connection = socket_reader.map(|_| ()).select(socket_writer.map(|_| ()));
        handle.spawn(connection.then(move |_| {
            connections.borrow_mut().remove(&cnt);
            println!("Connection {} closed.", addr);
            Ok(())
        }));

        Ok(())
    });

    // execute server
    core.run(srv).unwrap();
}
<|endoftext|>


<commit_before>//! # Schemes
//! A scheme is a primitive for handling filesystem syscalls in Redox.
//! Schemes accept paths from the kernel for `open`, and file descriptors that they generate
//! are then passed for operations like `close`, `read`, `write`, etc.
//!
//! The kernel validates paths and file descriptors before they are passed to schemes,
//! also stripping the scheme identifier of paths if necessary.

use alloc::arc::Arc;
use alloc::boxed::Box;
use collections::BTreeMap;
use core::sync::atomic::AtomicUsize;
use spin::{Once, RwLock, RwLockReadGuard, RwLockWriteGuard};

use syscall::error::*;
use syscall::scheme::Scheme;

use self::debug::DebugScheme;
use self::event::EventScheme;
use self::env::EnvScheme;
use self::initfs::InitFsScheme;
use self::irq::IrqScheme;
use self::memory::MemoryScheme;
use self::null::NullScheme;
use self::pipe::PipeScheme;
use self::root::RootScheme;
use self::sys::SysScheme;
use self::zero::ZeroScheme;

/// `debug:` - provides access to serial console
pub mod debug;

/// `event:` - allows reading of `Event`s which are registered using `fevent`
pub mod event;

/// `env:` - access and modify environmental variables
pub mod env;

/// `initfs:` - a readonly filesystem used for initializing the system
pub mod initfs;

/// `irq:` - allows userspace handling of IRQs
pub mod irq;

/// When compiled with "live" feature - `disk:` - embedded filesystem for live disk
#[cfg(feature="live")]
pub mod live;

/// `memory:` - a scheme for accessing physical memory
pub mod memory;

/// `null:` - a scheme that will discard all writes, and read no bytes
pub mod null;

/// `pipe:` - used internally by the kernel to implement `pipe`
pub mod pipe;

/// `:` - allows the creation of userspace schemes, tightly dependent on `user`
pub mod root;

/// `sys:` - system information, such as the context list and scheme list
pub mod sys;

/// A wrapper around userspace schemes, tightly dependent on `root`
pub mod user;

/// `zero:` - a scheme that will discard all writes, and always fill read buffers with zero
pub mod zero;

/// Limit on number of schemes
pub const SCHEME_MAX_SCHEMES: usize = 65536;

/// Unique identifier for a scheme namespace.
int_like!(SchemeNamespace, AtomicSchemeNamespace, usize, AtomicUsize);

/// Unique identifier for a scheme.
int_like!(SchemeId, AtomicSchemeId, usize, AtomicUsize);

pub const ATOMIC_SCHEMEID_INIT: AtomicSchemeId = AtomicSchemeId::default();

/// Unique identifier for a file descriptor.
int_like!(FileHandle, AtomicFileHandle, usize, AtomicUsize);

/// Scheme list type
pub struct SchemeList {
    map: BTreeMap<SchemeId, Arc<Box<Scheme + Send + Sync>>>,
    names: BTreeMap<SchemeNamespace, BTreeMap<Box<[u8]>, SchemeId>>,
    next_ns: usize,
    next_id: usize
}

impl SchemeList {
    /// Create a new scheme list.
    pub fn new() -> Self {
        let mut list = SchemeList {
            map: BTreeMap::new(),
            names: BTreeMap::new(),
            next_ns: 0,
            next_id: 1
        };
        list.new_root();
        list
    }

    /// Initialize a new namespace
    fn new_ns(&mut self) -> SchemeNamespace {
        let ns = SchemeNamespace(self.next_ns);
        self.next_ns += 1;
        self.names.insert(ns, BTreeMap::new());

        self.insert(ns, Box::new(*b""), |scheme_id| Arc::new(Box::new(RootScheme::new(ns, scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"event"), |_| Arc::new(Box::new(EventScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"env"), |_| Arc::new(Box::new(EnvScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"null"), |_| Arc::new(Box::new(NullScheme))).unwrap();
        self.insert(ns, Box::new(*b"sys"), |_| Arc::new(Box::new(SysScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"zero"), |_| Arc::new(Box::new(ZeroScheme))).unwrap();

        ns
    }

    /// Initialize the root namespace
    #[cfg(not(feature="live"))]
    fn new_root(&mut self) {
        // Do common namespace initialization
        let ns = self.new_ns();

        // Debug, Initfs and IRQ are only available in the root namespace. Pipe is special
        self.insert(ns, Box::new(*b"debug"), |scheme_id| Arc::new(Box::new(DebugScheme::new(scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"initfs"), |_| Arc::new(Box::new(InitFsScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"irq"), |scheme_id| Arc::new(Box::new(IrqScheme::new(scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"memory"), |_| Arc::new(Box::new(MemoryScheme))).unwrap();
        self.insert(ns, Box::new(*b"pipe"), |scheme_id| Arc::new(Box::new(PipeScheme::new(scheme_id)))).unwrap();
    }

    /// Initialize the root namespace - with live disk
    #[cfg(feature="live")]
    fn new_root(&mut self) {
        // Do common namespace initialization
        let ns = self.new_ns();

        // Debug, Disk, Initfs and IRQ are only available in the root namespace. Pipe is special
        self.insert(ns, Box::new(*b"debug"), |scheme_id| Arc::new(Box::new(DebugScheme::new(scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"disk"), |_| Arc::new(Box::new(self::live::DiskScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"initfs"), |_| Arc::new(Box::new(InitFsScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"irq"), |scheme_id| Arc::new(Box::new(IrqScheme::new(scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"pipe"), |scheme_id| Arc::new(Box::new(PipeScheme::new(scheme_id)))).unwrap();
    }

    pub fn make_ns(&mut self, from: SchemeNamespace, names: &[&[u8]]) -> Result<SchemeNamespace> {
        // Create an empty namespace
        let to = self.new_ns();

        // Copy requested scheme IDs
        for name in names.iter() {
            let id = if let Some((id, _scheme)) = self.get_name(from, name) {
                id
            } else {
                return Err(Error::new(ENODEV));
            };

            if let Some(ref mut names) = self.names.get_mut(&to) {
                assert!(names.insert(name.to_vec().into_boxed_slice(), id).is_none());
            } else {
                panic!("scheme namespace not found");
            }
        }

        Ok(to)
    }

    pub fn iter(&self) -> ::collections::btree_map::Iter<SchemeId, Arc<Box<Scheme + Send + Sync>>> {
        self.map.iter()
    }

    pub fn iter_name(&self, ns: SchemeNamespace) -> ::collections::btree_map::Iter<Box<[u8]>, SchemeId> {
        self.names[&ns].iter()
    }

    /// Get the nth scheme.
    pub fn get(&self, id: SchemeId) -> Option<&Arc<Box<Scheme + Send + Sync>>> {
        self.map.get(&id)
    }

    pub fn get_name(&self, ns: SchemeNamespace, name: &[u8]) -> Option<(SchemeId, &Arc<Box<Scheme + Send + Sync>>)> {
        if let Some(&id) = self.names[&ns].get(name) {
            self.get(id).map(|scheme| (id, scheme))
        } else {
            None
        }
    }

    /// Create a new scheme.
    pub fn insert<F>(&mut self, ns: SchemeNamespace, name: Box<[u8]>, scheme_fn: F) -> Result<SchemeId>
        where F: Fn(SchemeId) -> Arc<Box<Scheme + Send + Sync>>
    {
        if self.names[&ns].contains_key(&name) {
            return Err(Error::new(EEXIST));
        }

        if self.next_id >= SCHEME_MAX_SCHEMES {
            self.next_id = 1;
        }

        while self.map.contains_key(&SchemeId(self.next_id)) {
            self.next_id += 1;
        }

        /* Allow scheme list to grow if required
        if self.next_id >= SCHEME_MAX_SCHEMES {
            return Err(Error::new(EAGAIN));
        }
        */

        let id = SchemeId(self.next_id);
        self.next_id += 1;

        let scheme = scheme_fn(id);

        assert!(self.map.insert(id, scheme).is_none());
        if let Some(ref mut names) = self.names.get_mut(&ns) {
            assert!(names.insert(name, id).is_none());
        } else {
            panic!("scheme namespace not found");
        }
        Ok(id)
    }
}

/// Schemes list
static SCHEMES: Once<RwLock<SchemeList>> = Once::new();

/// Initialize schemes, called if needed
fn init_schemes() -> RwLock<SchemeList> {
    RwLock::new(SchemeList::new())
}

/// Get the global schemes list, const
pub fn schemes() -> RwLockReadGuard<'static, SchemeList> {
    SCHEMES.call_once(init_schemes).read()
}

/// Get the global schemes list, mutable
pub fn schemes_mut() -> RwLockWriteGuard<'static, SchemeList> {
    SCHEMES.call_once(init_schemes).write()
}
<commit_msg>Add memory scheme to live configuration<commit_after>//! # Schemes
//! A scheme is a primitive for handling filesystem syscalls in Redox.
//! Schemes accept paths from the kernel for `open`, and file descriptors that they generate
//! are then passed for operations like `close`, `read`, `write`, etc.
//!
//! The kernel validates paths and file descriptors before they are passed to schemes,
//! also stripping the scheme identifier of paths if necessary.

use alloc::arc::Arc;
use alloc::boxed::Box;
use collections::BTreeMap;
use core::sync::atomic::AtomicUsize;
use spin::{Once, RwLock, RwLockReadGuard, RwLockWriteGuard};

use syscall::error::*;
use syscall::scheme::Scheme;

use self::debug::DebugScheme;
use self::event::EventScheme;
use self::env::EnvScheme;
use self::initfs::InitFsScheme;
use self::irq::IrqScheme;
use self::memory::MemoryScheme;
use self::null::NullScheme;
use self::pipe::PipeScheme;
use self::root::RootScheme;
use self::sys::SysScheme;
use self::zero::ZeroScheme;

/// `debug:` - provides access to serial console
pub mod debug;

/// `event:` - allows reading of `Event`s which are registered using `fevent`
pub mod event;

/// `env:` - access and modify environmental variables
pub mod env;

/// `initfs:` - a readonly filesystem used for initializing the system
pub mod initfs;

/// `irq:` - allows userspace handling of IRQs
pub mod irq;

/// When compiled with "live" feature - `disk:` - embedded filesystem for live disk
#[cfg(feature="live")]
pub mod live;

/// `memory:` - a scheme for accessing physical memory
pub mod memory;

/// `null:` - a scheme that will discard all writes, and read no bytes
pub mod null;

/// `pipe:` - used internally by the kernel to implement `pipe`
pub mod pipe;

/// `:` - allows the creation of userspace schemes, tightly dependent on `user`
pub mod root;

/// `sys:` - system information, such as the context list and scheme list
pub mod sys;

/// A wrapper around userspace schemes, tightly dependent on `root`
pub mod user;

/// `zero:` - a scheme that will discard all writes, and always fill read buffers with zero
pub mod zero;

/// Limit on number of schemes
pub const SCHEME_MAX_SCHEMES: usize = 65536;

/// Unique identifier for a scheme namespace.
int_like!(SchemeNamespace, AtomicSchemeNamespace, usize, AtomicUsize);

/// Unique identifier for a scheme.
int_like!(SchemeId, AtomicSchemeId, usize, AtomicUsize);

pub const ATOMIC_SCHEMEID_INIT: AtomicSchemeId = AtomicSchemeId::default();

/// Unique identifier for a file descriptor.
int_like!(FileHandle, AtomicFileHandle, usize, AtomicUsize);

/// Scheme list type
pub struct SchemeList {
    map: BTreeMap<SchemeId, Arc<Box<Scheme + Send + Sync>>>,
    names: BTreeMap<SchemeNamespace, BTreeMap<Box<[u8]>, SchemeId>>,
    next_ns: usize,
    next_id: usize
}

impl SchemeList {
    /// Create a new scheme list.
    pub fn new() -> Self {
        let mut list = SchemeList {
            map: BTreeMap::new(),
            names: BTreeMap::new(),
            next_ns: 0,
            next_id: 1
        };
        list.new_root();
        list
    }

    /// Initialize a new namespace
    fn new_ns(&mut self) -> SchemeNamespace {
        let ns = SchemeNamespace(self.next_ns);
        self.next_ns += 1;
        self.names.insert(ns, BTreeMap::new());

        self.insert(ns, Box::new(*b""), |scheme_id| Arc::new(Box::new(RootScheme::new(ns, scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"event"), |_| Arc::new(Box::new(EventScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"env"), |_| Arc::new(Box::new(EnvScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"null"), |_| Arc::new(Box::new(NullScheme))).unwrap();
        self.insert(ns, Box::new(*b"sys"), |_| Arc::new(Box::new(SysScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"zero"), |_| Arc::new(Box::new(ZeroScheme))).unwrap();

        ns
    }

    /// Initialize the root namespace
    #[cfg(not(feature="live"))]
    fn new_root(&mut self) {
        // Do common namespace initialization
        let ns = self.new_ns();

        // Debug, Initfs and IRQ are only available in the root namespace. Pipe is special
        self.insert(ns, Box::new(*b"debug"), |scheme_id| Arc::new(Box::new(DebugScheme::new(scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"initfs"), |_| Arc::new(Box::new(InitFsScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"irq"), |scheme_id| Arc::new(Box::new(IrqScheme::new(scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"memory"), |_| Arc::new(Box::new(MemoryScheme))).unwrap();
        self.insert(ns, Box::new(*b"pipe"), |scheme_id| Arc::new(Box::new(PipeScheme::new(scheme_id)))).unwrap();
    }

    /// Initialize the root namespace - with live disk
    #[cfg(feature="live")]
    fn new_root(&mut self) {
        // Do common namespace initialization
        let ns = self.new_ns();

        // Debug, Disk, Initfs and IRQ are only available in the root namespace. Pipe is special
        self.insert(ns, Box::new(*b"debug"), |scheme_id| Arc::new(Box::new(DebugScheme::new(scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"disk"), |_| Arc::new(Box::new(self::live::DiskScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"initfs"), |_| Arc::new(Box::new(InitFsScheme::new()))).unwrap();
        self.insert(ns, Box::new(*b"irq"), |scheme_id| Arc::new(Box::new(IrqScheme::new(scheme_id)))).unwrap();
        self.insert(ns, Box::new(*b"memory"), |_| Arc::new(Box::new(MemoryScheme))).unwrap();
        self.insert(ns, Box::new(*b"pipe"), |scheme_id| Arc::new(Box::new(PipeScheme::new(scheme_id)))).unwrap();
    }

    pub fn make_ns(&mut self, from: SchemeNamespace, names: &[&[u8]]) -> Result<SchemeNamespace> {
        // Create an empty namespace
        let to = self.new_ns();

        // Copy requested scheme IDs
        for name in names.iter() {
            let id = if let Some((id, _scheme)) = self.get_name(from, name) {
                id
            } else {
                return Err(Error::new(ENODEV));
            };

            if let Some(ref mut names) = self.names.get_mut(&to) {
                assert!(names.insert(name.to_vec().into_boxed_slice(), id).is_none());
            } else {
                panic!("scheme namespace not found");
            }
        }

        Ok(to)
    }

    pub fn iter(&self) -> ::collections::btree_map::Iter<SchemeId, Arc<Box<Scheme + Send + Sync>>> {
        self.map.iter()
    }

    pub fn iter_name(&self, ns: SchemeNamespace) -> ::collections::btree_map::Iter<Box<[u8]>, SchemeId> {
        self.names[&ns].iter()
    }

    /// Get the nth scheme.
    pub fn get(&self, id: SchemeId) -> Option<&Arc<Box<Scheme + Send + Sync>>> {
        self.map.get(&id)
    }

    pub fn get_name(&self, ns: SchemeNamespace, name: &[u8]) -> Option<(SchemeId, &Arc<Box<Scheme + Send + Sync>>)> {
        if let Some(&id) = self.names[&ns].get(name) {
            self.get(id).map(|scheme| (id, scheme))
        } else {
            None
        }
    }

    /// Create a new scheme.
    pub fn insert<F>(&mut self, ns: SchemeNamespace, name: Box<[u8]>, scheme_fn: F) -> Result<SchemeId>
        where F: Fn(SchemeId) -> Arc<Box<Scheme + Send + Sync>>
    {
        if self.names[&ns].contains_key(&name) {
            return Err(Error::new(EEXIST));
        }

        if self.next_id >= SCHEME_MAX_SCHEMES {
            self.next_id = 1;
        }

        while self.map.contains_key(&SchemeId(self.next_id)) {
            self.next_id += 1;
        }

        /* Allow scheme list to grow if required
        if self.next_id >= SCHEME_MAX_SCHEMES {
            return Err(Error::new(EAGAIN));
        }
        */

        let id = SchemeId(self.next_id);
        self.next_id += 1;

        let scheme = scheme_fn(id);

        assert!(self.map.insert(id, scheme).is_none());
        if let Some(ref mut names) = self.names.get_mut(&ns) {
            assert!(names.insert(name, id).is_none());
        } else {
            panic!("scheme namespace not found");
        }
        Ok(id)
    }
}

/// Schemes list
static SCHEMES: Once<RwLock<SchemeList>> = Once::new();

/// Initialize schemes, called if needed
fn init_schemes() -> RwLock<SchemeList> {
    RwLock::new(SchemeList::new())
}

/// Get the global schemes list, const
pub fn schemes() -> RwLockReadGuard<'static, SchemeList> {
    SCHEMES.call_once(init_schemes).read()
}

/// Get the global schemes list, mutable
pub fn schemes_mut() -> RwLockWriteGuard<'static, SchemeList> {
    SCHEMES.call_once(init_schemes).write()
}
<|endoftext|>


<commit_before>use std::iter::Iterator;

use super::ast::*;
use super::ast::BuilderKind::*;
use super::ast::ScalarKind::*;
use super::ast::Type::*;
use super::ast::ExprKind::*;
use super::ast::LiteralKind::*;
use super::partial_types::*;

// TODO: These methods could take a mutable string as an argument, or even a fmt::Format.

/// A trait for printing types.
pub trait PrintableType: TypeBounds {
    fn print(&self) -> String;
}

/// Print implementation for full Types
impl PrintableType for Type {
    fn print(&self) -> String {
        match *self {
            Scalar(Bool) => "bool".to_string(),
            Scalar(I8) => "i8".to_string(),
            Scalar(I32) => "i32".to_string(),
            Scalar(I64) => "i64".to_string(),
            Scalar(F32) => "f32".to_string(),
            Scalar(F64) => "f64".to_string(),
            Vector(ref elem) => format!("vec[{}]", elem.print()),
            Dict(ref kt, ref vt) => format!("dict[{},{}]", kt.print(), vt.print()),
            Struct(ref elems) => join("{", ",", "}", elems.iter().map(|e| e.print())),
            Function(ref params, ref ret) => {
                let mut res = join("|", ",", "|(", params.iter().map(|e| e.print()));
                res.push_str(&ret.print());
                res.push_str(")");
                res
            }
            Builder(Appender(ref t)) => format!("appender[{}]", t.print()),
            Builder(DictMerger(ref kt, ref vt, op)) => {
                format!("dictmerger[{},{},{}]", kt.print(), vt.print(), op)
            }
            Builder(VecMerger(ref elem, op)) => format!("vecmerger[{},{}]", elem.print(), op),
            Builder(Merger(ref t, op)) => format!("merger[{},{}]", t.print(), op),
        }
    }
}

/// Print implementation for PartialTypes
impl PrintableType for PartialType {
    fn print(&self) -> String {
        use partial_types::PartialType::*;
        use partial_types::PartialBuilderKind::*;
        match *self {
            Unknown => "?".to_string(),
            Scalar(Bool) => "bool".to_string(),
            Scalar(I8) => "i8".to_string(),
            Scalar(I32) => "i32".to_string(),
            Scalar(I64) => "i64".to_string(),
            Scalar(F32) => "f32".to_string(),
            Scalar(F64) => "f64".to_string(),
            Vector(ref elem) => format!("vec[{}]", elem.print()),
            Dict(ref kt, ref vt) => format!("dict[{},{}]", kt.print(), vt.print()),
            Struct(ref elems) => join("{", ",", "}", elems.iter().map(|e| e.print())),
            Function(ref params, ref ret) => {
                let mut res = join("(", ",", ")=>", params.iter().map(|e| e.print()));
                res.push_str(&ret.print());
                res
            }
            Builder(Appender(ref elem)) => format!("appender[{}]", elem.print()),
            Builder(DictMerger(ref kt, ref vt, _, op)) => {
                format!("dictmerger[{},{},{}]", kt.print(), vt.print(), op)
            }
            Builder(VecMerger(ref elem, _, op)) => format!("vecmerger[{},{}]", elem.print(), op),
            Builder(Merger(ref t, op)) => format!("merger[{},{}]", t.print(), op),
        }
    }
}

/// Print a type.
pub fn print_type<T: PrintableType>(ty: &T) -> String {
    ty.print()
}

/// Print an expression concisely (without any type annotations).
pub fn print_expr<T: PrintableType>(expr: &Expr<T>) -> String {
    print_expr_impl(expr, false)
}

/// Print an expression with type annotations on all symbols.
pub fn print_typed_expr<T: PrintableType>(expr: &Expr<T>) -> String {
    print_expr_impl(expr, true)
}

fn print_iters<T: PrintableType>(iters: &Vec<Iter<T>>, typed: bool) -> String {
    let mut iter_strs = vec![];
    for iter in iters {
        if let Some(_) = iter.start {
            iter_strs.push(format!("iter({},{},{},{})",
                                   print_expr_impl(iter.data.as_ref(), typed),
                                   print_expr_impl(iter.start.as_ref().unwrap(), typed),
                                   print_expr_impl(iter.end.as_ref().unwrap(), typed),
                                   print_expr_impl(iter.stride.as_ref().unwrap(), typed)));
        } else {
            iter_strs.push(print_expr_impl(iter.data.as_ref(), typed));
        }
    }

    if iters.len() > 1 {
        format!("zip({})", iter_strs.join(","))
    } else {
        iter_strs[0].clone()
    }
}

/// Main work to print an expression.
fn print_expr_impl<T: PrintableType>(expr: &Expr<T>, typed: bool) -> String {
    match expr.kind {
        Literal(ref lit) => print_literal(lit),

        Ident(ref symbol) => {
            if typed {
                format!("{}:{}", symbol, expr.ty.print())
            } else {
                format!("{}", symbol)
            }
        }

        BinOp { kind, ref left, ref right } => {
            format!("({}{}{})",
                    print_expr_impl(left, typed),
                    kind,
                    print_expr_impl(right, typed))
        }

        Negate(ref e) => format!("(-{})", print_expr_impl(e, typed)),

        Cast { kind, ref child_expr } => {
            format!("({}({}))", kind, print_expr_impl(child_expr, typed))
        }

        ToVec { ref child_expr } => format!("toVec({})", print_expr_impl(child_expr, typed)),

        Let { ref name, ref value, ref body } => {
            if typed {
                format!("(let {}:{}=({});{})",
                        name,
                        value.ty.print(),
                        print_expr_impl(value, typed),
                        print_expr_impl(body, typed))
            } else {
                format!("(let {}=({});{})",
                        name,
                        print_expr_impl(value, typed),
                        print_expr_impl(body, typed))
            }
        }

        MakeStruct { ref elems } => {
            join("{",
                 ",",
                 "}",
                 elems.iter().map(|e| print_expr_impl(e, typed)))
        }

        MakeVector { ref elems } => {
            join("[",
                 ",",
                 "]",
                 elems.iter().map(|e| print_expr_impl(e, typed)))
        }

        Zip { ref vectors } => {
            join("zip(",
                 ",",
                 ")",
                 vectors.iter().map(|e| print_expr_impl(e, typed)))
        }

        GetField { ref expr, index } => format!("{}.${}", print_expr_impl(expr, typed), index),

        Length { ref data } => format!("len({})", print_expr_impl(data, typed)),
        Lookup { ref data, ref index } => {
            format!("lookup({},{})",
                    print_expr_impl(data, typed),
                    print_expr_impl(index, typed))
        }

        Slice { ref data, ref index, ref size } => {
            format!("slice({},{},{})",
                    print_expr_impl(data, typed),
                    print_expr_impl(index, typed),
                    print_expr_impl(size, typed))
        }

        Exp { ref value } => format!("exp({})",print_expr_impl(value, typed)),

        Lambda { ref params, ref body } => {
            let mut res = join("|",
                               ",",
                               "|",
                               params.iter().map(|e| print_parameter(e, typed)));
            res.push_str(&print_expr_impl(body, typed));
            res
        }

        NewBuilder(ref arg) => {
            match *arg {
                Some(ref e) => format!("{}({})", expr.ty.print(), print_expr_impl(e, typed)),
                None => expr.ty.print(),
            }
        }

        Res { ref builder } => format!("result({})", print_expr_impl(builder, typed)),

        Merge { ref builder, ref value } => {
            format!("merge({},{})",
                    print_expr_impl(builder, typed),
                    print_expr_impl(value, typed))
        }

        For { ref iters, ref builder, ref func } => {
            format!("for({},{},{})",
                    print_iters(iters, typed),
                    print_expr_impl(builder, typed),
                    print_expr_impl(func, typed))
        }

        If { ref cond, ref on_true, ref on_false } => {
            format!("if({},{},{})",
                    print_expr_impl(cond, typed),
                    print_expr_impl(on_true, typed),
                    print_expr_impl(on_false, typed))
        }

        Apply { ref func, ref params } => {
            let mut res = format!("({})", print_expr_impl(func, typed));
            res.push_str(&join("(",
                               ",",
                               ")",
                               params.iter().map(|e| print_expr_impl(e, typed))));
            res
        }
    }
}

/// Print a literal value.
pub fn print_literal(lit: &LiteralKind) -> String {
    match *lit {
        BoolLiteral(v) => format!("{}", v),
        I8Literal(v) => format!("{}", v),
        I32Literal(v) => format!("{}", v),
        I64Literal(v) => format!("{}L", v),
        F32Literal(v) => {
            let mut res = format!("{}", v);
            // Hack to disambiguate from integers.
            if !res.contains(".") {
                res.push_str(".0");
            }
            res.push_str("F");
            res
        }
        F64Literal(v) => {
            let mut res = format!("{}", v);
            // Hack to disambiguate from integers.
            if !res.contains(".") {
                res.push_str(".0");
            }
            res
        }
    }
}

/// Print a parameter, optionally showing its type.
fn print_parameter<T: PrintableType>(param: &Parameter<T>, typed: bool) -> String {
    if typed {
        format!("{}:{}", param.name, param.ty.print())
    } else {
        format!("{}", param.name)
    }
}

fn join<T: Iterator<Item = String>>(start: &str, sep: &str, end: &str, strings: T) -> String {
    let mut res = String::new();
    res.push_str(start);
    for (i, s) in strings.enumerate() {
        if i > 0 {
            res.push_str(sep);
        }
        res.push_str(&s);
    }
    res.push_str(end);
    res
}
<commit_msg>Improve pretty printing (#93)<commit_after>use std::iter;
use std::iter::Iterator;

use super::ast::*;
use super::ast::BuilderKind::*;
use super::ast::ScalarKind::*;
use super::ast::Type::*;
use super::ast::ExprKind::*;
use super::ast::LiteralKind::*;
use super::partial_types::*;

// TODO: These methods could take a mutable string as an argument, or even a fmt::Format.

/// A trait for printing types.
pub trait PrintableType: TypeBounds {
    fn print(&self) -> String;
}

/// Print implementation for full Types
impl PrintableType for Type {
    fn print(&self) -> String {
        match *self {
            Scalar(Bool) => "bool".to_string(),
            Scalar(I8) => "i8".to_string(),
            Scalar(I32) => "i32".to_string(),
            Scalar(I64) => "i64".to_string(),
            Scalar(F32) => "f32".to_string(),
            Scalar(F64) => "f64".to_string(),
            Vector(ref elem) => format!("vec[{}]", elem.print()),
            Dict(ref kt, ref vt) => format!("dict[{},{}]", kt.print(), vt.print()),
            Struct(ref elems) => join("{", ",", "}", elems.iter().map(|e| e.print())),
            Function(ref params, ref ret) => {
                let mut res = join("|", ",", "|(", params.iter().map(|e| e.print()));
                res.push_str(&ret.print());
                res.push_str(")");
                res
            }
            Builder(Appender(ref t)) => format!("appender[{}]", t.print()),
            Builder(DictMerger(ref kt, ref vt, op)) => {
                format!("dictmerger[{},{},{}]", kt.print(), vt.print(), op)
            }
            Builder(VecMerger(ref elem, op)) => format!("vecmerger[{},{}]", elem.print(), op),
            Builder(Merger(ref t, op)) => format!("merger[{},{}]", t.print(), op),
        }
    }
}

/// Print implementation for PartialTypes
impl PrintableType for PartialType {
    fn print(&self) -> String {
        use partial_types::PartialType::*;
        use partial_types::PartialBuilderKind::*;
        match *self {
            Unknown => "?".to_string(),
            Scalar(Bool) => "bool".to_string(),
            Scalar(I8) => "i8".to_string(),
            Scalar(I32) => "i32".to_string(),
            Scalar(I64) => "i64".to_string(),
            Scalar(F32) => "f32".to_string(),
            Scalar(F64) => "f64".to_string(),
            Vector(ref elem) => format!("vec[{}]", elem.print()),
            Dict(ref kt, ref vt) => format!("dict[{},{}]", kt.print(), vt.print()),
            Struct(ref elems) => join("{", ",", "}", elems.iter().map(|e| e.print())),
            Function(ref params, ref ret) => {
                let mut res = join("(", ",", ")=>", params.iter().map(|e| e.print()));
                res.push_str(&ret.print());
                res
            }
            Builder(Appender(ref elem)) => format!("appender[{}]", elem.print()),
            Builder(DictMerger(ref kt, ref vt, _, op)) => {
                format!("dictmerger[{},{},{}]", kt.print(), vt.print(), op)
            }
            Builder(VecMerger(ref elem, _, op)) => format!("vecmerger[{},{}]", elem.print(), op),
            Builder(Merger(ref t, op)) => format!("merger[{},{}]", t.print(), op),
        }
    }
}

/// Print a type.
pub fn print_type<T: PrintableType>(ty: &T) -> String {
    ty.print()
}

/// Print an expression concisely (without any type annotations).
pub fn print_expr<T: PrintableType>(expr: &Expr<T>) -> String {
    print_expr_impl(expr, false, 2)
}

/// Print an expression with type annotations on all symbols.
pub fn print_typed_expr<T: PrintableType>(expr: &Expr<T>) -> String {
    print_expr_impl(expr, true, 2)
}

fn print_iters<T: PrintableType>(iters: &Vec<Iter<T>>, typed: bool, indent: i32) -> String {
    let mut iter_strs = vec![];
    let less_indent_str: String = iter::repeat(" ").take((indent - 2) as usize).collect();
    let indent_str: String = iter::repeat(" ").take(indent as usize).collect();
    for iter in iters {
        if let Some(_) = iter.start {
            iter_strs.push(format!("iter({},{},{},{})",
                                   print_expr_impl(iter.data.as_ref(), typed, indent),
                                   print_expr_impl(iter.start.as_ref().unwrap(), typed, indent),
                                   print_expr_impl(iter.end.as_ref().unwrap(), typed, indent),
                                   print_expr_impl(iter.stride.as_ref().unwrap(), typed, indent)));
        } else {
            iter_strs.push(print_expr_impl(iter.data.as_ref(), typed, indent + 2));
        }
    }

    if iters.len() > 1 {
        format!("zip(\n{}{}\n{})", indent_str, iter_strs.join(&format!(",\n{}", indent_str)), less_indent_str)
    } else {
        iter_strs[0].clone()
    }
}

/// Main work to print an expression.
fn print_expr_impl<T: PrintableType>(expr: &Expr<T>, typed: bool, indent: i32) -> String {
    let less_indent_str: String = iter::repeat(" ").take((indent - 2) as usize).collect();
    let indent_str: String = iter::repeat(" ").take(indent as usize).collect();
    match expr.kind {
        Literal(ref lit) => print_literal(lit),

        Ident(ref symbol) => {
            if typed {
                format!("{}:{}", symbol, expr.ty.print())
            } else {
                format!("{}", symbol)
            }
        }

        BinOp { kind, ref left, ref right } => {
            format!("({}{}{})",
                    print_expr_impl(left, typed, indent),
                    kind,
                    print_expr_impl(right, typed, indent))
        }

        Negate(ref e) => format!("(-{})", print_expr_impl(e, typed, indent)),

        Cast { kind, ref child_expr } => {
            format!("({}({}))", kind, print_expr_impl(child_expr, typed, indent))
        }

        ToVec { ref child_expr } => format!("toVec({})", print_expr_impl(child_expr, typed, indent)),

        Let { ref name, ref value, ref body } => {
            if typed {
                format!("(let {}:{}=({});{})",
                        name,
                        value.ty.print(),
                        print_expr_impl(value, typed, indent),
                        print_expr_impl(body, typed, indent))
            } else {
                format!("(let {}=({});{})",
                        name,
                        print_expr_impl(value, typed, indent),
                        print_expr_impl(body, typed, indent))
            }
        }

        MakeStruct { ref elems } => {
            join("{",
                 ",",
                 "}",
                 elems.iter().map(|e| print_expr_impl(e, typed, indent)))
        }

        MakeVector { ref elems } => {
            join("[",
                 ",",
                 "]",
                 elems.iter().map(|e| print_expr_impl(e, typed, indent)))
        }

        Zip { ref vectors } => {
            join(&format!("zip(\n{}", indent_str),
                 &format!(",\n{}", indent_str),
                 ")",
                 vectors.iter().map(|e| print_expr_impl(e, typed, indent + 2)))
        }

        GetField { ref expr, index } => format!("{}.${}", print_expr_impl(expr, typed, indent), index),

        Length { ref data } => format!("len({})", print_expr_impl(data, typed, indent)),
        Lookup { ref data, ref index } => {
            format!("lookup({},{})",
                    print_expr_impl(data, typed, indent),
                    print_expr_impl(index, typed, indent))
        }

        Slice { ref data, ref index, ref size } => {
            format!("slice({},{},{})",
                    print_expr_impl(data, typed, indent),
                    print_expr_impl(index, typed, indent),
                    print_expr_impl(size, typed, indent))
        }

        Exp { ref value } => format!("exp({})",print_expr_impl(value, typed, indent)),

        Lambda { ref params, ref body } => {
            let mut res = join("|",
                               ",",
                               "|",
                               params.iter().map(|e| print_parameter(e, typed)));
            res.push_str(&format!("\n{}{}", indent_str, print_expr_impl(body, typed, indent + 2)));
            res
        }

        NewBuilder(ref arg) => {
            match *arg {
                Some(ref e) => format!("{}({})", expr.ty.print(), print_expr_impl(e, typed, indent)),
                None => expr.ty.print(),
            }
        }

        Res { ref builder } => {
            format!("result(\n{}{}\n{})",
                    indent_str,
                    print_expr_impl(builder, typed, indent + 2),
                    less_indent_str)
        }

        Merge { ref builder, ref value } => {
            format!("merge({},{})",
                    print_expr_impl(builder, typed, indent),
                    print_expr_impl(value, typed, indent))
        }

        For { ref iters, ref builder, ref func } => {
            format!("for(\n{}{},\n{}{},\n{}{}\n{})",
                    indent_str,
                    print_iters(iters, typed, indent + 2),
                    indent_str,
                    print_expr_impl(builder, typed, indent + 2),
                    indent_str,
                    print_expr_impl(func, typed, indent + 2),
                    less_indent_str)
        }

        If { ref cond, ref on_true, ref on_false } => {
            format!("if(\n{}{},\n{}{},\n{}{}\n{})",
                    indent_str,
                    print_expr_impl(cond, typed, indent + 2),
                    indent_str,
                    print_expr_impl(on_true, typed, indent + 2),
                    indent_str,
                    print_expr_impl(on_false, typed, indent + 2),
                    less_indent_str)
        }

        Apply { ref func, ref params } => {
            let mut res = format!("({})", print_expr_impl(func, typed, indent));
            res.push_str(&join("(",
                               ",",
                               ")",
                               params.iter().map(|e| print_expr_impl(e, typed, indent))));
            res
        }
    }
}

/// Print a literal value.
pub fn print_literal(lit: &LiteralKind) -> String {
    match *lit {
        BoolLiteral(v) => format!("{}", v),
        I8Literal(v) => format!("{}", v),
        I32Literal(v) => format!("{}", v),
        I64Literal(v) => format!("{}L", v),
        F32Literal(v) => {
            let mut res = format!("{}", v);
            // Hack to disambiguate from integers.
            if !res.contains(".") {
                res.push_str(".0");
            }
            res.push_str("F");
            res
        }
        F64Literal(v) => {
            let mut res = format!("{}", v);
            // Hack to disambiguate from integers.
            if !res.contains(".") {
                res.push_str(".0");
            }
            res
        }
    }
}

/// Print a parameter, optionally showing its type.
fn print_parameter<T: PrintableType>(param: &Parameter<T>, typed: bool) -> String {
    if typed {
        format!("{}:{}", param.name, param.ty.print())
    } else {
        format!("{}", param.name)
    }
}

fn join<T: Iterator<Item = String>>(start: &str, sep: &str, end: &str, strings: T) -> String {
    let mut res = String::new();
    res.push_str(start);
    for (i, s) in strings.enumerate() {
        if i > 0 {
            res.push_str(sep);
        }
        res.push_str(&s);
    }
    res.push_str(end);
    res
}
<|endoftext|>