<commit_before>#[cfg(any(feature = "alloc", feature = "std"))]
use base::dimension::Dynamic;
use base::dimension::{U1, U2, U3, U4, U5, U6};
#[cfg(any(feature = "std", feature = "alloc"))]
use base::matrix_vec::MatrixVec;
use base::storage::Owned;
use base::Matrix;

/*
 *
 *
 * Column-major matrices.
 *
 *
 */
/// A staticaly sized column-major matrix with `R` rows and `C` columns.
#[deprecated(note = "This matrix name contains a typo. Use MatrixMN instead.")]
pub type MatrixNM<N, R, C> = Matrix<N, R, C, Owned<N, R, C>>;

/// A staticaly sized column-major matrix with `R` rows and `C` columns.
pub type MatrixMN<N, R, C> = Matrix<N, R, C, Owned<N, R, C>>;

/// A staticaly sized column-major square matrix with `D` rows and columns.
pub type MatrixN<N, D> = MatrixMN<N, D, D>;

/// A dynamically sized column-major matrix.
#[cfg(any(feature = "std", feature = "alloc"))]
pub type DMatrix<N> = MatrixN<N, Dynamic>;

/// A stack-allocated, column-major, 1x1 square matrix.
pub type Matrix1<N> = MatrixN<N, U1>;
/// A stack-allocated, column-major, 2x2 square matrix.
pub type Matrix2<N> = MatrixN<N, U2>;
/// A stack-allocated, column-major, 3x3 square matrix.
pub type Matrix3<N> = MatrixN<N, U3>;
/// A stack-allocated, column-major, 4x4 square matrix.
pub type Matrix4<N> = MatrixN<N, U4>;
/// A stack-allocated, column-major, 5x5 square matrix.
pub type Matrix5<N> = MatrixN<N, U5>;
/// A stack-allocated, column-major, 6x6 square matrix.
pub type Matrix6<N> = MatrixN<N, U6>;

/// A stack-allocated, column-major, 1x2 square matrix.
pub type Matrix1x2<N> = MatrixMN<N, U1, U2>;
/// A stack-allocated, column-major, 1x3 square matrix.
pub type Matrix1x3<N> = MatrixMN<N, U1, U3>;
/// A stack-allocated, column-major, 1x4 square matrix.
pub type Matrix1x4<N> = MatrixMN<N, U1, U4>;
/// A stack-allocated, column-major, 1x5 square matrix.
pub type Matrix1x5<N> = MatrixMN<N, U1, U5>;
/// A stack-allocated, column-major, 1x6 square matrix.
pub type Matrix1x6<N> = MatrixMN<N, U1, U6>;

/// A stack-allocated, column-major, 2x3 square matrix.
pub type Matrix2x3<N> = MatrixMN<N, U2, U3>;
/// A stack-allocated, column-major, 2x4 square matrix.
pub type Matrix2x4<N> = MatrixMN<N, U2, U4>;
/// A stack-allocated, column-major, 2x5 square matrix.
pub type Matrix2x5<N> = MatrixMN<N, U2, U5>;
/// A stack-allocated, column-major, 2x6 square matrix.
pub type Matrix2x6<N> = MatrixMN<N, U2, U6>;

/// A stack-allocated, column-major, 3x4 square matrix.
pub type Matrix3x4<N> = MatrixMN<N, U3, U4>;
/// A stack-allocated, column-major, 3x5 square matrix.
pub type Matrix3x5<N> = MatrixMN<N, U3, U5>;
/// A stack-allocated, column-major, 3x6 square matrix.
pub type Matrix3x6<N> = MatrixMN<N, U3, U6>;

/// A stack-allocated, column-major, 4x5 square matrix.
pub type Matrix4x5<N> = MatrixMN<N, U4, U5>;
/// A stack-allocated, column-major, 4x6 square matrix.
pub type Matrix4x6<N> = MatrixMN<N, U4, U6>;

/// A stack-allocated, column-major, 5x6 square matrix.
pub type Matrix5x6<N> = MatrixMN<N, U5, U6>;

/// A stack-allocated, column-major, 2x1 square matrix.
pub type Matrix2x1<N> = MatrixMN<N, U2, U1>;
/// A stack-allocated, column-major, 3x1 square matrix.
pub type Matrix3x1<N> = MatrixMN<N, U3, U1>;
/// A stack-allocated, column-major, 4x1 square matrix.
pub type Matrix4x1<N> = MatrixMN<N, U4, U1>;
/// A stack-allocated, column-major, 5x1 square matrix.
pub type Matrix5x1<N> = MatrixMN<N, U5, U1>;
/// A stack-allocated, column-major, 6x1 square matrix.
pub type Matrix6x1<N> = MatrixMN<N, U6, U1>;

/// A stack-allocated, column-major, 3x2 square matrix.
pub type Matrix3x2<N> = MatrixMN<N, U3, U2>;
/// A stack-allocated, column-major, 4x2 square matrix.
pub type Matrix4x2<N> = MatrixMN<N, U4, U2>;
/// A stack-allocated, column-major, 5x2 square matrix.
pub type Matrix5x2<N> = MatrixMN<N, U5, U2>;
/// A stack-allocated, column-major, 6x2 square matrix.
pub type Matrix6x2<N> = MatrixMN<N, U6, U2>;

/// A stack-allocated, column-major, 4x3 square matrix.
pub type Matrix4x3<N> = MatrixMN<N, U4, U3>;
/// A stack-allocated, column-major, 5x3 square matrix.
pub type Matrix5x3<N> = MatrixMN<N, U5, U3>;
/// A stack-allocated, column-major, 6x3 square matrix.
pub type Matrix6x3<N> = MatrixMN<N, U6, U3>;

/// A stack-allocated, column-major, 5x4 square matrix.
pub type Matrix5x4<N> = MatrixMN<N, U5, U4>;
/// A stack-allocated, column-major, 6x4 square matrix.
pub type Matrix6x4<N> = MatrixMN<N, U6, U4>;

/// A stack-allocated, column-major, 6x5 square matrix.
pub type Matrix6x5<N> = MatrixMN<N, U6, U5>;

/*
 *
 *
 * Column vectors.
 *
 *
 */
/// A dynamically sized column vector.
#[cfg(any(feature = "std", feature = "alloc"))]
pub type DVector<N> = Matrix<N, Dynamic, U1, MatrixVec<N, Dynamic, U1>>;

/// A statically sized D-dimensional column vector.
pub type VectorN<N, D> = MatrixMN<N, D, U1>;

/// A stack-allocated, 1-dimensional column vector.
pub type Vector1<N> = VectorN<N, U1>;
/// A stack-allocated, 2-dimensional column vector.
pub type Vector2<N> = VectorN<N, U2>;
/// A stack-allocated, 3-dimensional column vector.
pub type Vector3<N> = VectorN<N, U3>;
/// A stack-allocated, 4-dimensional column vector.
pub type Vector4<N> = VectorN<N, U4>;
/// A stack-allocated, 5-dimensional column vector.
pub type Vector5<N> = VectorN<N, U5>;
/// A stack-allocated, 6-dimensional column vector.
pub type Vector6<N> = VectorN<N, U6>;

/*
 *
 *
 * Row vectors.
 *
 *
 */
/// A dynamically sized row vector.
#[cfg(any(feature = "std", feature = "alloc"))]
pub type RowDVector<N> = Matrix<N, U1, Dynamic, MatrixVec<N, U1, Dynamic>>;

/// A statically sized D-dimensional row vector.
pub type RowVectorN<N, D> = MatrixMN<N, U1, D>;

/// A stack-allocated, 1-dimensional row vector.
pub type RowVector1<N> = RowVectorN<N, U1>;
/// A stack-allocated, 2-dimensional row vector.
pub type RowVector2<N> = RowVectorN<N, U2>;
/// A stack-allocated, 3-dimensional row vector.
pub type RowVector3<N> = RowVectorN<N, U3>;
/// A stack-allocated, 4-dimensional row vector.
pub type RowVector4<N> = RowVectorN<N, U4>;
/// A stack-allocated, 5-dimensional row vector.
pub type RowVector5<N> = RowVectorN<N, U5>;
/// A stack-allocated, 6-dimensional row vector.
pub type RowVector6<N> = RowVectorN<N, U6>;
<commit_msg>Fixed error in comments<commit_after>#[cfg(any(feature = "alloc", feature = "std"))]
use base::dimension::Dynamic;
use base::dimension::{U1, U2, U3, U4, U5, U6};
#[cfg(any(feature = "std", feature = "alloc"))]
use base::matrix_vec::MatrixVec;
use base::storage::Owned;
use base::Matrix;

/*
 *
 *
 * Column-major matrices.
 *
 *
 */
/// A staticaly sized column-major matrix with `R` rows and `C` columns.
#[deprecated(note = "This matrix name contains a typo. Use MatrixMN instead.")]
pub type MatrixNM<N, R, C> = Matrix<N, R, C, Owned<N, R, C>>;

/// A staticaly sized column-major matrix with `R` rows and `C` columns.
pub type MatrixMN<N, R, C> = Matrix<N, R, C, Owned<N, R, C>>;

/// A staticaly sized column-major square matrix with `D` rows and columns.
pub type MatrixN<N, D> = MatrixMN<N, D, D>;

/// A dynamically sized column-major matrix.
#[cfg(any(feature = "std", feature = "alloc"))]
pub type DMatrix<N> = MatrixN<N, Dynamic>;

/// A stack-allocated, column-major, 1x1 square matrix.
pub type Matrix1<N> = MatrixN<N, U1>;
/// A stack-allocated, column-major, 2x2 square matrix.
pub type Matrix2<N> = MatrixN<N, U2>;
/// A stack-allocated, column-major, 3x3 square matrix.
pub type Matrix3<N> = MatrixN<N, U3>;
/// A stack-allocated, column-major, 4x4 square matrix.
pub type Matrix4<N> = MatrixN<N, U4>;
/// A stack-allocated, column-major, 5x5 square matrix.
pub type Matrix5<N> = MatrixN<N, U5>;
/// A stack-allocated, column-major, 6x6 square matrix.
pub type Matrix6<N> = MatrixN<N, U6>;

/// A stack-allocated, column-major, 1x2 matrix.
pub type Matrix1x2<N> = MatrixMN<N, U1, U2>;
/// A stack-allocated, column-major, 1x3 matrix.
pub type Matrix1x3<N> = MatrixMN<N, U1, U3>;
/// A stack-allocated, column-major, 1x4 matrix.
pub type Matrix1x4<N> = MatrixMN<N, U1, U4>;
/// A stack-allocated, column-major, 1x5 matrix.
pub type Matrix1x5<N> = MatrixMN<N, U1, U5>;
/// A stack-allocated, column-major, 1x6 matrix.
pub type Matrix1x6<N> = MatrixMN<N, U1, U6>;

/// A stack-allocated, column-major, 2x3 matrix.
pub type Matrix2x3<N> = MatrixMN<N, U2, U3>;
/// A stack-allocated, column-major, 2x4 matrix.
pub type Matrix2x4<N> = MatrixMN<N, U2, U4>;
/// A stack-allocated, column-major, 2x5 matrix.
pub type Matrix2x5<N> = MatrixMN<N, U2, U5>;
/// A stack-allocated, column-major, 2x6 matrix.
pub type Matrix2x6<N> = MatrixMN<N, U2, U6>;

/// A stack-allocated, column-major, 3x4 matrix.
pub type Matrix3x4<N> = MatrixMN<N, U3, U4>;
/// A stack-allocated, column-major, 3x5 matrix.
pub type Matrix3x5<N> = MatrixMN<N, U3, U5>;
/// A stack-allocated, column-major, 3x6 matrix.
pub type Matrix3x6<N> = MatrixMN<N, U3, U6>;

/// A stack-allocated, column-major, 4x5 matrix.
pub type Matrix4x5<N> = MatrixMN<N, U4, U5>;
/// A stack-allocated, column-major, 4x6 matrix.
pub type Matrix4x6<N> = MatrixMN<N, U4, U6>;

/// A stack-allocated, column-major, 5x6 matrix.
pub type Matrix5x6<N> = MatrixMN<N, U5, U6>;

/// A stack-allocated, column-major, 2x1 matrix.
pub type Matrix2x1<N> = MatrixMN<N, U2, U1>;
/// A stack-allocated, column-major, 3x1 matrix.
pub type Matrix3x1<N> = MatrixMN<N, U3, U1>;
/// A stack-allocated, column-major, 4x1 matrix.
pub type Matrix4x1<N> = MatrixMN<N, U4, U1>;
/// A stack-allocated, column-major, 5x1 matrix.
pub type Matrix5x1<N> = MatrixMN<N, U5, U1>;
/// A stack-allocated, column-major, 6x1 matrix.
pub type Matrix6x1<N> = MatrixMN<N, U6, U1>;

/// A stack-allocated, column-major, 3x2 matrix.
pub type Matrix3x2<N> = MatrixMN<N, U3, U2>;
/// A stack-allocated, column-major, 4x2 matrix.
pub type Matrix4x2<N> = MatrixMN<N, U4, U2>;
/// A stack-allocated, column-major, 5x2 matrix.
pub type Matrix5x2<N> = MatrixMN<N, U5, U2>;
/// A stack-allocated, column-major, 6x2 matrix.
pub type Matrix6x2<N> = MatrixMN<N, U6, U2>;

/// A stack-allocated, column-major, 4x3 matrix.
pub type Matrix4x3<N> = MatrixMN<N, U4, U3>;
/// A stack-allocated, column-major, 5x3 matrix.
pub type Matrix5x3<N> = MatrixMN<N, U5, U3>;
/// A stack-allocated, column-major, 6x3 matrix.
pub type Matrix6x3<N> = MatrixMN<N, U6, U3>;

/// A stack-allocated, column-major, 5x4 matrix.
pub type Matrix5x4<N> = MatrixMN<N, U5, U4>;
/// A stack-allocated, column-major, 6x4 matrix.
pub type Matrix6x4<N> = MatrixMN<N, U6, U4>;

/// A stack-allocated, column-major, 6x5 matrix.
pub type Matrix6x5<N> = MatrixMN<N, U6, U5>;

/*
 *
 *
 * Column vectors.
 *
 *
 */
/// A dynamically sized column vector.
#[cfg(any(feature = "std", feature = "alloc"))]
pub type DVector<N> = Matrix<N, Dynamic, U1, MatrixVec<N, Dynamic, U1>>;

/// A statically sized D-dimensional column vector.
pub type VectorN<N, D> = MatrixMN<N, D, U1>;

/// A stack-allocated, 1-dimensional column vector.
pub type Vector1<N> = VectorN<N, U1>;
/// A stack-allocated, 2-dimensional column vector.
pub type Vector2<N> = VectorN<N, U2>;
/// A stack-allocated, 3-dimensional column vector.
pub type Vector3<N> = VectorN<N, U3>;
/// A stack-allocated, 4-dimensional column vector.
pub type Vector4<N> = VectorN<N, U4>;
/// A stack-allocated, 5-dimensional column vector.
pub type Vector5<N> = VectorN<N, U5>;
/// A stack-allocated, 6-dimensional column vector.
pub type Vector6<N> = VectorN<N, U6>;

/*
 *
 *
 * Row vectors.
 *
 *
 */
/// A dynamically sized row vector.
#[cfg(any(feature = "std", feature = "alloc"))]
pub type RowDVector<N> = Matrix<N, U1, Dynamic, MatrixVec<N, U1, Dynamic>>;

/// A statically sized D-dimensional row vector.
pub type RowVectorN<N, D> = MatrixMN<N, U1, D>;

/// A stack-allocated, 1-dimensional row vector.
pub type RowVector1<N> = RowVectorN<N, U1>;
/// A stack-allocated, 2-dimensional row vector.
pub type RowVector2<N> = RowVectorN<N, U2>;
/// A stack-allocated, 3-dimensional row vector.
pub type RowVector3<N> = RowVectorN<N, U3>;
/// A stack-allocated, 4-dimensional row vector.
pub type RowVector4<N> = RowVectorN<N, U4>;
/// A stack-allocated, 5-dimensional row vector.
pub type RowVector5<N> = RowVectorN<N, U5>;
/// A stack-allocated, 6-dimensional row vector.
pub type RowVector6<N> = RowVectorN<N, U6>;
<|endoftext|>


<commit_before>use std::io::prelude::*;
use std::fs::File;
use std::process::Command;

extern crate clap;
use clap::{Arg, App};

extern crate isbfc;
use isbfc::Token::*;
use isbfc::parse;
use isbfc::optimize;

const BUFSIZE: i32 = 8192;

fn main() {
    let matches = App::new("isbfc")
        .version("0.0.1")
        .author("Ian D. Scott <ian@iandouglasscott.com>")
        .about("Brainfuck compiler")
        .arg(Arg::with_name("FILENAME")
             .help("Source file to compile")
             .required(true)
             .index(1))
        .get_matches();

    let path = matches.value_of("FILENAME").unwrap();
    let name = path.rsplitn(2, '.').last().unwrap();
    let mut file = File::open(&path).unwrap();
    let mut code = String::new();
    file.read_to_string(&mut code).unwrap();

    println!("Compiling...");
    let output = compile(&code);
    let mut asmfile = File::create(format!("{}.s", name)).unwrap();
    asmfile.write_all(&output.into_bytes()).unwrap();
    asm_and_link(&name);
}

fn compile(code: &str) -> String {
    let tokens = parse(code);
    let tokens = optimize(tokens);

    let mut output = String::new();
    let mut loops: Vec<i32> = Vec::new();
    let mut ifs: Vec<i32> = Vec::new();
    let mut loopnum = 0;
    let mut ifnum = 0;
    let mut outbuffpos = 0;
    let mut outbuffsize = 0;
    for token in tokens.iter() {
        match *token {
            Add(offset, value) => {
                let dest = if offset == 0 {
                    "%r12".to_string()
                } else {
                    format!("{}(%rbx)", (offset*8))
                };
                if value == 1 && dest == "%r12" {
                    output.push_str("    inc %r12\n");
                } else if value >= 1 {
                    output.push_str(&format!("    addq ${}, {}\n", value, dest));
                } else if value == -1 && dest == "%r12" {
                    output.push_str("    dec %r12\n");
                } else if value <= -1 {
                    output.push_str(&format!("    subq ${}, {}\n", -value, dest));
                }
            },
            MulCopy(src_idx, dest_idx, mul) => {
                let mut src = if src_idx == 0 {
                    "%r12".to_string()
                } else {
                    format!("{}(%rbx)", (src_idx*8))
                };
                let dest = if dest_idx == 0 {
                    "%r12".to_string()
                } else {
                    format!("{}(%rbx)", (dest_idx*8))
                };

                if mul != -1 && mul != 1 {
                    output.push_str(&format!(concat!(
                                "    movq {}, %rax\n",
                                "    movq ${}, %rdx\n",
                                "    mulq %rdx\n"), src, mul.abs()));
                    src = "%rax".to_string();
                } else if src != "%r12" && dest != "%r12" {
                    // x86 cannot move memory to memory
                    output.push_str(&format!("    movq {}, %rax\n", src));
                    src = "%rax".to_string();
                }

                if mul > 0 {
                    output.push_str(&format!("    addq {}, {}\n", src, dest));
                } else {
                    output.push_str(&format!("    subq {}, {}\n", src, dest));
                }
            },
            Set(offset, value) => {
                if offset == 0 && value == 0 {
                    output.push_str("    xor %r12, %r12\n");
                } else if offset == 0 {
                    output.push_str(&format!("    movq ${}, %r12\n", value));
                } else {
                    output.push_str(&format!("    movq ${}, {}(%rbx)\n", value, offset*8));
                }
            },
            Move(offset) => {
                if offset != 0 {
                    output.push_str("    movq %r12, (%rbx)\n");
                    if offset > 0 {
                        output.push_str(&format!("    addq ${}, %rbx\n", offset*8));
                    } else {
                        output.push_str(&format!("    subq ${}, %rbx\n", -offset*8));
                    }
                    output.push_str("    movq (%rbx), %r12\n");
                }
            },
            Loop => {
                loopnum += 1;
                loops.push(loopnum);
                output.push_str(&format!(concat!(
                        "    jmp endloop{}\n",
                        "    loop{}:\n"),
                        loopnum, loopnum));
            },
            EndLoop => {
                let curloop = loops.pop().unwrap();
                output.push_str(&format!(concat!(
                            "    endloop{}:\n",
                            "    test %r12, %r12\n",
                            "    jnz loop{}\n"),
                            curloop, curloop))
            },
            If(offset) => {
                ifnum += 1;
                ifs.push(ifnum);
                if offset == 0 {
                    output.push_str("    test %r12, %r12\n");
                } else {
                    output.push_str(&format!("    cmpq $0, {}(%rbx)\n", offset*8));
                }
                output.push_str(&format!("    jz endif{}\n", ifnum));
            },
            EndIf =>
                output.push_str(&format!("    endif{}:\n", ifs.pop().unwrap())),
            Scan(offset) => {
                // Slighly more optimal than normal loop and move
                loopnum += 1;
                output.push_str(&format!(concat!(
                            "    movq %r12, (%rbx)\n",
                            "    jmp endloop{}\n",
                            "    loop{}:\n"),
                            loopnum, loopnum));
                if offset > 0 {
                    output.push_str(&format!("    addq ${}, %rbx\n", offset*8));
                } else {
                    output.push_str(&format!("    subq ${}, %rbx\n", -offset*8));
                }
                output.push_str(&format!(concat!(
                            "    endloop{}:\n",
                            "    cmp $0, (%rbx)\n",
                            "    jnz loop{}\n",
                            "    movq (%rbx), %r12\n"),
                            loopnum, loopnum));
            },
            Input =>
                output.push_str(concat!("\n    xor %rax, %rax\n",
                                        "    xor %rdi, %rdi\n",
                                        "    movq %rbx, %rsi\n",
                                        "    movq $1, %rdx\n",
                                        "    syscall\n",
                                        "    movq (%rbx), %r12\n\n")),
            LoadOut(offset, add) => {
                let outaddr = format!("(strbuff+{})", outbuffpos);
                if offset == 0 {
                    output.push_str(&format!("    movq %r12, {}\n", outaddr));
                } else {
                    output.push_str(&format!("    movq {}(%rbx), %rax\n", offset*8));
                    output.push_str(&format!("    movq %rax, {}\n", outaddr));
                }
                if add > 0 {
                    output.push_str(&format!("    addb ${}, {}\n", add, outaddr));
                } else if add < 0 {
                    output.push_str(&format!("    subb ${}, {}\n", -add, outaddr));
                }
                outbuffpos += 1;
            },
            LoadOutSet(value) => {
                let outaddr = format!("(strbuff+{})", outbuffpos);
                output.push_str(&format!("    movq ${}, {}\n", value, outaddr));
                outbuffpos += 1;
            },
            Output => {
                output.push_str(&format!(concat!(
                            "    movq $1, %rax\n",
                            "    movq $1, %rdi\n",
                            "    movq $strbuff, %rsi\n",
                            "    movq ${}, %rdx\n",
                            "    syscall\n\n"),
                            outbuffpos));

                if outbuffsize < outbuffpos + 8 {
                    outbuffsize = outbuffpos + 8;
                }
                outbuffpos = 0;
            }
        }
    }

    // Exit syscall
    output.push_str(concat!("\n    movq $60, %rax\n",
                            "    movq $0, %rdi\n",
                            "    syscall\n"));

    format!(concat!(
            ".section .bss\n",
            "    .lcomm strbuff, {outbuffsize}\n",
            "    .lcomm mem, {}\n",
            "    .set startidx, mem + {}\n",
            ".section .text\n",
            ".global _start\n",
            "_start:\n",
            "    xor %r12, %r12\n",
            "    movq $startidx, %rbx\n\n{}"),
            BUFSIZE, BUFSIZE/2, output, outbuffsize=outbuffsize)
}

fn asm_and_link(name: &str) {;
    println!("Assembling...");
    let status = Command::new("as")
        .arg("-g")
        .arg(format!("{}.s", name))
        .arg("-o")
        .arg(format!("{}.o", name))
        .status().unwrap();
    if status.code() == Some(0) {
        println!("Linking...");
        Command::new("ld")
            .arg(format!("{}.o", name))
            .arg("-o")
            .arg(name)
            .spawn().unwrap();
    }
}
<commit_msg>Pass code to stdin of as, unless -S option is provided<commit_after>use std::io::prelude::*;
use std::fs::File;
use std::process::{Command, Stdio};

extern crate clap;
use clap::{Arg, App};

extern crate isbfc;
use isbfc::Token::*;
use isbfc::parse;
use isbfc::optimize;

const BUFSIZE: i32 = 8192;

fn main() {
    let matches = App::new("isbfc")
        .version("0.0.1")
        .author("Ian D. Scott <ian@iandouglasscott.com>")
        .about("Brainfuck compiler")
        .arg(Arg::with_name("output_asm")
             .short("S")
             .help("Assemble but do not link"))
        .arg(Arg::with_name("FILENAME")
             .help("Source file to compile")
             .required(true)
             .index(1))
        .get_matches();

    let path = matches.value_of("FILENAME").unwrap();
    let name = path.rsplitn(2, '.').last().unwrap();
    let mut file = File::open(&path).unwrap();
    let mut code = String::new();
    file.read_to_string(&mut code).unwrap();

    println!("Compiling...");
    let output = compile(&code);
    if matches.is_present("output_asm") {
        let mut asmfile = File::create(format!("{}.s", name)).unwrap();
        asmfile.write_all(&output.into_bytes()).unwrap();
    } else {
        asm_and_link(&output, &name);
    }
}

fn compile(code: &str) -> String {
    let tokens = parse(code);
    let tokens = optimize(tokens);

    let mut output = String::new();
    let mut loops: Vec<i32> = Vec::new();
    let mut ifs: Vec<i32> = Vec::new();
    let mut loopnum = 0;
    let mut ifnum = 0;
    let mut outbuffpos = 0;
    let mut outbuffsize = 0;
    for token in tokens.iter() {
        match *token {
            Add(offset, value) => {
                let dest = if offset == 0 {
                    "%r12".to_string()
                } else {
                    format!("{}(%rbx)", (offset*8))
                };
                if value == 1 && dest == "%r12" {
                    output.push_str("    inc %r12\n");
                } else if value >= 1 {
                    output.push_str(&format!("    addq ${}, {}\n", value, dest));
                } else if value == -1 && dest == "%r12" {
                    output.push_str("    dec %r12\n");
                } else if value <= -1 {
                    output.push_str(&format!("    subq ${}, {}\n", -value, dest));
                }
            },
            MulCopy(src_idx, dest_idx, mul) => {
                let mut src = if src_idx == 0 {
                    "%r12".to_string()
                } else {
                    format!("{}(%rbx)", (src_idx*8))
                };
                let dest = if dest_idx == 0 {
                    "%r12".to_string()
                } else {
                    format!("{}(%rbx)", (dest_idx*8))
                };

                if mul != -1 && mul != 1 {
                    output.push_str(&format!(concat!(
                                "    movq {}, %rax\n",
                                "    movq ${}, %rdx\n",
                                "    mulq %rdx\n"), src, mul.abs()));
                    src = "%rax".to_string();
                } else if src != "%r12" && dest != "%r12" {
                    // x86 cannot move memory to memory
                    output.push_str(&format!("    movq {}, %rax\n", src));
                    src = "%rax".to_string();
                }

                if mul > 0 {
                    output.push_str(&format!("    addq {}, {}\n", src, dest));
                } else {
                    output.push_str(&format!("    subq {}, {}\n", src, dest));
                }
            },
            Set(offset, value) => {
                if offset == 0 && value == 0 {
                    output.push_str("    xor %r12, %r12\n");
                } else if offset == 0 {
                    output.push_str(&format!("    movq ${}, %r12\n", value));
                } else {
                    output.push_str(&format!("    movq ${}, {}(%rbx)\n", value, offset*8));
                }
            },
            Move(offset) => {
                if offset != 0 {
                    output.push_str("    movq %r12, (%rbx)\n");
                    if offset > 0 {
                        output.push_str(&format!("    addq ${}, %rbx\n", offset*8));
                    } else {
                        output.push_str(&format!("    subq ${}, %rbx\n", -offset*8));
                    }
                    output.push_str("    movq (%rbx), %r12\n");
                }
            },
            Loop => {
                loopnum += 1;
                loops.push(loopnum);
                output.push_str(&format!(concat!(
                        "    jmp endloop{}\n",
                        "    loop{}:\n"),
                        loopnum, loopnum));
            },
            EndLoop => {
                let curloop = loops.pop().unwrap();
                output.push_str(&format!(concat!(
                            "    endloop{}:\n",
                            "    test %r12, %r12\n",
                            "    jnz loop{}\n"),
                            curloop, curloop))
            },
            If(offset) => {
                ifnum += 1;
                ifs.push(ifnum);
                if offset == 0 {
                    output.push_str("    test %r12, %r12\n");
                } else {
                    output.push_str(&format!("    cmpq $0, {}(%rbx)\n", offset*8));
                }
                output.push_str(&format!("    jz endif{}\n", ifnum));
            },
            EndIf =>
                output.push_str(&format!("    endif{}:\n", ifs.pop().unwrap())),
            Scan(offset) => {
                // Slighly more optimal than normal loop and move
                loopnum += 1;
                output.push_str(&format!(concat!(
                            "    movq %r12, (%rbx)\n",
                            "    jmp endloop{}\n",
                            "    loop{}:\n"),
                            loopnum, loopnum));
                if offset > 0 {
                    output.push_str(&format!("    addq ${}, %rbx\n", offset*8));
                } else {
                    output.push_str(&format!("    subq ${}, %rbx\n", -offset*8));
                }
                output.push_str(&format!(concat!(
                            "    endloop{}:\n",
                            "    cmp $0, (%rbx)\n",
                            "    jnz loop{}\n",
                            "    movq (%rbx), %r12\n"),
                            loopnum, loopnum));
            },
            Input =>
                output.push_str(concat!("\n    xor %rax, %rax\n",
                                        "    xor %rdi, %rdi\n",
                                        "    movq %rbx, %rsi\n",
                                        "    movq $1, %rdx\n",
                                        "    syscall\n",
                                        "    movq (%rbx), %r12\n\n")),
            LoadOut(offset, add) => {
                let outaddr = format!("(strbuff+{})", outbuffpos);
                if offset == 0 {
                    output.push_str(&format!("    movq %r12, {}\n", outaddr));
                } else {
                    output.push_str(&format!("    movq {}(%rbx), %rax\n", offset*8));
                    output.push_str(&format!("    movq %rax, {}\n", outaddr));
                }
                if add > 0 {
                    output.push_str(&format!("    addb ${}, {}\n", add, outaddr));
                } else if add < 0 {
                    output.push_str(&format!("    subb ${}, {}\n", -add, outaddr));
                }
                outbuffpos += 1;
            },
            LoadOutSet(value) => {
                let outaddr = format!("(strbuff+{})", outbuffpos);
                output.push_str(&format!("    movq ${}, {}\n", value, outaddr));
                outbuffpos += 1;
            },
            Output => {
                output.push_str(&format!(concat!(
                            "    movq $1, %rax\n",
                            "    movq $1, %rdi\n",
                            "    movq $strbuff, %rsi\n",
                            "    movq ${}, %rdx\n",
                            "    syscall\n\n"),
                            outbuffpos));

                if outbuffsize < outbuffpos + 8 {
                    outbuffsize = outbuffpos + 8;
                }
                outbuffpos = 0;
            }
        }
    }

    // Exit syscall
    output.push_str(concat!("\n    movq $60, %rax\n",
                            "    movq $0, %rdi\n",
                            "    syscall\n"));

    format!(concat!(
            ".section .bss\n",
            "    .lcomm strbuff, {outbuffsize}\n",
            "    .lcomm mem, {}\n",
            "    .set startidx, mem + {}\n",
            ".section .text\n",
            ".global _start\n",
            "_start:\n",
            "    xor %r12, %r12\n",
            "    movq $startidx, %rbx\n\n{}"),
            BUFSIZE, BUFSIZE/2, output, outbuffsize=outbuffsize)
}

fn asm_and_link(code: &str, name: &str) {;
    println!("Assembling...");

    let mut child = Command::new("as")
        .arg("-g")
        .arg("-") // Standard input
        .arg("-o")
        .arg(format!("{}.o", name))
        .stdin(Stdio::piped())
        .spawn().unwrap();

    child.stdin.take().unwrap().write_all(code.as_bytes()).unwrap();

    let status = child.wait().unwrap();
    if status.code() == Some(0) {
        println!("Linking...");
        Command::new("ld")
            .arg(format!("{}.o", name))
            .arg("-o")
            .arg(name)
            .spawn().unwrap();
    }
}
<|endoftext|>


<commit_before>//! Write systems to disk.

use error::Result;
use grafen::system::System;

use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::{Path, PathBuf};

/// Output a system to disk as a GROMOS formatted file.
/// The filename extension is adjusted to .gro.
///
/// # Errors
/// Returns an error if the file could not be written to.
pub fn write_gromos(system: &System, output_file: &Path, title: &str) -> Result<()> {
    let path = PathBuf::from(output_file).with_extension("gro");
    let file = File::create(&path)?;
    let mut writer = BufWriter::new(file);

    writer.write_fmt(format_args!("{}\n", title))?;
    writer.write_fmt(format_args!("{}\n", system.num_atoms()))?;

    for (i, residue) in system.residues.iter().enumerate() {
        for (j, atom) in residue.base.atoms.iter().enumerate() {
            // GROMOS files wrap atom and residue numbering after five digits
            // so we must output at most that. We also switch to indexing the
            // numbers from 1 instead of from 0.
            let residue_number = (i + 1) % 100_000;
            let atom_number = (j + 1) % 100_000;

            let position = residue.position + atom.position;
            let (x, y, z) = position.to_tuple();

            writer.write_fmt(format_args!("{:>5}{:<5}{:>5}{:>5}{:>8.3}{:>8.3}{:>8.3}\n",
                                        residue_number,
                                        residue.base.code,
                                        atom.code,
                                        atom_number,
                                        x, y, z))?;
        }
    }

    let (dx, dy, dz) = system.dimensions.to_tuple();
    writer.write_fmt(format_args!("{:12.8} {:12.8} {:12.8}\n", dx, dy, dz))?;

    Ok(())
}
<commit_msg>Fix atom numbering output in `write_gromos`<commit_after>//! Write systems to disk.

use error::Result;
use grafen::system::System;

use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::Path;

/// Output a system to disk as a GROMOS formatted file.
/// The filename extension is adjusted to .gro.
///
/// # Errors
/// Returns an error if the file could not be written to.
pub fn write_gromos(system: &System, output_file: &Path, title: &str) -> Result<()> {
    let path = output_file.with_extension("gro");
    let file = File::create(path)?;
    let mut writer = BufWriter::new(file);

    writer.write_fmt(format_args!("{}\n", title))?;
    writer.write_fmt(format_args!("{}\n", system.num_atoms()))?;

    // Absolute atom numbering.
    let mut j = 0;

    for (i, residue) in system.residues.iter().enumerate() {
        for atom in residue.base.atoms.iter() {
            // GROMOS files wrap atom and residue numbering after five digits
            // so we must output at most that. We also switch to indexing the
            // numbers from 1 instead of from 0.
            let residue_number = (i + 1) % 100_000;
            let atom_number = (j + 1) % 100_000;
            j += 1;

            let position = residue.position + atom.position;
            let (x, y, z) = position.to_tuple();

            writer.write_fmt(format_args!("{:>5}{:<5}{:>5}{:>5}{:>8.3}{:>8.3}{:>8.3}\n",
                                        residue_number,
                                        residue.base.code,
                                        atom.code,
                                        atom_number,
                                        x, y, z))?;
        }
    }

    let (dx, dy, dz) = system.dimensions.to_tuple();
    writer.write_fmt(format_args!("{:12.8} {:12.8} {:12.8}\n", dx, dy, dz))?;

    Ok(())
}
<|endoftext|>


<commit_before>extern crate termion;

use std::io::{self, Write};
use termion::cursor::{self, DetectCursorPos};
use termion::raw::IntoRawMode;

fn main() {
    let mut stdout = io::stdout().into_raw_mode().unwrap();

    write!(stdout, "{}", cursor::Save).unwrap();
    stdout.flush().unwrap();

    write!(stdout, "{}", cursor::Goto(999, 999)).unwrap();
    stdout.flush().unwrap();

    let (w, h) = stdout.cursor_pos().unwrap();

    write!(stdout, "{}", cursor::Restore).unwrap();
    stdout.flush().unwrap();

    drop(stdout);

    println!("export COLUMNS={}", w);
    println!("export LINES={}", h);
}
<commit_msg>Fix resize program by using get_tty<commit_after>extern crate extra;
extern crate termion;

use extra::option::OptionalExt;
use std::io::{self, Write};
use termion::cursor::{self, DetectCursorPos};
use termion::raw::IntoRawMode;

fn main() {
    let mut stderr = io::stderr();

    let mut terminal = termion::get_tty().try(&mut stderr)
                               .into_raw_mode().try(&mut stderr);

    write!(terminal, "{}", cursor::Save).unwrap();
    terminal.flush().unwrap();

    write!(terminal, "{}", cursor::Goto(999, 999)).unwrap();
    terminal.flush().unwrap();

    let (w, h) = terminal.cursor_pos().unwrap();

    write!(terminal, "{}", cursor::Restore).unwrap();
    terminal.flush().unwrap();

    drop(terminal);

    println!("export COLUMNS={};", w);
    println!("export LINES={};", h);
}
<|endoftext|>


<commit_before>use failure::{Error, ResultExt};
use mal::{self, MAL};
use regex::Regex;
use process;
use serde_json;
use std;
use std::collections::HashMap;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::process::ExitStatus;

#[derive(Fail, Debug)]
pub enum SeriesError {
    #[fail(display = "multiple series found")] MultipleSeriesFound,
    #[fail(display = "episode {} not found", _0)] EpisodeNotFound(u32),
    #[fail(display = "season {} information not found", _0)] SeasonInfoNotFound(u32),
}

#[derive(Debug)]
pub struct Series {
    pub name: String,
    pub data: SeriesData,
    pub episodes: HashMap<u32, PathBuf>,
    data_path: PathBuf,
}

impl Series {
    pub const DATA_FILE_NAME: &'static str = ".anitrack";

    pub fn from_path(path: &Path) -> Result<Series, Error> {
        let data_path = PathBuf::from(path).join(Series::DATA_FILE_NAME);
        let data = SeriesData::from_path_or_default(&data_path)?;

        let ep_data = EpisodeData::parse(path)?;

        Ok(Series {
            name: ep_data.series_name,
            data,
            episodes: ep_data.paths,
            data_path,
        })
    }

    pub fn play_episode(&self, ep_num: u32) -> Result<ExitStatus, Error> {
        let path = self.episodes.get(&ep_num)
            .ok_or(SeriesError::EpisodeNotFound(ep_num))?;

        let output = process::open_with_default(path).output()?;
        Ok(output.status)
    }

    pub fn set_season_data(&mut self, season: u32, info: SeasonInfo) {
        self.data.seasons.insert(season, info);
    }

    pub fn get_season_data(&self, season: u32) -> Result<&SeasonInfo, SeriesError> {
        self.data.seasons.get(&season).ok_or(SeriesError::SeasonInfoNotFound(season))
    }

    pub fn has_season_data(&self, season: u32) -> bool {
        self.data.seasons.contains_key(&season)
    }

    pub fn save_data(&self) -> Result<(), Error> {
        self.data.write_to(&self.data_path)
    }
}

#[derive(Debug, Default, Serialize, Deserialize)]
pub struct SeriesData {
    pub seasons: HashMap<u32, SeasonInfo>,
}

impl SeriesData {
    fn from_path(path: &Path) -> Result<SeriesData, Error> {
        let file = File::open(path)?;
        let data = serde_json::from_reader(file)?;

        Ok(data)
    }

    fn from_path_or_default(path: &Path) -> Result<SeriesData, Error> {
        if path.exists() {
            SeriesData::from_path(path)
        } else {
            Ok(SeriesData::default())
        }
    }

    fn write_to(&self, path: &Path) -> Result<(), Error> {
        let file = File::create(path)?;
        serde_json::to_writer_pretty(file, self)?;

        Ok(())
    }
}

#[derive(Fail, Debug)]
pub enum SeasonInfoError {
    #[fail(display = "no anime with id {} found with name [{}] on MAL", _0, _1)]
    UnknownAnimeID(u32, String),
}

#[derive(Debug, Serialize, Deserialize)]
pub struct SeasonInfo {
    pub series_id: u32,
    pub episodes: u32,
    pub search_title: String,
}

impl SeasonInfo {
    pub fn create_basic(id: u32, episodes: u32, search_title: String) -> SeasonInfo {
        SeasonInfo {
            series_id: id,
            episodes,
            search_title,
        }
    }

    pub fn request_mal_info(&self, mal: &MAL) -> Result<mal::SeriesInfo, Error> {
        mal.search(&self.search_title)
            .context("MAL search failed")?
            .into_iter()
            .find(|i| i.id == self.series_id)
            .ok_or(SeasonInfoError::UnknownAnimeID(self.series_id, self.search_title.clone()).into())
    }
}

#[derive(Fail, Debug)]
pub enum EpisodeDataError {
    #[fail(display = "no episodes found")]
    NoEpisodesFound,
}

#[derive(Debug)]
pub struct EpisodeData {
    pub series_name: String,
    pub paths: HashMap<u32, PathBuf>,
}

impl EpisodeData {
    pub fn parse(path: &Path) -> Result<EpisodeData, Error> {
        let mut series = None;
        let mut episodes = HashMap::new();

        for entry in std::fs::read_dir(path)? {
            let path = entry?.path();

            let info = match EpisodeInfo::parse(&path) {
                Some(info) => info,
                None => continue,
            };

            if let Some(ref series) = series {
                if series != &info.series {
                    bail!(SeriesError::MultipleSeriesFound);
                }
            } else {
                series = Some(info.series);
            }

            episodes.insert(info.episode, path);
        }

        let series = series.ok_or(EpisodeDataError::NoEpisodesFound)?;

        Ok(EpisodeData {
            series_name: series,
            paths: episodes,
        })
    }
}

#[derive(Debug)]
struct EpisodeInfo {
    series: String,
    episode: u32,
}

impl EpisodeInfo {
    fn parse(path: &Path) -> Option<EpisodeInfo> {
        if !path.is_file() {
            return None;
        }

        lazy_static! {
            static ref EP_FORMAT: Regex = Regex::new(r"(?:\[.+?\]\s*)?(?P<series>.+?)\s*-?\s*(?P<episode>\d+)\s*(?:\(|\[|\.)")
                .unwrap();
        }

        let filename = path.file_name()?.to_str().unwrap().replace('_', " ");
        let caps = EP_FORMAT.captures(&filename)?;

        let info = EpisodeInfo {
            series: caps["series"].into(),
            episode: caps["episode"].parse().ok()?,
        };

        Some(info)
    }
}<commit_msg>Sanitize file name for parsing<commit_after>use failure::{Error, ResultExt};
use mal::{self, MAL};
use regex::Regex;
use process;
use serde_json;
use std;
use std::collections::HashMap;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::process::ExitStatus;

#[derive(Fail, Debug)]
pub enum SeriesError {
    #[fail(display = "multiple series found")] MultipleSeriesFound,
    #[fail(display = "episode {} not found", _0)] EpisodeNotFound(u32),
    #[fail(display = "season {} information not found", _0)] SeasonInfoNotFound(u32),
}

#[derive(Debug)]
pub struct Series {
    pub name: String,
    pub data: SeriesData,
    pub episodes: HashMap<u32, PathBuf>,
    data_path: PathBuf,
}

impl Series {
    pub const DATA_FILE_NAME: &'static str = ".anitrack";

    pub fn from_path(path: &Path) -> Result<Series, Error> {
        let data_path = PathBuf::from(path).join(Series::DATA_FILE_NAME);
        let data = SeriesData::from_path_or_default(&data_path)?;

        let ep_data = EpisodeData::parse(path)?;

        Ok(Series {
            name: ep_data.series_name,
            data,
            episodes: ep_data.paths,
            data_path,
        })
    }

    pub fn play_episode(&self, ep_num: u32) -> Result<ExitStatus, Error> {
        let path = self.episodes.get(&ep_num)
            .ok_or(SeriesError::EpisodeNotFound(ep_num))?;

        let output = process::open_with_default(path).output()?;
        Ok(output.status)
    }

    pub fn set_season_data(&mut self, season: u32, info: SeasonInfo) {
        self.data.seasons.insert(season, info);
    }

    pub fn get_season_data(&self, season: u32) -> Result<&SeasonInfo, SeriesError> {
        self.data.seasons.get(&season).ok_or(SeriesError::SeasonInfoNotFound(season))
    }

    pub fn has_season_data(&self, season: u32) -> bool {
        self.data.seasons.contains_key(&season)
    }

    pub fn save_data(&self) -> Result<(), Error> {
        self.data.write_to(&self.data_path)
    }
}

#[derive(Debug, Default, Serialize, Deserialize)]
pub struct SeriesData {
    pub seasons: HashMap<u32, SeasonInfo>,
}

impl SeriesData {
    fn from_path(path: &Path) -> Result<SeriesData, Error> {
        let file = File::open(path)?;
        let data = serde_json::from_reader(file)?;

        Ok(data)
    }

    fn from_path_or_default(path: &Path) -> Result<SeriesData, Error> {
        if path.exists() {
            SeriesData::from_path(path)
        } else {
            Ok(SeriesData::default())
        }
    }

    fn write_to(&self, path: &Path) -> Result<(), Error> {
        let file = File::create(path)?;
        serde_json::to_writer_pretty(file, self)?;

        Ok(())
    }
}

#[derive(Fail, Debug)]
pub enum SeasonInfoError {
    #[fail(display = "no anime with id {} found with name [{}] on MAL", _0, _1)]
    UnknownAnimeID(u32, String),
}

#[derive(Debug, Serialize, Deserialize)]
pub struct SeasonInfo {
    pub series_id: u32,
    pub episodes: u32,
    pub search_title: String,
}

impl SeasonInfo {
    pub fn create_basic(id: u32, episodes: u32, search_title: String) -> SeasonInfo {
        SeasonInfo {
            series_id: id,
            episodes,
            search_title,
        }
    }

    pub fn request_mal_info(&self, mal: &MAL) -> Result<mal::SeriesInfo, Error> {
        mal.search(&self.search_title)
            .context("MAL search failed")?
            .into_iter()
            .find(|i| i.id == self.series_id)
            .ok_or(SeasonInfoError::UnknownAnimeID(self.series_id, self.search_title.clone()).into())
    }
}

#[derive(Fail, Debug)]
pub enum EpisodeDataError {
    #[fail(display = "no episodes found")]
    NoEpisodesFound,
}

#[derive(Debug)]
pub struct EpisodeData {
    pub series_name: String,
    pub paths: HashMap<u32, PathBuf>,
}

impl EpisodeData {
    pub fn parse(path: &Path) -> Result<EpisodeData, Error> {
        let mut series = None;
        let mut episodes = HashMap::new();

        for entry in std::fs::read_dir(path)? {
            let path = entry?.path();

            let info = match EpisodeInfo::parse(&path) {
                Some(info) => info,
                None => continue,
            };

            if let Some(ref series) = series {
                if series != &info.series {
                    bail!(SeriesError::MultipleSeriesFound);
                }
            } else {
                series = Some(info.series);
            }

            episodes.insert(info.episode, path);
        }

        let series = series.ok_or(EpisodeDataError::NoEpisodesFound)?;

        Ok(EpisodeData {
            series_name: series,
            paths: episodes,
        })
    }
}

#[derive(Debug)]
struct EpisodeInfo {
    series: String,
    episode: u32,
}

impl EpisodeInfo {
    fn parse(path: &Path) -> Option<EpisodeInfo> {
        if !path.is_file() {
            return None;
        }

        lazy_static! {
            static ref EP_FORMAT: Regex = Regex::new(r"(?:\[.+?\]\s*)?(?P<series>.+?)\s*-?\s*(?P<episode>\d+)\s*(?:\(|\[|\.)")
                .unwrap();
        }

        // Replace certain special characters with spaces since they can either
        // affect parsing or prevent finding results on MAL
        let filename = path.file_name()?.to_str()
            .unwrap()
            .replace('_', " ");

        let caps = EP_FORMAT.captures(&filename)?;

        let clean_name = {
            let raw = &caps["series"];
            raw.replace('.', " ")
               .replace(" -", ":") // Dashes typically represent a colon in file names
               .trim()
               .to_string()
        };

        let info = EpisodeInfo {
            series: clean_name,
            episode: caps["episode"].parse().ok()?,
        };

        Some(info)
    }
}<|endoftext|>


<commit_before>//! Module for parsing ISO Base Media Format aka video/mp4 streams.
//! Internal unit tests.

// 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 https://mozilla.org/MPL/2.0/.

use std::io::Cursor;
use super::*;
extern crate test_assembler;
use self::test_assembler::*;

enum BoxSize {
    Short(u32),
    Long(u64),
    UncheckedShort(u32),
    UncheckedLong(u64),
}

fn make_box_raw<F>(size: BoxSize, name: &[u8; 4], func: F) -> Cursor<Vec<u8>>
    where F: Fn(Section) -> Section
{
    let mut section = Section::new();
    section = match size {
        BoxSize::Short(size) | BoxSize::UncheckedShort(size) => section.B32(size),
        BoxSize::Long(_) | BoxSize::UncheckedLong(_) => section.B32(1),
    };
    section = section.append_bytes(name);
    section = match size {
        // The spec allows the 32-bit size to be 0 to indicate unknown
        // length streams.  It's not clear if this is valid when using a
        // 64-bit size, so prohibit it for now.
        BoxSize::Long(size) => {
            assert!(size > 0);
            section.B64(size)
        }
        BoxSize::UncheckedLong(size) => section.B64(size),
        _ => section,
    };
    section = func(section);
    match size {
        BoxSize::Short(size) => {
            if size > 0 {
                assert_eq!(size as u64, section.size())
            }
        }
        BoxSize::Long(size) => assert_eq!(size, section.size()),
        _ => (),
    }
    Cursor::new(section.get_contents().unwrap())
}

fn make_box<F>(size: u32, name: &[u8; 4], func: F) -> Cursor<Vec<u8>>
    where F: Fn(Section) -> Section
{
    make_box_raw(BoxSize::Short(size), name, func)
}

fn make_fullbox<F>(size: u32, name: &[u8; 4], version: u8, func: F) -> Cursor<Vec<u8>>
    where F: Fn(Section) -> Section
{
    make_box_raw(BoxSize::Short(size), name, |s| {
        func(s.B8(version)
              .B8(0)
              .B8(0)
              .B8(0))
    })
}

#[test]
fn read_box_header_short() {
    let mut stream = make_box_raw(BoxSize::Short(8), b"test", |s| s);
    let parsed = read_box_header(&mut stream).unwrap();
    assert_eq!(parsed.name, FourCC(*b"test"));
    assert_eq!(parsed.size, 8);
}

#[test]
fn read_box_header_long() {
    let mut stream = make_box_raw(BoxSize::Long(16), b"test", |s| s);
    let parsed = read_box_header(&mut stream).unwrap();
    assert_eq!(parsed.name, FourCC(*b"test"));
    assert_eq!(parsed.size, 16);
}

#[test]
fn read_box_header_short_unknown_size() {
    let mut stream = make_box_raw(BoxSize::Short(0), b"test", |s| s);
    match read_box_header(&mut stream) {
        Err(Error::Unsupported) => (),
        _ => panic!("unexpected result reading box with unknown size"),
    };
}

#[test]
fn read_box_header_short_invalid_size() {
    let mut stream = make_box_raw(BoxSize::UncheckedShort(2), b"test", |s| s);
    match read_box_header(&mut stream) {
        Err(Error::InvalidData) => (),
        _ => panic!("unexpected result reading box with invalid size"),
    };
}

#[test]
fn read_box_header_long_invalid_size() {
    let mut stream = make_box_raw(BoxSize::UncheckedLong(2), b"test", |s| s);
    match read_box_header(&mut stream) {
        Err(Error::InvalidData) => (),
        _ => panic!("unexpected result reading box with invalid size"),
    };
}

#[test]
fn read_ftyp() {
    let mut stream = make_box(24, b"ftyp", |s| {
        s.append_bytes(b"mp42")
         .B32(0) // minor version
         .append_bytes(b"isom")
         .append_bytes(b"mp42")
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_ftyp(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"ftyp"));
    assert_eq!(parsed.header.size, 24);
    assert_eq!(parsed.major_brand, FourCC(*b"mp42"));
    assert_eq!(parsed.minor_version, 0);
    assert_eq!(parsed.compatible_brands.len(), 2);
    assert_eq!(parsed.compatible_brands[0], FourCC(*b"isom"));
    assert_eq!(parsed.compatible_brands[1], FourCC(*b"mp42"));
}

#[test]
fn read_elst_v0() {
    let mut stream = make_fullbox(28, b"elst", 0, |s| {
        s.B32(1) // list count
          // first entry
         .B32(1234) // duration
         .B32(5678) // time
         .B16(12) // rate integer
         .B16(34) // rate fraction
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_elst(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"elst"));
    assert_eq!(parsed.header.size, 28);
    assert_eq!(parsed.edits.len(), 1);
    assert_eq!(parsed.edits[0].segment_duration, 1234);
    assert_eq!(parsed.edits[0].media_time, 5678);
    assert_eq!(parsed.edits[0].media_rate_integer, 12);
    assert_eq!(parsed.edits[0].media_rate_fraction, 34);
}

#[test]
fn read_elst_v1() {
    let mut stream = make_fullbox(56, b"elst", 1, |s| {
        s.B32(2) // list count
         // first entry
         .B64(1234) // duration
         .B64(5678) // time
         .B16(12) // rate integer
         .B16(34) // rate fraction
         // second entry
         .B64(1234) // duration
         .B64(5678) // time
         .B16(12) // rate integer
         .B16(34) // rate fraction
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_elst(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"elst"));
    assert_eq!(parsed.header.size, 56);
    assert_eq!(parsed.edits.len(), 2);
    assert_eq!(parsed.edits[1].segment_duration, 1234);
    assert_eq!(parsed.edits[1].media_time, 5678);
    assert_eq!(parsed.edits[1].media_rate_integer, 12);
    assert_eq!(parsed.edits[1].media_rate_fraction, 34);
}

#[test]
fn read_mdhd_v0() {
    let mut stream = make_fullbox(32, b"mdhd", 0, |s| {
        s.B32(0)
         .B32(0)
         .B32(1234) // timescale
         .B32(5678) // duration
         .B32(0)
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_mdhd(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"mdhd"));
    assert_eq!(parsed.header.size, 32);
    assert_eq!(parsed.timescale, 1234);
    assert_eq!(parsed.duration, 5678);
}

#[test]
fn read_mdhd_v1() {
    let mut stream = make_fullbox(44, b"mdhd", 1, |s| {
        s.B64(0)
         .B64(0)
         .B32(1234) // timescale
         .B64(5678) // duration
         .B32(0)
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_mdhd(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"mdhd"));
    assert_eq!(parsed.header.size, 44);
    assert_eq!(parsed.timescale, 1234);
    assert_eq!(parsed.duration, 5678);
}

#[test]
fn read_mdhd_unknown_duration() {
    let mut stream = make_fullbox(32, b"mdhd", 0, |s| {
        s.B32(0)
         .B32(0)
         .B32(1234) // timescale
         .B32(::std::u32::MAX) // duration
         .B32(0)
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_mdhd(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"mdhd"));
    assert_eq!(parsed.header.size, 32);
    assert_eq!(parsed.timescale, 1234);
    assert_eq!(parsed.duration, ::std::u64::MAX);
}

#[test]
fn read_mdhd_invalid_timescale() {
    let mut stream = make_fullbox(44, b"mdhd", 1, |s| {
        s.B64(0)
         .B64(0)
         .B32(0) // timescale
         .B64(5678) // duration
         .B32(0)
    });
    let header = read_box_header(&mut stream).unwrap();
    let r = super::parse_mdhd(&mut stream, &header, &mut super::Track::new(0));
    assert_eq!(r.is_err(), true);
}

#[test]
fn read_mvhd_v0() {
    let mut stream = make_fullbox(108, b"mvhd", 0, |s| {
        s.B32(0)
         .B32(0)
         .B32(1234)
         .B32(5678)
         .append_repeated(0, 80)
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_mvhd(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"mvhd"));
    assert_eq!(parsed.header.size, 108);
    assert_eq!(parsed.timescale, 1234);
    assert_eq!(parsed.duration, 5678);
}

#[test]
fn read_mvhd_v1() {
    let mut stream = make_fullbox(120, b"mvhd", 1, |s| {
        s.B64(0)
         .B64(0)
         .B32(1234)
         .B64(5678)
         .append_repeated(0, 80)
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_mvhd(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"mvhd"));
    assert_eq!(parsed.header.size, 120);
    assert_eq!(parsed.timescale, 1234);
    assert_eq!(parsed.duration, 5678);
}

#[test]
fn read_mvhd_invalid_timescale() {
    let mut stream = make_fullbox(120, b"mvhd", 1, |s| {
        s.B64(0)
         .B64(0)
         .B32(0)
         .B64(5678)
         .append_repeated(0, 80)
    });
    let header = read_box_header(&mut stream).unwrap();
    let r = super::parse_mvhd(&mut stream, &header);
    assert_eq!(r.is_err(), true);
}

#[test]
fn read_mvhd_unknown_duration() {
    let mut stream = make_fullbox(108, b"mvhd", 0, |s| {
        s.B32(0)
         .B32(0)
         .B32(1234)
         .B32(::std::u32::MAX)
         .append_repeated(0, 80)
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_mvhd(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"mvhd"));
    assert_eq!(parsed.header.size, 108);
    assert_eq!(parsed.timescale, 1234);
    assert_eq!(parsed.duration, ::std::u64::MAX);
}
<commit_msg>Add a test for parsing truncated (ftyp) boxes.<commit_after>//! Module for parsing ISO Base Media Format aka video/mp4 streams.
//! Internal unit tests.

// 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 https://mozilla.org/MPL/2.0/.

use std::io::Cursor;
use super::*;
extern crate test_assembler;
use self::test_assembler::*;

enum BoxSize {
    Short(u32),
    Long(u64),
    UncheckedShort(u32),
    UncheckedLong(u64),
}

fn make_box_raw<F>(size: BoxSize, name: &[u8; 4], func: F) -> Cursor<Vec<u8>>
    where F: Fn(Section) -> Section
{
    let mut section = Section::new();
    section = match size {
        BoxSize::Short(size) | BoxSize::UncheckedShort(size) => section.B32(size),
        BoxSize::Long(_) | BoxSize::UncheckedLong(_) => section.B32(1),
    };
    section = section.append_bytes(name);
    section = match size {
        // The spec allows the 32-bit size to be 0 to indicate unknown
        // length streams.  It's not clear if this is valid when using a
        // 64-bit size, so prohibit it for now.
        BoxSize::Long(size) => {
            assert!(size > 0);
            section.B64(size)
        }
        BoxSize::UncheckedLong(size) => section.B64(size),
        _ => section,
    };
    section = func(section);
    match size {
        BoxSize::Short(size) => {
            if size > 0 {
                assert_eq!(size as u64, section.size())
            }
        }
        BoxSize::Long(size) => assert_eq!(size, section.size()),
        _ => (),
    }
    Cursor::new(section.get_contents().unwrap())
}

fn make_box<F>(size: u32, name: &[u8; 4], func: F) -> Cursor<Vec<u8>>
    where F: Fn(Section) -> Section
{
    make_box_raw(BoxSize::Short(size), name, func)
}

fn make_fullbox<F>(size: u32, name: &[u8; 4], version: u8, func: F) -> Cursor<Vec<u8>>
    where F: Fn(Section) -> Section
{
    make_box_raw(BoxSize::Short(size), name, |s| {
        func(s.B8(version)
              .B8(0)
              .B8(0)
              .B8(0))
    })
}

#[test]
fn read_box_header_short() {
    let mut stream = make_box_raw(BoxSize::Short(8), b"test", |s| s);
    let parsed = read_box_header(&mut stream).unwrap();
    assert_eq!(parsed.name, FourCC(*b"test"));
    assert_eq!(parsed.size, 8);
}

#[test]
fn read_box_header_long() {
    let mut stream = make_box_raw(BoxSize::Long(16), b"test", |s| s);
    let parsed = read_box_header(&mut stream).unwrap();
    assert_eq!(parsed.name, FourCC(*b"test"));
    assert_eq!(parsed.size, 16);
}

#[test]
fn read_box_header_short_unknown_size() {
    let mut stream = make_box_raw(BoxSize::Short(0), b"test", |s| s);
    match read_box_header(&mut stream) {
        Err(Error::Unsupported) => (),
        _ => panic!("unexpected result reading box with unknown size"),
    };
}

#[test]
fn read_box_header_short_invalid_size() {
    let mut stream = make_box_raw(BoxSize::UncheckedShort(2), b"test", |s| s);
    match read_box_header(&mut stream) {
        Err(Error::InvalidData) => (),
        _ => panic!("unexpected result reading box with invalid size"),
    };
}

#[test]
fn read_box_header_long_invalid_size() {
    let mut stream = make_box_raw(BoxSize::UncheckedLong(2), b"test", |s| s);
    match read_box_header(&mut stream) {
        Err(Error::InvalidData) => (),
        _ => panic!("unexpected result reading box with invalid size"),
    };
}

#[test]
fn read_ftyp() {
    let mut stream = make_box(24, b"ftyp", |s| {
        s.append_bytes(b"mp42")
         .B32(0) // minor version
         .append_bytes(b"isom")
         .append_bytes(b"mp42")
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_ftyp(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"ftyp"));
    assert_eq!(parsed.header.size, 24);
    assert_eq!(parsed.major_brand, FourCC(*b"mp42"));
    assert_eq!(parsed.minor_version, 0);
    assert_eq!(parsed.compatible_brands.len(), 2);
    assert_eq!(parsed.compatible_brands[0], FourCC(*b"isom"));
    assert_eq!(parsed.compatible_brands[1], FourCC(*b"mp42"));
}

#[test]
#[ignore]
fn read_truncated_ftyp() {
    // We declare a 24 byte box, but only write 20 bytes.
    let mut stream = make_box_raw(BoxSize::UncheckedShort(24), b"ftyp", |s| {
        s.append_bytes(b"mp42")
            .B32(0) // minor version
            .append_bytes(b"isom")
    });
    let mut context = MediaContext::new();
    context.trace(true);
    match read_mp4(&mut stream, &mut context) {
        Err(Error::UnexpectedEOF) => (),
        Ok(_) => assert!(false, "expected an error result"),
        _ => assert!(false, "expected a different error result"),
    }
}

#[test]
fn read_elst_v0() {
    let mut stream = make_fullbox(28, b"elst", 0, |s| {
        s.B32(1) // list count
          // first entry
         .B32(1234) // duration
         .B32(5678) // time
         .B16(12) // rate integer
         .B16(34) // rate fraction
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_elst(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"elst"));
    assert_eq!(parsed.header.size, 28);
    assert_eq!(parsed.edits.len(), 1);
    assert_eq!(parsed.edits[0].segment_duration, 1234);
    assert_eq!(parsed.edits[0].media_time, 5678);
    assert_eq!(parsed.edits[0].media_rate_integer, 12);
    assert_eq!(parsed.edits[0].media_rate_fraction, 34);
}

#[test]
fn read_elst_v1() {
    let mut stream = make_fullbox(56, b"elst", 1, |s| {
        s.B32(2) // list count
         // first entry
         .B64(1234) // duration
         .B64(5678) // time
         .B16(12) // rate integer
         .B16(34) // rate fraction
         // second entry
         .B64(1234) // duration
         .B64(5678) // time
         .B16(12) // rate integer
         .B16(34) // rate fraction
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_elst(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"elst"));
    assert_eq!(parsed.header.size, 56);
    assert_eq!(parsed.edits.len(), 2);
    assert_eq!(parsed.edits[1].segment_duration, 1234);
    assert_eq!(parsed.edits[1].media_time, 5678);
    assert_eq!(parsed.edits[1].media_rate_integer, 12);
    assert_eq!(parsed.edits[1].media_rate_fraction, 34);
}

#[test]
fn read_mdhd_v0() {
    let mut stream = make_fullbox(32, b"mdhd", 0, |s| {
        s.B32(0)
         .B32(0)
         .B32(1234) // timescale
         .B32(5678) // duration
         .B32(0)
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_mdhd(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"mdhd"));
    assert_eq!(parsed.header.size, 32);
    assert_eq!(parsed.timescale, 1234);
    assert_eq!(parsed.duration, 5678);
}

#[test]
fn read_mdhd_v1() {
    let mut stream = make_fullbox(44, b"mdhd", 1, |s| {
        s.B64(0)
         .B64(0)
         .B32(1234) // timescale
         .B64(5678) // duration
         .B32(0)
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_mdhd(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"mdhd"));
    assert_eq!(parsed.header.size, 44);
    assert_eq!(parsed.timescale, 1234);
    assert_eq!(parsed.duration, 5678);
}

#[test]
fn read_mdhd_unknown_duration() {
    let mut stream = make_fullbox(32, b"mdhd", 0, |s| {
        s.B32(0)
         .B32(0)
         .B32(1234) // timescale
         .B32(::std::u32::MAX) // duration
         .B32(0)
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_mdhd(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"mdhd"));
    assert_eq!(parsed.header.size, 32);
    assert_eq!(parsed.timescale, 1234);
    assert_eq!(parsed.duration, ::std::u64::MAX);
}

#[test]
fn read_mdhd_invalid_timescale() {
    let mut stream = make_fullbox(44, b"mdhd", 1, |s| {
        s.B64(0)
         .B64(0)
         .B32(0) // timescale
         .B64(5678) // duration
         .B32(0)
    });
    let header = read_box_header(&mut stream).unwrap();
    let r = super::parse_mdhd(&mut stream, &header, &mut super::Track::new(0));
    assert_eq!(r.is_err(), true);
}

#[test]
fn read_mvhd_v0() {
    let mut stream = make_fullbox(108, b"mvhd", 0, |s| {
        s.B32(0)
         .B32(0)
         .B32(1234)
         .B32(5678)
         .append_repeated(0, 80)
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_mvhd(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"mvhd"));
    assert_eq!(parsed.header.size, 108);
    assert_eq!(parsed.timescale, 1234);
    assert_eq!(parsed.duration, 5678);
}

#[test]
fn read_mvhd_v1() {
    let mut stream = make_fullbox(120, b"mvhd", 1, |s| {
        s.B64(0)
         .B64(0)
         .B32(1234)
         .B64(5678)
         .append_repeated(0, 80)
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_mvhd(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"mvhd"));
    assert_eq!(parsed.header.size, 120);
    assert_eq!(parsed.timescale, 1234);
    assert_eq!(parsed.duration, 5678);
}

#[test]
fn read_mvhd_invalid_timescale() {
    let mut stream = make_fullbox(120, b"mvhd", 1, |s| {
        s.B64(0)
         .B64(0)
         .B32(0)
         .B64(5678)
         .append_repeated(0, 80)
    });
    let header = read_box_header(&mut stream).unwrap();
    let r = super::parse_mvhd(&mut stream, &header);
    assert_eq!(r.is_err(), true);
}

#[test]
fn read_mvhd_unknown_duration() {
    let mut stream = make_fullbox(108, b"mvhd", 0, |s| {
        s.B32(0)
         .B32(0)
         .B32(1234)
         .B32(::std::u32::MAX)
         .append_repeated(0, 80)
    });
    let header = read_box_header(&mut stream).unwrap();
    let parsed = super::read_mvhd(&mut stream, &header).unwrap();
    assert_eq!(parsed.header.name, FourCC(*b"mvhd"));
    assert_eq!(parsed.header.size, 108);
    assert_eq!(parsed.timescale, 1234);
    assert_eq!(parsed.duration, ::std::u64::MAX);
}
<|endoftext|>


<commit_before>use crate::error::ZmqError;
use bytes::{Buf, BufMut, Bytes, BytesMut};
use std::collections::HashMap;
use std::convert::TryFrom;
use std::fmt::Display;
use tokio_util::codec::{Decoder, Encoder};

use crate::SocketType;

#[derive(Debug, Copy, Clone)]
pub(crate) enum ZmqMechanism {
    NULL,
    PLAIN,
    CURVE,
}

impl Display for ZmqMechanism {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        return match self {
            ZmqMechanism::NULL => write!(f, "NULL"),
            ZmqMechanism::PLAIN => write!(f, "PLAIN"),
            ZmqMechanism::CURVE => write!(f, "CURVE"),
        };
    }
}

impl TryFrom<Vec<u8>> for ZmqMechanism {
    type Error = ZmqError;

    fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
        let mech = value.split(|x| *x == 0x0).next().unwrap_or(b"");
        // mechanism-char = "A"-"Z" | DIGIT
        //                  | "-" | "_" | "." | "+" | %x0
        // according to https://rfc.zeromq.org/spec:23/ZMTP/
        let mechanism = unsafe { String::from_utf8_unchecked(mech.to_vec()) };
        match mechanism.as_str() {
            "NULL" => Ok(ZmqMechanism::NULL),
            "PLAIN" => Ok(ZmqMechanism::PLAIN),
            "CURVE" => Ok(ZmqMechanism::CURVE),
            _ => Err(ZmqError::Other("Failed to parse ZmqMechanism")),
        }
    }
}

#[derive(Debug, Copy, Clone)]
pub(crate) struct ZmqGreeting {
    pub version: (u8, u8),
    pub mechanism: ZmqMechanism,
    pub as_server: bool,
}

impl Default for ZmqGreeting {
    fn default() -> Self {
        Self {
            version: (3, 0),
            mechanism: ZmqMechanism::NULL,
            as_server: false,
        }
    }
}

impl TryFrom<Bytes> for ZmqGreeting {
    type Error = ZmqError;

    fn try_from(value: Bytes) -> Result<Self, Self::Error> {
        if !(value[0] == 0xff && value[9] == 0x7f) {
            return Err(ZmqError::Codec("Failed to parse greeting"));
        }
        Ok(ZmqGreeting {
            version: (value[10], value[11]),
            mechanism: ZmqMechanism::try_from(value[12..32].to_vec())?,
            as_server: value[32] == 0x01,
        })
    }
}

impl From<ZmqGreeting> for BytesMut {
    fn from(greet: ZmqGreeting) -> Self {
        let mut data: [u8; 64] = [0; 64];
        data[0] = 0xff;
        data[9] = 0x7f;
        data[10] = greet.version.0;
        data[11] = greet.version.1;
        let mech = format!("{}", greet.mechanism);
        data[12..12 + mech.len()].copy_from_slice(mech.as_bytes());
        data[32] = greet.as_server.into();
        let mut bytes = BytesMut::new();
        bytes.extend_from_slice(&data);
        bytes
    }
}

#[derive(Debug, Clone)]
pub struct ZmqMessage {
    pub data: Bytes,
    pub more: bool,
}

#[derive(Debug, Clone)]
pub(crate) enum Message {
    Greeting(ZmqGreeting),
    Command(ZmqCommand),
    Message(ZmqMessage),
    MultipartMessage(Vec<ZmqMessage>),
}

#[derive(Debug, Copy, Clone)]
pub(crate) enum ZmqCommandName {
    READY,
}

impl From<ZmqCommandName> for String {
    fn from(c_name: ZmqCommandName) -> Self {
        match c_name {
            ZmqCommandName::READY => "READY".into(),
        }
    }
}

#[derive(Debug, Clone)]
pub(crate) struct ZmqCommand {
    pub name: ZmqCommandName,
    pub properties: HashMap<String, String>,
}

impl ZmqCommand {
    pub fn ready(socket: SocketType) -> Self {
        let mut properties = HashMap::new();
        properties.insert("Socket-Type".into(), format!("{}", socket));
        Self {
            name: ZmqCommandName::READY,
            properties,
        }
    }
}

impl TryFrom<BytesMut> for ZmqCommand {
    type Error = ZmqError;

    fn try_from(mut buf: BytesMut) -> Result<Self, Self::Error> {
        let command_len = buf.get_u8() as usize;
        // command-name-char = ALPHA according to https://rfc.zeromq.org/spec:23/ZMTP/
        let command_name =
            unsafe { String::from_utf8_unchecked(buf.split_to(command_len).to_vec()) };
        let command = match command_name.as_str() {
            "READY" => ZmqCommandName::READY,
            _ => return Err(ZmqError::Codec("Uknown command received")),
        };
        let mut properties = HashMap::new();

        while !buf.is_empty() {
            // Collect command properties
            let prop_len = buf.get_u8() as usize;
            let property = unsafe { String::from_utf8_unchecked(buf.split_to(prop_len).to_vec()) };

            let prop_val_len = buf.get_u32() as usize;
            let prop_value =
                unsafe { String::from_utf8_unchecked(buf.split_to(prop_val_len).to_vec()) };
            properties.insert(property, prop_value);
        }
        Ok(Self {
            name: command,
            properties,
        })
    }
}

impl From<ZmqCommand> for BytesMut {
    fn from(command: ZmqCommand) -> Self {
        let mut message_len = 0;

        let command_name: String = command.name.into();
        message_len += command_name.len() + 1;
        for (prop, val) in command.properties.iter() {
            message_len += prop.len() + 1;
            message_len += val.len() + 4;
        }

        let long_message = message_len > 255;

        let mut bytes = BytesMut::new();
        if long_message {
            bytes.reserve(message_len + 9);
            bytes.put_u8(0x06);
            bytes.put_u64(message_len as u64);
        } else {
            bytes.reserve(message_len + 2);
            bytes.put_u8(0x04);
            bytes.put_u8(message_len as u8);
        };
        bytes.put_u8(command_name.len() as u8);
        bytes.extend_from_slice(command_name.as_ref());
        for (prop, val) in command.properties.iter() {
            bytes.put_u8(prop.len() as u8);
            bytes.extend_from_slice(prop.as_ref());
            bytes.put_u32(val.len() as u32);
            bytes.extend_from_slice(val.as_ref());
        }
        bytes
    }
}

#[derive(Debug)]
enum DecoderState {
    Greeting,
    FrameHeader,
    Frame(bool, usize, bool),
}

pub(crate) struct ZmqCodec {
    state: DecoderState,
}

impl ZmqCodec {
    pub fn new() -> Self {
        Self {
            state: DecoderState::Greeting,
        }
    }
}

impl Decoder for ZmqCodec {
    type Item = Message;
    type Error = ZmqError;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        if src.is_empty() {
            return Ok(None);
        }
        match self.state {
            DecoderState::Greeting => {
                if src.len() >= 64 {
                    self.state = DecoderState::FrameHeader;
                    Ok(Some(Message::Greeting(ZmqGreeting::try_from(
                        src.split_to(64).freeze(),
                    )?)))
                } else {
                    if src[0] == 0xff {
                        src.reserve(64);
                        Ok(None)
                    } else {
                        Err(ZmqError::Codec("Bad first byte of greeting"))
                    }
                }
            }
            DecoderState::FrameHeader => {
                let flags = src[0];
                let command = (flags & 0b0000_0100) != 0;
                let long = (flags & 0b0000_0010) != 0;
                let more = ((flags & 0b0000_0001) != 0) | command;

                let frame_len = if !long && src.len() >= 2 {
                    src.advance(1);
                    src.get_u8() as usize
                } else if long && src.len() >= 9 {
                    src.advance(1);
                    src.get_u64() as usize
                } else {
                    return Ok(None);
                };
                self.state = DecoderState::Frame(command, frame_len, more);
                return self.decode(src);
            }
            DecoderState::Frame(command, frame_len, more) => {
                if src.len() < frame_len {
                    src.reserve(frame_len);
                    return Ok(None);
                }
                self.state = DecoderState::FrameHeader;
                let message = if command {
                    Message::Command(ZmqCommand::try_from(src.split_to(frame_len))?)
                } else {
                    Message::Message(ZmqMessage {
                        data: src.split_to(frame_len).freeze(),
                        more,
                    })
                };
                return Ok(Some(message));
            }
        }
    }
}

impl ZmqCodec {
    fn _encode_message(&mut self, message: ZmqMessage, dst: &mut BytesMut) {
        let mut flags: u8 = 0;
        if message.more {
            flags |= 0b0000_0001;
        }
        let len = message.data.len();
        if len > 255 {
            flags |= 0b0000_0010;
            dst.reserve(len + 9);
        } else {
            dst.reserve(len + 2);
        }
        dst.put_u8(flags);
        if len > 255 {
            dst.put_u64(len as u64);
        } else {
            dst.put_u8(len as u8);
        }
        dst.extend_from_slice(message.data.as_ref());
    }
}

impl Encoder for ZmqCodec {
    type Item = Message;
    type Error = ZmqError;

    fn encode(&mut self, message: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
        match message {
            Message::Greeting(payload) => dst.unsplit(payload.into()),
            Message::Message(message) => self._encode_message(message, dst),
            Message::Command(command) => dst.unsplit(command.into()),
            Message::MultipartMessage(parts) => {
                for part in parts {
                    self._encode_message(part, dst);
                }
            }
        }
        Ok(())
    }
}
<commit_msg>Refactored codec code a little bit<commit_after>use crate::error::ZmqError;
use bytes::{Buf, BufMut, Bytes, BytesMut};
use std::collections::HashMap;
use std::convert::TryFrom;
use std::fmt::Display;
use tokio_util::codec::{Decoder, Encoder};

use crate::SocketType;

#[derive(Debug, Copy, Clone)]
pub(crate) enum ZmqMechanism {
    NULL,
    PLAIN,
    CURVE,
}

impl Display for ZmqMechanism {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        return match self {
            ZmqMechanism::NULL => write!(f, "NULL"),
            ZmqMechanism::PLAIN => write!(f, "PLAIN"),
            ZmqMechanism::CURVE => write!(f, "CURVE"),
        };
    }
}

impl TryFrom<Vec<u8>> for ZmqMechanism {
    type Error = ZmqError;

    fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
        let mech = value.split(|x| *x == 0x0).next().unwrap_or(b"");
        // mechanism-char = "A"-"Z" | DIGIT
        //                  | "-" | "_" | "." | "+" | %x0
        // according to https://rfc.zeromq.org/spec:23/ZMTP/
        let mechanism = unsafe { String::from_utf8_unchecked(mech.to_vec()) };
        match mechanism.as_str() {
            "NULL" => Ok(ZmqMechanism::NULL),
            "PLAIN" => Ok(ZmqMechanism::PLAIN),
            "CURVE" => Ok(ZmqMechanism::CURVE),
            _ => Err(ZmqError::Other("Failed to parse ZmqMechanism")),
        }
    }
}

#[derive(Debug, Copy, Clone)]
pub(crate) struct ZmqGreeting {
    pub version: (u8, u8),
    pub mechanism: ZmqMechanism,
    pub as_server: bool,
}

impl Default for ZmqGreeting {
    fn default() -> Self {
        Self {
            version: (3, 0),
            mechanism: ZmqMechanism::NULL,
            as_server: false,
        }
    }
}

impl TryFrom<Bytes> for ZmqGreeting {
    type Error = ZmqError;

    fn try_from(value: Bytes) -> Result<Self, Self::Error> {
        if !(value[0] == 0xff && value[9] == 0x7f) {
            return Err(ZmqError::Codec("Failed to parse greeting"));
        }
        Ok(ZmqGreeting {
            version: (value[10], value[11]),
            mechanism: ZmqMechanism::try_from(value[12..32].to_vec())?,
            as_server: value[32] == 0x01,
        })
    }
}

impl From<ZmqGreeting> for BytesMut {
    fn from(greet: ZmqGreeting) -> Self {
        let mut data: [u8; 64] = [0; 64];
        data[0] = 0xff;
        data[9] = 0x7f;
        data[10] = greet.version.0;
        data[11] = greet.version.1;
        let mech = format!("{}", greet.mechanism);
        data[12..12 + mech.len()].copy_from_slice(mech.as_bytes());
        data[32] = greet.as_server.into();
        let mut bytes = BytesMut::new();
        bytes.extend_from_slice(&data);
        bytes
    }
}

#[derive(Debug, Clone)]
pub struct ZmqMessage {
    pub data: Bytes,
    pub more: bool,
}

#[derive(Debug, Clone)]
pub(crate) enum Message {
    Greeting(ZmqGreeting),
    Command(ZmqCommand),
    Message(ZmqMessage),
    MultipartMessage(Vec<ZmqMessage>),
}

#[derive(Debug, Copy, Clone)]
pub(crate) enum ZmqCommandName {
    READY,
}

impl From<ZmqCommandName> for String {
    fn from(c_name: ZmqCommandName) -> Self {
        match c_name {
            ZmqCommandName::READY => "READY".into(),
        }
    }
}

#[derive(Debug, Clone)]
pub(crate) struct ZmqCommand {
    pub name: ZmqCommandName,
    pub properties: HashMap<String, String>,
}

impl ZmqCommand {
    pub fn ready(socket: SocketType) -> Self {
        let mut properties = HashMap::new();
        properties.insert("Socket-Type".into(), format!("{}", socket));
        Self {
            name: ZmqCommandName::READY,
            properties,
        }
    }
}

impl TryFrom<BytesMut> for ZmqCommand {
    type Error = ZmqError;

    fn try_from(mut buf: BytesMut) -> Result<Self, Self::Error> {
        let command_len = buf.get_u8() as usize;
        // command-name-char = ALPHA according to https://rfc.zeromq.org/spec:23/ZMTP/
        let command_name =
            unsafe { String::from_utf8_unchecked(buf.split_to(command_len).to_vec()) };
        let command = match command_name.as_str() {
            "READY" => ZmqCommandName::READY,
            _ => return Err(ZmqError::Codec("Uknown command received")),
        };
        let mut properties = HashMap::new();

        while !buf.is_empty() {
            // Collect command properties
            let prop_len = buf.get_u8() as usize;
            let property = unsafe { String::from_utf8_unchecked(buf.split_to(prop_len).to_vec()) };

            let prop_val_len = buf.get_u32() as usize;
            let prop_value =
                unsafe { String::from_utf8_unchecked(buf.split_to(prop_val_len).to_vec()) };
            properties.insert(property, prop_value);
        }
        Ok(Self {
            name: command,
            properties,
        })
    }
}

impl From<ZmqCommand> for BytesMut {
    fn from(command: ZmqCommand) -> Self {
        let mut message_len = 0;

        let command_name: String = command.name.into();
        message_len += command_name.len() + 1;
        for (prop, val) in command.properties.iter() {
            message_len += prop.len() + 1;
            message_len += val.len() + 4;
        }

        let long_message = message_len > 255;

        let mut bytes = BytesMut::new();
        if long_message {
            bytes.reserve(message_len + 9);
            bytes.put_u8(0x06);
            bytes.put_u64(message_len as u64);
        } else {
            bytes.reserve(message_len + 2);
            bytes.put_u8(0x04);
            bytes.put_u8(message_len as u8);
        };
        bytes.put_u8(command_name.len() as u8);
        bytes.extend_from_slice(command_name.as_ref());
        for (prop, val) in command.properties.iter() {
            bytes.put_u8(prop.len() as u8);
            bytes.extend_from_slice(prop.as_ref());
            bytes.put_u32(val.len() as u32);
            bytes.extend_from_slice(val.as_ref());
        }
        bytes
    }
}

#[derive(Debug, Clone, Copy)]
struct Frame {
    command: bool,
    long: bool,
    more: bool,
}

#[derive(Debug)]
enum DecoderState {
    Greeting,
    FrameHeader,
    FrameLen(Frame),
    Frame(Frame),
}

pub(crate) struct ZmqCodec {
    state: DecoderState,
    waiting_for: usize, // Number of bytes needed to decode frame
}

impl ZmqCodec {
    pub fn new() -> Self {
        Self {
            state: DecoderState::Greeting,
            waiting_for: 64, // len of the greeting frame
        }
    }
}

impl Decoder for ZmqCodec {
    type Item = Message;
    type Error = ZmqError;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        dbg!(&src);
        if src.len() < self.waiting_for {
            src.reserve(self.waiting_for - src.len());
            return Ok(None);
        }
        match self.state {
            DecoderState::Greeting => {
                if src[0] != 0xff {
                    return Err(ZmqError::Codec("Bad first byte of greeting"))
                }
                self.state = DecoderState::FrameHeader;
                self.waiting_for = 1;
                Ok(Some(Message::Greeting(ZmqGreeting::try_from(
                    src.split_to(64).freeze(),
                )?)))
            }
            DecoderState::FrameHeader => {
                let flags = src.get_u8();
                let command = (flags & 0b0000_0100) != 0;
                let long = (flags & 0b0000_0010) != 0;
                let more = ((flags & 0b0000_0001) != 0) | command;

                let frame = Frame {
                    command,
                    long,
                    more,
                };

                self.state = DecoderState::FrameLen(frame);
                self.waiting_for = if frame.long { 8 } else { 1 };
                return self.decode(src);
            }
            DecoderState::FrameLen(frame) => {
                self.state = DecoderState::Frame(frame);
                self.waiting_for = if frame.long {
                    src.get_u64() as usize
                } else {
                    src.get_u8() as usize
                };
                return self.decode(src);
            }
            DecoderState::Frame(frame) => {
                let message = if frame.command {
                    Message::Command(ZmqCommand::try_from(src.split_to(self.waiting_for))?)
                } else {
                    Message::Message(ZmqMessage {
                        data: src.split_to(self.waiting_for).freeze(),
                        more: frame.more,
                    })
                };
                self.state = DecoderState::FrameHeader;
                self.waiting_for = 1;
                return Ok(Some(message));
            }
        }
    }
}

impl ZmqCodec {
    fn _encode_message(&mut self, message: ZmqMessage, dst: &mut BytesMut) {
        let mut flags: u8 = 0;
        if message.more {
            flags |= 0b0000_0001;
        }
        let len = message.data.len();
        if len > 255 {
            flags |= 0b0000_0010;
            dst.reserve(len + 9);
        } else {
            dst.reserve(len + 2);
        }
        dst.put_u8(flags);
        if len > 255 {
            dst.put_u64(len as u64);
        } else {
            dst.put_u8(len as u8);
        }
        dst.extend_from_slice(message.data.as_ref());
    }
}

impl Encoder for ZmqCodec {
    type Item = Message;
    type Error = ZmqError;

    fn encode(&mut self, message: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
        match message {
            Message::Greeting(payload) => dst.unsplit(payload.into()),
            Message::Message(message) => self._encode_message(message, dst),
            Message::Command(command) => dst.unsplit(command.into()),
            Message::MultipartMessage(parts) => {
                for part in parts {
                    self._encode_message(part, dst);
                }
            }
        }
        Ok(())
    }
}
<|endoftext|>