<commit_before>use hir::Semantics;
use ide_db::{
    base_db::FilePosition,
    defs::Definition,
    helpers::pick_best_token,
    search::{FileReference, ReferenceAccess, SearchScope},
    RootDatabase,
};
use syntax::{
    ast, match_ast, AstNode,
    SyntaxKind::{ASYNC_KW, AWAIT_KW, QUESTION, RETURN_KW, THIN_ARROW},
    SyntaxNode, SyntaxToken, TextRange, WalkEvent,
};

use crate::{display::TryToNav, references, NavigationTarget};

pub struct DocumentHighlight {
    pub range: TextRange,
    pub access: Option<ReferenceAccess>,
}

// Feature: Highlight Related
//
// Highlights exit points, yield points or the definition and all references of the item at the cursor location in the current file.
pub(crate) fn highlight_related(
    sema: &Semantics<RootDatabase>,
    position: FilePosition,
) -> Option<Vec<DocumentHighlight>> {
    let _p = profile::span("document_highlight");
    let syntax = sema.parse(position.file_id).syntax().clone();

    let token = pick_best_token(syntax.token_at_offset(position.offset), |kind| match kind {
        QUESTION => 2, // prefer `?` when the cursor is sandwiched like `await$0?`
        AWAIT_KW | ASYNC_KW | THIN_ARROW | RETURN_KW => 1,
        _ => 0,
    })?;

    match token.kind() {
        QUESTION | RETURN_KW | THIN_ARROW => highlight_exit_points(sema, token),
        AWAIT_KW | ASYNC_KW => highlight_yield_points(token),
        _ => highlight_references(sema, &syntax, position),
    }
}

fn highlight_references(
    sema: &Semantics<RootDatabase>,
    syntax: &SyntaxNode,
    FilePosition { offset, file_id }: FilePosition,
) -> Option<Vec<DocumentHighlight>> {
    let def = references::find_def(sema, syntax, offset)?;
    let usages = def.usages(sema).set_scope(Some(SearchScope::single_file(file_id))).all();

    let declaration = match def {
        Definition::ModuleDef(hir::ModuleDef::Module(module)) => {
            Some(NavigationTarget::from_module_to_decl(sema.db, module))
        }
        def => def.try_to_nav(sema.db),
    }
    .filter(|decl| decl.file_id == file_id)
    .and_then(|decl| {
        let range = decl.focus_range?;
        let access = references::decl_access(&def, syntax, range);
        Some(DocumentHighlight { range, access })
    });

    let file_refs = usages.references.get(&file_id).map_or(&[][..], Vec::as_slice);
    let mut res = Vec::with_capacity(file_refs.len() + 1);
    res.extend(declaration);
    res.extend(
        file_refs
            .iter()
            .map(|&FileReference { access, range, .. }| DocumentHighlight { range, access }),
    );
    Some(res)
}

fn highlight_exit_points(
    sema: &Semantics<RootDatabase>,
    token: SyntaxToken,
) -> Option<Vec<DocumentHighlight>> {
    fn hl(
        sema: &Semantics<RootDatabase>,
        body: Option<ast::Expr>,
    ) -> Option<Vec<DocumentHighlight>> {
        let mut highlights = Vec::new();
        let body = body?;
        walk(&body, |node| {
            match_ast! {
                match node {
                    ast::ReturnExpr(expr) => if let Some(token) = expr.return_token() {
                        highlights.push(DocumentHighlight {
                            access: None,
                            range: token.text_range(),
                        });
                    },
                    ast::TryExpr(try_) => if let Some(token) = try_.question_mark_token() {
                        highlights.push(DocumentHighlight {
                            access: None,
                            range: token.text_range(),
                        });
                    },
                    ast::Expr(expr) => match expr {
                        ast::Expr::MethodCallExpr(_) | ast::Expr::CallExpr(_) | ast::Expr::MacroCall(_) => {
                            if sema.type_of_expr(&expr).map_or(false, |ty| ty.is_never()) {
                                highlights.push(DocumentHighlight {
                                    access: None,
                                    range: expr.syntax().text_range(),
                                });
                            }
                        },
                        ast::Expr::EffectExpr(effect) => return effect.async_token().is_some() || effect.try_token().is_some(),
                        ast::Expr::ClosureExpr(_) => return true,
                        _ => (),
                    },
                    ast::Item(__) => return true,
                    // Don't look into const args
                    ast::Path(__) => return true,
                    _ => (),
                }
            }
            false
        });
        let tail = match body {
            ast::Expr::BlockExpr(b) => b.tail_expr(),
            e => Some(e),
        };
        if let Some(tail) = tail {
            highlights.push(DocumentHighlight { access: None, range: tail.syntax().text_range() });
        }
        Some(highlights)
    }
    for anc in token.ancestors() {
        return match_ast! {
            match anc {
                ast::Fn(fn_) => hl(sema, fn_.body().map(ast::Expr::BlockExpr)),
                ast::ClosureExpr(closure) => hl(sema, closure.body()),
                ast::EffectExpr(effect) => if effect.async_token().is_some() || effect.try_token().is_some() {
                    hl(sema, effect.block_expr().map(ast::Expr::BlockExpr))
                } else {
                    continue;
                },
                _ => continue,
            }
        };
    }
    None
}

fn highlight_yield_points(token: SyntaxToken) -> Option<Vec<DocumentHighlight>> {
    fn hl(
        async_token: Option<SyntaxToken>,
        body: Option<ast::Expr>,
    ) -> Option<Vec<DocumentHighlight>> {
        let mut highlights = Vec::new();
        highlights.push(DocumentHighlight { access: None, range: async_token?.text_range() });
        if let Some(body) = body {
            walk(&body, |node| {
                match_ast! {
                    match node {
                        ast::AwaitExpr(expr) => if let Some(token) = expr.await_token() {
                            highlights.push(DocumentHighlight {
                                access: None,
                                range: token.text_range(),
                            });
                        },
                        // All the following are different contexts so skip them
                        ast::EffectExpr(effect) => return effect.async_token().is_some() || effect.try_token().is_some(),
                        ast::ClosureExpr(__) => return true,
                        ast::Item(__) => return true,
                        // Don't look into const args
                        ast::Path(__) => return true,
                        _ => (),
                    }
                }
                false
            });
        }
        Some(highlights)
    }
    for anc in token.ancestors() {
        return match_ast! {
            match anc {
                ast::Fn(fn_) => hl(fn_.async_token(), fn_.body().map(ast::Expr::BlockExpr)),
                ast::EffectExpr(effect) => hl(effect.async_token(), effect.block_expr().map(ast::Expr::BlockExpr)),
                ast::ClosureExpr(closure) => hl(closure.async_token(), closure.body()),
                _ => continue,
            }
        };
    }
    None
}

/// Preorder walk the expression node skipping a node's subtrees if the callback returns `true` for the node.
fn walk(expr: &ast::Expr, mut cb: impl FnMut(SyntaxNode) -> bool) {
    let mut preorder = expr.syntax().preorder();
    while let Some(event) = preorder.next() {
        let node = match event {
            WalkEvent::Enter(node) => node,
            WalkEvent::Leave(_) => continue,
        };
        if cb(node) {
            preorder.skip_subtree();
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::fixture;

    use super::*;

    fn check(ra_fixture: &str) {
        let (analysis, pos, annotations) = fixture::annotations(ra_fixture);
        let hls = analysis.highlight_related(pos).unwrap().unwrap();

        let mut expected = annotations
            .into_iter()
            .map(|(r, access)| (r.range, (!access.is_empty()).then(|| access)))
            .collect::<Vec<_>>();

        let mut actual = hls
            .into_iter()
            .map(|hl| {
                (
                    hl.range,
                    hl.access.map(|it| {
                        match it {
                            ReferenceAccess::Read => "read",
                            ReferenceAccess::Write => "write",
                        }
                        .to_string()
                    }),
                )
            })
            .collect::<Vec<_>>();
        actual.sort_by_key(|(range, _)| range.start());
        expected.sort_by_key(|(range, _)| range.start());

        assert_eq!(expected, actual);
    }

    #[test]
    fn test_hl_module() {
        check(
            r#"
//- /lib.rs
mod foo$0;
 // ^^^
//- /foo.rs
struct Foo;
"#,
        );
    }

    #[test]
    fn test_hl_self_in_crate_root() {
        check(
            r#"
use self$0;
"#,
        );
    }

    #[test]
    fn test_hl_self_in_module() {
        check(
            r#"
//- /lib.rs
mod foo;
//- /foo.rs
use self$0;
"#,
        );
    }

    #[test]
    fn test_hl_local() {
        check(
            r#"
fn foo() {
    let mut bar = 3;
         // ^^^ write
    bar$0;
 // ^^^ read
}
"#,
        );
    }

    #[test]
    fn test_hl_yield_points() {
        check(
            r#"
pub async fn foo() {
 // ^^^^^
    let x = foo()
        .await$0
      // ^^^^^
        .await;
      // ^^^^^
    || { 0.await };
    (async { 0.await }).await
                     // ^^^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_yield_points2() {
        check(
            r#"
pub async$0 fn foo() {
 // ^^^^^
    let x = foo()
        .await
      // ^^^^^
        .await;
      // ^^^^^
    || { 0.await };
    (async { 0.await }).await
                     // ^^^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_yield_nested_fn() {
        check(
            r#"
async fn foo() {
    async fn foo2() {
 // ^^^^^
        async fn foo3() {
            0.await
        }
        0.await$0
       // ^^^^^
    }
    0.await
}
"#,
        );
    }

    #[test]
    fn test_hl_yield_nested_async_blocks() {
        check(
            r#"
async fn foo() {
    (async {
  // ^^^^^
        (async {
           0.await
        }).await$0 }
        // ^^^^^
    ).await;
}
"#,
        );
    }

    #[test]
    fn test_hl_exit_points() {
        check(
            r#"
fn foo() -> u32 {
    if true {
        return$0 0;
     // ^^^^^^
    }

    0?;
  // ^
    0xDEAD_BEEF
 // ^^^^^^^^^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_exit_points2() {
        check(
            r#"
fn foo() ->$0 u32 {
    if true {
        return 0;
     // ^^^^^^
    }

    0?;
  // ^
    0xDEAD_BEEF
 // ^^^^^^^^^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_prefer_ref_over_tail_exit() {
        check(
            r#"
fn foo() -> u32 {
// ^^^
    if true {
        return 0;
    }

    0?;

    foo$0()
 // ^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_never_call_is_exit_point() {
        check(
            r#"
struct Never;
impl Never {
    fn never(self) -> ! { loop {} }
}
macro_rules! never {
    () => { never() }
}
fn never() -> ! { loop {} }
fn foo() ->$0 u32 {
    never();
 // ^^^^^^^
    never!();
 // FIXME sema doesnt give us types for macrocalls

    Never.never();
 // ^^^^^^^^^^^^^

    0
 // ^
}
"#,
        );
    }
}
<commit_msg>Improve feature docs for highlight_related<commit_after>use hir::Semantics;
use ide_db::{
    base_db::FilePosition,
    defs::Definition,
    helpers::pick_best_token,
    search::{FileReference, ReferenceAccess, SearchScope},
    RootDatabase,
};
use syntax::{
    ast, match_ast, AstNode,
    SyntaxKind::{ASYNC_KW, AWAIT_KW, QUESTION, RETURN_KW, THIN_ARROW},
    SyntaxNode, SyntaxToken, TextRange, WalkEvent,
};

use crate::{display::TryToNav, references, NavigationTarget};

pub struct DocumentHighlight {
    pub range: TextRange,
    pub access: Option<ReferenceAccess>,
}

// Feature: Highlight Related
//
// Highlights constructs related to the thing under the cursor:
// - if on an identifier, highlights all references to that identifier in the current file
// - if on an `async` or `await token, highlights all yield points for that async context
// - if on a `return` token, `?` character or `->` return type arrow, highlights all exit points for that context
pub(crate) fn highlight_related(
    sema: &Semantics<RootDatabase>,
    position: FilePosition,
) -> Option<Vec<DocumentHighlight>> {
    let _p = profile::span("document_highlight");
    let syntax = sema.parse(position.file_id).syntax().clone();

    let token = pick_best_token(syntax.token_at_offset(position.offset), |kind| match kind {
        QUESTION => 2, // prefer `?` when the cursor is sandwiched like `await$0?`
        AWAIT_KW | ASYNC_KW | THIN_ARROW | RETURN_KW => 1,
        _ => 0,
    })?;

    match token.kind() {
        QUESTION | RETURN_KW | THIN_ARROW => highlight_exit_points(sema, token),
        AWAIT_KW | ASYNC_KW => highlight_yield_points(token),
        _ => highlight_references(sema, &syntax, position),
    }
}

fn highlight_references(
    sema: &Semantics<RootDatabase>,
    syntax: &SyntaxNode,
    FilePosition { offset, file_id }: FilePosition,
) -> Option<Vec<DocumentHighlight>> {
    let def = references::find_def(sema, syntax, offset)?;
    let usages = def.usages(sema).set_scope(Some(SearchScope::single_file(file_id))).all();

    let declaration = match def {
        Definition::ModuleDef(hir::ModuleDef::Module(module)) => {
            Some(NavigationTarget::from_module_to_decl(sema.db, module))
        }
        def => def.try_to_nav(sema.db),
    }
    .filter(|decl| decl.file_id == file_id)
    .and_then(|decl| {
        let range = decl.focus_range?;
        let access = references::decl_access(&def, syntax, range);
        Some(DocumentHighlight { range, access })
    });

    let file_refs = usages.references.get(&file_id).map_or(&[][..], Vec::as_slice);
    let mut res = Vec::with_capacity(file_refs.len() + 1);
    res.extend(declaration);
    res.extend(
        file_refs
            .iter()
            .map(|&FileReference { access, range, .. }| DocumentHighlight { range, access }),
    );
    Some(res)
}

fn highlight_exit_points(
    sema: &Semantics<RootDatabase>,
    token: SyntaxToken,
) -> Option<Vec<DocumentHighlight>> {
    fn hl(
        sema: &Semantics<RootDatabase>,
        body: Option<ast::Expr>,
    ) -> Option<Vec<DocumentHighlight>> {
        let mut highlights = Vec::new();
        let body = body?;
        walk(&body, |node| {
            match_ast! {
                match node {
                    ast::ReturnExpr(expr) => if let Some(token) = expr.return_token() {
                        highlights.push(DocumentHighlight {
                            access: None,
                            range: token.text_range(),
                        });
                    },
                    ast::TryExpr(try_) => if let Some(token) = try_.question_mark_token() {
                        highlights.push(DocumentHighlight {
                            access: None,
                            range: token.text_range(),
                        });
                    },
                    ast::Expr(expr) => match expr {
                        ast::Expr::MethodCallExpr(_) | ast::Expr::CallExpr(_) | ast::Expr::MacroCall(_) => {
                            if sema.type_of_expr(&expr).map_or(false, |ty| ty.is_never()) {
                                highlights.push(DocumentHighlight {
                                    access: None,
                                    range: expr.syntax().text_range(),
                                });
                            }
                        },
                        ast::Expr::EffectExpr(effect) => return effect.async_token().is_some() || effect.try_token().is_some(),
                        ast::Expr::ClosureExpr(_) => return true,
                        _ => (),
                    },
                    ast::Item(__) => return true,
                    // Don't look into const args
                    ast::Path(__) => return true,
                    _ => (),
                }
            }
            false
        });
        let tail = match body {
            ast::Expr::BlockExpr(b) => b.tail_expr(),
            e => Some(e),
        };
        if let Some(tail) = tail {
            highlights.push(DocumentHighlight { access: None, range: tail.syntax().text_range() });
        }
        Some(highlights)
    }
    for anc in token.ancestors() {
        return match_ast! {
            match anc {
                ast::Fn(fn_) => hl(sema, fn_.body().map(ast::Expr::BlockExpr)),
                ast::ClosureExpr(closure) => hl(sema, closure.body()),
                ast::EffectExpr(effect) => if effect.async_token().is_some() || effect.try_token().is_some() {
                    hl(sema, effect.block_expr().map(ast::Expr::BlockExpr))
                } else {
                    continue;
                },
                _ => continue,
            }
        };
    }
    None
}

fn highlight_yield_points(token: SyntaxToken) -> Option<Vec<DocumentHighlight>> {
    fn hl(
        async_token: Option<SyntaxToken>,
        body: Option<ast::Expr>,
    ) -> Option<Vec<DocumentHighlight>> {
        let mut highlights = Vec::new();
        highlights.push(DocumentHighlight { access: None, range: async_token?.text_range() });
        if let Some(body) = body {
            walk(&body, |node| {
                match_ast! {
                    match node {
                        ast::AwaitExpr(expr) => if let Some(token) = expr.await_token() {
                            highlights.push(DocumentHighlight {
                                access: None,
                                range: token.text_range(),
                            });
                        },
                        // All the following are different contexts so skip them
                        ast::EffectExpr(effect) => return effect.async_token().is_some() || effect.try_token().is_some(),
                        ast::ClosureExpr(__) => return true,
                        ast::Item(__) => return true,
                        // Don't look into const args
                        ast::Path(__) => return true,
                        _ => (),
                    }
                }
                false
            });
        }
        Some(highlights)
    }
    for anc in token.ancestors() {
        return match_ast! {
            match anc {
                ast::Fn(fn_) => hl(fn_.async_token(), fn_.body().map(ast::Expr::BlockExpr)),
                ast::EffectExpr(effect) => hl(effect.async_token(), effect.block_expr().map(ast::Expr::BlockExpr)),
                ast::ClosureExpr(closure) => hl(closure.async_token(), closure.body()),
                _ => continue,
            }
        };
    }
    None
}

/// Preorder walk the expression node skipping a node's subtrees if the callback returns `true` for the node.
fn walk(expr: &ast::Expr, mut cb: impl FnMut(SyntaxNode) -> bool) {
    let mut preorder = expr.syntax().preorder();
    while let Some(event) = preorder.next() {
        let node = match event {
            WalkEvent::Enter(node) => node,
            WalkEvent::Leave(_) => continue,
        };
        if cb(node) {
            preorder.skip_subtree();
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::fixture;

    use super::*;

    fn check(ra_fixture: &str) {
        let (analysis, pos, annotations) = fixture::annotations(ra_fixture);
        let hls = analysis.highlight_related(pos).unwrap().unwrap();

        let mut expected = annotations
            .into_iter()
            .map(|(r, access)| (r.range, (!access.is_empty()).then(|| access)))
            .collect::<Vec<_>>();

        let mut actual = hls
            .into_iter()
            .map(|hl| {
                (
                    hl.range,
                    hl.access.map(|it| {
                        match it {
                            ReferenceAccess::Read => "read",
                            ReferenceAccess::Write => "write",
                        }
                        .to_string()
                    }),
                )
            })
            .collect::<Vec<_>>();
        actual.sort_by_key(|(range, _)| range.start());
        expected.sort_by_key(|(range, _)| range.start());

        assert_eq!(expected, actual);
    }

    #[test]
    fn test_hl_module() {
        check(
            r#"
//- /lib.rs
mod foo$0;
 // ^^^
//- /foo.rs
struct Foo;
"#,
        );
    }

    #[test]
    fn test_hl_self_in_crate_root() {
        check(
            r#"
use self$0;
"#,
        );
    }

    #[test]
    fn test_hl_self_in_module() {
        check(
            r#"
//- /lib.rs
mod foo;
//- /foo.rs
use self$0;
"#,
        );
    }

    #[test]
    fn test_hl_local() {
        check(
            r#"
fn foo() {
    let mut bar = 3;
         // ^^^ write
    bar$0;
 // ^^^ read
}
"#,
        );
    }

    #[test]
    fn test_hl_yield_points() {
        check(
            r#"
pub async fn foo() {
 // ^^^^^
    let x = foo()
        .await$0
      // ^^^^^
        .await;
      // ^^^^^
    || { 0.await };
    (async { 0.await }).await
                     // ^^^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_yield_points2() {
        check(
            r#"
pub async$0 fn foo() {
 // ^^^^^
    let x = foo()
        .await
      // ^^^^^
        .await;
      // ^^^^^
    || { 0.await };
    (async { 0.await }).await
                     // ^^^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_yield_nested_fn() {
        check(
            r#"
async fn foo() {
    async fn foo2() {
 // ^^^^^
        async fn foo3() {
            0.await
        }
        0.await$0
       // ^^^^^
    }
    0.await
}
"#,
        );
    }

    #[test]
    fn test_hl_yield_nested_async_blocks() {
        check(
            r#"
async fn foo() {
    (async {
  // ^^^^^
        (async {
           0.await
        }).await$0 }
        // ^^^^^
    ).await;
}
"#,
        );
    }

    #[test]
    fn test_hl_exit_points() {
        check(
            r#"
fn foo() -> u32 {
    if true {
        return$0 0;
     // ^^^^^^
    }

    0?;
  // ^
    0xDEAD_BEEF
 // ^^^^^^^^^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_exit_points2() {
        check(
            r#"
fn foo() ->$0 u32 {
    if true {
        return 0;
     // ^^^^^^
    }

    0?;
  // ^
    0xDEAD_BEEF
 // ^^^^^^^^^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_prefer_ref_over_tail_exit() {
        check(
            r#"
fn foo() -> u32 {
// ^^^
    if true {
        return 0;
    }

    0?;

    foo$0()
 // ^^^
}
"#,
        );
    }

    #[test]
    fn test_hl_never_call_is_exit_point() {
        check(
            r#"
struct Never;
impl Never {
    fn never(self) -> ! { loop {} }
}
macro_rules! never {
    () => { never() }
}
fn never() -> ! { loop {} }
fn foo() ->$0 u32 {
    never();
 // ^^^^^^^
    never!();
 // FIXME sema doesnt give us types for macrocalls

    Never.never();
 // ^^^^^^^^^^^^^

    0
 // ^
}
"#,
        );
    }
}
<|endoftext|>


<commit_before>//! Handle process life-time and message passing for proc-macro client

use crossbeam_channel::{bounded, Receiver, Sender};
use ra_tt::Subtree;

use crate::msg::{ErrorCode, Message, Request, Response, ResponseError};
use crate::rpc::{ExpansionResult, ExpansionTask, ListMacrosResult, ListMacrosTask, ProcMacroKind};

use io::{BufRead, BufReader};
use std::{
    convert::{TryFrom, TryInto},
    io::{self, Write},
    path::{Path, PathBuf},
    process::{Child, Command, Stdio},
};

#[derive(Debug, Default)]
pub(crate) struct ProcMacroProcessSrv {
    inner: Option<Sender<Task>>,
}

#[derive(Debug)]
pub(crate) struct ProcMacroProcessThread {
    // XXX: drop order is significant
    sender: SenderGuard,
    handle: jod_thread::JoinHandle<()>,
}

#[derive(Debug)]
struct SenderGuard(pub Sender<Task>);

impl std::ops::Drop for SenderGuard {
    fn drop(&mut self) {
        self.0.send(Task::Close).unwrap();
    }
}

enum Task {
    Request { req: Request, result_tx: Sender<Response> },
    Close,
}

struct Process {
    path: PathBuf,
    child: Child,
}

impl Process {
    fn run(process_path: &Path) -> Result<Process, io::Error> {
        let child = Command::new(process_path.clone())
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .spawn()?;

        Ok(Process { path: process_path.into(), child })
    }

    fn restart(&mut self) -> Result<(), io::Error> {
        let _ = self.child.kill();
        self.child =
            Command::new(self.path.clone()).stdin(Stdio::piped()).stdout(Stdio::piped()).spawn()?;
        Ok(())
    }

    fn stdio(&mut self) -> Option<(impl Write, impl BufRead)> {
        let stdin = self.child.stdin.take()?;
        let stdout = self.child.stdout.take()?;
        let read = BufReader::new(stdout);

        Some((stdin, read))
    }
}

impl ProcMacroProcessSrv {
    pub fn run(
        process_path: &Path,
    ) -> Result<(ProcMacroProcessThread, ProcMacroProcessSrv), io::Error> {
        let process = Process::run(process_path)?;

        let (task_tx, task_rx) = bounded(0);
        let handle = jod_thread::spawn(move || {
            client_loop(task_rx, process);
        });

        let srv = ProcMacroProcessSrv { inner: Some(task_tx.clone()) };
        let thread = ProcMacroProcessThread { handle, sender: SenderGuard(task_tx) };

        Ok((thread, srv))
    }

    pub fn find_proc_macros(
        &self,
        dylib_path: &Path,
    ) -> Result<Vec<(String, ProcMacroKind)>, ra_tt::ExpansionError> {
        let task = ListMacrosTask { lib: dylib_path.to_path_buf() };

        let result: ListMacrosResult = self.send_task(Request::ListMacro(task))?;
        Ok(result.macros)
    }

    pub fn custom_derive(
        &self,
        dylib_path: &Path,
        subtree: &Subtree,
        derive_name: &str,
    ) -> Result<Subtree, ra_tt::ExpansionError> {
        let task = ExpansionTask {
            macro_body: subtree.clone(),
            macro_name: derive_name.to_string(),
            attributes: None,
            lib: dylib_path.to_path_buf(),
        };

        let result: ExpansionResult = self.send_task(Request::ExpansionMacro(task))?;
        Ok(result.expansion)
    }

    pub fn send_task<R>(&self, req: Request) -> Result<R, ra_tt::ExpansionError>
    where
        R: TryFrom<Response, Error = &'static str>,
    {
        let sender = match &self.inner {
            None => return Err(ra_tt::ExpansionError::Unknown("No sender is found.".to_string())),
            Some(it) => it,
        };

        let (result_tx, result_rx) = bounded(0);

        sender.send(Task::Request { req: req.into(), result_tx }).unwrap();

        let res = result_rx.recv().unwrap();
        match res {
            Response::Error(err) => {
                return Err(ra_tt::ExpansionError::ExpansionError(err.message));
            }
            _ => Ok(res.try_into().map_err(|err| {
                ra_tt::ExpansionError::Unknown(format!(
                    "Fail to get response, reason : {:#?} ",
                    err
                ))
            })?),
        }
    }
}

fn client_loop(task_rx: Receiver<Task>, mut process: Process) {
    let (mut stdin, mut stdout) = match process.stdio() {
        None => return,
        Some(it) => it,
    };

    loop {
        let task = match task_rx.recv() {
            Ok(task) => task,
            Err(_) => break,
        };

        let (req, result_tx) = match task {
            Task::Request { req, result_tx } => (req, result_tx),
            Task::Close => break,
        };

        let res = match send_request(&mut stdin, &mut stdout, req) {
            Ok(res) => res,
            Err(_err) => {
                let res = Response::Error(ResponseError {
                    code: ErrorCode::ServerErrorEnd,
                    message: "Server closed".into(),
                });
                result_tx.send(res.into()).unwrap();
                // Restart the process
                if process.restart().is_err() {
                    break;
                }
                let stdio = match process.stdio() {
                    None => break,
                    Some(it) => it,
                };
                stdin = stdio.0;
                stdout = stdio.1;
                continue;
            }
        };

        if let Some(res) = res {
            result_tx.send(res).unwrap();
        }
    }

    let _ = process.child.kill();
}

fn send_request(
    mut writer: &mut impl Write,
    mut reader: &mut impl BufRead,
    req: Request,
) -> Result<Option<Response>, io::Error> {
    req.write(&mut writer)?;
    Ok(Response::read(&mut reader)?)
}
<commit_msg>Add drop for process<commit_after>//! Handle process life-time and message passing for proc-macro client

use crossbeam_channel::{bounded, Receiver, Sender};
use ra_tt::Subtree;

use crate::msg::{ErrorCode, Message, Request, Response, ResponseError};
use crate::rpc::{ExpansionResult, ExpansionTask, ListMacrosResult, ListMacrosTask, ProcMacroKind};

use io::{BufRead, BufReader};
use std::{
    convert::{TryFrom, TryInto},
    io::{self, Write},
    path::{Path, PathBuf},
    process::{Child, Command, Stdio},
};

#[derive(Debug, Default)]
pub(crate) struct ProcMacroProcessSrv {
    inner: Option<Sender<Task>>,
}

#[derive(Debug)]
pub(crate) struct ProcMacroProcessThread {
    // XXX: drop order is significant
    sender: SenderGuard,
    handle: jod_thread::JoinHandle<()>,
}

#[derive(Debug)]
struct SenderGuard(pub Sender<Task>);

impl std::ops::Drop for SenderGuard {
    fn drop(&mut self) {
        self.0.send(Task::Close).unwrap();
    }
}

enum Task {
    Request { req: Request, result_tx: Sender<Response> },
    Close,
}

struct Process {
    path: PathBuf,
    child: Child,
}

impl Drop for Process {
    fn drop(&mut self) {
        let _ = self.child.kill();
    }
}

impl Process {
    fn run(process_path: &Path) -> Result<Process, io::Error> {
        let child = Command::new(process_path.clone())
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .spawn()?;

        Ok(Process { path: process_path.into(), child })
    }

    fn restart(&mut self) -> Result<(), io::Error> {
        let _ = self.child.kill();
        self.child =
            Command::new(self.path.clone()).stdin(Stdio::piped()).stdout(Stdio::piped()).spawn()?;
        Ok(())
    }

    fn stdio(&mut self) -> Option<(impl Write, impl BufRead)> {
        let stdin = self.child.stdin.take()?;
        let stdout = self.child.stdout.take()?;
        let read = BufReader::new(stdout);

        Some((stdin, read))
    }
}

impl ProcMacroProcessSrv {
    pub fn run(
        process_path: &Path,
    ) -> Result<(ProcMacroProcessThread, ProcMacroProcessSrv), io::Error> {
        let process = Process::run(process_path)?;

        let (task_tx, task_rx) = bounded(0);
        let handle = jod_thread::spawn(move || {
            client_loop(task_rx, process);
        });

        let srv = ProcMacroProcessSrv { inner: Some(task_tx.clone()) };
        let thread = ProcMacroProcessThread { handle, sender: SenderGuard(task_tx) };

        Ok((thread, srv))
    }

    pub fn find_proc_macros(
        &self,
        dylib_path: &Path,
    ) -> Result<Vec<(String, ProcMacroKind)>, ra_tt::ExpansionError> {
        let task = ListMacrosTask { lib: dylib_path.to_path_buf() };

        let result: ListMacrosResult = self.send_task(Request::ListMacro(task))?;
        Ok(result.macros)
    }

    pub fn custom_derive(
        &self,
        dylib_path: &Path,
        subtree: &Subtree,
        derive_name: &str,
    ) -> Result<Subtree, ra_tt::ExpansionError> {
        let task = ExpansionTask {
            macro_body: subtree.clone(),
            macro_name: derive_name.to_string(),
            attributes: None,
            lib: dylib_path.to_path_buf(),
        };

        let result: ExpansionResult = self.send_task(Request::ExpansionMacro(task))?;
        Ok(result.expansion)
    }

    pub fn send_task<R>(&self, req: Request) -> Result<R, ra_tt::ExpansionError>
    where
        R: TryFrom<Response, Error = &'static str>,
    {
        let sender = match &self.inner {
            None => return Err(ra_tt::ExpansionError::Unknown("No sender is found.".to_string())),
            Some(it) => it,
        };

        let (result_tx, result_rx) = bounded(0);

        sender.send(Task::Request { req: req.into(), result_tx }).unwrap();

        let res = result_rx.recv().unwrap();
        match res {
            Response::Error(err) => {
                return Err(ra_tt::ExpansionError::ExpansionError(err.message));
            }
            _ => Ok(res.try_into().map_err(|err| {
                ra_tt::ExpansionError::Unknown(format!(
                    "Fail to get response, reason : {:#?} ",
                    err
                ))
            })?),
        }
    }
}

fn client_loop(task_rx: Receiver<Task>, mut process: Process) {
    let (mut stdin, mut stdout) = match process.stdio() {
        None => return,
        Some(it) => it,
    };

    loop {
        let task = match task_rx.recv() {
            Ok(task) => task,
            Err(_) => break,
        };

        let (req, result_tx) = match task {
            Task::Request { req, result_tx } => (req, result_tx),
            Task::Close => break,
        };

        let res = match send_request(&mut stdin, &mut stdout, req) {
            Ok(res) => res,
            Err(_err) => {
                let res = Response::Error(ResponseError {
                    code: ErrorCode::ServerErrorEnd,
                    message: "Server closed".into(),
                });
                result_tx.send(res.into()).unwrap();
                // Restart the process
                if process.restart().is_err() {
                    break;
                }
                let stdio = match process.stdio() {
                    None => break,
                    Some(it) => it,
                };
                stdin = stdio.0;
                stdout = stdio.1;
                continue;
            }
        };

        if let Some(res) = res {
            result_tx.send(res).unwrap();
        }
    }
}

fn send_request(
    mut writer: &mut impl Write,
    mut reader: &mut impl BufRead,
    req: Request,
) -> Result<Option<Response>, io::Error> {
    req.write(&mut writer)?;
    Ok(Response::read(&mut reader)?)
}
<|endoftext|>


<commit_before>//   Copyright 2014 Colin Sherratt
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.

#![crate_name = "snowmew-position"]
#![license = "ASL2"]
#![crate_type = "lib"]
#![comment = "A position manager for snowmew"]
#![feature(macro_rules)]
#![feature(tuple_indexing)]

extern crate "snowmew-core" as snowmew;
extern crate cgmath;
extern crate opencl;
extern crate cow;
extern crate time;
extern crate serialize;

use std::default::Default;
use serialize::Encodable;

use cgmath::{Transform, Decomposed};
use cgmath::Quaternion;
use cgmath::{Vector3, Vector4};
use cgmath::{Matrix4, ToMatrix4, Matrix};

use opencl::hl::{Device, Context, CommandQueue, Kernel, Event};
use opencl::mem::CLBuffer;
use opencl::cl::{CL_MEM_READ_ONLY};

use snowmew::common::{Entity, Common, Duplicate, Delete};
use snowmew::input_integrator::InputIntegratorGameData;
use snowmew::debugger::DebuggerGameData;
use snowmew::table::{Dynamic, DynamicIterator};

const OPENCL_PROGRAM: &'static str = include_str!("position.c");

pub trait MatrixManager {
    fn size(&mut self, size: uint);
    fn set(&mut self, idx: uint, mat: Matrix4<f32>);
    fn get(&self, idx: uint) -> Matrix4<f32>;
}

impl<'r> MatrixManager for &'r mut Vec<Matrix4<f32>> {
    fn size(&mut self, size: uint) {
        if self.len() < size {
            let amount = self.len() - size;
            self.grow(amount, Matrix4::identity())
        }
    }
    fn set(&mut self, idx: uint, m: Matrix4<f32>) { self[idx] = m; }
    fn get(&self, idx: uint) -> Matrix4<f32> { self[idx] }
}

#[deriving(Encodable, Decodable)]
pub struct Delta {
    pub parent: Option<Entity>,
    pub delta: Decomposed<f32, Vector3<f32>, Quaternion<f32>>
}

impl Clone for Delta {
    fn clone(&self) -> Delta {
        Delta {
            parent: self.parent,
            delta: self.delta
        }
    }
}

impl Default for Delta {
    fn default() -> Delta {
        Delta {
            parent: None,
            delta: Transform::identity()
        }
    }
}

#[deriving(Clone, Encodable, Decodable)]
pub struct PositionData {
    delta: Dynamic<Delta>,
}

impl PositionData {
    pub fn new() -> PositionData {
        PositionData {
            delta: Dynamic::new()
        }
    }
}

pub trait Positions {
    fn get_position<'a>(&'a self) -> &'a PositionData;
    fn get_position_mut<'a>(&'a mut self) -> &'a mut PositionData;

    fn set_delta(&mut self,
                 key: Entity,
                 parent: Option<Entity>,
                 transform: Decomposed<f32, Vector3<f32>, Quaternion<f32>>) {
        self.get_position_mut().delta.insert(key,
            Delta {
                parent: parent,
                delta: transform
            }
        );
    }

    fn set_to_identity(&mut self, key: Entity) {
        self.set_delta(key, None, Transform::identity())
    }

    fn init_position(&mut self, key: Entity) {
        if self.get_position().delta.get(key).is_none() {
            self.set_to_identity(key);
        }
    }

    fn set_scale(&mut self, key: Entity, scale: f32) {
        self.init_position(key);
        self.get_position_mut()
            .delta.get_mut(key)
            .map(|d| d.delta.scale = scale);
    }

    fn set_displacement(&mut self, key: Entity, disp: Vector3<f32>) {
        self.init_position(key);
        self.get_position_mut()
            .delta.get_mut(key)
            .map(|d| d.delta.disp = disp);
    }

    fn set_rotation(&mut self, key: Entity, rot: Quaternion<f32>) {
        self.init_position(key);
        self.get_position_mut()
            .delta.get_mut(key)
            .map(|d| d.delta.rot = rot);
    }

    fn get_scale(&mut self, key: Entity) -> Option<&f32> {
        self.get_position()
            .delta.get(key)
            .map(|d| &d.delta.scale)
    }

    fn get_displacement(&mut self, key: Entity) -> Option<&Vector3<f32>> {
        self.get_position()
            .delta.get(key)
            .map(|d| &d.delta.disp)
    }

    fn get_rotation(&mut self, key: Entity) -> Option<&Quaternion<f32>> {
        self.get_position()
            .delta.get(key)
            .map(|d| &d.delta.rot)
    }

    fn get_mut_scale(&mut self, key: Entity) -> Option<&mut f32> {
        self.get_position_mut()
            .delta.get_mut(key)
            .map(|d| &mut d.delta.scale)
    }

    fn get_mut_displacement(&mut self, key: Entity) -> Option<&mut Vector3<f32>> {
        self.get_position_mut()
            .delta.get_mut(key)
            .map(|d| &mut d.delta.disp)
    }

    fn get_mut_rotation(&mut self, key: Entity) -> Option<&mut Quaternion<f32>> {
        self.get_position_mut()
            .delta.get_mut(key)
            .map(|d| &mut d.delta.rot)
    }

    fn get_transform(&self, key: Entity) -> Option<Decomposed<f32, Vector3<f32>, Quaternion<f32>>> {
        self.get_position()
            .delta.get(key)
            .map(|d| d.delta)
    }

    fn get_parent(&self, key: Entity) -> Option<&Option<Entity>> {
        self.get_position()
            .delta.get(key)
            .map(|d| &d.parent)
    }

    fn position(&self, key: Entity) -> Matrix4<f32> {
        self.get_position()
            .delta.get(key)
            .map(|d| {
                let matrix = d.delta.to_matrix4();
                match d.parent {
                    Some(p) => {
                        let parent = self.position(p);
                        parent.mul_m(&matrix)
                    }
                    None => {
                        matrix
                    }
                }
            })
            .unwrap_or_else(|| Matrix4::identity())
    }

    fn write_positions<MM: MatrixManager>(&self, mm: &mut MM) {
        mm.size(self.get_position().delta.highest_entity() as uint);
        for (key, _) in self.get_position().delta.iter() {
            mm.set(key as uint, self.position(key));
        }
    }

    fn position_iter(&self) -> PositionIter {
        let pos = self.get_position();
        PositionIter {
            pos: pos,
            iter: pos.delta.iter()
        }
    }

    /*fn write_positions_cl_vec4x4(&self, cq: &CommandQueue,
                        ctx: &mut CalcPositionsCl, out: &[CLBuffer<Vector4<f32>>, ..4]) -> Event {
        self.get_position().position.write_positions_cl_vec4x4(cq, ctx, out)
    }

    fn write_positions_cl_mat4(&self, cq: &CommandQueue,
                        ctx: &mut CalcPositionsCl, out: &[CLBuffer<Matrix4<f32>>]) -> Event {
        self.get_position().position.write_positions_cl_mat4(cq, ctx, out)
    }*/
}

pub struct PositionIter<'a> {
    pos: &'a PositionData,
    iter: DynamicIterator<'a, Delta>
}

impl<'a> Iterator<(Entity, Matrix4<f32>)> for PositionIter<'a> {
    fn next(&mut self) -> Option<(Entity, Matrix4<f32>)> {
        self.iter.next().map(|(id, _)| (id, self.pos.position(id)) )
    }
}

impl Duplicate for PositionData {
    fn duplicate(&mut self, src: Entity, dst: Entity) {
        let delta = self.delta.get(src).map(|&x| x);
        delta.map(|delta| self.delta.insert(dst, delta));
    }
}

impl Delete for PositionData {
    fn delete(&mut self, key: Entity) -> bool {
        self.delta.remove(key)
    }
}

impl Positions for PositionData {
    fn get_position<'a>(&'a self) -> &'a PositionData { self }
    fn get_position_mut<'a>(&'a mut self) -> &'a mut PositionData { self }
}

impl<T: Positions> Positions for InputIntegratorGameData<T> {
    fn get_position<'a>(&'a self) -> &'a PositionData { self.inner.get_position() }
    fn get_position_mut<'a>(&'a mut self) -> &'a mut PositionData { self.inner.get_position_mut() }
}

impl<T: Positions> Positions for DebuggerGameData<T> {
    fn get_position<'a>(&'a self) -> &'a PositionData { self.inner.get_position() }
    fn get_position_mut<'a>(&'a mut self) -> &'a mut PositionData { self.inner.get_position_mut() }
}
<commit_msg>added back delta iter<commit_after>//   Copyright 2014 Colin Sherratt
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.

#![crate_name = "snowmew-position"]
#![license = "ASL2"]
#![crate_type = "lib"]
#![comment = "A position manager for snowmew"]
#![feature(macro_rules)]
#![feature(tuple_indexing)]

extern crate "snowmew-core" as snowmew;
extern crate cgmath;
extern crate opencl;
extern crate cow;
extern crate time;
extern crate serialize;

use std::default::Default;
use serialize::Encodable;

use cgmath::{Transform, Decomposed};
use cgmath::Quaternion;
use cgmath::{Vector3, Vector4};
use cgmath::{Matrix4, ToMatrix4, Matrix};

use opencl::hl::{Device, Context, CommandQueue, Kernel, Event};
use opencl::mem::CLBuffer;
use opencl::cl::{CL_MEM_READ_ONLY};

use snowmew::common::{Entity, Common, Duplicate, Delete};
use snowmew::input_integrator::InputIntegratorGameData;
use snowmew::debugger::DebuggerGameData;
use snowmew::table::{Dynamic, DynamicIterator};

const OPENCL_PROGRAM: &'static str = include_str!("position.c");

pub trait MatrixManager {
    fn size(&mut self, size: uint);
    fn set(&mut self, idx: uint, mat: Matrix4<f32>);
    fn get(&self, idx: uint) -> Matrix4<f32>;
}

impl<'r> MatrixManager for &'r mut Vec<Matrix4<f32>> {
    fn size(&mut self, size: uint) {
        if self.len() < size {
            let amount = self.len() - size;
            self.grow(amount, Matrix4::identity())
        }
    }
    fn set(&mut self, idx: uint, m: Matrix4<f32>) { self[idx] = m; }
    fn get(&self, idx: uint) -> Matrix4<f32> { self[idx] }
}

#[deriving(Encodable, Decodable)]
pub struct Delta {
    pub parent: Option<Entity>,
    pub delta: Decomposed<f32, Vector3<f32>, Quaternion<f32>>
}

impl Clone for Delta {
    fn clone(&self) -> Delta {
        Delta {
            parent: self.parent,
            delta: self.delta
        }
    }
}

impl Default for Delta {
    fn default() -> Delta {
        Delta {
            parent: None,
            delta: Transform::identity()
        }
    }
}

#[deriving(Clone, Encodable, Decodable)]
pub struct PositionData {
    delta: Dynamic<Delta>,
}

impl PositionData {
    pub fn new() -> PositionData {
        PositionData {
            delta: Dynamic::new()
        }
    }
}

pub trait Positions {
    fn get_position<'a>(&'a self) -> &'a PositionData;
    fn get_position_mut<'a>(&'a mut self) -> &'a mut PositionData;

    fn set_delta(&mut self,
                 key: Entity,
                 parent: Option<Entity>,
                 transform: Decomposed<f32, Vector3<f32>, Quaternion<f32>>) {
        self.get_position_mut().delta.insert(key,
            Delta {
                parent: parent,
                delta: transform
            }
        );
    }

    fn set_to_identity(&mut self, key: Entity) {
        self.set_delta(key, None, Transform::identity())
    }

    fn init_position(&mut self, key: Entity) {
        if self.get_position().delta.get(key).is_none() {
            self.set_to_identity(key);
        }
    }

    fn set_scale(&mut self, key: Entity, scale: f32) {
        self.init_position(key);
        self.get_position_mut()
            .delta.get_mut(key)
            .map(|d| d.delta.scale = scale);
    }

    fn set_displacement(&mut self, key: Entity, disp: Vector3<f32>) {
        self.init_position(key);
        self.get_position_mut()
            .delta.get_mut(key)
            .map(|d| d.delta.disp = disp);
    }

    fn set_rotation(&mut self, key: Entity, rot: Quaternion<f32>) {
        self.init_position(key);
        self.get_position_mut()
            .delta.get_mut(key)
            .map(|d| d.delta.rot = rot);
    }

    fn get_scale(&mut self, key: Entity) -> Option<&f32> {
        self.get_position()
            .delta.get(key)
            .map(|d| &d.delta.scale)
    }

    fn get_displacement(&mut self, key: Entity) -> Option<&Vector3<f32>> {
        self.get_position()
            .delta.get(key)
            .map(|d| &d.delta.disp)
    }

    fn get_rotation(&mut self, key: Entity) -> Option<&Quaternion<f32>> {
        self.get_position()
            .delta.get(key)
            .map(|d| &d.delta.rot)
    }

    fn get_mut_scale(&mut self, key: Entity) -> Option<&mut f32> {
        self.get_position_mut()
            .delta.get_mut(key)
            .map(|d| &mut d.delta.scale)
    }

    fn get_mut_displacement(&mut self, key: Entity) -> Option<&mut Vector3<f32>> {
        self.get_position_mut()
            .delta.get_mut(key)
            .map(|d| &mut d.delta.disp)
    }

    fn get_mut_rotation(&mut self, key: Entity) -> Option<&mut Quaternion<f32>> {
        self.get_position_mut()
            .delta.get_mut(key)
            .map(|d| &mut d.delta.rot)
    }

    fn get_transform(&self, key: Entity) -> Option<Decomposed<f32, Vector3<f32>, Quaternion<f32>>> {
        self.get_position()
            .delta.get(key)
            .map(|d| d.delta)
    }

    fn get_parent(&self, key: Entity) -> Option<&Option<Entity>> {
        self.get_position()
            .delta.get(key)
            .map(|d| &d.parent)
    }

    fn position(&self, key: Entity) -> Matrix4<f32> {
        self.get_position()
            .delta.get(key)
            .map(|d| {
                let matrix = d.delta.to_matrix4();
                match d.parent {
                    Some(p) => {
                        let parent = self.position(p);
                        parent.mul_m(&matrix)
                    }
                    None => {
                        matrix
                    }
                }
            })
            .unwrap_or_else(|| Matrix4::identity())
    }

    fn write_positions<MM: MatrixManager>(&self, mm: &mut MM) {
        mm.size(self.get_position().delta.highest_entity() as uint);
        for (key, _) in self.get_position().delta.iter() {
            mm.set(key as uint, self.position(key));
        }
    }

    fn delta_iter(&self) -> DynamicIterator<Delta> {
        self.get_position().delta.iter()
    }

    fn position_iter(&self) -> PositionIter {
        let pos = self.get_position();
        PositionIter {
            pos: pos,
            iter: pos.delta.iter()
        }
    }

    /*fn write_positions_cl_vec4x4(&self, cq: &CommandQueue,
                        ctx: &mut CalcPositionsCl, out: &[CLBuffer<Vector4<f32>>, ..4]) -> Event {
        self.get_position().position.write_positions_cl_vec4x4(cq, ctx, out)
    }

    fn write_positions_cl_mat4(&self, cq: &CommandQueue,
                        ctx: &mut CalcPositionsCl, out: &[CLBuffer<Matrix4<f32>>]) -> Event {
        self.get_position().position.write_positions_cl_mat4(cq, ctx, out)
    }*/
}

pub struct PositionIter<'a> {
    pos: &'a PositionData,
    iter: DynamicIterator<'a, Delta>
}

impl<'a> Iterator<(Entity, Matrix4<f32>)> for PositionIter<'a> {
    fn next(&mut self) -> Option<(Entity, Matrix4<f32>)> {
        self.iter.next().map(|(id, _)| (id, self.pos.position(id)) )
    }
}

impl Duplicate for PositionData {
    fn duplicate(&mut self, src: Entity, dst: Entity) {
        let delta = self.delta.get(src).map(|&x| x);
        delta.map(|delta| self.delta.insert(dst, delta));
    }
}

impl Delete for PositionData {
    fn delete(&mut self, key: Entity) -> bool {
        self.delta.remove(key)
    }
}

impl Positions for PositionData {
    fn get_position<'a>(&'a self) -> &'a PositionData { self }
    fn get_position_mut<'a>(&'a mut self) -> &'a mut PositionData { self }
}

impl<T: Positions> Positions for InputIntegratorGameData<T> {
    fn get_position<'a>(&'a self) -> &'a PositionData { self.inner.get_position() }
    fn get_position_mut<'a>(&'a mut self) -> &'a mut PositionData { self.inner.get_position_mut() }
}

impl<T: Positions> Positions for DebuggerGameData<T> {
    fn get_position<'a>(&'a self) -> &'a PositionData { self.inner.get_position() }
    fn get_position_mut<'a>(&'a mut self) -> &'a mut PositionData { self.inner.get_position_mut() }
}
<|endoftext|>


<commit_before>use alga::general::{ClosedAdd, ClosedMul};
use num::{One, Zero};
use std::iter;
use std::marker::PhantomData;
use std::ops::{Add, Mul, Range};
use std::slice;

use allocator::Allocator;
use constraint::{AreMultipliable, DimEq, SameNumberOfRows, ShapeConstraint};
use sparse::{CsMatrix, CsStorage, CsVector};
use storage::{Storage, StorageMut};
use {DefaultAllocator, Dim, Matrix, MatrixMN, Real, Scalar, Vector, VectorN, U1};

impl<N: Scalar, R: Dim, C: Dim, S: CsStorage<N, R, C>> CsMatrix<N, R, C, S> {
    fn scatter<R2: Dim, C2: Dim>(
        &self,
        j: usize,
        beta: N,
        timestamps: &mut [usize],
        timestamp: usize,
        workspace: &mut [N],
        mut nz: usize,
        res: &mut CsMatrix<N, R2, C2>,
    ) -> usize
    where
        N: ClosedAdd + ClosedMul,
        DefaultAllocator: Allocator<usize, C2>,
    {
        for (i, val) in self.data.column_entries(j) {
            if timestamps[i] < timestamp {
                timestamps[i] = timestamp;
                res.data.i[nz] = i;
                nz += 1;
                workspace[i] = val * beta;
            } else {
                workspace[i] += val * beta;
            }
        }

        nz
    }
}

/*
impl<N: Scalar, R, S> CsVector<N, R, S> {
    pub fn axpy(&mut self, alpha: N, x: CsVector<N, R, S>, beta: N) {
        // First, compute the number of non-zero entries.
        let mut nnzero = 0;

        // Allocate a size large enough.
        self.data.set_column_len(0, nnzero);

        // Fill with the axpy.
        let mut i = self.len();
        let mut j = x.len();
        let mut k = nnzero - 1;
        let mut rid1 = self.data.row_index(0, i - 1);
        let mut rid2 = x.data.row_index(0, j - 1);

        while k > 0 {
            if rid1 == rid2 {
                self.data.set_row_index(0, k, rid1);
                self[k] = alpha * x[j] + beta * self[k];
                i -= 1;
                j -= 1;
            } else if rid1 < rid2 {
                self.data.set_row_index(0, k, rid1);
                self[k] = beta * self[i];
                i -= 1;
            } else {
                self.data.set_row_index(0, k, rid2);
                self[k] = alpha * x[j];
                j -= 1;
            }

            k -= 1;
        }
    }
}
*/

impl<N: Scalar + Zero + ClosedAdd + ClosedMul, D: Dim, S: StorageMut<N, D>> Vector<N, D, S> {
    pub fn axpy_cs<D2: Dim, S2>(&mut self, alpha: N, x: &CsVector<N, D2, S2>, beta: N)
    where
        S2: CsStorage<N, D2>,
        ShapeConstraint: DimEq<D, D2>,
    {
        if beta.is_zero() {
            for i in 0..x.len() {
                unsafe {
                    let k = x.data.row_index_unchecked(i);
                    let y = self.vget_unchecked_mut(k);
                    *y = alpha * *x.data.get_value_unchecked(i);
                }
            }
        } else {
            // Needed to be sure even components not present on `x` are multiplied.
            *self *= beta;

            for i in 0..x.len() {
                unsafe {
                    let k = x.data.row_index_unchecked(i);
                    let y = self.vget_unchecked_mut(k);
                    *y += alpha * *x.data.get_value_unchecked(i);
                }
            }
        }
    }

    /*
    pub fn gemv_sparse<R2: Dim, C2: Dim, S2>(&mut self, alpha: N, a: &CsMatrix<N, R2, C2, S2>, x: &DVector<N>, beta: N)
        where
            S2: CsStorage<N, R2, C2> {
        let col2 = a.column(0);
        let val = unsafe { *x.vget_unchecked(0) };
        self.axpy_sparse(alpha * val, &col2, beta);

        for j in 1..ncols2 {
            let col2 = a.column(j);
            let val = unsafe { *x.vget_unchecked(j) };

            self.axpy_sparse(alpha * val, &col2, N::one());
        }
    }
    */
}

impl<'a, 'b, N, R1, R2, C1, C2, S1, S2> Mul<&'b CsMatrix<N, R2, C2, S2>>
    for &'a CsMatrix<N, R1, C1, S1>
where
    N: Scalar + ClosedAdd + ClosedMul + Zero,
    R1: Dim,
    C1: Dim,
    R2: Dim,
    C2: Dim,
    S1: CsStorage<N, R1, C1>,
    S2: CsStorage<N, R2, C2>,
    ShapeConstraint: AreMultipliable<R1, C1, R2, C2>,
    DefaultAllocator: Allocator<usize, C2> + Allocator<usize, R1> + Allocator<N, R1>,
{
    type Output = CsMatrix<N, R1, C2>;

    fn mul(self, rhs: &'b CsMatrix<N, R2, C2, S2>) -> CsMatrix<N, R1, C2> {
        let (nrows1, ncols1) = self.data.shape();
        let (nrows2, ncols2) = rhs.data.shape();
        assert_eq!(
            ncols1.value(),
            nrows2.value(),
            "Mismatched dimensions for matrix multiplication."
        );

        let mut res = CsMatrix::new_uninitialized_generic(nrows1, ncols2, self.len() + rhs.len());
        let mut timestamps = VectorN::zeros_generic(nrows1, U1);
        let mut workspace = unsafe { VectorN::new_uninitialized_generic(nrows1, U1) };
        let mut nz = 0;

        for j in 0..ncols2.value() {
            res.data.p[j] = nz;
            let new_size_bound = nz + nrows1.value();
            res.data.i.resize(new_size_bound, 0);
            res.data.vals.resize(new_size_bound, N::zero());

            for (i, val) in rhs.data.column_entries(j) {
                nz = self.scatter(
                    i,
                    val,
                    timestamps.as_mut_slice(),
                    j + 1,
                    workspace.as_mut_slice(),
                    nz,
                    &mut res,
                );
            }

            // Keep the output sorted.
            let range = res.data.p[j]..nz;
            res.data.i[range.clone()].sort();

            for p in range {
                res.data.vals[p] = workspace[res.data.i[p]]
            }
        }

        res.data.i.truncate(nz);
        res.data.i.shrink_to_fit();
        res.data.vals.truncate(nz);
        res.data.vals.shrink_to_fit();
        res
    }
}

impl<'a, 'b, N, R1, R2, C1, C2, S1, S2> Add<&'b CsMatrix<N, R2, C2, S2>>
    for &'a CsMatrix<N, R1, C1, S1>
where
    N: Scalar + ClosedAdd + ClosedMul + One,
    R1: Dim,
    C1: Dim,
    R2: Dim,
    C2: Dim,
    S1: CsStorage<N, R1, C1>,
    S2: CsStorage<N, R2, C2>,
    ShapeConstraint: DimEq<R1, R2> + DimEq<C1, C2>,
    DefaultAllocator: Allocator<usize, C2> + Allocator<usize, R1> + Allocator<N, R1>,
{
    type Output = CsMatrix<N, R1, C2>;

    fn add(self, rhs: &'b CsMatrix<N, R2, C2, S2>) -> CsMatrix<N, R1, C2> {
        let (nrows1, ncols1) = self.data.shape();
        let (nrows2, ncols2) = rhs.data.shape();
        assert_eq!(
            (nrows1.value(), ncols1.value()),
            (nrows2.value(), ncols2.value()),
            "Mismatched dimensions for matrix sum."
        );

        let mut res = CsMatrix::new_uninitialized_generic(nrows1, ncols2, self.len() + rhs.len());
        let mut timestamps = VectorN::zeros_generic(nrows1, U1);
        let mut workspace = unsafe { VectorN::new_uninitialized_generic(nrows1, U1) };
        let mut nz = 0;

        for j in 0..ncols2.value() {
            res.data.p[j] = nz;

            nz = self.scatter(
                j,
                N::one(),
                timestamps.as_mut_slice(),
                j + 1,
                workspace.as_mut_slice(),
                nz,
                &mut res,
            );

            nz = rhs.scatter(
                j,
                N::one(),
                timestamps.as_mut_slice(),
                j + 1,
                workspace.as_mut_slice(),
                nz,
                &mut res,
            );

            // Keep the output sorted.
            let range = res.data.p[j]..nz;
            res.data.i[range.clone()].sort();

            for p in range {
                res.data.vals[p] = workspace[res.data.i[p]]
            }
        }

        res.data.i.truncate(nz);
        res.data.i.shrink_to_fit();
        res.data.vals.truncate(nz);
        res.data.vals.shrink_to_fit();
        res
    }
}
<commit_msg>Improve CsMatrix multiplaction performances.<commit_after>use alga::general::{ClosedAdd, ClosedMul};
use num::{One, Zero};
use std::iter;
use std::marker::PhantomData;
use std::ops::{Add, Mul, Range};
use std::slice;

use allocator::Allocator;
use constraint::{AreMultipliable, DimEq, SameNumberOfRows, ShapeConstraint};
use sparse::{CsMatrix, CsStorage, CsStorageMut, CsVector};
use storage::{Storage, StorageMut};
use {DefaultAllocator, Dim, Matrix, MatrixMN, Real, Scalar, Vector, VectorN, U1};

impl<N: Scalar, R: Dim, C: Dim, S: CsStorage<N, R, C>> CsMatrix<N, R, C, S> {
    fn scatter<R2: Dim, C2: Dim>(
        &self,
        j: usize,
        beta: N,
        timestamps: &mut [usize],
        timestamp: usize,
        workspace: &mut [N],
        mut nz: usize,
        res: &mut CsMatrix<N, R2, C2>,
    ) -> usize
    where
        N: ClosedAdd + ClosedMul,
        DefaultAllocator: Allocator<usize, C2>,
    {
        for (i, val) in self.data.column_entries(j) {
            if timestamps[i] < timestamp {
                timestamps[i] = timestamp;
                res.data.i[nz] = i;
                nz += 1;
                workspace[i] = val * beta;
            } else {
                workspace[i] += val * beta;
            }
        }

        nz
    }
}

/*
impl<N: Scalar, R, S> CsVector<N, R, S> {
    pub fn axpy(&mut self, alpha: N, x: CsVector<N, R, S>, beta: N) {
        // First, compute the number of non-zero entries.
        let mut nnzero = 0;

        // Allocate a size large enough.
        self.data.set_column_len(0, nnzero);

        // Fill with the axpy.
        let mut i = self.len();
        let mut j = x.len();
        let mut k = nnzero - 1;
        let mut rid1 = self.data.row_index(0, i - 1);
        let mut rid2 = x.data.row_index(0, j - 1);

        while k > 0 {
            if rid1 == rid2 {
                self.data.set_row_index(0, k, rid1);
                self[k] = alpha * x[j] + beta * self[k];
                i -= 1;
                j -= 1;
            } else if rid1 < rid2 {
                self.data.set_row_index(0, k, rid1);
                self[k] = beta * self[i];
                i -= 1;
            } else {
                self.data.set_row_index(0, k, rid2);
                self[k] = alpha * x[j];
                j -= 1;
            }

            k -= 1;
        }
    }
}
*/

impl<N: Scalar + Zero + ClosedAdd + ClosedMul, D: Dim, S: StorageMut<N, D>> Vector<N, D, S> {
    pub fn axpy_cs<D2: Dim, S2>(&mut self, alpha: N, x: &CsVector<N, D2, S2>, beta: N)
    where
        S2: CsStorage<N, D2>,
        ShapeConstraint: DimEq<D, D2>,
    {
        if beta.is_zero() {
            for i in 0..x.len() {
                unsafe {
                    let k = x.data.row_index_unchecked(i);
                    let y = self.vget_unchecked_mut(k);
                    *y = alpha * *x.data.get_value_unchecked(i);
                }
            }
        } else {
            // Needed to be sure even components not present on `x` are multiplied.
            *self *= beta;

            for i in 0..x.len() {
                unsafe {
                    let k = x.data.row_index_unchecked(i);
                    let y = self.vget_unchecked_mut(k);
                    *y += alpha * *x.data.get_value_unchecked(i);
                }
            }
        }
    }

    /*
    pub fn gemv_sparse<R2: Dim, C2: Dim, S2>(&mut self, alpha: N, a: &CsMatrix<N, R2, C2, S2>, x: &DVector<N>, beta: N)
        where
            S2: CsStorage<N, R2, C2> {
        let col2 = a.column(0);
        let val = unsafe { *x.vget_unchecked(0) };
        self.axpy_sparse(alpha * val, &col2, beta);

        for j in 1..ncols2 {
            let col2 = a.column(j);
            let val = unsafe { *x.vget_unchecked(j) };

            self.axpy_sparse(alpha * val, &col2, N::one());
        }
    }
    */
}

impl<'a, 'b, N, R1, R2, C1, C2, S1, S2> Mul<&'b CsMatrix<N, R2, C2, S2>>
    for &'a CsMatrix<N, R1, C1, S1>
where
    N: Scalar + ClosedAdd + ClosedMul + Zero,
    R1: Dim,
    C1: Dim,
    R2: Dim,
    C2: Dim,
    S1: CsStorage<N, R1, C1>,
    S2: CsStorage<N, R2, C2>,
    ShapeConstraint: AreMultipliable<R1, C1, R2, C2>,
    DefaultAllocator: Allocator<usize, C2> + Allocator<usize, R1> + Allocator<N, R1>,
{
    type Output = CsMatrix<N, R1, C2>;

    fn mul(self, rhs: &'b CsMatrix<N, R2, C2, S2>) -> CsMatrix<N, R1, C2> {
        let (nrows1, ncols1) = self.data.shape();
        let (nrows2, ncols2) = rhs.data.shape();
        assert_eq!(
            ncols1.value(),
            nrows2.value(),
            "Mismatched dimensions for matrix multiplication."
        );

        let mut res = CsMatrix::new_uninitialized_generic(nrows1, ncols2, self.len() + rhs.len());
        let mut workspace = VectorN::<N, R1>::zeros_generic(nrows1, U1);
        let mut nz = 0;

        for j in 0..ncols2.value() {
            res.data.p[j] = nz;
            let new_size_bound = nz + nrows1.value();
            res.data.i.resize(new_size_bound, 0);
            res.data.vals.resize(new_size_bound, N::zero());

            for (i, beta) in rhs.data.column_entries(j) {
                for (k, val) in self.data.column_entries(i) {
                    workspace[k] += val * beta;
                }
            }

            for (i, val) in workspace.as_mut_slice().iter_mut().enumerate() {
                if !val.is_zero() {
                    res.data.i[nz] = i;
                    res.data.vals[nz] = *val;
                    *val = N::zero();
                    nz += 1;
                }
            }
        }

        res.data.i.truncate(nz);
        res.data.i.shrink_to_fit();
        res.data.vals.truncate(nz);
        res.data.vals.shrink_to_fit();
        res
    }
}

impl<'a, 'b, N, R1, R2, C1, C2, S1, S2> Add<&'b CsMatrix<N, R2, C2, S2>>
    for &'a CsMatrix<N, R1, C1, S1>
where
    N: Scalar + ClosedAdd + ClosedMul + One,
    R1: Dim,
    C1: Dim,
    R2: Dim,
    C2: Dim,
    S1: CsStorage<N, R1, C1>,
    S2: CsStorage<N, R2, C2>,
    ShapeConstraint: DimEq<R1, R2> + DimEq<C1, C2>,
    DefaultAllocator: Allocator<usize, C2> + Allocator<usize, R1> + Allocator<N, R1>,
{
    type Output = CsMatrix<N, R1, C2>;

    fn add(self, rhs: &'b CsMatrix<N, R2, C2, S2>) -> CsMatrix<N, R1, C2> {
        let (nrows1, ncols1) = self.data.shape();
        let (nrows2, ncols2) = rhs.data.shape();
        assert_eq!(
            (nrows1.value(), ncols1.value()),
            (nrows2.value(), ncols2.value()),
            "Mismatched dimensions for matrix sum."
        );

        let mut res = CsMatrix::new_uninitialized_generic(nrows1, ncols2, self.len() + rhs.len());
        let mut timestamps = VectorN::zeros_generic(nrows1, U1);
        let mut workspace = unsafe { VectorN::new_uninitialized_generic(nrows1, U1) };
        let mut nz = 0;

        for j in 0..ncols2.value() {
            res.data.p[j] = nz;

            nz = self.scatter(
                j,
                N::one(),
                timestamps.as_mut_slice(),
                j + 1,
                workspace.as_mut_slice(),
                nz,
                &mut res,
            );

            nz = rhs.scatter(
                j,
                N::one(),
                timestamps.as_mut_slice(),
                j + 1,
                workspace.as_mut_slice(),
                nz,
                &mut res,
            );

            // Keep the output sorted.
            let range = res.data.p[j]..nz;
            res.data.i[range.clone()].sort();

            for p in range {
                res.data.vals[p] = workspace[res.data.i[p]]
            }
        }

        res.data.i.truncate(nz);
        res.data.i.shrink_to_fit();
        res.data.vals.truncate(nz);
        res.data.vals.shrink_to_fit();
        res
    }
}

impl<'a, 'b, N, R, C, S> Mul<N> for CsMatrix<N, R, C, S>
where
    N: Scalar + ClosedAdd + ClosedMul + Zero,
    R: Dim,
    C: Dim,
    S: CsStorageMut<N, R, C>,
{
    type Output = Self;

    fn mul(mut self, rhs: N) -> Self {
        for e in self.values_mut() {
            *e *= rhs
        }

        self
    }
}
<|endoftext|>


<commit_before>#![crate_name = "docopt_macros"]
#![crate_type = "dylib"]
#![feature(plugin_registrar, macro_rules, quote)]

//! This crate defines the `docopt!` macro. It is documented in the
//! documentation of the `docopt` crate.

extern crate syntax;
extern crate rustc;
extern crate docopt;

use std::collections::HashMap;

use rustc::plugin::Registry;
use syntax::ast;
use syntax::codemap;
use syntax::ext::base::{ExtCtxt, MacResult, MacItems, DummyResult};
use syntax::ext::build::AstBuilder;
use syntax::fold::Folder;
use syntax::parse::common::SeqSep;
use syntax::parse::parser::Parser;
use syntax::parse::token;
use syntax::print::pprust;
use syntax::ptr::P;

use docopt::{Docopt, ArgvMap};
use docopt::parse::{Options, Atom, Positional, Zero, One};

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("docopt", expand);
}

fn expand(cx: &mut ExtCtxt, span: codemap::Span, tts: &[ast::TokenTree])
         -> Box<MacResult+'static> {
    let parsed = match MacParser::new(cx, tts).parse() {
        Ok(parsed) => parsed,
        Err(_) => return DummyResult::any(span),
    };
    parsed.items(cx)
}

/// Parsed corresponds to the result of parsing a `docopt` macro call.
/// It can be used to write a corresponding struct.
struct Parsed {
    struct_info: StructInfo,
    doc: Docopt,
    /// Overrided type annotations for struct members. May be empty.
    /// When a type annotation for an atom doesn't exist, then one is
    /// inferred automatically. It is one of: `bool`, `uint`, `String` or
    /// `Vec<String>`.
    types: HashMap<Atom, P<ast::Ty>>,
}

impl Parsed {
    /// Returns a macro result suitable for expansion.
    /// Contains two items: one for the struct and one for the struct impls.
    fn items(&self, cx: &ExtCtxt) -> Box<MacResult+'static> {
        let mut its = vec!();
        its.push(self.struct_decl(cx));

        let struct_name = self.struct_info.name;
        let full_doc = self.doc.parser().full_doc.as_slice();
        its.push(quote_item!(cx,
            impl $struct_name {
                #[allow(dead_code)]
                fn docopt() -> docopt::Docopt {
                    // The unwrap is justified here because this code
                    // gen only happens if the Docopt usage string is valid.
                    docopt::Docopt::new($full_doc).unwrap()
                }
            }
        ).unwrap());
        MacItems::new(its.into_iter())
    }

    /// Returns an item for the struct definition.
    fn struct_decl(&self, cx: &ExtCtxt) -> P<ast::Item> {
        let name = self.struct_info.name.clone();
        let vis = if self.struct_info.public { ast::Public }
                  else { ast::Inherited };
        let def = ast::StructDef {
            fields: self.struct_fields(cx),
            ctor_id: None
        };

        let mut traits = vec!["Decodable".to_string()];
        traits.push_all(self.struct_info.deriving.as_slice());
        let attrs = vec![attribute(cx, "allow", vec!["non_snake_case"]),
                         attribute(cx, "deriving", traits)];

        let st = cx.item_struct(codemap::DUMMY_SP, name.clone(), def);
        cx.item(codemap::DUMMY_SP, name, attrs, st.node.clone()).map(|mut it| {
            it.vis = vis;
            it
        })
    }

    /// Returns a list of fields for the struct definition.
    /// Handles type annotations.
    fn struct_fields(&self, cx: &ExtCtxt) -> Vec<ast::StructField> {
        let mut fields: Vec<ast::StructField> = vec!();
        for (atom, opts) in self.doc.parser().descs.iter() {
            let name = ArgvMap::key_to_struct_field(atom.to_string().as_slice());
            let ty = match self.types.get(atom) {
                None => self.pat_type(cx, atom, opts),
                Some(ty) => ty.clone(),
            };
            fields.push(self.mk_struct_field(name.as_slice(), ty));
        }
        fields
    }

    /// Returns an inferred type for a usage pattern.
    /// This is only invoked when a type annotation is not present.
    fn pat_type(&self, cx: &ExtCtxt, atom: &Atom, opts: &Options) -> P<ast::Ty> {
        match (opts.repeats, &opts.arg) {
            (false, &Zero) => {
                match atom {
                    &Positional(_) => quote_ty!(cx, String),
                    _ => quote_ty!(cx, bool),
                }
            }
            (true, &Zero) => {
                match atom {
                    &Positional(_) => quote_ty!(cx, Vec<String>),
                    _ => quote_ty!(cx, uint),
                }
            }
            (false, &One(_)) => quote_ty!(cx, String),
            (true, &One(_)) => quote_ty!(cx, Vec<String>),
        }
    }

    /// Creates a struct field from a member name and type.
    fn mk_struct_field(&self, name: &str, ty: P<ast::Ty>) -> ast::StructField {
        codemap::dummy_spanned(ast::StructField_ {
            kind: ast::NamedField(ident(name), ast::Public),
            id: ast::DUMMY_NODE_ID,
            ty: ty,
            attrs: vec!(),
        })
    }
}

/// State for parsing a `docopt` macro invocation.
struct MacParser<'a, 'b:'a> {
    cx: &'a mut ExtCtxt<'b>,
    p: Parser<'b>,
}

impl<'a, 'b> MacParser<'a, 'b> {
    fn new(cx: &'a mut ExtCtxt<'b>, tts: &[ast::TokenTree]) -> MacParser<'a, 'b> {
        let p = cx.new_parser_from_tts(tts);
        MacParser { cx: cx, p: p }
    }

    /// Main entry point for parsing arguments to `docopt` macro.
    /// First looks for an identifier for the struct name.
    /// Second, a string containing the docopt usage patterns.
    /// Third, an optional list of type annotations.
    fn parse(&mut self) -> Result<Parsed, ()> {
        if self.p.token == token::Eof {
            self.cx.span_err(self.cx.call_site(), "macro expects arguments");
            return Err(());
        }
        let struct_info = try!(self.parse_struct_info());
        let docstr = try!(self.parse_str());

        let sep = SeqSep {
            sep: Some(token::Comma),
            trailing_sep_allowed: true,
        };
        let types = self.p.parse_seq_to_end(
            &token::Eof, sep, |p| MacParser::parse_type_annotation(p)
        ).into_iter()
         .map(|(ident, ty)| {
             let field_name = token::get_ident(ident).to_string();
             let key = ArgvMap::struct_field_to_key(field_name.as_slice());
             (Atom::new(key.as_slice()), ty)
          })
         .collect::<HashMap<Atom, P<ast::Ty>>>();
        self.p.expect(&token::Eof);

        // This config does not matter because we're only asking for the
        // usage patterns in the Docopt string. The configuration does not
        // affect the retrieval of usage patterns.
        let doc = match Docopt::new(docstr) {
            Ok(doc) => doc,
            Err(err) => {
                self.cx.span_err(self.cx.call_site(),
                                 format!("Invalid Docopt usage: {}",
                                         err).as_slice());
                return Err(());
            }
        };
        Ok(Parsed {
            struct_info: struct_info,
            doc: doc,
            types: types,
        })
    }

    /// Parses a single string literal. On failure, an error is logged and
    /// unit is returned.
    fn parse_str(&mut self) -> Result<String, ()> {
        fn lit_is_str(lit: &ast::Lit) -> bool {
            match lit.node {
                ast::LitStr(_, _) => true,
                _ => false,
            }
        }
        fn lit_to_string(lit: &ast::Lit) -> String {
            match lit.node {
                ast::LitStr(ref s, _) => s.to_string(),
                _ => panic!("BUG: expected string literal"),
            }
        }
        let exp = self.cx.expander().fold_expr(self.p.parse_expr());
        let s = match exp.node {
            ast::ExprLit(ref lit) if lit_is_str(&**lit) => {
                lit_to_string(&**lit)
            }
            _ => {
                let err = format!("Expected string literal but got {}",
                                  pprust::expr_to_string(&*exp));
                self.cx.span_err(exp.span, err.as_slice());
                return Err(());
            }
        };
        self.p.bump();
        Ok(s)
    }

    /// Parses a type annotation in a `docopt` invocation of the form
    /// `ident: Ty`.
    /// Note that this is a static method as it is used as a HOF.
    fn parse_type_annotation(p: &mut Parser) -> (ast::Ident, P<ast::Ty>) {
        let ident = p.parse_ident();
        p.expect(&token::Colon);
        let ty = p.parse_ty();
        (ident, ty)
    }

    /// Parses struct information, like visibility, name and deriving.
    fn parse_struct_info(&mut self) -> Result<StructInfo, ()> {
        let public = self.p.eat_keyword(token::keywords::Pub);
        let mut info = StructInfo {
            name: self.p.parse_ident(),
            public: public,
            deriving: vec![],
        };
        if self.p.eat(&token::Comma) { return Ok(info); }
        let deriving = self.p.parse_ident();
        if deriving.as_str() != "deriving" {
            let err = format!("Expected 'deriving' keyword but got '{}'",
                              deriving);
            self.cx.span_err(self.cx.call_site(), err.as_slice());
            return Err(());
        }
        while !self.p.eat(&token::Comma) {
            info.deriving.push(self.p.parse_ident().as_str().to_string());
        }
        Ok(info)
    }
}

struct StructInfo {
    name: ast::Ident,
    public: bool,
    deriving: Vec<String>,
}

// Convenience functions for building intermediate values.

fn ident(s: &str) -> ast::Ident {
    ast::Ident::new(token::intern(s))
}

fn attribute<S, T>(cx: &ExtCtxt, name: S, items: Vec<T>) -> ast::Attribute
            where S: Str, T: Str {
    let sp = codemap::DUMMY_SP;
    let its = items.into_iter().map(|s| meta_item(cx, s.as_slice())).collect();
    let mi = cx.meta_list(sp, intern(name.as_slice()), its);
    cx.attribute(sp, mi)
}

fn meta_item(cx: &ExtCtxt, s: &str) -> P<ast::MetaItem> {
    cx.meta_word(codemap::DUMMY_SP, intern(s))
}

fn intern(s: &str) -> token::InternedString {
    token::intern_and_get_ident(s)
}
<commit_msg>quote_ty is broken for now, so do things by hand. :-(<commit_after>#![crate_name = "docopt_macros"]
#![crate_type = "dylib"]
#![feature(plugin_registrar, macro_rules, quote)]

//! This crate defines the `docopt!` macro. It is documented in the
//! documentation of the `docopt` crate.

extern crate syntax;
extern crate rustc;
extern crate docopt;

use std::collections::HashMap;

use rustc::plugin::Registry;
use syntax::ast;
use syntax::codemap;
use syntax::ext::base::{ExtCtxt, MacResult, MacItems, DummyResult};
use syntax::ext::build::AstBuilder;
use syntax::fold::Folder;
use syntax::owned_slice::OwnedSlice;
use syntax::parse::common::SeqSep;
use syntax::parse::parser::Parser;
use syntax::parse::token;
use syntax::print::pprust;
use syntax::ptr::P;

use docopt::{Docopt, ArgvMap};
use docopt::parse::{Options, Atom, Positional, Zero, One};

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_macro("docopt", expand);
}

fn expand(cx: &mut ExtCtxt, span: codemap::Span, tts: &[ast::TokenTree])
         -> Box<MacResult+'static> {
    let parsed = match MacParser::new(cx, tts).parse() {
        Ok(parsed) => parsed,
        Err(_) => return DummyResult::any(span),
    };
    parsed.items(cx)
}

/// Parsed corresponds to the result of parsing a `docopt` macro call.
/// It can be used to write a corresponding struct.
struct Parsed {
    struct_info: StructInfo,
    doc: Docopt,
    /// Overrided type annotations for struct members. May be empty.
    /// When a type annotation for an atom doesn't exist, then one is
    /// inferred automatically. It is one of: `bool`, `uint`, `String` or
    /// `Vec<String>`.
    types: HashMap<Atom, P<ast::Ty>>,
}

impl Parsed {
    /// Returns a macro result suitable for expansion.
    /// Contains two items: one for the struct and one for the struct impls.
    fn items(&self, cx: &ExtCtxt) -> Box<MacResult+'static> {
        let mut its = vec!();
        its.push(self.struct_decl(cx));

        let struct_name = self.struct_info.name;
        let full_doc = self.doc.parser().full_doc.as_slice();
        its.push(quote_item!(cx,
            impl $struct_name {
                #[allow(dead_code)]
                fn docopt() -> docopt::Docopt {
                    // The unwrap is justified here because this code
                    // gen only happens if the Docopt usage string is valid.
                    docopt::Docopt::new($full_doc).unwrap()
                }
            }
        ).unwrap());
        MacItems::new(its.into_iter())
    }

    /// Returns an item for the struct definition.
    fn struct_decl(&self, cx: &ExtCtxt) -> P<ast::Item> {
        let name = self.struct_info.name.clone();
        let vis = if self.struct_info.public { ast::Public }
                  else { ast::Inherited };
        let def = ast::StructDef {
            fields: self.struct_fields(cx),
            ctor_id: None
        };

        let mut traits = vec!["Decodable".to_string()];
        traits.push_all(self.struct_info.deriving.as_slice());
        let attrs = vec![attribute(cx, "allow", vec!["non_snake_case"]),
                         attribute(cx, "deriving", traits)];

        let st = cx.item_struct(codemap::DUMMY_SP, name.clone(), def);
        cx.item(codemap::DUMMY_SP, name, attrs, st.node.clone()).map(|mut it| {
            it.vis = vis;
            it
        })
    }

    /// Returns a list of fields for the struct definition.
    /// Handles type annotations.
    fn struct_fields(&self, cx: &ExtCtxt) -> Vec<ast::StructField> {
        let mut fields: Vec<ast::StructField> = vec!();
        for (atom, opts) in self.doc.parser().descs.iter() {
            let name = ArgvMap::key_to_struct_field(atom.to_string().as_slice());
            let ty = match self.types.get(atom) {
                None => self.pat_type(cx, atom, opts),
                Some(ty) => ty.clone(),
            };
            fields.push(self.mk_struct_field(name.as_slice(), ty));
        }
        fields
    }

    /// Returns an inferred type for a usage pattern.
    /// This is only invoked when a type annotation is not present.
    fn pat_type(&self, cx: &ExtCtxt, atom: &Atom, opts: &Options) -> P<ast::Ty> {
        let sp = codemap::DUMMY_SP;
        match (opts.repeats, &opts.arg) {
            (false, &Zero) => {
                match atom {
                    &Positional(_) => cx.ty_ident(sp, ident("String")),
                    _ => cx.ty_ident(sp, ident("bool")),
                }
            }
            (true, &Zero) => {
                match atom {
                    &Positional(_) => ty_vec_string(cx),
                    _ => cx.ty_ident(sp, ident("uint")),
                }
            }
            (false, &One(_)) => cx.ty_ident(sp, ident("String")),
            (true, &One(_)) => ty_vec_string(cx),
        }
    }

    /// Creates a struct field from a member name and type.
    fn mk_struct_field(&self, name: &str, ty: P<ast::Ty>) -> ast::StructField {
        codemap::dummy_spanned(ast::StructField_ {
            kind: ast::NamedField(ident(name), ast::Public),
            id: ast::DUMMY_NODE_ID,
            ty: ty,
            attrs: vec!(),
        })
    }
}

/// State for parsing a `docopt` macro invocation.
struct MacParser<'a, 'b:'a> {
    cx: &'a mut ExtCtxt<'b>,
    p: Parser<'b>,
}

impl<'a, 'b> MacParser<'a, 'b> {
    fn new(cx: &'a mut ExtCtxt<'b>, tts: &[ast::TokenTree]) -> MacParser<'a, 'b> {
        let p = cx.new_parser_from_tts(tts);
        MacParser { cx: cx, p: p }
    }

    /// Main entry point for parsing arguments to `docopt` macro.
    /// First looks for an identifier for the struct name.
    /// Second, a string containing the docopt usage patterns.
    /// Third, an optional list of type annotations.
    fn parse(&mut self) -> Result<Parsed, ()> {
        if self.p.token == token::Eof {
            self.cx.span_err(self.cx.call_site(), "macro expects arguments");
            return Err(());
        }
        let struct_info = try!(self.parse_struct_info());
        let docstr = try!(self.parse_str());

        let sep = SeqSep {
            sep: Some(token::Comma),
            trailing_sep_allowed: true,
        };
        let types = self.p.parse_seq_to_end(
            &token::Eof, sep, |p| MacParser::parse_type_annotation(p)
        ).into_iter()
         .map(|(ident, ty)| {
             let field_name = token::get_ident(ident).to_string();
             let key = ArgvMap::struct_field_to_key(field_name.as_slice());
             (Atom::new(key.as_slice()), ty)
          })
         .collect::<HashMap<Atom, P<ast::Ty>>>();
        self.p.expect(&token::Eof);

        // This config does not matter because we're only asking for the
        // usage patterns in the Docopt string. The configuration does not
        // affect the retrieval of usage patterns.
        let doc = match Docopt::new(docstr) {
            Ok(doc) => doc,
            Err(err) => {
                self.cx.span_err(self.cx.call_site(),
                                 format!("Invalid Docopt usage: {}",
                                         err).as_slice());
                return Err(());
            }
        };
        Ok(Parsed {
            struct_info: struct_info,
            doc: doc,
            types: types,
        })
    }

    /// Parses a single string literal. On failure, an error is logged and
    /// unit is returned.
    fn parse_str(&mut self) -> Result<String, ()> {
        fn lit_is_str(lit: &ast::Lit) -> bool {
            match lit.node {
                ast::LitStr(_, _) => true,
                _ => false,
            }
        }
        fn lit_to_string(lit: &ast::Lit) -> String {
            match lit.node {
                ast::LitStr(ref s, _) => s.to_string(),
                _ => panic!("BUG: expected string literal"),
            }
        }
        let exp = self.cx.expander().fold_expr(self.p.parse_expr());
        let s = match exp.node {
            ast::ExprLit(ref lit) if lit_is_str(&**lit) => {
                lit_to_string(&**lit)
            }
            _ => {
                let err = format!("Expected string literal but got {}",
                                  pprust::expr_to_string(&*exp));
                self.cx.span_err(exp.span, err.as_slice());
                return Err(());
            }
        };
        self.p.bump();
        Ok(s)
    }

    /// Parses a type annotation in a `docopt` invocation of the form
    /// `ident: Ty`.
    /// Note that this is a static method as it is used as a HOF.
    fn parse_type_annotation(p: &mut Parser) -> (ast::Ident, P<ast::Ty>) {
        let ident = p.parse_ident();
        p.expect(&token::Colon);
        let ty = p.parse_ty();
        (ident, ty)
    }

    /// Parses struct information, like visibility, name and deriving.
    fn parse_struct_info(&mut self) -> Result<StructInfo, ()> {
        let public = self.p.eat_keyword(token::keywords::Pub);
        let mut info = StructInfo {
            name: self.p.parse_ident(),
            public: public,
            deriving: vec![],
        };
        if self.p.eat(&token::Comma) { return Ok(info); }
        let deriving = self.p.parse_ident();
        if deriving.as_str() != "deriving" {
            let err = format!("Expected 'deriving' keyword but got '{}'",
                              deriving);
            self.cx.span_err(self.cx.call_site(), err.as_slice());
            return Err(());
        }
        while !self.p.eat(&token::Comma) {
            info.deriving.push(self.p.parse_ident().as_str().to_string());
        }
        Ok(info)
    }
}

struct StructInfo {
    name: ast::Ident,
    public: bool,
    deriving: Vec<String>,
}

// Convenience functions for building intermediate values.

fn ident(s: &str) -> ast::Ident {
    ast::Ident::new(token::intern(s))
}

fn attribute<S, T>(cx: &ExtCtxt, name: S, items: Vec<T>) -> ast::Attribute
            where S: Str, T: Str {
    let sp = codemap::DUMMY_SP;
    let its = items.into_iter().map(|s| meta_item(cx, s.as_slice())).collect();
    let mi = cx.meta_list(sp, intern(name.as_slice()), its);
    cx.attribute(sp, mi)
}

fn meta_item(cx: &ExtCtxt, s: &str) -> P<ast::MetaItem> {
    cx.meta_word(codemap::DUMMY_SP, intern(s))
}

fn intern(s: &str) -> token::InternedString {
    token::intern_and_get_ident(s)
}

fn ty_vec_string(cx: &ExtCtxt) -> P<ast::Ty> {
    let sp = codemap::DUMMY_SP;
    let tystr = ast::AngleBracketedParameterData {
        lifetimes: vec![],
        types: OwnedSlice::from_vec(vec![cx.ty_ident(sp, ident("String"))]),
    };
    cx.ty_path(ast::Path {
        span: sp,
        global: false,
        segments: vec![ast::PathSegment {
            identifier: ident("Vec"),
            parameters: ast::PathParameters::AngleBracketedParameters(tystr),
        }]
    })
}
<|endoftext|>


<commit_before>use handlebars::Handlebars;
use serde_json::json;

#[test]
fn test_partial_with_blocks() {
    let hbs = Handlebars::new();

    let data = json!({
        "a": [
            {"b": 1},
            {"b": 2},
        ],
    });

    let template = "{{#*inline \"test\"}}{{b}};{{/inline}}{{#each a as |z|}}{{> test z}}{{/each}}";
    assert_eq!(
        hbs.render_template(template, &data).unwrap(),
        "1;2;"
    );
}
<commit_msg>(test) added case for #264<commit_after>use handlebars::Handlebars;
use serde_json::json;

#[test]
fn test_partial_with_blocks() {
    let hbs = Handlebars::new();

    let data = json!({
        "a": [
            {"b": 1},
            {"b": 2},
        ],
    });

    let template = "{{#*inline \"test\"}}{{b}};{{/inline}}{{#each a as |z|}}{{> test z}}{{/each}}";
    assert_eq!(hbs.render_template(template, &data).unwrap(), "1;2;");
}

#[test]
fn test_root_with_blocks() {
    let hbs = Handlebars::new();

    let data = json!({
        "a": [
            {"b": 1},
            {"b": 2},
        ],
        "b": 3,
    });

    let template =
        "{{#*inline \"test\"}}{{b}}:{{@root.b}};{{/inline}}{{#each a}}{{> test}}{{/each}}";
    assert_eq!(hbs.render_template(template, &data).unwrap(), "1:3;2:3;");
}
<|endoftext|>


<commit_before>use std::io;

fn main() {
    println!("Type something!");

    let input = io::stdin()
        .read_line()
        .ok()
        .expect("Failed to read line");

    println!("{}", input);
}

<commit_msg>switch to old_io temporarily<commit_after>#![feature(old_io)]
use std::old_io as io;

fn main() {
    println!("Type something!");

    let input = io::stdin()
        .read_line()
        .ok()
        .expect("Failed to read line");

    println!("{}", input);
}

<|endoftext|>


<commit_before>// * This file is part of the uutils coreutils package.
// *
// * (c) T. Jameson Little <t.jameson.little@gmail.com>
// * (c) Wiktor Kuropatwa <wiktor.kuropatwa@gmail.com>
// *     * 2015-02-23 ~ added Pollard rho method implementation
// * (c) kwantam <kwantam@gmail.com>
// *     * 2015-04-29 ~ sped up trial division by adding table of prime inverses
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

extern crate rand;

#[macro_use]
extern crate uucore;

use std::collections::HashMap;
use std::fmt;
use std::io::{stdin, BufRead};
use std::ops;

mod miller_rabin;
mod numeric;
mod rho;
mod table;

static SYNTAX: &str = "[OPTION] [NUMBER]...";
static SUMMARY: &str = "Print the prime factors of the given number(s).
 If none are specified, read from standard input.";
static LONG_HELP: &str = "";

struct Factors {
    f: HashMap<u64, u8>,
}

impl Factors {
    fn new() -> Factors {
        Factors { f: HashMap::new() }
    }

    fn add(&mut self, prime: u64, exp: u8) {
        assert!(exp > 0);
        let n = *self.f.get(&prime).unwrap_or(&0);
        self.f.insert(prime, exp + n);
    }

    fn push(&mut self, prime: u64) {
        self.add(prime, 1)
    }
}

impl ops::MulAssign<Factors> for Factors {
    fn mul_assign(&mut self, other: Factors) {
        for (prime, exp) in &other.f {
            self.add(*prime, *exp);
        }
    }
}

impl fmt::Display for Factors {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // TODO: Use a representation with efficient in-order iteration
        let mut primes: Vec<&u64> = self.f.keys().collect();
        primes.sort();

        for p in primes {
            for _ in 0..self.f[&p] {
                write!(f, " {}", p)?
            }
        }

        Ok(())
    }
}

fn factor(mut n: u64) -> Factors {
    let mut factors = Factors::new();

    if n < 2 {
        factors.push(n);
        return factors;
    }

    let z = n.trailing_zeros();
    if z > 0 {
        factors.add(2, z as u8);
        n >>= z;
    }

    if n == 1 {
        return factors;
    }

    let (f, n) = table::factor(n);
    factors *= f;

    if n >= table::NEXT_PRIME {
        factors *= rho::factor(n);
    }

    factors
}

fn print_factors(num: u64) {
    print!("{}:{}", num, factor(num));
    println!();
}

fn print_factors_str(num_str: &str) {
    if let Err(e) = num_str.parse::<u64>().and_then(|x| {
        print_factors(x);
        Ok(())
    }) {
        show_warning!("{}: {}", num_str, e);
    }
}

pub fn uumain(args: Vec<String>) -> i32 {
    let matches = app!(SYNTAX, SUMMARY, LONG_HELP).parse(args);

    if matches.free.is_empty() {
        let stdin = stdin();
        for line in stdin.lock().lines() {
            for number in line.unwrap().split_whitespace() {
                print_factors_str(number);
            }
        }
    } else {
        for num_str in &matches.free {
            print_factors_str(num_str);
        }
    }
    0
}
<commit_msg>factor::factor: Add integration tests<commit_after>// * This file is part of the uutils coreutils package.
// *
// * (c) T. Jameson Little <t.jameson.little@gmail.com>
// * (c) Wiktor Kuropatwa <wiktor.kuropatwa@gmail.com>
// *     * 2015-02-23 ~ added Pollard rho method implementation
// * (c) kwantam <kwantam@gmail.com>
// *     * 2015-04-29 ~ sped up trial division by adding table of prime inverses
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

extern crate rand;

#[macro_use]
extern crate uucore;

use std::collections::HashMap;
use std::fmt;
use std::io::{stdin, BufRead};
use std::ops;

mod miller_rabin;
mod numeric;
mod rho;
mod table;

static SYNTAX: &str = "[OPTION] [NUMBER]...";
static SUMMARY: &str = "Print the prime factors of the given number(s).
 If none are specified, read from standard input.";
static LONG_HELP: &str = "";

struct Factors {
    f: HashMap<u64, u8>,
}

impl Factors {
    fn new() -> Factors {
        Factors { f: HashMap::new() }
    }

    fn add(&mut self, prime: u64, exp: u8) {
        assert!(exp > 0);
        let n = *self.f.get(&prime).unwrap_or(&0);
        self.f.insert(prime, exp + n);
    }

    fn push(&mut self, prime: u64) {
        self.add(prime, 1)
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.f
            .iter()
            .fold(1, |acc, (p, exp)| acc * p.pow(*exp as u32))
    }
}

impl ops::MulAssign<Factors> for Factors {
    fn mul_assign(&mut self, other: Factors) {
        for (prime, exp) in &other.f {
            self.add(*prime, *exp);
        }
    }
}

impl fmt::Display for Factors {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // TODO: Use a representation with efficient in-order iteration
        let mut primes: Vec<&u64> = self.f.keys().collect();
        primes.sort();

        for p in primes {
            for _ in 0..self.f[&p] {
                write!(f, " {}", p)?
            }
        }

        Ok(())
    }
}

fn factor(mut n: u64) -> Factors {
    let mut factors = Factors::new();

    if n < 2 {
        factors.push(n);
        return factors;
    }

    let z = n.trailing_zeros();
    if z > 0 {
        factors.add(2, z as u8);
        n >>= z;
    }

    if n == 1 {
        return factors;
    }

    let (f, n) = table::factor(n);
    factors *= f;

    if n >= table::NEXT_PRIME {
        factors *= rho::factor(n);
    }

    factors
}

fn print_factors(num: u64) {
    print!("{}:{}", num, factor(num));
    println!();
}

fn print_factors_str(num_str: &str) {
    if let Err(e) = num_str.parse::<u64>().and_then(|x| {
        print_factors(x);
        Ok(())
    }) {
        show_warning!("{}: {}", num_str, e);
    }
}

pub fn uumain(args: Vec<String>) -> i32 {
    let matches = app!(SYNTAX, SUMMARY, LONG_HELP).parse(args);

    if matches.free.is_empty() {
        let stdin = stdin();
        for line in stdin.lock().lines() {
            for number in line.unwrap().split_whitespace() {
                print_factors_str(number);
            }
        }
    } else {
        for num_str in &matches.free {
            print_factors_str(num_str);
        }
    }
    0
}

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

    #[test]
    fn factor_recombines_small() {
        assert!((1..10_000)
            .map(|i| 2 * i + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_overflowing() {
        assert!((0..250)
            .map(|i| 2 * i + 2u64.pow(32) + 1)
            .all(|i| factor(i).product() == i));
    }
}
<|endoftext|>