<commit_before>//! Client Responses
use std::io::{self, Read};
use std::marker::PhantomData;
use std::net::Shutdown;

use buffer::BufReader;
use header;
use header::{ContentLength, TransferEncoding};
use header::Encoding::Chunked;
use net::{NetworkStream, HttpStream};
use http::{self, HttpReader, RawStatus};
use http::HttpReader::{SizedReader, ChunkedReader, EofReader};
use status;
use version;

/// A response for a client request to a remote server.
#[derive(Debug)]
pub struct Response<S = HttpStream> {
    /// The status from the server.
    pub status: status::StatusCode,
    /// The headers from the server.
    pub headers: header::Headers,
    /// The HTTP version of this response from the server.
    pub version: version::HttpVersion,
    status_raw: RawStatus,
    body: HttpReader<BufReader<Box<NetworkStream + Send>>>,

    _marker: PhantomData<S>,
}

impl Response {

    /// Creates a new response from a server.
    pub fn new(stream: Box<NetworkStream + Send>) -> ::Result<Response> {
        trace!("Response::new");
        let mut stream = BufReader::new(stream);

        let head = try!(http::parse_response(&mut stream));
        let raw_status = head.subject;
        let headers = head.headers;

        let status = status::StatusCode::from_u16(raw_status.0);
        debug!("version={:?}, status={:?}", head.version, status);
        debug!("headers={:?}", headers);

        if !http::should_keep_alive(head.version, &headers) {
            try!(stream.get_mut().close(Shutdown::Write));
        }

        let body = if headers.has::<TransferEncoding>() {
            match headers.get::<TransferEncoding>() {
                Some(&TransferEncoding(ref codings)) => {
                    if codings.len() > 1 {
                        trace!("TODO: #2 handle other codings: {:?}", codings);
                    };

                    if codings.contains(&Chunked) {
                        ChunkedReader(stream, None)
                    } else {
                        trace!("not chuncked. read till eof");
                        EofReader(stream)
                    }
                }
                None => unreachable!()
            }
        } else if headers.has::<ContentLength>() {
            match headers.get::<ContentLength>() {
                Some(&ContentLength(len)) => SizedReader(stream, len),
                None => unreachable!()
            }
        } else {
            trace!("neither Transfer-Encoding nor Content-Length");
            EofReader(stream)
        };

        Ok(Response {
            status: status,
            version: head.version,
            headers: headers,
            body: body,
            status_raw: raw_status,
            _marker: PhantomData,
        })
    }

    /// Get the raw status code and reason.
    pub fn status_raw(&self) -> &RawStatus {
        &self.status_raw
    }

    /// Consumes the Request to return the NetworkStream underneath.
    pub fn into_inner(self) -> Box<NetworkStream + Send> {
        self.body.into_inner().into_inner()
    }
}

impl Read for Response {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.body.read(buf)
    }
}

#[cfg(test)]
mod tests {
    use std::borrow::Cow::Borrowed;
    use std::io::{self, Read};
    use std::marker::PhantomData;

    use buffer::BufReader;
    use header::Headers;
    use header::TransferEncoding;
    use header::Encoding;
    use http::HttpReader::EofReader;
    use http::RawStatus;
    use mock::MockStream;
    use status;
    use version;

    use super::Response;

    fn read_to_string(mut r: Response) -> io::Result<String> {
        let mut s = String::new();
        try!(r.read_to_string(&mut s));
        Ok(s)
    }

    #[test]
    fn test_into_inner() {
        let res = Response {
            status: status::StatusCode::Ok,
            headers: Headers::new(),
            version: version::HttpVersion::Http11,
            body: EofReader(BufReader::new(Box::new(MockStream::new()))),
            status_raw: RawStatus(200, Borrowed("OK")),
            _marker: PhantomData,
        };

        let b = res.into_inner().downcast::<MockStream>().ok().unwrap();
        assert_eq!(b, Box::new(MockStream::new()));

    }

    #[test]
    fn test_parse_chunked_response() {
        let stream = MockStream::with_input(b"\
            HTTP/1.1 200 OK\r\n\
            Transfer-Encoding: chunked\r\n\
            \r\n\
            1\r\n\
            q\r\n\
            2\r\n\
            we\r\n\
            2\r\n\
            rt\r\n\
            0\r\n\
            \r\n"
        );

        let res = Response::new(Box::new(stream)).unwrap();

        // The status line is correct?
        assert_eq!(res.status, status::StatusCode::Ok);
        assert_eq!(res.version, version::HttpVersion::Http11);
        // The header is correct?
        match res.headers.get::<TransferEncoding>() {
            Some(encodings) => {
                assert_eq!(1, encodings.len());
                assert_eq!(Encoding::Chunked, encodings[0]);
            },
            None => panic!("Transfer-Encoding: chunked expected!"),
        };
        // The body is correct?
        assert_eq!(read_to_string(res).unwrap(), "qwert".to_owned());
    }

    /// Tests that when a chunk size is not a valid radix-16 number, an error
    /// is returned.
    #[test]
    fn test_invalid_chunk_size_not_hex_digit() {
        let stream = MockStream::with_input(b"\
            HTTP/1.1 200 OK\r\n\
            Transfer-Encoding: chunked\r\n\
            \r\n\
            X\r\n\
            1\r\n\
            0\r\n\
            \r\n"
        );

        let res = Response::new(Box::new(stream)).unwrap();

        assert!(read_to_string(res).is_err());
    }

    /// Tests that when a chunk size contains an invalid extension, an error is
    /// returned.
    #[test]
    fn test_invalid_chunk_size_extension() {
        let stream = MockStream::with_input(b"\
            HTTP/1.1 200 OK\r\n\
            Transfer-Encoding: chunked\r\n\
            \r\n\
            1 this is an invalid extension\r\n\
            1\r\n\
            0\r\n\
            \r\n"
        );

        let res = Response::new(Box::new(stream)).unwrap();

        assert!(read_to_string(res).is_err());
    }

    /// Tests that when a valid extension that contains a digit is appended to
    /// the chunk size, the chunk is correctly read.
    #[test]
    fn test_chunk_size_with_extension() {
        let stream = MockStream::with_input(b"\
            HTTP/1.1 200 OK\r\n\
            Transfer-Encoding: chunked\r\n\
            \r\n\
            1;this is an extension with a digit 1\r\n\
            1\r\n\
            0\r\n\
            \r\n"
        );

        let res = Response::new(Box::new(stream)).unwrap();

        assert_eq!(read_to_string(res).unwrap(), "1".to_owned());
    }
}
<commit_msg>fix(client): don't close stream until EOF<commit_after>//! Client Responses
use std::io::{self, Read};
use std::marker::PhantomData;
use std::net::Shutdown;

use buffer::BufReader;
use header;
use header::{ContentLength, TransferEncoding};
use header::Encoding::Chunked;
use net::{NetworkStream, HttpStream};
use http::{self, HttpReader, RawStatus};
use http::HttpReader::{SizedReader, ChunkedReader, EofReader};
use status;
use version;

/// A response for a client request to a remote server.
#[derive(Debug)]
pub struct Response<S = HttpStream> {
    /// The status from the server.
    pub status: status::StatusCode,
    /// The headers from the server.
    pub headers: header::Headers,
    /// The HTTP version of this response from the server.
    pub version: version::HttpVersion,
    status_raw: RawStatus,
    body: HttpReader<BufReader<Box<NetworkStream + Send>>>,

    _marker: PhantomData<S>,
}

impl Response {

    /// Creates a new response from a server.
    pub fn new(stream: Box<NetworkStream + Send>) -> ::Result<Response> {
        trace!("Response::new");
        let mut stream = BufReader::new(stream);

        let head = try!(http::parse_response(&mut stream));
        let raw_status = head.subject;
        let headers = head.headers;

        let status = status::StatusCode::from_u16(raw_status.0);
        debug!("version={:?}, status={:?}", head.version, status);
        debug!("headers={:?}", headers);

        let body = if headers.has::<TransferEncoding>() {
            match headers.get::<TransferEncoding>() {
                Some(&TransferEncoding(ref codings)) => {
                    if codings.len() > 1 {
                        trace!("TODO: #2 handle other codings: {:?}", codings);
                    };

                    if codings.contains(&Chunked) {
                        ChunkedReader(stream, None)
                    } else {
                        trace!("not chuncked. read till eof");
                        EofReader(stream)
                    }
                }
                None => unreachable!()
            }
        } else if headers.has::<ContentLength>() {
            match headers.get::<ContentLength>() {
                Some(&ContentLength(len)) => SizedReader(stream, len),
                None => unreachable!()
            }
        } else {
            trace!("neither Transfer-Encoding nor Content-Length");
            EofReader(stream)
        };

        Ok(Response {
            status: status,
            version: head.version,
            headers: headers,
            body: body,
            status_raw: raw_status,
            _marker: PhantomData,
        })
    }

    /// Get the raw status code and reason.
    pub fn status_raw(&self) -> &RawStatus {
        &self.status_raw
    }

    /// Consumes the Request to return the NetworkStream underneath.
    pub fn into_inner(self) -> Box<NetworkStream + Send> {
        self.body.into_inner().into_inner()
    }
}

impl Read for Response {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let count = try!(self.body.read(buf));

        if count == 0 {
            if !http::should_keep_alive(self.version, &self.headers) {
                try!(self.body.get_mut().get_mut().close(Shutdown::Both));
            }
        }

        Ok(count)
    }
}

#[cfg(test)]
mod tests {
    use std::borrow::Cow::Borrowed;
    use std::io::{self, Read};
    use std::marker::PhantomData;

    use buffer::BufReader;
    use header::Headers;
    use header::TransferEncoding;
    use header::Encoding;
    use http::HttpReader::EofReader;
    use http::RawStatus;
    use mock::MockStream;
    use status;
    use version;

    use super::Response;

    fn read_to_string(mut r: Response) -> io::Result<String> {
        let mut s = String::new();
        try!(r.read_to_string(&mut s));
        Ok(s)
    }

    #[test]
    fn test_into_inner() {
        let res = Response {
            status: status::StatusCode::Ok,
            headers: Headers::new(),
            version: version::HttpVersion::Http11,
            body: EofReader(BufReader::new(Box::new(MockStream::new()))),
            status_raw: RawStatus(200, Borrowed("OK")),
            _marker: PhantomData,
        };

        let b = res.into_inner().downcast::<MockStream>().ok().unwrap();
        assert_eq!(b, Box::new(MockStream::new()));

    }

    #[test]
    fn test_parse_chunked_response() {
        let stream = MockStream::with_input(b"\
            HTTP/1.1 200 OK\r\n\
            Transfer-Encoding: chunked\r\n\
            \r\n\
            1\r\n\
            q\r\n\
            2\r\n\
            we\r\n\
            2\r\n\
            rt\r\n\
            0\r\n\
            \r\n"
        );

        let res = Response::new(Box::new(stream)).unwrap();

        // The status line is correct?
        assert_eq!(res.status, status::StatusCode::Ok);
        assert_eq!(res.version, version::HttpVersion::Http11);
        // The header is correct?
        match res.headers.get::<TransferEncoding>() {
            Some(encodings) => {
                assert_eq!(1, encodings.len());
                assert_eq!(Encoding::Chunked, encodings[0]);
            },
            None => panic!("Transfer-Encoding: chunked expected!"),
        };
        // The body is correct?
        assert_eq!(read_to_string(res).unwrap(), "qwert".to_owned());
    }

    /// Tests that when a chunk size is not a valid radix-16 number, an error
    /// is returned.
    #[test]
    fn test_invalid_chunk_size_not_hex_digit() {
        let stream = MockStream::with_input(b"\
            HTTP/1.1 200 OK\r\n\
            Transfer-Encoding: chunked\r\n\
            \r\n\
            X\r\n\
            1\r\n\
            0\r\n\
            \r\n"
        );

        let res = Response::new(Box::new(stream)).unwrap();

        assert!(read_to_string(res).is_err());
    }

    /// Tests that when a chunk size contains an invalid extension, an error is
    /// returned.
    #[test]
    fn test_invalid_chunk_size_extension() {
        let stream = MockStream::with_input(b"\
            HTTP/1.1 200 OK\r\n\
            Transfer-Encoding: chunked\r\n\
            \r\n\
            1 this is an invalid extension\r\n\
            1\r\n\
            0\r\n\
            \r\n"
        );

        let res = Response::new(Box::new(stream)).unwrap();

        assert!(read_to_string(res).is_err());
    }

    /// Tests that when a valid extension that contains a digit is appended to
    /// the chunk size, the chunk is correctly read.
    #[test]
    fn test_chunk_size_with_extension() {
        let stream = MockStream::with_input(b"\
            HTTP/1.1 200 OK\r\n\
            Transfer-Encoding: chunked\r\n\
            \r\n\
            1;this is an extension with a digit 1\r\n\
            1\r\n\
            0\r\n\
            \r\n"
        );

        let res = Response::new(Box::new(stream)).unwrap();

        assert_eq!(read_to_string(res).unwrap(), "1".to_owned());
    }
}
<|endoftext|>


<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 codegen;
    use driver;
    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::generate(ctxt, fct);

            driver::dump_asm(&buffer, &ctxt.interner.str(fct.name));

            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) -> (CodeMemory, extern "C" fn(T) -> T) {
        let m = compile(code);
        let ptr = m.ptr();

        (m, unsafe { mem::transmute(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 (mem, f) = fct1("fn f(a: int) -> int { return a; }");

        assert_eq!(0i32, f(0));
        assert_eq!(1, f(1));
        assert_eq!(2, f(2));
    }

    #[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: emit_var_store/load extended for bool/str<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];

        match var.data_type {
            BuiltinType::Bool => emit_movb_reg_memq(&mut self.buf, src, RBP, var.offset),
            BuiltinType::Int => emit_movl_reg_memq(&mut self.buf, src, RBP, var.offset),
            BuiltinType::Str => emit_movq_reg_memq(&mut self.buf, src, RBP, var.offset),
            BuiltinType::Unit => {},
        }
    }

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

        match var.data_type {
            BuiltinType::Bool => emit_movb_memq_reg(&mut self.buf, RBP, var.offset, dest),
            BuiltinType::Int => emit_movl_memq_reg(&mut self.buf, RBP, var.offset, dest),
            BuiltinType::Str => emit_movq_memq_reg(&mut self.buf, RBP, var.offset, dest),
            BuiltinType::Unit => {},
        }
    }
}

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 codegen;
    use driver;
    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::generate(ctxt, fct);

            driver::dump_asm(&buffer, &ctxt.interner.str(fct.name));

            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) -> (CodeMemory, extern "C" fn(T) -> T) {
        let m = compile(code);
        let ptr = m.ptr();

        (m, unsafe { mem::transmute(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 (mem, f) = fct1("fn f(a: int) -> int { return a; }");

        assert_eq!(0i32, f(0));
        assert_eq!(1, f(1));
        assert_eq!(2, f(2));
    }

    #[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; }"));
        assert_eq!(true, run("fn f() -> bool { var a = true; return a; }"));
        assert_eq!(false, run("fn f() -> bool { var a = false; return a; }"));
    }
}
<|endoftext|>


<commit_before>//! Abstractions for default-sized and huge virtual memory pages.

use crate::structures::paging::page_table::PageTableLevel;
use crate::structures::paging::PageTableIndex;
use crate::VirtAddr;
use core::fmt;
#[cfg(feature = "step_trait")]
use core::iter::Step;
use core::marker::PhantomData;
use core::ops::{Add, AddAssign, Sub, SubAssign};

/// Trait for abstracting over the three possible page sizes on x86_64, 4KiB, 2MiB, 1GiB.
pub trait PageSize: Copy + Eq + PartialOrd + Ord {
    /// The page size in bytes.
    const SIZE: u64;

    /// A string representation of the page size for debug output.
    const SIZE_AS_DEBUG_STR: &'static str;
}

/// This trait is implemented for 4KiB and 2MiB pages, but not for 1GiB pages.
pub trait NotGiantPageSize: PageSize {}

/// A standard 4KiB page.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size4KiB {}

/// A “huge” 2MiB page.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size2MiB {}

/// A “giant” 1GiB page.
///
/// (Only available on newer x86_64 CPUs.)
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size1GiB {}

impl PageSize for Size4KiB {
    const SIZE: u64 = 4096;
    const SIZE_AS_DEBUG_STR: &'static str = "4KiB";
}

impl NotGiantPageSize for Size4KiB {}

impl PageSize for Size2MiB {
    const SIZE: u64 = Size4KiB::SIZE * 512;
    const SIZE_AS_DEBUG_STR: &'static str = "2MiB";
}

impl NotGiantPageSize for Size2MiB {}

impl PageSize for Size1GiB {
    const SIZE: u64 = Size2MiB::SIZE * 512;
    const SIZE_AS_DEBUG_STR: &'static str = "1GiB";
}

/// A virtual memory page.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(C)]
pub struct Page<S: PageSize = Size4KiB> {
    start_address: VirtAddr,
    size: PhantomData<S>,
}

impl<S: PageSize> Page<S> {
    /// The page size in bytes.
    pub const SIZE: u64 = S::SIZE;

    /// Returns the page that starts at the given virtual address.
    ///
    /// Returns an error if the address is not correctly aligned (i.e. is not a valid page start).
    #[inline]
    pub fn from_start_address(address: VirtAddr) -> Result<Self, AddressNotAligned> {
        if !address.is_aligned(S::SIZE) {
            return Err(AddressNotAligned);
        }
        Ok(Page::containing_address(address))
    }

    const_fn! {
        /// Returns the page that starts at the given virtual address.
        ///
        /// ## Safety
        ///
        /// The address must be correctly aligned.
        #[inline]
        pub unsafe fn from_start_address_unchecked(start_address: VirtAddr) -> Self {
            Page {
                start_address,
                size: PhantomData,
            }
        }
    }

    /// Returns the page that contains the given virtual address.
    #[inline]
    pub fn containing_address(address: VirtAddr) -> Self {
        Page {
            start_address: address.align_down(S::SIZE),
            size: PhantomData,
        }
    }

    const_fn! {
        /// Returns the start address of the page.
        #[inline]
        pub fn start_address(self) -> VirtAddr {
            self.start_address
        }
    }

    const_fn! {
        /// Returns the size the page (4KB, 2MB or 1GB).
        #[inline]
        pub fn size(self) -> u64 {
            S::SIZE
        }
    }

    const_fn! {
        /// Returns the level 4 page table index of this page.
        #[inline]
        pub fn p4_index(self) -> PageTableIndex {
            self.start_address().p4_index()
        }
    }

    const_fn! {
        /// Returns the level 3 page table index of this page.
        #[inline]
        pub fn p3_index(self) -> PageTableIndex {
            self.start_address().p3_index()
        }
    }

    const_fn! {
        /// Returns the table index of this page at the specified level.
        #[inline]
        pub fn page_table_index(self, level: PageTableLevel) -> PageTableIndex {
            self.start_address().page_table_index(level)
        }
    }

    const_fn! {
        /// Returns a range of pages, exclusive `end`.
        #[inline]
        pub fn range(start: Self, end: Self) -> PageRange<S> {
            PageRange { start, end }
        }
    }

    const_fn! {
        /// Returns a range of pages, inclusive `end`.
        #[inline]
        pub fn range_inclusive(start: Self, end: Self) -> PageRangeInclusive<S> {
            PageRangeInclusive { start, end }
        }
    }
}

impl<S: NotGiantPageSize> Page<S> {
    const_fn! {
        /// Returns the level 2 page table index of this page.
        #[inline]
        pub fn p2_index(self) -> PageTableIndex {
            self.start_address().p2_index()
        }
    }
}

impl Page<Size1GiB> {
    /// Returns the 1GiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices_1gib(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        Page::containing_address(VirtAddr::new(addr))
    }
}

impl Page<Size2MiB> {
    /// Returns the 2MiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices_2mib(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
        p2_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        addr.set_bits(21..30, u64::from(p2_index));
        Page::containing_address(VirtAddr::new(addr))
    }
}

impl Page<Size4KiB> {
    /// Returns the 4KiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
        p2_index: PageTableIndex,
        p1_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        addr.set_bits(21..30, u64::from(p2_index));
        addr.set_bits(12..21, u64::from(p1_index));
        Page::containing_address(VirtAddr::new(addr))
    }

    /// Returns the level 1 page table index of this page.
    #[inline]
    pub const fn p1_index(self) -> PageTableIndex {
        self.start_address.p1_index()
    }
}

impl<S: PageSize> fmt::Debug for Page<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_fmt(format_args!(
            "Page[{}]({:#x})",
            S::SIZE_AS_DEBUG_STR,
            self.start_address().as_u64()
        ))
    }
}

impl<S: PageSize> Add<u64> for Page<S> {
    type Output = Self;
    #[inline]
    fn add(self, rhs: u64) -> Self::Output {
        Page::containing_address(self.start_address() + rhs * S::SIZE)
    }
}

impl<S: PageSize> AddAssign<u64> for Page<S> {
    #[inline]
    fn add_assign(&mut self, rhs: u64) {
        *self = *self + rhs;
    }
}

impl<S: PageSize> Sub<u64> for Page<S> {
    type Output = Self;
    #[inline]
    fn sub(self, rhs: u64) -> Self::Output {
        Page::containing_address(self.start_address() - rhs * S::SIZE)
    }
}

impl<S: PageSize> SubAssign<u64> for Page<S> {
    #[inline]
    fn sub_assign(&mut self, rhs: u64) {
        *self = *self - rhs;
    }
}

impl<S: PageSize> Sub<Self> for Page<S> {
    type Output = u64;
    #[inline]
    fn sub(self, rhs: Self) -> Self::Output {
        (self.start_address - rhs.start_address) / S::SIZE
    }
}

#[cfg(feature = "step_trait")]
impl<S: PageSize> Step for Page<S> {
    fn steps_between(start: &Self, end: &Self) -> Option<usize> {
        Step::steps_between(&start.start_address, &end.start_address)
            .map(|steps| steps / S::SIZE as usize)
    }

    fn forward_checked(start: Self, count: usize) -> Option<Self> {
        let count = count.checked_mul(S::SIZE as usize)?;
        let start_address = Step::forward_checked(start.start_address, count)?;
        Some(Self {
            start_address,
            size: PhantomData,
        })
    }

    fn backward_checked(start: Self, count: usize) -> Option<Self> {
        let count = count.checked_mul(S::SIZE as usize)?;
        let start_address = Step::backward_checked(start.start_address, count)?;
        Some(Self {
            start_address,
            size: PhantomData,
        })
    }
}

/// A range of pages with exclusive upper bound.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
#[repr(C)]
pub struct PageRange<S: PageSize = Size4KiB> {
    /// The start of the range, inclusive.
    pub start: Page<S>,
    /// The end of the range, exclusive.
    pub end: Page<S>,
}

impl<S: PageSize> PageRange<S> {
    /// Returns wether this range contains no pages.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.start >= self.end
    }
}

impl<S: PageSize> Iterator for PageRange<S> {
    type Item = Page<S>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.start < self.end {
            let page = self.start;
            self.start += 1;
            Some(page)
        } else {
            None
        }
    }
}

impl PageRange<Size2MiB> {
    /// Converts the range of 2MiB pages to a range of 4KiB pages.
    #[inline]
    pub fn as_4kib_page_range(self) -> PageRange<Size4KiB> {
        PageRange {
            start: Page::containing_address(self.start.start_address()),
            end: Page::containing_address(self.end.start_address()),
        }
    }
}

impl<S: PageSize> fmt::Debug for PageRange<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("PageRange")
            .field("start", &self.start)
            .field("end", &self.end)
            .finish()
    }
}

/// A range of pages with inclusive upper bound.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
#[repr(C)]
pub struct PageRangeInclusive<S: PageSize = Size4KiB> {
    /// The start of the range, inclusive.
    pub start: Page<S>,
    /// The end of the range, inclusive.
    pub end: Page<S>,
}

impl<S: PageSize> PageRangeInclusive<S> {
    /// Returns wether this range contains no pages.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.start > self.end
    }
}

impl<S: PageSize> Iterator for PageRangeInclusive<S> {
    type Item = Page<S>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.start <= self.end {
            let page = self.start;

            // If the end of the inclusive range is the maximum page possible for size S,
            // incrementing start until it is greater than the end will cause an integer overflow.
            // So instead, in that case we decrement end rather than incrementing start.
            let max_page_addr = VirtAddr::new(u64::MAX) - (S::SIZE - 1);
            if self.start.start_address() < max_page_addr {
                self.start += 1;
            } else {
                self.end -= 1;
            }
            Some(page)
        } else {
            None
        }
    }
}

impl<S: PageSize> fmt::Debug for PageRangeInclusive<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("PageRangeInclusive")
            .field("start", &self.start)
            .field("end", &self.end)
            .finish()
    }
}

/// The given address was not sufficiently aligned.
#[derive(Debug)]
pub struct AddressNotAligned;

impl fmt::Display for AddressNotAligned {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "the given address was not sufficiently aligned")
    }
}

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

    #[test]
    pub fn test_page_ranges() {
        let page_size = Size4KiB::SIZE;
        let number = 1000;

        let start_addr = VirtAddr::new(0xdead_beaf);
        let start: Page = Page::containing_address(start_addr);
        let end = start + number;

        let mut range = Page::range(start, end);
        for i in 0..number {
            assert_eq!(
                range.next(),
                Some(Page::containing_address(start_addr + page_size * i))
            );
        }
        assert_eq!(range.next(), None);

        let mut range_inclusive = Page::range_inclusive(start, end);
        for i in 0..=number {
            assert_eq!(
                range_inclusive.next(),
                Some(Page::containing_address(start_addr + page_size * i))
            );
        }
        assert_eq!(range_inclusive.next(), None);
    }
}
<commit_msg>Added test to cover ranges ending in the max page<commit_after>//! Abstractions for default-sized and huge virtual memory pages.

use crate::structures::paging::page_table::PageTableLevel;
use crate::structures::paging::PageTableIndex;
use crate::VirtAddr;
use core::fmt;
#[cfg(feature = "step_trait")]
use core::iter::Step;
use core::marker::PhantomData;
use core::ops::{Add, AddAssign, Sub, SubAssign};

/// Trait for abstracting over the three possible page sizes on x86_64, 4KiB, 2MiB, 1GiB.
pub trait PageSize: Copy + Eq + PartialOrd + Ord {
    /// The page size in bytes.
    const SIZE: u64;

    /// A string representation of the page size for debug output.
    const SIZE_AS_DEBUG_STR: &'static str;
}

/// This trait is implemented for 4KiB and 2MiB pages, but not for 1GiB pages.
pub trait NotGiantPageSize: PageSize {}

/// A standard 4KiB page.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size4KiB {}

/// A “huge” 2MiB page.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size2MiB {}

/// A “giant” 1GiB page.
///
/// (Only available on newer x86_64 CPUs.)
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Size1GiB {}

impl PageSize for Size4KiB {
    const SIZE: u64 = 4096;
    const SIZE_AS_DEBUG_STR: &'static str = "4KiB";
}

impl NotGiantPageSize for Size4KiB {}

impl PageSize for Size2MiB {
    const SIZE: u64 = Size4KiB::SIZE * 512;
    const SIZE_AS_DEBUG_STR: &'static str = "2MiB";
}

impl NotGiantPageSize for Size2MiB {}

impl PageSize for Size1GiB {
    const SIZE: u64 = Size2MiB::SIZE * 512;
    const SIZE_AS_DEBUG_STR: &'static str = "1GiB";
}

/// A virtual memory page.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(C)]
pub struct Page<S: PageSize = Size4KiB> {
    start_address: VirtAddr,
    size: PhantomData<S>,
}

impl<S: PageSize> Page<S> {
    /// The page size in bytes.
    pub const SIZE: u64 = S::SIZE;

    /// Returns the page that starts at the given virtual address.
    ///
    /// Returns an error if the address is not correctly aligned (i.e. is not a valid page start).
    #[inline]
    pub fn from_start_address(address: VirtAddr) -> Result<Self, AddressNotAligned> {
        if !address.is_aligned(S::SIZE) {
            return Err(AddressNotAligned);
        }
        Ok(Page::containing_address(address))
    }

    const_fn! {
        /// Returns the page that starts at the given virtual address.
        ///
        /// ## Safety
        ///
        /// The address must be correctly aligned.
        #[inline]
        pub unsafe fn from_start_address_unchecked(start_address: VirtAddr) -> Self {
            Page {
                start_address,
                size: PhantomData,
            }
        }
    }

    /// Returns the page that contains the given virtual address.
    #[inline]
    pub fn containing_address(address: VirtAddr) -> Self {
        Page {
            start_address: address.align_down(S::SIZE),
            size: PhantomData,
        }
    }

    const_fn! {
        /// Returns the start address of the page.
        #[inline]
        pub fn start_address(self) -> VirtAddr {
            self.start_address
        }
    }

    const_fn! {
        /// Returns the size the page (4KB, 2MB or 1GB).
        #[inline]
        pub fn size(self) -> u64 {
            S::SIZE
        }
    }

    const_fn! {
        /// Returns the level 4 page table index of this page.
        #[inline]
        pub fn p4_index(self) -> PageTableIndex {
            self.start_address().p4_index()
        }
    }

    const_fn! {
        /// Returns the level 3 page table index of this page.
        #[inline]
        pub fn p3_index(self) -> PageTableIndex {
            self.start_address().p3_index()
        }
    }

    const_fn! {
        /// Returns the table index of this page at the specified level.
        #[inline]
        pub fn page_table_index(self, level: PageTableLevel) -> PageTableIndex {
            self.start_address().page_table_index(level)
        }
    }

    const_fn! {
        /// Returns a range of pages, exclusive `end`.
        #[inline]
        pub fn range(start: Self, end: Self) -> PageRange<S> {
            PageRange { start, end }
        }
    }

    const_fn! {
        /// Returns a range of pages, inclusive `end`.
        #[inline]
        pub fn range_inclusive(start: Self, end: Self) -> PageRangeInclusive<S> {
            PageRangeInclusive { start, end }
        }
    }
}

impl<S: NotGiantPageSize> Page<S> {
    const_fn! {
        /// Returns the level 2 page table index of this page.
        #[inline]
        pub fn p2_index(self) -> PageTableIndex {
            self.start_address().p2_index()
        }
    }
}

impl Page<Size1GiB> {
    /// Returns the 1GiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices_1gib(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        Page::containing_address(VirtAddr::new(addr))
    }
}

impl Page<Size2MiB> {
    /// Returns the 2MiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices_2mib(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
        p2_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        addr.set_bits(21..30, u64::from(p2_index));
        Page::containing_address(VirtAddr::new(addr))
    }
}

impl Page<Size4KiB> {
    /// Returns the 4KiB memory page with the specified page table indices.
    #[inline]
    pub fn from_page_table_indices(
        p4_index: PageTableIndex,
        p3_index: PageTableIndex,
        p2_index: PageTableIndex,
        p1_index: PageTableIndex,
    ) -> Self {
        use bit_field::BitField;

        let mut addr = 0;
        addr.set_bits(39..48, u64::from(p4_index));
        addr.set_bits(30..39, u64::from(p3_index));
        addr.set_bits(21..30, u64::from(p2_index));
        addr.set_bits(12..21, u64::from(p1_index));
        Page::containing_address(VirtAddr::new(addr))
    }

    /// Returns the level 1 page table index of this page.
    #[inline]
    pub const fn p1_index(self) -> PageTableIndex {
        self.start_address.p1_index()
    }
}

impl<S: PageSize> fmt::Debug for Page<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_fmt(format_args!(
            "Page[{}]({:#x})",
            S::SIZE_AS_DEBUG_STR,
            self.start_address().as_u64()
        ))
    }
}

impl<S: PageSize> Add<u64> for Page<S> {
    type Output = Self;
    #[inline]
    fn add(self, rhs: u64) -> Self::Output {
        Page::containing_address(self.start_address() + rhs * S::SIZE)
    }
}

impl<S: PageSize> AddAssign<u64> for Page<S> {
    #[inline]
    fn add_assign(&mut self, rhs: u64) {
        *self = *self + rhs;
    }
}

impl<S: PageSize> Sub<u64> for Page<S> {
    type Output = Self;
    #[inline]
    fn sub(self, rhs: u64) -> Self::Output {
        Page::containing_address(self.start_address() - rhs * S::SIZE)
    }
}

impl<S: PageSize> SubAssign<u64> for Page<S> {
    #[inline]
    fn sub_assign(&mut self, rhs: u64) {
        *self = *self - rhs;
    }
}

impl<S: PageSize> Sub<Self> for Page<S> {
    type Output = u64;
    #[inline]
    fn sub(self, rhs: Self) -> Self::Output {
        (self.start_address - rhs.start_address) / S::SIZE
    }
}

#[cfg(feature = "step_trait")]
impl<S: PageSize> Step for Page<S> {
    fn steps_between(start: &Self, end: &Self) -> Option<usize> {
        Step::steps_between(&start.start_address, &end.start_address)
            .map(|steps| steps / S::SIZE as usize)
    }

    fn forward_checked(start: Self, count: usize) -> Option<Self> {
        let count = count.checked_mul(S::SIZE as usize)?;
        let start_address = Step::forward_checked(start.start_address, count)?;
        Some(Self {
            start_address,
            size: PhantomData,
        })
    }

    fn backward_checked(start: Self, count: usize) -> Option<Self> {
        let count = count.checked_mul(S::SIZE as usize)?;
        let start_address = Step::backward_checked(start.start_address, count)?;
        Some(Self {
            start_address,
            size: PhantomData,
        })
    }
}

/// A range of pages with exclusive upper bound.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
#[repr(C)]
pub struct PageRange<S: PageSize = Size4KiB> {
    /// The start of the range, inclusive.
    pub start: Page<S>,
    /// The end of the range, exclusive.
    pub end: Page<S>,
}

impl<S: PageSize> PageRange<S> {
    /// Returns wether this range contains no pages.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.start >= self.end
    }
}

impl<S: PageSize> Iterator for PageRange<S> {
    type Item = Page<S>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.start < self.end {
            let page = self.start;
            self.start += 1;
            Some(page)
        } else {
            None
        }
    }
}

impl PageRange<Size2MiB> {
    /// Converts the range of 2MiB pages to a range of 4KiB pages.
    #[inline]
    pub fn as_4kib_page_range(self) -> PageRange<Size4KiB> {
        PageRange {
            start: Page::containing_address(self.start.start_address()),
            end: Page::containing_address(self.end.start_address()),
        }
    }
}

impl<S: PageSize> fmt::Debug for PageRange<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("PageRange")
            .field("start", &self.start)
            .field("end", &self.end)
            .finish()
    }
}

/// A range of pages with inclusive upper bound.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
#[repr(C)]
pub struct PageRangeInclusive<S: PageSize = Size4KiB> {
    /// The start of the range, inclusive.
    pub start: Page<S>,
    /// The end of the range, inclusive.
    pub end: Page<S>,
}

impl<S: PageSize> PageRangeInclusive<S> {
    /// Returns wether this range contains no pages.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.start > self.end
    }
}

impl<S: PageSize> Iterator for PageRangeInclusive<S> {
    type Item = Page<S>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.start <= self.end {
            let page = self.start;

            // If the end of the inclusive range is the maximum page possible for size S,
            // incrementing start until it is greater than the end will cause an integer overflow.
            // So instead, in that case we decrement end rather than incrementing start.
            let max_page_addr = VirtAddr::new(u64::MAX) - (S::SIZE - 1);
            if self.start.start_address() < max_page_addr {
                self.start += 1;
            } else {
                self.end -= 1;
            }
            Some(page)
        } else {
            None
        }
    }
}

impl<S: PageSize> fmt::Debug for PageRangeInclusive<S> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("PageRangeInclusive")
            .field("start", &self.start)
            .field("end", &self.end)
            .finish()
    }
}

/// The given address was not sufficiently aligned.
#[derive(Debug)]
pub struct AddressNotAligned;

impl fmt::Display for AddressNotAligned {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "the given address was not sufficiently aligned")
    }
}

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

    #[test]
    pub fn test_page_ranges() {
        let page_size = Size4KiB::SIZE;
        let number = 1000;

        let start_addr = VirtAddr::new(0xdead_beaf);
        let start: Page = Page::containing_address(start_addr);
        let end = start + number;

        let mut range = Page::range(start, end);
        for i in 0..number {
            assert_eq!(
                range.next(),
                Some(Page::containing_address(start_addr + page_size * i))
            );
        }
        assert_eq!(range.next(), None);

        let mut range_inclusive = Page::range_inclusive(start, end);
        for i in 0..=number {
            assert_eq!(
                range_inclusive.next(),
                Some(Page::containing_address(start_addr + page_size * i))
            );
        }
        assert_eq!(range_inclusive.next(), None);
    }

    #[test]
    pub fn test_max_page_overflow() {
        let page_size = Size4KiB::SIZE;
        let number = 1000;

        let start_addr = VirtAddr::new(u64::MAX).align_down(page_size) - number * page_size;
        let start: Page = Page::containing_address(start_addr);
        let end = start + number;

        let mut range_inclusive = Page::range_inclusive(start, end);
        for i in 0..=number {
            assert_eq!(
                range_inclusive.next(),
                Some(Page::containing_address(start_addr + page_size * i))
            );
        }
        assert_eq!(range_inclusive.next(), None);
    }
}
<|endoftext|>


<commit_before>use crate::PortMappingProtocol;
use std::net::SocketAddrV4;

// Content of the request.
pub const SEARCH_REQUEST: &'static str = "M-SEARCH * HTTP/1.1\r
Host:239.255.255.250:1900\r
ST:urn:schemas-upnp-org:device:InternetGatewayDevice:1\r
Man:\"ssdp:discover\"\r
MX:3\r\n\r\n";

pub const GET_EXTERNAL_IP_HEADER: &'static str =
    r#""urn:schemas-upnp-org:service:WANIPConnection:1#GetExternalIPAddress""#;

pub const ADD_ANY_PORT_MAPPING_HEADER: &'static str =
    r#""urn:schemas-upnp-org:service:WANIPConnection:1#AddAnyPortMapping""#;

pub const ADD_PORT_MAPPING_HEADER: &'static str = r#""urn:schemas-upnp-org:service:WANIPConnection:1#AddPortMapping""#;

pub const DELETE_PORT_MAPPING_HEADER: &'static str =
    r#""urn:schemas-upnp-org:service:WANIPConnection:1#DeletePortMapping""#;

pub const GET_GENERIC_PORT_MAPPING_ENTRY: &'static str =
    r#""urn:schemas-upnp-org:service:WANIPConnection:1#GetGenericPortMappingEntry""#;

const MESSAGE_HEAD: &'static str = r#"<?xml version="1.0"?>
<s:Envelope s:encodingStyle="http://schemas.xmlsoap.org/soap/encoding/" xmlns:s="http://schemas.xmlsoap.org/soap/envelope/">
<s:Body>"#;

const MESSAGE_TAIL: &'static str = r#"</s:Body>
</s:Envelope>"#;

fn format_message(body: String) -> String {
    format!("{}{}{}", MESSAGE_HEAD, body, MESSAGE_TAIL)
}

pub fn format_get_external_ip_message() -> String {
    format!(
        r#"<?xml version="1.0"?>
<s:Envelope s:encodingStyle="http://schemas.xmlsoap.org/soap/encoding/" xmlns:s="http://schemas.xmlsoap.org/soap/envelope/">
    <s:Body>
        <m:GetExternalIPAddress xmlns:m="urn:schemas-upnp-org:service:WANIPConnection:1">
        </m:GetExternalIPAddress>
    </s:Body>
</s:Envelope>"#
    )
}

pub fn format_add_any_port_mapping_message(
    schema: &Vec<String>,
    protocol: PortMappingProtocol,
    external_port: u16,
    local_addr: SocketAddrV4,
    lease_duration: u32,
    description: &str,
) -> String {
    let args = schema
        .iter()
        .map(|argument| {
            let value = match argument.as_str() {
                "NewEnabled" => 1.to_string(),
                "NewExternalPort" => external_port.to_string(),
                "NewInternalClient" => local_addr.ip().to_string(),
                "NewInternalPort" => local_addr.port().to_string(),
                "NewLeaseDuration" => lease_duration.to_string(),
                "NewPortMappingDescription" => description.to_string(),
                "NewProtocol" => protocol.to_string(),
                "NewRemoteHost" => "".to_string(),
                _ => panic!("Nope"),
            };
            format!("<{argument}>{value}</{argument}>", argument = argument, value = value)
        })
        .collect::<Vec<_>>()
        .join("\n");

    format_message(format!(
        r#"<u:AddAnyPortMapping xmlns:u="urn:schemas-upnp-org:service:WANIPConnection:1">
        {}
        </u:AddAnyPortMapping>"#,
        args,
    ))
}

pub fn format_add_port_mapping_message(
    schema: &Vec<String>,
    protocol: PortMappingProtocol,
    external_port: u16,
    local_addr: SocketAddrV4,
    lease_duration: u32,
    description: &str,
) -> String {
    let args = schema
        .iter()
        .map(|argument| {
            let value = match argument.as_str() {
                "NewEnabled" => 1.to_string(),
                "NewExternalPort" => external_port.to_string(),
                "NewInternalClient" => local_addr.ip().to_string(),
                "NewInternalPort" => local_addr.port().to_string(),
                "NewLeaseDuration" => lease_duration.to_string(),
                "NewPortMappingDescription" => description.to_string(),
                "NewProtocol" => protocol.to_string(),
                "NewRemoteHost" => "".to_string(),
                _ => panic!("Nope"),
            };
            format!("<{argument}>{value}</{argument}>", argument = argument, value = value)
        })
        .collect::<Vec<_>>()
        .join("\n");

    format_message(format!(
        r#"<u:AddPortMapping xmlns:u="urn:schemas-upnp-org:service:WANIPConnection:1">
        {}
        </u:AddPortMapping>"#,
        args,
    ))
}

pub fn format_delete_port_message(schema: &Vec<String>, protocol: PortMappingProtocol, external_port: u16) -> String {
    let args = schema
        .iter()
        .map(|argument| {
            let value = match argument.as_str() {
                "NewExternalPort" => external_port.to_string(),
                "NewProtocol" => protocol.to_string(),
                "NewRemoteHost" => "".to_string(),
                _ => panic!("Nope"),
            };
            format!("<{argument}>{value}</{argument}>", argument = argument, value = value)
        })
        .collect::<Vec<_>>()
        .join("\n");

    format_message(format!(
        r#"<u:DeletePortMapping xmlns:u="urn:schemas-upnp-org:service:WANIPConnection:1">
        {}
        </u:DeletePortMapping>"#,
        args,
    ))
}

pub fn formate_get_generic_port_mapping_entry_message(port_mapping_index: u32) -> String {
    format_message(format!(
        r#"<u:GetGenericPortMappingEntry xmlns:u="urn:schemas-upnp-org:service:WANIPConnection:1">
        <NewPortMappingIndex>{}</NewPortMappingIndex>
        </u:GetGenericPortMappingEntry>"#,
        port_mapping_index
    ))
}
<commit_msg>Warn about extraneous arguments, do not panic<commit_after>use crate::PortMappingProtocol;
use std::net::SocketAddrV4;

// Content of the request.
pub const SEARCH_REQUEST: &'static str = "M-SEARCH * HTTP/1.1\r
Host:239.255.255.250:1900\r
ST:urn:schemas-upnp-org:device:InternetGatewayDevice:1\r
Man:\"ssdp:discover\"\r
MX:3\r\n\r\n";

pub const GET_EXTERNAL_IP_HEADER: &'static str =
    r#""urn:schemas-upnp-org:service:WANIPConnection:1#GetExternalIPAddress""#;

pub const ADD_ANY_PORT_MAPPING_HEADER: &'static str =
    r#""urn:schemas-upnp-org:service:WANIPConnection:1#AddAnyPortMapping""#;

pub const ADD_PORT_MAPPING_HEADER: &'static str = r#""urn:schemas-upnp-org:service:WANIPConnection:1#AddPortMapping""#;

pub const DELETE_PORT_MAPPING_HEADER: &'static str =
    r#""urn:schemas-upnp-org:service:WANIPConnection:1#DeletePortMapping""#;

pub const GET_GENERIC_PORT_MAPPING_ENTRY: &'static str =
    r#""urn:schemas-upnp-org:service:WANIPConnection:1#GetGenericPortMappingEntry""#;

const MESSAGE_HEAD: &'static str = r#"<?xml version="1.0"?>
<s:Envelope s:encodingStyle="http://schemas.xmlsoap.org/soap/encoding/" xmlns:s="http://schemas.xmlsoap.org/soap/envelope/">
<s:Body>"#;

const MESSAGE_TAIL: &'static str = r#"</s:Body>
</s:Envelope>"#;

fn format_message(body: String) -> String {
    format!("{}{}{}", MESSAGE_HEAD, body, MESSAGE_TAIL)
}

pub fn format_get_external_ip_message() -> String {
    format!(
        r#"<?xml version="1.0"?>
<s:Envelope s:encodingStyle="http://schemas.xmlsoap.org/soap/encoding/" xmlns:s="http://schemas.xmlsoap.org/soap/envelope/">
    <s:Body>
        <m:GetExternalIPAddress xmlns:m="urn:schemas-upnp-org:service:WANIPConnection:1">
        </m:GetExternalIPAddress>
    </s:Body>
</s:Envelope>"#
    )
}

pub fn format_add_any_port_mapping_message(
    schema: &Vec<String>,
    protocol: PortMappingProtocol,
    external_port: u16,
    local_addr: SocketAddrV4,
    lease_duration: u32,
    description: &str,
) -> String {
    let args = schema
        .iter()
        .filter_map(|argument| {
            let value = match argument.as_str() {
                "NewEnabled" => 1.to_string(),
                "NewExternalPort" => external_port.to_string(),
                "NewInternalClient" => local_addr.ip().to_string(),
                "NewInternalPort" => local_addr.port().to_string(),
                "NewLeaseDuration" => lease_duration.to_string(),
                "NewPortMappingDescription" => description.to_string(),
                "NewProtocol" => protocol.to_string(),
                "NewRemoteHost" => "".to_string(),
                unknown => {
                    warn!("Unknown argument: {}", unknown);
                    return None;
                }
            };
            Some(format!(
                "<{argument}>{value}</{argument}>",
                argument = argument,
                value = value
            ))
        })
        .collect::<Vec<_>>()
        .join("\n");

    format_message(format!(
        r#"<u:AddAnyPortMapping xmlns:u="urn:schemas-upnp-org:service:WANIPConnection:1">
        {}
        </u:AddAnyPortMapping>"#,
        args,
    ))
}

pub fn format_add_port_mapping_message(
    schema: &Vec<String>,
    protocol: PortMappingProtocol,
    external_port: u16,
    local_addr: SocketAddrV4,
    lease_duration: u32,
    description: &str,
) -> String {
    let args = schema
        .iter()
        .filter_map(|argument| {
            let value = match argument.as_str() {
                "NewEnabled" => 1.to_string(),
                "NewExternalPort" => external_port.to_string(),
                "NewInternalClient" => local_addr.ip().to_string(),
                "NewInternalPort" => local_addr.port().to_string(),
                "NewLeaseDuration" => lease_duration.to_string(),
                "NewPortMappingDescription" => description.to_string(),
                "NewProtocol" => protocol.to_string(),
                "NewRemoteHost" => "".to_string(),
                unknown => {
                    warn!("Unknown argument: {}", unknown);
                    return None;
                }
            };
            Some(format!(
                "<{argument}>{value}</{argument}>",
                argument = argument,
                value = value
            ))
        })
        .collect::<Vec<_>>()
        .join("\n");

    format_message(format!(
        r#"<u:AddPortMapping xmlns:u="urn:schemas-upnp-org:service:WANIPConnection:1">
        {}
        </u:AddPortMapping>"#,
        args,
    ))
}

pub fn format_delete_port_message(schema: &Vec<String>, protocol: PortMappingProtocol, external_port: u16) -> String {
    let args = schema
        .iter()
        .filter_map(|argument| {
            let value = match argument.as_str() {
                "NewExternalPort" => external_port.to_string(),
                "NewProtocol" => protocol.to_string(),
                "NewRemoteHost" => "".to_string(),
                unknown => {
                    warn!("Unknown argument: {}", unknown);
                    return None;
                }
            };
            Some(format!(
                "<{argument}>{value}</{argument}>",
                argument = argument,
                value = value
            ))
        })
        .collect::<Vec<_>>()
        .join("\n");

    format_message(format!(
        r#"<u:DeletePortMapping xmlns:u="urn:schemas-upnp-org:service:WANIPConnection:1">
        {}
        </u:DeletePortMapping>"#,
        args,
    ))
}

pub fn formate_get_generic_port_mapping_entry_message(port_mapping_index: u32) -> String {
    format_message(format!(
        r#"<u:GetGenericPortMappingEntry xmlns:u="urn:schemas-upnp-org:service:WANIPConnection:1">
        <NewPortMappingIndex>{}</NewPortMappingIndex>
        </u:GetGenericPortMappingEntry>"#,
        port_mapping_index
    ))
}
<|endoftext|>


<commit_before>#![allow(clippy::useless_attribute)] //issue #2910

#[macro_use]
extern crate serde_derive;

/// Tests that we do not lint for unused underscores in a `MacroAttribute`
/// expansion
#[deny(clippy::used_underscore_binding)]
#[derive(Deserialize)]
struct MacroAttributesTest {
    _foo: u32,
}

#[test]
fn macro_attributes_test() {
    let _ = MacroAttributesTest { _foo: 0 };
}

fn main() {}
<commit_msg>Auto merge of #78399 - vn-ki:gsgdt-graphviz, r=oli-obk<commit_after>// edition:2018

use serde::Deserialize;

/// Tests that we do not lint for unused underscores in a `MacroAttribute`
/// expansion
#[deny(clippy::used_underscore_binding)]
#[derive(Deserialize)]
struct MacroAttributesTest {
    _foo: u32,
}

#[test]
fn macro_attributes_test() {
    let _ = MacroAttributesTest { _foo: 0 };
}

fn main() {}
<|endoftext|>


<commit_before>//! rev4 of the verification dataset from HSLuv
//! https://github.com/hsluv/hsluv/blob/master/snapshots/snapshot-rev4.json

use approx::assert_relative_eq;
use lazy_static::lazy_static;
use serde_json;

use palette::convert::IntoColorUnclamped;
use palette::white_point::D65;
use palette::{Hsluv, Lchuv, Luv, LuvHue, Xyz};
use std::collections::HashMap;

#[derive(Clone, Debug)]
struct HsluvExample {
    name: String,
    lchuv: Lchuv<D65, f64>,
    hsluv: Hsluv<D65, f64>,
    luv: Luv<D65, f64>,
    xyz: Xyz<D65, f64>,
}

type Examples = HashMap<String, HsluvExample>;

fn load_data() -> Examples {
    let filename = "tests/hsluv_dataset/hsluv_dataset.json";
    let data_str = std::fs::read_to_string(filename).unwrap();
    let raw_data: serde_json::Value = serde_json::from_str(&data_str).unwrap();

    let m = raw_data.as_object().expect("failed to parse dataset");
    let to_vec = |c: &serde_json::Value| {
        c.as_array()
            .unwrap()
            .iter()
            .flat_map(|x| x.as_f64())
            .collect()
    };

    m.iter()
        .map(|(k, v)| {
            let colors = v.as_object().unwrap();
            let luv_data: Vec<f64> = to_vec(&colors["luv"]);
            let lchuv_data: Vec<f64> = to_vec(&colors["lch"]);
            let hsluv_data: Vec<f64> = to_vec(&colors["hsluv"]);
            let xyz_data: Vec<f64> = to_vec(&colors["xyz"]);

            (
                k.clone(),
                HsluvExample {
                    name: k.clone(),
                    luv: Luv::new(luv_data[0], luv_data[1], luv_data[2]),
                    hsluv: Hsluv::new(hsluv_data[0], hsluv_data[1], hsluv_data[2]),
                    lchuv: Lchuv::new(lchuv_data[0], lchuv_data[1], lchuv_data[2]),
                    xyz: Xyz::new(xyz_data[0], xyz_data[1], xyz_data[2]),
                },
            )
        })
        .collect()
}

lazy_static! {
    static ref TEST_DATA: Examples = load_data();
}

#[test]
pub fn run_xyz_to_luv_tests() {
    for (_, v) in TEST_DATA.iter() {
        let to_luv: Luv<D65, f64> = v.xyz.into_color_unclamped();
        assert_relative_eq!(to_luv, v.luv, epsilon = 0.1);
    }
}

#[test]
pub fn run_luv_to_xyz_tests() {
    for (_, v) in TEST_DATA.iter() {
        let to_xyz: Xyz<D65, f64> = v.luv.into_color_unclamped();
        assert_relative_eq!(to_xyz, v.xyz, epsilon = 0.001);
    }
}

#[test]
pub fn run_lchuv_to_luv_tests() {
    for (_, v) in TEST_DATA.iter() {
        let to_luv: Luv<D65, f64> = v.lchuv.into_color_unclamped();
        assert_relative_eq!(to_luv, v.luv, epsilon = 0.001);
    }
}

#[test]
pub fn run_luv_to_lchuv_tests() {
    for (_, v) in TEST_DATA.iter() {
        let mut to_lchuv: Lchuv<D65, f64> = v.luv.into_color_unclamped();
        if to_lchuv.chroma < 1e-8 {
            to_lchuv.hue = LuvHue::from_degrees(0.0);
        }
        assert_relative_eq!(to_lchuv, v.lchuv, epsilon = 0.001);
    }
}

#[test]
pub fn run_lchuv_to_hsluv_tests() {
    for (_, v) in TEST_DATA.iter() {
        let mut to_hsluv: Hsluv<D65, f64> = v.lchuv.into_color_unclamped();
        if to_hsluv.l > 100.0 - 1e-5 {
            to_hsluv.saturation = 0.0;
        }
        assert_relative_eq!(to_hsluv, v.hsluv, epsilon = 1e-5);
    }
}

#[test]
pub fn run_hsluv_to_lchuv_tests() {
    for (_, v) in TEST_DATA.iter() {
        let to_lchuv: Lchuv<D65, f64> = v.hsluv.into_color_unclamped();
        assert_relative_eq!(to_lchuv, v.lchuv, epsilon = 1e-5);
    }
}
<commit_msg>Fix warning in hsluv_dataset test<commit_after>//! rev4 of the verification dataset from HSLuv
//! https://github.com/hsluv/hsluv/blob/master/snapshots/snapshot-rev4.json

use approx::assert_relative_eq;
use lazy_static::lazy_static;
use serde_json;

use palette::convert::IntoColorUnclamped;
use palette::white_point::D65;
use palette::{Hsluv, Lchuv, Luv, LuvHue, Xyz};
use std::collections::HashMap;

#[derive(Clone, Debug)]
struct HsluvExample {
    lchuv: Lchuv<D65, f64>,
    hsluv: Hsluv<D65, f64>,
    luv: Luv<D65, f64>,
    xyz: Xyz<D65, f64>,
}

type Examples = HashMap<String, HsluvExample>;

fn load_data() -> Examples {
    let filename = "tests/hsluv_dataset/hsluv_dataset.json";
    let data_str = std::fs::read_to_string(filename).unwrap();
    let raw_data: serde_json::Value = serde_json::from_str(&data_str).unwrap();

    let m = raw_data.as_object().expect("failed to parse dataset");
    let to_vec = |c: &serde_json::Value| {
        c.as_array()
            .unwrap()
            .iter()
            .flat_map(|x| x.as_f64())
            .collect()
    };

    m.iter()
        .map(|(k, v)| {
            let colors = v.as_object().unwrap();
            let luv_data: Vec<f64> = to_vec(&colors["luv"]);
            let lchuv_data: Vec<f64> = to_vec(&colors["lch"]);
            let hsluv_data: Vec<f64> = to_vec(&colors["hsluv"]);
            let xyz_data: Vec<f64> = to_vec(&colors["xyz"]);

            (
                k.clone(),
                HsluvExample {
                    luv: Luv::new(luv_data[0], luv_data[1], luv_data[2]),
                    hsluv: Hsluv::new(hsluv_data[0], hsluv_data[1], hsluv_data[2]),
                    lchuv: Lchuv::new(lchuv_data[0], lchuv_data[1], lchuv_data[2]),
                    xyz: Xyz::new(xyz_data[0], xyz_data[1], xyz_data[2]),
                },
            )
        })
        .collect()
}

lazy_static! {
    static ref TEST_DATA: Examples = load_data();
}

#[test]
pub fn run_xyz_to_luv_tests() {
    for (_, v) in TEST_DATA.iter() {
        let to_luv: Luv<D65, f64> = v.xyz.into_color_unclamped();
        assert_relative_eq!(to_luv, v.luv, epsilon = 0.1);
    }
}

#[test]
pub fn run_luv_to_xyz_tests() {
    for (_, v) in TEST_DATA.iter() {
        let to_xyz: Xyz<D65, f64> = v.luv.into_color_unclamped();
        assert_relative_eq!(to_xyz, v.xyz, epsilon = 0.001);
    }
}

#[test]
pub fn run_lchuv_to_luv_tests() {
    for (_, v) in TEST_DATA.iter() {
        let to_luv: Luv<D65, f64> = v.lchuv.into_color_unclamped();
        assert_relative_eq!(to_luv, v.luv, epsilon = 0.001);
    }
}

#[test]
pub fn run_luv_to_lchuv_tests() {
    for (_, v) in TEST_DATA.iter() {
        let mut to_lchuv: Lchuv<D65, f64> = v.luv.into_color_unclamped();
        if to_lchuv.chroma < 1e-8 {
            to_lchuv.hue = LuvHue::from_degrees(0.0);
        }
        assert_relative_eq!(to_lchuv, v.lchuv, epsilon = 0.001);
    }
}

#[test]
pub fn run_lchuv_to_hsluv_tests() {
    for (_, v) in TEST_DATA.iter() {
        let mut to_hsluv: Hsluv<D65, f64> = v.lchuv.into_color_unclamped();
        if to_hsluv.l > 100.0 - 1e-5 {
            to_hsluv.saturation = 0.0;
        }
        assert_relative_eq!(to_hsluv, v.hsluv, epsilon = 1e-5);
    }
}

#[test]
pub fn run_hsluv_to_lchuv_tests() {
    for (_, v) in TEST_DATA.iter() {
        let to_lchuv: Lchuv<D65, f64> = v.hsluv.into_color_unclamped();
        assert_relative_eq!(to_lchuv, v.lchuv, epsilon = 1e-5);
    }
}
<|endoftext|>