<commit_before>use codegen::buffer::*;
use codegen::info;
use codegen::x64::reg::*;
use codegen::x64::reg::Reg::*;
use codegen::x64::emit::*;

use parser::ast::ctxt::*;
use parser::ast::*;
use parser::ast::Expr::*;
use parser::ast::Stmt::*;
use parser::ast::visit::*;

use sym::BuiltinType;

pub struct CodeGen<'a, 'ast: 'a> {
    ctxt: &'a Context<'a, 'ast>,
    fct: &'ast Function,
    buf: Buffer,

    stacksize: u32,
    lbl_break: Option<Label>,
    lbl_continue: Option<Label>,
}

impl<'a, 'ast> CodeGen<'a, 'ast> where 'ast: 'a {
    pub fn new(ctxt: &'a Context<'a, 'ast>, fct: &'ast Function) -> CodeGen<'a, 'ast> {
        CodeGen {
            ctxt: ctxt,
            fct: fct,
            buf: Buffer::new(),

            stacksize: ctxt.function(fct.id, |fct| fct.stacksize),
            lbl_break: None,
            lbl_continue: None
        }
    }

    pub fn generate(mut self) -> Buffer {
        info::generate(self.ctxt, self.fct);

        self.emit_prolog();
        self.store_params();
        self.visit_fct(self.fct);
        self.emit_epilog();

        self.buf
    }

    fn store_params(&mut self) {
        let params_len = self.fct.params.len();
        if params_len == 0 { return; }

        let defs = self.ctxt.defs.borrow();

        for (reg, p) in PARAM_REGS.iter().zip(&self.fct.params) {
            let varid = *defs.get(&p.id).unwrap();
            self.emit_var_store(*reg, varid);
        }
    }

    fn emit_prolog(&mut self) {
        emit_pushq_reg(&mut self.buf, Reg::RBP);
        emit_movq_reg_reg(&mut self.buf, Reg::RSP, Reg::RBP);

        if self.stacksize > 0 {
            emit_subq_imm_reg(&mut self.buf, self.stacksize as i32, RSP);
        }
    }

    fn emit_epilog(&mut self) {
        if self.stacksize > 0 {
            emit_addq_imm_reg(&mut self.buf, self.stacksize as i32, RSP);
        }

        emit_popq_reg(&mut self.buf, Reg::RBP);
        emit_retq(&mut self.buf);
    }

    fn emit_stmt_return(&mut self, s: &'ast StmtReturnType) {
        if let Some(ref expr) = s.expr {
            self.visit_expr(expr);
        }

        self.emit_epilog();
    }

    fn emit_stmt_while(&mut self, s: &'ast StmtWhileType) {
        let lbl_start = self.buf.create_label();
        let lbl_end = self.buf.create_label();

        self.buf.define_label(lbl_start);

        // execute condition, when condition is false jump to
        // end of while
        self.visit_expr(&s.cond);
        emit_testl_reg_reg(&mut self.buf, RAX, RAX);
        emit_jz(&mut self.buf, lbl_end);

        self.save_label_state(lbl_end, lbl_start, |this| {
            // execute while body, then jump back to condition
            this.visit_stmt(&s.block);
            emit_jmp(&mut this.buf, lbl_start);
        });

        self.buf.define_label(lbl_end);
    }

    fn emit_stmt_loop(&mut self, s: &'ast StmtLoopType) {
        let lbl_start = self.buf.create_label();
        let lbl_end = self.buf.create_label();
        self.buf.define_label(lbl_start);

        self.save_label_state(lbl_end, lbl_start, |this| {
            this.visit_stmt(&s.block);
        });

        self.buf.define_label(lbl_end);

        emit_jmp(&mut self.buf, lbl_start);
    }

    fn save_label_state<F>(&mut self, lbl_break: Label, lbl_continue: Label, f: F)
            where F: FnOnce(&mut CodeGen<'a, 'ast>) {
        let old_lbl_break = self.lbl_break;
        let old_lbl_continue = self.lbl_continue;

        self.lbl_break = Some(lbl_break);
        self.lbl_continue = Some(lbl_continue);

        f(self);

        self.lbl_break = old_lbl_break;
        self.lbl_continue = old_lbl_continue;
    }

    fn emit_stmt_if(&mut self, s: &'ast StmtIfType) {
        let lbl_end = self.buf.create_label();
        let lbl_else = if let Some(_) = s.else_block {
            self.buf.create_label()
        } else {
            lbl_end
        };

        self.visit_expr(&s.cond);
        emit_testl_reg_reg(&mut self.buf, RAX, RAX);
        emit_jz(&mut self.buf, lbl_else);

        self.visit_stmt(&s.then_block);

        if let Some(ref else_block) = s.else_block {
            emit_jmp(&mut self.buf, lbl_end);
            self.buf.define_label(lbl_else);

            self.visit_stmt(else_block);
        }

        self.buf.define_label(lbl_end);
    }

    fn emit_stmt_break(&mut self, s: &'ast StmtBreakType)  {
        emit_jmp(&mut self.buf, self.lbl_break.unwrap());
    }

    fn emit_stmt_continue(&mut self, s: &'ast StmtContinueType) {
        emit_jmp(&mut self.buf, self.lbl_continue.unwrap());
    }

    fn emit_stmt_expr(&mut self, s: &'ast StmtExprType) {
        self.visit_expr(&s.expr);
    }

    fn emit_stmt_block(&mut self, s: &'ast StmtBlockType) {
        for stmt in &s.stmts {
            self.visit_stmt(stmt);
        }
    }

    fn emit_stmt_var(&mut self, s: &'ast StmtVarType) {
        if let Some(ref expr) = s.expr {
            self.visit_expr(expr);

            let defs = self.ctxt.defs.borrow();
            let varid = *defs.get(&s.id).unwrap();

            self.emit_var_store(RAX, varid);
        }
    }

    fn emit_expr_lit_int(&mut self, lit: &'ast ExprLitIntType) {
        emit_movl_imm_reg(&mut self.buf, lit.value as u32, Reg::RAX);
    }

    fn emit_expr_lit_bool(&mut self, lit: &'ast ExprLitBoolType) {
        let value : u32 = if lit.value { 1 } else { 0 };
        emit_movl_imm_reg(&mut self.buf, value, Reg::RAX);
    }

    fn emit_expr_ident(&mut self, e: &'ast ExprIdentType) {
        let defs = self.ctxt.defs.borrow();
        let varid = *defs.get(&e.id).unwrap();

        self.emit_var_load(varid, RAX);
    }

    fn emit_var_store(&mut self, src: Reg, var: VarInfoId) {
        let var_infos = self.ctxt.var_infos.borrow();
        let var = &var_infos[var.0];

        assert_eq!(BuiltinType::Int, var.data_type);
        emit_movl_reg_memq(&mut self.buf, src, RBP, var.offset);
    }

    fn emit_var_load(&mut self, var: VarInfoId, dest: Reg) {
        let var_infos = self.ctxt.var_infos.borrow();
        let var = &var_infos[var.0];

        assert_eq!(BuiltinType::Int, var.data_type);
        emit_movl_memq_reg(&mut self.buf, RBP, var.offset, dest);
    }
}

impl<'a, 'ast> visit::Visitor<'ast> for CodeGen<'a, 'ast> {
    fn visit_stmt(&mut self, s: &'ast Stmt) {
        match *s {
            StmtExpr(ref stmt) => self.emit_stmt_expr(stmt),
            StmtIf(ref stmt) => self.emit_stmt_if(stmt),
            StmtLoop(ref stmt) => self.emit_stmt_loop(stmt),
            StmtWhile(ref stmt) => self.emit_stmt_while(stmt),
            StmtReturn(ref stmt) => self.emit_stmt_return(stmt),
            StmtBreak(ref stmt) => self.emit_stmt_break(stmt),
            StmtContinue(ref stmt) => self.emit_stmt_continue(stmt),
            StmtBlock(ref stmt) => self.emit_stmt_block(stmt),
            StmtVar(ref stmt) => self.emit_stmt_var(stmt),
        }
    }

    fn visit_expr(&mut self, e: &'ast Expr) {
        match *e {
            ExprLitInt(ref expr) => self.emit_expr_lit_int(expr),
            ExprLitBool(ref expr) => self.emit_expr_lit_bool(expr),
            ExprIdent(ref expr) => self.emit_expr_ident(expr),
            _ => unreachable!()
        }
    }
}

#[cfg(test)]
mod tests {
    use std::mem;

    use mem::CodeMemory;
    use semck;
    use test;

    use super::*;

    fn run<T>(code: &'static str) -> T {
        test::parse(code, |ctxt| {
            // generate code for first function
            let fct = ctxt.ast.elements[0].to_function().unwrap();

            let buffer = CodeGen::new(&ctxt, fct).generate().finish();
            let mem = CodeMemory::new(&buffer);

            let compiled_fct : extern "C" fn() -> T = unsafe { mem::transmute(mem.ptr()) };

            compiled_fct()
        })
    }

    #[test]
    fn test_lit_int() {
        assert_eq!(1i32, run("fn f() -> int { return 1; }"));
        assert_eq!(2i32, run("fn f() -> int { return 2; }"));
        assert_eq!(3i32, run("fn f() -> int { return 3; }"));
    }

    #[test]
    fn test_lit_bool() {
        assert_eq!(true, run("fn f() -> bool { return true; }"));
        assert_eq!(false, run("fn f() -> bool { return false; }"));
    }

    #[test]
    fn test_if() {
        assert_eq!(1i32, run("fn f() -> int { if true { return 1; } else { return 2; } }"));
        assert_eq!(1i32, run("fn f() -> int { if true { return 1; } return 2; }"));
        assert_eq!(2i32, run("fn f() -> int { if false { return 1; } else { return 2; } }"));
        assert_eq!(2i32, run("fn f() -> int { if false { return 1; } return 2; }"));
    }

    #[test]
    fn test_ident_load_and_store() {
        assert_eq!(4711, run("fn f() -> int { var a = 4711; return a; }"));
    }
}
<commit_msg>codegen: made it possible to test fct params<commit_after>use codegen::buffer::*;
use codegen::info;
use codegen::x64::reg::*;
use codegen::x64::reg::Reg::*;
use codegen::x64::emit::*;

use parser::ast::ctxt::*;
use parser::ast::*;
use parser::ast::Expr::*;
use parser::ast::Stmt::*;
use parser::ast::visit::*;

use sym::BuiltinType;

pub struct CodeGen<'a, 'ast: 'a> {
    ctxt: &'a Context<'a, 'ast>,
    fct: &'ast Function,
    buf: Buffer,

    stacksize: u32,
    lbl_break: Option<Label>,
    lbl_continue: Option<Label>,
}

impl<'a, 'ast> CodeGen<'a, 'ast> where 'ast: 'a {
    pub fn new(ctxt: &'a Context<'a, 'ast>, fct: &'ast Function) -> CodeGen<'a, 'ast> {
        CodeGen {
            ctxt: ctxt,
            fct: fct,
            buf: Buffer::new(),

            stacksize: ctxt.function(fct.id, |fct| fct.stacksize),
            lbl_break: None,
            lbl_continue: None
        }
    }

    pub fn generate(mut self) -> Buffer {
        info::generate(self.ctxt, self.fct);

        self.emit_prolog();
        self.store_params();
        self.visit_fct(self.fct);
        self.emit_epilog();

        self.buf
    }

    fn store_params(&mut self) {
        let params_len = self.fct.params.len();
        if params_len == 0 { return; }

        let defs = self.ctxt.defs.borrow();

        for (reg, p) in PARAM_REGS.iter().zip(&self.fct.params) {
            let varid = *defs.get(&p.id).unwrap();
            self.emit_var_store(*reg, varid);
        }
    }

    fn emit_prolog(&mut self) {
        emit_pushq_reg(&mut self.buf, Reg::RBP);
        emit_movq_reg_reg(&mut self.buf, Reg::RSP, Reg::RBP);

        if self.stacksize > 0 {
            emit_subq_imm_reg(&mut self.buf, self.stacksize as i32, RSP);
        }
    }

    fn emit_epilog(&mut self) {
        if self.stacksize > 0 {
            emit_addq_imm_reg(&mut self.buf, self.stacksize as i32, RSP);
        }

        emit_popq_reg(&mut self.buf, Reg::RBP);
        emit_retq(&mut self.buf);
    }

    fn emit_stmt_return(&mut self, s: &'ast StmtReturnType) {
        if let Some(ref expr) = s.expr {
            self.visit_expr(expr);
        }

        self.emit_epilog();
    }

    fn emit_stmt_while(&mut self, s: &'ast StmtWhileType) {
        let lbl_start = self.buf.create_label();
        let lbl_end = self.buf.create_label();

        self.buf.define_label(lbl_start);

        // execute condition, when condition is false jump to
        // end of while
        self.visit_expr(&s.cond);
        emit_testl_reg_reg(&mut self.buf, RAX, RAX);
        emit_jz(&mut self.buf, lbl_end);

        self.save_label_state(lbl_end, lbl_start, |this| {
            // execute while body, then jump back to condition
            this.visit_stmt(&s.block);
            emit_jmp(&mut this.buf, lbl_start);
        });

        self.buf.define_label(lbl_end);
    }

    fn emit_stmt_loop(&mut self, s: &'ast StmtLoopType) {
        let lbl_start = self.buf.create_label();
        let lbl_end = self.buf.create_label();
        self.buf.define_label(lbl_start);

        self.save_label_state(lbl_end, lbl_start, |this| {
            this.visit_stmt(&s.block);
        });

        self.buf.define_label(lbl_end);

        emit_jmp(&mut self.buf, lbl_start);
    }

    fn save_label_state<F>(&mut self, lbl_break: Label, lbl_continue: Label, f: F)
            where F: FnOnce(&mut CodeGen<'a, 'ast>) {
        let old_lbl_break = self.lbl_break;
        let old_lbl_continue = self.lbl_continue;

        self.lbl_break = Some(lbl_break);
        self.lbl_continue = Some(lbl_continue);

        f(self);

        self.lbl_break = old_lbl_break;
        self.lbl_continue = old_lbl_continue;
    }

    fn emit_stmt_if(&mut self, s: &'ast StmtIfType) {
        let lbl_end = self.buf.create_label();
        let lbl_else = if let Some(_) = s.else_block {
            self.buf.create_label()
        } else {
            lbl_end
        };

        self.visit_expr(&s.cond);
        emit_testl_reg_reg(&mut self.buf, RAX, RAX);
        emit_jz(&mut self.buf, lbl_else);

        self.visit_stmt(&s.then_block);

        if let Some(ref else_block) = s.else_block {
            emit_jmp(&mut self.buf, lbl_end);
            self.buf.define_label(lbl_else);

            self.visit_stmt(else_block);
        }

        self.buf.define_label(lbl_end);
    }

    fn emit_stmt_break(&mut self, s: &'ast StmtBreakType)  {
        emit_jmp(&mut self.buf, self.lbl_break.unwrap());
    }

    fn emit_stmt_continue(&mut self, s: &'ast StmtContinueType) {
        emit_jmp(&mut self.buf, self.lbl_continue.unwrap());
    }

    fn emit_stmt_expr(&mut self, s: &'ast StmtExprType) {
        self.visit_expr(&s.expr);
    }

    fn emit_stmt_block(&mut self, s: &'ast StmtBlockType) {
        for stmt in &s.stmts {
            self.visit_stmt(stmt);
        }
    }

    fn emit_stmt_var(&mut self, s: &'ast StmtVarType) {
        if let Some(ref expr) = s.expr {
            self.visit_expr(expr);

            let defs = self.ctxt.defs.borrow();
            let varid = *defs.get(&s.id).unwrap();

            self.emit_var_store(RAX, varid);
        }
    }

    fn emit_expr_lit_int(&mut self, lit: &'ast ExprLitIntType) {
        emit_movl_imm_reg(&mut self.buf, lit.value as u32, Reg::RAX);
    }

    fn emit_expr_lit_bool(&mut self, lit: &'ast ExprLitBoolType) {
        let value : u32 = if lit.value { 1 } else { 0 };
        emit_movl_imm_reg(&mut self.buf, value, Reg::RAX);
    }

    fn emit_expr_ident(&mut self, e: &'ast ExprIdentType) {
        let defs = self.ctxt.defs.borrow();
        let varid = *defs.get(&e.id).unwrap();

        self.emit_var_load(varid, RAX);
    }

    fn emit_var_store(&mut self, src: Reg, var: VarInfoId) {
        let var_infos = self.ctxt.var_infos.borrow();
        let var = &var_infos[var.0];

        assert_eq!(BuiltinType::Int, var.data_type);
        emit_movl_reg_memq(&mut self.buf, src, RBP, var.offset);
    }

    fn emit_var_load(&mut self, var: VarInfoId, dest: Reg) {
        let var_infos = self.ctxt.var_infos.borrow();
        let var = &var_infos[var.0];

        assert_eq!(BuiltinType::Int, var.data_type);
        emit_movl_memq_reg(&mut self.buf, RBP, var.offset, dest);
    }
}

impl<'a, 'ast> visit::Visitor<'ast> for CodeGen<'a, 'ast> {
    fn visit_stmt(&mut self, s: &'ast Stmt) {
        match *s {
            StmtExpr(ref stmt) => self.emit_stmt_expr(stmt),
            StmtIf(ref stmt) => self.emit_stmt_if(stmt),
            StmtLoop(ref stmt) => self.emit_stmt_loop(stmt),
            StmtWhile(ref stmt) => self.emit_stmt_while(stmt),
            StmtReturn(ref stmt) => self.emit_stmt_return(stmt),
            StmtBreak(ref stmt) => self.emit_stmt_break(stmt),
            StmtContinue(ref stmt) => self.emit_stmt_continue(stmt),
            StmtBlock(ref stmt) => self.emit_stmt_block(stmt),
            StmtVar(ref stmt) => self.emit_stmt_var(stmt),
        }
    }

    fn visit_expr(&mut self, e: &'ast Expr) {
        match *e {
            ExprLitInt(ref expr) => self.emit_expr_lit_int(expr),
            ExprLitBool(ref expr) => self.emit_expr_lit_bool(expr),
            ExprIdent(ref expr) => self.emit_expr_ident(expr),
            _ => unreachable!()
        }
    }
}

#[cfg(test)]
mod tests {
    use std::mem;

    use mem::CodeMemory;
    use semck;
    use test;

    use super::*;

    fn compile(code: &'static str) -> CodeMemory {
        test::parse(code, |ctxt| {
            // generate code for first function
            let fct = ctxt.ast.elements[0].to_function().unwrap();

            let buffer = CodeGen::new(&ctxt, fct).generate().finish();

            CodeMemory::new(&buffer)
        })
    }

    fn run<T>(code: &'static str) -> T {
        let mem = compile(code);
        let compiled_fct : extern "C" fn() -> T = unsafe { mem::transmute(mem.ptr()) };

        compiled_fct()
    }

    fn fct1<T>(code: &'static str) -> extern "C" fn(T) -> T {
        let mem = compile(code);

        unsafe { mem::transmute(mem.ptr()) }
    }

    fn fct2<T>(code: &'static str) -> extern "C" fn(T, T) -> T {
        let mem = compile(code);

        unsafe { mem::transmute(mem.ptr()) }
    }

    #[test]
    fn test_lit_int() {
        assert_eq!(1i32, run("fn f() -> int { return 1; }"));
        assert_eq!(2i32, run("fn f() -> int { return 2; }"));
        assert_eq!(3i32, run("fn f() -> int { return 3; }"));
    }

    // #[test]
    // fn test_param() {
    //     let f = fct1("fn f(a: int) -> int { return a; }");
    //     assert_eq!(1i32, f(1));
    // }

    #[test]
    fn test_lit_bool() {
        assert_eq!(true, run("fn f() -> bool { return true; }"));
        assert_eq!(false, run("fn f() -> bool { return false; }"));
    }

    #[test]
    fn test_if() {
        assert_eq!(1i32, run("fn f() -> int { if true { return 1; } else { return 2; } }"));
        assert_eq!(1i32, run("fn f() -> int { if true { return 1; } return 2; }"));
        assert_eq!(2i32, run("fn f() -> int { if false { return 1; } else { return 2; } }"));
        assert_eq!(2i32, run("fn f() -> int { if false { return 1; } return 2; }"));
    }

    #[test]
    fn test_ident_load_and_store() {
        assert_eq!(4711, run("fn f() -> int { var a = 4711; return a; }"));
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#![no_main]

use std::fs::File;
use std::io::{Cursor, Read, Seek, SeekFrom};
use std::mem::size_of;
use std::os::unix::io::{AsRawFd, FromRawFd};
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;

use cros_fuzz::fuzz_target;
use devices::virtio::{Block, Interrupt, Queue, VirtioDevice};
use sys_util::{EventFd, GuestAddress, GuestMemory, SharedMemory};

const MEM_SIZE: u64 = 256 * 1024 * 1024;
const DESC_SIZE: u64 = 16; // Bytes in one virtio descriptor.
const QUEUE_SIZE: u16 = 16; // Max entries in the queue.
const CMD_SIZE: usize = 16; // Bytes in the command.

fuzz_target!(|bytes| {
    let size_u64 = size_of::<u64>();
    let mem = GuestMemory::new(&[(GuestAddress(0), MEM_SIZE)]).unwrap();

    // The fuzz data is interpreted as:
    // starting index 8 bytes
    // command location 8 bytes
    // command 16 bytes
    // descriptors circular buffer 16 bytes * 3
    if bytes.len() < 4 * size_u64 {
        // Need an index to start.
        return;
    }

    let mut data_image = Cursor::new(bytes);

    let first_index = read_u64(&mut data_image);
    if first_index > MEM_SIZE / DESC_SIZE {
        return;
    }
    let first_offset = first_index * DESC_SIZE;
    if first_offset as usize + size_u64 > bytes.len() {
        return;
    }

    let command_addr = read_u64(&mut data_image);
    if command_addr > MEM_SIZE - CMD_SIZE as u64 {
        return;
    }
    if mem
        .write_all_at_addr(
            &bytes[2 * size_u64..(2 * size_u64) + CMD_SIZE],
            GuestAddress(command_addr as u64),
        )
        .is_err()
    {
        return;
    }

    data_image.seek(SeekFrom::Start(first_offset)).unwrap();
    let desc_table = read_u64(&mut data_image);

    if mem
        .write_all_at_addr(&bytes[32..], GuestAddress(desc_table as u64))
        .is_err()
    {
        return;
    }

    let mut q = Queue::new(QUEUE_SIZE);
    q.ready = true;
    q.size = QUEUE_SIZE / 2;
    q.max_size = QUEUE_SIZE;

    let queue_evts: Vec<EventFd> = vec![EventFd::new().unwrap()];
    let queue_fd = queue_evts[0].as_raw_fd();
    let queue_evt = unsafe { EventFd::from_raw_fd(libc::dup(queue_fd)) };

    let shm = SharedMemory::anon().unwrap();
    let disk_file: File = shm.into();
    let mut block = Block::new(Box::new(disk_file), false, true, None).unwrap();

    block.activate(
        mem,
        Interrupt::new(
            Arc::new(AtomicUsize::new(0)),
            EventFd::new().unwrap(),
            EventFd::new().unwrap(),
            None,   // msix_config
            0xFFFF, // VIRTIO_MSI_NO_VECTOR
        ),
        vec![q],
        queue_evts,
    );

    queue_evt.write(77).unwrap(); // Rings the doorbell, any byte will do.
});

fn read_u64<T: Read>(readable: &mut T) -> u64 {
    let mut buf = [0u8; size_of::<u64>()];
    readable.read_exact(&mut buf[..]).unwrap();
    u64::from_le_bytes(buf)
}
<commit_msg>fuzz: add block_size to block_fuzzer constructor<commit_after>// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#![no_main]

use std::fs::File;
use std::io::{Cursor, Read, Seek, SeekFrom};
use std::mem::size_of;
use std::os::unix::io::{AsRawFd, FromRawFd};
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;

use cros_fuzz::fuzz_target;
use devices::virtio::{Block, Interrupt, Queue, VirtioDevice};
use sys_util::{EventFd, GuestAddress, GuestMemory, SharedMemory};

const MEM_SIZE: u64 = 256 * 1024 * 1024;
const DESC_SIZE: u64 = 16; // Bytes in one virtio descriptor.
const QUEUE_SIZE: u16 = 16; // Max entries in the queue.
const CMD_SIZE: usize = 16; // Bytes in the command.

fuzz_target!(|bytes| {
    let size_u64 = size_of::<u64>();
    let mem = GuestMemory::new(&[(GuestAddress(0), MEM_SIZE)]).unwrap();

    // The fuzz data is interpreted as:
    // starting index 8 bytes
    // command location 8 bytes
    // command 16 bytes
    // descriptors circular buffer 16 bytes * 3
    if bytes.len() < 4 * size_u64 {
        // Need an index to start.
        return;
    }

    let mut data_image = Cursor::new(bytes);

    let first_index = read_u64(&mut data_image);
    if first_index > MEM_SIZE / DESC_SIZE {
        return;
    }
    let first_offset = first_index * DESC_SIZE;
    if first_offset as usize + size_u64 > bytes.len() {
        return;
    }

    let command_addr = read_u64(&mut data_image);
    if command_addr > MEM_SIZE - CMD_SIZE as u64 {
        return;
    }
    if mem
        .write_all_at_addr(
            &bytes[2 * size_u64..(2 * size_u64) + CMD_SIZE],
            GuestAddress(command_addr as u64),
        )
        .is_err()
    {
        return;
    }

    data_image.seek(SeekFrom::Start(first_offset)).unwrap();
    let desc_table = read_u64(&mut data_image);

    if mem
        .write_all_at_addr(&bytes[32..], GuestAddress(desc_table as u64))
        .is_err()
    {
        return;
    }

    let mut q = Queue::new(QUEUE_SIZE);
    q.ready = true;
    q.size = QUEUE_SIZE / 2;
    q.max_size = QUEUE_SIZE;

    let queue_evts: Vec<EventFd> = vec![EventFd::new().unwrap()];
    let queue_fd = queue_evts[0].as_raw_fd();
    let queue_evt = unsafe { EventFd::from_raw_fd(libc::dup(queue_fd)) };

    let shm = SharedMemory::anon().unwrap();
    let disk_file: File = shm.into();
    let mut block = Block::new(Box::new(disk_file), false, true, 512, None).unwrap();

    block.activate(
        mem,
        Interrupt::new(
            Arc::new(AtomicUsize::new(0)),
            EventFd::new().unwrap(),
            EventFd::new().unwrap(),
            None,   // msix_config
            0xFFFF, // VIRTIO_MSI_NO_VECTOR
        ),
        vec![q],
        queue_evts,
    );

    queue_evt.write(77).unwrap(); // Rings the doorbell, any byte will do.
});

fn read_u64<T: Read>(readable: &mut T) -> u64 {
    let mut buf = [0u8; size_of::<u64>()];
    readable.read_exact(&mut buf[..]).unwrap();
    u64::from_le_bytes(buf)
}
<|endoftext|>


<commit_before>// http://cr.yp.to/chacha/chacha-20080128.pdf
// http://cr.yp.to/chacha.html

// convert $e.slice($i, $i + 4) into u32
macro_rules! to_le_u32 {
    ($e:ident[$i:expr]) => ({
        let i: usize = $i;
        let v1 = $e[i + 0] as u32;
        let v2 = $e[i + 1] as u32;
        let v3 = $e[i + 2] as u32;
        let v4 = $e[i + 3] as u32;
        v1 | (v2 << 8) | (v3 << 16) | (v4 << 24)
    })
}

pub struct ChaCha20 {
    // SECRET
    vals: [u32; 16],
}

impl ChaCha20 {
    // key: SECRET
    pub fn new(key: &[u8], nonce: &[u8]) -> ChaCha20 {
        assert_eq!(key.len(), 32);
        assert_eq!(nonce.len(), 8);

        let mut vals = [0u32; 16];

        // "expand 32-byte k"
        vals[0] = 0x61707865;
        vals[1] = 0x3320646e;
        vals[2] = 0x79622d32;
        vals[3] = 0x6b206574;

        for i in (0..8) {
            vals[4 + i] = to_le_u32!(key[4 * i]);
        }

        // counter
        vals[12] = 0;
        vals[13] = 0;

        vals[14] = to_le_u32!(nonce[0]);
        vals[15] = to_le_u32!(nonce[4]);

        ChaCha20 {
            vals: vals,
        }
    }

    fn round20(&self) -> [u32; 16] {
        // $e must be > 0 and < 32
        macro_rules! rot {
            ($a:expr, $e:expr) => ({
                let a: u32 = $a;
                let e: usize = $e;
                (a << e) | (a >> (32 - e))
            })
        }

        macro_rules! quarter_round {
            ($a:expr, $b:expr, $c:expr, $d:expr) => ({
                $a += $b;
                $d ^= $a;
                $d = rot!($d, 16);

                $c += $d;
                $b ^= $c;
                $b = rot!($b, 12);

                $a += $b;
                $d ^= $a;
                $d = rot!($d, 8);

                $c += $d;
                $b ^= $c;
                $b = rot!($b, 7);
            })
        }

        macro_rules! quarter_round_idx {
            ($e:expr, $a:expr, $b:expr, $c:expr, $d:expr) => (
                quarter_round!($e[$a], $e[$b], $e[$c], $e[$d])
            )
        }

        let mut vals = self.vals;
        for _ in (0..10) {
            // column round
            quarter_round_idx!(vals, 0, 4, 8, 12);
            quarter_round_idx!(vals, 1, 5, 9, 13);
            quarter_round_idx!(vals, 2, 6, 10, 14);
            quarter_round_idx!(vals, 3, 7, 11, 15);

            // diagonal round
            quarter_round_idx!(vals, 0, 5, 10, 15);
            quarter_round_idx!(vals, 1, 6, 11, 12);
            quarter_round_idx!(vals, 2, 7, 8, 13);
            quarter_round_idx!(vals, 3, 4, 9, 14);
        }

        for i in (0..16) {
            vals[i] += self.vals[i];
        }

        vals
    }

    pub fn next(&mut self) -> [u8; 64] {
        let next = self.round20();

        // in TLS, vals[13] never increases
        {
            self.vals[12] += 1;
            //let mut count = (self.vals[12] as u64) | (self.vals[13] as u64 << 32);
            //count += 1;
            //self.vals[12] = count as u32;
            //self.vals[13] = (count >> 32) as u32;
        }

        let next_bytes = {
            let mut next_bytes = [0u8; 64];
            for i in (0..16) {
                next_bytes[4 * i + 0] = next[i] as u8;
                next_bytes[4 * i + 1] = (next[i] >> 8) as u8;
                next_bytes[4 * i + 2] = (next[i] >> 16) as u8;
                next_bytes[4 * i + 3] = (next[i] >> 24) as u8;
            }
            next_bytes
        };

        next_bytes
    }

    // Do not use same nonce for more than 2^70 bytes.
    //
    // if data is 1 byte, it still produces 64 bytes then 63 bytes are just discarded.
    // so this is not suitable for "byte-streaming" mode.
    //
    // data: SECRET
    pub fn encrypt(&mut self, data: &[u8]) -> Vec<u8> {
        let mut ret: Vec<u8> = Vec::new();

        for chunk in data.chunks(64) {
            let next = self.next();
            let xor_iter = next.iter().zip(chunk.iter()).map(|(&x, &y)| x ^ y);
            ret.extend(xor_iter);
        }

        ret
    }
}

#[cfg(test)]
mod test {
    use std::iter::repeat;

    use super::ChaCha20;

    fn check_keystream(key: &[u8], nonce: &[u8], keystream: &[u8]) {
        let mut chacha = ChaCha20::new(key, nonce);
        let input: Vec<_> = repeat(0u8).take(keystream.len()).collect();
        let output = chacha.encrypt(&input);
        assert_eq!(&output[..], keystream);
    }

    #[test]
    fn test_chacha20() {
        // from https://tools.ietf.org/html/draft-agl-tls-chacha20poly1305-04

        let mut key = [0u8; 32];
        let mut nonce = [0u8; 8];
        let keystream = b"\x76\xb8\xe0\xad\xa0\xf1\x3d\x90\x40\x5d\x6a\xe5\x53\x86\xbd\x28\
                          \xbd\xd2\x19\xb8\xa0\x8d\xed\x1a\xa8\x36\xef\xcc\x8b\x77\x0d\xc7\
                          \xda\x41\x59\x7c\x51\x57\x48\x8d\x77\x24\xe0\x3f\xb8\xd8\x4a\x37\
                          \x6a\x43\xb8\xf4\x15\x18\xa1\x1c\xc3\x87\xb6\x69\xb2\xee\x65\x86";
        check_keystream(&key, &nonce, keystream);

        key[31] = 1;
        let keystream = b"\x45\x40\xf0\x5a\x9f\x1f\xb2\x96\xd7\x73\x6e\x7b\x20\x8e\x3c\x96\
                          \xeb\x4f\xe1\x83\x46\x88\xd2\x60\x4f\x45\x09\x52\xed\x43\x2d\x41\
                          \xbb\xe2\xa0\xb6\xea\x75\x66\xd2\xa5\xd1\xe7\xe2\x0d\x42\xaf\x2c\
                          \x53\xd7\x92\xb1\xc4\x3f\xea\x81\x7e\x9a\xd2\x75\xae\x54\x69\x63";
        check_keystream(&key, &nonce, keystream);

        key[31] = 0;
        nonce[7] = 1;
        let keystream = b"\xde\x9c\xba\x7b\xf3\xd6\x9e\xf5\xe7\x86\xdc\x63\x97\x3f\x65\x3a\
                          \x0b\x49\xe0\x15\xad\xbf\xf7\x13\x4f\xcb\x7d\xf1\x37\x82\x10\x31\
                          \xe8\x5a\x05\x02\x78\xa7\x08\x45\x27\x21\x4f\x73\xef\xc7\xfa\x5b\
                          \x52\x77\x06\x2e\xb7\xa0\x43\x3e\x44\x5f\x41\xe3";
        check_keystream(&key, &nonce, keystream);

        key[31] = 0;
        nonce[7] = 0;
        nonce[0] = 1;
        let keystream = b"\xef\x3f\xdf\xd6\xc6\x15\x78\xfb\xf5\xcf\x35\xbd\x3d\xd3\x3b\x80\
                          \x09\x63\x16\x34\xd2\x1e\x42\xac\x33\x96\x0b\xd1\x38\xe5\x0d\x32\
                          \x11\x1e\x4c\xaf\x23\x7e\xe5\x3c\xa8\xad\x64\x26\x19\x4a\x88\x54\
                          \x5d\xdc\x49\x7a\x0b\x46\x6e\x7d\x6b\xbd\xb0\x04\x1b\x2f\x58\x6b";
        check_keystream(&key, &nonce, keystream);

        for i in (0..0x20) {
            key[i] = i as u8;
        }
        for i in (0..0x08) {
            nonce[i] = i as u8;
        }
        let keystream = b"\xf7\x98\xa1\x89\xf1\x95\xe6\x69\x82\x10\x5f\xfb\x64\x0b\xb7\x75\
                          \x7f\x57\x9d\xa3\x16\x02\xfc\x93\xec\x01\xac\x56\xf8\x5a\xc3\xc1\
                          \x34\xa4\x54\x7b\x73\x3b\x46\x41\x30\x42\xc9\x44\x00\x49\x17\x69\
                          \x05\xd3\xbe\x59\xea\x1c\x53\xf1\x59\x16\x15\x5c\x2b\xe8\x24\x1a\
                          \x38\x00\x8b\x9a\x26\xbc\x35\x94\x1e\x24\x44\x17\x7c\x8a\xde\x66\
                          \x89\xde\x95\x26\x49\x86\xd9\x58\x89\xfb\x60\xe8\x46\x29\xc9\xbd\
                          \x9a\x5a\xcb\x1c\xc1\x18\xbe\x56\x3e\xb9\xb3\xa4\xa4\x72\xf8\x2e\
                          \x09\xa7\xe7\x78\x49\x2b\x56\x2e\xf7\x13\x0e\x88\xdf\xe0\x31\xc7\
                          \x9d\xb9\xd4\xf7\xc7\xa8\x99\x15\x1b\x9a\x47\x50\x32\xb6\x3f\xc3\
                          \x85\x24\x5f\xe0\x54\xe3\xdd\x5a\x97\xa5\xf5\x76\xfe\x06\x40\x25\
                          \xd3\xce\x04\x2c\x56\x6a\xb2\xc5\x07\xb1\x38\xdb\x85\x3e\x3d\x69\
                          \x59\x66\x09\x96\x54\x6c\xc9\xc4\xa6\xea\xfd\xc7\x77\xc0\x40\xd7\
                          \x0e\xaf\x46\xf7\x6d\xad\x39\x79\xe5\xc5\x36\x0c\x33\x17\x16\x6a\
                          \x1c\x89\x4c\x94\xa3\x71\x87\x6a\x94\xdf\x76\x28\xfe\x4e\xaa\xf2\
                          \xcc\xb2\x7d\x5a\xaa\xe0\xad\x7a\xd0\xf9\xd4\xb6\xad\x3b\x54\x09\
                          \x87\x46\xd4\x52\x4d\x38\x40\x7a\x6d\xeb\x3a\xb7\x8f\xab\x78\xc9";
        check_keystream(&key, &nonce, keystream);
    }
}
<commit_msg>chacha20: Fix for overflow semantics change<commit_after>// http://cr.yp.to/chacha/chacha-20080128.pdf
// http://cr.yp.to/chacha.html

use crypto::wrapping::*;

// convert $e.slice($i, $i + 4) into u32
macro_rules! to_le_u32 {
    ($e:ident[$i:expr]) => ({
        let i: usize = $i;
        let v1 = w8($e[i + 0]).to_w32();
        let v2 = w8($e[i + 1]).to_w32();
        let v3 = w8($e[i + 2]).to_w32();
        let v4 = w8($e[i + 3]).to_w32();
        v1 | (v2 << 8) | (v3 << 16) | (v4 << 24)
    })
}

pub struct ChaCha20 {
    // SECRET
    vals: [w32; 16],
}

impl ChaCha20 {
    // key: SECRET
    pub fn new(key: &[u8], nonce: &[u8]) -> ChaCha20 {
        assert_eq!(key.len(), 32);
        assert_eq!(nonce.len(), 8);

        let mut vals = [w32(0u32); 16];

        // "expand 32-byte k"
        vals[0] = w32(0x61707865);
        vals[1] = w32(0x3320646e);
        vals[2] = w32(0x79622d32);
        vals[3] = w32(0x6b206574);

        for i in (0..8) {
            vals[4 + i] = to_le_u32!(key[4 * i]);
        }

        // counter
        vals[12] = w32(0);
        vals[13] = w32(0);

        vals[14] = to_le_u32!(nonce[0]);
        vals[15] = to_le_u32!(nonce[4]);

        ChaCha20 {
            vals: vals,
        }
    }

    fn round20(&self) -> [w32; 16] {
        // $e must be > 0 and < 32
        macro_rules! rot {
            ($a:expr, $e:expr) => ({
                let a: w32 = $a;
                let e: usize = $e;
                (a << e) | (a >> (32 - e))
            })
        }

        macro_rules! quarter_round {
            ($a:expr, $b:expr, $c:expr, $d:expr) => ({
                $a = $a + $b;
                $d = $d ^ $a;
                $d = rot!($d, 16);

                $c = $c + $d;
                $b = $b ^ $c;
                $b = rot!($b, 12);

                $a = $a + $b;
                $d = $d ^ $a;
                $d = rot!($d, 8);

                $c = $c + $d;
                $b = $b ^ $c;
                $b = rot!($b, 7);
            })
        }

        macro_rules! quarter_round_idx {
            ($e:expr, $a:expr, $b:expr, $c:expr, $d:expr) => (
                quarter_round!($e[$a], $e[$b], $e[$c], $e[$d])
            )
        }

        let mut vals = self.vals;
        for _ in (0..10) {
            // column round
            quarter_round_idx!(vals, 0, 4, 8, 12);
            quarter_round_idx!(vals, 1, 5, 9, 13);
            quarter_round_idx!(vals, 2, 6, 10, 14);
            quarter_round_idx!(vals, 3, 7, 11, 15);

            // diagonal round
            quarter_round_idx!(vals, 0, 5, 10, 15);
            quarter_round_idx!(vals, 1, 6, 11, 12);
            quarter_round_idx!(vals, 2, 7, 8, 13);
            quarter_round_idx!(vals, 3, 4, 9, 14);
        }

        for i in (0..16) {
            vals[i] = vals[i] + self.vals[i];
        }

        vals
    }

    pub fn next(&mut self) -> [u8; 64] {
        let next = self.round20();

        // in TLS, vals[13] never increases
        {
            self.vals[12] = self.vals[12] + w32(1);
            // let mut count = (self.vals[12].to_w64()) | (self.vals[13].to_w64() << 32);
            // count += w64(1);
            // self.vals[12] = count.to_w32();
            // self.vals[13] = (count >> 32).to_w32();
        }

        let next_bytes = {
            let mut next_bytes = [0u8; 64];
            for i in (0..16) {
                next_bytes[4 * i + 0] = next[i].to_w8().0;
                next_bytes[4 * i + 1] = (next[i] >> 8).to_w8().0;
                next_bytes[4 * i + 2] = (next[i] >> 16).to_w8().0;
                next_bytes[4 * i + 3] = (next[i] >> 24).to_w8().0;
            }
            next_bytes
        };

        next_bytes
    }

    // Do not use same nonce for more than 2^70 bytes.
    //
    // if data is 1 byte, it still produces 64 bytes then 63 bytes are just discarded.
    // so this is not suitable for "byte-streaming" mode.
    //
    // data: SECRET
    pub fn encrypt(&mut self, data: &[u8]) -> Vec<u8> {
        let mut ret: Vec<u8> = Vec::new();

        for chunk in data.chunks(64) {
            let next = self.next();
            let xor_iter = next.iter().zip(chunk.iter()).map(|(&x, &y)| x ^ y);
            ret.extend(xor_iter);
        }

        ret
    }
}

#[cfg(test)]
mod test {
    use std::iter::repeat;

    use super::ChaCha20;

    fn check_keystream(key: &[u8], nonce: &[u8], keystream: &[u8]) {
        let mut chacha = ChaCha20::new(key, nonce);
        let input: Vec<_> = repeat(0u8).take(keystream.len()).collect();
        let output = chacha.encrypt(&input);
        assert_eq!(&output[..], keystream);
    }

    #[test]
    fn test_chacha20() {
        // from https://tools.ietf.org/html/draft-agl-tls-chacha20poly1305-04

        let mut key = [0u8; 32];
        let mut nonce = [0u8; 8];
        let keystream = b"\x76\xb8\xe0\xad\xa0\xf1\x3d\x90\x40\x5d\x6a\xe5\x53\x86\xbd\x28\
                          \xbd\xd2\x19\xb8\xa0\x8d\xed\x1a\xa8\x36\xef\xcc\x8b\x77\x0d\xc7\
                          \xda\x41\x59\x7c\x51\x57\x48\x8d\x77\x24\xe0\x3f\xb8\xd8\x4a\x37\
                          \x6a\x43\xb8\xf4\x15\x18\xa1\x1c\xc3\x87\xb6\x69\xb2\xee\x65\x86";
        check_keystream(&key, &nonce, keystream);

        key[31] = 1;
        let keystream = b"\x45\x40\xf0\x5a\x9f\x1f\xb2\x96\xd7\x73\x6e\x7b\x20\x8e\x3c\x96\
                          \xeb\x4f\xe1\x83\x46\x88\xd2\x60\x4f\x45\x09\x52\xed\x43\x2d\x41\
                          \xbb\xe2\xa0\xb6\xea\x75\x66\xd2\xa5\xd1\xe7\xe2\x0d\x42\xaf\x2c\
                          \x53\xd7\x92\xb1\xc4\x3f\xea\x81\x7e\x9a\xd2\x75\xae\x54\x69\x63";
        check_keystream(&key, &nonce, keystream);

        key[31] = 0;
        nonce[7] = 1;
        let keystream = b"\xde\x9c\xba\x7b\xf3\xd6\x9e\xf5\xe7\x86\xdc\x63\x97\x3f\x65\x3a\
                          \x0b\x49\xe0\x15\xad\xbf\xf7\x13\x4f\xcb\x7d\xf1\x37\x82\x10\x31\
                          \xe8\x5a\x05\x02\x78\xa7\x08\x45\x27\x21\x4f\x73\xef\xc7\xfa\x5b\
                          \x52\x77\x06\x2e\xb7\xa0\x43\x3e\x44\x5f\x41\xe3";
        check_keystream(&key, &nonce, keystream);

        key[31] = 0;
        nonce[7] = 0;
        nonce[0] = 1;
        let keystream = b"\xef\x3f\xdf\xd6\xc6\x15\x78\xfb\xf5\xcf\x35\xbd\x3d\xd3\x3b\x80\
                          \x09\x63\x16\x34\xd2\x1e\x42\xac\x33\x96\x0b\xd1\x38\xe5\x0d\x32\
                          \x11\x1e\x4c\xaf\x23\x7e\xe5\x3c\xa8\xad\x64\x26\x19\x4a\x88\x54\
                          \x5d\xdc\x49\x7a\x0b\x46\x6e\x7d\x6b\xbd\xb0\x04\x1b\x2f\x58\x6b";
        check_keystream(&key, &nonce, keystream);

        for i in (0..0x20) {
            key[i] = i as u8;
        }
        for i in (0..0x08) {
            nonce[i] = i as u8;
        }
        let keystream = b"\xf7\x98\xa1\x89\xf1\x95\xe6\x69\x82\x10\x5f\xfb\x64\x0b\xb7\x75\
                          \x7f\x57\x9d\xa3\x16\x02\xfc\x93\xec\x01\xac\x56\xf8\x5a\xc3\xc1\
                          \x34\xa4\x54\x7b\x73\x3b\x46\x41\x30\x42\xc9\x44\x00\x49\x17\x69\
                          \x05\xd3\xbe\x59\xea\x1c\x53\xf1\x59\x16\x15\x5c\x2b\xe8\x24\x1a\
                          \x38\x00\x8b\x9a\x26\xbc\x35\x94\x1e\x24\x44\x17\x7c\x8a\xde\x66\
                          \x89\xde\x95\x26\x49\x86\xd9\x58\x89\xfb\x60\xe8\x46\x29\xc9\xbd\
                          \x9a\x5a\xcb\x1c\xc1\x18\xbe\x56\x3e\xb9\xb3\xa4\xa4\x72\xf8\x2e\
                          \x09\xa7\xe7\x78\x49\x2b\x56\x2e\xf7\x13\x0e\x88\xdf\xe0\x31\xc7\
                          \x9d\xb9\xd4\xf7\xc7\xa8\x99\x15\x1b\x9a\x47\x50\x32\xb6\x3f\xc3\
                          \x85\x24\x5f\xe0\x54\xe3\xdd\x5a\x97\xa5\xf5\x76\xfe\x06\x40\x25\
                          \xd3\xce\x04\x2c\x56\x6a\xb2\xc5\x07\xb1\x38\xdb\x85\x3e\x3d\x69\
                          \x59\x66\x09\x96\x54\x6c\xc9\xc4\xa6\xea\xfd\xc7\x77\xc0\x40\xd7\
                          \x0e\xaf\x46\xf7\x6d\xad\x39\x79\xe5\xc5\x36\x0c\x33\x17\x16\x6a\
                          \x1c\x89\x4c\x94\xa3\x71\x87\x6a\x94\xdf\x76\x28\xfe\x4e\xaa\xf2\
                          \xcc\xb2\x7d\x5a\xaa\xe0\xad\x7a\xd0\xf9\xd4\xb6\xad\x3b\x54\x09\
                          \x87\x46\xd4\x52\x4d\x38\x40\x7a\x6d\xeb\x3a\xb7\x8f\xab\x78\xc9";
        check_keystream(&key, &nonce, keystream);
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::{
    error::LSPError,
    lsp::{
        Connection, DidChangeTextDocument, DidCloseTextDocument, DidOpenTextDocument,
        InitializeParams, LSPBridgeMessage, Message, Notification, Request, ServerCapabilities,
        ServerNotification, ServerRequest, ServerRequestId, TextDocumentSyncCapability,
        TextDocumentSyncKind,
    },
};

use relay_compiler::{errors::Error::DiagnosticsError, FileSource};

use relay_compiler::config::Config;

use crate::error::Result;
use crate::lsp::show_info_message;

use common::ConsoleLogger;
use common::PerfLogger;
use log::info;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::{mpsc, Notify};

use crate::lsp_compiler::LSPCompiler;

use crate::text_documents::initialize_compiler_if_contains_graphql;

/// Initializes an LSP connection, handling the `initize` message and `initialized` notification
/// handshake.
pub fn initialize(connection: &Connection) -> Result<InitializeParams> {
    let mut server_capabilities = ServerCapabilities::default();
    // Enable text document syncing so we can know when files are opened/changed/saved/closed
    server_capabilities.text_document_sync =
        Some(TextDocumentSyncCapability::Kind(TextDocumentSyncKind::Full));

    /* TODO: Re-enable auto-complete
    server_capabilities.completion_provider = Some(CompletionOptions {
        resolve_provider: Some(true),
        trigger_characters: None,
        work_done_progress_options: WorkDoneProgressOptions {
            work_done_progress: None,
        },
    });
    */

    let server_capabilities = serde_json::to_value(&server_capabilities)?;
    let params = connection.initialize(server_capabilities)?;
    let params: InitializeParams = serde_json::from_value(params)?;
    Ok(params)
}

/// Run the main server loop
pub async fn run(connection: Connection, _params: InitializeParams) -> Result<()> {
    show_info_message("Relay Language Server Started", &connection)?;
    info!("Running language server");

    let receiver = connection.receiver.clone();

    // A `Notify` instance used to signal that the compiler should be initialized.
    let compiler_notify = Arc::new(Notify::new());

    // Thread for the LSP message loop
    let compiler_notifier = compiler_notify.clone();

    // A channel to communicate between the LSP message loop and the compiler loop
    let (mut lsp_tx, lsp_rx) = mpsc::channel::<LSPBridgeMessage>(100);

    tokio::spawn(async move {
        // Cache for the extracted GraphQL sources
        for msg in receiver {
            match msg {
                Message::Request(_) => {
                    /* TODO: Re-enable auto-complete
                    // Auto-complete request
                    if req.method == Completion::METHOD {
                        let (request_id, params) = extract_request_params::<Completion>(req);
                        lsp_tx
                            .send(LSPBridgeMessage::CompletionRequest { request_id, params })
                            .await
                            .ok();
                    }
                    */
                }
                Message::Notification(notif) => {
                    match &notif.method {
                        method if method == DidOpenTextDocument::METHOD => {
                            let params = extract_notif_params::<DidOpenTextDocument>(notif);
                            initialize_compiler_if_contains_graphql(
                                &params,
                                compiler_notifier.clone(),
                            );
                            lsp_tx
                                .send(LSPBridgeMessage::DidOpenTextDocument(params))
                                .await
                                .ok();
                        }
                        method if method == DidChangeTextDocument::METHOD => {
                            let params = extract_notif_params::<DidChangeTextDocument>(notif);
                            lsp_tx
                                .send(LSPBridgeMessage::DidChangeTextDocument(params))
                                .await
                                .ok();
                        }
                        method if method == DidCloseTextDocument::METHOD => {
                            let params = extract_notif_params::<DidCloseTextDocument>(notif);
                            lsp_tx
                                .send(LSPBridgeMessage::DidCloseTextDocument(params))
                                .await
                                .ok();
                        }
                        _ => {
                            // Notifications we don't care about
                        }
                    }
                }
                Message::Response(_) => {
                    // Ignore responses for now
                }
            }
        }
    });

    info!("Waiting for compiler to initialize...");
    compiler_notify.notified().await;
    info!("Compiler has initialized");

    let config = load_config();
    let setup_event = ConsoleLogger.create_event("lsp_compiler_setup");
    let file_source = FileSource::connect(&config, &setup_event).await?;
    let (compiler_state, subscription) = file_source
        .subscribe(&setup_event, &ConsoleLogger)
        .await
        .unwrap();
    let schemas = LSPCompiler::build_schemas(&config, &compiler_state, &setup_event)
        .map_err(|errors| LSPError::CompilerError(DiagnosticsError { errors }))?;
    let mut lsp_compiler = LSPCompiler::new(
        &schemas,
        &config,
        subscription,
        compiler_state,
        lsp_rx,
        connection,
    );
    lsp_compiler.watch().await.unwrap();
    Ok(())
}

fn load_config() -> Config {
    // TODO(brandondail) don't hardcode the test project config here
    let home = std::env::var("HOME").unwrap();
    let config_path = PathBuf::from(format!(
        "{}/fbsource/fbcode/relay/config/config.example.json",
        home
    ));
    let config = Config::load(config_path).unwrap();
    config
}

fn extract_notif_params<N>(notif: ServerNotification) -> N::Params
where
    N: Notification,
{
    notif.extract(N::METHOD).unwrap()
}

#[allow(dead_code)]
fn extract_request_params<R>(req: ServerRequest) -> (ServerRequestId, R::Params)
where
    R: Request,
{
    req.extract(R::METHOD).unwrap()
}
<commit_msg>Handle shutdown request<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::{
    error::LSPError,
    lsp::{
        Connection, DidChangeTextDocument, DidCloseTextDocument, DidOpenTextDocument, Exit,
        InitializeParams, LSPBridgeMessage, Message, Notification, Request, ServerCapabilities,
        ServerNotification, ServerRequest, ServerRequestId, ServerResponse, Shutdown,
        TextDocumentSyncCapability, TextDocumentSyncKind,
    },
};

use relay_compiler::{errors::Error::DiagnosticsError, FileSource};

use relay_compiler::config::Config;

use crate::error::Result;
use crate::lsp::show_info_message;

use common::ConsoleLogger;
use common::PerfLogger;
use log::info;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::{mpsc, Notify};

use crate::lsp_compiler::LSPCompiler;

use crate::text_documents::initialize_compiler_if_contains_graphql;

/// Initializes an LSP connection, handling the `initize` message and `initialized` notification
/// handshake.
pub fn initialize(connection: &Connection) -> Result<InitializeParams> {
    let mut server_capabilities = ServerCapabilities::default();
    // Enable text document syncing so we can know when files are opened/changed/saved/closed
    server_capabilities.text_document_sync =
        Some(TextDocumentSyncCapability::Kind(TextDocumentSyncKind::Full));

    /* TODO: Re-enable auto-complete
    server_capabilities.completion_provider = Some(CompletionOptions {
        resolve_provider: Some(true),
        trigger_characters: None,
        work_done_progress_options: WorkDoneProgressOptions {
            work_done_progress: None,
        },
    });
    */

    let server_capabilities = serde_json::to_value(&server_capabilities)?;
    let params = connection.initialize(server_capabilities)?;
    let params: InitializeParams = serde_json::from_value(params)?;
    Ok(params)
}

/// Run the main server loop
pub async fn run(connection: Connection, _params: InitializeParams) -> Result<()> {
    show_info_message("Relay Language Server Started", &connection)?;
    info!("Running language server");

    let receiver = connection.receiver.clone();
    let sender = connection.sender.clone();

    // A `Notify` instance used to signal that the compiler should be initialized.
    let compiler_notify = Arc::new(Notify::new());

    // Thread for the LSP message loop
    let compiler_notifier = compiler_notify.clone();

    // A channel to communicate between the LSP message loop and the compiler loop
    let (mut lsp_tx, lsp_rx) = mpsc::channel::<LSPBridgeMessage>(100);

    let mut has_shutdown = false;
    tokio::spawn(async move {
        // Cache for the extracted GraphQL sources
        for msg in receiver {
            match msg {
                Message::Request(req) => {
                    /* TODO: Re-enable auto-complete
                    // Auto-complete request
                    if req.method == Completion::METHOD {
                        let (request_id, params) = extract_request_params::<Completion>(req);
                        lsp_tx
                            .send(LSPBridgeMessage::CompletionRequest { request_id, params })
                            .await
                            .ok();
                    }
                    */
                    if req.method == Shutdown::METHOD {
                        has_shutdown = true;
                        let (request_id, _) = extract_request_params::<Shutdown>(req);
                        let response = ServerResponse {
                            id: request_id,
                            error: None,
                            result: None,
                        };
                        sender.send(Message::Response(response)).ok();
                        // TODO: We should exit when receiving Exit notification according to the protocal,
                        // but the notification is never received.
                        std::process::exit(0);
                    }
                }
                Message::Notification(notif) => {
                    match &notif.method {
                        method if method == DidOpenTextDocument::METHOD => {
                            let params = extract_notif_params::<DidOpenTextDocument>(notif);
                            initialize_compiler_if_contains_graphql(
                                &params,
                                compiler_notifier.clone(),
                            );
                            lsp_tx
                                .send(LSPBridgeMessage::DidOpenTextDocument(params))
                                .await
                                .ok();
                        }
                        method if method == DidChangeTextDocument::METHOD => {
                            let params = extract_notif_params::<DidChangeTextDocument>(notif);
                            lsp_tx
                                .send(LSPBridgeMessage::DidChangeTextDocument(params))
                                .await
                                .ok();
                        }
                        method if method == DidCloseTextDocument::METHOD => {
                            let params = extract_notif_params::<DidCloseTextDocument>(notif);
                            lsp_tx
                                .send(LSPBridgeMessage::DidCloseTextDocument(params))
                                .await
                                .ok();
                        }
                        method if method == Exit::METHOD => {
                            std::process::exit(!has_shutdown as i32)
                        }
                        _ => {
                            // Notifications we don't care about
                        }
                    }
                }
                Message::Response(_) => {
                    // Ignore responses for now
                }
            }
        }
    });

    info!("Waiting for compiler to initialize...");
    compiler_notify.notified().await;
    info!("Compiler has initialized");

    let config = load_config();
    let setup_event = ConsoleLogger.create_event("lsp_compiler_setup");
    let file_source = FileSource::connect(&config, &setup_event).await?;
    let (compiler_state, subscription) = file_source
        .subscribe(&setup_event, &ConsoleLogger)
        .await
        .unwrap();
    let schemas = LSPCompiler::build_schemas(&config, &compiler_state, &setup_event)
        .map_err(|errors| LSPError::CompilerError(DiagnosticsError { errors }))?;
    let mut lsp_compiler = LSPCompiler::new(
        &schemas,
        &config,
        subscription,
        compiler_state,
        lsp_rx,
        connection,
    );
    lsp_compiler.watch().await.unwrap();
    Ok(())
}

fn load_config() -> Config {
    // TODO(brandondail) don't hardcode the test project config here
    let home = std::env::var("HOME").unwrap();
    let config_path = PathBuf::from(format!(
        "{}/fbsource/fbcode/relay/config/config.example.json",
        home
    ));
    let config = Config::load(config_path).unwrap();
    config
}

fn extract_notif_params<N>(notif: ServerNotification) -> N::Params
where
    N: Notification,
{
    notif.extract(N::METHOD).unwrap()
}

fn extract_request_params<R>(req: ServerRequest) -> (ServerRequestId, R::Params)
where
    R: Request,
{
    req.extract(R::METHOD).unwrap()
}
<|endoftext|>


<commit_before>use bt_messages::Message;
use buffered_reader::BufferedReader;
use chunk::{Position, Piece};
use peer::{Peer, SendPeerMessage};
use std::net::TcpStream;
use std::sync::{Arc, RwLock};
use std::ops::{Deref, DerefMut};
use rand::{Rng, thread_rng};

const BLOCK_LENGTH:usize = 16384; //block length in bytes

pub struct GlobalState {
    gpc: Vec<u16>,
    pub owned: Vec<u8>,
    pub owned_pieces: Vec<Piece>,
    pub request_map: Vec<u8>,
    s_request_map: Vec<u8>,
    piece_length: usize,
    peer_list: Vec<(Arc<RwLock<Peer>>, TcpStream)>
}

impl GlobalState {
    pub fn new (piece_length: usize) -> GlobalState {
        GlobalState {
            gpc: vec![],
            owned: vec![],
            owned_pieces: vec![],
            request_map: vec![],
            s_request_map: vec![],
            peer_list: vec![],
            piece_length: piece_length
        }
    }

    pub fn add_new_peer (&mut self, peer: Arc<RwLock<Peer>>, stream: TcpStream) {
        self.peer_list.push((peer, stream));
    }

    /// Increases the value of gpc[piece_index] by n
    #[inline]
    pub fn gpc_incr (&mut self, piece_index: usize, n: u16) {
        //starting to regret making the bitfield variable in size... maybe i can preallocate. will come back and re-eval
        let len = self.gpc.len();
        if piece_index >= len {
            self.gpc.extend((0..piece_index+1 - len).map(|_| 0));
        }
        self.gpc[piece_index] += n;
    }

    /// Returns the index of the rarest piece that isn't owned or currently being requested.
    /// TODO: Approximations may yield optimized results
    #[inline]
    pub fn rarest (&self) -> Option<usize> {
        //there's probably a faster way. doing this naively for the sake of forward progress
        let mut most_rare = (None, u16::max_value());
        for (index, byte) in self.unclaimed_fields().iter().enumerate() {
            for i in 0..8 { //cast up so i don't have to deal with overflows
                let n = 1 & (((*byte as u16) << i) >> (8-i));
                if n == 1 {
                    let true_index = index*8+1;
                    let population = self.gpc[true_index];
                    let (_, mr_pop) = most_rare;
                    if population < mr_pop {
                        most_rare = (Some(true_index), population);
                    }
                }
            }
        }
        let (index, _) = most_rare;
        index
    }

    /// Returns the index of the rarest piece that is owned by the peer and isn't both owned
    /// and currently being requested
    /// TODO: Approximations may yield optimized results
    #[inline]
    pub fn rarest_wrt_peer (&self, peer_bitfield: &Vec<u8>) -> Option<usize> {
        //there's probably a faster way. doing this naively for the sake of forward progress
        let mut most_rare = (None, u16::max_value());
        let eligible = and_slice_vbr_len(&self.unclaimed_fields(), &peer_bitfield);
        for (index, byte) in eligible.iter().enumerate() {
            for i in 0..8 { //cast up so i don't have to deal with overflows
                let n = 1 & (((*byte as u16) << i) >> (8-i));
                //println!("n: {}", n);
                if n == 1 {
                    let true_index = index*8+1;
                    let population = self.gpc[true_index];
                    let (_, mr_pop) = most_rare;
                    if population < mr_pop {
                        most_rare = (Some(true_index), population)
                    }
                }
            }
        }
        let (index, _) = most_rare;
        index
    }

    /// returns a complete bitfield of pieces that aren't owned or being requested
    /// this is done almost as strictly as possible - and might be a little of a waste as it isn't
    /// really necessary to get a complete picture to get a request chunk
    /// additionally the definitions of owned and request_map are not strict yet - currently they
    /// are growable, and impelementers should take note of that
    #[inline]
    pub fn unclaimed_fields (&self) -> Vec<u8> {
        nand_slice_vbr_len(&self.owned, &self.request_map)
    }

    #[inline]
    pub fn req (&mut self, peer_bitfield: &Vec<u8>) {
        let index = self.rarest_wrt_peer(peer_bitfield);
        println!("REQ: {:?}", index);
    }
}

pub trait Spin {
    fn spin (&mut self);
}

impl Spin for GlobalState {
    fn spin (&mut self) {

        //NOTE: this shuffles the peer_list
        thread_rng().shuffle(&mut self.peer_list);
        for tup in self.peer_list.iter() {
            let (ref rw_lock_peer, _) = *tup;

            {
                let peer = match rw_lock_peer.try_read() {
                    Ok(a) => a,
                    Err(_) => continue//do nothing. it's locked
                };
                //TODO: owned_pieces is not sufficient. it should be the union of owned_pieces and
                //requests
                let want = Piece::complement(&peer.deref().state.pieces, &self.owned_pieces);
                println!("want: {:?}", want);
            }
        }
    }
}

pub struct DefaultHandler;

/// Handles messages. This is a cheap way to force reactive style
pub trait Handler {
    type MessageType;
    fn handle(&mut self, message: Self::MessageType, peer: &mut Peer, global_state: &mut GlobalState);
}

/// The default algorithm
impl Handler for DefaultHandler {
    type MessageType = Message;
    #[inline]
    fn handle (&mut self, message: Message, peer: &mut Peer, global: &mut GlobalState) {
        println!("{:?}", message);
        match message {
            Message::Have{piece_index: index} => {
                let i = index as usize;
                global.gpc_incr(i, 1);
                //peer.state.set_have(i);
                let piece = Piece::from(BLOCK_LENGTH, i, 0, BLOCK_LENGTH);
                let i_index = Piece::add_to_boundary_vec(&mut global.owned_pieces, piece);
                Piece::compact_if_possible(&mut global.owned_pieces, i_index);
            },
            Message::Choke => {
                peer.state.set_us_choked(true);
            },
            Message::Unchoke => {
                peer.state.set_us_choked(false);
            },
            Message::Interested => {
                peer.state.set_us_interested(true);
            },
            Message::Bitfield(bitfield) => {
                for (index, byte) in bitfield.iter().enumerate() {
                    for i in 0..8 { //cast up so i don't have to deal with overflows
                        let n = 1 & (((*byte as u16) << i) >> (8-i));
                        global.gpc_incr(index*8+i, n);
                    }
                }

                peer.state.set_pieces_from_bitfield(&bitfield);
                //peer.state.set_bitfield(bitfield);
            },
            _ => {
            }
        };
    }
}

/// Zip-maps a generic func over two byte slices with variable lengths
#[inline]
pub fn apply_bitwise_slice_vbr_len <F, T:Clone> (lhs: &[T], rhs: &[T], default: T, func: F) -> Vec<T>
    where F: Fn((&T, &T)) -> T {
    if lhs.len() == rhs.len() {
        bitwise_byte_slice(lhs, rhs, func)
    } else {
        if lhs.len() < rhs.len() {
            let mut a = lhs.to_owned();
            a.extend((0..rhs.len()-lhs.len()).map(|_| default.clone()));
            bitwise_byte_slice(&a, rhs, func)
        } else {
            let mut b = rhs.to_owned();
            b.extend((0..lhs.len()-rhs.len()).map(|_| default.clone()));
            bitwise_byte_slice(lhs, &b, func)
        }
    }
}

/// Zip-maps a generic func (intended bitwise) over two byte slices
#[inline]
pub fn bitwise_byte_slice <F, T> (lhs: &[T], rhs: &[T], func: F) -> Vec<T>
    where F: Fn((&T, &T)) -> T {
    assert!(lhs.len() == rhs.len());
    lhs.iter().zip(rhs)
              .map(func)
              .collect::<Vec<T>>()
}

#[inline]
pub fn nand_slice_vbr_len (lhs: &[u8], rhs: &[u8]) -> Vec<u8> {
    apply_bitwise_slice_vbr_len(lhs, rhs, 0, |(a, b)| !a & !b)
}

#[inline]
pub fn and_slice_vbr_len(lhs: &[u8], rhs: &[u8]) -> Vec<u8> {
    apply_bitwise_slice_vbr_len(lhs, rhs, 255, |(a, b)| a & b)
}
<commit_msg>adding exclude<commit_after>use bt_messages::Message;
use buffered_reader::BufferedReader;
use chunk::{Position, Piece};
use peer::{Peer, SendPeerMessage};
use std::net::TcpStream;
use std::sync::{Arc, RwLock};
use std::ops::{Deref, DerefMut};
use rand::{Rng, thread_rng};

const BLOCK_LENGTH:usize = 16384; //block length in bytes

pub struct GlobalState {
    gpc: Vec<u16>,
    pub owned: Vec<u8>,
    pub owned_pieces: Vec<Piece>,
    pub request_map: Vec<u8>,
    pub exclude: Vec<Piece>,
    s_request_map: Vec<u8>,
    piece_length: usize,
    peer_list: Vec<(Arc<RwLock<Peer>>, TcpStream)>
}

impl GlobalState {
    pub fn new (piece_length: usize) -> GlobalState {
        GlobalState {
            gpc: vec![],
            owned: vec![],
            owned_pieces: vec![],
            exclude: vec![],
            request_map: vec![],
            s_request_map: vec![],
            peer_list: vec![],
            piece_length: piece_length
        }
    }

    pub fn add_new_peer (&mut self, peer: Arc<RwLock<Peer>>, stream: TcpStream) {
        self.peer_list.push((peer, stream));
    }

    /// Increases the value of gpc[piece_index] by n
    #[inline]
    pub fn gpc_incr (&mut self, piece_index: usize, n: u16) {
        //starting to regret making the bitfield variable in size... maybe i can preallocate. will come back and re-eval
        let len = self.gpc.len();
        if piece_index >= len {
            self.gpc.extend((0..piece_index+1 - len).map(|_| 0));
        }
        self.gpc[piece_index] += n;
    }

    /// Returns the index of the rarest piece that isn't owned or currently being requested.
    /// TODO: Approximations may yield optimized results
    #[inline]
    pub fn rarest (&self) -> Option<usize> {
        //there's probably a faster way. doing this naively for the sake of forward progress
        let mut most_rare = (None, u16::max_value());
        for (index, byte) in self.unclaimed_fields().iter().enumerate() {
            for i in 0..8 { //cast up so i don't have to deal with overflows
                let n = 1 & (((*byte as u16) << i) >> (8-i));
                if n == 1 {
                    let true_index = index*8+1;
                    let population = self.gpc[true_index];
                    let (_, mr_pop) = most_rare;
                    if population < mr_pop {
                        most_rare = (Some(true_index), population);
                    }
                }
            }
        }
        let (index, _) = most_rare;
        index
    }

    /// Returns the index of the rarest piece that is owned by the peer and isn't both owned
    /// and currently being requested
    /// TODO: Approximations may yield optimized results
    #[inline]
    pub fn rarest_wrt_peer (&self, peer_bitfield: &Vec<u8>) -> Option<usize> {
        //there's probably a faster way. doing this naively for the sake of forward progress
        let mut most_rare = (None, u16::max_value());
        let eligible = and_slice_vbr_len(&self.unclaimed_fields(), &peer_bitfield);
        for (index, byte) in eligible.iter().enumerate() {
            for i in 0..8 { //cast up so i don't have to deal with overflows
                let n = 1 & (((*byte as u16) << i) >> (8-i));
                //println!("n: {}", n);
                if n == 1 {
                    let true_index = index*8+1;
                    let population = self.gpc[true_index];
                    let (_, mr_pop) = most_rare;
                    if population < mr_pop {
                        most_rare = (Some(true_index), population)
                    }
                }
            }
        }
        let (index, _) = most_rare;
        index
    }

    /// returns a complete bitfield of pieces that aren't owned or being requested
    /// this is done almost as strictly as possible - and might be a little of a waste as it isn't
    /// really necessary to get a complete picture to get a request chunk
    /// additionally the definitions of owned and request_map are not strict yet - currently they
    /// are growable, and impelementers should take note of that
    #[inline]
    pub fn unclaimed_fields (&self) -> Vec<u8> {
        nand_slice_vbr_len(&self.owned, &self.request_map)
    }

    #[inline]
    pub fn req (&mut self, peer_bitfield: &Vec<u8>) {
        let index = self.rarest_wrt_peer(peer_bitfield);
        println!("REQ: {:?}", index);
    }
}

pub trait Spin {
    fn spin (&mut self);
}

impl Spin for GlobalState {
    fn spin (&mut self) {

        //NOTE: this shuffles the peer_list
        thread_rng().shuffle(&mut self.peer_list);
        for tup in self.peer_list.iter() {
            let (ref rw_lock_peer, _) = *tup;
            {
                let peer = match rw_lock_peer.try_read() {
                    Ok(a) => a,
                    Err(_) => continue//do nothing. it's locked
                };
                //TODO: owned_pieces is not sufficient. it should be the union of owned_pieces and
                //requests
                let want = Piece::complement(&peer.deref().state.pieces, &self.exclude);
                let req_piece = match want.len() {
                    0 => continue,
                    _ => select_piece(&want, BLOCK_LENGTH)
                };

                println!("{:?}", req_piece);

                let index = Piece::add_to_boundary_vec(&mut self.exclude, req_piece);
                Piece::compact_if_possible(&mut self.exclude, index);
            }
        }
    }
}

fn select_piece (pieces: &[Piece], piece_length: usize) -> Piece {
    pieces.first().unwrap().to_owned()
}

pub struct DefaultHandler;

/// Handles messages. This is a cheap way to force reactive style
pub trait Handler {
    type MessageType;
    fn handle(&mut self, message: Self::MessageType, peer: &mut Peer, global_state: &mut GlobalState);
}

/// The default algorithm
impl Handler for DefaultHandler {
    type MessageType = Message;
    #[inline]
    fn handle (&mut self, message: Message, peer: &mut Peer, global: &mut GlobalState) {
        println!("{:?}", message);
        match message {
            Message::Have{piece_index: index} => {
                let i = index as usize;
                global.gpc_incr(i, 1);
                //peer.state.set_have(i);
                let piece = Piece::from(BLOCK_LENGTH, i, 0, BLOCK_LENGTH);
                let i_index = Piece::add_to_boundary_vec(&mut global.owned_pieces, piece);
                Piece::compact_if_possible(&mut global.owned_pieces, i_index);
            },
            Message::Choke => {
                peer.state.set_us_choked(true);
            },
            Message::Unchoke => {
                peer.state.set_us_choked(false);
            },
            Message::Interested => {
                peer.state.set_us_interested(true);
            },
            Message::Bitfield(bitfield) => {
                for (index, byte) in bitfield.iter().enumerate() {
                    for i in 0..8 { //cast up so i don't have to deal with overflows
                        let n = 1 & (((*byte as u16) << i) >> (8-i));
                        global.gpc_incr(index*8+i, n);
                    }
                }

                peer.state.set_pieces_from_bitfield(&bitfield);
                //peer.state.set_bitfield(bitfield);
            },
            _ => {
            }
        };
    }
}

/// Zip-maps a generic func over two byte slices with variable lengths
#[inline]
pub fn apply_bitwise_slice_vbr_len <F, T:Clone> (lhs: &[T], rhs: &[T], default: T, func: F) -> Vec<T>
    where F: Fn((&T, &T)) -> T {
    if lhs.len() == rhs.len() {
        bitwise_byte_slice(lhs, rhs, func)
    } else {
        if lhs.len() < rhs.len() {
            let mut a = lhs.to_owned();
            a.extend((0..rhs.len()-lhs.len()).map(|_| default.clone()));
            bitwise_byte_slice(&a, rhs, func)
        } else {
            let mut b = rhs.to_owned();
            b.extend((0..lhs.len()-rhs.len()).map(|_| default.clone()));
            bitwise_byte_slice(lhs, &b, func)
        }
    }
}

/// Zip-maps a generic func (intended bitwise) over two byte slices
#[inline]
pub fn bitwise_byte_slice <F, T> (lhs: &[T], rhs: &[T], func: F) -> Vec<T>
    where F: Fn((&T, &T)) -> T {
    assert!(lhs.len() == rhs.len());
    lhs.iter().zip(rhs)
              .map(func)
              .collect::<Vec<T>>()
}

#[inline]
pub fn nand_slice_vbr_len (lhs: &[u8], rhs: &[u8]) -> Vec<u8> {
    apply_bitwise_slice_vbr_len(lhs, rhs, 0, |(a, b)| !a & !b)
}

#[inline]
pub fn and_slice_vbr_len(lhs: &[u8], rhs: &[u8]) -> Vec<u8> {
    apply_bitwise_slice_vbr_len(lhs, rhs, 255, |(a, b)| a & b)
}
<|endoftext|>


<commit_before>use seed::prelude::*;
use crate::Msg;

pub mod perfects_utils {
	use num_bigint::{BigUint,ToBigUint};
	use num_traits::{ToPrimitive};
    use seed::prelude::*;
    use crate::Msg;

	#[derive(Clone)]
	pub struct Perfect {
		pub n: u64,
		pub p: u64,
		pub digits: u64,
		pub discovery: String,
	}

	#[derive(Clone)]
	pub struct PerfectDownload {
		pub n: u64,
		pub p: u64,
	}

	pub fn perfects() -> Vec<Perfect> {
        vec![
			/*Perfect {n: 51, p: 82589934, digits: 49724095, discovery: "2018 Laroche, Woltman, Kurowski, Blosser, et al."},
			Perfect {n: 50, p: 77232917, digits: 46498850, discovery: "2017 Jonathan Pace, George Woltman, Scott Kurowski, Aaron Blosser, et al.." },
			Perfect {n: 49, p: 74207281, digits: 44677235, discovery: "2016 Cooper, Woltman, Kurowski, Blosser et al." },
			Perfect {n: 48, p: 57885161, digits: 34850340, discovery: "2013 Cooper, Woltman, Kurowski, et al." },
			Perfect {n: 47, p: 43112609, digits: 25956377, discovery: "2008 Smith, Woltman, Kurowski, et al." },
			Perfect {n: 46, p: 42643801, digits: 42643801, discovery: "2009 Strindmo, Woltman, Kurowski, et al." },
			Perfect {n: 45, p: 37156667, digits: 22370543, discovery: "2008 Elvenich, Woltman, Kurowski, et al." },
			Perfect {n: 44, p: 32582657, digits: 19616714, discovery: "2006 Cooper, Boone, Woltman, Kurowski, et al." },
			Perfect {n: 43, p: 30402457, digits: 18304103, discovery: "2005 Cooper, Boone, Woltman, Kurowski, et al." },
			Perfect {n: 42, p: 25964951, digits: 15632458, discovery: "2005 Nowak, Woltman, Kurowski, et. al." },
			Perfect {n: 41, p: 24036583, digits: 14471465, discovery: "2004 Findley, Woltman, Kurowski, et. al." },
			Perfect {n: 40, p: 20996011, digits: 12640858, discovery: "2003 Shafer, Woltman, Kurowski, et. al." },
			Perfect {n: 39, p: 13466917, digits: 8107892, discovery: "2001 Cameron, Woltman, Kurowski, et. al." },
			Perfect {n: 38, p: 6972593, digits: 4197919, discovery: "1999 Hajratwala, Woltman, Kurowski, et. al." },
			Perfect {n: 37, p: 3021377, digits: 1819050, discovery: "1998 Clarkson, Woltman, Kurowski, et. al." },
			Perfect {n: 36, p: 2976221, digits: 1791864, discovery: "1997 Spence, Woltman, et. al." },
			Perfect {n: 35, p: 1398269, digits: 841842, discovery: "1996 Armengaud, Woltman, et. al." },
			Perfect {n: 34, p: 1257787, digits: 757263, discovery: "1996 Slowinski&Gage" },
			Perfect {n: 33, p: 859433, digits: 517430, discovery: "1994 Slowinski&Gage" },
			Perfect {n: 32, p: 756839, digits: 455663, discovery: "1992 Slowinski&Gage" },
			Perfect {n: 31, p: 216091, digits: 130100, discovery: "1985 Slowinski" },
			Perfect {n: 30, p: 132049, digits: 79502, discovery: "1983 Slowinski" },
			Perfect {n: 29, p: 110503, digits: 66530, discovery: "1988 Colquitt&Welsh" },
			Perfect {n: 28, p: 86243, digits: 51924, discovery: "1982 Slowinski" },*/
			Perfect {n: 27, p: 44497, digits: 26790, discovery: String::from("1979 Nelson&Slowinski") },
			Perfect {n: 26, p: 23209, digits: 13973, discovery: String::from("1979 Noll") },
			Perfect {n: 25, p: 21701, digits: 13066, discovery: String::from("1978 Noll&Nickel") },
			Perfect {n: 24, p: 19937, digits: 12003, discovery: String::from("1971 Tuckerman") },
			Perfect {n: 23, p: 11213, digits: 6751, discovery: String::from("1963 Gillies") },
			Perfect {n: 22, p: 9941, digits: 5985, discovery: String::from("1963 Gillies") },
			Perfect {n: 21, p: 9689, digits: 5834, discovery: String::from("1963 Gillies") },
			Perfect {n: 20, p: 4423, digits: 2663, discovery: String::from("1961 Hurwitz") },
			Perfect {n: 19, p: 4253, digits: 2561, discovery: String::from("1961 Hurwitz") },
			Perfect {n: 18, p: 3217, digits: 1937, discovery: String::from("1957 Riesel") },
			Perfect {n: 17, p: 2281, digits: 1373, discovery: String::from("1952 Robinson") },
			Perfect {n: 16, p: 2203, digits: 1327, discovery: String::from("1952 Robinson") },
			Perfect {n: 14, p: 607, digits: 366, discovery: String::from("1952 Robinson") },
			Perfect {n: 13, p: 521, digits: 314, discovery: String::from("1952 Robinson") },
			Perfect {n: 12, p: 127, digits: 77, discovery: String::from("1876 Lucas") },
			Perfect {n: 11, p: 107, digits: 65, discovery: String::from("1914 Powers") },
			Perfect {n: 10, p: 89, digits: 54, discovery: String::from("1911 Powers") },
			Perfect {n: 9, p: 61, digits: 37, discovery: String::from("1883 Pervushin") },
			Perfect {n: 8, p: 31, digits: 19, discovery: String::from("1772 Euler") },
			Perfect {n: 7, p: 19, digits: 12, discovery: String::from("1588 Cataldi") },
			Perfect {n: 6, p: 17, digits: 10, discovery: String::from("1588 Cataldi") },
			Perfect {n: 5, p: 13, digits: 8, discovery: String::from("1456 ?") },
			Perfect {n: 4, p: 7, digits: 4, discovery: String::from("?") },
			Perfect {n: 3, p: 5, digits: 3, discovery: String::from("?") },
			Perfect {n: 2, p: 3, digits: 2, discovery: String::from("?") },
			Perfect {n: 1, p: 2, digits: 1, discovery: String::from("?") },
		]
	}

	pub fn save_as_file(filename:String, filecontent:String) -> seed::dom_types::Node<Msg> {
        let href:String = vec!["data:text/plain,",&filecontent].into_iter().collect();
        a![attrs!{At::Download => &filename, At::Href => &href}, "TXT"]
    }

	pub fn generate_file(n:u64, p:u64) -> seed::dom_types::Node<Msg> {
		let download_filename:String = format!("P{}.txt",&n.to_string());
		let two:BigUint = 2.to_biguint().unwrap();
		let power:BigUint = two.clone() << (p.to_usize().unwrap()-1-1);
        let perfect_value:BigUint = power.clone() * ((power.clone()*two.clone()) -1.to_biguint().unwrap());
        save_as_file(download_filename, perfect_value.to_string())
    }
}

pub fn render(model: &crate::Model) -> seed::dom_types::Node<Msg> {
    let mut html = vec![];
    let perfects = perfects_utils::perfects().to_owned();

    for perfect in perfects.iter() {
		let perfectp:u64 = perfect.p;
		let perfect_download = perfects_utils::PerfectDownload {n: perfect.n, p: perfect.p};
		let equation:String = format!("2{} × (2{}-1)",&(perfectp-1).to_string(),&(perfectp).to_string());

        html.push(
            tr![
                td![perfect.n.to_string()],//rank
				td![El::from_html(&equation)],//perfect number as a formula

                td![perfect.digits.to_string()],//digits in length
                td![perfect.discovery],//discovery
				if model.download.n == perfect.n {
					perfects_utils::generate_file(model.download.n,model.download.p)
				} else {
					td![button!["Generate",mouse_ev("mouseup", move |event| Msg::GenerateDownload(event, perfect_download))]]
				},
            ]
        );
    }

    div![
        h1!["The Perfect Numbers"],
        br![],
        br![],
        br![],
        table![
            attrs!{At::Class => "perfecttable text"},
            tbody![
                tr![
                    td![
                        b!["No."]
                    ],
                    td![
                        b!["Prime"]
                    ],
                    td![
                        b!["Digits"]
                    ],
                    td![
                        b!["Discovered"]
                    ],
                    td![
                        b!["Download"]
                    ]
                ],
                html
            ]
        ]
    ]
}

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

    #[test]
    fn perfect_test<'a>() {
        let mut perfects:Vec<perfects_utils::Perfect> = perfects_utils::perfects();
		perfects.reverse();
        assert_eq!(perfects[49].digits, 46498850);
        assert_eq!(perfects[39].digits, 12640858);
    }
}<commit_msg>Now displays all 51 perfects<commit_after>use seed::prelude::*;
use crate::Msg;

pub mod perfects_utils {
	use num_bigint::{BigUint,ToBigUint};
	use num_traits::{ToPrimitive};
    use seed::prelude::*;
    use crate::Msg;

	#[derive(Clone)]
	pub struct Perfect {
		pub n: u64,
		pub p: u64,
		pub digits: u64,
		pub discovery: String,
	}

	#[derive(Clone)]
	pub struct PerfectDownload {
		pub n: u64,
		pub p: u64,
	}

	pub fn perfects() -> Vec<Perfect> {
        vec![
			Perfect {n: 51, p: 82589934, digits: 49724095, discovery: String::from("2018 Laroche, Woltman, Kurowski, Blosser, et al.")},
			Perfect {n: 50, p: 77232917, digits: 46498850, discovery: String::from("2017 Jonathan Pace, George Woltman, Scott Kurowski, Aaron Blosser, et al..") },
			Perfect {n: 49, p: 74207281, digits: 44677235, discovery: String::from("2016 Cooper, Woltman, Kurowski, Blosser et al.") },
			Perfect {n: 48, p: 57885161, digits: 34850340, discovery: String::from("2013 Cooper, Woltman, Kurowski, et al.") },
			Perfect {n: 47, p: 43112609, digits: 25956377, discovery: String::from("2008 Smith, Woltman, Kurowski, et al.") },
			Perfect {n: 46, p: 42643801, digits: 42643801, discovery: String::from("2009 Strindmo, Woltman, Kurowski, et al.") },
			Perfect {n: 45, p: 37156667, digits: 22370543, discovery: String::from("2008 Elvenich, Woltman, Kurowski, et al.") },
			Perfect {n: 44, p: 32582657, digits: 19616714, discovery: String::from("2006 Cooper, Boone, Woltman, Kurowski, et al.") },
			Perfect {n: 43, p: 30402457, digits: 18304103, discovery: String::from("2005 Cooper, Boone, Woltman, Kurowski, et al.") },
			Perfect {n: 42, p: 25964951, digits: 15632458, discovery: String::from("2005 Nowak, Woltman, Kurowski, et. al.") },
			Perfect {n: 41, p: 24036583, digits: 14471465, discovery: String::from("2004 Findley, Woltman, Kurowski, et. al.") },
			Perfect {n: 40, p: 20996011, digits: 12640858, discovery: String::from("2003 Shafer, Woltman, Kurowski, et. al.") },
			Perfect {n: 39, p: 13466917, digits: 8107892, discovery: String::from("2001 Cameron, Woltman, Kurowski, et. al.") },
			Perfect {n: 38, p: 6972593, digits: 4197919, discovery: String::from("1999 Hajratwala, Woltman, Kurowski, et. al.") },
			Perfect {n: 37, p: 3021377, digits: 1819050, discovery: String::from("1998 Clarkson, Woltman, Kurowski, et. al.") },
			Perfect {n: 36, p: 2976221, digits: 1791864, discovery: String::from("1997 Spence, Woltman, et. al.") },
			Perfect {n: 35, p: 1398269, digits: 841842, discovery: String::from("1996 Armengaud, Woltman, et. al.") },
			Perfect {n: 34, p: 1257787, digits: 757263, discovery: String::from("1996 Slowinski&Gage") },
			Perfect {n: 33, p: 859433, digits: 517430, discovery: String::from("1994 Slowinski&Gage") },
			Perfect {n: 32, p: 756839, digits: 455663, discovery: String::from("1992 Slowinski&Gage") },
			Perfect {n: 31, p: 216091, digits: 130100, discovery: String::from("1985 Slowinski") },
			Perfect {n: 30, p: 132049, digits: 79502, discovery: String::from("1983 Slowinski") },
			Perfect {n: 29, p: 110503, digits: 66530, discovery: String::from("1988 Colquitt&Welsh") },
			Perfect {n: 28, p: 86243, digits: 51924, discovery: String::from("1982 Slowinski") },
			Perfect {n: 27, p: 44497, digits: 26790, discovery: String::from("1979 Nelson&Slowinski") },
			Perfect {n: 26, p: 23209, digits: 13973, discovery: String::from("1979 Noll") },
			Perfect {n: 25, p: 21701, digits: 13066, discovery: String::from("1978 Noll&Nickel") },
			Perfect {n: 24, p: 19937, digits: 12003, discovery: String::from("1971 Tuckerman") },
			Perfect {n: 23, p: 11213, digits: 6751, discovery: String::from("1963 Gillies") },
			Perfect {n: 22, p: 9941, digits: 5985, discovery: String::from("1963 Gillies") },
			Perfect {n: 21, p: 9689, digits: 5834, discovery: String::from("1963 Gillies") },
			Perfect {n: 20, p: 4423, digits: 2663, discovery: String::from("1961 Hurwitz") },
			Perfect {n: 19, p: 4253, digits: 2561, discovery: String::from("1961 Hurwitz") },
			Perfect {n: 18, p: 3217, digits: 1937, discovery: String::from("1957 Riesel") },
			Perfect {n: 17, p: 2281, digits: 1373, discovery: String::from("1952 Robinson") },
			Perfect {n: 16, p: 2203, digits: 1327, discovery: String::from("1952 Robinson") },
			Perfect {n: 14, p: 607, digits: 366, discovery: String::from("1952 Robinson") },
			Perfect {n: 13, p: 521, digits: 314, discovery: String::from("1952 Robinson") },
			Perfect {n: 12, p: 127, digits: 77, discovery: String::from("1876 Lucas") },
			Perfect {n: 11, p: 107, digits: 65, discovery: String::from("1914 Powers") },
			Perfect {n: 10, p: 89, digits: 54, discovery: String::from("1911 Powers") },
			Perfect {n: 9, p: 61, digits: 37, discovery: String::from("1883 Pervushin") },
			Perfect {n: 8, p: 31, digits: 19, discovery: String::from("1772 Euler") },
			Perfect {n: 7, p: 19, digits: 12, discovery: String::from("1588 Cataldi") },
			Perfect {n: 6, p: 17, digits: 10, discovery: String::from("1588 Cataldi") },
			Perfect {n: 5, p: 13, digits: 8, discovery: String::from("1456 ?") },
			Perfect {n: 4, p: 7, digits: 4, discovery: String::from("?") },
			Perfect {n: 3, p: 5, digits: 3, discovery: String::from("?") },
			Perfect {n: 2, p: 3, digits: 2, discovery: String::from("?") },
			Perfect {n: 1, p: 2, digits: 1, discovery: String::from("?") },
		]
	}

	pub fn save_as_file(filename:String, filecontent:String) -> seed::dom_types::Node<Msg> {
        let href:String = vec!["data:text/plain,",&filecontent].into_iter().collect();
        a![attrs!{At::Download => &filename, At::Href => &href}, "TXT"]
    }

	pub fn generate_file(n:u64, p:u64) -> seed::dom_types::Node<Msg> {
		let download_filename:String = format!("P{}.txt",&n.to_string());
		let two:BigUint = 2.to_biguint().unwrap();
		let power:BigUint = two.clone() << (p.to_usize().unwrap()-1-1);
        let perfect_value:BigUint = power.clone() * ((power.clone()*two.clone()) -1.to_biguint().unwrap());
        save_as_file(download_filename, perfect_value.to_string())
    }
}

pub fn render(model: &crate::Model) -> seed::dom_types::Node<Msg> {
    let mut html = vec![];
    let perfects = perfects_utils::perfects().to_owned();

    for perfect in perfects.iter() {
		let perfectp:u64 = perfect.p;
		let perfect_download = perfects_utils::PerfectDownload {n: perfect.n, p: perfect.p};
		let equation:String = format!("2{} × (2{}-1)",&(perfectp-1).to_string(),&(perfectp).to_string());

        html.push(
            tr![
                td![perfect.n.to_string()],//rank
				td![El::from_html(&equation)],//perfect number as a formula

                td![perfect.digits.to_string()],//digits in length
                td![perfect.discovery],//discovery
				if model.download.n == perfect.n {
					perfects_utils::generate_file(model.download.n,model.download.p)
				} else {
					td![button!["Generate",mouse_ev("mouseup", move |event| Msg::GenerateDownload(event, perfect_download))]]
				},
            ]
        );
    }

    div![
        h1!["The Perfect Numbers"],
        br![],
        br![],
        br![],
        table![
            attrs!{At::Class => "perfecttable text"},
            tbody![
                tr![
                    td![
                        b!["No."]
                    ],
                    td![
                        b!["Prime"]
                    ],
                    td![
                        b!["Digits"]
                    ],
                    td![
                        b!["Discovered"]
                    ],
                    td![
                        b!["Download"]
                    ]
                ],
                html
            ]
        ]
    ]
}

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

    #[test]
    fn perfect_test<'a>() {
        let mut perfects:Vec<perfects_utils::Perfect> = perfects_utils::perfects();
		perfects.reverse();
        assert_eq!(perfects[49].digits, 46498850);
        assert_eq!(perfects[39].digits, 12640858);
    }
}<|endoftext|>