<commit_before>use std::iter::once;

use hir::Semantics;
use ide_db::{base_db::FileRange, RootDatabase};
use syntax::{algo, AstNode, NodeOrToken, SyntaxElement, SyntaxKind, SyntaxNode};
use text_edit::{TextEdit, TextEditBuilder};

pub enum Direction {
    Up,
    Down,
}

// Feature: Move Item
//
// Move item under cursor or selection up and down.
//
// |===
// | Editor  | Action Name
//
// | VS Code | **Rust Analyzer: Move item up**
// | VS Code | **Rust Analyzer: Move item down**
// |===
pub(crate) fn move_item(
    db: &RootDatabase,
    range: FileRange,
    direction: Direction,
) -> Option<TextEdit> {
    let sema = Semantics::new(db);
    let file = sema.parse(range.file_id);

    let item = file.syntax().covering_element(range.range);
    find_ancestors(item, direction)
}

fn find_ancestors(item: SyntaxElement, direction: Direction) -> Option<TextEdit> {
    let root = match item {
        NodeOrToken::Node(node) => node,
        NodeOrToken::Token(token) => token.parent()?,
    };

    let ancestor = once(root.clone())
        .chain(root.ancestors())
        .filter_map(|ancestor| kind_priority(ancestor.kind()).map(|priority| (priority, ancestor)))
        .max_by_key(|(priority, _)| *priority)
        .map(|(_, ancestor)| ancestor)?;

    move_in_direction(&ancestor, direction)
}

fn kind_priority(kind: SyntaxKind) -> Option<i32> {
    match kind {
        SyntaxKind::MATCH_ARM => Some(4),

        SyntaxKind::LET_STMT | SyntaxKind::EXPR_STMT | SyntaxKind::MATCH_EXPR => Some(3),

        SyntaxKind::TRAIT
        | SyntaxKind::IMPL
        | SyntaxKind::MACRO_CALL
        | SyntaxKind::MACRO_DEF
        | SyntaxKind::STRUCT
        | SyntaxKind::ENUM
        | SyntaxKind::MODULE
        | SyntaxKind::USE
        | SyntaxKind::FN
        | SyntaxKind::CONST
        | SyntaxKind::TYPE_ALIAS => Some(2),

        _ => None,
    }
}

fn move_in_direction(node: &SyntaxNode, direction: Direction) -> Option<TextEdit> {
    let sibling = match direction {
        Direction::Up => node.prev_sibling(),
        Direction::Down => node.next_sibling(),
    }?;

    Some(replace_nodes(&sibling, node))
}

fn replace_nodes(first: &SyntaxNode, second: &SyntaxNode) -> TextEdit {
    let mut edit = TextEditBuilder::default();

    algo::diff(first, second).into_text_edit(&mut edit);
    algo::diff(second, first).into_text_edit(&mut edit);

    edit.finish()
}

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

    use crate::Direction;

    fn check(ra_fixture: &str, expect: Expect, direction: Direction) {
        let (analysis, range) = fixture::range(ra_fixture);
        let edit = analysis.move_item(range, direction).unwrap().unwrap_or_default();
        let mut file = analysis.file_text(range.file_id).unwrap().to_string();
        edit.apply(&mut file);
        expect.assert_eq(&file);
    }

    #[test]
    fn test_moves_match_arm_up() {
        check(
            r#"
fn main() {
    match true {
        true => {
            println!("Hello, world");
        },
        false =>$0$0 {
            println!("Test");
        }
    };
}
            "#,
            expect![[r#"
fn main() {
    match true {
        false => {
            println!("Test");
        },
        true => {
            println!("Hello, world");
        }
    };
}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn test_moves_match_arm_down() {
        check(
            r#"
fn main() {
    match true {
        true =>$0$0 {
            println!("Hello, world");
        },
        false => {
            println!("Test");
        }
    };
}
            "#,
            expect![[r#"
fn main() {
    match true {
        false => {
            println!("Test");
        },
        true => {
            println!("Hello, world");
        }
    };
}
            "#]],
            Direction::Down,
        );
    }

    #[test]
    fn test_nowhere_to_move() {
        check(
            r#"
fn main() {
    match true {
        true =>$0$0 {
            println!("Hello, world");
        },
        false => {
            println!("Test");
        }
    };
}
            "#,
            expect![[r#"
fn main() {
    match true {
        true => {
            println!("Hello, world");
        },
        false => {
            println!("Test");
        }
    };
}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn test_moves_let_stmt_up() {
        check(
            r#"
fn main() {
    let test = 123;
    let test2$0$0 = 456;
}
            "#,
            expect![[r#"
fn main() {
    let test2 = 456;
    let test = 123;
}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn test_prioritizes_match_arm() {
        check(
            r#"
fn main() {
    match true {
        true => {
            let test = 123;$0$0
            let test2 = 456;
        },
        false => {
            println!("Test");
        }
    };
}
            "#,
            expect![[r#"
fn main() {
    match true {
        false => {
            println!("Test");
        },
        true => {
            let test = 123;
            let test2 = 456;
        }
    };
}
            "#]],
            Direction::Down,
        );
    }

    #[test]
    fn test_moves_expr_up() {
        check(
            r#"
fn main() {
    println!("Hello, world");
    println!("All I want to say is...");$0$0
}
            "#,
            expect![[r#"
fn main() {
    println!("All I want to say is...");
    println!("Hello, world");
}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn test_nowhere_to_move_stmt() {
        check(
            r#"
fn main() {
    println!("All I want to say is...");$0$0
    println!("Hello, world");
}
            "#,
            expect![[r#"
fn main() {
    println!("All I want to say is...");
    println!("Hello, world");
}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn test_move_item() {
        check(
            r#"
fn main() {}

fn foo() {}$0$0
            "#,
            expect![[r#"
fn foo() {}

fn main() {}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn test_move_impl_up() {
        check(
            r#"
struct Yay;

trait Wow {}

impl Wow for Yay {}$0$0
            "#,
            expect![[r#"
struct Yay;

impl Wow for Yay {}

trait Wow {}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn test_move_use_up() {
        check(
            r#"
use std::vec::Vec;
use std::collections::HashMap$0$0;
            "#,
            expect![[r#"
use std::collections::HashMap;
use std::vec::Vec;
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn moves_match_expr_up() {
        check(
            r#"
fn main() {
    let test = 123;

    $0match test {
        456 => {},
        _ => {}
    }$0;
}
            "#,
            expect![[r#"
fn main() {
    match test {
        456 => {},
        _ => {}
    };

    let test = 123;
}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn handles_empty_file() {
        check(r#"$0$0"#, expect![[r#""#]], Direction::Up);
    }
}
<commit_msg>Remove prioritization, add more movable items<commit_after>use std::iter::once;

use hir::Semantics;
use ide_db::{base_db::FileRange, RootDatabase};
use syntax::{algo, AstNode, NodeOrToken, SyntaxElement, SyntaxKind, SyntaxNode};
use text_edit::{TextEdit, TextEditBuilder};

pub enum Direction {
    Up,
    Down,
}

// Feature: Move Item
//
// Move item under cursor or selection up and down.
//
// |===
// | Editor  | Action Name
//
// | VS Code | **Rust Analyzer: Move item up**
// | VS Code | **Rust Analyzer: Move item down**
// |===
pub(crate) fn move_item(
    db: &RootDatabase,
    range: FileRange,
    direction: Direction,
) -> Option<TextEdit> {
    let sema = Semantics::new(db);
    let file = sema.parse(range.file_id);

    let item = file.syntax().covering_element(range.range);
    find_ancestors(item, direction)
}

fn find_ancestors(item: SyntaxElement, direction: Direction) -> Option<TextEdit> {
    let root = match item {
        NodeOrToken::Node(node) => node,
        NodeOrToken::Token(token) => token.parent()?,
    };

    let movable = [
        SyntaxKind::MATCH_ARM,
        SyntaxKind::PARAM,
        SyntaxKind::LET_STMT,
        SyntaxKind::EXPR_STMT,
        SyntaxKind::MATCH_EXPR,
        SyntaxKind::MACRO_CALL,
        SyntaxKind::TYPE_ALIAS,
        SyntaxKind::TRAIT,
        SyntaxKind::IMPL,
        SyntaxKind::MACRO_DEF,
        SyntaxKind::STRUCT,
        SyntaxKind::UNION,
        SyntaxKind::ENUM,
        SyntaxKind::FN,
        SyntaxKind::MODULE,
        SyntaxKind::USE,
        SyntaxKind::STATIC,
        SyntaxKind::CONST,
        SyntaxKind::MACRO_RULES,
    ];

    let ancestor = once(root.clone())
        .chain(root.ancestors())
        .find(|ancestor| movable.contains(&ancestor.kind()))?;

    move_in_direction(&ancestor, direction)
}

fn move_in_direction(node: &SyntaxNode, direction: Direction) -> Option<TextEdit> {
    let sibling = match direction {
        Direction::Up => node.prev_sibling(),
        Direction::Down => node.next_sibling(),
    }?;

    Some(replace_nodes(node, &sibling))
}

fn replace_nodes(first: &SyntaxNode, second: &SyntaxNode) -> TextEdit {
    let mut edit = TextEditBuilder::default();

    algo::diff(first, second).into_text_edit(&mut edit);
    algo::diff(second, first).into_text_edit(&mut edit);

    edit.finish()
}

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

    use crate::Direction;

    fn check(ra_fixture: &str, expect: Expect, direction: Direction) {
        let (analysis, range) = fixture::range(ra_fixture);
        let edit = analysis.move_item(range, direction).unwrap().unwrap_or_default();
        let mut file = analysis.file_text(range.file_id).unwrap().to_string();
        edit.apply(&mut file);
        expect.assert_eq(&file);
    }

    #[test]
    fn test_moves_match_arm_up() {
        check(
            r#"
fn main() {
    match true {
        true => {
            println!("Hello, world");
        },
        false =>$0$0 {
            println!("Test");
        }
    };
}
            "#,
            expect![[r#"
fn main() {
    match true {
        false => {
            println!("Test");
        },
        true => {
            println!("Hello, world");
        }
    };
}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn test_moves_match_arm_down() {
        check(
            r#"
fn main() {
    match true {
        true =>$0$0 {
            println!("Hello, world");
        },
        false => {
            println!("Test");
        }
    };
}
            "#,
            expect![[r#"
fn main() {
    match true {
        false => {
            println!("Test");
        },
        true => {
            println!("Hello, world");
        }
    };
}
            "#]],
            Direction::Down,
        );
    }

    #[test]
    fn test_nowhere_to_move() {
        check(
            r#"
fn main() {
    match true {
        true =>$0$0 {
            println!("Hello, world");
        },
        false => {
            println!("Test");
        }
    };
}
            "#,
            expect![[r#"
fn main() {
    match true {
        true => {
            println!("Hello, world");
        },
        false => {
            println!("Test");
        }
    };
}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn test_moves_let_stmt_up() {
        check(
            r#"
fn main() {
    let test = 123;
    let test2$0$0 = 456;
}
            "#,
            expect![[r#"
fn main() {
    let test2 = 456;
    let test = 123;
}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn test_moves_expr_up() {
        check(
            r#"
fn main() {
    println!("Hello, world");
    println!("All I want to say is...");$0$0
}
            "#,
            expect![[r#"
fn main() {
    println!("All I want to say is...");
    println!("Hello, world");
}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn test_nowhere_to_move_stmt() {
        check(
            r#"
fn main() {
    println!("All I want to say is...");$0$0
    println!("Hello, world");
}
            "#,
            expect![[r#"
fn main() {
    println!("All I want to say is...");
    println!("Hello, world");
}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn test_move_item() {
        check(
            r#"
fn main() {}

fn foo() {}$0$0
            "#,
            expect![[r#"
fn foo() {}

fn main() {}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn test_move_impl_up() {
        check(
            r#"
struct Yay;

trait Wow {}

impl Wow for Yay {}$0$0
            "#,
            expect![[r#"
struct Yay;

impl Wow for Yay {}

trait Wow {}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn test_move_use_up() {
        check(
            r#"
use std::vec::Vec;
use std::collections::HashMap$0$0;
            "#,
            expect![[r#"
use std::collections::HashMap;
use std::vec::Vec;
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn moves_match_expr_up() {
        check(
            r#"
fn main() {
    let test = 123;

    $0match test {
        456 => {},
        _ => {}
    };$0
}
            "#,
            expect![[r#"
fn main() {
    match test {
        456 => {},
        _ => {}
    };

    let test = 123;
}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn moves_param_up() {
        check(
            r#"
fn test(one: i32, two$0$0: u32) {}

fn main() {
    test(123, 456);
}
            "#,
            expect![[r#"
fn test(two: u32, one: i32) {}

fn main() {
    test(123, 456);
}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn test_prioritizes_trait_items() {
        check(
            r#"
struct Test;

trait Yay {
    type One;

    type Two;

    fn inner();
}

impl Yay for Test {
    type One = i32;

    type Two = u32;

    fn inner() {$0$0
        println!("Mmmm");
    }
}
            "#,
            expect![[r#"
struct Test;

trait Yay {
    type One;

    type Two;

    fn inner();
}

impl Yay for Test {
    type One = i32;

    fn inner() {
        println!("Mmmm");
    }

    type Two = u32;
}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn test_weird_nesting() {
        check(
            r#"
fn test() {
    mod hello {
        fn inner() {}
    }

    mod hi {$0$0
        fn inner() {}
    }
}
            "#,
            expect![[r#"
fn test() {
    mod hi {
        fn inner() {}
    }

    mod hello {
        fn inner() {}
    }
}
            "#]],
            Direction::Up,
        );
    }

    #[test]
    fn handles_empty_file() {
        check(r#"$0$0"#, expect![[r#""#]], Direction::Up);
    }
}
<|endoftext|>


<commit_before>extern crate env_logger;
extern crate log;

use std::sync::{Once, ONCE_INIT};

pub struct LoggerUtils {}

struct SimpleLogger;

use self::log::{LogMetadata, LogRecord, LogLevelFilter};

impl log::Log for SimpleLogger {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        true
    }
    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            println!("{}:{} - {}", record.level(), record.target(), record.args());
        }
    }
}

impl LoggerUtils {
    pub fn init() {
        lazy_static! {
            static ref LOGGER_INIT: Once = ONCE_INIT;
        }

        LOGGER_INIT.call_once(|| {
            log::set_logger(|max_log_level| {
                max_log_level.set(LogLevelFilter::Info);
                Box::new(SimpleLogger)
            }).unwrap();
        });
    }
}

#[macro_export]
macro_rules! try_log {
    ($expr:expr) => (match $expr {
        Ok(val) => val,
        Err(err) => {
            error!("Error from try_log macro - {}", err);
            return Err(From::from(err))
        }
    })
}
<commit_msg>Update log format.<commit_after>extern crate env_logger;
extern crate log;

use std::sync::{Once, ONCE_INIT};

pub struct LoggerUtils {}

struct SimpleLogger;

use self::log::{LogMetadata, LogRecord, LogLevelFilter};

impl log::Log for SimpleLogger {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        true
    }
    fn log(&self, record: &LogRecord) {
        if self.enabled(record.metadata()) {
            println!("{:>5}|{:<25}|{:>30}:{:<4}| {}", record.level(), record.target(), record.location().file(), record.location().line(), record.args());
        }
    }
}

impl LoggerUtils {
    pub fn init() {
        lazy_static! {
            static ref LOGGER_INIT: Once = ONCE_INIT;
        }

        LOGGER_INIT.call_once(|| {
            log::set_logger(|max_log_level| {
                max_log_level.set(LogLevelFilter::Info);
                Box::new(SimpleLogger)
            }).unwrap();
        });
    }
}

#[macro_export]
macro_rules! try_log {
    ($expr:expr) => (match $expr {
        Ok(val) => val,
        Err(err) => {
            error!("try_log! | {}", err);
            return Err(From::from(err))
        }
    })
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 Fabian Schuiki

//! Packages

#![allow(dead_code)]

use std::cell::RefCell;

use common::NodeId;
use common::name::Name;
use common::source::{Span, Spanned, INVALID_SPAN};

use arenas::{Alloc, AllocInto};
use syntax::ast;

/// A placeholder for an HIR node.
pub struct Slot<'t, T>(RefCell<SlotState<'t, T>>)
where
    T: FromAst<'t> + 't;

#[derive(Copy, Clone)]
enum SlotState<'t, T>
where
    T: FromAst<'t> + 't,
{
    Fresh(&'t AnyScope, T::Input, T::Arena),
    ReadyOk(&'t T),
    ReadyErr,
}

impl<'t, T> Slot<'t, T>
where
    T: FromAst<'t>,
    T::Arena: AllocInto<'t, T> + Clone,
{
    /// Create a new slot.
    pub fn new(scope: &'t AnyScope, ast: T::Input, arena: T::Arena) -> Slot<'t, T> {
        Slot(RefCell::new(SlotState::Fresh(scope, ast, arena)))
    }

    /// Poll the slot, creating the HIR node from the AST the first time.
    pub fn poll(&self) -> Result<&'t T, ()> {
        match *self.0.borrow() {
            SlotState::ReadyOk(x) => return Ok(x),
            SlotState::ReadyErr => return Err(()),
            _ => (),
        }
        let (scope, ast, arena) = match self.0.replace(SlotState::ReadyErr) {
            SlotState::Fresh(scope, ast, arena) => (scope, ast, arena),
            _ => unreachable!(),
        };
        let node = T::from_ast(scope, ast, arena.clone()).map(|x| arena.alloc(x) as &T);
        self.0.replace(match node {
            Ok(x) => SlotState::ReadyOk(x),
            Err(()) => SlotState::ReadyErr,
        });
        node
    }
}

impl<'t, T> Node for Slot<'t, T>
where
    T: FromAst<'t> + Node,
    T::Arena: AllocInto<'t, T> + Clone,
{
    fn span(&self) -> Span {
        self.poll().map(Node::span).unwrap_or(INVALID_SPAN)
    }
}

pub struct Package2<'t> {
    id: NodeId,
    span: Span,
    name: Spanned<Name>,
    scope: &'t AnyScope,
    decls: Vec<Box<Node + 't>>,
}

impl<'t> FromAst<'t> for Package2<'t> {
    type Input = &'t ast::PkgDecl;
    type Arena = &'t Alloc<Self>;

    fn alloc_slot(
        scope: &'t AnyScope,
        ast: Self::Input,
        arena: Self::Arena,
    ) -> Result<Slot<'t, Self>, ()> {
        // TODO: register the package name in the scope
        Ok(Slot::new(scope, ast, arena))
    }

    fn from_ast(scope: &'t AnyScope, ast: Self::Input, _arena: Self::Arena) -> Result<Self, ()> {
        // TODO: create a new scope for the package
        let decls = ast.decls
            .iter()
            .flat_map(|decl| -> Option<Box<Node>> {
                match *decl {
                    // ast::DeclItem::PkgDecl(ref decl) => {
                    //     Some(Box::new(Package2::alloc_slot(scope, decl, arena).ok()?))
                    // }
                    // ast::DeclItem::TypeDecl(ref decl) => {
                    //     Some(Box::new(TypeDecl2::alloc_slot(scope, decl, arena).ok()?))
                    // }
                    _ => None,
                }
            })
            .collect::<Vec<_>>();
        Ok(Package2 {
            id: NodeId::alloc(),
            span: ast.span,
            name: ast.name,
            scope: scope,
            decls: decls,
        })
    }
}

impl<'t> Node for Package2<'t> {
    fn span(&self) -> Span {
        self.span
    }
}

pub struct TypeDecl2 {
    id: NodeId,
    span: Span,
    name: Spanned<Name>,
}

impl<'t> FromAst<'t> for TypeDecl2 {
    type Input = &'t ast::TypeDecl;
    type Arena = &'t Alloc<Self>;

    fn alloc_slot(
        scope: &'t AnyScope,
        ast: Self::Input,
        arena: Self::Arena,
    ) -> Result<Slot<'t, Self>, ()> {
        // TODO: register the type name in the scope
        Ok(Slot::new(scope, ast, arena))
    }

    fn from_ast(_scope: &'t AnyScope, ast: Self::Input, _arena: Self::Arena) -> Result<Self, ()> {
        Ok(TypeDecl2 {
            id: NodeId::alloc(),
            span: ast.span,
            name: ast.name,
        })
    }
}

impl Node for TypeDecl2 {
    fn span(&self) -> Span {
        self.span
    }
}

pub trait AnyScope {}

pub trait Node {
    /// The source file location of this node.
    fn span(&self) -> Span;
}

/// Construct something from an AST node.
pub trait FromAst<'t>: Sized {
    type Input;
    type Arena;

    fn alloc_slot(
        scope: &'t AnyScope,
        ast: Self::Input,
        arena: Self::Arena,
    ) -> Result<Slot<'t, Self>, ()>;

    fn from_ast(scope: &'t AnyScope, ast: Self::Input, arena: Self::Arena) -> Result<Self, ()>;
}
<commit_msg>vhdl: wrap allocation arena in context<commit_after>// Copyright (c) 2018 Fabian Schuiki

//! Packages

#![allow(dead_code)]

use std::cell::RefCell;

use common::NodeId;
use common::name::Name;
use common::source::{Span, Spanned, INVALID_SPAN};

use arenas::{Alloc, AllocInto};
use syntax::ast;

make_arenas!(
    pub struct Arenas<'t> {
        package:   Package2<'t>,
        type_decl: TypeDecl2,
    }
);

/// A placeholder for an HIR node.
pub struct Slot<'t, T>(RefCell<SlotState<'t, T>>)
where
    T: FromAst<'t> + 't;

#[derive(Copy, Clone)]
enum SlotState<'t, T>
where
    T: FromAst<'t> + 't,
{
    Fresh(&'t AnyScope, T::Input, T::Arena),
    ReadyOk(&'t T),
    ReadyErr,
}

impl<'t, T> Slot<'t, T>
where
    T: FromAst<'t>,
    T::Arena: AllocInto<'t, T> + Clone,
{
    /// Create a new slot.
    pub fn new(scope: &'t AnyScope, ast: T::Input, arena: T::Arena) -> Slot<'t, T> {
        Slot(RefCell::new(SlotState::Fresh(scope, ast, arena)))
    }

    /// Poll the slot, creating the HIR node from the AST the first time.
    pub fn poll(&self) -> Result<&'t T, ()> {
        match *self.0.borrow() {
            SlotState::ReadyOk(x) => return Ok(x),
            SlotState::ReadyErr => return Err(()),
            _ => (),
        }
        let (scope, ast, arena) = match self.0.replace(SlotState::ReadyErr) {
            SlotState::Fresh(scope, ast, arena) => (scope, ast, arena),
            _ => unreachable!(),
        };
        let node = T::from_ast(scope, ast, arena.clone()).map(|x| arena.alloc(x) as &T);
        self.0.replace(match node {
            Ok(x) => SlotState::ReadyOk(x),
            Err(()) => SlotState::ReadyErr,
        });
        node
    }
}

impl<'t, T> Node for Slot<'t, T>
where
    T: FromAst<'t> + Node,
    T::Arena: AllocInto<'t, T> + Clone,
{
    fn span(&self) -> Span {
        self.poll().map(Node::span).unwrap_or(INVALID_SPAN)
    }
}

pub struct Package2<'t> {
    id: NodeId,
    span: Span,
    name: Spanned<Name>,
    scope: &'t AnyScope,
    decls: Vec<Box<Node + 't>>,
}

impl<'t> FromAst<'t> for Package2<'t> {
    type Input = &'t ast::PkgDecl;
    type Arena = Context<'t>;

    fn alloc_slot(
        scope: &'t AnyScope,
        ast: Self::Input,
        arena: Self::Arena,
    ) -> Result<Slot<'t, Self>, ()> {
        // TODO: register the package name in the scope
        Ok(Slot::new(scope, ast, arena))
    }

    fn from_ast(scope: &'t AnyScope, ast: Self::Input, _arena: Self::Arena) -> Result<Self, ()> {
        // TODO: create a new scope for the package
        let decls = ast.decls
            .iter()
            .flat_map(|decl| -> Option<Box<Node>> {
                match *decl {
                    // ast::DeclItem::PkgDecl(ref decl) => {
                    //     Some(Box::new(Package2::alloc_slot(scope, decl, arena).ok()?))
                    // }
                    // ast::DeclItem::TypeDecl(ref decl) => {
                    //     Some(Box::new(TypeDecl2::alloc_slot(scope, decl, arena).ok()?))
                    // }
                    _ => None,
                }
            })
            .collect::<Vec<_>>();
        Ok(Package2 {
            id: NodeId::alloc(),
            span: ast.span,
            name: ast.name,
            scope: scope,
            decls: decls,
        })
    }
}

impl<'t> Node for Package2<'t> {
    fn span(&self) -> Span {
        self.span
    }
}

pub struct TypeDecl2 {
    id: NodeId,
    span: Span,
    name: Spanned<Name>,
}

impl<'t> FromAst<'t> for TypeDecl2 {
    type Input = &'t ast::TypeDecl;
    type Arena = &'t Alloc<Self>;

    fn alloc_slot(
        scope: &'t AnyScope,
        ast: Self::Input,
        arena: Self::Arena,
    ) -> Result<Slot<'t, Self>, ()> {
        // TODO: register the type name in the scope
        Ok(Slot::new(scope, ast, arena))
    }

    fn from_ast(_scope: &'t AnyScope, ast: Self::Input, _arena: Self::Arena) -> Result<Self, ()> {
        Ok(TypeDecl2 {
            id: NodeId::alloc(),
            span: ast.span,
            name: ast.name,
        })
    }
}

impl Node for TypeDecl2 {
    fn span(&self) -> Span {
        self.span
    }
}

pub trait AnyScope {}

pub trait Node {
    /// The source file location of this node.
    fn span(&self) -> Span;
}

/// Construct something from an AST node.
pub trait FromAst<'t>: Sized {
    type Input;
    type Arena;

    fn alloc_slot(
        scope: &'t AnyScope,
        ast: Self::Input,
        arena: Self::Arena,
    ) -> Result<Slot<'t, Self>, ()>;

    fn from_ast(scope: &'t AnyScope, ast: Self::Input, arena: Self::Arena) -> Result<Self, ()>;
}

#[derive(Copy, Clone)]
pub struct Context<'t> {
    arenas: &'t Arenas<'t>,
}

impl<'t, T> AllocInto<'t, T> for Context<'t>
where
    Arenas<'t>: Alloc<T>,
{
    fn alloc(&self, value: T) -> &'t mut T {
        self.arenas.alloc(value)
    }
}
<|endoftext|>


<commit_before>extern crate structopt;

use crate::exceptions::IOError;
use std::path::PathBuf;
use std::process::{Command, ExitStatus};
use std::fs;
use structopt::StructOpt;

#[derive(Clone, Debug, StructOpt)]
#[structopt(name = "performance", about = "performance regression testing runner")]
enum Performance {
    #[structopt(name = "measure")]
    Measure {
        #[structopt(parse(from_os_str))]
        #[structopt(short)]
        projects_dir: PathBuf,
        #[structopt(short)]
        branch_name: bool,
    },
    #[structopt(name = "compare")]
    Compare {
        #[structopt(parse(from_os_str))]
        #[structopt(short)]
        results_dir: PathBuf,
    },
}

#[derive(Debug, Clone)]
struct Metric<'a> {
    name: &'a str,
    prepare: &'a str,
    cmd: &'a str,
}

impl Metric<'_> {
    fn outfile(&self, project: &str, branch: &str) -> String {
        [branch, "_", self.name, "_", project, ".json"].join("")
    }
}

// calls hyperfine via system command, and returns result of runs
pub fn measure(projects_directory: &PathBuf, dbt_branch: &str) -> Result<Vec<ExitStatus>, IOError> {
    // to add a new metric to the test suite, simply define it in this list:
    // TODO read from some config file?
    let metrics: Vec<Metric> = vec![
        Metric {
            name: "parse",
            prepare: "rm -rf target/",
            cmd: "dbt parse --no-version-check",
        },
    ];

    fs::read_dir(projects_directory)
        .or_else(|e| Err(IOError::ReadErr(projects_directory.to_path_buf(), Some(e))))?
        .map(|entry| {
            let path = entry
                .or_else(|e| Err(IOError::ReadErr(projects_directory.to_path_buf(), Some(e))))?
                .path();

            let project_name: String = path
                .file_name()
                .ok_or_else(|| IOError::MissingFilenameErr(path.clone().to_path_buf()))
                .and_then(|x| x.to_str().ok_or_else(|| IOError::FilenameNotUnicodeErr(path.clone().to_path_buf())))?
                .to_owned();

            // each project-metric pair we will run
            let pairs = metrics
                .clone()
                .into_iter()
                .map(|metric| (path.clone(), project_name.clone(), metric))
                .collect::<Vec<(PathBuf, String, Metric)>>();

            Ok(pairs)
        })
        .collect::<Result<Vec<Vec<(PathBuf, String, Metric)>>, IOError>>()?
        .concat()
        .into_iter()
        // run hyperfine on each pairing
        .map(|(path, project_name, metric)| {
            Command::new("hyperfine")
                .current_dir(&path)
                // warms filesystem caches by running the command first without counting it.
                // alternatively we could clear them before each run
                .arg("--warmup")
                .arg("1")
                .arg("--prepare")
                .arg(metric.prepare)
                .arg([metric.cmd, " --profiles-dir ", "../../project_config/"].join(""))
                .arg("--export-json")
                .arg(
                    ["../../results/", &metric.outfile(&project_name, &dbt_branch)].join(""),
                )
                // this prevents capture dbt output. Noisy, but good for debugging when tests fail.
                .arg("--show-output")
                .status() // use spawn() here instead for more information
                .or_else(|e| Err(IOError::CommandErr(Some(e))))
        }
    )
    .collect()
}
<commit_msg>more reference handling<commit_after>extern crate structopt;

use crate::exceptions::IOError;
use std::path::PathBuf;
use std::process::{Command, ExitStatus};
use std::fs;
use structopt::StructOpt;

#[derive(Clone, Debug, StructOpt)]
#[structopt(name = "performance", about = "performance regression testing runner")]
enum Performance {
    #[structopt(name = "measure")]
    Measure {
        #[structopt(parse(from_os_str))]
        #[structopt(short)]
        projects_dir: PathBuf,
        #[structopt(short)]
        branch_name: bool,
    },
    #[structopt(name = "compare")]
    Compare {
        #[structopt(parse(from_os_str))]
        #[structopt(short)]
        results_dir: PathBuf,
    },
}

#[derive(Debug, Clone)]
struct Metric<'a> {
    name: &'a str,
    prepare: &'a str,
    cmd: &'a str,
}

impl Metric<'_> {
    fn outfile(&self, project: &str, branch: &str) -> String {
        [branch, "_", self.name, "_", project, ".json"].join("")
    }
}

// calls hyperfine via system command, and returns result of runs
pub fn measure<'a>(projects_directory: &PathBuf, dbt_branch: &str) -> Result<Vec<ExitStatus>, IOError> {
    // to add a new metric to the test suite, simply define it in this list:
    // TODO read from some config file?
    let metrics: Vec<Metric> = vec![
        Metric {
            name: "parse",
            prepare: "rm -rf target/",
            cmd: "dbt parse --no-version-check",
        },
    ];

    fs::read_dir(projects_directory)
        .or_else(|e| Err(IOError::ReadErr(projects_directory.to_path_buf(), Some(e))))?
        .map(|entry| {
            let path = entry
                .or_else(|e| Err(IOError::ReadErr(projects_directory.to_path_buf(), Some(e))))?
                .path();

            let project_name: String = path
                .file_name()
                .ok_or_else(|| IOError::MissingFilenameErr(path.clone().to_path_buf()))
                .and_then(|x| x.to_str().ok_or_else(|| IOError::FilenameNotUnicodeErr(path.clone().to_path_buf())))?
                .to_owned();

            // each project-metric pair we will run
            let pairs = metrics
                .iter()
                .map(|metric| (path.clone(), project_name.clone(), metric))
                .collect::<Vec<(PathBuf, String, &Metric<'a>)>>();

            Ok(pairs)
        })
        .collect::<Result<Vec<Vec<(PathBuf, String, &Metric<'a>)>>, IOError>>()?
        .concat()
        .iter()
        // run hyperfine on each pairing
        .map(|(path, project_name, metric)| {
            Command::new("hyperfine")
                .current_dir(path)
                // warms filesystem caches by running the command first without counting it.
                // alternatively we could clear them before each run
                .arg("--warmup")
                .arg("1")
                .arg("--prepare")
                .arg(metric.prepare)
                .arg([metric.cmd, " --profiles-dir ", "../../project_config/"].join(""))
                .arg("--export-json")
                .arg(
                    ["../../results/", &metric.outfile(project_name, dbt_branch)].join(""),
                )
                // this prevents capture dbt output. Noisy, but good for debugging when tests fail.
                .arg("--show-output")
                .status() // use spawn() here instead for more information
                .or_else(|e| Err(IOError::CommandErr(Some(e))))
        }
    )
    .collect()
}
<|endoftext|>


<commit_before>#![feature(collections)]
#![feature(link_args)]
#![feature(start)]
#![feature(unique)]

extern crate cairo;
extern crate collections;
extern crate graphics;
extern crate os;
extern crate syscall;

use cairo::bindings::*;
use cairo::cairo::Cairo;
use cairo::CairoObj;
use collections::btree_map::BTreeMap;
use graphics::{Client,Event,Window};
use os::{Mutex,Result};

fn run() -> Result<()> {
    let client = Mutex::new(Client::new())?;
    let mut windows = BTreeMap::new();
    for i in 0 .. 5 {
        let mut window = Window::new(&client, i as f64 * 100.0, i as f64 * 100.0, 150.0, 120.0)?;

        {
            let surface = window.create_surface();
            let cr = Cairo::new(surface);
            unsafe {
                let pat = CairoObj::wrap(cairo_pattern_create_linear(0.0, 0.0, 0.0, 100.0));
                cairo_pattern_add_color_stop_rgba(pat.as_ptr(), 1.0, 0.0, 0.0, 0.0, 1.0);
                cairo_pattern_add_color_stop_rgba(pat.as_ptr(), 0.0, 1.0, 1.0, 1.0, 1.0);
                cairo_rectangle(cr.as_ptr(), 0.0, 0.0, 150.0, 120.0);
                cairo_set_source(cr.as_ptr(), pat.as_ptr());
                cairo_fill(cr.as_ptr());
            }
        }

        window.invalidate()?;
        windows.insert(window.id(), window);
    }

    let checkpoint_id = {
        let mut client = client.lock().unwrap();
        client.checkpoint()?
    };

    while !windows.is_empty() {
        let e = {
            let mut client = client.lock().unwrap();
            client.wait_for_event()?
        };

        println!("[Client] {:?}", e);
        match e {
            Event::Checkpoint { id } if id == checkpoint_id => {
                println!("System ready");
            },

            Event::KeyPress { window_id, code } if code as u32 == 27 => {
                windows.remove(&window_id);
            },

            _ => { }
        }
    }

    Ok(())
}

#[cfg(target_arch="x86_64")]
#[link_args = "-T ../../libsyscall/arch/amd64/link.ld"]
extern {
}

#[start]
#[no_mangle]
pub fn start(_: isize, _: *const *const u8) -> isize {
    run().map(|()| 0).unwrap_or_else(|num| -(num as isize))
}
<commit_msg>Type some words<commit_after>#![feature(collections)]
#![feature(link_args)]
#![feature(start)]
#![feature(unique)]

extern crate cairo;
extern crate collections;
extern crate graphics;
extern crate os;
extern crate syscall;

use cairo::bindings::*;
use cairo::cairo::Cairo;
use cairo::CairoObj;
use collections::btree_map::BTreeMap;
use graphics::{Client,Event,Window};
use os::{Mutex,Result};
use std::mem;

struct DemoWindow<'a> {
    window: Window<'a>,
    text: String,
}

impl<'a> DemoWindow<'a> {
    fn new(window: Window<'a>) -> Result<Self> {
        let mut demo_window = DemoWindow { window, text: "".into() };
        demo_window.invalidate()?;
        Ok(demo_window)
    }

    fn invalidate(&mut self) -> Result<()> {
        {
            let surface = self.window.create_surface();
            let cr = Cairo::new(surface);
            unsafe {
                let pat = CairoObj::wrap(cairo_pattern_create_linear(0.0, 0.0, 0.0, 100.0));
                cairo_pattern_add_color_stop_rgba(pat.as_ptr(), 1.0, 0.0, 0.0, 0.0, 1.0);
                cairo_pattern_add_color_stop_rgba(pat.as_ptr(), 0.0, 1.0, 1.0, 1.0, 1.0);
                cairo_rectangle(cr.as_ptr(), 0.0, 0.0, 150.0, 120.0);
                cairo_set_source(cr.as_ptr(), pat.as_ptr());
                cairo_fill(cr.as_ptr());

                let mut text = Vec::<u8>::with_capacity(self.text.len() + 1);
                text.extend(self.text.as_bytes());
                text.push(0);

                let text_ptr = (&text[..]).as_ptr() as *const i8;
                let mut extents = mem::uninitialized();
                cairo_text_extents(cr.as_ptr(), text_ptr, &mut extents);
                cairo_set_source_rgb(cr.as_ptr(), 0.0, 0.0, 0.0);
                cairo_move_to(cr.as_ptr(), 0.0, extents.height);
                cairo_show_text(cr.as_ptr(), text_ptr);
            }
        }

        self.window.invalidate()?;
        Ok(())
    }

    fn handle_keypress(&mut self, code: char) -> Result<()> {
        self.text.push(code);
        self.invalidate()?;
        Ok(())
    }
}

fn run() -> Result<()> {
    let client = Mutex::new(Client::new())?;
    let mut windows = BTreeMap::new();
    for i in 0 .. 5 {
        let window = Window::new(&client, i as f64 * 100.0, i as f64 * 100.0, 150.0, 120.0)?;
        windows.insert(window.id(), DemoWindow::new(window)?);
    }

    let checkpoint_id = {
        let mut client = client.lock().unwrap();
        client.checkpoint()?
    };

    while !windows.is_empty() {
        let e = {
            let mut client = client.lock().unwrap();
            client.wait_for_event()?
        };

        println!("[Client] {:?}", e);
        match e {
            Event::Checkpoint { id } if id == checkpoint_id => {
                println!("System ready");
            },

            Event::KeyPress { window_id, code } if code == '\u{1b}' => {
                windows.remove(&window_id);
            },

            Event::KeyPress { window_id, code } => {
                if let Some(demo_window) = windows.get_mut(&window_id) {
                    demo_window.handle_keypress(code)?;
                }
            }

            _ => { }
        }
    }

    Ok(())
}

#[cfg(target_arch="x86_64")]
#[link_args = "-T ../../libsyscall/arch/amd64/link.ld"]
extern {
}

#[start]
#[no_mangle]
pub fn start(_: isize, _: *const *const u8) -> isize {
    run().map(|()| 0).unwrap_or_else(|num| -(num as isize))
}
<|endoftext|>


<commit_before>#![deny(unused_import_braces, unused_imports,
        unused_comparisons, unused_must_use,
        unused_variables, non_shorthand_field_patterns,
        unreachable_code)]

extern crate trie;
extern crate sputnikvm;
extern crate sha3;
extern crate block;
extern crate rlp;
extern crate bigint;

use bigint::{H256, U256, M256, Address};
use sputnikvm::{ValidTransaction, HeaderParams, Memory, TransactionVM, VM,
                AccountCommitment, Patch, AccountState, AccountChange};
use sputnikvm::errors::{PreExecutionError, RequireError};
use sha3::{Keccak256, Digest};
use trie::{FixedSecureTrie, DatabaseGuard, MemoryDatabase, Database, DatabaseOwned};
use block::{Account, Transaction};
use std::collections::HashMap;
use std::cmp::min;

pub struct Stateful<D> {
    database: D,
    root: H256,
}

impl<D> Stateful<D> {
    pub fn new(database: D, root: H256) -> Self {
        Self {
            database,
            root
        }
    }
}

impl<D: Default> Default for Stateful<D> {
    fn default() -> Self {
        Self {
            database: D::default(),
            root: MemoryDatabase::new().create_empty().root(),
        }
    }
}

impl<D: DatabaseOwned> Stateful<D> {
    fn is_empty_hash(hash: H256) -> bool {
        hash == H256::from(Keccak256::digest(&[]).as_slice())
    }

    pub fn code(&self, hash: H256) -> Option<Vec<u8>> {
        let code_hashes = self.database.create_guard();

        if Self::is_empty_hash(hash) {
            Some(Vec::new())
        } else {
            code_hashes.get(hash)
        }
    }

    pub fn call<M: Memory + Default, P: Patch>(
        &self, transaction: ValidTransaction, block: HeaderParams,
        most_recent_block_hashes: &[H256]
    ) -> TransactionVM<M, P> {
        assert!(U256::from(most_recent_block_hashes.len()) >=
                min(block.number, U256::from(256)));

        let mut vm = TransactionVM::new(transaction, block.clone());
        let state = self.database.create_fixed_secure_trie(self.root);
        let code_hashes = self.database.create_guard();

        loop {
            match vm.fire() {
                Ok(()) => break,
                Err(RequireError::Account(address)) => {
                    let account: Option<Account> = state.get(&address);

                    match account {
                        Some(account) => {
                            let code = if Self::is_empty_hash(account.code_hash) {
                                Vec::new()
                            } else {
                                code_hashes.get(account.code_hash).unwrap()
                            };

                            vm.commit_account(AccountCommitment::Full {
                                nonce: account.nonce,
                                address: address,
                                balance: account.balance,
                                code: code,
                            }).unwrap();
                        },
                        None => {
                            vm.commit_account(AccountCommitment::Nonexist(address)).unwrap();
                        },
                    }
                },
                Err(RequireError::AccountCode(address)) => {
                    let account: Option<Account> = state.get(&address);

                    match account {
                        Some(account) => {
                            let code = if Self::is_empty_hash(account.code_hash) {
                                Vec::new()
                            } else {
                                code_hashes.get(account.code_hash).unwrap()
                            };

                            vm.commit_account(AccountCommitment::Code {
                                address: address,
                                code: code,
                            }).unwrap();
                        },
                        None => {
                            vm.commit_account(AccountCommitment::Nonexist(address)).unwrap();
                        },
                    }
                },
                Err(RequireError::AccountStorage(address, index)) => {
                    let account: Option<Account> = state.get(&address);

                    match account {
                        Some(account) => {
                            let storage = self.database.create_fixed_secure_trie(account.storage_root);
                            let value = storage.get(&H256::from(index)).unwrap_or(M256::zero());

                            vm.commit_account(AccountCommitment::Storage {
                                address: address,
                                index, value
                            }).unwrap();
                        },
                        None => {
                            vm.commit_account(AccountCommitment::Nonexist(address)).unwrap();
                        },
                    }
                },
                Err(RequireError::Blockhash(number)) => {
                    let index = (block.number - number).as_usize();
                    vm.commit_blockhash(number, most_recent_block_hashes[index]).unwrap();
                },
            }
        }

        vm
    }

    pub fn transit(
        &mut self, accounts: &[AccountChange]
    ) {
        let mut state = self.database.create_fixed_secure_trie(self.root);
        let mut code_hashes = self.database.create_guard();

        for account in accounts {
            match account.clone() {
                AccountChange::Full {
                    nonce, address, balance, changing_storage, code
                } => {
                    let changing_storage: HashMap<U256, M256> = changing_storage.into();

                    let mut account: Account = state.get(&address).unwrap();

                    let mut storage_trie = self.database.create_fixed_secure_trie(account.storage_root);
                    for (key, value) in changing_storage {
                        if value == M256::zero() {
                            storage_trie.remove(&H256::from(key));
                        } else {
                            storage_trie.insert(H256::from(key), value);
                        }
                    }

                    account.balance = balance;
                    account.nonce = nonce;
                    account.storage_root = storage_trie.root();
                    assert!(account.code_hash == H256::from(Keccak256::digest(&code).as_slice()));

                    state.insert(address, account);
                },
                AccountChange::IncreaseBalance(address, value) => {
                    match state.get(&address) {
                        Some(mut account) => {
                            account.balance = account.balance + value;
                            state.insert(address, account);
                        },
                        None => {
                            let account = Account {
                                nonce: U256::zero(),
                                balance: value,
                                storage_root: self.database.create_empty().root(),
                                code_hash: H256::from(Keccak256::digest(&[]).as_slice())
                            };
                            state.insert(address, account);
                        }
                    }
                },
                AccountChange::DecreaseBalance(address, value) => {
                    let mut account: Account = state.get(&address).unwrap();

                    account.balance = account.balance - value;

                    state.insert(address, account);
                },
                AccountChange::Create {
                    nonce, address, balance, storage, code, exists
                } => {
                    if !exists {
                        state.remove(&address);
                    } else {
                        let storage: HashMap<U256, M256> = storage.into();

                        let mut storage_trie = self.database.create_fixed_secure_empty();
                        for (key, value) in storage {
                            if value == M256::zero() {
                                storage_trie.remove(&H256::from(key));
                            } else {
                                storage_trie.insert(H256::from(key), value);
                            }
                        }

                        let code_hash = H256::from(Keccak256::digest(&code).as_slice());
                        code_hashes.set(code_hash, code);

                        let account = Account {
                            nonce: nonce,
                            balance: balance,
                            storage_root: storage_trie.root(),
                            code_hash
                        };

                        state.insert(address, account);
                    }
                },
            }
        }

        self.root = state.root();
    }

    pub fn execute<M: Memory + Default, P: Patch>(
        &mut self, transaction: ValidTransaction, block: HeaderParams,
        most_recent_block_hashes: &[H256]
    ) -> TransactionVM<M, P> {
        let vm = self.call::<_, P>(transaction, block, most_recent_block_hashes);
        let mut accounts = Vec::new();
        for account in vm.accounts() {
            accounts.push(account.clone());
        }
        self.transit(&accounts);
        vm
    }

    pub fn to_valid<P: Patch>(
        &self, transaction: Transaction,
    ) -> Result<ValidTransaction, PreExecutionError> {
        let state = self.database.create_fixed_secure_trie(self.root);
        let code_hashes = self.database.create_guard();
        let mut account_state = AccountState::default();

        loop {
            match ValidTransaction::from_transaction::<P>(&transaction, &account_state) {
                Ok(val) => return val,
                Err(RequireError::Account(address)) => {
                    let account: Option<Account> = state.get(&address);

                    match account {
                        Some(account) => {
                            let code = if Self::is_empty_hash(account.code_hash) {
                                Vec::new()
                            } else {
                                code_hashes.get(account.code_hash).unwrap()
                            };

                            account_state.commit(AccountCommitment::Full {
                                nonce: account.nonce,
                                address: address,
                                balance: account.balance,
                                code: code,
                            }).unwrap();
                        },
                        None => {
                            account_state.commit(AccountCommitment::Nonexist(address)).unwrap();
                        },
                    }
                },
                Err(RequireError::AccountCode(address)) => {
                    let account: Option<Account> = state.get(&address);

                    match account {
                        Some(account) => {
                            let code = if Self::is_empty_hash(account.code_hash) {
                                Vec::new()
                            } else {
                                code_hashes.get(account.code_hash).unwrap()
                            };

                            account_state.commit(AccountCommitment::Code {
                                address: address,
                                code: code,
                            }).unwrap();
                        },
                        None => {
                            account_state.commit(AccountCommitment::Nonexist(address)).unwrap();
                        },
                    }
                },
                Err(RequireError::AccountStorage(address, index)) => {
                    let account: Option<Account> = state.get(&address);

                    match account {
                        Some(account) => {
                            let storage = self.database.create_fixed_secure_trie(account.storage_root);
                            let value = storage.get(&H256::from(index)).unwrap_or(M256::zero());

                            account_state.commit(AccountCommitment::Storage {
                                address: address,
                                index, value
                            }).unwrap();
                        },
                        None => {
                            account_state.commit(AccountCommitment::Nonexist(address)).unwrap();
                        },
                    }
                },
                Err(RequireError::Blockhash(_)) => {
                    panic!()
                },
            }
        }
    }

    pub fn root(&self) -> H256 {
        self.root
    }

    pub fn state_of<'a>(&'a self, root: H256) -> FixedSecureTrie<<D as Database<'a>>::Guard, Address, Account> {
        self.database.create_fixed_secure_trie::<Address, Account>(root)
    }

    pub fn state<'a>(&'a self) -> FixedSecureTrie<<D as Database<'a>>::Guard, Address, Account> {
        self.state_of(self.root())
    }

    pub fn storage_state_of<'a>(&'a self, root: H256) -> FixedSecureTrie<<D as Database<'a>>::Guard, H256, M256> {
        self.database.create_fixed_secure_trie::<H256, M256>(root)
    }

    pub fn storage_state<'a>(&'a self, address: Address) -> Option<FixedSecureTrie<<D as Database<'a>>::Guard, H256, M256>> {
        let state = self.state();
        let account = state.get(&address);

        match account {
            Some(account) => {
                Some(self.storage_state_of(account.storage_root))
            },
            None => None,
        }
    }
}

pub type MemoryStateful = Stateful<MemoryDatabase>;
<commit_msg>Make stateful only take a reference of the database<commit_after>#![deny(unused_import_braces, unused_imports,
        unused_comparisons, unused_must_use,
        unused_variables, non_shorthand_field_patterns,
        unreachable_code)]

extern crate trie;
extern crate sputnikvm;
extern crate sha3;
extern crate block;
extern crate rlp;
extern crate bigint;

use bigint::{H256, U256, M256, Address};
use sputnikvm::{ValidTransaction, HeaderParams, Memory, TransactionVM, VM,
                AccountCommitment, Patch, AccountState, AccountChange};
use sputnikvm::errors::{PreExecutionError, RequireError};
use sha3::{Keccak256, Digest};
use trie::{FixedSecureTrie, DatabaseGuard, MemoryDatabase, Database, DatabaseOwned};
use block::{Account, Transaction};
use std::collections::HashMap;
use std::cmp::min;

pub struct Stateful<'a, D: 'a> {
    database: &'a D,
    root: H256,
}

impl<'a, D> Stateful<'a, D> {
    pub fn new(database: &'a D, root: H256) -> Self {
        Self {
            database,
            root
        }
    }

    pub fn empty(database: &'a D) -> Self {
        Self::new(database, MemoryDatabase::new().create_empty().root())
    }
}

impl<'b, D: DatabaseOwned> Stateful<'b, D> {
    fn is_empty_hash(hash: H256) -> bool {
        hash == H256::from(Keccak256::digest(&[]).as_slice())
    }

    pub fn code(&self, hash: H256) -> Option<Vec<u8>> {
        let code_hashes = self.database.create_guard();

        if Self::is_empty_hash(hash) {
            Some(Vec::new())
        } else {
            code_hashes.get(hash)
        }
    }

    pub fn call<M: Memory + Default, P: Patch>(
        &self, transaction: ValidTransaction, block: HeaderParams,
        most_recent_block_hashes: &[H256]
    ) -> TransactionVM<M, P> {
        assert!(U256::from(most_recent_block_hashes.len()) >=
                min(block.number, U256::from(256)));

        let mut vm = TransactionVM::new(transaction, block.clone());
        let state = self.database.create_fixed_secure_trie(self.root);
        let code_hashes = self.database.create_guard();

        loop {
            match vm.fire() {
                Ok(()) => break,
                Err(RequireError::Account(address)) => {
                    let account: Option<Account> = state.get(&address);

                    match account {
                        Some(account) => {
                            let code = if Self::is_empty_hash(account.code_hash) {
                                Vec::new()
                            } else {
                                code_hashes.get(account.code_hash).unwrap()
                            };

                            vm.commit_account(AccountCommitment::Full {
                                nonce: account.nonce,
                                address: address,
                                balance: account.balance,
                                code: code,
                            }).unwrap();
                        },
                        None => {
                            vm.commit_account(AccountCommitment::Nonexist(address)).unwrap();
                        },
                    }
                },
                Err(RequireError::AccountCode(address)) => {
                    let account: Option<Account> = state.get(&address);

                    match account {
                        Some(account) => {
                            let code = if Self::is_empty_hash(account.code_hash) {
                                Vec::new()
                            } else {
                                code_hashes.get(account.code_hash).unwrap()
                            };

                            vm.commit_account(AccountCommitment::Code {
                                address: address,
                                code: code,
                            }).unwrap();
                        },
                        None => {
                            vm.commit_account(AccountCommitment::Nonexist(address)).unwrap();
                        },
                    }
                },
                Err(RequireError::AccountStorage(address, index)) => {
                    let account: Option<Account> = state.get(&address);

                    match account {
                        Some(account) => {
                            let storage = self.database.create_fixed_secure_trie(account.storage_root);
                            let value = storage.get(&H256::from(index)).unwrap_or(M256::zero());

                            vm.commit_account(AccountCommitment::Storage {
                                address: address,
                                index, value
                            }).unwrap();
                        },
                        None => {
                            vm.commit_account(AccountCommitment::Nonexist(address)).unwrap();
                        },
                    }
                },
                Err(RequireError::Blockhash(number)) => {
                    let index = (block.number - number).as_usize();
                    vm.commit_blockhash(number, most_recent_block_hashes[index]).unwrap();
                },
            }
        }

        vm
    }

    pub fn transit(
        &mut self, accounts: &[AccountChange]
    ) {
        let mut state = self.database.create_fixed_secure_trie(self.root);
        let mut code_hashes = self.database.create_guard();

        for account in accounts {
            match account.clone() {
                AccountChange::Full {
                    nonce, address, balance, changing_storage, code
                } => {
                    let changing_storage: HashMap<U256, M256> = changing_storage.into();

                    let mut account: Account = state.get(&address).unwrap();

                    let mut storage_trie = self.database.create_fixed_secure_trie(account.storage_root);
                    for (key, value) in changing_storage {
                        if value == M256::zero() {
                            storage_trie.remove(&H256::from(key));
                        } else {
                            storage_trie.insert(H256::from(key), value);
                        }
                    }

                    account.balance = balance;
                    account.nonce = nonce;
                    account.storage_root = storage_trie.root();
                    assert!(account.code_hash == H256::from(Keccak256::digest(&code).as_slice()));

                    state.insert(address, account);
                },
                AccountChange::IncreaseBalance(address, value) => {
                    match state.get(&address) {
                        Some(mut account) => {
                            account.balance = account.balance + value;
                            state.insert(address, account);
                        },
                        None => {
                            let account = Account {
                                nonce: U256::zero(),
                                balance: value,
                                storage_root: self.database.create_empty().root(),
                                code_hash: H256::from(Keccak256::digest(&[]).as_slice())
                            };
                            state.insert(address, account);
                        }
                    }
                },
                AccountChange::DecreaseBalance(address, value) => {
                    let mut account: Account = state.get(&address).unwrap();

                    account.balance = account.balance - value;

                    state.insert(address, account);
                },
                AccountChange::Create {
                    nonce, address, balance, storage, code, exists
                } => {
                    if !exists {
                        state.remove(&address);
                    } else {
                        let storage: HashMap<U256, M256> = storage.into();

                        let mut storage_trie = self.database.create_fixed_secure_empty();
                        for (key, value) in storage {
                            if value == M256::zero() {
                                storage_trie.remove(&H256::from(key));
                            } else {
                                storage_trie.insert(H256::from(key), value);
                            }
                        }

                        let code_hash = H256::from(Keccak256::digest(&code).as_slice());
                        code_hashes.set(code_hash, code);

                        let account = Account {
                            nonce: nonce,
                            balance: balance,
                            storage_root: storage_trie.root(),
                            code_hash
                        };

                        state.insert(address, account);
                    }
                },
            }
        }

        self.root = state.root();
    }

    pub fn execute<M: Memory + Default, P: Patch>(
        &mut self, transaction: ValidTransaction, block: HeaderParams,
        most_recent_block_hashes: &[H256]
    ) -> TransactionVM<M, P> {
        let vm = self.call::<_, P>(transaction, block, most_recent_block_hashes);
        let mut accounts = Vec::new();
        for account in vm.accounts() {
            accounts.push(account.clone());
        }
        self.transit(&accounts);
        vm
    }

    pub fn to_valid<P: Patch>(
        &self, transaction: Transaction,
    ) -> Result<ValidTransaction, PreExecutionError> {
        let state = self.database.create_fixed_secure_trie(self.root);
        let code_hashes = self.database.create_guard();
        let mut account_state = AccountState::default();

        loop {
            match ValidTransaction::from_transaction::<P>(&transaction, &account_state) {
                Ok(val) => return val,
                Err(RequireError::Account(address)) => {
                    let account: Option<Account> = state.get(&address);

                    match account {
                        Some(account) => {
                            let code = if Self::is_empty_hash(account.code_hash) {
                                Vec::new()
                            } else {
                                code_hashes.get(account.code_hash).unwrap()
                            };

                            account_state.commit(AccountCommitment::Full {
                                nonce: account.nonce,
                                address: address,
                                balance: account.balance,
                                code: code,
                            }).unwrap();
                        },
                        None => {
                            account_state.commit(AccountCommitment::Nonexist(address)).unwrap();
                        },
                    }
                },
                Err(RequireError::AccountCode(address)) => {
                    let account: Option<Account> = state.get(&address);

                    match account {
                        Some(account) => {
                            let code = if Self::is_empty_hash(account.code_hash) {
                                Vec::new()
                            } else {
                                code_hashes.get(account.code_hash).unwrap()
                            };

                            account_state.commit(AccountCommitment::Code {
                                address: address,
                                code: code,
                            }).unwrap();
                        },
                        None => {
                            account_state.commit(AccountCommitment::Nonexist(address)).unwrap();
                        },
                    }
                },
                Err(RequireError::AccountStorage(address, index)) => {
                    let account: Option<Account> = state.get(&address);

                    match account {
                        Some(account) => {
                            let storage = self.database.create_fixed_secure_trie(account.storage_root);
                            let value = storage.get(&H256::from(index)).unwrap_or(M256::zero());

                            account_state.commit(AccountCommitment::Storage {
                                address: address,
                                index, value
                            }).unwrap();
                        },
                        None => {
                            account_state.commit(AccountCommitment::Nonexist(address)).unwrap();
                        },
                    }
                },
                Err(RequireError::Blockhash(_)) => {
                    panic!()
                },
            }
        }
    }

    pub fn root(&self) -> H256 {
        self.root
    }

    pub fn state_of<'a>(&'a self, root: H256) -> FixedSecureTrie<<D as Database<'a>>::Guard, Address, Account> {
        self.database.create_fixed_secure_trie::<Address, Account>(root)
    }

    pub fn state<'a>(&'a self) -> FixedSecureTrie<<D as Database<'a>>::Guard, Address, Account> {
        self.state_of(self.root())
    }

    pub fn storage_state_of<'a>(&'a self, root: H256) -> FixedSecureTrie<<D as Database<'a>>::Guard, H256, M256> {
        self.database.create_fixed_secure_trie::<H256, M256>(root)
    }

    pub fn storage_state<'a>(&'a self, address: Address) -> Option<FixedSecureTrie<<D as Database<'a>>::Guard, H256, M256>> {
        let state = self.state();
        let account = state.get(&address);

        match account {
            Some(account) => {
                Some(self.storage_state_of(account.storage_root))
            },
            None => None,
        }
    }
}

pub type MemoryStateful<'a> = Stateful<'a, MemoryDatabase>;
<|endoftext|>


<commit_before>// Copyright (C) 2019 Sebastian Dröge <sebastian@centricular.com>
// Copyright (C) 2019 Jordan Petridis <jordan@centricular.com>
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Library General Public
// License as published by the Free Software Foundation; either
// version 2 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Library General Public License for more details.
//
// You should have received a copy of the GNU Library General Public
// License along with this library; if not, write to the
// Free Software Foundation, Inc., 51 Franklin Street, Suite 500,
// Boston, MA 02110-1335, USA.

use glib;
use glib::prelude::*;
use glib::subclass;
use glib::subclass::prelude::*;
use gst;
use gst::prelude::*;
use gst::subclass::prelude::*;
use gst_video::{self, ValidVideoTimeCode};

use std::io::Write;
use std::sync::Mutex;

lazy_static! {
    static ref CAT: gst::DebugCategory = {
        gst::DebugCategory::new(
            "sccenc",
            gst::DebugColorFlags::empty(),
            "Scc Encoder Element",
        )
    };
}

#[derive(Debug)]
struct State {
    need_headers: bool,
    expected_timecode: Option<ValidVideoTimeCode>,
    internal_buffer: Vec<gst::Buffer>,
    framerate: Option<gst::Fraction>,
}

impl Default for State {
    fn default() -> Self {
        Self {
            need_headers: true,
            expected_timecode: None,
            internal_buffer: Vec::with_capacity(64),
            framerate: None,
        }
    }
}

impl State {
    // Write the header to the buffer and set need_headers to false
    fn generate_headers(&mut self, buffer: &mut Vec<u8>) {
        assert!(buffer.is_empty());
        self.need_headers = false;
        buffer.extend_from_slice(b"Scenarist_SCC V1.0\r\n\r\n");
    }

    fn encode_payload(outbuf: &mut Vec<u8>, slice: &[u8]) {
        write!(outbuf, "{:02x}{:02x}", slice[0], slice[1]).unwrap();
    }

    fn generate_caption(
        &mut self,
        element: &gst::Element,
        buffer: gst::Buffer,
    ) -> Result<Option<gst::Buffer>, gst::FlowError> {
        // Arbitrary number that was chosen to keep in order
        // to batch pushes of smaller buffers
        const MAXIMUM_PACKETES_PER_LINE: usize = 16;

        assert!(self.internal_buffer.len() < MAXIMUM_PACKETES_PER_LINE);

        if buffer.get_size() != 2 {
            gst_element_error!(
                element,
                gst::StreamError::Format,
                ["Wrongly sized CEA608 packet: {}", buffer.get_size()]
            );

            return Err(gst::FlowError::Error);
        };

        let mut timecode = buffer
            .get_meta::<gst_video::VideoTimeCodeMeta>()
            .ok_or_else(|| {
                gst_element_error!(
                    element,
                    gst::StreamError::Format,
                    ["Stream with timecodes on each buffer required"]
                );

                // If we neeed to skip a buffer, increment the frame if it exists
                // to avoid getting out of sync
                if let Some(ref mut timecode) = self.expected_timecode {
                    timecode.increment_frame();
                }

                gst::FlowError::Error
            })?
            .get_tc();

        if self.expected_timecode.is_none() {
            self.expected_timecode = Some(timecode.clone());
        }

        // if the timecode is different from the expected one,
        // flush the previous line into the buffer, and push
        // the new packet to the, now empty, internal buffer
        if Some(&timecode) != self.expected_timecode.as_ref() {
            let res = self.write_line(element);
            assert!(self.internal_buffer.is_empty());
            self.internal_buffer.push(buffer);

            timecode.increment_frame();
            self.expected_timecode = Some(timecode);
            return res;
        } else if let Some(ref mut timecode) = self.expected_timecode {
            timecode.increment_frame();
        }

        self.internal_buffer.push(buffer);

        if self.internal_buffer.len() == MAXIMUM_PACKETES_PER_LINE {
            return self.write_line(element);
        }

        Ok(None)
    }

    // Flush the internal buffers into a line
    fn write_line(
        &mut self,
        element: &gst::Element,
    ) -> Result<Option<gst::Buffer>, gst::FlowError> {
        let mut outbuf = Vec::new();
        let mut line_start = true;

        assert!(!self.internal_buffer.is_empty());

        if self.need_headers {
            self.generate_headers(&mut outbuf);
        }

        let first_buf = self.internal_buffer.first().unwrap();
        for buffer in self.internal_buffer.iter() {
            let map = buffer.map_readable().ok_or_else(|| {
                gst_element_error!(
                    element,
                    gst::StreamError::Format,
                    ["Failed to map buffer readable"]
                );

                gst::FlowError::Error
            })?;

            // If its the first packet in the line, write the timecode first
            // else, separate the packets with a space
            if line_start {
                let timecode = buffer
                    .get_meta::<gst_video::VideoTimeCodeMeta>()
                    // Checked already before the buffer has been pushed to the
                    // internal_buffer
                    .expect("Buffer without timecode")
                    .get_tc();

                let _ = write!(outbuf, "{}\t", timecode);
                line_start = false;
            } else {
                outbuf.push(b' ');
            }

            Self::encode_payload(&mut outbuf, &*map);
        }

        outbuf.extend_from_slice(b"\r\n\r\n".as_ref());

        let buffer = {
            let mut buffer = gst::Buffer::from_mut_slice(outbuf);
            let buf_mut = buffer.get_mut().unwrap();

            // Something is seriously wrong else
            assert!(self.framerate.is_some());
            let framerate = self.framerate.unwrap();

            let dur = gst::SECOND
                .mul_div_floor(
                    self.internal_buffer.len() as u64 * *framerate.denom() as u64,
                    *framerate.numer() as u64,
                )
                .unwrap_or(gst::CLOCK_TIME_NONE);
            buf_mut.set_duration(dur);

            // Copy the metadata of the first buffer
            first_buf
                .copy_into(buf_mut, *gst::BUFFER_COPY_METADATA, 0, None)
                .expect("Failed to copy buffer metadata");
            buf_mut.set_pts(first_buf.get_pts());
            buffer
        };

        // Clear the internal buffer
        drop(first_buf);
        self.internal_buffer.clear();

        Ok(Some(buffer))
    }
}

struct SccEnc {
    srcpad: gst::Pad,
    sinkpad: gst::Pad,
    state: Mutex<State>,
}

impl SccEnc {
    fn set_pad_functions(sinkpad: &gst::Pad, srcpad: &gst::Pad) {
        sinkpad.set_chain_function(|pad, parent, buffer| {
            SccEnc::catch_panic_pad_function(
                parent,
                || Err(gst::FlowError::Error),
                |enc, element| enc.sink_chain(pad, element, buffer),
            )
        });
        sinkpad.set_event_function(|pad, parent, event| {
            SccEnc::catch_panic_pad_function(
                parent,
                || false,
                |enc, element| enc.sink_event(pad, element, event),
            )
        });

        srcpad.set_event_function(|pad, parent, event| {
            SccEnc::catch_panic_pad_function(
                parent,
                || false,
                |enc, element| enc.src_event(pad, element, event),
            )
        });
        srcpad.set_query_function(|pad, parent, query| {
            SccEnc::catch_panic_pad_function(
                parent,
                || false,
                |enc, element| enc.src_query(pad, element, query),
            )
        });
    }

    fn sink_chain(
        &self,
        pad: &gst::Pad,
        element: &gst::Element,
        buffer: gst::Buffer,
    ) -> Result<gst::FlowSuccess, gst::FlowError> {
        gst_log!(CAT, obj: pad, "Handling buffer {:?}", buffer);

        let mut state = self.state.lock().unwrap();
        let res = state.generate_caption(element, buffer)?;

        if let Some(outbuf) = res {
            gst_trace!(CAT, obj: pad, "Pushing buffer {:?} to the pad", &outbuf);

            drop(state);
            self.srcpad.push(outbuf)?;
        };

        Ok(gst::FlowSuccess::Ok)
    }

    fn sink_event(&self, pad: &gst::Pad, element: &gst::Element, event: gst::Event) -> bool {
        use gst::EventView;

        gst_log!(CAT, obj: pad, "Handling event {:?}", event);

        match event.view() {
            EventView::Caps(ev) => {
                let caps = ev.get_caps();
                let s = caps.get_structure(0).unwrap();
                let framerate = s.get::<gst::Fraction>("framerate");
                if framerate.is_none() {
                    gst_error!(CAT, obj: pad, "Caps without framerate");
                    return false;
                };

                let mut state = self.state.lock().unwrap();
                state.framerate = framerate;

                // We send our own caps downstream
                let caps = gst::Caps::builder("application/x-scc").build();
                self.srcpad.push_event(gst::Event::new_caps(&caps).build())
            }
            EventView::Eos(_) => {
                let mut state = self.state.lock().unwrap();

                let outbuf = state.write_line(element);

                if let Ok(Some(buffer)) = outbuf {
                    gst_trace!(CAT, obj: pad, "Pushing buffer {:?} to the pad", &buffer);

                    drop(state);
                    if self.srcpad.push(buffer).is_err() {
                        gst_error!(CAT, obj: pad, "Failed to push buffer to the pad");
                        return false;
                    }
                } else {
                    gst_error!(CAT, obj: pad, "Failed to write a line after EOS");
                    return false;
                }
                pad.event_default(element, event)
            }
            _ => pad.event_default(element, event),
        }
    }

    fn src_event(&self, pad: &gst::Pad, element: &gst::Element, event: gst::Event) -> bool {
        use gst::EventView;

        gst_log!(CAT, obj: pad, "Handling event {:?}", event);
        match event.view() {
            EventView::Seek(_) => {
                gst_log!(CAT, obj: pad, "Dropping seek event");
                false
            }
            _ => pad.event_default(element, event),
        }
    }

    fn src_query(&self, pad: &gst::Pad, element: &gst::Element, query: &mut gst::QueryRef) -> bool {
        use gst::QueryView;

        gst_log!(CAT, obj: pad, "Handling query {:?}", query);

        match query.view_mut() {
            QueryView::Seeking(mut q) => {
                // We don't support any seeking at all
                let fmt = q.get_format();
                q.set(
                    false,
                    gst::GenericFormattedValue::Other(fmt, -1),
                    gst::GenericFormattedValue::Other(fmt, -1),
                );
                true
            }
            _ => pad.query_default(element, query),
        }
    }
}

impl ObjectSubclass for SccEnc {
    const NAME: &'static str = "RsSccEnc";
    type ParentType = gst::Element;
    type Instance = gst::subclass::ElementInstanceStruct<Self>;
    type Class = subclass::simple::ClassStruct<Self>;

    glib_object_subclass!();

    fn new_with_class(klass: &subclass::simple::ClassStruct<Self>) -> Self {
        let templ = klass.get_pad_template("sink").unwrap();
        let sinkpad = gst::Pad::new_from_template(&templ, "sink");
        let templ = klass.get_pad_template("src").unwrap();
        let srcpad = gst::Pad::new_from_template(&templ, "src");

        SccEnc::set_pad_functions(&sinkpad, &srcpad);

        Self {
            srcpad,
            sinkpad,
            state: Mutex::new(State::default()),
        }
    }

    fn class_init(klass: &mut subclass::simple::ClassStruct<Self>) {
        klass.set_metadata(
            "Scc Encoder",
            "Encoder/ClosedCaption",
            "Encodes SCC Closed Caption Files",
            "Sebastian Dröge <sebastian@centricular.com>, Jordan Petridis <jordan@centricular.com>",
        );

        let framerates =
            gst::List::new(&[&gst::Fraction::new(30000, 1001), &gst::Fraction::new(30, 1)]);
        let caps = gst::Caps::builder("closedcaption/x-cea-608")
            .field("format", &"raw")
            .field("framerate", &framerates)
            .build();
        let sink_pad_template = gst::PadTemplate::new(
            "sink",
            gst::PadDirection::Sink,
            gst::PadPresence::Always,
            &caps,
        )
        .unwrap();
        klass.add_pad_template(sink_pad_template);

        let caps = gst::Caps::builder("application/x-scc").build();
        let src_pad_template = gst::PadTemplate::new(
            "src",
            gst::PadDirection::Src,
            gst::PadPresence::Always,
            &caps,
        )
        .unwrap();
        klass.add_pad_template(src_pad_template);
    }
}

impl ObjectImpl for SccEnc {
    glib_object_impl!();

    fn constructed(&self, obj: &glib::Object) {
        self.parent_constructed(obj);

        let element = obj.downcast_ref::<gst::Element>().unwrap();
        element.add_pad(&self.sinkpad).unwrap();
        element.add_pad(&self.srcpad).unwrap();
    }
}

impl ElementImpl for SccEnc {
    fn change_state(
        &self,
        element: &gst::Element,
        transition: gst::StateChange,
    ) -> Result<gst::StateChangeSuccess, gst::StateChangeError> {
        gst_trace!(CAT, obj: element, "Changing state {:?}", transition);

        match transition {
            gst::StateChange::ReadyToPaused | gst::StateChange::PausedToReady => {
                // Reset the whole state
                let mut state = self.state.lock().unwrap();
                *state = State::default();
            }
            _ => (),
        }

        self.parent_change_state(element, transition)
    }
}

pub fn register(plugin: &gst::Plugin) -> Result<(), glib::BoolError> {
    gst::Element::register(plugin, "sccenc", 0, SccEnc::get_type())
}
<commit_msg>scc_enc: Don't assert if we're trying to write a line without anything buffered<commit_after>// Copyright (C) 2019 Sebastian Dröge <sebastian@centricular.com>
// Copyright (C) 2019 Jordan Petridis <jordan@centricular.com>
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Library General Public
// License as published by the Free Software Foundation; either
// version 2 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Library General Public License for more details.
//
// You should have received a copy of the GNU Library General Public
// License along with this library; if not, write to the
// Free Software Foundation, Inc., 51 Franklin Street, Suite 500,
// Boston, MA 02110-1335, USA.

use glib;
use glib::prelude::*;
use glib::subclass;
use glib::subclass::prelude::*;
use gst;
use gst::prelude::*;
use gst::subclass::prelude::*;
use gst_video::{self, ValidVideoTimeCode};

use std::io::Write;
use std::sync::Mutex;

lazy_static! {
    static ref CAT: gst::DebugCategory = {
        gst::DebugCategory::new(
            "sccenc",
            gst::DebugColorFlags::empty(),
            "Scc Encoder Element",
        )
    };
}

#[derive(Debug)]
struct State {
    need_headers: bool,
    expected_timecode: Option<ValidVideoTimeCode>,
    internal_buffer: Vec<gst::Buffer>,
    framerate: Option<gst::Fraction>,
}

impl Default for State {
    fn default() -> Self {
        Self {
            need_headers: true,
            expected_timecode: None,
            internal_buffer: Vec::with_capacity(64),
            framerate: None,
        }
    }
}

impl State {
    // Write the header to the buffer and set need_headers to false
    fn generate_headers(&mut self, buffer: &mut Vec<u8>) {
        assert!(buffer.is_empty());
        self.need_headers = false;
        buffer.extend_from_slice(b"Scenarist_SCC V1.0\r\n\r\n");
    }

    fn encode_payload(outbuf: &mut Vec<u8>, slice: &[u8]) {
        write!(outbuf, "{:02x}{:02x}", slice[0], slice[1]).unwrap();
    }

    fn generate_caption(
        &mut self,
        element: &gst::Element,
        buffer: gst::Buffer,
    ) -> Result<Option<gst::Buffer>, gst::FlowError> {
        // Arbitrary number that was chosen to keep in order
        // to batch pushes of smaller buffers
        const MAXIMUM_PACKETES_PER_LINE: usize = 16;

        assert!(self.internal_buffer.len() < MAXIMUM_PACKETES_PER_LINE);

        if buffer.get_size() != 2 {
            gst_element_error!(
                element,
                gst::StreamError::Format,
                ["Wrongly sized CEA608 packet: {}", buffer.get_size()]
            );

            return Err(gst::FlowError::Error);
        };

        let mut timecode = buffer
            .get_meta::<gst_video::VideoTimeCodeMeta>()
            .ok_or_else(|| {
                gst_element_error!(
                    element,
                    gst::StreamError::Format,
                    ["Stream with timecodes on each buffer required"]
                );

                // If we neeed to skip a buffer, increment the frame if it exists
                // to avoid getting out of sync
                if let Some(ref mut timecode) = self.expected_timecode {
                    timecode.increment_frame();
                }

                gst::FlowError::Error
            })?
            .get_tc();

        if self.expected_timecode.is_none() {
            self.expected_timecode = Some(timecode.clone());
        }

        // if the timecode is different from the expected one,
        // flush the previous line into the buffer, and push
        // the new packet to the, now empty, internal buffer
        if Some(&timecode) != self.expected_timecode.as_ref() {
            let outbuf = self.write_line(element)?;

            assert!(self.internal_buffer.is_empty());
            self.internal_buffer.push(buffer);

            timecode.increment_frame();
            self.expected_timecode = Some(timecode);

            return Ok(outbuf);
        } else if let Some(ref mut timecode) = self.expected_timecode {
            timecode.increment_frame();
        }

        self.internal_buffer.push(buffer);

        if self.internal_buffer.len() == MAXIMUM_PACKETES_PER_LINE {
            return self.write_line(element);
        }

        Ok(None)
    }

    // Flush the internal buffers into a line
    fn write_line(
        &mut self,
        element: &gst::Element,
    ) -> Result<Option<gst::Buffer>, gst::FlowError> {
        let mut outbuf = Vec::new();
        let mut line_start = true;

        if self.internal_buffer.is_empty() {
            return Ok(None);
        }

        if self.need_headers {
            self.generate_headers(&mut outbuf);
        }

        let first_buf = self.internal_buffer.first().unwrap();
        for buffer in self.internal_buffer.iter() {
            let map = buffer.map_readable().ok_or_else(|| {
                gst_element_error!(
                    element,
                    gst::StreamError::Format,
                    ["Failed to map buffer readable"]
                );

                gst::FlowError::Error
            })?;

            // If its the first packet in the line, write the timecode first
            // else, separate the packets with a space
            if line_start {
                let timecode = buffer
                    .get_meta::<gst_video::VideoTimeCodeMeta>()
                    // Checked already before the buffer has been pushed to the
                    // internal_buffer
                    .expect("Buffer without timecode")
                    .get_tc();

                let _ = write!(outbuf, "{}\t", timecode);
                line_start = false;
            } else {
                outbuf.push(b' ');
            }

            Self::encode_payload(&mut outbuf, &*map);
        }

        outbuf.extend_from_slice(b"\r\n\r\n".as_ref());

        let buffer = {
            let mut buffer = gst::Buffer::from_mut_slice(outbuf);
            let buf_mut = buffer.get_mut().unwrap();

            // Something is seriously wrong else
            assert!(self.framerate.is_some());
            let framerate = self.framerate.unwrap();

            let dur = gst::SECOND
                .mul_div_floor(
                    self.internal_buffer.len() as u64 * *framerate.denom() as u64,
                    *framerate.numer() as u64,
                )
                .unwrap_or(gst::CLOCK_TIME_NONE);
            buf_mut.set_duration(dur);

            // Copy the metadata of the first buffer
            first_buf
                .copy_into(buf_mut, *gst::BUFFER_COPY_METADATA, 0, None)
                .expect("Failed to copy buffer metadata");
            buf_mut.set_pts(first_buf.get_pts());
            buffer
        };

        // Clear the internal buffer
        drop(first_buf);
        self.internal_buffer.clear();

        Ok(Some(buffer))
    }
}

struct SccEnc {
    srcpad: gst::Pad,
    sinkpad: gst::Pad,
    state: Mutex<State>,
}

impl SccEnc {
    fn set_pad_functions(sinkpad: &gst::Pad, srcpad: &gst::Pad) {
        sinkpad.set_chain_function(|pad, parent, buffer| {
            SccEnc::catch_panic_pad_function(
                parent,
                || Err(gst::FlowError::Error),
                |enc, element| enc.sink_chain(pad, element, buffer),
            )
        });
        sinkpad.set_event_function(|pad, parent, event| {
            SccEnc::catch_panic_pad_function(
                parent,
                || false,
                |enc, element| enc.sink_event(pad, element, event),
            )
        });

        srcpad.set_event_function(|pad, parent, event| {
            SccEnc::catch_panic_pad_function(
                parent,
                || false,
                |enc, element| enc.src_event(pad, element, event),
            )
        });
        srcpad.set_query_function(|pad, parent, query| {
            SccEnc::catch_panic_pad_function(
                parent,
                || false,
                |enc, element| enc.src_query(pad, element, query),
            )
        });
    }

    fn sink_chain(
        &self,
        pad: &gst::Pad,
        element: &gst::Element,
        buffer: gst::Buffer,
    ) -> Result<gst::FlowSuccess, gst::FlowError> {
        gst_log!(CAT, obj: pad, "Handling buffer {:?}", buffer);

        let mut state = self.state.lock().unwrap();
        let res = state.generate_caption(element, buffer)?;

        if let Some(outbuf) = res {
            gst_trace!(CAT, obj: pad, "Pushing buffer {:?} to the pad", &outbuf);

            drop(state);
            self.srcpad.push(outbuf)?;
        };

        Ok(gst::FlowSuccess::Ok)
    }

    fn sink_event(&self, pad: &gst::Pad, element: &gst::Element, event: gst::Event) -> bool {
        use gst::EventView;

        gst_log!(CAT, obj: pad, "Handling event {:?}", event);

        match event.view() {
            EventView::Caps(ev) => {
                let caps = ev.get_caps();
                let s = caps.get_structure(0).unwrap();
                let framerate = s.get::<gst::Fraction>("framerate");
                if framerate.is_none() {
                    gst_error!(CAT, obj: pad, "Caps without framerate");
                    return false;
                };

                let mut state = self.state.lock().unwrap();
                state.framerate = framerate;

                // We send our own caps downstream
                let caps = gst::Caps::builder("application/x-scc").build();
                self.srcpad.push_event(gst::Event::new_caps(&caps).build())
            }
            EventView::Eos(_) => {
                let mut state = self.state.lock().unwrap();

                let outbuf = state.write_line(element);

                if let Ok(Some(buffer)) = outbuf {
                    gst_trace!(CAT, obj: pad, "Pushing buffer {:?} to the pad", &buffer);

                    drop(state);
                    if self.srcpad.push(buffer).is_err() {
                        gst_error!(CAT, obj: pad, "Failed to push buffer to the pad");
                        return false;
                    }
                } else if let Err(err) = outbuf {
                    gst_error!(CAT, obj: pad, "Failed to write a line after EOS: {:?}", err);
                    return false;
                }
                pad.event_default(element, event)
            }
            _ => pad.event_default(element, event),
        }
    }

    fn src_event(&self, pad: &gst::Pad, element: &gst::Element, event: gst::Event) -> bool {
        use gst::EventView;

        gst_log!(CAT, obj: pad, "Handling event {:?}", event);
        match event.view() {
            EventView::Seek(_) => {
                gst_log!(CAT, obj: pad, "Dropping seek event");
                false
            }
            _ => pad.event_default(element, event),
        }
    }

    fn src_query(&self, pad: &gst::Pad, element: &gst::Element, query: &mut gst::QueryRef) -> bool {
        use gst::QueryView;

        gst_log!(CAT, obj: pad, "Handling query {:?}", query);

        match query.view_mut() {
            QueryView::Seeking(mut q) => {
                // We don't support any seeking at all
                let fmt = q.get_format();
                q.set(
                    false,
                    gst::GenericFormattedValue::Other(fmt, -1),
                    gst::GenericFormattedValue::Other(fmt, -1),
                );
                true
            }
            _ => pad.query_default(element, query),
        }
    }
}

impl ObjectSubclass for SccEnc {
    const NAME: &'static str = "RsSccEnc";
    type ParentType = gst::Element;
    type Instance = gst::subclass::ElementInstanceStruct<Self>;
    type Class = subclass::simple::ClassStruct<Self>;

    glib_object_subclass!();

    fn new_with_class(klass: &subclass::simple::ClassStruct<Self>) -> Self {
        let templ = klass.get_pad_template("sink").unwrap();
        let sinkpad = gst::Pad::new_from_template(&templ, "sink");
        let templ = klass.get_pad_template("src").unwrap();
        let srcpad = gst::Pad::new_from_template(&templ, "src");

        SccEnc::set_pad_functions(&sinkpad, &srcpad);

        Self {
            srcpad,
            sinkpad,
            state: Mutex::new(State::default()),
        }
    }

    fn class_init(klass: &mut subclass::simple::ClassStruct<Self>) {
        klass.set_metadata(
            "Scc Encoder",
            "Encoder/ClosedCaption",
            "Encodes SCC Closed Caption Files",
            "Sebastian Dröge <sebastian@centricular.com>, Jordan Petridis <jordan@centricular.com>",
        );

        let framerates =
            gst::List::new(&[&gst::Fraction::new(30000, 1001), &gst::Fraction::new(30, 1)]);
        let caps = gst::Caps::builder("closedcaption/x-cea-608")
            .field("format", &"raw")
            .field("framerate", &framerates)
            .build();
        let sink_pad_template = gst::PadTemplate::new(
            "sink",
            gst::PadDirection::Sink,
            gst::PadPresence::Always,
            &caps,
        )
        .unwrap();
        klass.add_pad_template(sink_pad_template);

        let caps = gst::Caps::builder("application/x-scc").build();
        let src_pad_template = gst::PadTemplate::new(
            "src",
            gst::PadDirection::Src,
            gst::PadPresence::Always,
            &caps,
        )
        .unwrap();
        klass.add_pad_template(src_pad_template);
    }
}

impl ObjectImpl for SccEnc {
    glib_object_impl!();

    fn constructed(&self, obj: &glib::Object) {
        self.parent_constructed(obj);

        let element = obj.downcast_ref::<gst::Element>().unwrap();
        element.add_pad(&self.sinkpad).unwrap();
        element.add_pad(&self.srcpad).unwrap();
    }
}

impl ElementImpl for SccEnc {
    fn change_state(
        &self,
        element: &gst::Element,
        transition: gst::StateChange,
    ) -> Result<gst::StateChangeSuccess, gst::StateChangeError> {
        gst_trace!(CAT, obj: element, "Changing state {:?}", transition);

        match transition {
            gst::StateChange::ReadyToPaused | gst::StateChange::PausedToReady => {
                // Reset the whole state
                let mut state = self.state.lock().unwrap();
                *state = State::default();
            }
            _ => (),
        }

        self.parent_change_state(element, transition)
    }
}

pub fn register(plugin: &gst::Plugin) -> Result<(), glib::BoolError> {
    gst::Element::register(plugin, "sccenc", 0, SccEnc::get_type())
}
<|endoftext|>