<commit_before>use itertools::Itertools;
use std::iter::successors;
use syntax::{
    algo::neighbor,
    ast::{self, AstNode},
    Direction,
};

use crate::{AssistContext, AssistId, AssistKind, Assists, TextRange};

// Assist: merge_match_arms
//
// Merges the current match arm with the following if their bodies are identical.
//
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         $0Action::Move(..) => foo(),
//         Action::Stop => foo(),
//     }
// }
// ```
// ->
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         Action::Move(..) | Action::Stop => foo(),
//     }
// }
// ```
pub(crate) fn merge_match_arms(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let current_arm = ctx.find_node_at_offset::<ast::MatchArm>()?;
    // Don't try to handle arms with guards for now - can add support for this later
    if current_arm.guard().is_some() {
        return None;
    }
    let current_expr = current_arm.expr()?;
    let current_text_range = current_arm.syntax().text_range();
    let current_arm_types = get_arm_types(&ctx, &current_arm);

    // We check if the following match arms match this one. We could, but don't,
    // compare to the previous match arm as well.
    let arms_to_merge = successors(Some(current_arm), |it| neighbor(it, Direction::Next))
        .take_while(|arm| match arm.expr() {
            Some(expr) if arm.guard().is_none() => {
                let same_text = expr.syntax().text() == current_expr.syntax().text();
                if !same_text {
                    return false;
                }

                return are_same_types(&current_arm_types, arm, ctx);
            }
            _ => false,
        })
        .collect::<Vec<_>>();

    if arms_to_merge.len() <= 1 {
        return None;
    }

    acc.add(
        AssistId("merge_match_arms", AssistKind::RefactorRewrite),
        "Merge match arms",
        current_text_range,
        |edit| {
            let pats = if arms_to_merge.iter().any(contains_placeholder) {
                "_".into()
            } else {
                arms_to_merge
                    .iter()
                    .filter_map(ast::MatchArm::pat)
                    .map(|x| x.syntax().to_string())
                    .collect::<Vec<String>>()
                    .join(" | ")
            };

            let arm = format!("{} => {},", pats, current_expr.syntax().text());

            if let [first, .., last] = &*arms_to_merge {
                let start = first.syntax().text_range().start();
                let end = last.syntax().text_range().end();

                edit.replace(TextRange::new(start, end), arm);
            }
        },
    )
}

fn contains_placeholder(a: &ast::MatchArm) -> bool {
    matches!(a.pat(), Some(ast::Pat::WildcardPat(..)))
}

fn are_same_types(
    current_arm_types: &Vec<Option<hir::TypeInfo>>,
    arm: &ast::MatchArm,
    ctx: &AssistContext,
) -> bool {
    let arm_types = get_arm_types(&ctx, &arm);
    for i in 0..arm_types.len() {
        let other_arm_type = &arm_types[i];
        let current_arm_type = &current_arm_types[i];
        if let (Some(other_arm_type), Some(current_arm_type)) = (other_arm_type, current_arm_type) {
            return &other_arm_type.original == &current_arm_type.original;
        }
    }

    return true;
}

fn get_arm_types(ctx: &AssistContext, arm: &ast::MatchArm) -> Vec<Option<hir::TypeInfo>> {
    match arm.pat() {
        Some(ast::Pat::TupleStructPat(tp)) => tp
            .fields()
            .into_iter()
            .map(|field| {
                let pat_type = ctx.sema.type_of_pat(&field);
                pat_type
            })
            .collect_vec(),
        _ => Vec::new(),
    }
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn merge_match_arms_single_patterns() {
        check_assist(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X { A, B, C }

fn main() {
    let x = X::A;
    let y = match x {
        X::A => { 1i32$0 }
        X::B => { 1i32 }
        X::C => { 2i32 }
    }
}
"#,
            r#"
#[derive(Debug)]
enum X { A, B, C }

fn main() {
    let x = X::A;
    let y = match x {
        X::A | X::B => { 1i32 },
        X::C => { 2i32 }
    }
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_multiple_patterns() {
        check_assist(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A | X::B => {$0 1i32 },
        X::C | X::D => { 1i32 },
        X::E => { 2i32 },
    }
}
"#,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A | X::B | X::C | X::D => { 1i32 },
        X::E => { 2i32 },
    }
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_placeholder_pattern() {
        check_assist(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A => { 1i32 },
        X::B => { 2i$032 },
        _ => { 2i32 }
    }
}
"#,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A => { 1i32 },
        _ => { 2i32 },
    }
}
"#,
        );
    }

    #[test]
    fn merges_all_subsequent_arms() {
        check_assist(
            merge_match_arms,
            r#"
enum X { A, B, C, D, E }

fn main() {
    match X::A {
        X::A$0 => 92,
        X::B => 92,
        X::C => 92,
        X::D => 62,
        _ => panic!(),
    }
}
"#,
            r#"
enum X { A, B, C, D, E }

fn main() {
    match X::A {
        X::A | X::B | X::C => 92,
        X::D => 62,
        _ => panic!(),
    }
}
"#,
        )
    }

    #[test]
    fn merge_match_arms_rejects_guards() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X {
    A(i32),
    B,
    C
}

fn main() {
    let x = X::A;
    let y = match x {
        X::A(a) if a > 5 => { $01i32 },
        X::B => { 1i32 },
        X::C => { 2i32 }
    }
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_different_type() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<f64, f32>::Ok(0f64) {
        Ok(x) => $0x.classify(),
        Err(x) => x.classify()
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_different_type_multiple_fields() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<(f64, f64), (f32, f32)>::Ok((0f64, 0f64)) {
        Ok(x) => $0x.1.classify(),
        Err(x) => x.1.classify()
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_same_type_multiple_fields() {
        check_assist(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<(f64, f64), (f64, f64)>::Ok((0f64, 0f64)) {
        Ok(x) => $0x.1.classify(),
        Err(x) => x.1.classify()
    };
}
"#,
            r#"
fn func() {
    match Result::<(f64, f64), (f64, f64)>::Ok((0f64, 0f64)) {
        Ok(x) | Err(x) => x.1.classify(),
    };
}
"#,
        );
    }
}
<commit_msg>simplified chain<commit_after>use itertools::Itertools;
use std::iter::successors;
use syntax::{
    algo::neighbor,
    ast::{self, AstNode},
    Direction,
};

use crate::{AssistContext, AssistId, AssistKind, Assists, TextRange};

// Assist: merge_match_arms
//
// Merges the current match arm with the following if their bodies are identical.
//
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         $0Action::Move(..) => foo(),
//         Action::Stop => foo(),
//     }
// }
// ```
// ->
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         Action::Move(..) | Action::Stop => foo(),
//     }
// }
// ```
pub(crate) fn merge_match_arms(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let current_arm = ctx.find_node_at_offset::<ast::MatchArm>()?;
    // Don't try to handle arms with guards for now - can add support for this later
    if current_arm.guard().is_some() {
        return None;
    }
    let current_expr = current_arm.expr()?;
    let current_text_range = current_arm.syntax().text_range();
    let current_arm_types = get_arm_types(&ctx, &current_arm);

    // We check if the following match arms match this one. We could, but don't,
    // compare to the previous match arm as well.
    let arms_to_merge = successors(Some(current_arm), |it| neighbor(it, Direction::Next))
        .take_while(|arm| match arm.expr() {
            Some(expr) if arm.guard().is_none() => {
                let same_text = expr.syntax().text() == current_expr.syntax().text();
                if !same_text {
                    return false;
                }

                return are_same_types(&current_arm_types, arm, ctx);
            }
            _ => false,
        })
        .collect::<Vec<_>>();

    if arms_to_merge.len() <= 1 {
        return None;
    }

    acc.add(
        AssistId("merge_match_arms", AssistKind::RefactorRewrite),
        "Merge match arms",
        current_text_range,
        |edit| {
            let pats = if arms_to_merge.iter().any(contains_placeholder) {
                "_".into()
            } else {
                arms_to_merge
                    .iter()
                    .filter_map(ast::MatchArm::pat)
                    .map(|x| x.syntax().to_string())
                    .collect::<Vec<String>>()
                    .join(" | ")
            };

            let arm = format!("{} => {},", pats, current_expr.syntax().text());

            if let [first, .., last] = &*arms_to_merge {
                let start = first.syntax().text_range().start();
                let end = last.syntax().text_range().end();

                edit.replace(TextRange::new(start, end), arm);
            }
        },
    )
}

fn contains_placeholder(a: &ast::MatchArm) -> bool {
    matches!(a.pat(), Some(ast::Pat::WildcardPat(..)))
}

fn are_same_types(
    current_arm_types: &Vec<Option<hir::TypeInfo>>,
    arm: &ast::MatchArm,
    ctx: &AssistContext,
) -> bool {
    let arm_types = get_arm_types(&ctx, &arm);
    for i in 0..arm_types.len() {
        let other_arm_type = &arm_types[i];
        let current_arm_type = &current_arm_types[i];
        if let (Some(other_arm_type), Some(current_arm_type)) = (other_arm_type, current_arm_type) {
            return &other_arm_type.original == &current_arm_type.original;
        }
    }

    return true;
}

fn get_arm_types(ctx: &AssistContext, arm: &ast::MatchArm) -> Vec<Option<hir::TypeInfo>> {
    match arm.pat() {
        Some(ast::Pat::TupleStructPat(tp)) => {
            tp.fields().into_iter().map(|field| ctx.sema.type_of_pat(&field)).collect_vec()
        }
        _ => Vec::new(),
    }
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn merge_match_arms_single_patterns() {
        check_assist(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X { A, B, C }

fn main() {
    let x = X::A;
    let y = match x {
        X::A => { 1i32$0 }
        X::B => { 1i32 }
        X::C => { 2i32 }
    }
}
"#,
            r#"
#[derive(Debug)]
enum X { A, B, C }

fn main() {
    let x = X::A;
    let y = match x {
        X::A | X::B => { 1i32 },
        X::C => { 2i32 }
    }
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_multiple_patterns() {
        check_assist(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A | X::B => {$0 1i32 },
        X::C | X::D => { 1i32 },
        X::E => { 2i32 },
    }
}
"#,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A | X::B | X::C | X::D => { 1i32 },
        X::E => { 2i32 },
    }
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_placeholder_pattern() {
        check_assist(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A => { 1i32 },
        X::B => { 2i$032 },
        _ => { 2i32 }
    }
}
"#,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A => { 1i32 },
        _ => { 2i32 },
    }
}
"#,
        );
    }

    #[test]
    fn merges_all_subsequent_arms() {
        check_assist(
            merge_match_arms,
            r#"
enum X { A, B, C, D, E }

fn main() {
    match X::A {
        X::A$0 => 92,
        X::B => 92,
        X::C => 92,
        X::D => 62,
        _ => panic!(),
    }
}
"#,
            r#"
enum X { A, B, C, D, E }

fn main() {
    match X::A {
        X::A | X::B | X::C => 92,
        X::D => 62,
        _ => panic!(),
    }
}
"#,
        )
    }

    #[test]
    fn merge_match_arms_rejects_guards() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X {
    A(i32),
    B,
    C
}

fn main() {
    let x = X::A;
    let y = match x {
        X::A(a) if a > 5 => { $01i32 },
        X::B => { 1i32 },
        X::C => { 2i32 }
    }
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_different_type() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<f64, f32>::Ok(0f64) {
        Ok(x) => $0x.classify(),
        Err(x) => x.classify()
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_different_type_multiple_fields() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<(f64, f64), (f32, f32)>::Ok((0f64, 0f64)) {
        Ok(x) => $0x.1.classify(),
        Err(x) => x.1.classify()
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_same_type_multiple_fields() {
        check_assist(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<(f64, f64), (f64, f64)>::Ok((0f64, 0f64)) {
        Ok(x) => $0x.1.classify(),
        Err(x) => x.1.classify()
    };
}
"#,
            r#"
fn func() {
    match Result::<(f64, f64), (f64, f64)>::Ok((0f64, 0f64)) {
        Ok(x) | Err(x) => x.1.classify(),
    };
}
"#,
        );
    }
}
<|endoftext|>


<commit_before>#[derive(Default, Debug)]
pub struct RegStatus {
    // CU
    coproc_usability: [bool; 4],

    // RP
    low_power: bool,

    // FR
    fpregs_extend: bool,

    // RE
    reverse_endian: bool,

    // DS
    diag_status: DiagnosticStatus,

    // IM
    interrupt_mask: InterruptMask,

    // KX
    kernel_mode_64bit: bool,

    // SX
    supervisor_mode_64bit: bool,

    // UX
    user_mode_64bit: bool,

    // KSU
    mode: Mode,

    // ERL
    error_level: ErrorLevel,

    // EXL
    exception_level: ExceptionLevel,

    // IE
    interrupt_enabled: bool,
}

impl RegStatus {
    pub fn write(&mut self, data: u32) {
        let coproc = data >> 28;
        self.coproc_usability[3] = (coproc & 0b1000) != 0;
        self.coproc_usability[2] = (coproc & 0b0100) != 0;
        self.coproc_usability[1] = (coproc & 0b0010) != 0;
        self.coproc_usability[0] = (coproc & 0b0001) != 0;

        self.low_power = ((data >> 27) & 0b1) != 0;
        self.fpregs_extend = ((data >> 26) & 0b1) != 0;
        self.reverse_endian = ((data >> 25) & 0b1) != 0;

        let diag_status = ((data >> 16) & 0b111111111) as u16;
        self.diag_status.write(diag_status);
        let interrupt_mask = ((data >> 8) & 0xff) as u8;
        self.interrupt_mask.write(interrupt_mask);

        self.kernel_mode_64bit = ((data >> 7) & 0b1) != 0;
        self.supervisor_mode_64bit = ((data >> 6) & 0b1) != 0;
        self.user_mode_64bit = ((data >> 5) & 0b1) != 0;

        self.mode = data.into();

        self.error_level = data.into();
        self.exception_level = data.into();

        self.interrupt_enabled = (data & 0b1) != 0;
    }
}

#[derive(Debug, Default)]
struct DiagnosticStatus {
    // ITS
    instruction_trace_support: bool,

    // BEV
    tlb_exception_vector_location: ExceptionVectorLocation,

    // TS
    tlb_shutdown: bool,

    // SR
    soft_reset_or_nmi_occurred: bool,

    // CH
    condition_bit: bool,
}

impl DiagnosticStatus {
    fn write(&mut self, data: u16) {
        self.instruction_trace_support = 0b10000000 != 0;
        self.tlb_exception_vector_location = data.into();
        self.tlb_shutdown = 0b100000 != 0;
        self.soft_reset_or_nmi_occurred = 0b10000 != 0;
        self.condition_bit = 0b100 != 0;
    }
}

#[derive(Debug, Default)]
struct InterruptMask {
    // IM(7)
    timer_interrupt: bool,

    // IM(6:2)
    external_interrupt: [bool; 5],

    // IM(1:0)
    software_interrupt: [bool; 2],
}

impl InterruptMask {
    fn write(&mut self, data: u8) {
        self.timer_interrupt = (data & 0b10000000) != 0;
        self.external_interrupt[4] = (data & 0b01000000) != 0;
        self.external_interrupt[3] = (data & 0b00100000) != 0;
        self.external_interrupt[2] = (data & 0b00010000) != 0;
        self.external_interrupt[1] = (data & 0b00001000) != 0;
        self.external_interrupt[0] = (data & 0b00000100) != 0;
        self.software_interrupt[1] = (data & 0b00000010) != 0;
        self.software_interrupt[0] = (data & 0b00000001) != 0;
    }
}

#[derive(Debug)]
enum Mode {
    // 10 User
    User,

    // 01 Supervisor
    Supervisor,

    // 00 Kernel
    Kernel,
}

impl Default for Mode {
    fn default() -> Self {
        Mode::Kernel
    }
}

impl From<u32> for Mode {
    fn from(f: u32) -> Self {
        match (f >> 3) & 0b11 {
            0b00 => Mode::Kernel,
            0b01 => Mode::Supervisor,
            0b10 => Mode::User,
            _ => panic!("Invalid mode {:#b}", f),
        }
    }
}

#[derive(Debug)]
enum ExceptionVectorLocation {
    Normal,
    Bootstrap,
}

impl Default for ExceptionVectorLocation {
    fn default() -> Self {
        ExceptionVectorLocation::Normal
    }
}

impl From<u16> for ExceptionVectorLocation {
    fn from(f: u16) -> Self {
        if f & 0b001000000 == 0 {
            ExceptionVectorLocation::Normal
        } else {
            ExceptionVectorLocation::Bootstrap
        }
    }
}

#[derive(Debug)]
enum ErrorLevel {
    Normal,
    Error,
}

impl Default for ErrorLevel {
    fn default() -> Self {
        ErrorLevel::Normal
    }
}

impl From<u32> for ErrorLevel {
    fn from(f: u32) -> Self {
        if ((f >> 2) & 0b1) != 0 {
            ErrorLevel::Normal
        } else {
            ErrorLevel::Error
        }
    }
}

#[derive(Debug)]
enum ExceptionLevel {
    Normal,
    Exception,
}

impl Default for ExceptionLevel {
    fn default() -> Self {
        ExceptionLevel::Normal
    }
}

impl From<u32> for ExceptionLevel {
    fn from(f: u32) -> Self {
        if ((f >> 2) & 0b1) != 0 {
            ExceptionLevel::Normal
        } else {
            ExceptionLevel::Exception
        }
    }
}
<commit_msg>Tidy up DiagnosticStatus, InterruptMask and ExceptionVectorLocation<commit_after>#[derive(Default, Debug)]
pub struct RegStatus {
    // CU
    coproc_usability: [bool; 4],

    // RP
    low_power: bool,

    // FR
    fpregs_extend: bool,

    // RE
    reverse_endian: bool,

    // DS
    diag_status: DiagnosticStatus,

    // IM
    interrupt_mask: InterruptMask,

    // KX
    kernel_mode_64bit: bool,

    // SX
    supervisor_mode_64bit: bool,

    // UX
    user_mode_64bit: bool,

    // KSU
    mode: Mode,

    // ERL
    error_level: ErrorLevel,

    // EXL
    exception_level: ExceptionLevel,

    // IE
    interrupt_enabled: bool,
}

impl RegStatus {
    pub fn write(&mut self, data: u32) {
        let coproc = data >> 28;
        self.coproc_usability[3] = (coproc & 0b1000) != 0;
        self.coproc_usability[2] = (coproc & 0b0100) != 0;
        self.coproc_usability[1] = (coproc & 0b0010) != 0;
        self.coproc_usability[0] = (coproc & 0b0001) != 0;

        self.low_power = ((data >> 27) & 0b1) != 0;
        self.fpregs_extend = ((data >> 26) & 0b1) != 0;
        self.reverse_endian = ((data >> 25) & 0b1) != 0;

        let diag_status = ((data >> 16) & 0b111111111) as u16;
        self.diag_status = diag_status.into();
        let interrupt_mask = ((data >> 8) & 0xff) as u8;
        self.interrupt_mask = interrupt_mask.into();

        self.kernel_mode_64bit = ((data >> 7) & 0b1) != 0;
        self.supervisor_mode_64bit = ((data >> 6) & 0b1) != 0;
        self.user_mode_64bit = ((data >> 5) & 0b1) != 0;

        self.mode = data.into();

        self.error_level = data.into();
        self.exception_level = data.into();

        self.interrupt_enabled = (data & 0b1) != 0;
    }
}

#[derive(Debug, Default)]
struct DiagnosticStatus {
    // ITS
    instruction_trace_support: bool,

    // BEV
    tlb_exception_vector_location: TLBExceptionVectorLocation,

    // TS
    tlb_shutdown: bool,

    // SR
    soft_reset_or_nmi_occurred: bool,

    // CH
    condition_bit: bool,
}

impl From<u16> for DiagnosticStatus {
    fn from(f: u16) -> Self {
        DiagnosticStatus {
            instruction_trace_support: 0b10000000 != 0,
            tlb_exception_vector_location: f.into(),
            tlb_shutdown: 0b100000 != 0,
            soft_reset_or_nmi_occurred: 0b10000 != 0,
            condition_bit: 0b100 != 0,
        }
    }
}

#[derive(Debug, Default)]
struct InterruptMask {
    // IM(7)
    timer_interrupt: bool,

    // IM(6:2)
    external_interrupt: [bool; 5],

    // IM(1:0)
    software_interrupt: [bool; 2],
}

impl From<u8> for InterruptMask {
    fn from(data: u8) -> Self {
        InterruptMask {
            timer_interrupt: (data & 0b10000000) != 0,
            external_interrupt: [(data & 0b00000100) != 0,
                                 (data & 0b00001000) != 0,
                                 (data & 0b00010000) != 0,
                                 (data & 0b00100000) != 0,
                                 (data & 0b01000000) != 0],
            software_interrupt: [(data & 0b00000001) != 0, (data & 0b00000010) != 0],
        }
    }
}

#[derive(Debug)]
enum Mode {
    // 10 User
    User,

    // 01 Supervisor
    Supervisor,

    // 00 Kernel
    Kernel,
}

impl Default for Mode {
    fn default() -> Self {
        Mode::Kernel
    }
}

impl From<u32> for Mode {
    fn from(f: u32) -> Self {
        match (f >> 3) & 0b11 {
            0b00 => Mode::Kernel,
            0b01 => Mode::Supervisor,
            0b10 => Mode::User,
            _ => panic!("Invalid mode {:#b}", f),
        }
    }
}

#[derive(Debug)]
enum TLBExceptionVectorLocation {
    Normal,
    Bootstrap,
}

impl Default for TLBExceptionVectorLocation {
    fn default() -> Self {
        TLBExceptionVectorLocation::Normal
    }
}

impl From<u16> for TLBExceptionVectorLocation {
    fn from(f: u16) -> Self {
        match f & 0b001000000 {
            0 => TLBExceptionVectorLocation::Normal,
            1 => TLBExceptionVectorLocation::Bootstrap,
            _ => unreachable!(),
        }
    }
}

#[derive(Debug)]
enum ErrorLevel {
    Normal,
    Error,
}

impl Default for ErrorLevel {
    fn default() -> Self {
        ErrorLevel::Normal
    }
}

impl From<u32> for ErrorLevel {
    fn from(f: u32) -> Self {
        if ((f >> 2) & 0b1) != 0 {
            ErrorLevel::Normal
        } else {
            ErrorLevel::Error
        }
    }
}

#[derive(Debug)]
enum ExceptionLevel {
    Normal,
    Exception,
}

impl Default for ExceptionLevel {
    fn default() -> Self {
        ExceptionLevel::Normal
    }
}

impl From<u32> for ExceptionLevel {
    fn from(f: u32) -> Self {
        if ((f >> 2) & 0b1) != 0 {
            ExceptionLevel::Normal
        } else {
            ExceptionLevel::Exception
        }
    }
}
<|endoftext|>


<commit_before>use definitions::{ExternalInterrupt, InternalInterrupt};
use definitions::Config;
use definitions::Instruction;
use definitions::Program;
use definitions::Target;
use definitions::Value;
use definitions::typedef::*;
use error::*;
use std::collections::{BTreeMap, LinkedList};
use std::sync::mpsc::{Receiver, SyncSender, TrySendError};
use std::thread::{self, JoinHandle};

pub fn start(program: Program, sender: SyncSender<Frame>, receiver: Receiver<ExternalInterrupt>)
    -> JoinHandle<()> {
    thread::spawn(
        move || {
            if let Err(ref e) = VM::default().exec(program, sender, receiver) {
                println!("error: {}", e);

                for e in e.iter().skip(1) {
                    println!("caused by: {}", e);
                }

                // The backtrace is not always generated. Try to run this example
                // with `RUST_BACKTRACE=1`.
                if let Some(backtrace) = e.backtrace() {
                    println!("backtrace: {:?}", backtrace);
                }
            }
        },
    )
}

/// Since rusts `std:::cmp::Ordering` doesn't implement serialization, we have
/// to do this
#[derive(Serialize, Deserialize, Debug, PartialEq)]
enum Ordering {
    Less,
    Greater,
    Equal,
}

/// The whole state of the VM
#[derive(Serialize, Deserialize, Default, Debug)]
struct VM {
    /// The instructions that are currently executed
    image_data: Vec<Instruction>,
    /// The current program counter
    pc: Address,
    stack: LinkedList<Value>,
    val_index: BTreeMap<Address, Value>,
    input_register: Integer,
    framebuffer: Frame,
    framebuffer_invalid: bool,
    /// A register for holding information about a recent comparison
    cmp_register: Option<Ordering>,
    /// A stack to hold the return addresses of function calls
    call_stack: LinkedList<Address>,
    /// A boolean used for locking the program counter
    pc_locked: bool,
    /// The configuration of the VM
    config: Config,
    halted: bool,
}

impl VM {
    // # Maintainance functions

    /// Executes the given program
    pub fn exec(
        &mut self, program: Program, sender: SyncSender<Frame>,
        receiver: Receiver<ExternalInterrupt>
    ) -> VMResult<()> {
        self.reset();
        self.load_program(program)?;
        self.build_framebuffer();

        while (self.pc < self.image_data.len()) && !self.halted {
            self.external_interrupt(&receiver)?;

            self.do_cycle()?;

            self.flush_framebuffer(&sender)?;

            thread::yield_now();
        }

        Ok(())
    }

    /// Run one instruction cycle
    fn do_cycle(&mut self) -> VMResult<()> {
        let current_instruction = self.image_data[self.pc].clone();
        self.handle_instruction(current_instruction)?;
        self.advance_pc();

        Ok(())
    }

    /// Handles a single instruction
    fn handle_instruction(&mut self, instruction: Instruction) -> VMResult<()> {
        match instruction {
            Instruction::Halt => self.halt(),
            Instruction::Int(interrupt) => self.int(&interrupt),

            Instruction::Add(dest, src) => self.add(&dest, &src)?,
            Instruction::Sub(dest, src) => self.sub(&dest, &src)?,
            Instruction::Div(dest, src) => self.div(&dest, &src)?,
            Instruction::Mul(dest, src) => self.mul(&dest, &src)?,
            Instruction::Rem(dest, src) => self.rem(&dest, &src)?,

            Instruction::Cmp(target_a, target_b) => self.cmp(&target_a, &target_b)?,
            Instruction::Jmp(addr) => self.jmp(&addr),
            Instruction::JmpLt(addr) => self.jmp_lt(&addr),
            Instruction::JmpGt(addr) => self.jmp_gt(&addr),
            Instruction::JmpEq(addr) => self.jmp_eq(&addr),
            Instruction::JmpLtEq(addr) => self.jmp_lt_eq(&addr),
            Instruction::JmpGtEq(addr) => self.jmp_gt_eq(&addr),

            Instruction::Push(dest, value) => self.push(&dest, value)?,
            Instruction::Mov(dest, src) => self.mov(&dest, &src)?,
            Instruction::Swp(target_a, target_b) => self.swp(&target_a, &target_b)?,

            Instruction::Call(addr) => self.call(&addr),
            Instruction::Ret => self.ret()?,
        }

        Ok(())
    }

    /// Loads the instructions of the given program to the VM's state
    fn load_program(&mut self, program: Program) -> VMResult<()> {
        let orig_program = Program::default();
        if program.preamble != orig_program.preamble {
            bail!("invalid preamble");
        } else if program.version != orig_program.version {
            bail!("invalid version");
        } else {
            self.image_data = program.instructions;
            self.config = program.config;
            Ok(())
        }
    }

    /// Aborts the execution of the current image
    fn halt(&mut self) {
        self.halted = true;
    }

    /// Handles an internal interrupt
    fn int(&mut self, interrupt: &InternalInterrupt) {
        match interrupt {
            &InternalInterrupt::FlushFramebuffer => self.invalidate_framebuffer(),
        }
    }

    /// Handles incoming interrupts or moves along
    fn external_interrupt(&mut self, receiver: &Receiver<ExternalInterrupt>) -> VMResult<()> {
        match receiver.try_recv() {
            Ok(interrupt) => {
                match interrupt {
                    ExternalInterrupt::Halt => self.halt(),
                    ExternalInterrupt::KeyDown(value) => {
                        self.input_register = value;
                    }
                    ExternalInterrupt::KeyUp => self.input_register = 0,
                    _ => {}
                }

                Ok(())
            }
            _ => Ok(()),
        }
    }

    /// Flushes the internal framebuffer using the given sender
    fn flush_framebuffer(&mut self, sender: &SyncSender<Frame>) -> VMResult<()> {
        if self.framebuffer_invalid {
            if let Err(TrySendError::Disconnected(..)) = sender.try_send(self.framebuffer.clone()) {
                bail!("output channel disconnected");
            }

            self.framebuffer_invalid = false;
        }

        Ok(())
    }

    /// Allocates all the needed space in the framebuffer
    fn build_framebuffer(&mut self) {
        let ref resolution = self.config.display.resolution;
        let allocation_space = resolution.width * resolution.height;
        self.framebuffer = vec![Color::default(); allocation_space];
    }

    /// Resets the VM to a clean state
    fn reset(&mut self) {
        *self = VM::default();
    }

    /// Locks the program counter in place
    fn lock_pc(&mut self) {
        self.pc_locked = true;
    }

    /// Advances the program counter
    fn advance_pc(&mut self) {
        if self.pc_locked {
            self.pc_locked = false;
        } else {
            self.pc += 1;
        }
    }

    /// Invalidates the frambuffer causing it to be resent to the display
    /// receiver
    fn invalidate_framebuffer(&mut self) {
        self.framebuffer_invalid = true;
    }

    /// Return the value at the specified target. The value of the target will
    /// be consumed
    fn pop(&mut self, target: &Target) -> VMResult<Value> {
        match target {
            &Target::ValueIndex(index) => {
                if let Some(value) = self.val_index.remove(&index) {
                    Ok(value)
                } else {
                    bail!("no value found at index {}", index);
                }
            }
            &Target::Stack => {
                if let Some(value) = self.stack.pop_front() {
                    Ok(value)
                } else {
                    bail!("unable to pop value off an empty stack");
                }
            }
            &Target::Framebuffer(index) => {
                if let Some(&(r, g, b)) = self.framebuffer.get(index) {
                    Ok(Value::Color(r, g, b))
                } else {
                    bail!("no value found in framebuffer at index {}", index);
                }
            }
            &Target::InputRegister => Ok(Value::Integer(self.input_register)),
        }
    }

    // # Instruction functions

    /// Adds the value of the src target to the value of the dest target
    fn add(&mut self, dest: &Target, src: &Target) -> VMResult<()> {
        let dest_value = self.pop(dest)?;
        let src_value = self.pop(src)?;

        self.push(dest, dest_value + src_value)?;

        Ok(())
    }

    /// Subtracts the value of the src target from the value of the dest target
    fn sub(&mut self, dest: &Target, src: &Target) -> VMResult<()> {
        let dest_value = self.pop(dest)?;
        let src_value = self.pop(src)?;

        self.push(dest, dest_value - src_value)?;

        Ok(())
    }

    /// Divides the value of the dest target through the value of the src target
    fn div(&mut self, dest: &Target, src: &Target) -> VMResult<()> {
        let dest_value = self.pop(dest)?;
        let src_value = self.pop(src)?;

        self.push(dest, dest_value / src_value)?;

        Ok(())
    }

    /// Multiplies the value of the dest target with the value of the src target
    fn mul(&mut self, dest: &Target, src: &Target) -> VMResult<()> {
        let dest_value = self.pop(dest)?;
        let src_value = self.pop(src)?;

        self.push(dest, dest_value * src_value)?;

        Ok(())
    }

    /// Applies the modulo operator on the value of the dest target using the
    /// value of the src target
    fn rem(&mut self, dest: &Target, src: &Target) -> VMResult<()> {
        let dest_value = self.pop(dest)?;
        let src_value = self.pop(src)?;

        self.push(dest, dest_value * src_value)?;

        Ok(())
    }

    /// Compares the top values of the two targets and saves the result to
    /// `self.cmp_register`
    fn cmp(&mut self, target_a: &Target, target_b: &Target) -> VMResult<()> {
        let target_a_value = self.pop(target_a)?;
        let target_b_value = self.pop(target_b)?;

        if target_a_value < target_b_value {
            self.cmp_register = Some(Ordering::Less);
        } else if target_a_value > target_b_value {
            self.cmp_register = Some(Ordering::Greater);
        } else if target_a_value == target_b_value {
            self.cmp_register = Some(Ordering::Equal);
        }

        Ok(())
    }

    /// Jumps unconditionally to the specified address
    fn jmp(&mut self, addr: &Address) {
        self.pc = *addr;
        self.lock_pc();
    }

    /// Jumps if the last compare got the result `Some(Ordering::Less)`
    fn jmp_lt(&mut self, addr: &Address) {
        if self.cmp_register == Some(Ordering::Less) {
            self.jmp(addr);
        }
    }

    /// Jumps if the last compare got the result `Some(Ordering::Greater)`
    fn jmp_gt(&mut self, addr: &Address) {
        if self.cmp_register == Some(Ordering::Greater) {
            self.jmp(addr);
        }
    }

    /// Jumps if the last compare got the result `Some(Ordering::Equal)`
    fn jmp_eq(&mut self, addr: &Address) {
        if self.cmp_register == Some(Ordering::Equal) {
            self.jmp(addr);
        }
    }

    /// Jumps if the last compare got the result `Some(Ordering::Less)` or
    /// `Some(Ordering::Equal)`
    fn jmp_lt_eq(&mut self, addr: &Address) {
        if (self.cmp_register == Some(Ordering::Less)) ||
           (self.cmp_register == Some(Ordering::Equal)) {
            self.jmp(addr);
        }
    }

    /// Jumps if the last compare got the result `Some(Ordering::Greater)` or
    /// `Some(Ordering::Equal)`
    fn jmp_gt_eq(&mut self, addr: &Address) {
        if (self.cmp_register == Some(Ordering::Greater)) ||
           (self.cmp_register == Some(Ordering::Equal)) {
            self.jmp(addr);
        }
    }

    /// Pushes the given value to the given target
    fn push(&mut self, dest: &Target, value: Value) -> VMResult<()> {
        match dest {
            &Target::ValueIndex(index) => {
                self.val_index.entry(index).or_insert(value);
                Ok(())
            }
            &Target::Stack => {
                self.stack.push_front(value);
                Ok(())
            }
            &Target::Framebuffer(index) => {
                if let Value::Color(r, g, b) = value {
                    self.framebuffer[index] = (r, g, b);
                    Ok(())
                } else {
                    bail!("unable push a non-color value to the framebuffer");
                }
            }
            &Target::InputRegister => {
                bail!("unable push to read-only input register");
            }
        }
    }

    /// Moves the top value of the src target to the dest target
    fn mov(&mut self, dest: &Target, src: &Target) -> VMResult<()> {
        let src_value = self.pop(src)?;
        self.push(dest, src_value)?;

        Ok(())
    }

    /// Swaps the top values of the targets
    fn swp(&mut self, target_a: &Target, target_b: &Target) -> VMResult<()> {
        let a_value = self.pop(target_a)?;
        let b_value = self.pop(target_b)?;

        self.push(target_a, b_value)?;
        self.push(target_b, a_value)?;

        Ok(())
    }

    /// Calls the function at the specified address saving the return address
    /// to the call stack
    fn call(&mut self, addr: &Address) {
        self.call_stack.push_front(self.pc + 1);
        self.jmp(addr);
    }

    /// Returns from an ongoing function call
    fn ret(&mut self) -> VMResult<()> {
        if let Some(retur_addr) = self.call_stack.pop_front() {
            self.jmp(&retur_addr);
        } else {
            bail!("unable to return from an empty call stack");
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use definitions::ImageBuilder;
    use rand;

    #[test]
    fn halt() {
        let mut vm = VM::default();
        vm.handle_instruction(Instruction::Halt).expect("failed to handle instruction");

        println!("{:#?}", vm);

        assert!(vm.halted);
    }

    #[test]
    fn add_stack() {
        for _ in 0..3000 {
            let val_a = rand::random::<Integer>() / 2;
            let val_b = rand::random::<Integer>() / 2;

            let mut vm = VM::default();

            let mut builder = ImageBuilder::new();
            builder.push(Target::Stack, Value::Integer(val_a));
            builder.push(Target::Stack, Value::Integer(val_b));
            builder.add(Target::Stack, Target::Stack);

            let program = builder.gen_program();
            vm.load_program(program).unwrap();

            vm.do_cycle().unwrap();
            vm.do_cycle().unwrap();
            vm.do_cycle().unwrap();

            let stack_value = vm.pop(&Target::Stack).unwrap();

            assert_eq!(stack_value, Value::Integer(val_a + val_b));
        }
    }
}
<commit_msg>Added subtraction test<commit_after>use definitions::{ExternalInterrupt, InternalInterrupt};
use definitions::Config;
use definitions::Instruction;
use definitions::Program;
use definitions::Target;
use definitions::Value;
use definitions::typedef::*;
use error::*;
use std::collections::{BTreeMap, LinkedList};
use std::sync::mpsc::{Receiver, SyncSender, TrySendError};
use std::thread::{self, JoinHandle};

pub fn start(program: Program, sender: SyncSender<Frame>, receiver: Receiver<ExternalInterrupt>)
    -> JoinHandle<()> {
    thread::spawn(
        move || {
            if let Err(ref e) = VM::default().exec(program, sender, receiver) {
                println!("error: {}", e);

                for e in e.iter().skip(1) {
                    println!("caused by: {}", e);
                }

                // The backtrace is not always generated. Try to run this example
                // with `RUST_BACKTRACE=1`.
                if let Some(backtrace) = e.backtrace() {
                    println!("backtrace: {:?}", backtrace);
                }
            }
        },
    )
}

/// Since rusts `std:::cmp::Ordering` doesn't implement serialization, we have
/// to do this
#[derive(Serialize, Deserialize, Debug, PartialEq)]
enum Ordering {
    Less,
    Greater,
    Equal,
}

/// The whole state of the VM
#[derive(Serialize, Deserialize, Default, Debug)]
struct VM {
    /// The instructions that are currently executed
    image_data: Vec<Instruction>,
    /// The current program counter
    pc: Address,
    stack: LinkedList<Value>,
    val_index: BTreeMap<Address, Value>,
    input_register: Integer,
    framebuffer: Frame,
    framebuffer_invalid: bool,
    /// A register for holding information about a recent comparison
    cmp_register: Option<Ordering>,
    /// A stack to hold the return addresses of function calls
    call_stack: LinkedList<Address>,
    /// A boolean used for locking the program counter
    pc_locked: bool,
    /// The configuration of the VM
    config: Config,
    halted: bool,
}

impl VM {
    // # Maintainance functions

    /// Executes the given program
    pub fn exec(
        &mut self, program: Program, sender: SyncSender<Frame>,
        receiver: Receiver<ExternalInterrupt>
    ) -> VMResult<()> {
        self.reset();
        self.load_program(program)?;
        self.build_framebuffer();

        while (self.pc < self.image_data.len()) && !self.halted {
            self.external_interrupt(&receiver)?;

            self.do_cycle()?;

            self.flush_framebuffer(&sender)?;

            thread::yield_now();
        }

        Ok(())
    }

    /// Run one instruction cycle
    fn do_cycle(&mut self) -> VMResult<()> {
        let current_instruction = self.image_data[self.pc].clone();
        self.handle_instruction(current_instruction)?;
        self.advance_pc();

        Ok(())
    }

    /// Handles a single instruction
    fn handle_instruction(&mut self, instruction: Instruction) -> VMResult<()> {
        match instruction {
            Instruction::Halt => self.halt(),
            Instruction::Int(interrupt) => self.int(&interrupt),

            Instruction::Add(dest, src) => self.add(&dest, &src)?,
            Instruction::Sub(dest, src) => self.sub(&dest, &src)?,
            Instruction::Div(dest, src) => self.div(&dest, &src)?,
            Instruction::Mul(dest, src) => self.mul(&dest, &src)?,
            Instruction::Rem(dest, src) => self.rem(&dest, &src)?,

            Instruction::Cmp(target_a, target_b) => self.cmp(&target_a, &target_b)?,
            Instruction::Jmp(addr) => self.jmp(&addr),
            Instruction::JmpLt(addr) => self.jmp_lt(&addr),
            Instruction::JmpGt(addr) => self.jmp_gt(&addr),
            Instruction::JmpEq(addr) => self.jmp_eq(&addr),
            Instruction::JmpLtEq(addr) => self.jmp_lt_eq(&addr),
            Instruction::JmpGtEq(addr) => self.jmp_gt_eq(&addr),

            Instruction::Push(dest, value) => self.push(&dest, value)?,
            Instruction::Mov(dest, src) => self.mov(&dest, &src)?,
            Instruction::Swp(target_a, target_b) => self.swp(&target_a, &target_b)?,

            Instruction::Call(addr) => self.call(&addr),
            Instruction::Ret => self.ret()?,
        }

        Ok(())
    }

    /// Loads the instructions of the given program to the VM's state
    fn load_program(&mut self, program: Program) -> VMResult<()> {
        let orig_program = Program::default();
        if program.preamble != orig_program.preamble {
            bail!("invalid preamble");
        } else if program.version != orig_program.version {
            bail!("invalid version");
        } else {
            self.image_data = program.instructions;
            self.config = program.config;
            Ok(())
        }
    }

    /// Aborts the execution of the current image
    fn halt(&mut self) {
        self.halted = true;
    }

    /// Handles an internal interrupt
    fn int(&mut self, interrupt: &InternalInterrupt) {
        match interrupt {
            &InternalInterrupt::FlushFramebuffer => self.invalidate_framebuffer(),
        }
    }

    /// Handles incoming interrupts or moves along
    fn external_interrupt(&mut self, receiver: &Receiver<ExternalInterrupt>) -> VMResult<()> {
        match receiver.try_recv() {
            Ok(interrupt) => {
                match interrupt {
                    ExternalInterrupt::Halt => self.halt(),
                    ExternalInterrupt::KeyDown(value) => {
                        self.input_register = value;
                    }
                    ExternalInterrupt::KeyUp => self.input_register = 0,
                    _ => {}
                }

                Ok(())
            }
            _ => Ok(()),
        }
    }

    /// Flushes the internal framebuffer using the given sender
    fn flush_framebuffer(&mut self, sender: &SyncSender<Frame>) -> VMResult<()> {
        if self.framebuffer_invalid {
            if let Err(TrySendError::Disconnected(..)) = sender.try_send(self.framebuffer.clone()) {
                bail!("output channel disconnected");
            }

            self.framebuffer_invalid = false;
        }

        Ok(())
    }

    /// Allocates all the needed space in the framebuffer
    fn build_framebuffer(&mut self) {
        let ref resolution = self.config.display.resolution;
        let allocation_space = resolution.width * resolution.height;
        self.framebuffer = vec![Color::default(); allocation_space];
    }

    /// Resets the VM to a clean state
    fn reset(&mut self) {
        *self = VM::default();
    }

    /// Locks the program counter in place
    fn lock_pc(&mut self) {
        self.pc_locked = true;
    }

    /// Advances the program counter
    fn advance_pc(&mut self) {
        if self.pc_locked {
            self.pc_locked = false;
        } else {
            self.pc += 1;
        }
    }

    /// Invalidates the frambuffer causing it to be resent to the display
    /// receiver
    fn invalidate_framebuffer(&mut self) {
        self.framebuffer_invalid = true;
    }

    /// Return the value at the specified target. The value of the target will
    /// be consumed
    fn pop(&mut self, target: &Target) -> VMResult<Value> {
        match target {
            &Target::ValueIndex(index) => {
                if let Some(value) = self.val_index.remove(&index) {
                    Ok(value)
                } else {
                    bail!("no value found at index {}", index);
                }
            }
            &Target::Stack => {
                if let Some(value) = self.stack.pop_front() {
                    Ok(value)
                } else {
                    bail!("unable to pop value off an empty stack");
                }
            }
            &Target::Framebuffer(index) => {
                if let Some(&(r, g, b)) = self.framebuffer.get(index) {
                    Ok(Value::Color(r, g, b))
                } else {
                    bail!("no value found in framebuffer at index {}", index);
                }
            }
            &Target::InputRegister => Ok(Value::Integer(self.input_register)),
        }
    }

    // # Instruction functions

    /// Adds the value of the src target to the value of the dest target
    fn add(&mut self, dest: &Target, src: &Target) -> VMResult<()> {
        let dest_value = self.pop(dest)?;
        let src_value = self.pop(src)?;

        self.push(dest, dest_value + src_value)?;

        Ok(())
    }

    /// Subtracts the value of the src target from the value of the dest target
    fn sub(&mut self, dest: &Target, src: &Target) -> VMResult<()> {
        let dest_value = self.pop(dest)?;
        let src_value = self.pop(src)?;

        self.push(dest, dest_value - src_value)?;

        Ok(())
    }

    /// Divides the value of the dest target through the value of the src target
    fn div(&mut self, dest: &Target, src: &Target) -> VMResult<()> {
        let dest_value = self.pop(dest)?;
        let src_value = self.pop(src)?;

        self.push(dest, dest_value / src_value)?;

        Ok(())
    }

    /// Multiplies the value of the dest target with the value of the src target
    fn mul(&mut self, dest: &Target, src: &Target) -> VMResult<()> {
        let dest_value = self.pop(dest)?;
        let src_value = self.pop(src)?;

        self.push(dest, dest_value * src_value)?;

        Ok(())
    }

    /// Applies the modulo operator on the value of the dest target using the
    /// value of the src target
    fn rem(&mut self, dest: &Target, src: &Target) -> VMResult<()> {
        let dest_value = self.pop(dest)?;
        let src_value = self.pop(src)?;

        self.push(dest, dest_value * src_value)?;

        Ok(())
    }

    /// Compares the top values of the two targets and saves the result to
    /// `self.cmp_register`
    fn cmp(&mut self, target_a: &Target, target_b: &Target) -> VMResult<()> {
        let target_a_value = self.pop(target_a)?;
        let target_b_value = self.pop(target_b)?;

        if target_a_value < target_b_value {
            self.cmp_register = Some(Ordering::Less);
        } else if target_a_value > target_b_value {
            self.cmp_register = Some(Ordering::Greater);
        } else if target_a_value == target_b_value {
            self.cmp_register = Some(Ordering::Equal);
        }

        Ok(())
    }

    /// Jumps unconditionally to the specified address
    fn jmp(&mut self, addr: &Address) {
        self.pc = *addr;
        self.lock_pc();
    }

    /// Jumps if the last compare got the result `Some(Ordering::Less)`
    fn jmp_lt(&mut self, addr: &Address) {
        if self.cmp_register == Some(Ordering::Less) {
            self.jmp(addr);
        }
    }

    /// Jumps if the last compare got the result `Some(Ordering::Greater)`
    fn jmp_gt(&mut self, addr: &Address) {
        if self.cmp_register == Some(Ordering::Greater) {
            self.jmp(addr);
        }
    }

    /// Jumps if the last compare got the result `Some(Ordering::Equal)`
    fn jmp_eq(&mut self, addr: &Address) {
        if self.cmp_register == Some(Ordering::Equal) {
            self.jmp(addr);
        }
    }

    /// Jumps if the last compare got the result `Some(Ordering::Less)` or
    /// `Some(Ordering::Equal)`
    fn jmp_lt_eq(&mut self, addr: &Address) {
        if (self.cmp_register == Some(Ordering::Less)) ||
           (self.cmp_register == Some(Ordering::Equal)) {
            self.jmp(addr);
        }
    }

    /// Jumps if the last compare got the result `Some(Ordering::Greater)` or
    /// `Some(Ordering::Equal)`
    fn jmp_gt_eq(&mut self, addr: &Address) {
        if (self.cmp_register == Some(Ordering::Greater)) ||
           (self.cmp_register == Some(Ordering::Equal)) {
            self.jmp(addr);
        }
    }

    /// Pushes the given value to the given target
    fn push(&mut self, dest: &Target, value: Value) -> VMResult<()> {
        match dest {
            &Target::ValueIndex(index) => {
                self.val_index.entry(index).or_insert(value);
                Ok(())
            }
            &Target::Stack => {
                self.stack.push_front(value);
                Ok(())
            }
            &Target::Framebuffer(index) => {
                if let Value::Color(r, g, b) = value {
                    self.framebuffer[index] = (r, g, b);
                    Ok(())
                } else {
                    bail!("unable push a non-color value to the framebuffer");
                }
            }
            &Target::InputRegister => {
                bail!("unable push to read-only input register");
            }
        }
    }

    /// Moves the top value of the src target to the dest target
    fn mov(&mut self, dest: &Target, src: &Target) -> VMResult<()> {
        let src_value = self.pop(src)?;
        self.push(dest, src_value)?;

        Ok(())
    }

    /// Swaps the top values of the targets
    fn swp(&mut self, target_a: &Target, target_b: &Target) -> VMResult<()> {
        let a_value = self.pop(target_a)?;
        let b_value = self.pop(target_b)?;

        self.push(target_a, b_value)?;
        self.push(target_b, a_value)?;

        Ok(())
    }

    /// Calls the function at the specified address saving the return address
    /// to the call stack
    fn call(&mut self, addr: &Address) {
        self.call_stack.push_front(self.pc + 1);
        self.jmp(addr);
    }

    /// Returns from an ongoing function call
    fn ret(&mut self) -> VMResult<()> {
        if let Some(retur_addr) = self.call_stack.pop_front() {
            self.jmp(&retur_addr);
        } else {
            bail!("unable to return from an empty call stack");
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use definitions::ImageBuilder;
    use rand;

    #[test]
    fn halt() {
        let mut vm = VM::default();
        vm.handle_instruction(Instruction::Halt).expect("failed to handle instruction");

        println!("{:#?}", vm);

        assert!(vm.halted);
    }

    #[test]
    fn add_stack() {
        for _ in 0..3000 {
            let val_a = rand::random::<Integer>() / 2;
            let val_b = rand::random::<Integer>() / 2;

            let mut vm = VM::default();

            let mut builder = ImageBuilder::new();
            builder.push(Target::Stack, Value::Integer(val_a));
            builder.push(Target::Stack, Value::Integer(val_b));
            builder.add(Target::Stack, Target::Stack);

            let program = builder.gen_program();
            vm.load_program(program).unwrap();

            vm.do_cycle().unwrap();
            vm.do_cycle().unwrap();
            vm.do_cycle().unwrap();

            let stack_value = vm.pop(&Target::Stack).unwrap();

            assert_eq!(stack_value, Value::Integer(val_b + val_a));
        }
    }

    #[test]
    fn sub_stack() {
        for _ in 0..3000 {
            let val_a = rand::random::<Integer>() / 2;
            let val_b = rand::random::<Integer>() / 2;

            let mut vm = VM::default();

            let mut builder = ImageBuilder::new();
            builder.push(Target::Stack, Value::Integer(val_a));
            builder.push(Target::Stack, Value::Integer(val_b));
            builder.sub(Target::Stack, Target::Stack);

            let program = builder.gen_program();
            vm.load_program(program).unwrap();

            vm.do_cycle().unwrap();
            vm.do_cycle().unwrap();
            vm.do_cycle().unwrap();

            let stack_value = vm.pop(&Target::Stack).unwrap();

            assert_eq!(stack_value, Value::Integer(val_b - val_a));
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;

use condition::Condition;
use dispatcher::demux::Demultiplexer;
use dispatcher::handlers;
use dispatcher::request::{Request, RequestHandler};
use reactor::{Event, EventDemultiplexer, EventHandler, Reactor};

pub struct RequestReactor {
    handlers: BTreeMap<RequestHandler, Box<EventHandler<Request, Handler=RequestHandler>>>,
    demultiplexer: Demultiplexer<Request>,
    exit_condition: Condition
}

impl RequestReactor {
    fn new(demultiplexer: Demultiplexer<Request>) -> RequestReactor {
        let exit_condition = Condition::new(false);
        let exit_handler = Box::new(handlers::exit::ExitHandler::new(exit_condition.clone()));
        let event_handler = Box::new(handlers::event::EventHandler::new());

        let mut reactor = RequestReactor {
            demultiplexer: demultiplexer,
            exit_condition: exit_condition,
            handlers: BTreeMap::new()
        };

        reactor.register_handler(exit_handler);
        reactor.register_handler(event_handler);
        reactor
    }
}

impl Reactor for RequestReactor {
    type Event = Request;
    type Handler = RequestHandler;
    fn handle_events(&mut self) {
        while !self.exit_condition.is_active() {
            if let Some(request) = self.demultiplexer.select() {
                let mut handler = self.handlers.get_mut(&request.handler()).unwrap();
                handler.handle_event(request);
            } else {
                break;
            }
        }
    }
    fn register_handler(&mut self, handler: Box<EventHandler<Self::Event, Handler=RequestHandler>>) {
        self.handlers.insert(handler.handler(), handler);
    }
    fn remove_handler(&mut self, handler: &EventHandler<Self::Event, Handler=RequestHandler>){}
}
<commit_msg>RequestReactor: make new() public<commit_after>use std::collections::BTreeMap;

use condition::Condition;
use dispatcher::demux::Demultiplexer;
use dispatcher::handlers;
use dispatcher::request::{Request, RequestHandler};
use reactor::{Event, EventDemultiplexer, EventHandler, Reactor};

pub struct RequestReactor {
    handlers: BTreeMap<RequestHandler, Box<EventHandler<Request, Handler=RequestHandler>>>,
    demultiplexer: Demultiplexer<Request>,
    exit_condition: Condition
}

impl RequestReactor {
    pub fn new(demultiplexer: Demultiplexer<Request>) -> RequestReactor {
        let exit_condition = Condition::new(false);
        let exit_handler = Box::new(handlers::exit::ExitHandler::new(exit_condition.clone()));
        let event_handler = Box::new(handlers::event::EventHandler::new());

        let mut reactor = RequestReactor {
            demultiplexer: demultiplexer,
            exit_condition: exit_condition,
            handlers: BTreeMap::new()
        };

        reactor.register_handler(exit_handler);
        reactor.register_handler(event_handler);
        reactor
    }
}

impl Reactor for RequestReactor {
    type Event = Request;
    type Handler = RequestHandler;
    fn handle_events(&mut self) {
        while !self.exit_condition.is_active() {
            if let Some(request) = self.demultiplexer.select() {
                let mut handler = self.handlers.get_mut(&request.handler()).unwrap();
                handler.handle_event(request);
            } else {
                break;
            }
        }
    }
    fn register_handler(&mut self, handler: Box<EventHandler<Self::Event, Handler=RequestHandler>>) {
        self.handlers.insert(handler.handler(), handler);
    }
    fn remove_handler(&mut self, handler: &EventHandler<Self::Event, Handler=RequestHandler>){}
}
<|endoftext|>


<commit_before>pub use alloc::Allocator;
use super::Command;
use super::util::floatX;
use super::v8;
use super::s16;
use super::PDF;
use super::StaticCommand;
use super::HistogramLiteral;
use super::HistogramCommand;
use super::HistogramDistance;
use super::HistogramPair;
use super::ContextType;
use super::HuffmanTree;
use super::ZopfliNode;

/*
struct CombiningAllocator<T1, T2, AllocT1:Allocator<T1>, AllocT2:Allocator<T2>>(AllocT1, AllocT2);

impl <T1, T2, AllocT1:Allocator<T1>, AllocT2:Allocator<T2>> CombiningAllocator<T1, T2, AllocT1, AllocT2> {
  pub fn new(a: AllocT1, b: AllocT2) -> Self {
    CombiningAllocator(a, b)
  }
}

impl <T1, T2, AllocT1:Allocator<T1>, AllocT2:Allocator<T2>> Allocator<T1> for CombiningAllocator<T1, T2, AllocT1, AllocT2> {

}

impl <T1, T2, AllocT1:Allocator<T1>, AllocT2:Allocator<T2>> Allocator<T2> for CombiningAllocator<T1, T2, AllocT1, AllocT2> {

}
*/

//pub trait BrotliAlloc:Allocator<u8> + Allocator<u16> + Allocator<i32> + Allocator<u32> + Allocator<u64> + Allocator<Command> + Allocator<super::util::floatX> + Allocator<v8> + Allocator<s16> + Allocator<PDF> + Allocator<StaticCommand> + Allocator<HistogramLiteral> + Allocator<HistogramCommand> + Allocator<HistogramDistance> + Allocator<HistogramPair> + Allocator<ContextType> + Allocator<HuffmanTree> + Allocator<ZopfliNode>{}
pub struct CombiningAllocator<AllocU8:Allocator<u8>,
                          AllocU16:Allocator<u16>,
                          AllocI32:Allocator<i32>,
                          AllocU32:Allocator<u32>,
                          AllocU64:Allocator<u64>,
                          AllocCommand:Allocator<super::Command>,
                          AllocFloatX:Allocator<super::util::floatX>,
                          AllocV8:Allocator<super::v8>,
                          AllocS16:Allocator<super::s16>,
                          AllocPDF:Allocator<super::PDF>,
                          AllocStaticCommand:Allocator<super::StaticCommand>,
                          AllocHistogramLiteral:Allocator<super::HistogramLiteral>,
                          AllocHistogramCommand:Allocator<super::HistogramCommand>,
                          AllocHistogramDistance:Allocator<super::HistogramDistance>,
                          AllocHistogramPair:Allocator<super::HistogramPair>,
                          AllocContextType:Allocator<super::ContextType>,
                          AllocHuffmanTree:Allocator<super::HuffmanTree>,
                          AllocZopfliNode:Allocator<super::ZopfliNode>,
                          >{
  alloc_u8: AllocU8,
  alloc_u16: AllocU16,
  alloc_i32: AllocI32,
  alloc_u32: AllocU32,
  alloc_u64: AllocU64,
  alloc_c: AllocCommand,
  alloc_f: AllocFloatX,
  alloc_f32x8: AllocV8,
  alloc_i16x16: AllocS16,
  alloc_pdf: AllocPDF,
  alloc_sc: AllocStaticCommand,
  alloc_hl: AllocHistogramLiteral,
  alloc_hc: AllocHistogramCommand,
  alloc_hd: AllocHistogramDistance,
  alloc_hp: AllocHistogramPair,
  alloc_ct: AllocContextType,
  alloc_ht: AllocHuffmanTree,
  alloc_zn: AllocZopfliNode,
}

impl<AllocU8:Allocator<u8>,
     AllocU16:Allocator<u16>,
     AllocI32:Allocator<i32>,
     AllocU32:Allocator<u32>,
     AllocU64:Allocator<u64>,
     AllocCommand:Allocator<super::Command>,
     AllocFloatX:Allocator<::enc::util::floatX>,
     AllocV8:Allocator<super::v8>,
     AllocS16:Allocator<super::s16>,
     AllocPDF:Allocator<super::PDF>,
     AllocStaticCommand:Allocator<super::StaticCommand>,
     AllocHistogramLiteral:Allocator<super::HistogramLiteral>,
     AllocHistogramCommand:Allocator<super::HistogramCommand>,
     AllocHistogramDistance:Allocator<super::HistogramDistance>,
     AllocHistogramPair:Allocator<super::HistogramPair>,
     AllocContextType:Allocator<super::ContextType>,
     AllocHuffmanTree:Allocator<super::HuffmanTree>,
     AllocZopfliNode:Allocator<super::ZopfliNode>,
     > CombiningAllocator<AllocU8,
                          AllocU16,
                          AllocI32,
                          AllocU32,
                          AllocU64,
                          AllocCommand,
                          AllocFloatX,
                          AllocV8,
                          AllocS16,
                          AllocPDF,
                          AllocStaticCommand,
                          AllocHistogramLiteral,
                          AllocHistogramCommand,
                          AllocHistogramDistance,
                          AllocHistogramPair,
                          AllocContextType,
                          AllocHuffmanTree,
                          AllocZopfliNode,
                          > {
  pub fn new(alloc_u8: AllocU8,
         alloc_u16: AllocU16,
         alloc_i32: AllocI32,
         alloc_u32: AllocU32,
         alloc_u64: AllocU64,
         alloc_c: AllocCommand,
         alloc_f: AllocFloatX,
         alloc_f32x8: AllocV8,
         alloc_i16x16: AllocS16,
         alloc_pdf: AllocPDF,
         alloc_sc: AllocStaticCommand,
         alloc_hl: AllocHistogramLiteral,
         alloc_hc: AllocHistogramCommand,
         alloc_hd: AllocHistogramDistance,
         alloc_hp: AllocHistogramPair,
         alloc_ct: AllocContextType,
         alloc_ht: AllocHuffmanTree,
         alloc_zn: AllocZopfliNode) -> Self {
    CombiningAllocator{
      alloc_u8:     alloc_u8,
      alloc_u16:    alloc_u16,
      alloc_i32:   alloc_i32,
      alloc_u32:   alloc_u32,
      alloc_u64:   alloc_u64,
      alloc_c:   alloc_c,
      alloc_f:   alloc_f,
      alloc_f32x8:   alloc_f32x8,
      alloc_i16x16:  alloc_i16x16,
      alloc_pdf:   alloc_pdf,
      alloc_sc:   alloc_sc,
      alloc_hl:   alloc_hl,
      alloc_hc:   alloc_hc,
      alloc_hd:   alloc_hd,
      alloc_hp:   alloc_hp,
      alloc_ct:   alloc_ct,
      alloc_ht:   alloc_ht,
      alloc_zn:   alloc_zn,
    }
  }
}

impl<AllocU8:Allocator<u8>,
     AllocU16:Allocator<u16>,
     AllocI32:Allocator<i32>,
     AllocU32:Allocator<u32>,
     AllocU64:Allocator<u64>,
     AllocCommand:Allocator<super::Command>,
     AllocFloatX:Allocator<super::util::floatX>,
     AllocV8:Allocator<super::v8>,
     AllocS16:Allocator<super::s16>,
     AllocPDF:Allocator<super::PDF>,
     AllocStaticCommand:Allocator<super::StaticCommand>,
     AllocHistogramLiteral:Allocator<super::HistogramLiteral>,
     AllocHistogramCommand:Allocator<super::HistogramCommand>,
     AllocHistogramDistance:Allocator<super::HistogramDistance>,
     AllocHistogramPair:Allocator<super::HistogramPair>,
     AllocContextType:Allocator<super::ContextType>,
     AllocHuffmanTree:Allocator<super::HuffmanTree>,
     AllocZopfliNode:Allocator<super::ZopfliNode>,
     > Allocator<u8> for CombiningAllocator<AllocU8,
                                            AllocU16,
                                            AllocI32,
                                            AllocU32,
                                            AllocU64,
                                            AllocCommand,
                                            AllocFloatX,
                                            AllocV8,
                                            AllocS16,
                                            AllocPDF,
                                            AllocStaticCommand,
                                            AllocHistogramLiteral,
                                            AllocHistogramCommand,
                                            AllocHistogramDistance,
                                            AllocHistogramPair,
                                            AllocContextType,
                                            AllocHuffmanTree,
                                            AllocZopfliNode,
                                            > {
  type AllocatedMemory = AllocU8::AllocatedMemory;
  fn alloc_cell(&mut self, size: usize) -> <Self as Allocator<u8>>::AllocatedMemory {
    self.alloc_u8.alloc_cell(size)
  }
  fn free_cell(&mut self, data: <Self as Allocator<u8>>::AllocatedMemory) {
    self.alloc_u8.free_cell(data)
  }
}

macro_rules! implement_allocator {
  ($bound_name: tt,
   $type_name: tt,
   $local_name: tt) => {
    impl<AllocU8:Allocator<u8>,
         AllocU16:Allocator<u16>,
         AllocI32:Allocator<i32>,
         AllocU32:Allocator<u32>,
         AllocU64:Allocator<u64>,
         AllocCommand:Allocator<super::Command>,
         AllocFloatX:Allocator<super::util::floatX>,
         AllocV8:Allocator<super::v8>,
         AllocS16:Allocator<super::s16>,
         AllocPDF:Allocator<super::PDF>,
         AllocStaticCommand:Allocator<super::StaticCommand>,
         AllocHistogramLiteral:Allocator<super::HistogramLiteral>,
         AllocHistogramCommand:Allocator<super::HistogramCommand>,
         AllocHistogramDistance:Allocator<super::HistogramDistance>,
         AllocHistogramPair:Allocator<super::HistogramPair>,
         AllocContextType:Allocator<super::ContextType>,
         AllocHuffmanTree:Allocator<super::HuffmanTree>,
         AllocZopfliNode:Allocator<super::ZopfliNode>,
         > Allocator<$type_name> for CombiningAllocator<AllocU8,
                                                 AllocU16,
                                                 AllocI32,
                                                 AllocU32,
                                                 AllocU64,
                                                 AllocCommand,
                                                 AllocFloatX,
                                                 AllocV8,
                                                 AllocS16,
                                                 AllocPDF,
                                                 AllocStaticCommand,
                                                 AllocHistogramLiteral,
                                                 AllocHistogramCommand,
                                                 AllocHistogramDistance,
                                                 AllocHistogramPair,
                                                 AllocContextType,
                                                 AllocHuffmanTree,
                                                 AllocZopfliNode,
                                                 > {
      type AllocatedMemory = $bound_name::AllocatedMemory;
      fn alloc_cell(&mut self, size: usize) -> <Self as Allocator<$type_name>>::AllocatedMemory {
        self.$local_name.alloc_cell(size)
      }
      fn free_cell(&mut self, data: <Self as Allocator<$type_name>>::AllocatedMemory) {
        self.$local_name.free_cell(data)
      }
    }
};
}

implement_allocator!(AllocU16, u16, alloc_u16);
implement_allocator!(AllocI32, i32, alloc_i32);
implement_allocator!(AllocU32, u32, alloc_u32);
implement_allocator!(AllocU64, u64, alloc_u64);
implement_allocator!(AllocCommand, Command, alloc_c);
implement_allocator!(AllocFloatX, floatX, alloc_f);
implement_allocator!(AllocV8, v8, alloc_f32x8);
implement_allocator!(AllocS16, s16, alloc_i16x16);
implement_allocator!(AllocPDF, PDF, alloc_pdf);
implement_allocator!(AllocStaticCommand, StaticCommand, alloc_sc);
implement_allocator!(AllocHistogramLiteral, HistogramLiteral, alloc_hl);
implement_allocator!(AllocHistogramCommand, HistogramCommand, alloc_hc);
implement_allocator!(AllocHistogramDistance, HistogramDistance, alloc_hd);
implement_allocator!(AllocHistogramPair, HistogramPair, alloc_hp);
implement_allocator!(AllocContextType, ContextType, alloc_ct);
implement_allocator!(AllocHuffmanTree, HuffmanTree, alloc_ht);
implement_allocator!(AllocZopfliNode, ZopfliNode, alloc_zn);
<commit_msg>make the combined_alloc more friendly for older versions of rust<commit_after>pub use alloc::Allocator;
use super::command::Command;
use super::util::floatX;
use super::v8;
use super::s16;
use super::PDF;
use super::interface::StaticCommand;
use super::histogram::{
  HistogramLiteral,
  HistogramCommand,
  HistogramDistance,
  ContextType,
};
use super::cluster::HistogramPair;

use super::entropy_encode::HuffmanTree;
use super::hash_to_binary_tree::ZopfliNode;

/*
struct CombiningAllocator<T1, T2, AllocT1:Allocator<T1>, AllocT2:Allocator<T2>>(AllocT1, AllocT2);

impl <T1, T2, AllocT1:Allocator<T1>, AllocT2:Allocator<T2>> CombiningAllocator<T1, T2, AllocT1, AllocT2> {
  pub fn new(a: AllocT1, b: AllocT2) -> Self {
    CombiningAllocator(a, b)
  }
}

impl <T1, T2, AllocT1:Allocator<T1>, AllocT2:Allocator<T2>> Allocator<T1> for CombiningAllocator<T1, T2, AllocT1, AllocT2> {

}

impl <T1, T2, AllocT1:Allocator<T1>, AllocT2:Allocator<T2>> Allocator<T2> for CombiningAllocator<T1, T2, AllocT1, AllocT2> {

}
*/

//pub trait BrotliAlloc:Allocator<u8> + Allocator<u16> + Allocator<i32> + Allocator<u32> + Allocator<u64> + Allocator<Command> + Allocator<super::util::floatX> + Allocator<v8> + Allocator<s16> + Allocator<PDF> + Allocator<StaticCommand> + Allocator<HistogramLiteral> + Allocator<HistogramCommand> + Allocator<HistogramDistance> + Allocator<HistogramPair> + Allocator<ContextType> + Allocator<HuffmanTree> + Allocator<ZopfliNode>{}
pub struct CombiningAllocator<AllocU8:Allocator<u8>,
                          AllocU16:Allocator<u16>,
                          AllocI32:Allocator<i32>,
                          AllocU32:Allocator<u32>,
                          AllocU64:Allocator<u64>,
                          AllocCommand:Allocator<Command>,
                          AllocFloatX:Allocator<floatX>,
                          AllocV8:Allocator<v8>,
                          AllocS16:Allocator<s16>,
                          AllocPDF:Allocator<PDF>,
                          AllocStaticCommand:Allocator<StaticCommand>,
                          AllocHistogramLiteral:Allocator<HistogramLiteral>,
                          AllocHistogramCommand:Allocator<HistogramCommand>,
                          AllocHistogramDistance:Allocator<HistogramDistance>,
                          AllocHistogramPair:Allocator<HistogramPair>,
                          AllocContextType:Allocator<ContextType>,
                          AllocHuffmanTree:Allocator<HuffmanTree>,
                          AllocZopfliNode:Allocator<ZopfliNode>,
                          >{
  alloc_u8: AllocU8,
  alloc_u16: AllocU16,
  alloc_i32: AllocI32,
  alloc_u32: AllocU32,
  alloc_u64: AllocU64,
  alloc_c: AllocCommand,
  alloc_f: AllocFloatX,
  alloc_f32x8: AllocV8,
  alloc_i16x16: AllocS16,
  alloc_pdf: AllocPDF,
  alloc_sc: AllocStaticCommand,
  alloc_hl: AllocHistogramLiteral,
  alloc_hc: AllocHistogramCommand,
  alloc_hd: AllocHistogramDistance,
  alloc_hp: AllocHistogramPair,
  alloc_ct: AllocContextType,
  alloc_ht: AllocHuffmanTree,
  alloc_zn: AllocZopfliNode,
}

impl<AllocU8:Allocator<u8>,
     AllocU16:Allocator<u16>,
     AllocI32:Allocator<i32>,
     AllocU32:Allocator<u32>,
     AllocU64:Allocator<u64>,
     AllocCommand:Allocator<Command>,
     AllocFloatX:Allocator<floatX>,
     AllocV8:Allocator<v8>,
     AllocS16:Allocator<s16>,
     AllocPDF:Allocator<PDF>,
     AllocStaticCommand:Allocator<StaticCommand>,
     AllocHistogramLiteral:Allocator<HistogramLiteral>,
     AllocHistogramCommand:Allocator<HistogramCommand>,
     AllocHistogramDistance:Allocator<HistogramDistance>,
     AllocHistogramPair:Allocator<HistogramPair>,
     AllocContextType:Allocator<ContextType>,
     AllocHuffmanTree:Allocator<HuffmanTree>,
     AllocZopfliNode:Allocator<ZopfliNode>,
     > CombiningAllocator<AllocU8,
                          AllocU16,
                          AllocI32,
                          AllocU32,
                          AllocU64,
                          AllocCommand,
                          AllocFloatX,
                          AllocV8,
                          AllocS16,
                          AllocPDF,
                          AllocStaticCommand,
                          AllocHistogramLiteral,
                          AllocHistogramCommand,
                          AllocHistogramDistance,
                          AllocHistogramPair,
                          AllocContextType,
                          AllocHuffmanTree,
                          AllocZopfliNode,
                          > {
  pub fn new(alloc_u8: AllocU8,
         alloc_u16: AllocU16,
         alloc_i32: AllocI32,
         alloc_u32: AllocU32,
         alloc_u64: AllocU64,
         alloc_c: AllocCommand,
         alloc_f: AllocFloatX,
         alloc_f32x8: AllocV8,
         alloc_i16x16: AllocS16,
         alloc_pdf: AllocPDF,
         alloc_sc: AllocStaticCommand,
         alloc_hl: AllocHistogramLiteral,
         alloc_hc: AllocHistogramCommand,
         alloc_hd: AllocHistogramDistance,
         alloc_hp: AllocHistogramPair,
         alloc_ct: AllocContextType,
         alloc_ht: AllocHuffmanTree,
         alloc_zn: AllocZopfliNode) -> Self {
    CombiningAllocator{
      alloc_u8:     alloc_u8,
      alloc_u16:    alloc_u16,
      alloc_i32:   alloc_i32,
      alloc_u32:   alloc_u32,
      alloc_u64:   alloc_u64,
      alloc_c:   alloc_c,
      alloc_f:   alloc_f,
      alloc_f32x8:   alloc_f32x8,
      alloc_i16x16:  alloc_i16x16,
      alloc_pdf:   alloc_pdf,
      alloc_sc:   alloc_sc,
      alloc_hl:   alloc_hl,
      alloc_hc:   alloc_hc,
      alloc_hd:   alloc_hd,
      alloc_hp:   alloc_hp,
      alloc_ct:   alloc_ct,
      alloc_ht:   alloc_ht,
      alloc_zn:   alloc_zn,
    }
  }
}

impl<AllocU8:Allocator<u8>,
     AllocU16:Allocator<u16>,
     AllocI32:Allocator<i32>,
     AllocU32:Allocator<u32>,
     AllocU64:Allocator<u64>,
     AllocCommand:Allocator<Command>,
     AllocFloatX:Allocator<floatX>,
     AllocV8:Allocator<v8>,
     AllocS16:Allocator<s16>,
     AllocPDF:Allocator<PDF>,
     AllocStaticCommand:Allocator<StaticCommand>,
     AllocHistogramLiteral:Allocator<HistogramLiteral>,
     AllocHistogramCommand:Allocator<HistogramCommand>,
     AllocHistogramDistance:Allocator<HistogramDistance>,
     AllocHistogramPair:Allocator<HistogramPair>,
     AllocContextType:Allocator<ContextType>,
     AllocHuffmanTree:Allocator<HuffmanTree>,
     AllocZopfliNode:Allocator<ZopfliNode>,
     > Allocator<u8> for CombiningAllocator<AllocU8,
                                            AllocU16,
                                            AllocI32,
                                            AllocU32,
                                            AllocU64,
                                            AllocCommand,
                                            AllocFloatX,
                                            AllocV8,
                                            AllocS16,
                                            AllocPDF,
                                            AllocStaticCommand,
                                            AllocHistogramLiteral,
                                            AllocHistogramCommand,
                                            AllocHistogramDistance,
                                            AllocHistogramPair,
                                            AllocContextType,
                                            AllocHuffmanTree,
                                            AllocZopfliNode,
                                            > {
  type AllocatedMemory = AllocU8::AllocatedMemory;
  fn alloc_cell(&mut self, size: usize) -> <Self as Allocator<u8>>::AllocatedMemory {
    self.alloc_u8.alloc_cell(size)
  }
  fn free_cell(&mut self, data: <Self as Allocator<u8>>::AllocatedMemory) {
    self.alloc_u8.free_cell(data)
  }
}

macro_rules! implement_allocator {
  ($bound_name: ty,
   $type_name: ty,
   $sub_type_name: ty,
   $local_name: ident) => {
    impl<AllocU8:Allocator<u8>,
         AllocU16:Allocator<u16>,
         AllocI32:Allocator<i32>,
         AllocU32:Allocator<u32>,
         AllocU64:Allocator<u64>,
         AllocCommand:Allocator<Command>,
         AllocFloatX:Allocator<floatX>,
         AllocV8:Allocator<v8>,
         AllocS16:Allocator<s16>,
         AllocPDF:Allocator<PDF>,
         AllocStaticCommand:Allocator<StaticCommand>,
         AllocHistogramLiteral:Allocator<HistogramLiteral>,
         AllocHistogramCommand:Allocator<HistogramCommand>,
         AllocHistogramDistance:Allocator<HistogramDistance>,
         AllocHistogramPair:Allocator<HistogramPair>,
         AllocContextType:Allocator<ContextType>,
         AllocHuffmanTree:Allocator<HuffmanTree>,
         AllocZopfliNode:Allocator<ZopfliNode>,
         > Allocator<$type_name> for CombiningAllocator<AllocU8,
                                                 AllocU16,
                                                 AllocI32,
                                                 AllocU32,
                                                 AllocU64,
                                                 AllocCommand,
                                                 AllocFloatX,
                                                 AllocV8,
                                                 AllocS16,
                                                 AllocPDF,
                                                 AllocStaticCommand,
                                                 AllocHistogramLiteral,
                                                 AllocHistogramCommand,
                                                 AllocHistogramDistance,
                                                 AllocHistogramPair,
                                                 AllocContextType,
                                                 AllocHuffmanTree,
                                                 AllocZopfliNode,
                                                 > {
      type AllocatedMemory = $sub_type_name;
      fn alloc_cell(&mut self, size: usize) -> <Self as Allocator<$type_name>>::AllocatedMemory {
        self.$local_name.alloc_cell(size)
      }
      fn free_cell(&mut self, data: <Self as Allocator<$type_name>>::AllocatedMemory) {
        self.$local_name.free_cell(data)
      }
    }
};
}

implement_allocator!(AllocU16, u16, AllocU16::AllocatedMemory, alloc_u16);

implement_allocator!(AllocI32, i32, AllocI32::AllocatedMemory, alloc_i32);
implement_allocator!(AllocU32, u32, AllocU32::AllocatedMemory, alloc_u32);
implement_allocator!(AllocU64, u64, AllocU64::AllocatedMemory, alloc_u64);
implement_allocator!(AllocCommand, Command, AllocCommand::AllocatedMemory, alloc_c);
implement_allocator!(AllocFloatX, floatX, AllocFloatX::AllocatedMemory, alloc_f);
implement_allocator!(AllocV8, v8, AllocV8::AllocatedMemory, alloc_f32x8);
implement_allocator!(AllocS16, s16, AllocS16::AllocatedMemory, alloc_i16x16);
implement_allocator!(AllocPDF, PDF, AllocPDF::AllocatedMemory, alloc_pdf);
implement_allocator!(AllocStaticCommand, StaticCommand, AllocStaticCommand::AllocatedMemory, alloc_sc);
implement_allocator!(AllocHistogramLiteral, HistogramLiteral, AllocHistogramLiteral::AllocatedMemory, alloc_hl);
implement_allocator!(AllocHistogramCommand, HistogramCommand, AllocHistogramCommand::AllocatedMemory, alloc_hc);
implement_allocator!(AllocHistogramDistance, HistogramDistance, AllocHistogramDistance::AllocatedMemory, alloc_hd);
implement_allocator!(AllocHistogramPair, HistogramPair, AllocHistogramPair::AllocatedMemory, alloc_hp);
implement_allocator!(AllocContextType, ContextType, AllocContextType::AllocatedMemory, alloc_ct);
implement_allocator!(AllocHuffmanTree, HuffmanTree, AllocHuffmanTree::AllocatedMemory, alloc_ht);
implement_allocator!(AllocZopfliNode, ZopfliNode, AllocZopfliNode::AllocatedMemory, alloc_zn);

<|endoftext|>


<commit_before>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the terms of the
// Mozilla Public License, v. 2.0. If a copy of the MPL was not
// distributed with this file, You can obtain one at
// http://mozilla.org/MPL/2.0/.

use std::os::unix::io::{RawFd, AsRawFd};
use std::io::{Read, Write, Error, ErrorKind};

use super::super::frame::{self, FrameState};
use super::super::stream::{SRecv, SSend, SStream, StreamShutdown};

#[derive(Clone)]
pub struct Plain<T> {
    inner: T,
    state: FrameState,
    buffer: Vec<u8>,
    scratch: Vec<u8>,
    tx_queue: Vec<Vec<u8>>,
    rx_queue: Vec<Vec<u8>>,
}

impl<T: Read + Write + AsRawFd + StreamShutdown> Plain<T> {
    pub fn new(stream: T) -> Plain<T> {
        Plain {
            inner: stream,
            state: FrameState::Start,
            buffer: Vec::with_capacity(3),
            scratch: Vec::new(),
            tx_queue: Vec::new(),
            rx_queue: Vec::new()
        }
    }
}

impl<T> Plain<T> {
    fn buf_with_scratch(&mut self, buf: &[u8], len: usize) -> Vec<u8> {
        let mut new_buf = Vec::<u8>::with_capacity(self.scratch.len() + len);
        for byte in self.scratch.iter() {
            new_buf.push(*byte);
        }
        self.scratch = Vec::<u8>::new();
        for x in 0..len {
            new_buf.push(buf[x]);
        }
        new_buf
    }

    fn read_for_frame_start(&mut self, buf: &[u8], offset: &mut usize, len: usize) {
        for _ in *offset..len {
            if buf[*offset] == frame::START {
                self.buffer.push(buf[*offset]);
                self.state = FrameState::PayloadLen;
                *offset += 1;
                break;
            }
            *offset += 1;
        }
    }

    fn read_payload_len(&mut self, buf: &[u8], offset: &mut usize, len: usize) {
        for _ in *offset..len {
            self.buffer.push(buf[*offset]);
            if self.buffer.len() == 3 {
                let len = self.payload_len() + 1;
                self.buffer.reserve_exact(len);
                self.state = FrameState::Payload;
                *offset += 1;
                break;
            }
            *offset += 1;
        }
    }

    fn read_payload(&mut self, buf: &[u8], offset: &mut usize, len: usize) {
        for _ in *offset..len {
            self.buffer.push(buf[*offset]);
            if self.buffer.len() == self.payload_len() + 3 {
                self.state = FrameState::End;
                *offset += 1;
                break;
            }
            *offset += 1;
        }
    }

    fn read_for_frame_end(&mut self, buf: &[u8], offset: usize, len: usize) -> Result<Vec<u8>, ()> {
        if offset < len {
            let expected_end_byte = buf[offset];
            if expected_end_byte == frame::END {
                let mut payload = Vec::<u8>::with_capacity(self.payload_len());
                for x in 3..self.buffer.len() {
                    payload.push(self.buffer[x]);
                }

                self.state = FrameState::Start;
                self.buffer = Vec::<u8>::with_capacity(3);

                // If there is anything left in buf, we need to put it in our
                // scratch space because we're exiting here
                let mut offset = offset;
                offset += 1;
                self.scratch = Vec::<u8>::with_capacity(len - offset);
                for x in offset..len {
                    self.scratch.push(buf[x]);
                }
                return Ok(payload);
            }

            // If we're here, the frame was wrong. Maybe our fault, who knows?
            // Either way, we're going to reset and try to start again from the start byte.
            // We need to dump whatever is left in the buffer into our scratch because it
            // might be in there?
            self.state = FrameState::Start;
            self.buffer = Vec::<u8>::with_capacity(3);
            self.scratch = Vec::<u8>::with_capacity(len - offset);
            for x in offset..len {
                self.scratch.push(buf[x]);
            }
        }
        Err(())
    }

    fn payload_len(&self) -> usize {
        let mask = 0xFFFFu16;
        let len = ((self.buffer[1] as u16) << 8) & mask;
        (len | self.buffer[2] as u16) as usize
    }

    fn vec_from_slice(&self, slice: &[u8]) -> Vec<u8> {
        let mut buf = Vec::<u8>::with_capacity(slice.len());
        for byte in slice.iter() {
            buf.push(*byte);
        }
        buf
    }
}

impl<T: Read + AsRawFd> SRecv for Plain<T> {
    fn recv(&mut self) -> Result<(), Error> {
        loop {
            let mut buf = Vec::<u8>::with_capacity(1024);
            unsafe {
                buf.set_len(1024);
            }
            let result = self.inner.read(&mut buf[..]);
            if result.is_err() {
                let err = result.unwrap_err();
                if err.kind() == ErrorKind::WouldBlock {
                    return Ok(());
                }
                return Err(err);
            }
            let num_read = result.unwrap();

            buf = self.buf_with_scratch(&buf[..], num_read);
            let len = buf.len();
            let mut seek_pos = 0usize;

            if self.state == FrameState::Start {
                self.read_for_frame_start(&buf[..], &mut seek_pos, len);
            }

            if self.state == FrameState::PayloadLen {
                self.read_payload_len(&buf[..], &mut seek_pos, len);
            }

            if self.state == FrameState::Payload {
                self.read_payload(&buf[..], &mut seek_pos, len);
            }

            if self.state == FrameState::End {
                let result = self.read_for_frame_end(&buf[..], seek_pos, len);
                if result.is_ok() {
                    self.rx_queue.push(result.unwrap());
                }
            }
        }
    }

    fn drain_rx_queue(&mut self) -> Vec<Vec<u8>> {
        let buf = self.rx_queue.clone();
        self.rx_queue = Vec::new();
        buf
    }
}

impl<T: Write + AsRawFd> SSend for Plain<T> {
    fn send(&mut self, buf: &[u8]) -> Result<usize, Error> {
        let mut total_written = 0usize;
        self.tx_queue.push(frame::from_slice(buf));
        for x in 0..self.tx_queue.len() {
            let b = self.tx_queue.remove(x);
            let result = self.inner.write(&b[..]);
            if result.is_err() {
                let err = result.unwrap_err();
                if err.kind() == ErrorKind::WouldBlock {
                    self.tx_queue.insert(x, b);
                    return Ok(total_written);
                }
                return Err(err);
            }

            let num_written = result.unwrap();
            total_written += num_written;
            if num_written < b.len() {
                let remainder = self.vec_from_slice(&b[(b.len() - num_written)..b.len()]);
                self.tx_queue.insert(x, remainder);
                return Ok(total_written);
            }
        }
        Ok(total_written)
    }
}

impl<T: StreamShutdown> StreamShutdown for Plain<T> {
    fn shutdown(&mut self) -> Result<(), Error> {
        self.inner.shutdown()
    }
}

impl<T: AsRawFd> AsRawFd for Plain<T> {
    fn as_raw_fd(&self) -> RawFd {
        self.inner.as_raw_fd()
    }
}

impl<T: 'static + Read + Write + AsRawFd + Clone + Send + StreamShutdown> SStream for Plain<T> {}
<commit_msg>Fixed a bug that caused panic when writing calls received EAGAIN/WOULDBLOCK. Bad maths.<commit_after>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the terms of the
// Mozilla Public License, v. 2.0. If a copy of the MPL was not
// distributed with this file, You can obtain one at
// http://mozilla.org/MPL/2.0/.

use std::os::unix::io::{RawFd, AsRawFd};
use std::io::{Read, Write, Error, ErrorKind};

use super::super::frame::{self, FrameState};
use super::super::stream::{SRecv, SSend, SStream, StreamShutdown};

#[derive(Clone)]
pub struct Plain<T> {
    inner: T,
    state: FrameState,
    buffer: Vec<u8>,
    scratch: Vec<u8>,
    tx_queue: Vec<Vec<u8>>,
    rx_queue: Vec<Vec<u8>>,
}

impl<T: Read + Write + AsRawFd + StreamShutdown> Plain<T> {
    pub fn new(stream: T) -> Plain<T> {
        Plain {
            inner: stream,
            state: FrameState::Start,
            buffer: Vec::with_capacity(3),
            scratch: Vec::new(),
            tx_queue: Vec::new(),
            rx_queue: Vec::new()
        }
    }
}

impl<T> Plain<T> {
    fn buf_with_scratch(&mut self, buf: &[u8], len: usize) -> Vec<u8> {
        let mut new_buf = Vec::<u8>::with_capacity(self.scratch.len() + len);
        for byte in self.scratch.iter() {
            new_buf.push(*byte);
        }
        self.scratch = Vec::<u8>::new();
        for x in 0..len {
            new_buf.push(buf[x]);
        }
        new_buf
    }

    fn read_for_frame_start(&mut self, buf: &[u8], offset: &mut usize, len: usize) {
        for _ in *offset..len {
            if buf[*offset] == frame::START {
                self.buffer.push(buf[*offset]);
                self.state = FrameState::PayloadLen;
                *offset += 1;
                break;
            }
            *offset += 1;
        }
    }

    fn read_payload_len(&mut self, buf: &[u8], offset: &mut usize, len: usize) {
        for _ in *offset..len {
            self.buffer.push(buf[*offset]);
            if self.buffer.len() == 3 {
                let len = self.payload_len() + 1;
                self.buffer.reserve_exact(len);
                self.state = FrameState::Payload;
                *offset += 1;
                break;
            }
            *offset += 1;
        }
    }

    fn read_payload(&mut self, buf: &[u8], offset: &mut usize, len: usize) {
        for _ in *offset..len {
            self.buffer.push(buf[*offset]);
            if self.buffer.len() == self.payload_len() + 3 {
                self.state = FrameState::End;
                *offset += 1;
                break;
            }
            *offset += 1;
        }
    }

    fn read_for_frame_end(&mut self, buf: &[u8], offset: usize, len: usize) -> Result<Vec<u8>, ()> {
        if offset < len {
            let expected_end_byte = buf[offset];
            if expected_end_byte == frame::END {
                let mut payload = Vec::<u8>::with_capacity(self.payload_len());
                for x in 3..self.buffer.len() {
                    payload.push(self.buffer[x]);
                }

                self.state = FrameState::Start;
                self.buffer = Vec::<u8>::with_capacity(3);

                // If there is anything left in buf, we need to put it in our
                // scratch space because we're exiting here
                let mut offset = offset;
                offset += 1;
                self.scratch = Vec::<u8>::with_capacity(len - offset);
                for x in offset..len {
                    self.scratch.push(buf[x]);
                }
                return Ok(payload);
            }

            // If we're here, the frame was wrong. Maybe our fault, who knows?
            // Either way, we're going to reset and try to start again from the start byte.
            // We need to dump whatever is left in the buffer into our scratch because it
            // might be in there?
            self.state = FrameState::Start;
            self.buffer = Vec::<u8>::with_capacity(3);
            self.scratch = Vec::<u8>::with_capacity(len - offset);
            for x in offset..len {
                self.scratch.push(buf[x]);
            }
        }
        Err(())
    }

    fn payload_len(&self) -> usize {
        let mask = 0xFFFFu16;
        let len = ((self.buffer[1] as u16) << 8) & mask;
        (len | self.buffer[2] as u16) as usize
    }
}

impl<T: Read + AsRawFd> SRecv for Plain<T> {
    fn recv(&mut self) -> Result<(), Error> {
        loop {
            let mut buf = Vec::<u8>::with_capacity(1024);
            unsafe {
                buf.set_len(1024);
            }
            let result = self.inner.read(&mut buf[..]);
            if result.is_err() {
                let err = result.unwrap_err();
                if err.kind() == ErrorKind::WouldBlock {
                    return Ok(());
                }
                return Err(err);
            }
            let num_read = result.unwrap();

            buf = self.buf_with_scratch(&buf[..], num_read);
            let len = buf.len();
            let mut seek_pos = 0usize;

            if self.state == FrameState::Start {
                self.read_for_frame_start(&buf[..], &mut seek_pos, len);
            }

            if self.state == FrameState::PayloadLen {
                self.read_payload_len(&buf[..], &mut seek_pos, len);
            }

            if self.state == FrameState::Payload {
                self.read_payload(&buf[..], &mut seek_pos, len);
            }

            if self.state == FrameState::End {
                let result = self.read_for_frame_end(&buf[..], seek_pos, len);
                if result.is_ok() {
                    self.rx_queue.push(result.unwrap());
                }
            }
        }
    }

    fn drain_rx_queue(&mut self) -> Vec<Vec<u8>> {
        let buf = self.rx_queue.clone();
        self.rx_queue = Vec::new();
        buf
    }
}

impl<T: Write + AsRawFd> SSend for Plain<T> {
    fn send(&mut self, buf: &[u8]) -> Result<usize, Error> {
        // Insert our message into the back of the queue
        let new_frame = frame::from_slice(buf);
        self.tx_queue.push(new_frame);

        // Counter for total bytes written
        let mut total_written = 0usize;

        // If there is anything in our tx_queue, we need to finish
        // writing those first
        let queue_len = self.tx_queue.len();
        for _ in 0..queue_len {
            let frame = self.tx_queue.remove(0);
            let write_result = self.inner.write(&frame[..]);
            if write_result.is_err() {
                let err = write_result.unwrap_err();
                if err.kind() == ErrorKind::WouldBlock {
                    // Internal buffer _completely_ full, nothing was written from frame
                    self.tx_queue.insert(0, frame);
                    return Ok(total_written);
                }
                return Err(err);
            }

            // Something was written to the buffer
            let num_written = write_result.unwrap();
            total_written += num_written;

            // If we wrote less than we expected, we filled up the buffer,
            // and need to insert the remaining bytes back into the queue to be finished
            // written out next time the socket sends.
            let frame_len = frame.len();
            if num_written < frame_len {
                let mut remaining_bytes = Vec::<u8>::with_capacity(frame_len - num_written);
                for offset in num_written..frame_len {
                    remaining_bytes.push(frame[offset]);
                }

                // Place unsent bytes back into the front of the queue
                self.tx_queue.insert(0, remaining_bytes);
                return Ok(total_written)
            }
        }

        Ok(total_written)
    }
}

impl<T: StreamShutdown> StreamShutdown for Plain<T> {
    fn shutdown(&mut self) -> Result<(), Error> {
        self.inner.shutdown()
    }
}

impl<T: AsRawFd> AsRawFd for Plain<T> {
    fn as_raw_fd(&self) -> RawFd {
        self.inner.as_raw_fd()
    }
}

impl<T: 'static + Read + Write + AsRawFd + Clone + Send + StreamShutdown> SStream for Plain<T> {}
<|endoftext|>