<commit_before>use std::iter;

use hir::Semantics;
use ide_db::{helpers::pick_best_token, RootDatabase};
use itertools::Itertools;
use syntax::{ast, ted, AstNode, NodeOrToken, SyntaxKind, SyntaxNode, WalkEvent, T};

use crate::FilePosition;

pub struct ExpandedMacro {
    pub name: String,
    pub expansion: String,
}

// Feature: Expand Macro Recursively
//
// Shows the full macro expansion of the macro at current cursor.
//
// |===
// | Editor  | Action Name
//
// | VS Code | **Rust Analyzer: Expand macro recursively**
// |===
//
// image::https://user-images.githubusercontent.com/48062697/113020648-b3973180-917a-11eb-84a9-ecb921293dc5.gif[]
pub(crate) fn expand_macro(db: &RootDatabase, position: FilePosition) -> Option<ExpandedMacro> {
    let sema = Semantics::new(db);
    let file = sema.parse(position.file_id);

    let tok = pick_best_token(file.syntax().token_at_offset(position.offset), |kind| match kind {
        SyntaxKind::IDENT => 1,
        _ => 0,
    })?;

    // due to how Rust Analyzer works internally, we need to special case derive attributes,
    // otherwise they might not get found, e.g. here with the cursor at $0 `#[attr]` would expand:
    // ```
    // #[attr]
    // #[derive($0Foo)]
    // struct Bar;
    // ```

    let derive = sema.descend_into_macros(tok.clone()).iter().find_map(|descended| {
        let attr = descended.ancestors().find_map(ast::Attr::cast)?;
        let (path, tt) = attr.as_simple_call()?;
        if path == "derive" {
            let mut tt = tt.syntax().children_with_tokens().skip(1).join("");
            tt.pop();
            let expansions = sema.expand_derive_macro(&attr)?;
            return Some(ExpandedMacro {
                name: tt,
                expansion: expansions.into_iter().map(insert_whitespaces).join(""),
            });
        } else {
            None
        }
    });

    if derive.is_some() {
        return derive;
    }

    // FIXME: Intermix attribute and bang! expansions
    // currently we only recursively expand one of the two types
    let mut expanded = None;
    let mut name = None;
    for node in tok.ancestors() {
        if let Some(item) = ast::Item::cast(node.clone()) {
            if let Some(def) = sema.resolve_attr_macro_call(&item) {
                name = def.name(db).map(|name| name.to_string());
                expanded = expand_attr_macro_recur(&sema, &item);
                break;
            }
        }
        if let Some(mac) = ast::MacroCall::cast(node) {
            name = Some(mac.path()?.segment()?.name_ref()?.to_string());
            expanded = expand_macro_recur(&sema, &mac);
            break;
        }
    }

    // FIXME:
    // macro expansion may lose all white space information
    // But we hope someday we can use ra_fmt for that
    let expansion = insert_whitespaces(expanded?);
    Some(ExpandedMacro { name: name.unwrap_or_else(|| "???".to_owned()), expansion })
}

fn expand_macro_recur(
    sema: &Semantics<RootDatabase>,
    macro_call: &ast::MacroCall,
) -> Option<SyntaxNode> {
    let expanded = sema.expand(macro_call)?.clone_for_update();
    expand(sema, expanded, ast::MacroCall::cast, expand_macro_recur)
}

fn expand_attr_macro_recur(sema: &Semantics<RootDatabase>, item: &ast::Item) -> Option<SyntaxNode> {
    let expanded = sema.expand_attr_macro(item)?.clone_for_update();
    expand(sema, expanded, ast::Item::cast, expand_attr_macro_recur)
}

fn expand<T: AstNode>(
    sema: &Semantics<RootDatabase>,
    expanded: SyntaxNode,
    f: impl FnMut(SyntaxNode) -> Option<T>,
    exp: impl Fn(&Semantics<RootDatabase>, &T) -> Option<SyntaxNode>,
) -> Option<SyntaxNode> {
    let children = expanded.descendants().filter_map(f);
    let mut replacements = Vec::new();

    for child in children {
        if let Some(new_node) = exp(sema, &child) {
            // check if the whole original syntax is replaced
            if expanded == *child.syntax() {
                return Some(new_node);
            }
            replacements.push((child, new_node));
        }
    }

    replacements.into_iter().rev().for_each(|(old, new)| ted::replace(old.syntax(), new));
    Some(expanded)
}

// FIXME: It would also be cool to share logic here and in the mbe tests,
// which are pretty unreadable at the moment.
fn insert_whitespaces(syn: SyntaxNode) -> String {
    use SyntaxKind::*;
    let mut res = String::new();

    let mut indent = 0;
    let mut last: Option<SyntaxKind> = None;

    for event in syn.preorder_with_tokens() {
        let token = match event {
            WalkEvent::Enter(NodeOrToken::Token(token)) => token,
            WalkEvent::Leave(NodeOrToken::Node(node))
                if matches!(node.kind(), ATTR | MATCH_ARM | STRUCT | ENUM | UNION | FN | IMPL) =>
            {
                res.push('\n');
                res.extend(iter::repeat(" ").take(2 * indent));
                continue;
            }
            _ => continue,
        };
        let is_next = |f: fn(SyntaxKind) -> bool, default| -> bool {
            token.next_token().map(|it| f(it.kind())).unwrap_or(default)
        };
        let is_last =
            |f: fn(SyntaxKind) -> bool, default| -> bool { last.map(f).unwrap_or(default) };

        match token.kind() {
            k if is_text(k) && is_next(|it| !it.is_punct(), true) => {
                res.push_str(token.text());
                res.push(' ');
            }
            L_CURLY if is_next(|it| it != R_CURLY, true) => {
                indent += 1;
                if is_last(is_text, false) {
                    res.push(' ');
                }
                res.push_str("{\n");
                res.extend(iter::repeat(" ").take(2 * indent));
            }
            R_CURLY if is_last(|it| it != L_CURLY, true) => {
                indent = indent.saturating_sub(1);
                res.push('\n');
                res.extend(iter::repeat(" ").take(2 * indent));
                res.push_str("}");
            }
            R_CURLY => {
                res.push_str("}\n");
                res.extend(iter::repeat(" ").take(2 * indent));
            }
            LIFETIME_IDENT if is_next(|it| it == IDENT || it == MUT_KW, true) => {
                res.push_str(token.text());
                res.push(' ');
            }
            T![;] => {
                res.push_str(";\n");
                res.extend(iter::repeat(" ").take(2 * indent));
            }
            T![->] => res.push_str(" -> "),
            T![=] => res.push_str(" = "),
            T![=>] => res.push_str(" => "),
            _ => res.push_str(token.text()),
        }

        last = Some(token.kind());
    }

    return res;

    fn is_text(k: SyntaxKind) -> bool {
        k.is_keyword() || k.is_literal() || k == IDENT
    }
}

#[cfg(test)]
mod tests {
    use expect_test::{expect, Expect};

    use crate::fixture;

    #[track_caller]
    fn check(ra_fixture: &str, expect: Expect) {
        let (analysis, pos) = fixture::position(ra_fixture);
        let expansion = analysis.expand_macro(pos).unwrap().unwrap();
        let actual = format!("{}\n{}", expansion.name, expansion.expansion);
        expect.assert_eq(&actual);
    }

    #[test]
    fn macro_expand_recursive_expansion() {
        check(
            r#"
macro_rules! bar {
    () => { fn  b() {} }
}
macro_rules! foo {
    () => { bar!(); }
}
macro_rules! baz {
    () => { foo!(); }
}
f$0oo!();
"#,
            expect![[r#"
                foo
                fn b(){}

            "#]],
        );
    }

    #[test]
    fn macro_expand_multiple_lines() {
        check(
            r#"
macro_rules! foo {
    () => {
        fn some_thing() -> u32 {
            let a = 0;
            a + 10
        }
    }
}
f$0oo!();
        "#,
            expect![[r#"
                foo
                fn some_thing() -> u32 {
                  let a = 0;
                  a+10
                }
            "#]],
        );
    }

    #[test]
    fn macro_expand_match_ast() {
        check(
            r#"
macro_rules! match_ast {
    (match $node:ident { $($tt:tt)* }) => { match_ast!(match ($node) { $($tt)* }) };
    (match ($node:expr) {
        $( ast::$ast:ident($it:ident) => $res:block, )*
        _ => $catch_all:expr $(,)?
    }) => {{
        $( if let Some($it) = ast::$ast::cast($node.clone()) $res else )*
        { $catch_all }
    }};
}

fn main() {
    mat$0ch_ast! {
        match container {
            ast::TraitDef(it) => {},
            ast::ImplDef(it) => {},
            _ => { continue },
        }
    }
}
"#,
            expect![[r#"
       match_ast
       {
         if let Some(it) = ast::TraitDef::cast(container.clone()){}
         else if let Some(it) = ast::ImplDef::cast(container.clone()){}
         else {
           {
             continue
           }
         }
       }"#]],
        );
    }

    #[test]
    fn macro_expand_match_ast_inside_let_statement() {
        check(
            r#"
macro_rules! match_ast {
    (match $node:ident { $($tt:tt)* }) => { match_ast!(match ($node) { $($tt)* }) };
    (match ($node:expr) {}) => {{}};
}

fn main() {
    let p = f(|it| {
        let res = mat$0ch_ast! { match c {}};
        Some(res)
    })?;
}
"#,
            expect![[r#"
                match_ast
                {}
            "#]],
        );
    }

    #[test]
    fn macro_expand_inner_macro_fail_to_expand() {
        check(
            r#"
macro_rules! bar {
    (BAD) => {};
}
macro_rules! foo {
    () => {bar!()};
}

fn main() {
    let res = fo$0o!();
}
"#,
            expect![[r#"
                foo
            "#]],
        );
    }

    #[test]
    fn macro_expand_with_dollar_crate() {
        check(
            r#"
#[macro_export]
macro_rules! bar {
    () => {0};
}
macro_rules! foo {
    () => {$crate::bar!()};
}

fn main() {
    let res = fo$0o!();
}
"#,
            expect![[r#"
                foo
                0 "#]],
        );
    }

    #[test]
    fn macro_expand_derive() {
        check(
            r#"
#[rustc_builtin_macro]
pub macro Clone {}

#[doc = ""]
#[derive(C$0lone)]
struct Foo {}
"#,
            expect![[r#"
                Clone
                impl< >crate::clone::Clone for Foo< >{}

            "#]],
        );
    }

    #[test]
    fn macro_expand_derive2() {
        check(
            r#"
#[rustc_builtin_macro]
pub macro Clone {}
#[rustc_builtin_macro]
pub macro Copy {}

#[derive(Cop$0y)]
#[derive(Clone)]
struct Foo {}
"#,
            expect![[r#"
                Copy
                impl< >crate::marker::Copy for Foo< >{}

            "#]],
        );
    }

    #[test]
    fn macro_expand_derive_multi() {
        check(
            r#"
#[rustc_builtin_macro]
pub macro Clone {}
#[rustc_builtin_macro]
pub macro Copy {}

#[derive(Cop$0y, Clone)]
struct Foo {}
"#,
            expect![[r#"
                Copy, Clone
                impl< >crate::marker::Copy for Foo< >{}

                impl< >crate::clone::Clone for Foo< >{}

            "#]],
        );
    }
}
<commit_msg>fix test<commit_after>use std::iter;

use hir::Semantics;
use ide_db::{helpers::pick_best_token, RootDatabase};
use itertools::Itertools;
use syntax::{ast, ted, AstNode, NodeOrToken, SyntaxKind, SyntaxNode, WalkEvent, T};

use crate::FilePosition;

pub struct ExpandedMacro {
    pub name: String,
    pub expansion: String,
}

// Feature: Expand Macro Recursively
//
// Shows the full macro expansion of the macro at current cursor.
//
// |===
// | Editor  | Action Name
//
// | VS Code | **Rust Analyzer: Expand macro recursively**
// |===
//
// image::https://user-images.githubusercontent.com/48062697/113020648-b3973180-917a-11eb-84a9-ecb921293dc5.gif[]
pub(crate) fn expand_macro(db: &RootDatabase, position: FilePosition) -> Option<ExpandedMacro> {
    let sema = Semantics::new(db);
    let file = sema.parse(position.file_id);

    let tok = pick_best_token(file.syntax().token_at_offset(position.offset), |kind| match kind {
        SyntaxKind::IDENT => 1,
        _ => 0,
    })?;

    // due to how Rust Analyzer works internally, we need to special case derive attributes,
    // otherwise they might not get found, e.g. here with the cursor at $0 `#[attr]` would expand:
    // ```
    // #[attr]
    // #[derive($0Foo)]
    // struct Bar;
    // ```

    let derive = sema.descend_into_macros(tok.clone()).iter().find_map(|descended| {
        let attr = descended.ancestors().find_map(ast::Attr::cast)?;
        let (path, tt) = attr.as_simple_call()?;
        if path == "derive" {
            let mut tt = tt.syntax().children_with_tokens().skip(1).join("");
            tt.pop();
            let expansions = sema.expand_derive_macro(&attr)?;
            return Some(ExpandedMacro {
                name: tt,
                expansion: expansions.into_iter().map(insert_whitespaces).join(""),
            });
        } else {
            None
        }
    });

    if derive.is_some() {
        return derive;
    }

    // FIXME: Intermix attribute and bang! expansions
    // currently we only recursively expand one of the two types
    let mut expanded = None;
    let mut name = None;
    for node in tok.ancestors() {
        if let Some(item) = ast::Item::cast(node.clone()) {
            if let Some(def) = sema.resolve_attr_macro_call(&item) {
                name = def.name(db).map(|name| name.to_string());
                expanded = expand_attr_macro_recur(&sema, &item);
                break;
            }
        }
        if let Some(mac) = ast::MacroCall::cast(node) {
            name = Some(mac.path()?.segment()?.name_ref()?.to_string());
            expanded = expand_macro_recur(&sema, &mac);
            break;
        }
    }

    // FIXME:
    // macro expansion may lose all white space information
    // But we hope someday we can use ra_fmt for that
    let expansion = insert_whitespaces(expanded?);
    Some(ExpandedMacro { name: name.unwrap_or_else(|| "???".to_owned()), expansion })
}

fn expand_macro_recur(
    sema: &Semantics<RootDatabase>,
    macro_call: &ast::MacroCall,
) -> Option<SyntaxNode> {
    let expanded = sema.expand(macro_call)?.clone_for_update();
    expand(sema, expanded, ast::MacroCall::cast, expand_macro_recur)
}

fn expand_attr_macro_recur(sema: &Semantics<RootDatabase>, item: &ast::Item) -> Option<SyntaxNode> {
    let expanded = sema.expand_attr_macro(item)?.clone_for_update();
    expand(sema, expanded, ast::Item::cast, expand_attr_macro_recur)
}

fn expand<T: AstNode>(
    sema: &Semantics<RootDatabase>,
    expanded: SyntaxNode,
    f: impl FnMut(SyntaxNode) -> Option<T>,
    exp: impl Fn(&Semantics<RootDatabase>, &T) -> Option<SyntaxNode>,
) -> Option<SyntaxNode> {
    let children = expanded.descendants().filter_map(f);
    let mut replacements = Vec::new();

    for child in children {
        if let Some(new_node) = exp(sema, &child) {
            // check if the whole original syntax is replaced
            if expanded == *child.syntax() {
                return Some(new_node);
            }
            replacements.push((child, new_node));
        }
    }

    replacements.into_iter().rev().for_each(|(old, new)| ted::replace(old.syntax(), new));
    Some(expanded)
}

// FIXME: It would also be cool to share logic here and in the mbe tests,
// which are pretty unreadable at the moment.
fn insert_whitespaces(syn: SyntaxNode) -> String {
    use SyntaxKind::*;
    let mut res = String::new();

    let mut indent = 0;
    let mut last: Option<SyntaxKind> = None;

    for event in syn.preorder_with_tokens() {
        let token = match event {
            WalkEvent::Enter(NodeOrToken::Token(token)) => token,
            WalkEvent::Leave(NodeOrToken::Node(node))
                if matches!(node.kind(), ATTR | MATCH_ARM | STRUCT | ENUM | UNION | FN | IMPL) =>
            {
                res.push('\n');
                res.extend(iter::repeat(" ").take(2 * indent));
                continue;
            }
            _ => continue,
        };
        let is_next = |f: fn(SyntaxKind) -> bool, default| -> bool {
            token.next_token().map(|it| f(it.kind())).unwrap_or(default)
        };
        let is_last =
            |f: fn(SyntaxKind) -> bool, default| -> bool { last.map(f).unwrap_or(default) };

        match token.kind() {
            k if is_text(k) && is_next(|it| !it.is_punct(), true) => {
                res.push_str(token.text());
                res.push(' ');
            }
            L_CURLY if is_next(|it| it != R_CURLY, true) => {
                indent += 1;
                if is_last(is_text, false) {
                    res.push(' ');
                }
                res.push_str("{\n");
                res.extend(iter::repeat(" ").take(2 * indent));
            }
            R_CURLY if is_last(|it| it != L_CURLY, true) => {
                indent = indent.saturating_sub(1);
                res.push('\n');
                res.extend(iter::repeat(" ").take(2 * indent));
                res.push_str("}");
            }
            R_CURLY => {
                res.push_str("}\n");
                res.extend(iter::repeat(" ").take(2 * indent));
            }
            LIFETIME_IDENT if is_next(|it| it == IDENT || it == MUT_KW, true) => {
                res.push_str(token.text());
                res.push(' ');
            }
            T![;] => {
                res.push_str(";\n");
                res.extend(iter::repeat(" ").take(2 * indent));
            }
            T![->] => res.push_str(" -> "),
            T![=] => res.push_str(" = "),
            T![=>] => res.push_str(" => "),
            _ => res.push_str(token.text()),
        }

        last = Some(token.kind());
    }

    return res;

    fn is_text(k: SyntaxKind) -> bool {
        k.is_keyword() || k.is_literal() || k == IDENT
    }
}

#[cfg(test)]
mod tests {
    use expect_test::{expect, Expect};

    use crate::fixture;

    #[track_caller]
    fn check(ra_fixture: &str, expect: Expect) {
        let (analysis, pos) = fixture::position(ra_fixture);
        let expansion = analysis.expand_macro(pos).unwrap().unwrap();
        let actual = format!("{}\n{}", expansion.name, expansion.expansion);
        expect.assert_eq(&actual);
    }

    #[test]
    fn macro_expand_recursive_expansion() {
        check(
            r#"
macro_rules! bar {
    () => { fn  b() {} }
}
macro_rules! foo {
    () => { bar!(); }
}
macro_rules! baz {
    () => { foo!(); }
}
f$0oo!();
"#,
            expect![[r#"
                foo
                fn b(){}

            "#]],
        );
    }

    #[test]
    fn macro_expand_multiple_lines() {
        check(
            r#"
macro_rules! foo {
    () => {
        fn some_thing() -> u32 {
            let a = 0;
            a + 10
        }
    }
}
f$0oo!();
        "#,
            expect![[r#"
                foo
                fn some_thing() -> u32 {
                  let a = 0;
                  a+10
                }
            "#]],
        );
    }

    #[test]
    fn macro_expand_match_ast() {
        check(
            r#"
macro_rules! match_ast {
    (match $node:ident { $($tt:tt)* }) => { match_ast!(match ($node) { $($tt)* }) };
    (match ($node:expr) {
        $( ast::$ast:ident($it:ident) => $res:block, )*
        _ => $catch_all:expr $(,)?
    }) => {{
        $( if let Some($it) = ast::$ast::cast($node.clone()) $res else )*
        { $catch_all }
    }};
}

fn main() {
    mat$0ch_ast! {
        match container {
            ast::TraitDef(it) => {},
            ast::ImplDef(it) => {},
            _ => { continue },
        }
    }
}
"#,
            expect![[r#"
       match_ast
       {
         if let Some(it) = ast::TraitDef::cast(container.clone()){}
         else if let Some(it) = ast::ImplDef::cast(container.clone()){}
         else {
           {
             continue
           }
         }
       }"#]],
        );
    }

    #[test]
    fn macro_expand_match_ast_inside_let_statement() {
        check(
            r#"
macro_rules! match_ast {
    (match $node:ident { $($tt:tt)* }) => { match_ast!(match ($node) { $($tt)* }) };
    (match ($node:expr) {}) => {{}};
}

fn main() {
    let p = f(|it| {
        let res = mat$0ch_ast! { match c {}};
        Some(res)
    })?;
}
"#,
            expect![[r#"
                match_ast
                {}
            "#]],
        );
    }

    #[test]
    fn macro_expand_inner_macro_fail_to_expand() {
        check(
            r#"
macro_rules! bar {
    (BAD) => {};
}
macro_rules! foo {
    () => {bar!()};
}

fn main() {
    let res = fo$0o!();
}
"#,
            expect![[r#"
                foo
            "#]],
        );
    }

    #[test]
    fn macro_expand_with_dollar_crate() {
        check(
            r#"
#[macro_export]
macro_rules! bar {
    () => {0};
}
macro_rules! foo {
    () => {$crate::bar!()};
}

fn main() {
    let res = fo$0o!();
}
"#,
            expect![[r#"
                foo
                0 "#]],
        );
    }

    #[test]
    fn macro_expand_derive() {
        check(
            r#"
//- proc_macros: identity

#[rustc_builtin_macro]
pub macro Clone {}

#[proc_macros::identity]
#[derive(C$0lone)]
struct Foo {}
"#,
            expect![[r#"
                Clone
                impl< >crate::clone::Clone for Foo< >{}

            "#]],
        );
    }

    #[test]
    fn macro_expand_derive2() {
        check(
            r#"
#[rustc_builtin_macro]
pub macro Clone {}
#[rustc_builtin_macro]
pub macro Copy {}

#[derive(Cop$0y)]
#[derive(Clone)]
struct Foo {}
"#,
            expect![[r#"
                Copy
                impl< >crate::marker::Copy for Foo< >{}

            "#]],
        );
    }

    #[test]
    fn macro_expand_derive_multi() {
        check(
            r#"
#[rustc_builtin_macro]
pub macro Clone {}
#[rustc_builtin_macro]
pub macro Copy {}

#[derive(Cop$0y, Clone)]
struct Foo {}
"#,
            expect![[r#"
                Copy, Clone
                impl< >crate::marker::Copy for Foo< >{}

                impl< >crate::clone::Clone for Foo< >{}

            "#]],
        );
    }
}
<|endoftext|>


<commit_before>use std::fmt;

/// An 8 bit register. This struct simply wraps an u8 value, it is also meant to offer simple
/// arithmetic operations with wrapping in case of overflows.
pub struct Register8 {
    value: u8,
}

impl fmt::Debug for Register8 {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "0x{:02X}", self.value)
    }
}

impl Register8 {
    pub fn new(v: u8) -> Register8 {
        Register8 { value: v }
    }

    pub fn get(&self) -> u8 {
        self.value
    }

    pub fn set(&mut self, value: u8) {
        self.value = value;
    }

    /// adds the value with wrapping semantics
    pub fn add(&mut self, value: u8) {
        self.value.wrapping_add(value);
    }

    /// adds the value with wrapping semantics
    pub fn sub(&mut self, value: u8) {
        self.value.wrapping_sub(value);
    }
}

/// A 16 bit register. It is composed of 2 Register8 structs, making the high and low parts of the
/// 16 bits. The struct also offers the ability to get/set individual 8 bit registers or act as a
/// whole u16 number.
pub struct Register16 {
    pub low: Register8,
    pub high: Register8,
}

impl fmt::Debug for Register16 {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "0x{:04X}", self.get())
    }
}

impl Register16 {
    pub fn new(v: u16) -> Register16 {
        let low = (v & 0x00FF) as u8;
        let high = ((v & 0xFF00) >> 8) as u8;

        Register16 {
            low: Register8::new(low),
            high: Register8::new(high),
        }
    }

    /// adds the value with wrapping semantics
    pub fn add(&mut self, value: u16) {
        let new_value = self.get().wrapping_add(value);
        self.set(new_value);
    }

    /// adds the value with wrapping semantics
    pub fn sub(&mut self, value: u16) {
        let new_value = self.get().wrapping_sub(value);
        self.set(new_value);
    }

    pub fn get(&self) -> u16 {
        ((self.high.get() as u16) << 8) | self.low.get() as u16
    }

    pub fn set(&mut self, v: u16) {
        let low = (v & 0x00FF) as u8;
        let high = (v >> 8) as u8;

        self.low.set(low);
        self.high.set(high);
    }
}

/// The main registers in a gameboy CPU.
/// Registers are 16 bit wide and can be accessed as L(low) and H(high).
/// Register AF is used for (A) accumulator and (F) flags
/// This struct also offers convenience methods to get and set individual flags.
#[derive(Debug)]
pub struct Registers {
    pub af: Register16,
    pub bc: Register16,
    pub de: Register16,
    pub hl: Register16,

    pub pc: Register16, // program counter
    pub sp: Register16, // stack pointer
}

impl Registers {
    /// true if bit 7 of flags is 1
    pub fn zero_flag(&self) -> bool {
        (self.af.low.get() & 0x80) == 0x80
    }

    /// true if bit 6 of flags is 1
    pub fn add_sub_flag(&self) -> bool {
        (self.af.low.get() & 0x40) == 0x40
    }

    /// true if bit 5 of flags is 1
    pub fn half_carry_flag(&self) -> bool {
        (self.af.low.get() & 0x20) == 0x20
    }

    /// true if bit 4 of flags is 1
    pub fn carry_flag(&self) -> bool {
        (self.af.low.get() & 0x10) == 0x10
    }

    /// sets bit 7 of flags to the specified bool value
    pub fn set_zero_flag(&mut self, flag: bool) {
        let f = self.af.low.get();
        if flag {
            self.af.low.set(f | 0x80);
        } else {
            self.af.low.set(f & 0x7F);
        }
    }

    /// sets bit 6 of flags to the specified bool value
    pub fn set_add_sub_flag(&mut self, flag: bool) {
        let f = {
            self.af.low.get()
        };
        if flag {
            self.af.low.set(f | 0x40);
        } else {
            self.af.low.set(f & 0xBF);
        }
    }

    /// sets bit 5 of flags to the specified bool value
    pub fn set_half_carry_flag(&mut self, flag: bool) {
        let f = {
            self.af.low.get()
        };
        if flag {
            self.af.low.set(f | 0x20);
        } else {
            self.af.low.set(f & 0xDF);
        }
    }

    /// sets bit 4 of flags to the specified bool value
    pub fn set_carry_flag(&mut self, flag: bool) {
        let f = {
            self.af.low.get()
        };
        if flag {
            self.af.low.set(f | 0x10);
        } else {
            self.af.low.set(f & 0xEF);
        }
    }
}
<commit_msg>Improved flags implementation<commit_after>use std::fmt;

/// An 8 bit register. This struct simply wraps an u8 value, it is also meant to offer simple
/// arithmetic operations with wrapping in case of overflows.
pub struct Register8 {
    value: u8,
}

impl fmt::Debug for Register8 {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "0x{:02X}", self.value)
    }
}

impl Register8 {
    pub fn new(v: u8) -> Register8 {
        Register8 { value: v }
    }

    pub fn get(&self) -> u8 {
        self.value
    }

    pub fn set(&mut self, value: u8) {
        self.value = value;
    }

    /// adds the value with wrapping semantics
    pub fn add(&mut self, value: u8) {
        self.value.wrapping_add(value);
    }

    /// adds the value with wrapping semantics
    pub fn sub(&mut self, value: u8) {
        self.value.wrapping_sub(value);
    }
}

/// A 16 bit register. It is composed of 2 Register8 structs, making the high and low parts of the
/// 16 bits. The struct also offers the ability to get/set individual 8 bit registers or act as a
/// whole u16 number.
pub struct Register16 {
    pub low: Register8,
    pub high: Register8,
}

impl fmt::Debug for Register16 {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "0x{:04X}", self.get())
    }
}

impl Register16 {
    pub fn new(v: u16) -> Register16 {
        let low = (v & 0x00FF) as u8;
        let high = ((v & 0xFF00) >> 8) as u8;

        Register16 {
            low: Register8::new(low),
            high: Register8::new(high),
        }
    }

    /// adds the value with wrapping semantics
    pub fn add(&mut self, value: u16) {
        let new_value = self.get().wrapping_add(value);
        self.set(new_value);
    }

    /// adds the value with wrapping semantics
    pub fn sub(&mut self, value: u16) {
        let new_value = self.get().wrapping_sub(value);
        self.set(new_value);
    }

    pub fn get(&self) -> u16 {
        ((self.high.get() as u16) << 8) | self.low.get() as u16
    }

    pub fn set(&mut self, v: u16) {
        let low = (v & 0x00FF) as u8;
        let high = (v >> 8) as u8;

        self.low.set(low);
        self.high.set(high);
    }
}

/// The main registers in a gameboy CPU.
/// Registers are 16 bit wide and can be accessed as L(low) and H(high).
/// Register AF is used for (A) accumulator and (F) flags.
/// This struct also offers convenience methods to get and set individual flags.
#[derive(Debug)]
pub struct Registers {
    pub af: Register16,
    pub bc: Register16,
    pub de: Register16,
    pub hl: Register16,

    pub pc: Register16, // program counter
    pub sp: Register16, // stack pointer
}

/// The four flags and their respective bit values. Bits 0-3 are unused.
pub enum Flags {
    /// This bit is set when the result of a math operation
    /// is zero or two values match when using the CP
    /// instruction.
    Zero        = 0b10000000,

    /// This bit is set if a subtraction was performed in the
    /// last math instruction.
    Sub         = 0b01000000,

    /// This bit is set if a carry occurred from the lower
    /// nibble in the last math operation.
    Carry       = 0b00100000,

    /// This bit is set if a carry occurred from the last
    /// math operation or if register A is the smaller value
    /// when executing the CP instruction.
    HalfCarry   = 0b00010000,
}

impl Registers {

    /// Clears all flag values by resetting the F register to 0.
    pub fn clear_all_flags(&mut self) {
        self.af.low.set(0);
    }

    /// Sets the selected flag to 1.
    pub fn set_flag(&mut self, flag: Flags) {
        let new_flags = self.af.low.get() | (flag as u8);
        self.af.low.set(new_flags);
    }

    /// Clears the selected flag.
    pub fn clear_flag(&mut self, flag: Flags) {
        let new_flags = self.af.low.get() & !(flag as u8);
        self.af.low.set(new_flags);
    }

    /// Returns true if the selected flag is set.
    pub fn is_set(&self, flag: Flags) -> bool {
        let flag_value = flag as u8;
        (self.af.low.get() & flag_value) == flag_value
    }
}
<|endoftext|>


<commit_before>// Copyright 2013-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OsRng;

#[cfg(all(unix, not(target_os = "ios")))]
mod imp {
    extern crate libc;

    use self::OsRngInner::*;

    use io::{IoResult, File};
    use path::Path;
    use rand::Rng;
    use rand::reader::ReaderRng;
    use result::Result::{Ok, Err};
    use slice::SlicePrelude;
    use mem;
    use os::errno;

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm")))]
    fn getrandom(buf: &mut [u8]) -> libc::c_long {
        extern "C" {
            fn syscall(number: libc::c_long, ...) -> libc::c_long;
        }

        #[cfg(target_arch = "x86_64")]
        const NR_GETRANDOM: libc::c_long = 318;
        #[cfg(target_arch = "x86")]
        const NR_GETRANDOM: libc::c_long = 355;
        #[cfg(target_arch = "arm")]
        const NR_GETRANDOM: libc::c_long = 384;

        unsafe {
            syscall(NR_GETRANDOM, buf.as_mut_ptr(), buf.len(), 0u)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm"))))]
    fn getrandom(_buf: &mut [u8]) -> libc::c_long { -1 }

    fn getrandom_fill_bytes(v: &mut [u8]) {
        let mut read = 0;
        let len = v.len();
        while read < len {
            let result = getrandom(v[mut read..]);
            if result == -1 {
                let err = errno() as libc::c_int;
                if err == libc::EINTR {
                    continue;
                } else {
                    panic!("unexpected getrandom error: {}", err);
                }
            } else {
                read += result as uint;
            }
        }
    }

    fn getrandom_next_u32() -> u32 {
        let mut buf: [u8, ..4] = [0u8, ..4];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8, ..4], u32>(buf) }
    }

    fn getrandom_next_u64() -> u64 {
        let mut buf: [u8, ..8] = [0u8, ..8];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8, ..8], u64>(buf) }
    }

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm")))]
    fn is_getrandom_available() -> bool {
        use sync::atomic::{AtomicBool, INIT_ATOMIC_BOOL, Relaxed};

        static GETRANDOM_CHECKED: AtomicBool = INIT_ATOMIC_BOOL;
        static GETRANDOM_AVAILABLE: AtomicBool = INIT_ATOMIC_BOOL;

        if !GETRANDOM_CHECKED.load(Relaxed) {
            let mut buf: [u8, ..0] = [];
            let result = getrandom(&mut buf);
            let available = if result == -1 {
                let err = errno() as libc::c_int;
                err != libc::ENOSYS
            } else {
                true
            };
            GETRANDOM_AVAILABLE.store(available, Relaxed);
            GETRANDOM_CHECKED.store(true, Relaxed);
            available
        } else {
            GETRANDOM_AVAILABLE.load(Relaxed)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm"))))]
    fn is_getrandom_available() -> bool { false }

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed
    /// This does not block.
    pub struct OsRng {
        inner: OsRngInner,
    }

    enum OsRngInner {
        OsGetrandomRng,
        OsReaderRng(ReaderRng<File>),
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            if is_getrandom_available() {
                return Ok(OsRng { inner: OsGetrandomRng });
            }

            let reader = try!(File::open(&Path::new("/dev/urandom")));
            let reader_rng = ReaderRng::new(reader);

            Ok(OsRng { inner: OsReaderRng(reader_rng) })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u32(),
                OsReaderRng(ref mut rng) => rng.next_u32(),
            }
        }
        fn next_u64(&mut self) -> u64 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u64(),
                OsReaderRng(ref mut rng) => rng.next_u64(),
            }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            match self.inner {
                OsGetrandomRng => getrandom_fill_bytes(v),
                OsReaderRng(ref mut rng) => rng.fill_bytes(v)
            }
        }
    }
}

#[cfg(target_os = "ios")]
mod imp {
    extern crate libc;

    use io::{IoResult};
    use kinds::marker;
    use mem;
    use os;
    use rand::Rng;
    use result::Result::{Ok};
    use self::libc::{c_int, size_t};
    use slice::{SlicePrelude};

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed
    /// This does not block.
    pub struct OsRng {
        marker: marker::NoCopy
    }

    #[repr(C)]
    struct SecRandom;

    #[allow(non_upper_case_globals)]
    static kSecRandomDefault: *const SecRandom = 0 as *const SecRandom;

    #[link(name = "Security", kind = "framework")]
    extern "C" {
        fn SecRandomCopyBytes(rnd: *const SecRandom,
                              count: size_t, bytes: *mut u8) -> c_int;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            Ok(OsRng {marker: marker::NoCopy} )
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                SecRandomCopyBytes(kSecRandomDefault, v.len() as size_t, v.as_mut_ptr())
            };
            if ret == -1 {
                panic!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }
}

#[cfg(windows)]
mod imp {
    extern crate libc;

    use io::{IoResult, IoError};
    use mem;
    use ops::Drop;
    use os;
    use rand::Rng;
    use result::Result::{Ok, Err};
    use self::libc::{DWORD, BYTE, LPCSTR, BOOL};
    use self::libc::types::os::arch::extra::{LONG_PTR};
    use slice::{SlicePrelude};

    type HCRYPTPROV = LONG_PTR;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed
    /// This does not block.
    pub struct OsRng {
        hcryptprov: HCRYPTPROV
    }

    static PROV_RSA_FULL: DWORD = 1;
    static CRYPT_SILENT: DWORD = 64;
    static CRYPT_VERIFYCONTEXT: DWORD = 0xF0000000;

    #[allow(non_snake_case)]
    extern "system" {
        fn CryptAcquireContextA(phProv: *mut HCRYPTPROV,
                                pszContainer: LPCSTR,
                                pszProvider: LPCSTR,
                                dwProvType: DWORD,
                                dwFlags: DWORD) -> BOOL;
        fn CryptGenRandom(hProv: HCRYPTPROV,
                          dwLen: DWORD,
                          pbBuffer: *mut BYTE) -> BOOL;
        fn CryptReleaseContext(hProv: HCRYPTPROV, dwFlags: DWORD) -> BOOL;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            let mut hcp = 0;
            let ret = unsafe {
                CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            if ret == 0 {
                Err(IoError::last_error())
            } else {
                Ok(OsRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                CryptGenRandom(self.hcryptprov, v.len() as DWORD,
                               v.as_mut_ptr())
            };
            if ret == 0 {
                panic!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }

    impl Drop for OsRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                panic!("couldn't release context: {}", os::last_os_error());
            }
        }
    }
}

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

    use super::OsRng;
    use rand::Rng;
    use task;

    #[test]
    fn test_os_rng() {
        let mut r = OsRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8, .. 1000];
        r.fill_bytes(&mut v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in range(0u, 20) {
            let (tx, rx) = channel();
            txs.push(tx);
            task::spawn(move|| {
                // wait until all the tasks are ready to go.
                rx.recv();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OsRng::new().unwrap();
                task::deschedule();
                let mut v = [0u8, .. 1000];

                for _ in range(0u, 100) {
                    r.next_u32();
                    task::deschedule();
                    r.next_u64();
                    task::deschedule();
                    r.fill_bytes(&mut v);
                    task::deschedule();
                }
            })
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(())
        }
    }
}
<commit_msg>Remove internal uses of `marker::NoCopy`<commit_after>// Copyright 2013-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OsRng;

#[cfg(all(unix, not(target_os = "ios")))]
mod imp {
    extern crate libc;

    use self::OsRngInner::*;

    use io::{IoResult, File};
    use path::Path;
    use rand::Rng;
    use rand::reader::ReaderRng;
    use result::Result::{Ok, Err};
    use slice::SlicePrelude;
    use mem;
    use os::errno;

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm")))]
    fn getrandom(buf: &mut [u8]) -> libc::c_long {
        extern "C" {
            fn syscall(number: libc::c_long, ...) -> libc::c_long;
        }

        #[cfg(target_arch = "x86_64")]
        const NR_GETRANDOM: libc::c_long = 318;
        #[cfg(target_arch = "x86")]
        const NR_GETRANDOM: libc::c_long = 355;
        #[cfg(target_arch = "arm")]
        const NR_GETRANDOM: libc::c_long = 384;

        unsafe {
            syscall(NR_GETRANDOM, buf.as_mut_ptr(), buf.len(), 0u)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm"))))]
    fn getrandom(_buf: &mut [u8]) -> libc::c_long { -1 }

    fn getrandom_fill_bytes(v: &mut [u8]) {
        let mut read = 0;
        let len = v.len();
        while read < len {
            let result = getrandom(v[mut read..]);
            if result == -1 {
                let err = errno() as libc::c_int;
                if err == libc::EINTR {
                    continue;
                } else {
                    panic!("unexpected getrandom error: {}", err);
                }
            } else {
                read += result as uint;
            }
        }
    }

    fn getrandom_next_u32() -> u32 {
        let mut buf: [u8, ..4] = [0u8, ..4];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8, ..4], u32>(buf) }
    }

    fn getrandom_next_u64() -> u64 {
        let mut buf: [u8, ..8] = [0u8, ..8];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8, ..8], u64>(buf) }
    }

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm")))]
    fn is_getrandom_available() -> bool {
        use sync::atomic::{AtomicBool, INIT_ATOMIC_BOOL, Relaxed};

        static GETRANDOM_CHECKED: AtomicBool = INIT_ATOMIC_BOOL;
        static GETRANDOM_AVAILABLE: AtomicBool = INIT_ATOMIC_BOOL;

        if !GETRANDOM_CHECKED.load(Relaxed) {
            let mut buf: [u8, ..0] = [];
            let result = getrandom(&mut buf);
            let available = if result == -1 {
                let err = errno() as libc::c_int;
                err != libc::ENOSYS
            } else {
                true
            };
            GETRANDOM_AVAILABLE.store(available, Relaxed);
            GETRANDOM_CHECKED.store(true, Relaxed);
            available
        } else {
            GETRANDOM_AVAILABLE.load(Relaxed)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm"))))]
    fn is_getrandom_available() -> bool { false }

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed
    /// This does not block.
    pub struct OsRng {
        inner: OsRngInner,
    }

    enum OsRngInner {
        OsGetrandomRng,
        OsReaderRng(ReaderRng<File>),
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            if is_getrandom_available() {
                return Ok(OsRng { inner: OsGetrandomRng });
            }

            let reader = try!(File::open(&Path::new("/dev/urandom")));
            let reader_rng = ReaderRng::new(reader);

            Ok(OsRng { inner: OsReaderRng(reader_rng) })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u32(),
                OsReaderRng(ref mut rng) => rng.next_u32(),
            }
        }
        fn next_u64(&mut self) -> u64 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u64(),
                OsReaderRng(ref mut rng) => rng.next_u64(),
            }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            match self.inner {
                OsGetrandomRng => getrandom_fill_bytes(v),
                OsReaderRng(ref mut rng) => rng.fill_bytes(v)
            }
        }
    }
}

#[cfg(target_os = "ios")]
mod imp {
    extern crate libc;

    use io::{IoResult};
    use kinds::marker;
    use mem;
    use os;
    use rand::Rng;
    use result::Result::{Ok};
    use self::libc::{c_int, size_t};
    use slice::{SlicePrelude};

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed
    /// This does not block.
    #[allow(missing_copy_implementations)]
    pub struct OsRng;

    #[repr(C)]
    struct SecRandom;

    #[allow(non_upper_case_globals)]
    static kSecRandomDefault: *const SecRandom = 0 as *const SecRandom;

    #[link(name = "Security", kind = "framework")]
    extern "C" {
        fn SecRandomCopyBytes(rnd: *const SecRandom,
                              count: size_t, bytes: *mut u8) -> c_int;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            Ok(OsRng)
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                SecRandomCopyBytes(kSecRandomDefault, v.len() as size_t, v.as_mut_ptr())
            };
            if ret == -1 {
                panic!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }
}

#[cfg(windows)]
mod imp {
    extern crate libc;

    use io::{IoResult, IoError};
    use mem;
    use ops::Drop;
    use os;
    use rand::Rng;
    use result::Result::{Ok, Err};
    use self::libc::{DWORD, BYTE, LPCSTR, BOOL};
    use self::libc::types::os::arch::extra::{LONG_PTR};
    use slice::{SlicePrelude};

    type HCRYPTPROV = LONG_PTR;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed
    /// This does not block.
    pub struct OsRng {
        hcryptprov: HCRYPTPROV
    }

    static PROV_RSA_FULL: DWORD = 1;
    static CRYPT_SILENT: DWORD = 64;
    static CRYPT_VERIFYCONTEXT: DWORD = 0xF0000000;

    #[allow(non_snake_case)]
    extern "system" {
        fn CryptAcquireContextA(phProv: *mut HCRYPTPROV,
                                pszContainer: LPCSTR,
                                pszProvider: LPCSTR,
                                dwProvType: DWORD,
                                dwFlags: DWORD) -> BOOL;
        fn CryptGenRandom(hProv: HCRYPTPROV,
                          dwLen: DWORD,
                          pbBuffer: *mut BYTE) -> BOOL;
        fn CryptReleaseContext(hProv: HCRYPTPROV, dwFlags: DWORD) -> BOOL;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            let mut hcp = 0;
            let ret = unsafe {
                CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            if ret == 0 {
                Err(IoError::last_error())
            } else {
                Ok(OsRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                CryptGenRandom(self.hcryptprov, v.len() as DWORD,
                               v.as_mut_ptr())
            };
            if ret == 0 {
                panic!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }

    impl Drop for OsRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                panic!("couldn't release context: {}", os::last_os_error());
            }
        }
    }
}

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

    use super::OsRng;
    use rand::Rng;
    use task;

    #[test]
    fn test_os_rng() {
        let mut r = OsRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8, .. 1000];
        r.fill_bytes(&mut v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in range(0u, 20) {
            let (tx, rx) = channel();
            txs.push(tx);
            task::spawn(move|| {
                // wait until all the tasks are ready to go.
                rx.recv();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OsRng::new().unwrap();
                task::deschedule();
                let mut v = [0u8, .. 1000];

                for _ in range(0u, 100) {
                    r.next_u32();
                    task::deschedule();
                    r.next_u64();
                    task::deschedule();
                    r.fill_bytes(&mut v);
                    task::deschedule();
                }
            })
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(())
        }
    }
}
<|endoftext|>


<commit_before>use std::marker::PhantomData;
use std::ops::MulAssign;
use approx::ApproxEq;
use slab::Slab;
use alga::linear::Transformation;
use alga::linear::ProjectiveTransformation;
use na::{self, DVector, Dim, Dynamic, Real, U1, Unit, VectorSliceMutN};
use ncollide::query::ContactKinematic;
use ncollide::query::closest_points_internal;

use object::{BodyHandle, BodyPart, BodySet};
use joint::JointConstraint;
use solver::helper;
use solver::{ForceDirection, IntegrationParameters,
             MultibodyJointLimitsNonlinearConstraintGenerator, NonlinearConstraintGenerator,
             NonlinearUnilateralConstraint};
use math::{Isometry, Point, Rotation, Vector};

pub struct NonlinearSORProx<N: Real> {
    _phantom: PhantomData<N>,
}

impl<N: Real> NonlinearSORProx<N> {
    pub fn new() -> Self {
        NonlinearSORProx {
            _phantom: PhantomData,
        }
    }

    pub fn solve(
        &self,
        params: &IntegrationParameters<N>,
        bodies: &mut BodySet<N>,
        constraints: &mut [NonlinearUnilateralConstraint<N>],
        multibody_limits: &[MultibodyJointLimitsNonlinearConstraintGenerator],
        joints_constraints: &Slab<Box<JointConstraint<N>>>, // FIXME: ugly, use a slice of refs instead.
        jacobians: &mut [N],
        max_iter: usize,
    ) {
        for _ in 0..max_iter {
            for constraint in constraints.iter_mut() {
                // FIXME: specialize for SPATIAL_DIM.
                let dim1 = Dynamic::new(constraint.ndofs1);
                let dim2 = Dynamic::new(constraint.ndofs2);
                self.solve_unilateral(params, bodies, constraint, jacobians, dim1, dim2);
            }

            for generator in multibody_limits {
                self.solve_generic(params, bodies, generator, jacobians)
            }

            for joint in &*joints_constraints {
                self.solve_generic(params, bodies, &**joint.1, jacobians)
            }
        }
    }

    fn solve_generic<Gen: ?Sized + NonlinearConstraintGenerator<N>>(
        &self,
        params: &IntegrationParameters<N>,
        bodies: &mut BodySet<N>,
        generator: &Gen,
        jacobians: &mut [N],
    ) {
        let nconstraints = generator.num_position_constraints(bodies);

        for i in 0..nconstraints {
            if let Some(constraint) = generator.position_constraint(params, i, bodies, jacobians) {
                let dim1 = Dynamic::new(constraint.dim1);
                let dim2 = Dynamic::new(constraint.dim2);

                let rhs = if constraint.is_angular {
                    na::sup(
                        &((constraint.rhs + params.allowed_angular_error) * params.erp),
                        &(-params.max_angular_correction),
                    )
                } else {
                    na::sup(
                        &((constraint.rhs + params.allowed_linear_error) * params.erp),
                        &(-params.max_linear_correction),
                    )
                };

                if rhs < N::zero() {
                    let impulse = -rhs * constraint.r;

                    VectorSliceMutN::new_generic_mut(&mut jacobians[constraint.wj_id1..], dim1, U1)
                        .mul_assign(impulse);

                    VectorSliceMutN::new_generic_mut(&mut jacobians[constraint.wj_id2..], dim2, U1)
                        .mul_assign(impulse);

                    // FIXME: the body update should be performed lazily, especially because
                    // we dont actually need to update the kinematic of a multibody until
                    // we have to solve a contact involvoing one of its links.
                    bodies.body_mut(constraint.body1).apply_displacement(
                        &jacobians[constraint.wj_id1..constraint.wj_id1 + constraint.dim1],
                    );
                    bodies.body_mut(constraint.body2).apply_displacement(
                        &jacobians[constraint.wj_id2..constraint.wj_id2 + constraint.dim2],
                    );
                }
            }
        }
    }

    fn solve_unilateral<D1: Dim, D2: Dim>(
        &self,
        params: &IntegrationParameters<N>,
        bodies: &mut BodySet<N>,
        constraint: &mut NonlinearUnilateralConstraint<N>,
        jacobians: &mut [N],
        dim1: D1,
        dim2: D2,
    ) {
        if self.update_contact_constraint(params, bodies, constraint, jacobians) {
            let impulse = -constraint.rhs * constraint.r;

            VectorSliceMutN::new_generic_mut(jacobians, dim1, U1).mul_assign(impulse);
            VectorSliceMutN::new_generic_mut(&mut jacobians[dim1.value()..], dim2, U1)
                .mul_assign(impulse);

            if dim1.value() != 0 {
                bodies
                    .body_mut(constraint.body1)
                    .apply_displacement(&jacobians[0..dim1.value()]);
            }
            if dim2.value() != 0 {
                bodies
                    .body_mut(constraint.body2)
                    .apply_displacement(&jacobians[dim1.value()..dim1.value() + dim2.value()]);
            }
        }
    }

    fn update_contact_constraint(
        &self,
        params: &IntegrationParameters<N>,
        bodies: &BodySet<N>,
        constraint: &mut NonlinearUnilateralConstraint<N>,
        jacobians: &mut [N],
    ) -> bool {
        let body1 = bodies.body_part(constraint.body1);
        let body2 = bodies.body_part(constraint.body2);
        let pos1 = body1.position();
        let pos2 = body2.position();

        if let Some(contact) = constraint.kinematic.contact(&pos1, &pos2, &constraint.normal1) {
            println!("Updated contact: {:?}", contact);
            constraint.rhs = na::sup(
                &((-contact.depth + params.allowed_linear_error) * params.erp),
                &(-params.max_linear_correction),
            );

            if constraint.rhs >= N::zero() {
                return false;
            }

            // XXX: should use constraint_pair_geometry to properly handle multibodies.
            let mut inv_r = N::zero();

            if constraint.ndofs1 != 0 {
                helper::fill_constraint_geometry(
                    &body1,
                    constraint.ndofs1,
                    &contact.world1,
                    &ForceDirection::Linear(-contact.normal),
                    constraint.ndofs1 + constraint.ndofs2,
                    0,
                    jacobians,
                    &mut inv_r,
                );
            }

            if constraint.ndofs2 != 0 {
                helper::fill_constraint_geometry(
                    &body2,
                    constraint.ndofs2,
                    &contact.world2,
                    &ForceDirection::Linear(contact.normal),
                    (constraint.ndofs1 * 2) + constraint.ndofs2,
                    constraint.ndofs1,
                    jacobians,
                    &mut inv_r,
                );
            }

            // May happen sometimes for self-collisions (e.g. a multibody with itself).
            if inv_r.is_zero() {
                constraint.r = N::one();
            } else {
                constraint.r = N::one() / inv_r
            }

            true
        } else {
            false
        }
    }
}
<commit_msg>Remove debug print.<commit_after>use std::marker::PhantomData;
use std::ops::MulAssign;
use approx::ApproxEq;
use slab::Slab;
use alga::linear::Transformation;
use alga::linear::ProjectiveTransformation;
use na::{self, DVector, Dim, Dynamic, Real, U1, Unit, VectorSliceMutN};
use ncollide::query::ContactKinematic;
use ncollide::query::closest_points_internal;

use object::{BodyHandle, BodyPart, BodySet};
use joint::JointConstraint;
use solver::helper;
use solver::{ForceDirection, IntegrationParameters,
             MultibodyJointLimitsNonlinearConstraintGenerator, NonlinearConstraintGenerator,
             NonlinearUnilateralConstraint};
use math::{Isometry, Point, Rotation, Vector};

pub struct NonlinearSORProx<N: Real> {
    _phantom: PhantomData<N>,
}

impl<N: Real> NonlinearSORProx<N> {
    pub fn new() -> Self {
        NonlinearSORProx {
            _phantom: PhantomData,
        }
    }

    pub fn solve(
        &self,
        params: &IntegrationParameters<N>,
        bodies: &mut BodySet<N>,
        constraints: &mut [NonlinearUnilateralConstraint<N>],
        multibody_limits: &[MultibodyJointLimitsNonlinearConstraintGenerator],
        joints_constraints: &Slab<Box<JointConstraint<N>>>, // FIXME: ugly, use a slice of refs instead.
        jacobians: &mut [N],
        max_iter: usize,
    ) {
        for _ in 0..max_iter {
            for constraint in constraints.iter_mut() {
                // FIXME: specialize for SPATIAL_DIM.
                let dim1 = Dynamic::new(constraint.ndofs1);
                let dim2 = Dynamic::new(constraint.ndofs2);
                self.solve_unilateral(params, bodies, constraint, jacobians, dim1, dim2);
            }

            for generator in multibody_limits {
                self.solve_generic(params, bodies, generator, jacobians)
            }

            for joint in &*joints_constraints {
                self.solve_generic(params, bodies, &**joint.1, jacobians)
            }
        }
    }

    fn solve_generic<Gen: ?Sized + NonlinearConstraintGenerator<N>>(
        &self,
        params: &IntegrationParameters<N>,
        bodies: &mut BodySet<N>,
        generator: &Gen,
        jacobians: &mut [N],
    ) {
        let nconstraints = generator.num_position_constraints(bodies);

        for i in 0..nconstraints {
            if let Some(constraint) = generator.position_constraint(params, i, bodies, jacobians) {
                let dim1 = Dynamic::new(constraint.dim1);
                let dim2 = Dynamic::new(constraint.dim2);

                let rhs = if constraint.is_angular {
                    na::sup(
                        &((constraint.rhs + params.allowed_angular_error) * params.erp),
                        &(-params.max_angular_correction),
                    )
                } else {
                    na::sup(
                        &((constraint.rhs + params.allowed_linear_error) * params.erp),
                        &(-params.max_linear_correction),
                    )
                };

                if rhs < N::zero() {
                    let impulse = -rhs * constraint.r;

                    VectorSliceMutN::new_generic_mut(&mut jacobians[constraint.wj_id1..], dim1, U1)
                        .mul_assign(impulse);

                    VectorSliceMutN::new_generic_mut(&mut jacobians[constraint.wj_id2..], dim2, U1)
                        .mul_assign(impulse);

                    // FIXME: the body update should be performed lazily, especially because
                    // we dont actually need to update the kinematic of a multibody until
                    // we have to solve a contact involvoing one of its links.
                    bodies.body_mut(constraint.body1).apply_displacement(
                        &jacobians[constraint.wj_id1..constraint.wj_id1 + constraint.dim1],
                    );
                    bodies.body_mut(constraint.body2).apply_displacement(
                        &jacobians[constraint.wj_id2..constraint.wj_id2 + constraint.dim2],
                    );
                }
            }
        }
    }

    fn solve_unilateral<D1: Dim, D2: Dim>(
        &self,
        params: &IntegrationParameters<N>,
        bodies: &mut BodySet<N>,
        constraint: &mut NonlinearUnilateralConstraint<N>,
        jacobians: &mut [N],
        dim1: D1,
        dim2: D2,
    ) {
        if self.update_contact_constraint(params, bodies, constraint, jacobians) {
            let impulse = -constraint.rhs * constraint.r;

            VectorSliceMutN::new_generic_mut(jacobians, dim1, U1).mul_assign(impulse);
            VectorSliceMutN::new_generic_mut(&mut jacobians[dim1.value()..], dim2, U1)
                .mul_assign(impulse);

            if dim1.value() != 0 {
                bodies
                    .body_mut(constraint.body1)
                    .apply_displacement(&jacobians[0..dim1.value()]);
            }
            if dim2.value() != 0 {
                bodies
                    .body_mut(constraint.body2)
                    .apply_displacement(&jacobians[dim1.value()..dim1.value() + dim2.value()]);
            }
        }
    }

    fn update_contact_constraint(
        &self,
        params: &IntegrationParameters<N>,
        bodies: &BodySet<N>,
        constraint: &mut NonlinearUnilateralConstraint<N>,
        jacobians: &mut [N],
    ) -> bool {
        let body1 = bodies.body_part(constraint.body1);
        let body2 = bodies.body_part(constraint.body2);
        let pos1 = body1.position();
        let pos2 = body2.position();

        if let Some(contact) = constraint
            .kinematic
            .contact(&pos1, &pos2, &constraint.normal1)
        {
            constraint.rhs = na::sup(
                &((-contact.depth + params.allowed_linear_error) * params.erp),
                &(-params.max_linear_correction),
            );

            if constraint.rhs >= N::zero() {
                return false;
            }

            // XXX: should use constraint_pair_geometry to properly handle multibodies.
            let mut inv_r = N::zero();

            if constraint.ndofs1 != 0 {
                helper::fill_constraint_geometry(
                    &body1,
                    constraint.ndofs1,
                    &contact.world1,
                    &ForceDirection::Linear(-contact.normal),
                    constraint.ndofs1 + constraint.ndofs2,
                    0,
                    jacobians,
                    &mut inv_r,
                );
            }

            if constraint.ndofs2 != 0 {
                helper::fill_constraint_geometry(
                    &body2,
                    constraint.ndofs2,
                    &contact.world2,
                    &ForceDirection::Linear(contact.normal),
                    (constraint.ndofs1 * 2) + constraint.ndofs2,
                    constraint.ndofs1,
                    jacobians,
                    &mut inv_r,
                );
            }

            // May happen sometimes for self-collisions (e.g. a multibody with itself).
            if inv_r.is_zero() {
                constraint.r = N::one();
            } else {
                constraint.r = N::one() / inv_r
            }

            true
        } else {
            false
        }
    }
}
<|endoftext|>


<commit_before>extern mod std;

/* For future reference, if necessary (from pcwalton, #rust, 2012-10-17 20:38)
   basically to allocate something in the same lifetime as yourself, you need to use an arena
   take in the arena, and then allocate into it. for example:
   impl Foo { fn make_foo(&self, arena_self_is_in: &self/Arena) -> &self/Foo { ... } }
   you need to pass in the arena that self is in
   so that the function can allocate into it
*/

struct Tree {
    mut leftTree: Option<~Tree>,
    mut rightTree: Option<~Tree>,
    key: int
}

impl Tree {
    fn insert(&self, i: int) {
        let mut node = None;
        if i > self.key {
            node <-> self.rightTree;
            match node {
                Some(ref t) => t.insert(i),
                None => node = Some(~Tree {leftTree: None, rightTree: None, key: i})
            }
            node <-> self.rightTree;
        }
        if i < self.key {
            node <-> self.leftTree;
            match node {
                Some (ref t) => t.insert(i),
                None => node = Some(~Tree {leftTree: None, rightTree: None, key: i})
            }
            node <-> self.leftTree;
        }

    }

    pure fn treeDepth (&self) -> int {
        let leftCount = match self.leftTree {
            Some (ref t) => t.treeDepth(),
            None => 0
        };
        let rightCount = match self.rightTree {
            Some (ref t) => t.treeDepth(),
            None => 0
        };
        1 + int::max(leftCount, rightCount)
    }

    pure fn contains(i: int) -> bool {
        if i == self.key {
            return true;
        }

        if i > self.key {
            match self.rightTree {
                Some(~t) => t.contains(i),
                None => false
            }
        }
        else {
            match self.leftTree {
                Some(~t) => t.contains(i),
                None => false
            }
        }
    }

    fn print() {
        let mut node = None;
        node <-> self.leftTree;
        match node {
            Some(ref t) => t.print(),
            None => ()
        };
        node <-> self.leftTree;

        io::print(fmt!("%d ", self.key));

        node <-> self.rightTree;
        match node {
            Some(ref t) => t.print(),
            None => ()
        };
        node <-> self.rightTree;
    }

}

fn fillTree (tree: &Tree, numNodes: int) {
    let r = rand::Rng();
    for int::range(0, numNodes) |_index| {
        tree.insert(r.gen_int_range(0, numNodes));
    }
}

fn main () {
    let root: Tree = Tree {leftTree: None, rightTree: None, key: 5};

    root.insert(8);
    fillTree(&root, 1000);

    root.print();
    io::println("");

    if root.contains(4) {
        io::println("Contains 4")
    }
    else {
        io::println("Doesn't contain 4")
    }

    root.insert(4);

    if root.contains(4) {
        io::println("Contains 4 now ...")
    }
    else {
        io::println("Doesn't contain 4 yet")
    }

    io::println(fmt!("Tree depth: %d", root.treeDepth()));

}
<commit_msg>Lots more code for binary tree - delete is not working correctly yet; incorrect algorithm :(<commit_after>extern mod std;

/* For future reference, if necessary (from pcwalton, #rust, 2012-10-17 20:38)
   basically to allocate something in the same lifetime as yourself, you need to use an arena
   take in the arena, and then allocate into it. for example:
   impl Foo { fn make_foo(&self, arena_self_is_in: &self/Arena) -> &self/Foo { ... } }
   you need to pass in the arena that self is in
   so that the function can allocate into it
*/

struct Tree {
    mut leftTree: Option<~Tree>,
    mut rightTree: Option<~Tree>,
    key: int
}

impl Tree {
    fn insert(&self, i: int) {
        let mut node = None;
        if i > self.key {
            node <-> self.rightTree;
            match node {
                Some(ref t) => t.insert(i),
                None => node = Some(~Tree {leftTree: None, rightTree: None, key: i})
            }
            node <-> self.rightTree;
        }
        if i < self.key {
            node <-> self.leftTree;
            match node {
                Some (ref t) => t.insert(i),
                None => node = Some(~Tree {leftTree: None, rightTree: None, key: i})
            }
            node <-> self.leftTree;
        }
    }

    fn insertTree(&self, t: &self/Tree) {
        if self.key == t.key {
            let mut node = None;
            node <-> t.leftTree;
            match node {
                Some (ref lt) => {
                    let mut leftTree = None;
                    leftTree <-> self.leftTree;
                    match leftTree {
                        Some (ref slt) => slt.insertTree(*lt),
                        None => leftTree = Some (copy *lt)
                    };
                    leftTree <-> self.leftTree;
                },
                None => ()
            };
            node <-> t.rightTree;
            match node {
                Some (ref rt) => {
                    let mut rightTree = None;
                    rightTree <-> self.rightTree;
                    match rightTree {
                        Some (ref srt) => srt.insertTree(*rt),
                        None => rightTree = Some (copy *rt)
                    };
                    rightTree <-> self.rightTree;
                },
                None => ()
            };
        }
    }

    fn delete(&self, key: int) {
        assert self.key != key;

        let mut node = None;
        node <-> self.leftTree;
        match node {
            Some (ref t) => {
                if t.key == key {
                    // self.leftTree --> top node is the one to delete

                    // need to remove the left tree root
                    let mut tlt = None;
                    tlt <-> t.leftTree;
                    self.leftTree = tlt;

                    // and the right tree root
                    let mut trt = None;
                    trt <-> t.rightTree;
                    match trt {
                        Some (ref tr) => {
                            // the deleted node has a right tree
                            // we need to insert that into self.rightTree
                            let mut srt = None;
                            srt <-> self.rightTree;
                            match srt {
                                Some (ref rt) => rt.insertTree(*tr),
                                None => self.rightTree = tlt
                            };
                        },
                        None => {
                            // the deleted node didn't have a right tree
                            // we don't need to do anything here
                            ()
                        }
                    };

                    return ();
                }
                else {
                    // don't remove the left tree, put it back instead (later)
                    if t.key < key {
                        t.delete(key);
                    }
                }
            },
            None => {
                // self.leftTree doesn't exists
                ()
            }
        };
        node <-> self.leftTree;

        node = None;
        node <-> self.rightTree;
        match node {
            Some (ref t) => {
                if t.key == key {
                    // self.rightTree --> top node is the one to delete

                    // need to remove the right tree root
                    let mut trt = None;
                    trt <-> t.rightTree;
                    self.rightTree = trt;

                    // and the left tree root
                    let mut tlt = None;
                    tlt <-> t.leftTree;
                    match tlt {
                        Some (ref tl) => {
                            // the deleted node has a left tree
                            // we need to insert that into self.leftTree
                            let mut slt = None;
                            slt <-> self.leftTree;
                            match slt {
                                Some (ref lt) => lt.insertTree(*tl),
                                None => self.leftTree = tlt
                            };
                        },
                        None => {
                            // the deleted node didn't have a right tree
                            // we don't need to do anything here
                            ()
                        }
                    };

                    return ();
                    // need to remove the right tree root
                    //let mut trt = None;
                    //trt <-> t.rightTree;
                    //self.rightTree = trt;
                    //return ();
                }
                else {
                    // don't remove the right tree, put it back instead
                    if t.key > key {
                        t.delete(key);
                    }
                }
            },
            None => {
                ()
            }
        };
        node <-> self.rightTree;
    }

    pure fn treeDepth (&self) -> int {
        let leftCount = match self.leftTree {
            Some (ref t) => t.treeDepth(),
            None => 0
        };
        let rightCount = match self.rightTree {
            Some (ref t) => t.treeDepth(),
            None => 0
        };
        1 + int::max(leftCount, rightCount)
    }

    pure fn contains(i: int) -> bool {
        if i == self.key {
            return true;
        }

        if i > self.key {
            match self.rightTree {
                Some(~t) => t.contains(i),
                None => false
            }
        }
        else {
            match self.leftTree {
                Some(~t) => t.contains(i),
                None => false
            }
        }
    }

    fn print() {
        let mut node = None;
        node <-> self.leftTree;
        match node {
            Some(ref t) => t.print(),
            None => ()
        };
        node <-> self.leftTree;

        io::print(fmt!("%d ", self.key));

        node <-> self.rightTree;
        match node {
            Some(ref t) => t.print(),
            None => ()
        };
        node <-> self.rightTree;
    }
}

impl Tree: ToStr {
    pure fn to_str() -> ~str {
        return fmt!("%d", self.key);
    }
}

fn fillTreeRandom (tree: &Tree, numNodes: int) {
    let r = rand::Rng();
    for int::range(0, numNodes) |_index| {
        tree.insert(r.gen_int_range(0, numNodes));
    }
}

fn fillTreeRec (tree: &Tree, key: int, depth: int) {
    tree.insert(key);
    if depth >= 0 {
        let div = float::pow_with_uint(2, depth as uint) as int;
        //io::println(fmt!("key: %d, depth: %d, div: %d", key, depth, div));
        fillTreeRec(tree, key - div, depth - 1);
        fillTreeRec(tree, key + div, depth - 1);
    }
}

fn main () {
    let root: Tree = Tree {leftTree: None, rightTree: None, key: 5};

    root.insert(8);
    //fillTreeRandom(&root, 1000);
    fillTreeRec(&root, 8, 2);

    root.print();
    io::println("");

    if root.contains(4) {
        io::println("Contains 4")
    }
    else {
        io::println("Doesn't contain 4")
    }

    root.delete(4);
    root.delete(8);

    if root.contains(4) {
        io::println("Still contains 4 now ...")
    }
    else {
        io::println("Doesn't contain 4 anymore")
    }
    root.print();
    io::println("");

    io::println(fmt!("Tree depth: %d", root.treeDepth()));

}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: MIT
// Copyright (C) 2018-present iced project and contributors

#![allow(unknown_lints)]
#![warn(absolute_paths_not_starting_with_crate)]
#![warn(anonymous_parameters)]
#![warn(elided_lifetimes_in_paths)]
#![warn(explicit_outlives_requirements)]
#![warn(invalid_html_tags)]
#![warn(keyword_idents)]
#![warn(macro_use_extern_crate)]
#![warn(meta_variable_misuse)]
#![warn(missing_copy_implementations)]
#![warn(missing_debug_implementations)]
#![warn(non_ascii_idents)]
#![warn(trivial_casts)]
#![warn(trivial_numeric_casts)]
#![warn(unsafe_op_in_unsafe_fn)]
#![warn(unused_extern_crates)]
#![warn(unused_import_braces)]
#![warn(unused_lifetimes)]
#![warn(unused_must_use)]
#![warn(unused_qualifications)]
#![warn(unused_results)]
#![allow(clippy::cast_lossless)]
#![allow(clippy::collapsible_else_if)]
#![allow(clippy::collapsible_if)]
#![allow(clippy::field_reassign_with_default)]
#![allow(clippy::manual_range_contains)]
#![allow(clippy::match_ref_pats)]
#![allow(clippy::ptr_eq)]
#![allow(clippy::too_many_arguments)]
#![allow(clippy::type_complexity)]
#![allow(clippy::upper_case_acronyms)]
#![allow(clippy::wrong_self_convention)]
#![warn(clippy::cloned_instead_of_copied)]
#![warn(clippy::dbg_macro)]
#![warn(clippy::debug_assert_with_mut_call)]
#![warn(clippy::default_trait_access)]
#![warn(clippy::doc_markdown)]
#![warn(clippy::empty_line_after_outer_attr)]
#![warn(clippy::expect_used)]
#![warn(clippy::explicit_into_iter_loop)]
#![warn(clippy::explicit_iter_loop)]
#![warn(clippy::fallible_impl_from)]
#![warn(clippy::get_unwrap)]
#![warn(clippy::implicit_saturating_sub)]
#![warn(clippy::large_digit_groups)]
#![warn(clippy::let_unit_value)]
#![warn(clippy::match_bool)]
#![warn(clippy::match_on_vec_items)]
#![warn(clippy::match_wild_err_arm)]
#![warn(clippy::missing_errors_doc)]
#![warn(clippy::missing_inline_in_public_items)]
#![warn(clippy::needless_borrow)]
#![warn(clippy::print_stderr)]
#![warn(clippy::print_stdout)]
#![warn(clippy::rc_buffer)]
#![warn(clippy::redundant_closure)]
#![warn(clippy::redundant_closure_for_method_calls)]
#![warn(clippy::same_functions_in_if_condition)]
#![warn(clippy::todo)]
#![warn(clippy::unimplemented)]
#![warn(clippy::unreadable_literal)]
#![warn(clippy::unused_self)]
#![warn(clippy::unwrap_in_result)]
#![warn(clippy::unwrap_used)]
#![warn(clippy::useless_let_if_seq)]
#![warn(clippy::useless_transmute)]
#![warn(clippy::zero_sized_map_values)]

pub mod lua;
pub mod lua_api;
pub mod macros;
<commit_msg>Remove `clippy::let_unit_value`<commit_after>// SPDX-License-Identifier: MIT
// Copyright (C) 2018-present iced project and contributors

#![allow(unknown_lints)]
#![warn(absolute_paths_not_starting_with_crate)]
#![warn(anonymous_parameters)]
#![warn(elided_lifetimes_in_paths)]
#![warn(explicit_outlives_requirements)]
#![warn(invalid_html_tags)]
#![warn(keyword_idents)]
#![warn(macro_use_extern_crate)]
#![warn(meta_variable_misuse)]
#![warn(missing_copy_implementations)]
#![warn(missing_debug_implementations)]
#![warn(non_ascii_idents)]
#![warn(trivial_casts)]
#![warn(trivial_numeric_casts)]
#![warn(unsafe_op_in_unsafe_fn)]
#![warn(unused_extern_crates)]
#![warn(unused_import_braces)]
#![warn(unused_lifetimes)]
#![warn(unused_must_use)]
#![warn(unused_qualifications)]
#![warn(unused_results)]
#![allow(clippy::cast_lossless)]
#![allow(clippy::collapsible_else_if)]
#![allow(clippy::collapsible_if)]
#![allow(clippy::field_reassign_with_default)]
#![allow(clippy::manual_range_contains)]
#![allow(clippy::match_ref_pats)]
#![allow(clippy::ptr_eq)]
#![allow(clippy::too_many_arguments)]
#![allow(clippy::type_complexity)]
#![allow(clippy::upper_case_acronyms)]
#![allow(clippy::wrong_self_convention)]
#![warn(clippy::cloned_instead_of_copied)]
#![warn(clippy::dbg_macro)]
#![warn(clippy::debug_assert_with_mut_call)]
#![warn(clippy::default_trait_access)]
#![warn(clippy::doc_markdown)]
#![warn(clippy::empty_line_after_outer_attr)]
#![warn(clippy::expect_used)]
#![warn(clippy::explicit_into_iter_loop)]
#![warn(clippy::explicit_iter_loop)]
#![warn(clippy::fallible_impl_from)]
#![warn(clippy::get_unwrap)]
#![warn(clippy::implicit_saturating_sub)]
#![warn(clippy::large_digit_groups)]
#![warn(clippy::match_bool)]
#![warn(clippy::match_on_vec_items)]
#![warn(clippy::match_wild_err_arm)]
#![warn(clippy::missing_errors_doc)]
#![warn(clippy::missing_inline_in_public_items)]
#![warn(clippy::needless_borrow)]
#![warn(clippy::print_stderr)]
#![warn(clippy::print_stdout)]
#![warn(clippy::rc_buffer)]
#![warn(clippy::redundant_closure)]
#![warn(clippy::redundant_closure_for_method_calls)]
#![warn(clippy::same_functions_in_if_condition)]
#![warn(clippy::todo)]
#![warn(clippy::unimplemented)]
#![warn(clippy::unreadable_literal)]
#![warn(clippy::unused_self)]
#![warn(clippy::unwrap_in_result)]
#![warn(clippy::unwrap_used)]
#![warn(clippy::useless_let_if_seq)]
#![warn(clippy::useless_transmute)]
#![warn(clippy::zero_sized_map_values)]

pub mod lua;
pub mod lua_api;
pub mod macros;
<|endoftext|>


<commit_before>use image::{ImageBuffer, RgbImage};
use bitreader::{BitReader};
use itertools::Itertools;
use discrete_transforms::dct_2d::Dct2D;

#[derive(Debug)]
pub struct CoverImage {
    cover_image: RgbImage,
    message: &'static str,
    output_image_path: &'static str,
    message_as_bits: Vec<u8>,
    tiles: Vec<u8>,
    transformed_tiles: Vec<f64>,
    modified_pixels: Vec<f64>
}

impl CoverImage {
    pub fn new() -> CoverImage {
        CoverImage {
            cover_image: ImageBuffer::new(0, 0),
            message: "",
            output_image_path: "",
            message_as_bits: vec![],
            tiles: vec![],
            transformed_tiles: vec![],
            modified_pixels: vec![]
        }
    }

    pub fn set_cover_image(&mut self, input_path: &'static str) -> &mut Self {
        self.cover_image = image::open(input_path).unwrap().as_rgb8().unwrap().clone();

        self
    }

    pub fn set_message(&mut self, message: &'static str) -> &mut Self {
        self.message = message;

        self
    }

    pub fn set_output_image_path(&mut self, output_path: &'static str) -> &mut Self {
        self.output_image_path = output_path;

        self
    }

    pub fn encode(&mut self) {
        self.get_message_as_bits();
        self.tile_image();
        self.encode_message();

        self.cover_image.save(self.output_image_path).unwrap();
    }

    fn get_message_as_bits(&mut self) {
        let mut reader = BitReader::new(&self.message.as_bytes());

        for _ in 0..self.message.len() * 8 {
            self.message_as_bits.push(reader.read_u8(1).unwrap());
        }
    }

    fn tile_image(&mut self) {
        let (width, height) = self.cover_image.dimensions();
        let mut count = 0;

        for row_index in 0..(height / 8) as u32 {
            for col_index in 0..(width / 8) as u32 {
                for channel in 0..3 {
                    if count == self.message_as_bits.len() {
                        break
                    }

                    for row in 0..8 {
                        for column in 0..8 {
                            let pixel = self.cover_image.get_pixel(column + (col_index * 8),
                                                                      row + (row_index * 8));

                            self.tiles.push(pixel.data[channel]);
                        }
                    }
                    count += 1;
                }
            }
        }
    }

    fn encode_message(&mut self) {
        let mut dct = Dct2D::new();
        let mut count = 0;

        for chunk in &self.tiles.iter().chunks(64) {
            let input = chunk.map(|x| *x as f64).collect_vec();

            dct.set_input(input);

            let mut transformed_tile = dct.forward();

            // If the coefficient is even and message bit is odd add 1
            if transformed_tile[0].round() as u32 % 2 == 0 && self.message_as_bits[count] == 1 {
                transformed_tile[0] = transformed_tile[0] + 1.0;
            }

            dct.set_input(transformed_tile);

            self.modified_pixels.extend(&dct.inverse());

            count += 1;
        }

        let (width, height) = self.cover_image.dimensions();
        let mut count = 0;
        let mut pixel_count = 0;

        for row_index in 0..(height / 8) as u32 {
            for col_index in 0..(width / 8) as u32 {
                for channel in 0..3 {
                    if count == self.message_as_bits.len() {
                        break
                    }

                    for row in 0..8 {
                        for column in 0..8 {
                            let pixel = self.cover_image.get_pixel_mut(column + (col_index * 8),
                                                                          row + (row_index * 8));

                            pixel.data[channel] = self.modified_pixels[pixel_count] as u8;

                            pixel_count += 1;
                        }
                    }

                    count += 1;
                }
            }
        }
    }
}

#[test]
fn get_message_as_bits_test() {
    let test_bits = vec![0, 1, 1, 0, 1, 0, 0, 0];
    let mut cover_image = CoverImage::new();

    cover_image.set_message("h");
    cover_image.get_message_as_bits();

    assert_eq!(cover_image.message_as_bits, test_bits);
}

#[test]
fn tile_image_length_test() {
    let mut cover_image = CoverImage::new();

    cover_image.set_message("h");
    cover_image.set_cover_image("src/testing.jpg");

    // 64 = length of 1 tile, message.len() * 8 = number of tiles needed; 1 tile per bit
    let length = 64 * cover_image.message.len() * 8;

    cover_image.get_message_as_bits();
    cover_image.tile_image();

    let tile_vec_length = cover_image.tiles.len();

    assert_eq!(length, tile_vec_length);
}

#[test]
fn test_random_pixel_in_tile() {
    let mut cover_image = CoverImage::new();

    cover_image.set_message("h");
    cover_image.set_cover_image("src/testing.jpg");
    cover_image.get_message_as_bits();
    cover_image.tile_image();

    let pixel = cover_image.cover_image.get_pixel(8, 8);

    let r_pixel = pixel.data[0];
    let g_pixel = pixel.data[1];
    let b_pixel = pixel.data[2];

    let r_pixel_from_tile = cover_image.tiles[63];
    let g_pixel_from_tile = cover_image.tiles[127];
    let b_pixel_from_tile = cover_image.tiles[191];

    assert_eq!(r_pixel, r_pixel_from_tile);
    assert_eq!(g_pixel, g_pixel_from_tile);
    assert_eq!(b_pixel, b_pixel_from_tile);
}<commit_msg>Begin refactor to fix encoder<commit_after>use image::{ImageBuffer, RgbImage};
use bitreader::{BitReader};
use itertools::Itertools;
use discrete_transforms::dct_2d::Dct2D;

#[derive(Debug)]
pub struct CoverImage {
    cover_image: RgbImage,
    message: &'static str,
    output_image_path: &'static str,
    message_as_bits: Vec<u8>,
    tiles: Vec<u8>,
    dct_tiles: Vec<f64>,
    modified_pixels: Vec<f64>
}

impl CoverImage {
    pub fn new() -> CoverImage {
        CoverImage {
            cover_image: ImageBuffer::new(0, 0),
            message: "",
            output_image_path: "",
            message_as_bits: vec![],
            tiles: vec![],
            dct_tiles: vec![],
            modified_pixels: vec![]
        }
    }

    pub fn set_cover_image(&mut self, input_path: &'static str) -> &mut Self {
        self.cover_image = image::open(input_path).unwrap().as_rgb8().unwrap().clone();

        self
    }

    pub fn set_message(&mut self, message: &'static str) -> &mut Self {
        self.message = message;

        self
    }

    pub fn set_output_image_path(&mut self, output_path: &'static str) -> &mut Self {
        self.output_image_path = output_path;

        self
    }

    pub fn encode(&mut self) {
        self.get_message_as_bits();
        self.tile_image();
        self.encode_message();

        self.cover_image.save(self.output_image_path).unwrap();
    }

    fn get_message_as_bits(&mut self) {
        let mut reader = BitReader::new(&self.message.as_bytes());

        for _ in 0..self.message.len() * 8 {
            self.message_as_bits.push(reader.read_u8(1).unwrap());
        }
    }

    fn tile_image(&mut self) {
        let (width, height) = self.cover_image.dimensions();
        let mut count = 0;

        for row_index in 0..(height / 8) as u32 {
            for col_index in 0..(width / 8) as u32 {
                for channel in 0..3 {
                    if count == self.message_as_bits.len() {
                        break
                    }

                    for row in 0..8 {
                        for column in 0..8 {
                            let pixel = self.cover_image.get_pixel(column + (col_index * 8),
                                                                      row + (row_index * 8));

                            self.tiles.push(pixel.data[channel]);
                        }
                    }

                    count += 1;
                }
            }
        }
    }

    fn encode_message(&mut self) {
        let mut dct = Dct2D::new();

        for chunk in &self.tiles.iter().chunks(64) {
            let input = chunk.map(|x| *x as f64).collect_vec();

            dct.set_input(input);

            self.dct_tiles.extend(&dct.forward());
        }
    }
}

#[test]
fn get_message_as_bits_test() {
    let test_bits = vec![0, 1, 1, 0, 1, 0, 0, 0];
    let mut cover_image = CoverImage::new();

    cover_image.set_message("h");
    cover_image.get_message_as_bits();

    assert_eq!(cover_image.message_as_bits, test_bits);
}

#[test]
fn tile_image_length_test() {
    let mut cover_image = CoverImage::new();

    cover_image.set_message("h");
    cover_image.set_cover_image("src/testing.jpg");

    // 64 = length of 1 tile, message.len() * 8 = number of tiles needed; 1 tile per bit
    let length = 64 * cover_image.message.len() * 8;

    cover_image.get_message_as_bits();
    cover_image.tile_image();

    let tile_vec_length = cover_image.tiles.len();

    assert_eq!(length, tile_vec_length);
}

#[test]
fn test_random_pixel_in_tile() {
    let mut cover_image = CoverImage::new();

    cover_image.set_message("h");
    cover_image.set_cover_image("src/testing.jpg");
    cover_image.get_message_as_bits();
    cover_image.tile_image();

    let pixel = cover_image.cover_image.get_pixel(8, 8);

    let r_pixel = pixel.data[0];
    let g_pixel = pixel.data[1];
    let b_pixel = pixel.data[2];

    let r_pixel_from_tile = cover_image.tiles[63];
    let g_pixel_from_tile = cover_image.tiles[127];
    let b_pixel_from_tile = cover_image.tiles[191];

    assert_eq!(r_pixel, r_pixel_from_tile);
    assert_eq!(g_pixel, g_pixel_from_tile);
    assert_eq!(b_pixel, b_pixel_from_tile);
}

#[test]
fn delete() {
    let mut cover_image = CoverImage::new();

    cover_image.set_message("h");
    cover_image.set_cover_image("src/testing.jpg");
    cover_image.set_output_image_path("output.jpg");
    cover_image.encode();
}<|endoftext|>