<commit_before>use std::mem;

use ra_syntax::{
    TextUnit, TextRange, SyntaxNodeRef, File, AstNode, SyntaxKind,
    ast,
    algo::{
        find_covering_node, find_leaf_at_offset, LeafAtOffset,
    },
    text_utils::{intersect, contains_offset_nonstrict},
    SyntaxKind::*,
};

use {LocalEdit, EditBuilder, find_node_at_offset};

pub fn join_lines(file: &File, range: TextRange) -> LocalEdit {
    let range = if range.is_empty() {
        let syntax = file.syntax();
        let text = syntax.text().slice(range.start()..);
        let pos = match text.find('\n') {
            None => return LocalEdit {
                edit: EditBuilder::new().finish(),
                cursor_position: None
            },
            Some(pos) => pos
        };
        TextRange::offset_len(
            range.start() + pos,
            TextUnit::of_char('\n'),
        )
    } else {
        range
    };

    let node = find_covering_node(file.syntax(), range);
    let mut edit = EditBuilder::new();
    for node in node.descendants() {
        let text = match node.leaf_text() {
            Some(text) => text,
            None => continue,
        };
        let range = match intersect(range, node.range()) {
            Some(range) => range,
            None => continue,
        } - node.range().start();
        for (pos, _) in text[range].bytes().enumerate().filter(|&(_, b)| b == b'\n') {
            let pos: TextUnit = (pos as u32).into();
            let off = node.range().start() + range.start() + pos;
            if !edit.invalidates_offset(off) {
                remove_newline(&mut edit, node, text.as_str(), off);
            }
        }
    }

    LocalEdit {
        edit: edit.finish(),
        cursor_position: None,
    }
}

pub fn on_enter(file: &File, offset: TextUnit) -> Option<LocalEdit> {
    let comment = find_leaf_at_offset(file.syntax(), offset).left_biased().and_then(|it| ast::Comment::cast(it))?;

    if let ast::CommentFlavor::Multiline = comment.flavor() {
        return None;
    }

    let prefix = comment.prefix();
    if offset < comment.syntax().range().start() + TextUnit::of_str(prefix) + TextUnit::from(1) {
        return None;
    }

    let indent = node_indent(file, comment.syntax())?;
    let inserted = format!("\n{}{} ", indent, prefix);
    let cursor_position = offset + TextUnit::of_str(&inserted);
    let mut edit = EditBuilder::new();
    edit.insert(offset, inserted);
    Some(LocalEdit {
        edit: edit.finish(),
        cursor_position: Some(cursor_position),
    })
}

fn node_indent<'a>(file: &'a File, node: SyntaxNodeRef) -> Option<&'a str> {
    let ws = match find_leaf_at_offset(file.syntax(), node.range().start()) {
        LeafAtOffset::Between(l, r) => {
            assert!(r == node);
            l
        }
        LeafAtOffset::Single(n) => {
            assert!(n == node);
            return Some("")
        }
        LeafAtOffset::None => unreachable!(),
    };
    if ws.kind() != WHITESPACE {
        return None;
    }
    let text = ws.leaf_text().unwrap();
    let pos = text.as_str().rfind('\n').map(|it| it + 1).unwrap_or(0);
    Some(&text[pos..])
}

pub fn on_eq_typed(file: &File, offset: TextUnit) -> Option<LocalEdit> {
    let let_stmt: ast::LetStmt = find_node_at_offset(file.syntax(), offset)?;
    if let_stmt.has_semi() {
        return None;
    }
    if let Some(expr) = let_stmt.initializer() {
        let expr_range = expr.syntax().range();
        if contains_offset_nonstrict(expr_range, offset) && offset != expr_range.start() {
            return None;
        }
        if file.syntax().text().slice(offset..expr_range.start()).contains('\n') {
            return None;
        }
    } else {
        return None;
    }
    let offset = let_stmt.syntax().range().end();
    let mut edit = EditBuilder::new();
    edit.insert(offset, ";".to_string());
    Some(LocalEdit {
        edit: edit.finish(),
        cursor_position: None,
    })
}

fn remove_newline(
    edit: &mut EditBuilder,
    node: SyntaxNodeRef,
    node_text: &str,
    offset: TextUnit,
) {
    if node.kind() != WHITESPACE || node_text.bytes().filter(|&b| b == b'\n').count() != 1 {
        // The node is either the first or the last in the file
        let suff = &node_text[TextRange::from_to(
            offset - node.range().start() + TextUnit::of_char('\n'),
            TextUnit::of_str(node_text),
        )];
        let spaces = suff.bytes().take_while(|&b| b == b' ').count();

        edit.replace(
            TextRange::offset_len(offset, ((spaces + 1) as u32).into()),
            " ".to_string(),
        );
        return;
    }

    // Special case that turns something like:
    //
    // ```
    // my_function({<|>
    //    <some-expr>
    // })
    // ```
    //
    // into `my_function(<some-expr>)`
    if join_single_expr_block(edit, node).is_some() {
        return
    }

    // The node is between two other nodes
    let prev = node.prev_sibling().unwrap();
    let next = node.next_sibling().unwrap();
    if is_trailing_comma(prev.kind(), next.kind()) {
        // Removes: trailing comma, newline (incl. surrounding whitespace)
        edit.delete(TextRange::from_to(prev.range().start(), node.range().end()));
    } else if prev.kind() == COMMA && next.kind() == R_CURLY {
        // Removes: comma, newline (incl. surrounding whitespace)
        // Adds: a single whitespace
        edit.replace(
            TextRange::from_to(prev.range().start(), node.range().end()),
            " ".to_string()
        );
    } else if let (Some(_), Some(next)) = (ast::Comment::cast(prev), ast::Comment::cast(next)) {
        // Removes: newline (incl. surrounding whitespace), start of the next comment
        let comment_text = next.text();
        if let Some(newline_pos) = comment_text.find('\n') {
            // Special case for multiline comments: join the comment content but
            // keep the leading `/*`

            let newline_offset = next.syntax().range().start()
                                + TextUnit::from(newline_pos as u32)
                                + TextUnit::of_char('\n');

            edit.insert(newline_offset, "/*".to_string());
            edit.delete(TextRange::from_to(
                node.range().start(),
                next.syntax().range().start() + TextUnit::of_str(next.prefix())
            ));
        } else {
            // Single-line comments
            edit.delete(TextRange::from_to(
                node.range().start(),
                next.syntax().range().start() + TextUnit::of_str(next.prefix())
            ));
        }
    } else {
        // Remove newline but add a computed amount of whitespace characters
        edit.replace(
            node.range(),
            compute_ws(prev, next).to_string(),
        );
    }
}

fn is_trailing_comma(left: SyntaxKind, right: SyntaxKind) -> bool {
    match (left, right) {
       (COMMA, R_PAREN) | (COMMA, R_BRACK) => true,
       _ => false
    }
}

fn join_single_expr_block(
    edit: &mut EditBuilder,
    node: SyntaxNodeRef,
) -> Option<()> {
    let block = ast::Block::cast(node.parent()?)?;
    let block_expr = ast::BlockExpr::cast(block.syntax().parent()?)?;
    let expr = single_expr(block)?;
    edit.replace(
        block_expr.syntax().range(),
        expr.syntax().text().to_string(),
    );
    Some(())
}

fn single_expr(block: ast::Block) -> Option<ast::Expr> {
    let mut res = None;
    for child in block.syntax().children() {
        if let Some(expr) = ast::Expr::cast(child) {
            if expr.syntax().text().contains('\n') {
                return None;
            }
            if mem::replace(&mut res, Some(expr)).is_some() {
                return None;
            }
        } else {
            match child.kind() {
                WHITESPACE | L_CURLY | R_CURLY => (),
                _ => return None,
            }
        }
    }
    res
}

fn compute_ws(left: SyntaxNodeRef, right: SyntaxNodeRef) -> &'static str {
    match left.kind() {
        L_PAREN | L_BRACK => return "",
        _ => (),
    }
    match right.kind() {
        R_PAREN | R_BRACK => return "",
        DOT => return "",
        _ => (),
    }
    " "
}

#[cfg(test)]
mod tests {
    use super::*;
    use test_utils::{check_action, extract_range, extract_offset, add_cursor};

    fn check_join_lines(before: &str, after: &str) {
        check_action(before, after, |file, offset| {
            let range = TextRange::offset_len(offset, 0.into());
            let res = join_lines(file, range);
            Some(res)
        })
    }

    #[test]
    fn test_join_lines_comma() {
        check_join_lines(r"
fn foo() {
    <|>foo(1,
    )
}
", r"
fn foo() {
    <|>foo(1)
}
");
    }

    #[test]
    fn test_join_lines_lambda_block() {
        check_join_lines(r"
pub fn reparse(&self, edit: &AtomEdit) -> File {
    <|>self.incremental_reparse(edit).unwrap_or_else(|| {
        self.full_reparse(edit)
    })
}
", r"
pub fn reparse(&self, edit: &AtomEdit) -> File {
    <|>self.incremental_reparse(edit).unwrap_or_else(|| self.full_reparse(edit))
}
");
    }

    #[test]
    fn test_join_lines_block() {
        check_join_lines(r"
fn foo() {
    foo(<|>{
        92
    })
}", r"
fn foo() {
    foo(<|>92)
}");
    }

    #[test]
    fn test_join_lines_normal_comments() {
        check_join_lines(r"
fn foo() {
    // Hello<|>
    // world!
}
", r"
fn foo() {
    // Hello<|> world!
}
");
    }

    #[test]
    fn test_join_lines_doc_comments() {
        check_join_lines(r"
fn foo() {
    /// Hello<|>
    /// world!
}
", r"
fn foo() {
    /// Hello<|> world!
}
");
    }

    #[test]
    fn test_join_lines_mod_comments() {
        check_join_lines(r"
fn foo() {
    //! Hello<|>
    //! world!
}
", r"
fn foo() {
    //! Hello<|> world!
}
");
    }

    #[test]
    fn test_join_lines_multiline_comments() {
        check_join_lines(r"
fn foo() {
    // Hello<|>
    /* world! */
}
", r"
fn foo() {
    // Hello<|> world! */
}
");
    }

    fn check_join_lines_sel(before: &str, after: &str) {
        let (sel, before) = extract_range(before);
        let file = File::parse(&before);
        let result = join_lines(&file, sel);
        let actual = result.edit.apply(&before);
        assert_eq_text!(after, &actual);
    }

    #[test]
    fn test_join_lines_selection_fn_args() {
        check_join_lines_sel(r"
fn foo() {
    <|>foo(1,
        2,
        3,
    <|>)
}
    ", r"
fn foo() {
    foo(1, 2, 3)
}
    ");
    }

    #[test]
    fn test_join_lines_selection_struct() {
        check_join_lines_sel(r"
struct Foo <|>{
    f: u32,
}<|>
    ", r"
struct Foo { f: u32 }
    ");
    }

    #[test]
    fn test_join_lines_selection_dot_chain() {
        check_join_lines_sel(r"
fn foo() {
    join(<|>type_params.type_params()
            .filter_map(|it| it.name())
            .map(|it| it.text())<|>)
}", r"
fn foo() {
    join(type_params.type_params().filter_map(|it| it.name()).map(|it| it.text()))
}");
    }

    #[test]
    fn test_join_lines_selection_lambda_block_body() {
        check_join_lines_sel(r"
pub fn handle_find_matching_brace() {
    params.offsets
        .map(|offset| <|>{
            world.analysis().matching_brace(&file, offset).unwrap_or(offset)
        }<|>)
        .collect();
}", r"
pub fn handle_find_matching_brace() {
    params.offsets
        .map(|offset| world.analysis().matching_brace(&file, offset).unwrap_or(offset))
        .collect();
}");
    }

    #[test]
    fn test_on_eq_typed() {
        fn do_check(before: &str, after: &str) {
            let (offset, before) = extract_offset(before);
            let file = File::parse(&before);
            let result = on_eq_typed(&file, offset).unwrap();
            let actual = result.edit.apply(&before);
            assert_eq_text!(after, &actual);
        }

        //     do_check(r"
        // fn foo() {
        //     let foo =<|>
        // }
        // ", r"
        // fn foo() {
        //     let foo =;
        // }
        // ");
        do_check(r"
fn foo() {
    let foo =<|> 1 + 1
}
", r"
fn foo() {
    let foo = 1 + 1;
}
");
        //     do_check(r"
        // fn foo() {
        //     let foo =<|>
        //     let bar = 1;
        // }
        // ", r"
        // fn foo() {
        //     let foo =;
        //     let bar = 1;
        // }
        // ");
    }

    #[test]
    fn test_on_enter() {
        fn apply_on_enter(before: &str) -> Option<String> {
            let (offset, before) = extract_offset(before);
            let file = File::parse(&before);
            let result = on_enter(&file, offset)?;
            let actual = result.edit.apply(&before);
            let actual = add_cursor(&actual, result.cursor_position.unwrap());
            Some(actual)
        }

        fn do_check(before: &str, after: &str) {
            let actual = apply_on_enter(before).unwrap();
            assert_eq_text!(after, &actual);
        }

        fn do_check_noop(text: &str) {
            assert!(apply_on_enter(text).is_none())
        }

        do_check(r"
/// Some docs<|>
fn foo() {
}
", r"
/// Some docs
/// <|>
fn foo() {
}
");
        do_check(r"
impl S {
    /// Some<|> docs.
    fn foo() {}
}
", r"
impl S {
    /// Some
    /// <|> docs.
    fn foo() {}
}
");
        do_check_noop(r"<|>//! docz");
    }
}
<commit_msg>Remove smart multiline comment join<commit_after>use std::mem;

use ra_syntax::{
    TextUnit, TextRange, SyntaxNodeRef, File, AstNode, SyntaxKind,
    ast,
    algo::{
        find_covering_node, find_leaf_at_offset, LeafAtOffset,
    },
    text_utils::{intersect, contains_offset_nonstrict},
    SyntaxKind::*,
};

use {LocalEdit, EditBuilder, find_node_at_offset};

pub fn join_lines(file: &File, range: TextRange) -> LocalEdit {
    let range = if range.is_empty() {
        let syntax = file.syntax();
        let text = syntax.text().slice(range.start()..);
        let pos = match text.find('\n') {
            None => return LocalEdit {
                edit: EditBuilder::new().finish(),
                cursor_position: None
            },
            Some(pos) => pos
        };
        TextRange::offset_len(
            range.start() + pos,
            TextUnit::of_char('\n'),
        )
    } else {
        range
    };

    let node = find_covering_node(file.syntax(), range);
    let mut edit = EditBuilder::new();
    for node in node.descendants() {
        let text = match node.leaf_text() {
            Some(text) => text,
            None => continue,
        };
        let range = match intersect(range, node.range()) {
            Some(range) => range,
            None => continue,
        } - node.range().start();
        for (pos, _) in text[range].bytes().enumerate().filter(|&(_, b)| b == b'\n') {
            let pos: TextUnit = (pos as u32).into();
            let off = node.range().start() + range.start() + pos;
            if !edit.invalidates_offset(off) {
                remove_newline(&mut edit, node, text.as_str(), off);
            }
        }
    }

    LocalEdit {
        edit: edit.finish(),
        cursor_position: None,
    }
}

pub fn on_enter(file: &File, offset: TextUnit) -> Option<LocalEdit> {
    let comment = find_leaf_at_offset(file.syntax(), offset).left_biased().and_then(|it| ast::Comment::cast(it))?;

    if let ast::CommentFlavor::Multiline = comment.flavor() {
        return None;
    }

    let prefix = comment.prefix();
    if offset < comment.syntax().range().start() + TextUnit::of_str(prefix) + TextUnit::from(1) {
        return None;
    }

    let indent = node_indent(file, comment.syntax())?;
    let inserted = format!("\n{}{} ", indent, prefix);
    let cursor_position = offset + TextUnit::of_str(&inserted);
    let mut edit = EditBuilder::new();
    edit.insert(offset, inserted);
    Some(LocalEdit {
        edit: edit.finish(),
        cursor_position: Some(cursor_position),
    })
}

fn node_indent<'a>(file: &'a File, node: SyntaxNodeRef) -> Option<&'a str> {
    let ws = match find_leaf_at_offset(file.syntax(), node.range().start()) {
        LeafAtOffset::Between(l, r) => {
            assert!(r == node);
            l
        }
        LeafAtOffset::Single(n) => {
            assert!(n == node);
            return Some("")
        }
        LeafAtOffset::None => unreachable!(),
    };
    if ws.kind() != WHITESPACE {
        return None;
    }
    let text = ws.leaf_text().unwrap();
    let pos = text.as_str().rfind('\n').map(|it| it + 1).unwrap_or(0);
    Some(&text[pos..])
}

pub fn on_eq_typed(file: &File, offset: TextUnit) -> Option<LocalEdit> {
    let let_stmt: ast::LetStmt = find_node_at_offset(file.syntax(), offset)?;
    if let_stmt.has_semi() {
        return None;
    }
    if let Some(expr) = let_stmt.initializer() {
        let expr_range = expr.syntax().range();
        if contains_offset_nonstrict(expr_range, offset) && offset != expr_range.start() {
            return None;
        }
        if file.syntax().text().slice(offset..expr_range.start()).contains('\n') {
            return None;
        }
    } else {
        return None;
    }
    let offset = let_stmt.syntax().range().end();
    let mut edit = EditBuilder::new();
    edit.insert(offset, ";".to_string());
    Some(LocalEdit {
        edit: edit.finish(),
        cursor_position: None,
    })
}

fn remove_newline(
    edit: &mut EditBuilder,
    node: SyntaxNodeRef,
    node_text: &str,
    offset: TextUnit,
) {
    if node.kind() != WHITESPACE || node_text.bytes().filter(|&b| b == b'\n').count() != 1 {
        // The node is either the first or the last in the file
        let suff = &node_text[TextRange::from_to(
            offset - node.range().start() + TextUnit::of_char('\n'),
            TextUnit::of_str(node_text),
        )];
        let spaces = suff.bytes().take_while(|&b| b == b' ').count();

        edit.replace(
            TextRange::offset_len(offset, ((spaces + 1) as u32).into()),
            " ".to_string(),
        );
        return;
    }

    // Special case that turns something like:
    //
    // ```
    // my_function({<|>
    //    <some-expr>
    // })
    // ```
    //
    // into `my_function(<some-expr>)`
    if join_single_expr_block(edit, node).is_some() {
        return
    }

    // The node is between two other nodes
    let prev = node.prev_sibling().unwrap();
    let next = node.next_sibling().unwrap();
    if is_trailing_comma(prev.kind(), next.kind()) {
        // Removes: trailing comma, newline (incl. surrounding whitespace)
        edit.delete(TextRange::from_to(prev.range().start(), node.range().end()));
    } else if prev.kind() == COMMA && next.kind() == R_CURLY {
        // Removes: comma, newline (incl. surrounding whitespace)
        // Adds: a single whitespace
        edit.replace(
            TextRange::from_to(prev.range().start(), node.range().end()),
            " ".to_string()
        );
    } else if let (Some(_), Some(next)) = (ast::Comment::cast(prev), ast::Comment::cast(next)) {
        // Removes: newline (incl. surrounding whitespace), start of the next comment
        edit.delete(TextRange::from_to(
            node.range().start(),
            next.syntax().range().start() + TextUnit::of_str(next.prefix())
        ));
    } else {
        // Remove newline but add a computed amount of whitespace characters
        edit.replace(
            node.range(),
            compute_ws(prev, next).to_string(),
        );
    }
}

fn is_trailing_comma(left: SyntaxKind, right: SyntaxKind) -> bool {
    match (left, right) {
       (COMMA, R_PAREN) | (COMMA, R_BRACK) => true,
       _ => false
    }
}

fn join_single_expr_block(
    edit: &mut EditBuilder,
    node: SyntaxNodeRef,
) -> Option<()> {
    let block = ast::Block::cast(node.parent()?)?;
    let block_expr = ast::BlockExpr::cast(block.syntax().parent()?)?;
    let expr = single_expr(block)?;
    edit.replace(
        block_expr.syntax().range(),
        expr.syntax().text().to_string(),
    );
    Some(())
}

fn single_expr(block: ast::Block) -> Option<ast::Expr> {
    let mut res = None;
    for child in block.syntax().children() {
        if let Some(expr) = ast::Expr::cast(child) {
            if expr.syntax().text().contains('\n') {
                return None;
            }
            if mem::replace(&mut res, Some(expr)).is_some() {
                return None;
            }
        } else {
            match child.kind() {
                WHITESPACE | L_CURLY | R_CURLY => (),
                _ => return None,
            }
        }
    }
    res
}

fn compute_ws(left: SyntaxNodeRef, right: SyntaxNodeRef) -> &'static str {
    match left.kind() {
        L_PAREN | L_BRACK => return "",
        _ => (),
    }
    match right.kind() {
        R_PAREN | R_BRACK => return "",
        DOT => return "",
        _ => (),
    }
    " "
}

#[cfg(test)]
mod tests {
    use super::*;
    use test_utils::{check_action, extract_range, extract_offset, add_cursor};

    fn check_join_lines(before: &str, after: &str) {
        check_action(before, after, |file, offset| {
            let range = TextRange::offset_len(offset, 0.into());
            let res = join_lines(file, range);
            Some(res)
        })
    }

    #[test]
    fn test_join_lines_comma() {
        check_join_lines(r"
fn foo() {
    <|>foo(1,
    )
}
", r"
fn foo() {
    <|>foo(1)
}
");
    }

    #[test]
    fn test_join_lines_lambda_block() {
        check_join_lines(r"
pub fn reparse(&self, edit: &AtomEdit) -> File {
    <|>self.incremental_reparse(edit).unwrap_or_else(|| {
        self.full_reparse(edit)
    })
}
", r"
pub fn reparse(&self, edit: &AtomEdit) -> File {
    <|>self.incremental_reparse(edit).unwrap_or_else(|| self.full_reparse(edit))
}
");
    }

    #[test]
    fn test_join_lines_block() {
        check_join_lines(r"
fn foo() {
    foo(<|>{
        92
    })
}", r"
fn foo() {
    foo(<|>92)
}");
    }

    #[test]
    fn test_join_lines_normal_comments() {
        check_join_lines(r"
fn foo() {
    // Hello<|>
    // world!
}
", r"
fn foo() {
    // Hello<|> world!
}
");
    }

    #[test]
    fn test_join_lines_doc_comments() {
        check_join_lines(r"
fn foo() {
    /// Hello<|>
    /// world!
}
", r"
fn foo() {
    /// Hello<|> world!
}
");
    }

    #[test]
    fn test_join_lines_mod_comments() {
        check_join_lines(r"
fn foo() {
    //! Hello<|>
    //! world!
}
", r"
fn foo() {
    //! Hello<|> world!
}
");
    }

    #[test]
    fn test_join_lines_multiline_comments_1() {
        check_join_lines(r"
fn foo() {
    // Hello<|>
    /* world! */
}
", r"
fn foo() {
    // Hello<|> world! */
}
");
    }

    #[test]
    fn test_join_lines_multiline_comments_2() {
        check_join_lines(r"
fn foo() {
    // The<|>
    /* quick
    brown
    fox! */
}
", r"
fn foo() {
    // The<|> quick
    brown
    fox! */
}
");
    }

    fn check_join_lines_sel(before: &str, after: &str) {
        let (sel, before) = extract_range(before);
        let file = File::parse(&before);
        let result = join_lines(&file, sel);
        let actual = result.edit.apply(&before);
        assert_eq_text!(after, &actual);
    }

    #[test]
    fn test_join_lines_selection_fn_args() {
        check_join_lines_sel(r"
fn foo() {
    <|>foo(1,
        2,
        3,
    <|>)
}
    ", r"
fn foo() {
    foo(1, 2, 3)
}
    ");
    }

    #[test]
    fn test_join_lines_selection_struct() {
        check_join_lines_sel(r"
struct Foo <|>{
    f: u32,
}<|>
    ", r"
struct Foo { f: u32 }
    ");
    }

    #[test]
    fn test_join_lines_selection_dot_chain() {
        check_join_lines_sel(r"
fn foo() {
    join(<|>type_params.type_params()
            .filter_map(|it| it.name())
            .map(|it| it.text())<|>)
}", r"
fn foo() {
    join(type_params.type_params().filter_map(|it| it.name()).map(|it| it.text()))
}");
    }

    #[test]
    fn test_join_lines_selection_lambda_block_body() {
        check_join_lines_sel(r"
pub fn handle_find_matching_brace() {
    params.offsets
        .map(|offset| <|>{
            world.analysis().matching_brace(&file, offset).unwrap_or(offset)
        }<|>)
        .collect();
}", r"
pub fn handle_find_matching_brace() {
    params.offsets
        .map(|offset| world.analysis().matching_brace(&file, offset).unwrap_or(offset))
        .collect();
}");
    }

    #[test]
    fn test_on_eq_typed() {
        fn do_check(before: &str, after: &str) {
            let (offset, before) = extract_offset(before);
            let file = File::parse(&before);
            let result = on_eq_typed(&file, offset).unwrap();
            let actual = result.edit.apply(&before);
            assert_eq_text!(after, &actual);
        }

        //     do_check(r"
        // fn foo() {
        //     let foo =<|>
        // }
        // ", r"
        // fn foo() {
        //     let foo =;
        // }
        // ");
        do_check(r"
fn foo() {
    let foo =<|> 1 + 1
}
", r"
fn foo() {
    let foo = 1 + 1;
}
");
        //     do_check(r"
        // fn foo() {
        //     let foo =<|>
        //     let bar = 1;
        // }
        // ", r"
        // fn foo() {
        //     let foo =;
        //     let bar = 1;
        // }
        // ");
    }

    #[test]
    fn test_on_enter() {
        fn apply_on_enter(before: &str) -> Option<String> {
            let (offset, before) = extract_offset(before);
            let file = File::parse(&before);
            let result = on_enter(&file, offset)?;
            let actual = result.edit.apply(&before);
            let actual = add_cursor(&actual, result.cursor_position.unwrap());
            Some(actual)
        }

        fn do_check(before: &str, after: &str) {
            let actual = apply_on_enter(before).unwrap();
            assert_eq_text!(after, &actual);
        }

        fn do_check_noop(text: &str) {
            assert!(apply_on_enter(text).is_none())
        }

        do_check(r"
/// Some docs<|>
fn foo() {
}
", r"
/// Some docs
/// <|>
fn foo() {
}
");
        do_check(r"
impl S {
    /// Some<|> docs.
    fn foo() {}
}
", r"
impl S {
    /// Some
    /// <|> docs.
    fn foo() {}
}
");
        do_check_noop(r"<|>//! docz");
    }
}
<|endoftext|>


<commit_before>generate_error_module!(
    generate_error_types!(TagError, TagErrorKind,
        TagTypeError     => "Entry Header Tag Type wrong",
        HeaderReadError  => "Error while reading entry header",
        HeaderWriteError => "Error while writing entry header",
        NotATag          => "String is not a tag"
    );
);

pub use self::error::TagError;
pub use self::error::TagErrorKind;

<commit_msg>Export MapErrInto trait<commit_after>generate_error_module!(
    generate_error_types!(TagError, TagErrorKind,
        TagTypeError     => "Entry Header Tag Type wrong",
        HeaderReadError  => "Error while reading entry header",
        HeaderWriteError => "Error while writing entry header",
        NotATag          => "String is not a tag"
    );
);

pub use self::error::TagError;
pub use self::error::TagErrorKind;
pub use self::error::MapErrInto;

<|endoftext|>


<commit_before>use boring::ssl::ConnectConfiguration;
use std::net::{IpAddr, SocketAddr};

use tokio::io::AsyncWriteExt;
use tokio::net::{TcpListener, TcpStream};
use tracing::{debug, error, info, warn};

use crate::config::Config;
use crate::proxy::Error;
use crate::workload::{Protocol, Workload, WorkloadInformation};
use crate::{identity, socket};

pub struct Outbound {
    cfg: Config,
    cert_manager: identity::SecretManager,
    workloads: WorkloadInformation,
    listener: TcpListener,
}

impl Outbound {
    pub async fn new(
        cfg: Config,
        cert_manager: identity::SecretManager,
        workloads: WorkloadInformation,
    ) -> Result<Outbound, Error> {
        let listener: TcpListener = TcpListener::bind(cfg.outbound_addr)
            .await
            .map_err(Error::Bind)?;
        match socket::set_transparent(&listener) {
            Err(_e) => info!("running without transparent mode"),
            _ => info!("running with transparent mode"),
        };

        Ok(Outbound {
            cfg,
            cert_manager,
            workloads,
            listener,
        })
    }

    pub(super) async fn run(self) {
        let addr = self.listener.local_addr().unwrap();
        info!("outbound listener established {}", addr);

        loop {
            // Asynchronously wait for an inbound socket.
            let socket = self.listener.accept().await;
            match socket {
                Ok((stream, remote)) => {
                    info!("accepted outbound connection from {}", remote);
                    let cfg = self.cfg.clone();
                    let oc = OutboundConnection {
                        cert_manager: self.cert_manager.clone(),
                        workloads: self.workloads.clone(),
                        cfg,
                    };
                    tokio::spawn(async move {
                        let res = oc.proxy(stream).await;
                        match res {
                            Ok(_) => info!("outbound proxy complete"),
                            Err(ref e) => warn!("outbound proxy failed: {}", e),
                        };
                    });
                }
                Err(e) => error!("Failed TCP handshake {}", e),
            }
        }
    }
}

struct OutboundConnection {
    cert_manager: identity::SecretManager,
    workloads: WorkloadInformation,
    // TODO: Config may be excessively large, maybe we store a scoped OutboundConfig intended for cloning.
    cfg: Config,
}

impl OutboundConnection {
    async fn proxy(&self, mut stream: TcpStream) -> Result<(), Error> {
        let remote_addr =
            super::to_canonical_ip(stream.peer_addr().expect("must receive peer addr"));
        let orig = socket::orig_dst_addr(&stream).expect("must have original dst enabled");
        let req = self.build_request(remote_addr, orig).await;
        debug!("request from {} to {}", req.source.name, orig);
        match req.protocol {
            Protocol::Hbone => {
                info!(
                    "Proxying to {} using HBONE via {} type {:#?}",
                    req.destination, req.gateway, req.request_type
                );

                // Using the raw connection API, instead of client, is a bit annoying, but the only reasonable
                // way to work around https://github.com/hyperium/hyper/issues/2863
                // Eventually we will need to implement our own smarter pooling, TLS handshaking, etc anyways.
                let mut builder = hyper::client::conn::Builder::new();
                let builder = builder
                    .http2_only(true)
                    .http2_initial_stream_window_size(self.cfg.window_size)
                    .http2_max_frame_size(self.cfg.frame_size)
                    .http2_initial_connection_window_size(self.cfg.connection_window_size);

                let request = hyper::Request::builder()
                    .uri(&req.destination.to_string())
                    .method(hyper::Method::CONNECT)
                    .version(hyper::Version::HTTP_2)
                    .body(hyper::Body::empty())
                    .unwrap();

                let mut request_sender = if self.cfg.tls {
                    let id = req.source.identity();
                    let cert = self.cert_manager.fetch_certificate(id).await?;
                    let connector = cert.connector()?.configure()?;
                    let tcp_stream = TcpStream::connect(req.gateway).await?;
                    let tls_stream = connect_tls(connector, tcp_stream).await?;
                    let (request_sender, connection) = builder
                        .handshake(tls_stream)
                        .await
                        .map_err(Error::HttpHandshake)?;
                    // spawn a task to poll the connection and drive the HTTP state
                    tokio::spawn(async move {
                        if let Err(e) = connection.await {
                            error!("Error in HBONE connection handshake: {:?}", e);
                        }
                    });
                    request_sender
                } else {
                    let tcp_stream = TcpStream::connect(req.gateway).await?;
                    let (request_sender, connection) = builder
                        .handshake::<TcpStream, hyper::Body>(tcp_stream)
                        .await?;
                    // spawn a task to poll the connection and drive the HTTP state
                    tokio::spawn(async move {
                        if let Err(e) = connection.await {
                            error!("Error in connection: {}", e);
                        }
                    });
                    request_sender
                };

                let response = request_sender.send_request(request).await?;

                let code = response.status();
                match hyper::upgrade::on(response).await {
                    Ok(mut upgraded) => {
                        super::copy_hbone("hbone client", &mut upgraded, &mut stream)
                            .await
                            .expect("hbone client copy");
                    }
                    Err(e) => eprintln!("upgrade error: {}, {}", e, code),
                }
                info!("request complete");
                Ok(())
            }
            Protocol::Tcp => {
                info!(
                    "Proxying to {} using TCP via {} type {:?}",
                    req.destination, req.gateway, req.request_type
                );
                let mut outbound = TcpStream::connect(req.gateway).await?;

                let (mut ri, mut wi) = stream.split();
                let (mut ro, mut wo) = outbound.split();

                let client_to_server = async {
                    tokio::io::copy(&mut ri, &mut wo).await?;
                    wo.shutdown().await
                };

                let server_to_client = async {
                    tokio::io::copy(&mut ro, &mut wi).await?;
                    wi.shutdown().await
                };

                tokio::try_join!(client_to_server, server_to_client)?;

                Ok(())
            }
        }
    }

    async fn build_request(&self, downstream: IpAddr, target: SocketAddr) -> Request {
        let (source_workload, us, is_vip) = {
            let source_workload = self
                .workloads
                .fetch_workload(&downstream)
                .await
                .expect("todo: source must be found");

            // TODO: we want a single lock for source and upstream probably...?
            let (us, is_vip) = self.workloads.find_upstream(target).await;
            (source_workload, us, is_vip)
        };
        let mut req = Request {
            protocol: us.workload.protocol,
            source: source_workload.clone(), // TODO drop clone
            destination: SocketAddr::from((us.workload.workload_ip, us.port)),
            gateway: us
                .workload
                .gateway_ip
                .expect("todo: refactor gateway ip handling"),
            direction: Direction::Outbound, // TODO set this
            request_type: RequestType::Direct,
        };
        if source_workload.waypoint_address.is_some() {
            // Source has a remote proxy. We should delegate everything to that proxy - do not even resolve VIP.
            // TODO: add client skipping
            req.request_type = RequestType::ToClientWaypoint;
            // Let the client remote know we are on the outbound path. The remote proxy should strictly
            // validate the identity when we declare this
            req.direction = Direction::Outbound;
            // Load balancing decision is deferred to remote proxy
            req.destination = target;
            // Send to the remote proxy
            req.gateway = SocketAddr::from((source_workload.waypoint_address.unwrap(), 15001));
            // Always use HBONE here
            req.protocol = Protocol::Hbone;
        } else if us.workload.waypoint_address.is_some() {
            // TODO: even in this case, we are picking a single upstream pod and deciding if it has a remote proxy.
            // Typically this is all or nothing, but if not we should probably send to remote proxy if *any* upstream has one.
            if is_vip {
                // Use the original VIP, not translated
                req.destination = target
            }
            req.request_type = RequestType::ToServerWaypoint;
            // Always use HBONE here
            req.protocol = Protocol::Hbone;
            // Let the client remote know we are on the inbound path.
            req.direction = Direction::Inbound;
            req.gateway = SocketAddr::from((us.workload.waypoint_address.unwrap(), 15006));
        } else if !us.workload.node.is_empty()
            && self.cfg.local_node == Some(us.workload.node)
            && req.protocol == Protocol::Hbone
        {
            // Sending to a node on the same node (ourselves). Requests from the node proxy are not captured,
            // so we need to explicitly send it to ourselves.
            // In the future this could be optimized to avoid a full network traversal.
            req.request_type = RequestType::DirectLocal;
            req.gateway = "127.0.0.1:15088".parse().unwrap();
        } else if us.workload.name.is_empty() {
            req.request_type = RequestType::Passthrough;
        } else {
            req.request_type = RequestType::Direct;
        }
        req
    }
}

#[derive(Debug)]
struct Request {
    protocol: Protocol,
    direction: Direction,
    source: Workload,
    destination: SocketAddr,
    gateway: SocketAddr,
    request_type: RequestType,
}

#[derive(Debug)]
enum Direction {
    Inbound,
    Outbound,
}

#[derive(Debug)]
enum RequestType {
    ToClientWaypoint,
    ToServerWaypoint,
    Direct,
    DirectLocal,
    Passthrough,
}

async fn connect_tls(
    mut connector: ConnectConfiguration,
    stream: TcpStream,
) -> Result<tokio_boring::SslStream<TcpStream>, tokio_boring::HandshakeError<TcpStream>> {
    connector.set_verify_hostname(false);
    connector.set_use_server_name_indication(false);
    let addr = stream.local_addr();
    connector.set_verify_callback(boring::ssl::SslVerifyMode::PEER, move |_, x509| {
        info!("TLS callback for {:?}: {:?}", addr, x509.error());
        true
    });
    tokio_boring::connect(connector, "", stream).await
}
<commit_msg>Fix calls to pods on the same node (#59)<commit_after>use boring::ssl::ConnectConfiguration;
use std::net::{IpAddr, SocketAddr};

use tokio::io::AsyncWriteExt;
use tokio::net::{TcpListener, TcpStream};
use tracing::{debug, error, info, warn};

use crate::config::Config;
use crate::proxy::Error;
use crate::workload::{Protocol, Workload, WorkloadInformation};
use crate::{identity, socket};

pub struct Outbound {
    cfg: Config,
    cert_manager: identity::SecretManager,
    workloads: WorkloadInformation,
    listener: TcpListener,
}

impl Outbound {
    pub async fn new(
        cfg: Config,
        cert_manager: identity::SecretManager,
        workloads: WorkloadInformation,
    ) -> Result<Outbound, Error> {
        let listener: TcpListener = TcpListener::bind(cfg.outbound_addr)
            .await
            .map_err(Error::Bind)?;
        match socket::set_transparent(&listener) {
            Err(_e) => info!("running without transparent mode"),
            _ => info!("running with transparent mode"),
        };

        Ok(Outbound {
            cfg,
            cert_manager,
            workloads,
            listener,
        })
    }

    pub(super) async fn run(self) {
        let addr = self.listener.local_addr().unwrap();
        info!("outbound listener established {}", addr);

        loop {
            // Asynchronously wait for an inbound socket.
            let socket = self.listener.accept().await;
            match socket {
                Ok((stream, remote)) => {
                    info!("accepted outbound connection from {}", remote);
                    let cfg = self.cfg.clone();
                    let oc = OutboundConnection {
                        cert_manager: self.cert_manager.clone(),
                        workloads: self.workloads.clone(),
                        cfg,
                    };
                    tokio::spawn(async move {
                        let res = oc.proxy(stream).await;
                        match res {
                            Ok(_) => info!("outbound proxy complete"),
                            Err(ref e) => warn!("outbound proxy failed: {}", e),
                        };
                    });
                }
                Err(e) => error!("Failed TCP handshake {}", e),
            }
        }
    }
}

struct OutboundConnection {
    cert_manager: identity::SecretManager,
    workloads: WorkloadInformation,
    // TODO: Config may be excessively large, maybe we store a scoped OutboundConfig intended for cloning.
    cfg: Config,
}

impl OutboundConnection {
    async fn proxy(&self, mut stream: TcpStream) -> Result<(), Error> {
        let remote_addr =
            super::to_canonical_ip(stream.peer_addr().expect("must receive peer addr"));
        let orig = socket::orig_dst_addr(&stream).expect("must have original dst enabled");
        let req = self.build_request(remote_addr, orig).await;
        debug!("request from {} to {}", req.source.name, orig);
        match req.protocol {
            Protocol::Hbone => {
                info!(
                    "Proxying to {} using HBONE via {} type {:#?}",
                    req.destination, req.gateway, req.request_type
                );

                // Using the raw connection API, instead of client, is a bit annoying, but the only reasonable
                // way to work around https://github.com/hyperium/hyper/issues/2863
                // Eventually we will need to implement our own smarter pooling, TLS handshaking, etc anyways.
                let mut builder = hyper::client::conn::Builder::new();
                let builder = builder
                    .http2_only(true)
                    .http2_initial_stream_window_size(self.cfg.window_size)
                    .http2_max_frame_size(self.cfg.frame_size)
                    .http2_initial_connection_window_size(self.cfg.connection_window_size);

                let request = hyper::Request::builder()
                    .uri(&req.destination.to_string())
                    .method(hyper::Method::CONNECT)
                    .version(hyper::Version::HTTP_2)
                    .body(hyper::Body::empty())
                    .unwrap();

                let mut request_sender = if self.cfg.tls {
                    let id = req.source.identity();
                    let cert = self.cert_manager.fetch_certificate(id).await?;
                    let connector = cert.connector()?.configure()?;
                    let tcp_stream = TcpStream::connect(req.gateway).await?;
                    let tls_stream = connect_tls(connector, tcp_stream).await?;
                    let (request_sender, connection) = builder
                        .handshake(tls_stream)
                        .await
                        .map_err(Error::HttpHandshake)?;
                    // spawn a task to poll the connection and drive the HTTP state
                    tokio::spawn(async move {
                        if let Err(e) = connection.await {
                            error!("Error in HBONE connection handshake: {:?}", e);
                        }
                    });
                    request_sender
                } else {
                    let tcp_stream = TcpStream::connect(req.gateway).await?;
                    let (request_sender, connection) = builder
                        .handshake::<TcpStream, hyper::Body>(tcp_stream)
                        .await?;
                    // spawn a task to poll the connection and drive the HTTP state
                    tokio::spawn(async move {
                        if let Err(e) = connection.await {
                            error!("Error in connection: {}", e);
                        }
                    });
                    request_sender
                };

                let response = request_sender.send_request(request).await?;

                let code = response.status();
                match hyper::upgrade::on(response).await {
                    Ok(mut upgraded) => {
                        super::copy_hbone("hbone client", &mut upgraded, &mut stream)
                            .await
                            .expect("hbone client copy");
                    }
                    Err(e) => eprintln!("upgrade error: {}, {}", e, code),
                }
                info!("request complete");
                Ok(())
            }
            Protocol::Tcp => {
                info!(
                    "Proxying to {} using TCP via {} type {:?}",
                    req.destination, req.gateway, req.request_type
                );
                let mut outbound = TcpStream::connect(req.gateway).await?;

                let (mut ri, mut wi) = stream.split();
                let (mut ro, mut wo) = outbound.split();

                let client_to_server = async {
                    tokio::io::copy(&mut ri, &mut wo).await?;
                    wo.shutdown().await
                };

                let server_to_client = async {
                    tokio::io::copy(&mut ro, &mut wi).await?;
                    wi.shutdown().await
                };

                tokio::try_join!(client_to_server, server_to_client)?;

                Ok(())
            }
        }
    }

    async fn build_request(&self, downstream: IpAddr, target: SocketAddr) -> Request {
        let (source_workload, us, is_vip) = {
            let source_workload = self
                .workloads
                .fetch_workload(&downstream)
                .await
                .expect("todo: source must be found");

            // TODO: we want a single lock for source and upstream probably...?
            let (us, is_vip) = self.workloads.find_upstream(target).await;
            (source_workload, us, is_vip)
        };
        let mut req = Request {
            protocol: us.workload.protocol,
            source: source_workload.clone(), // TODO drop clone
            destination: SocketAddr::from((us.workload.workload_ip, us.port)),
            gateway: us
                .workload
                .gateway_ip
                .expect("todo: refactor gateway ip handling"),
            direction: Direction::Outbound, // TODO set this
            request_type: RequestType::Direct,
        };
        if source_workload.waypoint_address.is_some() {
            // Source has a remote proxy. We should delegate everything to that proxy - do not even resolve VIP.
            // TODO: add client skipping
            req.request_type = RequestType::ToClientWaypoint;
            // Let the client remote know we are on the outbound path. The remote proxy should strictly
            // validate the identity when we declare this
            req.direction = Direction::Outbound;
            // Load balancing decision is deferred to remote proxy
            req.destination = target;
            // Send to the remote proxy
            req.gateway = SocketAddr::from((source_workload.waypoint_address.unwrap(), 15001));
            // Always use HBONE here
            req.protocol = Protocol::Hbone;
        } else if us.workload.waypoint_address.is_some() {
            // TODO: even in this case, we are picking a single upstream pod and deciding if it has a remote proxy.
            // Typically this is all or nothing, but if not we should probably send to remote proxy if *any* upstream has one.
            if is_vip {
                // Use the original VIP, not translated
                req.destination = target
            }
            req.request_type = RequestType::ToServerWaypoint;
            // Always use HBONE here
            req.protocol = Protocol::Hbone;
            // Let the client remote know we are on the inbound path.
            req.direction = Direction::Inbound;
            req.gateway = SocketAddr::from((us.workload.waypoint_address.unwrap(), 15006));
        } else if !us.workload.node.is_empty()
            && self.cfg.local_node == Some(us.workload.node)
            && req.protocol == Protocol::Hbone
        {
            // Sending to a node on the same node (ourselves).
            // In the future this could be optimized to avoid a full network traversal.
            req.request_type = RequestType::DirectLocal;
            // We would want to send to 127.0.0.1:15008 in theory. However, the inbound listener
            // expects to lookup the desired certificate based on the destination IP. If we send directly,
            // we would try to lookup an IP for 127.0.0.1.
            // Instead, we send to the actual IP, but iptables in the pod ensures traffic is redirected to 15008.
            req.gateway = SocketAddr::from((req.gateway.ip(), 15088));
        } else if us.workload.name.is_empty() {
            req.request_type = RequestType::Passthrough;
        } else {
            req.request_type = RequestType::Direct;
        }
        req
    }
}

#[derive(Debug)]
struct Request {
    protocol: Protocol,
    direction: Direction,
    source: Workload,
    destination: SocketAddr,
    gateway: SocketAddr,
    request_type: RequestType,
}

#[derive(Debug)]
enum Direction {
    Inbound,
    Outbound,
}

#[derive(Debug)]
enum RequestType {
    ToClientWaypoint,
    ToServerWaypoint,
    Direct,
    DirectLocal,
    Passthrough,
}

async fn connect_tls(
    mut connector: ConnectConfiguration,
    stream: TcpStream,
) -> Result<tokio_boring::SslStream<TcpStream>, tokio_boring::HandshakeError<TcpStream>> {
    connector.set_verify_hostname(false);
    connector.set_use_server_name_indication(false);
    let addr = stream.local_addr();
    connector.set_verify_callback(boring::ssl::SslVerifyMode::PEER, move |_, x509| {
        info!("TLS callback for {:?}: {:?}", addr, x509.error());
        true
    });
    tokio_boring::connect(connector, "", stream).await
}
<|endoftext|>


<commit_before>use std::mem;
use std::ptr;
use std::marker::PhantomData;
use gl;
use gl::types::*;
use context::Context;
use framebuffer::FramebufferBinding;
use types::{DrawingMode, GLError};
use index_data::{IndexData, IndexDatum, IndexDatumType};

pub struct Buffer {
    gl_id: GLuint
}

impl Buffer {
    pub fn gl_id(&self) -> GLuint {
        self.gl_id
    }
}

impl Drop for Buffer {
    fn drop(&mut self) {
        unsafe {
            gl::DeleteBuffers(1, &self.gl_id as *const GLuint);
        }
    }
}

impl Context {
    pub fn gen_buffer(&self) -> Buffer {
        unsafe {
            let mut id : GLuint = 0;

            gl::GenBuffers(1, &mut id as *mut GLuint);
            dbg_gl_sanity_check! {
                GLError::InvalidValue => "`n` is negative",
                _ => "Unknown error"
            }

            Buffer { gl_id: id }
        }
    }
}

gl_enum! {
    pub gl_enum BufferDataUsage {
        StreamDraw as STREAM_DRAW = gl::STREAM_DRAW,
        StaticDraw as STATIC_DRAW = gl::STATIC_DRAW,
        DynamicDraw as DYNAMIC_DRAW = gl::DYNAMIC_DRAW
    }
}

gl_enum! {
    pub gl_enum BufferBindingTarget {
        ArrayBuffer as ARRAY_BUFFER = gl::ARRAY_BUFFER,
        ElementArrayBuffer as ELEMENT_ARRAY_BUFFER = gl::ELEMENT_ARRAY_BUFFER
    }
}

pub trait BufferBinding {
    fn target(&self) -> BufferBindingTarget;

    fn buffer_bytes(&mut self, bytes: &[u8], usage: BufferDataUsage) {
        let ptr = bytes.as_ptr();
        let size = bytes.len() * mem::size_of::<u8>();
        unsafe {
            gl::BufferData(self.target().gl_enum(),
                           size as GLsizeiptr,
                           ptr as *const GLvoid,
                           usage.gl_enum());
            dbg_gl_error! {
                GLError::InvalidEnum => "Invalid `target` or `usage`",
                GLError::InvalidValue => "`size` is negative",
                GLError::InvalidOperation => "Object 0 is bound to buffer target",
                GLError::OutOfMemory => "Unable to create a large enough buffer",
                _ => "Unknown error"
            }
        }
    }
}

pub struct ArrayBufferBinding<'a> {
    phantom: PhantomData<&'a mut Buffer>
}

impl<'a> BufferBinding for ArrayBufferBinding<'a> {
    fn target(&self) -> BufferBindingTarget {
        ARRAY_BUFFER
    }
}

unsafe fn _draw_elements(mode: DrawingMode,
                         count: usize,
                         index_type: IndexDatumType,
                         indicies: *const GLvoid)
{
    let gl_index_type: GLenum = match index_type {
        IndexDatumType::UnsignedByte => gl::UNSIGNED_BYTE,
        IndexDatumType::UnsignedShort => gl::UNSIGNED_SHORT
    };
    gl::DrawElements(mode.gl_enum(),
                     count as GLsizei,
                     gl_index_type,
                     indicies);
    dbg_gl_error! {
        GLError::InvalidEnum => "`mode` or `type` is not an accepted value",
        GLError::InvalidValue => "`count` is negative",
        GLError::InvalidFramebufferOperation => "The current framebuffer is not framebuffer-complete",
        _ => "Unknown error"
    }
}

impl<'a> ArrayBufferBinding<'a> {
    pub unsafe fn vertex_attrib_pointer(&self,
                                        attrib: super::ProgramAttrib,
                                        components: i8,
                                        gl_type: super::DataType,
                                        normalized: bool,
                                        stride: usize,
                                        offset: usize)
    {
        debug_assert!(1 <= components && components <= 4);

        let gl_normalized = if normalized { gl::TRUE } else { gl::FALSE };
        gl::VertexAttribPointer(attrib.gl_index,
                                components as GLint,
                                gl_type.gl_enum(),
                                gl_normalized,
                                stride as GLsizei,
                                offset as *const GLvoid);
        dbg_gl_error! {
            GLError::InvalidEnum => "Illegal vertex attribute type",
            GLError::InvalidValue => "`stride` is negative, `size` is not in range, or `index` is >= GL_MAX_VERTEX_ATTRIBS",
            GLError::InvalidFramebufferOperation => "Currently bound framebuffer is not framebuffer complete",
            _ => "Unknown error"
        }
    }
}

impl<'a> FramebufferBinding<'a> {
    pub unsafe fn draw_arrays_range(&mut self,
                                    _ab: &ArrayBufferBinding,
                                    mode: DrawingMode,
                                    first: u32,
                                    count: usize)
    {
        gl::DrawArrays(mode.gl_enum(), first as GLint, count as GLsizei);
        dbg_gl_sanity_check! {
            GLError::InvalidEnum => "`mode` is not an accepted value",
            GLError::InvalidValue => "`count` is negative",
            _ => "Unknown error"
        }
    }

    pub unsafe fn draw_n_elements_buffered(&mut self,
                                           _ab: &ArrayBufferBinding,
                                           _eab: &ElementArrayBufferBinding,
                                           mode: DrawingMode,
                                           count: usize,
                                           index_type: IndexDatumType)
    {
        _draw_elements(mode, count, index_type, ptr::null());
    }

    pub unsafe fn draw_n_elements<I>(&mut self,
                                     _ab: &ArrayBufferBinding,
                                     mode: DrawingMode,
                                     count: usize,
                                     indicies: &[I])
        where I: IndexDatum, [I]: IndexData
    {
        debug_assert!(count <= indicies.len());

        let ptr = indicies.index_bytes().as_ptr();
        let index_type = I::index_datum_type();
        _draw_elements(mode, count, index_type, mem::transmute(ptr));
    }

    pub unsafe fn draw_elements<I>(&mut self,
                                   _ab: &ArrayBufferBinding,
                                   mode: DrawingMode,
                                   indicies: &[I])
        where I: IndexDatum, [I]: IndexData
    {
        self.draw_n_elements(_ab, mode, indicies.len(), indicies);
    }
}

pub struct ElementArrayBufferBinding<'a> {
    phantom: PhantomData<&'a mut Buffer>
}

impl<'a> BufferBinding for ElementArrayBufferBinding<'a> {
    fn target(&self) -> BufferBindingTarget {
        ELEMENT_ARRAY_BUFFER
    }
}

unsafe fn _bind_buffer(target: BufferBindingTarget, buffer: &mut Buffer) {
    gl::BindBuffer(target as GLuint, buffer.gl_id());
    dbg_gl_sanity_check! {
        GLError::InvalidEnum => "`target` is not an allowed value",
        _ => "Unknown error"
    }
}

pub struct ArrayBufferBinder;
impl ArrayBufferBinder {
    pub fn bind<'a>(&'a mut self, buffer: &mut Buffer)
        -> ArrayBufferBinding<'a>
    {
        let binding = ArrayBufferBinding { phantom: PhantomData };
        unsafe {
            _bind_buffer(binding.target(), buffer);
        }
        binding
    }
}

pub struct ElementArrayBufferBinder;
impl ElementArrayBufferBinder {
    pub fn bind<'a>(&'a mut self, buffer: &mut Buffer)
        -> ElementArrayBufferBinding<'a>
    {
        let binding = ElementArrayBufferBinding { phantom: PhantomData };
        unsafe {
            _bind_buffer(binding.target(), buffer);
        }
        binding
    }
}
<commit_msg>Mark `buffer::_bind_buffer()` as safe<commit_after>use std::mem;
use std::ptr;
use std::marker::PhantomData;
use gl;
use gl::types::*;
use context::Context;
use framebuffer::FramebufferBinding;
use types::{DrawingMode, GLError};
use index_data::{IndexData, IndexDatum, IndexDatumType};

pub struct Buffer {
    gl_id: GLuint
}

impl Buffer {
    pub fn gl_id(&self) -> GLuint {
        self.gl_id
    }
}

impl Drop for Buffer {
    fn drop(&mut self) {
        unsafe {
            gl::DeleteBuffers(1, &self.gl_id as *const GLuint);
        }
    }
}

impl Context {
    pub fn gen_buffer(&self) -> Buffer {
        unsafe {
            let mut id : GLuint = 0;

            gl::GenBuffers(1, &mut id as *mut GLuint);
            dbg_gl_sanity_check! {
                GLError::InvalidValue => "`n` is negative",
                _ => "Unknown error"
            }

            Buffer { gl_id: id }
        }
    }
}

gl_enum! {
    pub gl_enum BufferDataUsage {
        StreamDraw as STREAM_DRAW = gl::STREAM_DRAW,
        StaticDraw as STATIC_DRAW = gl::STATIC_DRAW,
        DynamicDraw as DYNAMIC_DRAW = gl::DYNAMIC_DRAW
    }
}

gl_enum! {
    pub gl_enum BufferBindingTarget {
        ArrayBuffer as ARRAY_BUFFER = gl::ARRAY_BUFFER,
        ElementArrayBuffer as ELEMENT_ARRAY_BUFFER = gl::ELEMENT_ARRAY_BUFFER
    }
}

pub trait BufferBinding {
    fn target(&self) -> BufferBindingTarget;

    fn buffer_bytes(&mut self, bytes: &[u8], usage: BufferDataUsage) {
        let ptr = bytes.as_ptr();
        let size = bytes.len() * mem::size_of::<u8>();
        unsafe {
            gl::BufferData(self.target().gl_enum(),
                           size as GLsizeiptr,
                           ptr as *const GLvoid,
                           usage.gl_enum());
            dbg_gl_error! {
                GLError::InvalidEnum => "Invalid `target` or `usage`",
                GLError::InvalidValue => "`size` is negative",
                GLError::InvalidOperation => "Object 0 is bound to buffer target",
                GLError::OutOfMemory => "Unable to create a large enough buffer",
                _ => "Unknown error"
            }
        }
    }
}

pub struct ArrayBufferBinding<'a> {
    phantom: PhantomData<&'a mut Buffer>
}

impl<'a> BufferBinding for ArrayBufferBinding<'a> {
    fn target(&self) -> BufferBindingTarget {
        ARRAY_BUFFER
    }
}

unsafe fn _draw_elements(mode: DrawingMode,
                         count: usize,
                         index_type: IndexDatumType,
                         indicies: *const GLvoid)
{
    let gl_index_type: GLenum = match index_type {
        IndexDatumType::UnsignedByte => gl::UNSIGNED_BYTE,
        IndexDatumType::UnsignedShort => gl::UNSIGNED_SHORT
    };
    gl::DrawElements(mode.gl_enum(),
                     count as GLsizei,
                     gl_index_type,
                     indicies);
    dbg_gl_error! {
        GLError::InvalidEnum => "`mode` or `type` is not an accepted value",
        GLError::InvalidValue => "`count` is negative",
        GLError::InvalidFramebufferOperation => "The current framebuffer is not framebuffer-complete",
        _ => "Unknown error"
    }
}

impl<'a> ArrayBufferBinding<'a> {
    pub unsafe fn vertex_attrib_pointer(&self,
                                        attrib: super::ProgramAttrib,
                                        components: i8,
                                        gl_type: super::DataType,
                                        normalized: bool,
                                        stride: usize,
                                        offset: usize)
    {
        debug_assert!(1 <= components && components <= 4);

        let gl_normalized = if normalized { gl::TRUE } else { gl::FALSE };
        gl::VertexAttribPointer(attrib.gl_index,
                                components as GLint,
                                gl_type.gl_enum(),
                                gl_normalized,
                                stride as GLsizei,
                                offset as *const GLvoid);
        dbg_gl_error! {
            GLError::InvalidEnum => "Illegal vertex attribute type",
            GLError::InvalidValue => "`stride` is negative, `size` is not in range, or `index` is >= GL_MAX_VERTEX_ATTRIBS",
            GLError::InvalidFramebufferOperation => "Currently bound framebuffer is not framebuffer complete",
            _ => "Unknown error"
        }
    }
}

impl<'a> FramebufferBinding<'a> {
    pub unsafe fn draw_arrays_range(&mut self,
                                    _ab: &ArrayBufferBinding,
                                    mode: DrawingMode,
                                    first: u32,
                                    count: usize)
    {
        gl::DrawArrays(mode.gl_enum(), first as GLint, count as GLsizei);
        dbg_gl_sanity_check! {
            GLError::InvalidEnum => "`mode` is not an accepted value",
            GLError::InvalidValue => "`count` is negative",
            _ => "Unknown error"
        }
    }

    pub unsafe fn draw_n_elements_buffered(&mut self,
                                           _ab: &ArrayBufferBinding,
                                           _eab: &ElementArrayBufferBinding,
                                           mode: DrawingMode,
                                           count: usize,
                                           index_type: IndexDatumType)
    {
        _draw_elements(mode, count, index_type, ptr::null());
    }

    pub unsafe fn draw_n_elements<I>(&mut self,
                                     _ab: &ArrayBufferBinding,
                                     mode: DrawingMode,
                                     count: usize,
                                     indicies: &[I])
        where I: IndexDatum, [I]: IndexData
    {
        debug_assert!(count <= indicies.len());

        let ptr = indicies.index_bytes().as_ptr();
        let index_type = I::index_datum_type();
        _draw_elements(mode, count, index_type, mem::transmute(ptr));
    }

    pub unsafe fn draw_elements<I>(&mut self,
                                   _ab: &ArrayBufferBinding,
                                   mode: DrawingMode,
                                   indicies: &[I])
        where I: IndexDatum, [I]: IndexData
    {
        self.draw_n_elements(_ab, mode, indicies.len(), indicies);
    }
}

pub struct ElementArrayBufferBinding<'a> {
    phantom: PhantomData<&'a mut Buffer>
}

impl<'a> BufferBinding for ElementArrayBufferBinding<'a> {
    fn target(&self) -> BufferBindingTarget {
        ELEMENT_ARRAY_BUFFER
    }
}

fn _bind_buffer(target: BufferBindingTarget, buffer: &mut Buffer) {
    unsafe {
        gl::BindBuffer(target as GLuint, buffer.gl_id());
        dbg_gl_sanity_check! {
            GLError::InvalidEnum => "`target` is not an allowed value",
            _ => "Unknown error"
        }
    }
}

pub struct ArrayBufferBinder;
impl ArrayBufferBinder {
    pub fn bind<'a>(&'a mut self, buffer: &mut Buffer)
        -> ArrayBufferBinding<'a>
    {
        let binding = ArrayBufferBinding { phantom: PhantomData };
        _bind_buffer(binding.target(), buffer);
        binding
    }
}

pub struct ElementArrayBufferBinder;
impl ElementArrayBufferBinder {
    pub fn bind<'a>(&'a mut self, buffer: &mut Buffer)
        -> ElementArrayBufferBinding<'a>
    {
        let binding = ElementArrayBufferBinding { phantom: PhantomData };
        _bind_buffer(binding.target(), buffer);
        binding
    }
}
<|endoftext|>


<commit_before>//!
//!	Code to write a sequence table to disk.
//!
extern crate bincode;

use std::io::prelude::*;
use std::io::SeekFrom;
use std::io::Result;
use std::io::Error;
use byteorder::{ByteOrder, LittleEndian, WriteBytesExt};
use std::mem::size_of;
use bincode::rustc_serialize::{encode_into, encode};
use flate2::{Compression,Compress,Flush};
use std::cmp::max;

use ::seqtable::SeqTableParams;
use ::seqtable::SeqStore;

#[derive(Debug, RustcEncodable, RustcDecodable, PartialEq)]
pub struct SeqBlock {
    pub enc_size: u64,
    pub comp_size: u64,
    pub offset: u64,
}

#[derive(Debug, RustcEncodable, RustcDecodable, PartialEq)]
pub struct SeqInfo {
    pub name: String,
    pub length: u64,
    pub blocks: Vec<SeqBlock>,
}

pub struct SeqTableWriter<W: Write + Seek> {
    tailoffset: u64,
    writer: W,
    infotable: Vec<SeqInfo>,
    block_length: u64,
    max_buffer_size: u64,
}

// TODO: add UnMap instance
impl<W: Write + Seek> SeqTableWriter<W> {
    pub fn new(mut writer: W, params: SeqTableParams) -> Result<SeqTableWriter<W>> {
        let blen = 1024u64;

        // write parameters to file
        try!(writer.write_u8(params.cut_length));
        // offsets are not required, but are stored anyway for reference
        try!(writer.write_u8(params.plus_offset));
        try!(writer.write_u8(params.minus_offset));
        try!(writer.write_u16::<LittleEndian>(params.read_length));
        try!(writer.write_u64::<LittleEndian>(blen));

        // write temporary blank value to be filled in later
        // with offset of sequence table at end of file
        try!(writer.write_u64::<LittleEndian>(0));
        // same of max decoder size
        try!(writer.write_u64::<LittleEndian>(0));

        //
        let offset = 3 * size_of::<u64>() + 4 * size_of::<u8>();
        Ok(SeqTableWriter{
            tailoffset: offset as u64,
            writer: writer,
            infotable: Vec::new(),
            block_length: blen,
            max_buffer_size: 0,
        })
    }

    pub fn create_sequence<'a>(&'a mut self, name: String) -> SequenceWriter<'a, W> {
        // create entry in info-table
        // return new sequence writter at current offset

        self.infotable.push( SeqInfo { name: name.clone(), length: 0, blocks: Vec::new() } );
        SequenceWriter {
            offset: &mut self.tailoffset,
            writer: &mut self.writer,
            info: self.infotable.last_mut().unwrap(),
            block: Vec::new(),
            compressor: Compress::new(Compression::Best, false),
            output: vec![0u8; self.block_length as usize],
            max_buffer_size: &mut self.max_buffer_size,
        }
    }
}

impl<W: Write + Seek> Drop for SeqTableWriter<W> {
    fn drop(&mut self) {
        // at this point the writer is at the end of the file
        // and the tailoffset holds what will be the position
        // of the info table

        // write info table
        encode_into(&self.infotable, &mut self.writer, bincode::SizeLimit::Infinite).unwrap();

        // seek to start & fill info table offset and max decoder size in header
        let offset = size_of::<u64>() + 4 * size_of::<u8>();
        self.writer.seek(SeekFrom::Start(offset as u64)).unwrap();
        self.writer.write_u64::<LittleEndian>(self.tailoffset as u64).unwrap();
        self.writer.write_u64::<LittleEndian>(self.max_buffer_size).unwrap();
    }
}

pub struct SequenceWriter<'a, W: 'a + Write> {
    offset: &'a mut u64,
    writer: &'a mut W,
    info: &'a mut SeqInfo,
    block: Vec<(u16, u16)>,
    compressor: Compress,
    output: Vec<u8>,
    max_buffer_size: &'a mut u64,
}

impl<'a, W: 'a + Write> SequenceWriter<'a, W> {
    fn write_block(&mut self) -> Result<()> {
        for &(p, m) in &self.block {
            println!("write: ({},{})", p, m);
        }

        //   compress block and write to disk
        let binvec: Vec<u8> = encode(&self.block, bincode::SizeLimit::Infinite).unwrap();
        *self.max_buffer_size = max(*self.max_buffer_size, binvec.len() as u64);
        let mut total = self.compressor.total_out();
        self.compressor.compress(&binvec, &mut self.output, Flush::Finish);
        total = self.compressor.total_out() - total;

        try!(self.writer.write(&self.output[0..(total as usize)]));

        //   add block to info-table && length
        self.info.blocks.push(SeqBlock { enc_size: binvec.len() as u64, comp_size: total, offset: *self.offset });
        *self.offset = *self.offset + total;

        self.info.length = self.info.length + self.block.len() as u64;
        self.block.clear();

        Ok(())
    }
}

impl<'a, W: 'a + Write> SeqStore for SequenceWriter<'a, W> {
    fn write(&mut self, plus: u16, minus: u16) {
        self.block.push((plus, minus));

        if self.block.len() == 1024 {
                self.write_block().ok().expect("Failed to write block");
        }
    }
}

impl<'a, W: 'a + Write> Drop for SequenceWriter<'a, W> {
    fn drop(&mut self) {
        // write last block
        if self.block.len() > 0 {
            self.write_block().ok().expect("Failed to write block");
        }
    }
}
<commit_msg>Fix output buffer size & reset compressor<commit_after>//!
//!	Code to write a sequence table to disk.
//!
extern crate bincode;

use std::io::prelude::*;
use std::io::SeekFrom;
use std::io::Result;
use std::io::Error;
use byteorder::{ByteOrder, LittleEndian, WriteBytesExt};
use std::mem::size_of;
use bincode::rustc_serialize::{encode_into, encode};
use flate2::{Compression,Compress,Flush};
use std::cmp::max;

use ::seqtable::SeqTableParams;
use ::seqtable::SeqStore;

#[derive(Debug, RustcEncodable, RustcDecodable, PartialEq)]
pub struct SeqBlock {
    pub enc_size: u64,
    pub comp_size: u64,
    pub offset: u64,
}

#[derive(Debug, RustcEncodable, RustcDecodable, PartialEq)]
pub struct SeqInfo {
    pub name: String,
    pub length: u64,
    pub blocks: Vec<SeqBlock>,
}

pub struct SeqTableWriter<W: Write + Seek> {
    tailoffset: u64,
    writer: W,
    infotable: Vec<SeqInfo>,
    block_length: u64,
    max_buffer_size: u64,
}

// TODO: add UnMap instance
impl<W: Write + Seek> SeqTableWriter<W> {
    pub fn new(mut writer: W, params: SeqTableParams) -> Result<SeqTableWriter<W>> {
        let blen = 1024u64;

        // write parameters to file
        try!(writer.write_u8(params.cut_length));
        // offsets are not required, but are stored anyway for reference
        try!(writer.write_u8(params.plus_offset));
        try!(writer.write_u8(params.minus_offset));
        try!(writer.write_u16::<LittleEndian>(params.read_length));
        try!(writer.write_u64::<LittleEndian>(blen));

        // write temporary blank value to be filled in later
        // with offset of sequence table at end of file
        try!(writer.write_u64::<LittleEndian>(0));
        // same of max decoder size
        try!(writer.write_u64::<LittleEndian>(0));

        //
        let offset = 3 * size_of::<u64>() + 3 * size_of::<u8>() + size_of::<u16>();
        Ok(SeqTableWriter{
            tailoffset: offset as u64,
            writer: writer,
            infotable: Vec::new(),
            block_length: blen,
            max_buffer_size: 0,
        })
    }

    pub fn create_sequence<'a>(&'a mut self, name: String) -> SequenceWriter<'a, W> {
        // create entry in info-table
        // return new sequence writter at current offset

        self.infotable.push( SeqInfo { name: name.clone(), length: 0, blocks: Vec::new() } );
        SequenceWriter {
            offset: &mut self.tailoffset,
            writer: &mut self.writer,
            info: self.infotable.last_mut().unwrap(),
            block: Vec::new(),
            compressor: Compress::new(Compression::Best, false),
            output: vec![0u8; self.block_length as usize * size_of::<(u16,u16)>()],
            block_length: self.block_length,
            max_buffer_size: &mut self.max_buffer_size,
        }
    }
}

impl<W: Write + Seek> Drop for SeqTableWriter<W> {
    fn drop(&mut self) {
        // at this point the writer is at the end of the file
        // and the tailoffset holds what will be the position
        // of the info table

        // write info table
        encode_into(&self.infotable, &mut self.writer, bincode::SizeLimit::Infinite).unwrap();

        // seek to start & fill info table offset and max decoder size in header
        let offset = size_of::<u64>() + 3 * size_of::<u8>() + size_of::<u16>();
        self.writer.seek(SeekFrom::Start(offset as u64)).unwrap();
        self.writer.write_u64::<LittleEndian>(self.tailoffset as u64).unwrap();
        self.writer.write_u64::<LittleEndian>(self.max_buffer_size).unwrap();
    }
}

pub struct SequenceWriter<'a, W: 'a + Write> {
    offset: &'a mut u64,
    writer: &'a mut W,
    info: &'a mut SeqInfo,
    block: Vec<(u16, u16)>,
    compressor: Compress,
    output: Vec<u8>,
    block_length: u64,
    max_buffer_size: &'a mut u64,
}

impl<'a, W: 'a + Write> SequenceWriter<'a, W> {
    fn write_block(&mut self) -> Result<()> {
        /*for &(p, m) in &self.block {
            println!("write: ({},{})", p, m);
        }*/

        //   compress block and write to disk
        let binvec: Vec<u8> = encode(&self.block, bincode::SizeLimit::Infinite).unwrap();
        *self.max_buffer_size = max(*self.max_buffer_size, binvec.len() as u64);

        self.compressor.reset();
        let mut total = self.compressor.total_out();
        self.compressor.compress(&binvec, &mut self.output, Flush::Finish);
        total = self.compressor.total_out() - total;

        try!(self.writer.write(&self.output[0..(total as usize)]));

        //   add block to info-table && length
        self.info.blocks.push(SeqBlock { enc_size: binvec.len() as u64, comp_size: total, offset: *self.offset });
        *self.offset = *self.offset + total;

        self.info.length = self.info.length + self.block.len() as u64;
        self.block.clear();

        Ok(())
    }
}

impl<'a, W: 'a + Write> SeqStore for SequenceWriter<'a, W> {
    fn write(&mut self, plus: u16, minus: u16) {
        self.block.push((plus, minus));

        if self.block.len() == self.block_length as usize {
                self.write_block().ok().expect("Failed to write block");
        }
    }
}

impl<'a, W: 'a + Write> Drop for SequenceWriter<'a, W> {
    fn drop(&mut self) {
        // write last block
        if self.block.len() > 0 {
            self.write_block().ok().expect("Failed to write block");
        }
    }
}
<|endoftext|>


<commit_before>use std::ops::{Add, Sub};
use typenum::uint::{Unsigned, UInt};
use typenum::bit::Bit;
use typenum::consts::{U0, U1};

/// This enum serves to represent the type of a tensor. A tensor can have any number of indices, and each one can be either
/// covariant (a lower index), or contravariant (an upper index). For example, a vector is a tensor with only one contravariant index.
#[derive(Clone, Copy, PartialEq, Debug)]
pub enum IndexType {
    Covariant,
    Contravariant,
}

/// Trait identifying a type as representing a tensor variance. It is implemented for
/// `CovariantIndex`, `ContravariantIndex` and tuples (Index, Variance).
pub trait Variance {
    type Rank: Unsigned + Add<U1>;
    fn rank() -> usize {
        Self::Rank::to_usize()
    }
    fn variance() -> Vec<IndexType>;
}

impl Variance for () {
    type Rank = U0;

    fn variance() -> Vec<IndexType> {
        vec![]
    }
}

/// Trait identifying a type as representing a tensor index. It is implemented
/// for `CovariantIndex` and `ContravariantIndex`.
pub trait TensorIndex: Variance {
    fn index_type() -> IndexType;
}

/// Type representing a contravariant (upper) tensor index.
pub struct ContravariantIndex;
impl TensorIndex for ContravariantIndex {
    fn index_type() -> IndexType {
        IndexType::Contravariant
    }
}

/// Type representing a covariant (lower) tensor index.
pub struct CovariantIndex;
impl TensorIndex for CovariantIndex {
    fn index_type() -> IndexType {
        IndexType::Covariant
    }
}

impl Variance for ContravariantIndex {
    type Rank = U1;
    fn variance() -> Vec<IndexType> {
        vec![IndexType::Contravariant]
    }
}

impl Variance for CovariantIndex {
    type Rank = U1;
    fn variance() -> Vec<IndexType> {
        vec![IndexType::Covariant]
    }
}

impl<T, U> Variance for (T, U)
    where U: Variance,
          <U::Rank as Add<U1>>::Output: Unsigned + Add<U1>,
          T: TensorIndex
{
    type Rank = <U::Rank as Add<U1>>::Output;

    fn variance() -> Vec<IndexType> {
        let mut result = vec![T::index_type()];
        result.append(&mut U::variance());
        result
    }
}

/// Operator trait used for concatenating two variances.
pub trait Concat<T> {
    type Output;
}

impl<T, U> Concat<U> for T
    where T: TensorIndex,
          U: TensorIndex
{
    type Output = (T, U);
}

impl<T, U, V> Concat<V> for (T, U)
    where T: TensorIndex,
          V: TensorIndex,
          U: Variance + Concat<V>,
          <U as Concat<V>>::Output: Variance
{
    type Output = (T, <U as Concat<V>>::Output);
}

impl<T, U, V> Concat<(U, V)> for T
    where T: TensorIndex,
          U: TensorIndex,
          V: Variance
{
    type Output = (T, (U, V));
}

impl<T, U, V, W> Concat<(V, W)> for (T, U)
    where T: TensorIndex,
          U: Variance + Concat<(V, W)>,
          V: TensorIndex,
          W: Variance
{
    type Output = (T, <U as Concat<(V, W)>>::Output);
}

/// Indexing operator trait: Output is equal to the index type at the given position
///
/// Warning: Indices are numbered starting from 0!
pub trait Index<T: Unsigned> {
    type Output;
}

impl Index<U0> for CovariantIndex {
    type Output = CovariantIndex;
}

impl Index<U0> for ContravariantIndex {
    type Output = ContravariantIndex;
}

impl<T, V, U, B> Index<UInt<U, B>> for (V, T)
    where V: TensorIndex,
          U: Unsigned,
          B: Bit,
          UInt<U, B>: Sub<U1>,
          <UInt<U, B> as Sub<U1>>::Output: Unsigned,
          T: Variance + Index<<UInt<U, B> as Sub<U1>>::Output>
{
    type Output = <T as Index<<UInt<U, B> as Sub<U1>>::Output>>::Output;
}

impl<T, V> Index<U0> for (V, T)
    where V: TensorIndex,
          T: Variance
{
    type Output = V;
}

/// An operator trait, removing the indicated index from a variance
pub trait RemoveIndex<T: Unsigned> {
    type Output;
}

impl RemoveIndex<U0> for CovariantIndex {
    type Output = ();
}

impl RemoveIndex<U0> for ContravariantIndex {
    type Output = ();
}

impl<U, V> RemoveIndex<U0> for (U, V)
    where U: TensorIndex,
          V: Variance
{
    type Output = V;
}

impl<T, B, U, V> RemoveIndex<UInt<T, B>> for (U, V)
    where T: Unsigned,
          B: Bit,
          U: TensorIndex,
          UInt<T, B>: Sub<U1>,
          <UInt<T, B> as Sub<U1>>::Output: Unsigned,
          V: Variance + RemoveIndex<<UInt<T, B> as Sub<U1>>::Output>
{
    type Output = (U, <V as RemoveIndex<<UInt<T, B> as Sub<U1>>::Output>>::Output);
}

#[cfg(test)]
mod test {
    use super::*;
    use typenum::consts::{U0, U1, U2};

    #[test]
    fn test_variance() {
        assert_eq!(<(CovariantIndex, ContravariantIndex) as Variance>::variance(), vec![IndexType::Covariant, IndexType::Contravariant]);
    }

    #[test]
    fn test_variance_concat() {
        assert_eq!(<<CovariantIndex as Concat<ContravariantIndex>>::Output as Variance>::variance(),
            vec![IndexType::Covariant, IndexType::Contravariant]);

        assert_eq!(<<(CovariantIndex, CovariantIndex) as Concat<ContravariantIndex>>::Output as Variance>::variance(),
            vec![IndexType::Covariant, IndexType::Covariant, IndexType::Contravariant]);

        assert_eq!(<<CovariantIndex as Concat<(CovariantIndex, ContravariantIndex)>>::Output as Variance>::variance(),
            vec![IndexType::Covariant, IndexType::Covariant, IndexType::Contravariant]);

        assert_eq!(<<(ContravariantIndex, CovariantIndex) as Concat<(CovariantIndex, ContravariantIndex)>>::Output as Variance>::variance(),
            vec![IndexType::Contravariant, IndexType::Covariant, IndexType::Covariant, IndexType::Contravariant]);
    }

    #[test]
    fn test_index() {
        assert_eq!(<CovariantIndex as Index<U0>>::Output::index_type(),
            IndexType::Covariant);

        assert_eq!(<(CovariantIndex, ContravariantIndex) as Index<U0>>::Output::index_type(),
            IndexType::Covariant);

        assert_eq!(<(CovariantIndex, ContravariantIndex) as Index<U1>>::Output::index_type(),
            IndexType::Contravariant);

        assert_eq!(<(ContravariantIndex, (CovariantIndex, CovariantIndex)) as Index<U0>>::Output::index_type(),
            IndexType::Contravariant);

        assert_eq!(<(ContravariantIndex, (CovariantIndex, CovariantIndex)) as Index<U2>>::Output::index_type(),
            IndexType::Covariant);
    }

    #[test]
    fn test_remove() {
        assert_eq!(<CovariantIndex as RemoveIndex<U0>>::Output::variance(),
            vec![]);

        assert_eq!(<(CovariantIndex, ContravariantIndex) as RemoveIndex<U0>>::Output::variance(),
            vec![IndexType::Contravariant]);

        assert_eq!(<(CovariantIndex, ContravariantIndex) as RemoveIndex<U1>>::Output::variance(),
            vec![IndexType::Covariant]);

        assert_eq!(<(ContravariantIndex, (CovariantIndex, CovariantIndex)) as RemoveIndex<U1>>::Output::variance(),
            vec![IndexType::Contravariant, IndexType::Covariant]);
    }
}<commit_msg>Added tensor contraction<commit_after>use std::ops::{Add, Sub};
use typenum::uint::{Unsigned, UInt};
use typenum::bit::Bit;
use typenum::consts::{U0, U1};
use typenum::{Cmp, Same, Greater};

/// This enum serves to represent the type of a tensor. A tensor can have any number of indices, and each one can be either
/// covariant (a lower index), or contravariant (an upper index). For example, a vector is a tensor with only one contravariant index.
#[derive(Clone, Copy, PartialEq, Debug)]
pub enum IndexType {
    Covariant,
    Contravariant,
}

/// Trait identifying a type as representing a tensor variance. It is implemented for
/// `CovariantIndex`, `ContravariantIndex` and tuples (Index, Variance).
pub trait Variance {
    type Rank: Unsigned + Add<U1>;
    fn rank() -> usize {
        Self::Rank::to_usize()
    }
    fn variance() -> Vec<IndexType>;
}

impl Variance for () {
    type Rank = U0;

    fn variance() -> Vec<IndexType> {
        vec![]
    }
}

/// Trait identifying a type as representing a tensor index. It is implemented
/// for `CovariantIndex` and `ContravariantIndex`.
pub trait TensorIndex: Variance {
    fn index_type() -> IndexType;
}

/// Type representing a contravariant (upper) tensor index.
pub struct ContravariantIndex;
impl TensorIndex for ContravariantIndex {
    fn index_type() -> IndexType {
        IndexType::Contravariant
    }
}

/// Type representing a covariant (lower) tensor index.
pub struct CovariantIndex;
impl TensorIndex for CovariantIndex {
    fn index_type() -> IndexType {
        IndexType::Covariant
    }
}

impl Variance for ContravariantIndex {
    type Rank = U1;
    fn variance() -> Vec<IndexType> {
        vec![IndexType::Contravariant]
    }
}

impl Variance for CovariantIndex {
    type Rank = U1;
    fn variance() -> Vec<IndexType> {
        vec![IndexType::Covariant]
    }
}

/// Trait representing the other index type
pub trait OtherIndex: TensorIndex {
    type Output: TensorIndex;
}

impl OtherIndex for CovariantIndex {
    type Output = ContravariantIndex;
}

impl OtherIndex for ContravariantIndex {
    type Output = CovariantIndex;
}

// Back to implementing Variance

impl<T, U> Variance for (T, U)
    where U: Variance,
          <U::Rank as Add<U1>>::Output: Unsigned + Add<U1>,
          T: TensorIndex
{
    type Rank = <U::Rank as Add<U1>>::Output;

    fn variance() -> Vec<IndexType> {
        let mut result = vec![T::index_type()];
        result.append(&mut U::variance());
        result
    }
}

/// Operator trait used for concatenating two variances.
pub trait Concat<T> {
    type Output;
}

impl<T, U> Concat<U> for T
    where T: TensorIndex,
          U: TensorIndex
{
    type Output = (T, U);
}

impl<T, U, V> Concat<V> for (T, U)
    where T: TensorIndex,
          V: TensorIndex,
          U: Variance + Concat<V>,
          <U as Concat<V>>::Output: Variance
{
    type Output = (T, <U as Concat<V>>::Output);
}

impl<T, U, V> Concat<(U, V)> for T
    where T: TensorIndex,
          U: TensorIndex,
          V: Variance
{
    type Output = (T, (U, V));
}

impl<T, U, V, W> Concat<(V, W)> for (T, U)
    where T: TensorIndex,
          U: Variance + Concat<(V, W)>,
          V: TensorIndex,
          W: Variance
{
    type Output = (T, <U as Concat<(V, W)>>::Output);
}

/// Indexing operator trait: Output is equal to the index type at the given position
///
/// Warning: Indices are numbered starting from 0!
pub trait Index<T: Unsigned> {
    type Output;
}

impl Index<U0> for CovariantIndex {
    type Output = CovariantIndex;
}

impl Index<U0> for ContravariantIndex {
    type Output = ContravariantIndex;
}

impl<T, V, U, B> Index<UInt<U, B>> for (V, T)
    where V: TensorIndex,
          U: Unsigned,
          B: Bit,
          UInt<U, B>: Sub<U1>,
          <UInt<U, B> as Sub<U1>>::Output: Unsigned,
          T: Variance + Index<<UInt<U, B> as Sub<U1>>::Output>
{
    type Output = <T as Index<<UInt<U, B> as Sub<U1>>::Output>>::Output;
}

impl<T, V> Index<U0> for (V, T)
    where V: TensorIndex,
          T: Variance
{
    type Output = V;
}

/// An operator trait, removing the indicated index from a variance
pub trait RemoveIndex<T: Unsigned> {
    type Output;
}

impl RemoveIndex<U0> for CovariantIndex {
    type Output = ();
}

impl RemoveIndex<U0> for ContravariantIndex {
    type Output = ();
}

impl<U, V> RemoveIndex<U0> for (U, V)
    where U: TensorIndex,
          V: Variance
{
    type Output = V;
}

impl<T, B, U, V> RemoveIndex<UInt<T, B>> for (U, V)
    where T: Unsigned,
          B: Bit,
          U: TensorIndex,
          UInt<T, B>: Sub<U1>,
          <UInt<T, B> as Sub<U1>>::Output: Unsigned,
          V: Variance + RemoveIndex<<UInt<T, B> as Sub<U1>>::Output>
{
    type Output = (U, <V as RemoveIndex<<UInt<T, B> as Sub<U1>>::Output>>::Output);
}

/// An operator trait representing tensor contraction
pub trait Contract<Ul: Unsigned, Uh: Unsigned> {
    type Output;
}

// this is quite possibly the worst impl I have ever written
impl<Ul, Uh, V> Contract<Ul, Uh> for V
    where Ul: Unsigned,
          Uh: Unsigned + Sub<U1> + Cmp<Ul>,
          <Uh as Sub<U1>>::Output: Unsigned,
          <Uh as Cmp<Ul>>::Output: Same<Greater>,
          V: Index<Ul> + Index<Uh> + RemoveIndex<Ul>,
          <V as Index<Ul>>::Output: OtherIndex,
          <V as Index<Uh>>::Output: TensorIndex + Same<<<V as Index<Ul>>::Output as OtherIndex>::Output>,
          <V as RemoveIndex<Ul>>::Output: RemoveIndex<<Uh as Sub<U1>>::Output>,
          <<V as RemoveIndex<Ul>>::Output as RemoveIndex<<Uh as Sub<U1>>::Output>>::Output: Variance
{
    type Output = <<V as RemoveIndex<Ul>>::Output as RemoveIndex<<Uh as Sub<U1>>::Output>>::Output;
}

#[cfg(test)]
mod test {
    use super::*;
    use typenum::consts::{U0, U1, U2};

    #[test]
    fn test_variance() {
        assert_eq!(<(CovariantIndex, ContravariantIndex) as Variance>::variance(), vec![IndexType::Covariant, IndexType::Contravariant]);
    }

    #[test]
    fn test_variance_concat() {
        assert_eq!(<<CovariantIndex as Concat<ContravariantIndex>>::Output as Variance>::variance(),
            vec![IndexType::Covariant, IndexType::Contravariant]);

        assert_eq!(<<(CovariantIndex, CovariantIndex) as Concat<ContravariantIndex>>::Output as Variance>::variance(),
            vec![IndexType::Covariant, IndexType::Covariant, IndexType::Contravariant]);

        assert_eq!(<<CovariantIndex as Concat<(CovariantIndex, ContravariantIndex)>>::Output as Variance>::variance(),
            vec![IndexType::Covariant, IndexType::Covariant, IndexType::Contravariant]);

        assert_eq!(<<(ContravariantIndex, CovariantIndex) as Concat<(CovariantIndex, ContravariantIndex)>>::Output as Variance>::variance(),
            vec![IndexType::Contravariant, IndexType::Covariant, IndexType::Covariant, IndexType::Contravariant]);
    }

    #[test]
    fn test_index() {
        assert_eq!(<CovariantIndex as Index<U0>>::Output::index_type(),
            IndexType::Covariant);

        assert_eq!(<(CovariantIndex, ContravariantIndex) as Index<U0>>::Output::index_type(),
            IndexType::Covariant);

        assert_eq!(<(CovariantIndex, ContravariantIndex) as Index<U1>>::Output::index_type(),
            IndexType::Contravariant);

        assert_eq!(<(ContravariantIndex, (CovariantIndex, CovariantIndex)) as Index<U0>>::Output::index_type(),
            IndexType::Contravariant);

        assert_eq!(<(ContravariantIndex, (CovariantIndex, CovariantIndex)) as Index<U2>>::Output::index_type(),
            IndexType::Covariant);
    }

    #[test]
    fn test_remove() {
        assert_eq!(<CovariantIndex as RemoveIndex<U0>>::Output::variance(),
            vec![]);

        assert_eq!(<(CovariantIndex, ContravariantIndex) as RemoveIndex<U0>>::Output::variance(),
            vec![IndexType::Contravariant]);

        assert_eq!(<(CovariantIndex, ContravariantIndex) as RemoveIndex<U1>>::Output::variance(),
            vec![IndexType::Covariant]);

        assert_eq!(<(ContravariantIndex, (CovariantIndex, CovariantIndex)) as RemoveIndex<U1>>::Output::variance(),
            vec![IndexType::Contravariant, IndexType::Covariant]);
    }

    #[test]
    fn test_contract() {
        assert_eq!(<(CovariantIndex, ContravariantIndex) as Contract<U0, U1>>::Output::variance(),
            vec![]);

        assert_eq!(<(ContravariantIndex, CovariantIndex) as Contract<U0, U1>>::Output::variance(),
            vec![]);

        assert_eq!(<(ContravariantIndex, (CovariantIndex, CovariantIndex)) as Contract<U0, U1>>::Output::variance(),
            vec![IndexType::Covariant]);

        assert_eq!(<(ContravariantIndex, (CovariantIndex, CovariantIndex)) as Contract<U0, U2>>::Output::variance(),
            vec![IndexType::Covariant]);
    }
}<|endoftext|>