<commit_before>use join_to_string::join;

use ra_syntax::{
    algo::{find_covering_node, find_leaf_at_offset},
    ast::{self, AstNode, AttrsOwner, NameOwner, TypeParamsOwner},
    Direction, File,
    SyntaxKind::{COMMA, WHITESPACE},
    SyntaxNodeRef, TextRange, TextUnit,
};

use crate::{find_node_at_offset, Edit, EditBuilder};

#[derive(Debug)]
pub struct LocalEdit {
    pub edit: Edit,
    pub cursor_position: Option<TextUnit>,
}

pub fn flip_comma<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce() -> LocalEdit + 'a> {
    let syntax = file.syntax();

    let comma = find_leaf_at_offset(syntax, offset).find(|leaf| leaf.kind() == COMMA)?;
    let prev = non_trivia_sibling(comma, Direction::Prev)?;
    let next = non_trivia_sibling(comma, Direction::Next)?;
    Some(move || {
        let mut edit = EditBuilder::new();
        edit.replace(prev.range(), next.text().to_string());
        edit.replace(next.range(), prev.text().to_string());
        LocalEdit {
            edit: edit.finish(),
            cursor_position: None,
        }
    })
}

pub fn add_derive<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce() -> LocalEdit + 'a> {
    let nominal = find_node_at_offset::<ast::NominalDef>(file.syntax(), offset)?;
    Some(move || {
        let derive_attr = nominal
            .attrs()
            .filter_map(|x| x.as_call())
            .filter(|(name, _arg)| name == "derive")
            .map(|(_name, arg)| arg)
            .next();
        let mut edit = EditBuilder::new();
        let offset = match derive_attr {
            None => {
                let node_start = nominal.syntax().range().start();
                edit.insert(node_start, "#[derive()]\n".to_string());
                node_start + TextUnit::of_str("#[derive(")
            }
            Some(tt) => tt.syntax().range().end() - TextUnit::of_char(')'),
        };
        LocalEdit {
            edit: edit.finish(),
            cursor_position: Some(offset),
        }
    })
}

pub fn add_impl<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce() -> LocalEdit + 'a> {
    let nominal = find_node_at_offset::<ast::NominalDef>(file.syntax(), offset)?;
    let name = nominal.name()?;

    Some(move || {
        let type_params = nominal.type_param_list();
        let mut edit = EditBuilder::new();
        let start_offset = nominal.syntax().range().end();
        let mut buf = String::new();
        buf.push_str("\n\nimpl");
        if let Some(type_params) = type_params {
            type_params.syntax().text().push_to(&mut buf);
        }
        buf.push_str(" ");
        buf.push_str(name.text().as_str());
        if let Some(type_params) = type_params {
            let lifetime_params = type_params
                .lifetime_params()
                .filter_map(|it| it.lifetime())
                .map(|it| it.text());
            let type_params = type_params
                .type_params()
                .filter_map(|it| it.name())
                .map(|it| it.text());
            join(lifetime_params.chain(type_params))
                .surround_with("<", ">")
                .to_buf(&mut buf);
        }
        buf.push_str(" {\n");
        let offset = start_offset + TextUnit::of_str(&buf);
        buf.push_str("\n}");
        edit.insert(start_offset, buf);
        LocalEdit {
            edit: edit.finish(),
            cursor_position: Some(offset),
        }
    })
}

pub fn introduce_variable<'a>(
    file: &'a File,
    range: TextRange,
) -> Option<impl FnOnce() -> LocalEdit + 'a> {
    let node = find_covering_node(file.syntax(), range);
    let expr = node.ancestors().filter_map(ast::Expr::cast).next()?;
    let anchor_stmt = expr
        .syntax()
        .ancestors()
        .filter_map(ast::Stmt::cast)
        .next()?;
    let indent = anchor_stmt.syntax().prev_sibling()?;
    if indent.kind() != WHITESPACE {
        return None;
    }
    Some(move || {
        let mut buf = String::new();
        let mut edit = EditBuilder::new();

        buf.push_str("let var_name = ");
        expr.syntax().text().push_to(&mut buf);
        if expr.syntax().range().start() == anchor_stmt.syntax().range().start() {
            edit.replace(expr.syntax().range(), buf);
        } else {
            buf.push_str(";");
            indent.text().push_to(&mut buf);
            edit.replace(expr.syntax().range(), "var_name".to_string());
            edit.insert(anchor_stmt.syntax().range().start(), buf);
        }
        let cursor_position = anchor_stmt.syntax().range().start() + TextUnit::of_str("let ");
        LocalEdit {
            edit: edit.finish(),
            cursor_position: Some(cursor_position),
        }
    })
}

fn non_trivia_sibling(node: SyntaxNodeRef, direction: Direction) -> Option<SyntaxNodeRef> {
    node.siblings(direction)
        .skip(1)
        .find(|node| !node.kind().is_trivia())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_utils::{check_action, check_action_range};

    #[test]
    fn test_swap_comma() {
        check_action(
            "fn foo(x: i32,<|> y: Result<(), ()>) {}",
            "fn foo(y: Result<(), ()>,<|> x: i32) {}",
            |file, off| flip_comma(file, off).map(|f| f()),
        )
    }

    #[test]
    fn test_add_derive() {
        check_action(
            "struct Foo { a: i32, <|>}",
            "#[derive(<|>)]\nstruct Foo { a: i32, }",
            |file, off| add_derive(file, off).map(|f| f()),
        );
        check_action(
            "struct Foo { <|> a: i32, }",
            "#[derive(<|>)]\nstruct Foo {  a: i32, }",
            |file, off| add_derive(file, off).map(|f| f()),
        );
        check_action(
            "#[derive(Clone)]\nstruct Foo { a: i32<|>, }",
            "#[derive(Clone<|>)]\nstruct Foo { a: i32, }",
            |file, off| add_derive(file, off).map(|f| f()),
        );
    }

    #[test]
    fn test_add_impl() {
        check_action(
            "struct Foo {<|>}\n",
            "struct Foo {}\n\nimpl Foo {\n<|>\n}\n",
            |file, off| add_impl(file, off).map(|f| f()),
        );
        check_action(
            "struct Foo<T: Clone> {<|>}",
            "struct Foo<T: Clone> {}\n\nimpl<T: Clone> Foo<T> {\n<|>\n}",
            |file, off| add_impl(file, off).map(|f| f()),
        );
        check_action(
            "struct Foo<'a, T: Foo<'a>> {<|>}",
            "struct Foo<'a, T: Foo<'a>> {}\n\nimpl<'a, T: Foo<'a>> Foo<'a, T> {\n<|>\n}",
            |file, off| add_impl(file, off).map(|f| f()),
        );
    }

    #[test]
    fn test_intrdoduce_var_simple() {
        check_action_range(
            "
fn foo() {
    foo(<|>1 + 1<|>);
}",
            "
fn foo() {
    let <|>var_name = 1 + 1;
    foo(var_name);
}",
            |file, range| introduce_variable(file, range).map(|f| f()),
        );
    }
    #[test]
    fn test_intrdoduce_var_expr_stmt() {
        check_action_range(
            "
fn foo() {
    <|>1 + 1<|>;
}",
            "
fn foo() {
    let <|>var_name = 1 + 1;
}",
            |file, range| introduce_variable(file, range).map(|f| f()),
        );
    }

}
<commit_msg>Introduce var works with tail exprs<commit_after>use join_to_string::join;

use ra_syntax::{
    algo::{find_covering_node, find_leaf_at_offset},
    ast::{self, AstNode, AttrsOwner, NameOwner, TypeParamsOwner},
    Direction, File,
    SyntaxKind::{COMMA, WHITESPACE},
    SyntaxNodeRef, TextRange, TextUnit,
};

use crate::{find_node_at_offset, Edit, EditBuilder};

#[derive(Debug)]
pub struct LocalEdit {
    pub edit: Edit,
    pub cursor_position: Option<TextUnit>,
}

pub fn flip_comma<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce() -> LocalEdit + 'a> {
    let syntax = file.syntax();

    let comma = find_leaf_at_offset(syntax, offset).find(|leaf| leaf.kind() == COMMA)?;
    let prev = non_trivia_sibling(comma, Direction::Prev)?;
    let next = non_trivia_sibling(comma, Direction::Next)?;
    Some(move || {
        let mut edit = EditBuilder::new();
        edit.replace(prev.range(), next.text().to_string());
        edit.replace(next.range(), prev.text().to_string());
        LocalEdit {
            edit: edit.finish(),
            cursor_position: None,
        }
    })
}

pub fn add_derive<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce() -> LocalEdit + 'a> {
    let nominal = find_node_at_offset::<ast::NominalDef>(file.syntax(), offset)?;
    Some(move || {
        let derive_attr = nominal
            .attrs()
            .filter_map(|x| x.as_call())
            .filter(|(name, _arg)| name == "derive")
            .map(|(_name, arg)| arg)
            .next();
        let mut edit = EditBuilder::new();
        let offset = match derive_attr {
            None => {
                let node_start = nominal.syntax().range().start();
                edit.insert(node_start, "#[derive()]\n".to_string());
                node_start + TextUnit::of_str("#[derive(")
            }
            Some(tt) => tt.syntax().range().end() - TextUnit::of_char(')'),
        };
        LocalEdit {
            edit: edit.finish(),
            cursor_position: Some(offset),
        }
    })
}

pub fn add_impl<'a>(file: &'a File, offset: TextUnit) -> Option<impl FnOnce() -> LocalEdit + 'a> {
    let nominal = find_node_at_offset::<ast::NominalDef>(file.syntax(), offset)?;
    let name = nominal.name()?;

    Some(move || {
        let type_params = nominal.type_param_list();
        let mut edit = EditBuilder::new();
        let start_offset = nominal.syntax().range().end();
        let mut buf = String::new();
        buf.push_str("\n\nimpl");
        if let Some(type_params) = type_params {
            type_params.syntax().text().push_to(&mut buf);
        }
        buf.push_str(" ");
        buf.push_str(name.text().as_str());
        if let Some(type_params) = type_params {
            let lifetime_params = type_params
                .lifetime_params()
                .filter_map(|it| it.lifetime())
                .map(|it| it.text());
            let type_params = type_params
                .type_params()
                .filter_map(|it| it.name())
                .map(|it| it.text());
            join(lifetime_params.chain(type_params))
                .surround_with("<", ">")
                .to_buf(&mut buf);
        }
        buf.push_str(" {\n");
        let offset = start_offset + TextUnit::of_str(&buf);
        buf.push_str("\n}");
        edit.insert(start_offset, buf);
        LocalEdit {
            edit: edit.finish(),
            cursor_position: Some(offset),
        }
    })
}

pub fn introduce_variable<'a>(
    file: &'a File,
    range: TextRange,
) -> Option<impl FnOnce() -> LocalEdit + 'a> {
    let node = find_covering_node(file.syntax(), range);
    let expr = node.ancestors().filter_map(ast::Expr::cast).next()?;

    let anchor_stmt = ahchor_stmt(expr)?;
    let indent = anchor_stmt.prev_sibling()?;
    if indent.kind() != WHITESPACE {
        return None;
    }
    return Some(move || {
        let mut buf = String::new();
        let mut edit = EditBuilder::new();

        buf.push_str("let var_name = ");
        expr.syntax().text().push_to(&mut buf);
        if expr.syntax().range().start() == anchor_stmt.range().start() {
            edit.replace(expr.syntax().range(), buf);
        } else {
            buf.push_str(";");
            indent.text().push_to(&mut buf);
            edit.replace(expr.syntax().range(), "var_name".to_string());
            edit.insert(anchor_stmt.range().start(), buf);
        }
        let cursor_position = anchor_stmt.range().start() + TextUnit::of_str("let ");
        LocalEdit {
            edit: edit.finish(),
            cursor_position: Some(cursor_position),
        }
    });

    /// Statement or last in the block expression, which will follow
    /// the freshly introduced var.
    fn ahchor_stmt(expr: ast::Expr) -> Option<SyntaxNodeRef> {
        expr.syntax().ancestors().find(|&node| {
            if ast::Stmt::cast(node).is_some() {
                return true;
            }
            if let Some(expr) = node
                .parent()
                .and_then(ast::Block::cast)
                .and_then(|it| it.expr())
            {
                if expr.syntax() == node {
                    return true;
                }
            }
            false
        })
    }
}

fn non_trivia_sibling(node: SyntaxNodeRef, direction: Direction) -> Option<SyntaxNodeRef> {
    node.siblings(direction)
        .skip(1)
        .find(|node| !node.kind().is_trivia())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_utils::{check_action, check_action_range};

    #[test]
    fn test_swap_comma() {
        check_action(
            "fn foo(x: i32,<|> y: Result<(), ()>) {}",
            "fn foo(y: Result<(), ()>,<|> x: i32) {}",
            |file, off| flip_comma(file, off).map(|f| f()),
        )
    }

    #[test]
    fn test_add_derive() {
        check_action(
            "struct Foo { a: i32, <|>}",
            "#[derive(<|>)]\nstruct Foo { a: i32, }",
            |file, off| add_derive(file, off).map(|f| f()),
        );
        check_action(
            "struct Foo { <|> a: i32, }",
            "#[derive(<|>)]\nstruct Foo {  a: i32, }",
            |file, off| add_derive(file, off).map(|f| f()),
        );
        check_action(
            "#[derive(Clone)]\nstruct Foo { a: i32<|>, }",
            "#[derive(Clone<|>)]\nstruct Foo { a: i32, }",
            |file, off| add_derive(file, off).map(|f| f()),
        );
    }

    #[test]
    fn test_add_impl() {
        check_action(
            "struct Foo {<|>}\n",
            "struct Foo {}\n\nimpl Foo {\n<|>\n}\n",
            |file, off| add_impl(file, off).map(|f| f()),
        );
        check_action(
            "struct Foo<T: Clone> {<|>}",
            "struct Foo<T: Clone> {}\n\nimpl<T: Clone> Foo<T> {\n<|>\n}",
            |file, off| add_impl(file, off).map(|f| f()),
        );
        check_action(
            "struct Foo<'a, T: Foo<'a>> {<|>}",
            "struct Foo<'a, T: Foo<'a>> {}\n\nimpl<'a, T: Foo<'a>> Foo<'a, T> {\n<|>\n}",
            |file, off| add_impl(file, off).map(|f| f()),
        );
    }

    #[test]
    fn test_intrdoduce_var_simple() {
        check_action_range(
            "
fn foo() {
    foo(<|>1 + 1<|>);
}",
            "
fn foo() {
    let <|>var_name = 1 + 1;
    foo(var_name);
}",
            |file, range| introduce_variable(file, range).map(|f| f()),
        );
    }

    #[test]
    fn test_intrdoduce_var_expr_stmt() {
        check_action_range(
            "
fn foo() {
    <|>1 + 1<|>;
}",
            "
fn foo() {
    let <|>var_name = 1 + 1;
}",
            |file, range| introduce_variable(file, range).map(|f| f()),
        );
    }

    #[test]
    fn test_intrdoduce_var_last_expr() {
        check_action_range(
            "
fn foo() {
    bar(<|>1 + 1<|>)
}",
            "
fn foo() {
    let <|>var_name = 1 + 1;
    bar(var_name)
}",
            |file, range| introduce_variable(file, range).map(|f| f()),
        );
    }

}
<|endoftext|>


<commit_before>//! FIXME: write short doc here

use std::sync::Arc;

use ra_syntax::ast;

use crate::{
    db::{AstDatabase, DefDatabase, HirDatabase},
    Const, Enum, EnumVariant, FieldSource, Function, HasSource, MacroDef, Module, Static, Struct,
    StructField, Trait, TypeAlias, Union,
};

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum DocDef {
    Module(Module),
    StructField(StructField),
    Struct(Struct),
    Enum(Enum),
    EnumVariant(EnumVariant),
    Static(Static),
    Const(Const),
    Function(Function),
    Union(Union),
    Trait(Trait),
    TypeAlias(TypeAlias),
    MacroDef(MacroDef),
}

impl_froms!(
    DocDef: Module,
    StructField,
    Struct,
    Enum,
    EnumVariant,
    Static,
    Const,
    Function,
    Union,
    Trait,
    TypeAlias,
    MacroDef
);

/// Holds documentation
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Documentation(Arc<str>);

impl Documentation {
    fn new(s: &str) -> Documentation {
        Documentation(s.into())
    }

    pub fn as_str(&self) -> &str {
        &*self.0
    }
}

impl Into<String> for Documentation {
    fn into(self) -> String {
        self.as_str().to_owned()
    }
}

pub trait Docs {
    fn docs(&self, db: &impl HirDatabase) -> Option<Documentation>;
}

pub(crate) fn docs_from_ast(node: &impl ast::DocCommentsOwner) -> Option<Documentation> {
    node.doc_comment_text().map(|it| Documentation::new(&it))
}

pub(crate) fn documentation_query(
    db: &(impl DefDatabase + AstDatabase),
    def: DocDef,
) -> Option<Documentation> {
    match def {
        DocDef::Module(it) => docs_from_ast(&it.declaration_source(db)?.ast),
        DocDef::StructField(it) => match it.source(db).ast {
            FieldSource::Named(named) => docs_from_ast(&named),
            FieldSource::Pos(..) => None,
        },
        DocDef::Struct(it) => docs_from_ast(&it.source(db).ast),
        DocDef::Enum(it) => docs_from_ast(&it.source(db).ast),
        DocDef::EnumVariant(it) => docs_from_ast(&it.source(db).ast),
        DocDef::Static(it) => docs_from_ast(&it.source(db).ast),
        DocDef::Const(it) => docs_from_ast(&it.source(db).ast),
        DocDef::Function(it) => docs_from_ast(&it.source(db).ast),
        DocDef::Union(it) => docs_from_ast(&it.source(db).ast),
        DocDef::Trait(it) => docs_from_ast(&it.source(db).ast),
        DocDef::TypeAlias(it) => docs_from_ast(&it.source(db).ast),
        DocDef::MacroDef(it) => docs_from_ast(&it.source(db).ast),
    }
}

impl<T: Into<DocDef> + Copy> Docs for T {
    fn docs(&self, db: &impl HirDatabase) -> Option<Documentation> {
        db.documentation((*self).into())
    }
}
<commit_msg>refactor<commit_after>//! FIXME: write short doc here

use std::sync::Arc;

use ra_syntax::ast;

use crate::{
    db::{AstDatabase, DefDatabase, HirDatabase},
    Adt, Const, Enum, EnumVariant, FieldSource, Function, HasSource, MacroDef, Module, Static,
    Struct, StructField, Trait, TypeAlias, Union,
};

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum DocDef {
    Module(Module),
    StructField(StructField),
    Adt(Adt),
    EnumVariant(EnumVariant),
    Static(Static),
    Const(Const),
    Function(Function),
    Trait(Trait),
    TypeAlias(TypeAlias),
    MacroDef(MacroDef),
}

impl_froms!(
    DocDef: Module,
    StructField,
    Adt(Struct, Enum, Union),
    EnumVariant,
    Static,
    Const,
    Function,
    Trait,
    TypeAlias,
    MacroDef
);

/// Holds documentation
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Documentation(Arc<str>);

impl Documentation {
    fn new(s: &str) -> Documentation {
        Documentation(s.into())
    }

    pub fn as_str(&self) -> &str {
        &*self.0
    }
}

impl Into<String> for Documentation {
    fn into(self) -> String {
        self.as_str().to_owned()
    }
}

pub trait Docs {
    fn docs(&self, db: &impl HirDatabase) -> Option<Documentation>;
}

pub(crate) fn docs_from_ast(node: &impl ast::DocCommentsOwner) -> Option<Documentation> {
    node.doc_comment_text().map(|it| Documentation::new(&it))
}

pub(crate) fn documentation_query(
    db: &(impl DefDatabase + AstDatabase),
    def: DocDef,
) -> Option<Documentation> {
    match def {
        DocDef::Module(it) => docs_from_ast(&it.declaration_source(db)?.ast),
        DocDef::StructField(it) => match it.source(db).ast {
            FieldSource::Named(named) => docs_from_ast(&named),
            FieldSource::Pos(..) => None,
        },
        DocDef::Adt(it) => match it {
            Adt::Struct(it) => docs_from_ast(&it.source(db).ast),
            Adt::Enum(it) => docs_from_ast(&it.source(db).ast),
            Adt::Union(it) => docs_from_ast(&it.source(db).ast),
        },
        DocDef::EnumVariant(it) => docs_from_ast(&it.source(db).ast),
        DocDef::Static(it) => docs_from_ast(&it.source(db).ast),
        DocDef::Const(it) => docs_from_ast(&it.source(db).ast),
        DocDef::Function(it) => docs_from_ast(&it.source(db).ast),
        DocDef::Trait(it) => docs_from_ast(&it.source(db).ast),
        DocDef::TypeAlias(it) => docs_from_ast(&it.source(db).ast),
        DocDef::MacroDef(it) => docs_from_ast(&it.source(db).ast),
    }
}

impl<T: Into<DocDef> + Copy> Docs for T {
    fn docs(&self, db: &impl HirDatabase) -> Option<Documentation> {
        db.documentation((*self).into())
    }
}
<|endoftext|>


<commit_before>// based on define_flags by lexs
#[macro_escape];

macro_rules! define_flags (
    (
        $name:ident {
            $($flag:ident = $value:expr),+
        }
    ) => {
        define_flags!($name: u32 { $($flag = $value),+ })
    };

    (
        $name:ident: $t:ty {
            $($flag:ident $(= $v:expr)*),*
        }
    ) => {
        #[packed]
        pub struct $name($t);

        impl core::ops::BitOr<$name, $name> for $name {
            #[inline(always)]
            fn bitor(&self, other: &$name) -> $name {
                match (self, other) {
                    (&$name(flags1), &$name(flags2)) => $name(flags1 | flags2)
                }
            }
        }

        impl core::ops::BitAnd<$name, bool> for $name {
            #[inline(always)]
            fn bitand(&self, other: &$name) -> bool {
                match (self, other) {
                    (&$name(flags1), &$name(flags2)) => flags1 & flags2 != 0
                }
            }
        }

        impl core::ops::Not<$name> for $name {
            #[inline(always)]
            fn not(&self) -> $name {
                match self {
                    &$name(flags) => $name(!flags)
                }
            }
        }

        define_flags_rec!($name, 1, $( $flag $(= $v)* ),+)
    };
)
macro_rules! define_flags_rec (
    // full list (original behavior)
    (
        $name:ident,
        $default:expr,
        $($flag:ident = $value:expr),*
    ) => (
        $( pub static $flag: $name = $name($value); )+
    );
    // ----------
    // only one default value
    (
        $name:ident,
        $default:expr,
        $flag:ident
    ) => (
        pub static $flag: $name = $name($default);
    );
    // only one value
    (
        $name:ident,
        $default:expr,
        $flag:ident = $value:expr
    ) => (
        pub static $flag: $name = $name($value);
    );
    // ----------
    // without value (default)
    (
        $name:ident,
        $default:expr,
        $flag:ident,
        $($f:ident $(= $v:expr)*),*
    ) => (
        pub static $flag: $name = $name($default);
        define_flags_rec!($name, $default << 1, $($f $(= $v)*),+)
    );
    // with value
    (
        $name:ident,
        $default:expr,
        $flag:ident = $value:expr,
        $($f:ident $(= $v:expr)*),*
    ) => (
        pub static $flag: $name = $name($value);
        define_flags_rec!($name, $value << 1, $( $f $(= $v)* ),+)
    );
)
<commit_msg>Implement get, map with define_flags<commit_after>// based on define_flags by lexs
#[macro_escape];

macro_rules! define_flags (
    (
        $name:ident {
            $($flag:ident = $value:expr),+
        }
    ) => {
        define_flags!($name: u32 { $($flag = $value),+ })
    };

    (
        $name:ident: $T:ty {
            $($flag:ident $(= $v:expr)*),*
        }
    ) => {
        #[packed]
        pub struct $name($T);

        impl $name {
            /// Return the value.
            pub fn get(self) -> $T {
                match self { $name(x) => x }
            }

            /// Maps by applying a function to a contained value.
            pub fn map(self, f: |$T| -> $T) -> $name {
                match self { $name(x) => $name(f(x)) }
            }
        }

        impl core::ops::BitOr<$name, $name> for $name {
            #[inline(always)]
            fn bitor(&self, other: &$name) -> $name {
                match (self, other) {
                    (&$name(flags1), &$name(flags2)) => $name(flags1 | flags2)
                }
            }
        }

        impl core::ops::BitAnd<$name, bool> for $name {
            #[inline(always)]
            fn bitand(&self, other: &$name) -> bool {
                match (self, other) {
                    (&$name(flags1), &$name(flags2)) => flags1 & flags2 != 0
                }
            }
        }

        impl core::ops::Not<$name> for $name {
            #[inline(always)]
            fn not(&self) -> $name {
                match self {
                    &$name(flags) => $name(!flags)
                }
            }
        }

        define_flags_rec!($name, 1, $( $flag $(= $v)* ),+)
    };
)
macro_rules! define_flags_rec (
    // full list (original behavior)
    (
        $name:ident,
        $default:expr,
        $($flag:ident = $value:expr),*
    ) => (
        $( pub static $flag: $name = $name($value); )+
    );
    // ----------
    // only one default value
    (
        $name:ident,
        $default:expr,
        $flag:ident
    ) => (
        pub static $flag: $name = $name($default);
    );
    // only one value
    (
        $name:ident,
        $default:expr,
        $flag:ident = $value:expr
    ) => (
        pub static $flag: $name = $name($value);
    );
    // ----------
    // without value (default)
    (
        $name:ident,
        $default:expr,
        $flag:ident,
        $($f:ident $(= $v:expr)*),*
    ) => (
        pub static $flag: $name = $name($default);
        define_flags_rec!($name, $default << 1, $($f $(= $v)*),+)
    );
    // with value
    (
        $name:ident,
        $default:expr,
        $flag:ident = $value:expr,
        $($f:ident $(= $v:expr)*),*
    ) => (
        pub static $flag: $name = $name($value);
        define_flags_rec!($name, $value << 1, $( $f $(= $v)* ),+)
    );
)
<|endoftext|>


<commit_before>//! This crate supplies custom derive implementations for the
//! [juniper](https://github.com/graphql-rust/juniper) crate.
//!
//! You should not depend on juniper_codegen directly.
//! You only need the `juniper` crate.

#![recursion_limit = "1024"]

extern crate proc_macro;
extern crate proc_macro2;
#[macro_use]
extern crate quote;
#[macro_use]
extern crate syn;
#[macro_use]
extern crate lazy_static;
extern crate regex;

mod derive_enum;
mod derive_input_object;
mod derive_object;
mod derive_scalar_value;
mod util;

use proc_macro::TokenStream;

#[proc_macro_derive(GraphQLEnum, attributes(graphql))]
pub fn derive_enum(input: TokenStream) -> TokenStream {
    let ast = syn::parse::<syn::DeriveInput>(input).unwrap();
    let gen = derive_enum::impl_enum(&ast, false);
    gen.into()
}

#[proc_macro_derive(GraphQLEnumInternal, attributes(graphql))]
#[doc(hidden)]
pub fn derive_enum_internal(input: TokenStream) -> TokenStream {
    let ast = syn::parse::<syn::DeriveInput>(input).unwrap();
    let gen = derive_enum::impl_enum(&ast, true);
    gen.into()
}

#[proc_macro_derive(GraphQLInputObject, attributes(graphql))]
pub fn derive_input_object(input: TokenStream) -> TokenStream {
    let ast = syn::parse::<syn::DeriveInput>(input).unwrap();
    let gen = derive_input_object::impl_input_object(&ast, false);
    gen.into()
}

#[proc_macro_derive(GraphQLInputObjectInternal, attributes(graphql))]
#[doc(hidden)]
pub fn derive_input_object_internal(input: TokenStream) -> TokenStream {
    let ast = syn::parse::<syn::DeriveInput>(input).unwrap();
    let gen = derive_input_object::impl_input_object(&ast, true);
    gen.into()
}

#[proc_macro_derive(GraphQLObject, attributes(graphql))]
pub fn derive_object(input: TokenStream) -> TokenStream {
    let ast = syn::parse::<syn::DeriveInput>(input).unwrap();
    let gen = derive_object::impl_object(&ast);
    gen.into()
}

#[proc_macro_derive(GraphQLScalarValue)]
pub fn derive_scalar_value(input: TokenStream) -> TokenStream {
    let ast = syn::parse::<syn::DeriveInput>(input).unwrap();
    let gen = derive_scalar_value::impl_scalar_value(&ast, false);
    gen.into()
}

#[proc_macro_derive(GraphQLScalarValueInternal)]
#[doc(hidden)]
pub fn derive_scalar_value_internal(input: TokenStream) -> TokenStream {
    let ast = syn::parse::<syn::DeriveInput>(input).unwrap();
    let gen = derive_scalar_value::impl_scalar_value(&ast, true);
    gen.into()
}
<commit_msg>Keep ScalarValue but mark it as deprecated<commit_after>//! This crate supplies custom derive implementations for the
//! [juniper](https://github.com/graphql-rust/juniper) crate.
//!
//! You should not depend on juniper_codegen directly.
//! You only need the `juniper` crate.

#![recursion_limit = "1024"]

extern crate proc_macro;
extern crate proc_macro2;
#[macro_use]
extern crate quote;
#[macro_use]
extern crate syn;
#[macro_use]
extern crate lazy_static;
extern crate regex;

mod derive_enum;
mod derive_input_object;
mod derive_object;
mod derive_scalar_value;
mod util;

use proc_macro::TokenStream;

#[proc_macro_derive(GraphQLEnum, attributes(graphql))]
pub fn derive_enum(input: TokenStream) -> TokenStream {
    let ast = syn::parse::<syn::DeriveInput>(input).unwrap();
    let gen = derive_enum::impl_enum(&ast, false);
    gen.into()
}

#[proc_macro_derive(GraphQLEnumInternal, attributes(graphql))]
#[doc(hidden)]
pub fn derive_enum_internal(input: TokenStream) -> TokenStream {
    let ast = syn::parse::<syn::DeriveInput>(input).unwrap();
    let gen = derive_enum::impl_enum(&ast, true);
    gen.into()
}

#[proc_macro_derive(GraphQLInputObject, attributes(graphql))]
pub fn derive_input_object(input: TokenStream) -> TokenStream {
    let ast = syn::parse::<syn::DeriveInput>(input).unwrap();
    let gen = derive_input_object::impl_input_object(&ast, false);
    gen.into()
}

#[proc_macro_derive(GraphQLInputObjectInternal, attributes(graphql))]
#[doc(hidden)]
pub fn derive_input_object_internal(input: TokenStream) -> TokenStream {
    let ast = syn::parse::<syn::DeriveInput>(input).unwrap();
    let gen = derive_input_object::impl_input_object(&ast, true);
    gen.into()
}

#[proc_macro_derive(GraphQLObject, attributes(graphql))]
pub fn derive_object(input: TokenStream) -> TokenStream {
    let ast = syn::parse::<syn::DeriveInput>(input).unwrap();
    let gen = derive_object::impl_object(&ast);
    gen.into()
}

#[proc_macro_derive(GraphQLScalarValue)]
pub fn derive_scalar_value(input: TokenStream) -> TokenStream {
    let ast = syn::parse::<syn::DeriveInput>(input).unwrap();
    let gen = derive_scalar_value::impl_scalar_value(&ast, false);
    gen.into()
}

#[deprecated(note = "ScalarValue has been renamed to GraphQLScalarValue")]
#[proc_macro_derive(ScalarValue)]
pub fn derive_scalar_value_deprecated(input: TokenStream) -> TokenStream {
    derive_scalar_value(input)
}

#[proc_macro_derive(GraphQLScalarValueInternal)]
#[doc(hidden)]
pub fn derive_scalar_value_internal(input: TokenStream) -> TokenStream {
    let ast = syn::parse::<syn::DeriveInput>(input).unwrap();
    let gen = derive_scalar_value::impl_scalar_value(&ast, true);
    gen.into()
}
<|endoftext|>


<commit_before>use std::io;
use std::net::{SocketAddr, Ipv4Addr, SocketAddrV4};
use std::str;
use std::time::Duration;

use futures::{Future, IntoFuture, Stream};
use futures::future;
use tokio_core::reactor::Handle;
use tokio_core::net::UdpSocket;
use tokio_timer::Timer;
use hyper;
use xml::EventReader;
use xml::reader::XmlEvent;

use gateway::Gateway;
use errors::SearchError;
use search::{SEARCH_REQUEST, parse_result};

/// Search gateway, bind to all interfaces and use a timeout of 3 seconds.
///
/// Bind to all interfaces.
/// The request will timeout after 3 seconds.
pub fn search_gateway(handle: &Handle) -> Box<Future<Item = Gateway, Error = SearchError>> {
    search_gateway_timeout(Duration::from_secs(3), handle)
}

/// Search gateway, bind to all interfaces and use the given duration for the timeout.
///
/// Bind to all interfaces.
/// The request will timeout after the given duration.
pub fn search_gateway_timeout(
    timeout: Duration,
    handle: &Handle,
) -> Box<Future<Item = Gateway, Error = SearchError>> {
    search_gateway_from_timeout(Ipv4Addr::new(0, 0, 0, 0), timeout, handle)
}

/// Search gateway, bind to the given interface and use a time of 3 seconds.
///
/// Bind to the given interface.
/// The request will timeout after 3 seconds.
pub fn search_gateway_from(
    ip: Ipv4Addr,
    handle: &Handle,
) -> Box<Future<Item = Gateway, Error = SearchError>> {
    search_gateway_from_timeout(ip, Duration::from_secs(3), handle)
}

/// Search gateway, bind to the given interface and use the given duration for the timeout.
///
/// Bind to the given interface.
/// The request will timeout after the given duration.
pub fn search_gateway_from_timeout(
    ip: Ipv4Addr,
    timeout: Duration,
    handle: &Handle,
) -> Box<Future<Item = Gateway, Error = SearchError>> {
    let addr = SocketAddr::V4(SocketAddrV4::new(ip, 0));
    let handle = handle.clone();
    let task = UdpSocket::bind(&addr, &handle)
        .into_future()
        .and_then(|socket| {
            socket.send_dgram(
                SEARCH_REQUEST.as_bytes(),
                "239.255.255.250:1900".parse().unwrap(),
            )
        })
        .and_then(|(socket, _)| socket.recv_dgram(Vec::new()))
        .map_err(|err| SearchError::from(err))
        .and_then(|(_sock, buf, n, _addr)| {
            str::from_utf8(&buf[..n])
                .map_err(|err| SearchError::from(err))
                .and_then(|text| {
                    parse_result(text).ok_or(SearchError::InvalidResponse)
                })
        })
        .and_then(move |location| {
            get_control_url(&location, &handle).and_then(move |control_url| {
                Ok(Gateway {
                    addr: location.0,
                    control_url: control_url,
                })
            })
        });
    let timeout = Timer::default().timeout(task, timeout);
    Box::new(timeout)
}

pub fn get_control_url(
    location: &(SocketAddrV4, String),
    handle: &Handle,
) -> Box<Future<Item = String, Error = SearchError>> {
    let client = hyper::Client::new(handle);
    let uri = match format!("http://{}{}", location.0, location.1).parse() {
        Ok(uri) => uri,
        Err(err) => return Box::new(future::err(SearchError::from(err))),
    };
    let future = client
        .get(uri)
        .and_then(|resp| resp.body().concat2())
        .then(|result| match result {
            Ok(body) => parse_control_url(body.as_ref()),
            Err(err) => Err(SearchError::from(err)),
        });
    Box::new(future)
}

fn parse_control_url<R>(resp: R) -> Result<String, SearchError>
where
    R: io::Read,
{

    let parser = EventReader::new(resp);
    let mut chain = Vec::<String>::with_capacity(4);

    struct Service {
        service_type: String,
        control_url: String,
    }

    let mut service = Service {
        service_type: "".to_string(),
        control_url: "".to_string(),
    };

    for e in parser.into_iter() {
        match try!(e) {
            XmlEvent::StartElement { name, .. } => {
                chain.push(name.borrow().to_repr());
                let tail = if chain.len() >= 3 {
                    chain.iter().skip(chain.len() - 3)
                } else {
                    continue;
                };

                if vec!["device", "serviceList", "service"]
                    .iter()
                    .zip(tail)
                    .all(|(l, r)| l == r)
                {
                    service.service_type.clear();
                    service.control_url.clear();
                }
            }
            XmlEvent::EndElement { .. } => {
                let top = chain.pop();
                let tail = if top == Some("service".to_string()) && chain.len() >= 2 {
                    chain.iter().skip(chain.len() - 2)
                } else {
                    continue;
                };

                if vec!["device", "serviceList"].iter().zip(tail).all(
                    |(l, r)| {
                        l == r
                    },
                )
                {
                    if ("urn:schemas-upnp-org:service:WANIPConnection:1" == service.service_type ||
                            "urn:schemas-upnp-org:service:WANPPPConnection:1" ==
                                service.service_type) &&
                        service.control_url.len() != 0
                    {
                        return Ok(service.control_url);
                    }
                }
            }
            XmlEvent::Characters(text) => {
                let tail = if chain.len() >= 4 {
                    chain.iter().skip(chain.len() - 4)
                } else {
                    continue;
                };

                if vec!["device", "serviceList", "service", "serviceType"]
                    .iter()
                    .zip(tail.clone())
                    .all(|(l, r)| l == r)
                {
                    service.service_type.push_str(&text);
                }
                if vec!["device", "serviceList", "service", "controlURL"]
                    .iter()
                    .zip(tail)
                    .all(|(l, r)| l == r)
                {
                    service.control_url.push_str(&text);
                }
            }
            _ => (),
        }
    }
    Err(SearchError::InvalidResponse)
}
<commit_msg>Return an async gateway from async search<commit_after>use std::io;
use std::net::{SocketAddr, Ipv4Addr, SocketAddrV4};
use std::str;
use std::time::Duration;

use futures::{Future, IntoFuture, Stream};
use futures::future;
use tokio_core::reactor::Handle;
use tokio_core::net::UdpSocket;
use tokio_timer::Timer;
use hyper;
use xml::EventReader;
use xml::reader::XmlEvent;

use async::Gateway;
use errors::SearchError;
use search::{SEARCH_REQUEST, parse_result};

/// Search gateway, bind to all interfaces and use a timeout of 3 seconds.
///
/// Bind to all interfaces.
/// The request will timeout after 3 seconds.
pub fn search_gateway(handle: &Handle) -> Box<Future<Item = Gateway, Error = SearchError>> {
    search_gateway_timeout(Duration::from_secs(3), handle)
}

/// Search gateway, bind to all interfaces and use the given duration for the timeout.
///
/// Bind to all interfaces.
/// The request will timeout after the given duration.
pub fn search_gateway_timeout(
    timeout: Duration,
    handle: &Handle,
) -> Box<Future<Item = Gateway, Error = SearchError>> {
    search_gateway_from_timeout(Ipv4Addr::new(0, 0, 0, 0), timeout, handle)
}

/// Search gateway, bind to the given interface and use a time of 3 seconds.
///
/// Bind to the given interface.
/// The request will timeout after 3 seconds.
pub fn search_gateway_from(
    ip: Ipv4Addr,
    handle: &Handle,
) -> Box<Future<Item = Gateway, Error = SearchError>> {
    search_gateway_from_timeout(ip, Duration::from_secs(3), handle)
}

/// Search gateway, bind to the given interface and use the given duration for the timeout.
///
/// Bind to the given interface.
/// The request will timeout after the given duration.
pub fn search_gateway_from_timeout(
    ip: Ipv4Addr,
    timeout: Duration,
    handle: &Handle,
) -> Box<Future<Item = Gateway, Error = SearchError>> {
    let addr = SocketAddr::V4(SocketAddrV4::new(ip, 0));
    let handle = handle.clone();
    let task = UdpSocket::bind(&addr, &handle)
        .into_future()
        .and_then(|socket| {
            socket.send_dgram(
                SEARCH_REQUEST.as_bytes(),
                "239.255.255.250:1900".parse().unwrap(),
            )
        })
        .and_then(|(socket, _)| socket.recv_dgram(Vec::new()))
        .map_err(|err| SearchError::from(err))
        .and_then(|(_sock, buf, n, _addr)| {
            str::from_utf8(&buf[..n])
                .map_err(|err| SearchError::from(err))
                .and_then(|text| {
                    parse_result(text).ok_or(SearchError::InvalidResponse)
                })
        })
        .and_then(move |location| {
            get_control_url(&location, &handle).and_then(move |control_url| {
                Ok(Gateway::new(location.0, control_url, handle))
            })
        });
    let timeout = Timer::default().timeout(task, timeout);
    Box::new(timeout)
}

pub fn get_control_url(
    location: &(SocketAddrV4, String),
    handle: &Handle,
) -> Box<Future<Item = String, Error = SearchError>> {
    let client = hyper::Client::new(handle);
    let uri = match format!("http://{}{}", location.0, location.1).parse() {
        Ok(uri) => uri,
        Err(err) => return Box::new(future::err(SearchError::from(err))),
    };
    let future = client
        .get(uri)
        .and_then(|resp| resp.body().concat2())
        .then(|result| match result {
            Ok(body) => parse_control_url(body.as_ref()),
            Err(err) => Err(SearchError::from(err)),
        });
    Box::new(future)
}

fn parse_control_url<R>(resp: R) -> Result<String, SearchError>
where
    R: io::Read,
{

    let parser = EventReader::new(resp);
    let mut chain = Vec::<String>::with_capacity(4);

    struct Service {
        service_type: String,
        control_url: String,
    }

    let mut service = Service {
        service_type: "".to_string(),
        control_url: "".to_string(),
    };

    for e in parser.into_iter() {
        match try!(e) {
            XmlEvent::StartElement { name, .. } => {
                chain.push(name.borrow().to_repr());
                let tail = if chain.len() >= 3 {
                    chain.iter().skip(chain.len() - 3)
                } else {
                    continue;
                };

                if vec!["device", "serviceList", "service"]
                    .iter()
                    .zip(tail)
                    .all(|(l, r)| l == r)
                {
                    service.service_type.clear();
                    service.control_url.clear();
                }
            }
            XmlEvent::EndElement { .. } => {
                let top = chain.pop();
                let tail = if top == Some("service".to_string()) && chain.len() >= 2 {
                    chain.iter().skip(chain.len() - 2)
                } else {
                    continue;
                };

                if vec!["device", "serviceList"].iter().zip(tail).all(
                    |(l, r)| {
                        l == r
                    },
                )
                {
                    if ("urn:schemas-upnp-org:service:WANIPConnection:1" == service.service_type ||
                            "urn:schemas-upnp-org:service:WANPPPConnection:1" ==
                                service.service_type) &&
                        service.control_url.len() != 0
                    {
                        return Ok(service.control_url);
                    }
                }
            }
            XmlEvent::Characters(text) => {
                let tail = if chain.len() >= 4 {
                    chain.iter().skip(chain.len() - 4)
                } else {
                    continue;
                };

                if vec!["device", "serviceList", "service", "serviceType"]
                    .iter()
                    .zip(tail.clone())
                    .all(|(l, r)| l == r)
                {
                    service.service_type.push_str(&text);
                }
                if vec!["device", "serviceList", "service", "controlURL"]
                    .iter()
                    .zip(tail)
                    .all(|(l, r)| l == r)
                {
                    service.control_url.push_str(&text);
                }
            }
            _ => (),
        }
    }
    Err(SearchError::InvalidResponse)
}
<|endoftext|>


<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

//! This module provides `Pruner` which manages a thread pruning old data in the background and is
//! meant to be triggered by other threads as they commit new data to the DB.

mod db_pruner;
pub(crate) mod db_sub_pruner;
pub(crate) mod event_store;
mod ledger_store;
pub(crate) mod state_store;
pub(crate) mod transaction_store;
pub mod utils;
pub(crate) mod worker;

use crate::metrics::{PRUNER_BATCH_SIZE, PRUNER_WINDOW};

use aptos_config::config::StoragePrunerConfig;
use aptos_infallible::Mutex;

use crate::{pruner::PrunerIndex::LedgerPrunerIndex, EventStore, LedgerStore, TransactionStore};
use aptos_types::transaction::Version;
use schemadb::DB;
use std::{
    sync::{
        mpsc::{channel, Sender},
        Arc,
    },
    thread::JoinHandle,
};
use worker::{Command, Worker};

/// The `Pruner` is meant to be part of a `AptosDB` instance and runs in the background to prune old
/// data.
///
/// It creates a worker thread on construction and joins it on destruction. When destructed, it
/// quits the worker thread eagerly without waiting for all pending work to be done.
#[derive(Debug)]
pub(crate) struct Pruner {
    /// DB version window, which dictates how many versions of state store
    /// to keep.
    state_store_prune_window: Version,
    /// DB version window, which dictates how many version of other stores like transaction, ledger
    /// info, events etc to keep.
    ledger_prune_window: Version,
    /// The worker thread handle, created upon Pruner instance construction and joined upon its
    /// destruction. It only becomes `None` after joined in `drop()`.
    worker_thread: Option<JoinHandle<()>>,
    /// The sender side of the channel talking to the worker thread.
    command_sender: Mutex<Sender<Command>>,
    /// A way for the worker thread to inform the `Pruner` the pruning progress. If it
    /// sets value to `V`, all versions before `V` can no longer be accessed. This is protected by Mutex
    /// as this is accessed both by the Pruner thread and the worker thread.
    #[allow(dead_code)]
    least_readable_version: Arc<Mutex<Vec<Version>>>,
    /// We send a batch of version to the underlying pruners for performance reason. This tracks the
    /// last version we sent to the pruner.
    last_version_sent_to_pruners: Arc<Mutex<Version>>,
    /// Ideal batch size of the versions to be sent to the pruner
    pruning_batch_size: usize,
    /// latest version
    latest_version: Arc<Mutex<Version>>,
}

pub enum PrunerIndex {
    #[cfg(test)]
    StateStorePrunerIndex,
    LedgerPrunerIndex,
}

impl Pruner {
    /// Creates a worker thread that waits on a channel for pruning commands.
    pub fn new(
        db: Arc<DB>,
        storage_pruner_config: StoragePrunerConfig,
        transaction_store: Arc<TransactionStore>,
        ledger_store: Arc<LedgerStore>,
        event_store: Arc<EventStore>,
    ) -> Self {
        let (command_sender, command_receiver) = channel();

        let least_readable_version = Arc::new(Mutex::new(vec![0, 0, 0, 0, 0]));
        let worker_progress_clone = Arc::clone(&least_readable_version);

        PRUNER_WINDOW
            .with_label_values(&["state_pruner"])
            .set((storage_pruner_config.state_store_prune_window.unwrap_or(0)) as i64);

        PRUNER_WINDOW
            .with_label_values(&["ledger_pruner"])
            .set((storage_pruner_config.ledger_prune_window.unwrap_or(0)) as i64);

        PRUNER_BATCH_SIZE.set(storage_pruner_config.pruning_batch_size as i64);

        let worker = Worker::new(
            db,
            transaction_store,
            ledger_store,
            event_store,
            command_receiver,
            least_readable_version,
            storage_pruner_config.pruning_batch_size as u64,
        );
        let worker_thread = std::thread::Builder::new()
            .name("aptosdb_pruner".into())
            .spawn(move || worker.work())
            .expect("Creating pruner thread should succeed.");

        Self {
            state_store_prune_window: storage_pruner_config
                .state_store_prune_window
                .expect("State store prune window must be specified"),
            ledger_prune_window: storage_pruner_config
                .ledger_prune_window
                .expect("Default prune window must be specified"),
            worker_thread: Some(worker_thread),
            command_sender: Mutex::new(command_sender),
            least_readable_version: worker_progress_clone,
            last_version_sent_to_pruners: Arc::new(Mutex::new(0)),
            pruning_batch_size: storage_pruner_config.pruning_batch_size,
            latest_version: Arc::new(Mutex::new(0)),
        }
    }

    pub fn get_state_store_pruner_window(&self) -> Version {
        self.state_store_prune_window
    }

    pub fn get_ledger_pruner_window(&self) -> Version {
        self.ledger_prune_window
    }

    pub fn get_least_readable_ledger_version(&self) -> Version {
        self.least_readable_version.lock()[LedgerPrunerIndex as usize]
    }
    /// Sends pruning command to the worker thread when necessary.
    pub fn maybe_wake_pruner(&self, latest_version: Version) {
        *self.latest_version.lock() = latest_version;
        if latest_version
            >= *self.last_version_sent_to_pruners.lock() + self.pruning_batch_size as u64
        {
            self.wake_pruner(latest_version);
            *self.last_version_sent_to_pruners.lock() = latest_version;
        }
    }

    fn wake_pruner(&self, latest_version: Version) {
        let least_readable_state_store_version =
            latest_version.saturating_sub(self.state_store_prune_window);
        let least_readable_ledger_version = latest_version.saturating_sub(self.ledger_prune_window);

        self.command_sender
            .lock()
            .send(Command::Prune {
                target_db_versions: vec![
                    least_readable_state_store_version,
                    least_readable_ledger_version,
                ],
            })
            .expect("Receiver should not destruct prematurely.");
    }

    /// (For tests only.) Notifies the worker thread and waits for it to finish its job by polling
    /// an internal counter.
    #[cfg(test)]
    pub fn wake_and_wait(
        &self,
        latest_version: Version,
        pruner_index: usize,
    ) -> anyhow::Result<()> {
        use std::{
            thread::sleep,
            time::{Duration, Instant},
        };

        self.maybe_wake_pruner(latest_version);

        if latest_version > self.state_store_prune_window
            || latest_version > self.ledger_prune_window
        {
            let least_readable_state_store_version = latest_version - self.state_store_prune_window;
            // Assuming no big pruning chunks will be issued by a test.
            const TIMEOUT: Duration = Duration::from_secs(10);
            let end = Instant::now() + TIMEOUT;

            while Instant::now() < end {
                if *self
                    .least_readable_version
                    .lock()
                    .get(pruner_index)
                    .unwrap()
                    >= least_readable_state_store_version
                {
                    return Ok(());
                }
                sleep(Duration::from_millis(1));
            }
            anyhow::bail!("Timeout waiting for pruner worker.");
        }
        Ok(())
    }
}

impl Drop for Pruner {
    fn drop(&mut self) {
        self.command_sender
            .lock()
            .send(Command::Quit)
            .expect("Receiver should not destruct.");
        self.worker_thread
            .take()
            .expect("Worker thread must exist.")
            .join()
            .expect("Worker thread should join peacefully.");
    }
}
<commit_msg>[Storage] Remove test compilation from the PrunerIndex.<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

//! This module provides `Pruner` which manages a thread pruning old data in the background and is
//! meant to be triggered by other threads as they commit new data to the DB.

mod db_pruner;
pub(crate) mod db_sub_pruner;
pub(crate) mod event_store;
mod ledger_store;
pub(crate) mod state_store;
pub(crate) mod transaction_store;
pub mod utils;
pub(crate) mod worker;

use crate::metrics::{PRUNER_BATCH_SIZE, PRUNER_WINDOW};

use aptos_config::config::StoragePrunerConfig;
use aptos_infallible::Mutex;

use crate::{pruner::PrunerIndex::LedgerPrunerIndex, EventStore, LedgerStore, TransactionStore};
use aptos_types::transaction::Version;
use schemadb::DB;
use std::{
    sync::{
        mpsc::{channel, Sender},
        Arc,
    },
    thread::JoinHandle,
};
use worker::{Command, Worker};

/// The `Pruner` is meant to be part of a `AptosDB` instance and runs in the background to prune old
/// data.
///
/// It creates a worker thread on construction and joins it on destruction. When destructed, it
/// quits the worker thread eagerly without waiting for all pending work to be done.
#[derive(Debug)]
pub(crate) struct Pruner {
    /// DB version window, which dictates how many versions of state store
    /// to keep.
    state_store_prune_window: Version,
    /// DB version window, which dictates how many version of other stores like transaction, ledger
    /// info, events etc to keep.
    ledger_prune_window: Version,
    /// The worker thread handle, created upon Pruner instance construction and joined upon its
    /// destruction. It only becomes `None` after joined in `drop()`.
    worker_thread: Option<JoinHandle<()>>,
    /// The sender side of the channel talking to the worker thread.
    command_sender: Mutex<Sender<Command>>,
    /// A way for the worker thread to inform the `Pruner` the pruning progress. If it
    /// sets value to `V`, all versions before `V` can no longer be accessed. This is protected by Mutex
    /// as this is accessed both by the Pruner thread and the worker thread.
    #[allow(dead_code)]
    least_readable_version: Arc<Mutex<Vec<Version>>>,
    /// We send a batch of version to the underlying pruners for performance reason. This tracks the
    /// last version we sent to the pruner.
    last_version_sent_to_pruners: Arc<Mutex<Version>>,
    /// Ideal batch size of the versions to be sent to the pruner
    pruning_batch_size: usize,
    /// latest version
    latest_version: Arc<Mutex<Version>>,
}

pub enum PrunerIndex {
    #[allow(dead_code)]
    StateStorePrunerIndex,
    LedgerPrunerIndex,
}

impl Pruner {
    /// Creates a worker thread that waits on a channel for pruning commands.
    pub fn new(
        db: Arc<DB>,
        storage_pruner_config: StoragePrunerConfig,
        transaction_store: Arc<TransactionStore>,
        ledger_store: Arc<LedgerStore>,
        event_store: Arc<EventStore>,
    ) -> Self {
        let (command_sender, command_receiver) = channel();

        let least_readable_version = Arc::new(Mutex::new(vec![0, 0, 0, 0, 0]));
        let worker_progress_clone = Arc::clone(&least_readable_version);

        PRUNER_WINDOW
            .with_label_values(&["state_pruner"])
            .set((storage_pruner_config.state_store_prune_window.unwrap_or(0)) as i64);

        PRUNER_WINDOW
            .with_label_values(&["ledger_pruner"])
            .set((storage_pruner_config.ledger_prune_window.unwrap_or(0)) as i64);

        PRUNER_BATCH_SIZE.set(storage_pruner_config.pruning_batch_size as i64);

        let worker = Worker::new(
            db,
            transaction_store,
            ledger_store,
            event_store,
            command_receiver,
            least_readable_version,
            storage_pruner_config.pruning_batch_size as u64,
        );
        let worker_thread = std::thread::Builder::new()
            .name("aptosdb_pruner".into())
            .spawn(move || worker.work())
            .expect("Creating pruner thread should succeed.");

        Self {
            state_store_prune_window: storage_pruner_config
                .state_store_prune_window
                .expect("State store prune window must be specified"),
            ledger_prune_window: storage_pruner_config
                .ledger_prune_window
                .expect("Default prune window must be specified"),
            worker_thread: Some(worker_thread),
            command_sender: Mutex::new(command_sender),
            least_readable_version: worker_progress_clone,
            last_version_sent_to_pruners: Arc::new(Mutex::new(0)),
            pruning_batch_size: storage_pruner_config.pruning_batch_size,
            latest_version: Arc::new(Mutex::new(0)),
        }
    }

    pub fn get_state_store_pruner_window(&self) -> Version {
        self.state_store_prune_window
    }

    pub fn get_ledger_pruner_window(&self) -> Version {
        self.ledger_prune_window
    }

    pub fn get_least_readable_ledger_version(&self) -> Version {
        self.least_readable_version.lock()[LedgerPrunerIndex as usize]
    }
    /// Sends pruning command to the worker thread when necessary.
    pub fn maybe_wake_pruner(&self, latest_version: Version) {
        *self.latest_version.lock() = latest_version;
        if latest_version
            >= *self.last_version_sent_to_pruners.lock() + self.pruning_batch_size as u64
        {
            self.wake_pruner(latest_version);
            *self.last_version_sent_to_pruners.lock() = latest_version;
        }
    }

    fn wake_pruner(&self, latest_version: Version) {
        let least_readable_state_store_version =
            latest_version.saturating_sub(self.state_store_prune_window);
        let least_readable_ledger_version = latest_version.saturating_sub(self.ledger_prune_window);

        self.command_sender
            .lock()
            .send(Command::Prune {
                target_db_versions: vec![
                    least_readable_state_store_version,
                    least_readable_ledger_version,
                ],
            })
            .expect("Receiver should not destruct prematurely.");
    }

    /// (For tests only.) Notifies the worker thread and waits for it to finish its job by polling
    /// an internal counter.
    #[cfg(test)]
    pub fn wake_and_wait(
        &self,
        latest_version: Version,
        pruner_index: usize,
    ) -> anyhow::Result<()> {
        use std::{
            thread::sleep,
            time::{Duration, Instant},
        };

        self.maybe_wake_pruner(latest_version);

        if latest_version > self.state_store_prune_window
            || latest_version > self.ledger_prune_window
        {
            let least_readable_state_store_version = latest_version - self.state_store_prune_window;
            // Assuming no big pruning chunks will be issued by a test.
            const TIMEOUT: Duration = Duration::from_secs(10);
            let end = Instant::now() + TIMEOUT;

            while Instant::now() < end {
                if *self
                    .least_readable_version
                    .lock()
                    .get(pruner_index)
                    .unwrap()
                    >= least_readable_state_store_version
                {
                    return Ok(());
                }
                sleep(Duration::from_millis(1));
            }
            anyhow::bail!("Timeout waiting for pruner worker.");
        }
        Ok(())
    }
}

impl Drop for Pruner {
    fn drop(&mut self) {
        self.command_sender
            .lock()
            .send(Command::Quit)
            .expect("Receiver should not destruct.");
        self.worker_thread
            .take()
            .expect("Worker thread must exist.")
            .join()
            .expect("Worker thread should join peacefully.");
    }
}
<|endoftext|>


<commit_before>//! Handle `cargo add` arguments

#![allow(clippy::bool_assert_comparison)]

use cargo_edit::{
    find, get_features_from_registry, get_manifest_from_path, get_manifest_from_url, registry_url,
    workspace_members, Dependency,
};
use cargo_edit::{get_latest_dependency, CrateName};
use cargo_metadata::Package;
use clap::Parser;
use std::path::PathBuf;

use crate::errors::*;

#[derive(Debug, Parser)]
#[clap(bin_name = "cargo")]
pub enum Command {
    /// Add dependency to a Cargo.toml manifest file.
    #[clap(name = "add")]
    #[clap(after_help = "\
This command allows you to add a dependency to a Cargo.toml manifest file. If <crate> is a github \
or gitlab repository URL, or a local path, `cargo add` will try to automatically get the crate \
name and set the appropriate `--git` or `--path` value.

Please note that Cargo treats versions like '1.2.3' as '^1.2.3' (and that '^1.2.3' is specified \
as '>=1.2.3 and <2.0.0'). By default, `cargo add` will use this format, as it is the one that the \
crates.io registry suggests. One goal of `cargo add` is to prevent you from using wildcard \
dependencies (version set to '*').")]
    Add(Args),
}

#[derive(Debug, Parser)]
#[clap(about, version)]
pub struct Args {
    /// Crates to be added.
    #[clap(value_name = "CRATE", required = true)]
    pub crates: Vec<String>,

    /// Rename a dependency in Cargo.toml,
    /// https://doc.rust-lang.org/cargo/reference/specifying-dependencies.html#renaming-dependencies-in-cargotoml.
    /// Only works when specifying a single dependency.
    #[clap(long, short)]
    pub rename: Option<String>,

    /// Add crate as development dependency.
    #[clap(long, short = 'D', conflicts_with = "build")]
    pub dev: bool,

    /// Add crate as build dependency.
    #[clap(long, short = 'B', conflicts_with = "dev")]
    pub build: bool,

    /// Specify the version to grab from the registry(crates.io).
    /// You can also specify version as part of name, e.g
    /// `cargo add bitflags@0.3.2`.
    #[clap(long, value_name = "URI", conflicts_with = "git")]
    pub vers: Option<String>,

    /// Specify a git repository to download the crate from.
    #[clap(
        long,
        value_name = "URI",
        conflicts_with = "vers",
        conflicts_with = "path"
    )]
    pub git: Option<String>,

    /// Specify a git branch to download the crate from.
    #[clap(
        long,
        value_name = "BRANCH",
        conflicts_with = "vers",
        conflicts_with = "path"
    )]
    pub branch: Option<String>,

    /// Specify the path the crate should be loaded from.
    #[clap(long, parse(from_os_str), conflicts_with = "git")]
    pub path: Option<PathBuf>,

    /// Add as dependency to the given target platform.
    #[clap(long, conflicts_with = "dev", conflicts_with = "build")]
    pub target: Option<String>,

    /// Add as an optional dependency (for use in features).
    #[clap(long, conflicts_with = "dev")]
    pub optional: bool,

    /// Path to the manifest to add a dependency to.
    #[clap(
        long,
        value_name = "PATH",
        parse(from_os_str),
        conflicts_with = "pkgid"
    )]
    pub manifest_path: Option<PathBuf>,

    /// Package id of the crate to add this dependency to.
    #[clap(
        long = "package",
        short = 'p',
        value_name = "PKGID",
        conflicts_with = "manifest-path"
    )]
    pub pkgid: Option<String>,

    /// Choose method of semantic version upgrade.  Must be one of "none" (exact version, `=`
    /// modifier), "patch" (`~` modifier), "minor" (`^` modifier), "all" (`>=`), or "default" (no
    /// modifier).
    #[clap(
        long,
        value_name = "METHOD",
        possible_value = "none",
        possible_value = "patch",
        possible_value = "minor",
        possible_value = "all",
        possible_value = "default",
        default_value = "default"
    )]
    pub upgrade: String,

    /// Include prerelease versions when fetching from crates.io (e.g.
    /// '0.6.0-alpha').
    #[clap(long)]
    pub allow_prerelease: bool,

    /// Space-separated list of features to add. For an alternative approach to
    /// enabling features, consider installing the `cargo-feature` utility.
    #[clap(long)]
    pub features: Option<Vec<String>>,

    /// Set `default-features = false` for the added dependency.
    #[clap(long)]
    pub no_default_features: bool,

    /// Do not print any output in case of success.
    #[clap(long)]
    pub quiet: bool,

    /// Run without accessing the network
    #[clap(long)]
    pub offline: bool,

    /// Sort dependencies even if currently unsorted
    #[clap(long, short)]
    pub sort: bool,

    /// Registry to use
    #[clap(long, conflicts_with = "git", conflicts_with = "path")]
    pub registry: Option<String>,
}

fn parse_version_req(s: &str) -> Result<&str> {
    semver::VersionReq::parse(s).chain_err(|| "Invalid dependency version requirement")?;
    Ok(s)
}

impl Args {
    /// Get dependency section
    pub fn get_section(&self) -> Vec<String> {
        if self.dev {
            vec!["dev-dependencies".to_owned()]
        } else if self.build {
            vec!["build-dependencies".to_owned()]
        } else if let Some(ref target) = self.target {
            assert!(!target.is_empty(), "Target specification may not be empty");

            vec![
                "target".to_owned(),
                target.clone(),
                "dependencies".to_owned(),
            ]
        } else {
            vec!["dependencies".to_owned()]
        }
    }

    fn parse_single_dependency(
        &self,
        crate_name: &str,
        workspace_members: &[Package],
    ) -> Result<Dependency> {
        let crate_name = CrateName::new(crate_name);
        let manifest_path = find(&self.manifest_path)?;
        let registry_url = registry_url(&manifest_path, self.registry.as_deref())?;

        let mut dependency = if let Some(mut dependency) = crate_name.parse_as_version()? {
            // crate specifier includes a version (e.g. `docopt@0.8`)
            if let Some(ref url) = self.git {
                let url = url.clone();
                let version = dependency.version().unwrap().to_string();
                return Err(ErrorKind::GitUrlWithVersion(url, version).into());
            }

            if let Some(ref path) = self.path {
                let manifest = get_manifest_from_path(path)?;
                let dep_path = dunce::canonicalize(path)?;

                dependency = dependency.set_available_features(manifest.features()?);
                dependency = dependency.set_path(dep_path);
            } else {
                let features = get_features_from_registry(
                    &dependency.name,
                    dependency
                        .version()
                        .expect("version populated by `parse_as_version`"),
                    &registry_url,
                )?;
                dependency = dependency.set_available_features(features);
            }

            dependency
        } else if let Some(mut dependency) = crate_name.parse_crate_name_from_uri()? {
            // dev-dependencies do not need the version populated
            if !self.dev {
                let dep_path = dependency.path().map(ToOwned::to_owned);
                if let Some(dep_path) = dep_path {
                    if let Some(package) = workspace_members.iter().find(|p| {
                        p.manifest_path.parent().map(|p| p.as_std_path())
                            == Some(dep_path.as_path())
                    }) {
                        let v = format!(
                            "{prefix}{version}",
                            prefix = self.get_upgrade_prefix(),
                            version = package.version
                        );

                        dependency = dependency.set_version(&v);
                    }
                }
            }
            dependency
        } else {
            let mut dependency = Dependency::new(crate_name.name());

            if let Some(repo) = &self.git {
                assert!(self.vers.is_none());
                assert!(self.path.is_none());
                assert!(self.registry.is_none());
                let features = get_manifest_from_url(repo)?
                    .map(|m| m.features())
                    .transpose()?
                    .unwrap_or_else(Vec::new);

                dependency = dependency
                    .set_git(repo, self.branch.clone())
                    .set_available_features(features);
            } else {
                if let Some(path) = &self.path {
                    assert!(self.registry.is_none());
                    let manifest = get_manifest_from_path(path)?;

                    dependency = dependency.set_available_features(manifest.features()?);
                    dependency = dependency.set_path(dunce::canonicalize(path)?);
                }
                if let Some(version) = &self.vers {
                    dependency = dependency.set_version(parse_version_req(version)?);
                }

                if self.git.is_none() && self.path.is_none() && self.vers.is_none() {
                    // Only special-case workspaces when the user doesn't provide any extra
                    // information, otherwise, trust the user.
                    if let Some(package) = workspace_members
                        .iter()
                        .find(|p| p.name == crate_name.name())
                    {
                        dependency = dependency.set_path(
                            package
                                .manifest_path
                                .parent()
                                .expect("at least parent dir")
                                .as_std_path()
                                .to_owned(),
                        );
                        // dev-dependencies do not need the version populated
                        if !self.dev {
                            let v = format!(
                                "{prefix}{version}",
                                prefix = self.get_upgrade_prefix(),
                                version = package.version
                            );
                            dependency = dependency.set_version(&v);
                        }
                    } else {
                        dependency = get_latest_dependency(
                            crate_name.name(),
                            self.allow_prerelease,
                            &manifest_path,
                            Some(&registry_url),
                        )?;
                        let v = format!(
                            "{prefix}{version}",
                            prefix = self.get_upgrade_prefix(),
                            // If version is unavailable `get_latest_dependency` must have
                            // returned `Err(FetchVersionError::GetVersion)`
                            version = dependency.version().unwrap_or_else(|| unreachable!())
                        );
                        dependency = dependency.set_version(&v);
                    }
                }
            }

            dependency
        };

        if let Some(registry) = &self.registry {
            dependency = dependency.set_registry(registry);
        }

        Ok(dependency)
    }

    /// Build dependencies from arguments
    pub fn parse_dependencies(
        &self,
        requested_features: Option<Vec<String>>,
    ) -> Result<Vec<Dependency>> {
        let workspace_members = workspace_members(self.manifest_path.as_deref())?;

        if self.crates.len() > 1
            && (self.git.is_some() || self.path.is_some() || self.vers.is_some())
        {
            return Err(ErrorKind::MultipleCratesWithGitOrPathOrVers.into());
        }

        if self.crates.len() > 1 && self.rename.is_some() {
            return Err(ErrorKind::MultipleCratesWithRename.into());
        }

        if self.crates.len() > 1 && self.features.is_some() {
            return Err(ErrorKind::MultipleCratesWithFeatures.into());
        }

        self.crates
            .iter()
            .map(|crate_name| {
                self.parse_single_dependency(crate_name, &workspace_members)
                    .map(|x| {
                        let mut x = x
                            .set_optional(self.optional)
                            .set_features(requested_features.to_owned())
                            .set_default_features(!self.no_default_features);
                        if let Some(ref rename) = self.rename {
                            x = x.set_rename(rename);
                        }
                        x
                    })
            })
            .collect()
    }

    fn get_upgrade_prefix(&self) -> &'static str {
        match self.upgrade.as_ref() {
            "default" => "",
            "none" => "=",
            "patch" => "~",
            "minor" => "^",
            "all" => ">=",
            _ => unreachable!(),
        }
    }
}

#[cfg(test)]
impl Default for Args {
    fn default() -> Args {
        Args {
            crates: vec!["demo".to_owned()],
            rename: None,
            dev: false,
            build: false,
            vers: None,
            git: None,
            branch: None,
            path: None,
            target: None,
            optional: false,
            manifest_path: None,
            pkgid: None,
            upgrade: "minor".to_string(),
            allow_prerelease: false,
            features: None,
            no_default_features: false,
            quiet: false,
            offline: true,
            sort: false,
            registry: None,
        }
    }
}

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

    #[test]
    fn test_dependency_parsing() {
        let args = Args {
            vers: Some("0.4.2".to_owned()),
            ..Args::default()
        };

        assert_eq!(
            args.parse_dependencies(None).unwrap(),
            vec![Dependency::new("demo").set_version("0.4.2")]
        );
    }

    #[test]
    #[cfg(feature = "test-external-apis")]
    fn test_repo_as_arg_parsing() {
        let github_url = "https://github.com/killercup/cargo-edit/";
        let args_github = Args {
            crates: vec![github_url.to_owned()],
            ..Args::default()
        };
        assert_eq!(
            args_github.parse_dependencies(None).unwrap(),
            vec![Dependency::new("cargo-edit").set_git(github_url, None)]
        );

        let gitlab_url = "https://gitlab.com/Polly-lang/Polly.git";
        let args_gitlab = Args {
            crates: vec![gitlab_url.to_owned()],
            ..Args::default()
        };
        assert_eq!(
            args_gitlab.parse_dependencies(None).unwrap(),
            vec![Dependency::new("polly").set_git(gitlab_url, None)]
        );
    }

    #[test]
    fn test_path_as_arg_parsing() {
        let self_path = dunce::canonicalize(std::env::current_dir().unwrap()).unwrap();
        let args_path = Args {
            // Hacky to `display` but should generally work
            crates: vec![self_path.display().to_string()],
            ..Args::default()
        };
        assert_eq!(
            args_path.parse_dependencies(None).unwrap()[0]
                .path()
                .unwrap(),
            self_path
        );
    }

    #[test]
    fn verify_app() {
        use clap::IntoApp;
        Command::into_app().debug_assert()
    }
}
<commit_msg>fix(add): Report empty error rather than panic<commit_after>//! Handle `cargo add` arguments

#![allow(clippy::bool_assert_comparison)]

use cargo_edit::{
    find, get_features_from_registry, get_manifest_from_path, get_manifest_from_url, registry_url,
    workspace_members, Dependency,
};
use cargo_edit::{get_latest_dependency, CrateName};
use cargo_metadata::Package;
use clap::Parser;
use std::path::PathBuf;

use crate::errors::*;

#[derive(Debug, Parser)]
#[clap(bin_name = "cargo")]
pub enum Command {
    /// Add dependency to a Cargo.toml manifest file.
    #[clap(name = "add")]
    #[clap(after_help = "\
This command allows you to add a dependency to a Cargo.toml manifest file. If <crate> is a github \
or gitlab repository URL, or a local path, `cargo add` will try to automatically get the crate \
name and set the appropriate `--git` or `--path` value.

Please note that Cargo treats versions like '1.2.3' as '^1.2.3' (and that '^1.2.3' is specified \
as '>=1.2.3 and <2.0.0'). By default, `cargo add` will use this format, as it is the one that the \
crates.io registry suggests. One goal of `cargo add` is to prevent you from using wildcard \
dependencies (version set to '*').")]
    Add(Args),
}

#[derive(Debug, Parser)]
#[clap(about, version)]
pub struct Args {
    /// Crates to be added.
    #[clap(value_name = "CRATE", required = true)]
    pub crates: Vec<String>,

    /// Rename a dependency in Cargo.toml,
    /// https://doc.rust-lang.org/cargo/reference/specifying-dependencies.html#renaming-dependencies-in-cargotoml.
    /// Only works when specifying a single dependency.
    #[clap(long, short)]
    pub rename: Option<String>,

    /// Add crate as development dependency.
    #[clap(long, short = 'D', conflicts_with = "build")]
    pub dev: bool,

    /// Add crate as build dependency.
    #[clap(long, short = 'B', conflicts_with = "dev")]
    pub build: bool,

    /// Specify the version to grab from the registry(crates.io).
    /// You can also specify version as part of name, e.g
    /// `cargo add bitflags@0.3.2`.
    #[clap(long, value_name = "URI", conflicts_with = "git")]
    pub vers: Option<String>,

    /// Specify a git repository to download the crate from.
    #[clap(
        long,
        value_name = "URI",
        conflicts_with = "vers",
        conflicts_with = "path"
    )]
    pub git: Option<String>,

    /// Specify a git branch to download the crate from.
    #[clap(
        long,
        value_name = "BRANCH",
        conflicts_with = "vers",
        conflicts_with = "path"
    )]
    pub branch: Option<String>,

    /// Specify the path the crate should be loaded from.
    #[clap(long, parse(from_os_str), conflicts_with = "git")]
    pub path: Option<PathBuf>,

    /// Add as dependency to the given target platform.
    #[clap(
        long,
        conflicts_with = "dev",
        conflicts_with = "build",
        forbid_empty_values = true
    )]
    pub target: Option<String>,

    /// Add as an optional dependency (for use in features).
    #[clap(long, conflicts_with = "dev")]
    pub optional: bool,

    /// Path to the manifest to add a dependency to.
    #[clap(
        long,
        value_name = "PATH",
        parse(from_os_str),
        conflicts_with = "pkgid"
    )]
    pub manifest_path: Option<PathBuf>,

    /// Package id of the crate to add this dependency to.
    #[clap(
        long = "package",
        short = 'p',
        value_name = "PKGID",
        conflicts_with = "manifest-path"
    )]
    pub pkgid: Option<String>,

    /// Choose method of semantic version upgrade.  Must be one of "none" (exact version, `=`
    /// modifier), "patch" (`~` modifier), "minor" (`^` modifier), "all" (`>=`), or "default" (no
    /// modifier).
    #[clap(
        long,
        value_name = "METHOD",
        possible_value = "none",
        possible_value = "patch",
        possible_value = "minor",
        possible_value = "all",
        possible_value = "default",
        default_value = "default"
    )]
    pub upgrade: String,

    /// Include prerelease versions when fetching from crates.io (e.g.
    /// '0.6.0-alpha').
    #[clap(long)]
    pub allow_prerelease: bool,

    /// Space-separated list of features to add. For an alternative approach to
    /// enabling features, consider installing the `cargo-feature` utility.
    #[clap(long)]
    pub features: Option<Vec<String>>,

    /// Set `default-features = false` for the added dependency.
    #[clap(long)]
    pub no_default_features: bool,

    /// Do not print any output in case of success.
    #[clap(long)]
    pub quiet: bool,

    /// Run without accessing the network
    #[clap(long)]
    pub offline: bool,

    /// Sort dependencies even if currently unsorted
    #[clap(long, short)]
    pub sort: bool,

    /// Registry to use
    #[clap(long, conflicts_with = "git", conflicts_with = "path")]
    pub registry: Option<String>,
}

fn parse_version_req(s: &str) -> Result<&str> {
    semver::VersionReq::parse(s).chain_err(|| "Invalid dependency version requirement")?;
    Ok(s)
}

impl Args {
    /// Get dependency section
    pub fn get_section(&self) -> Vec<String> {
        if self.dev {
            vec!["dev-dependencies".to_owned()]
        } else if self.build {
            vec!["build-dependencies".to_owned()]
        } else if let Some(ref target) = self.target {
            assert!(!target.is_empty(), "Target specification may not be empty");

            vec![
                "target".to_owned(),
                target.clone(),
                "dependencies".to_owned(),
            ]
        } else {
            vec!["dependencies".to_owned()]
        }
    }

    fn parse_single_dependency(
        &self,
        crate_name: &str,
        workspace_members: &[Package],
    ) -> Result<Dependency> {
        let crate_name = CrateName::new(crate_name);
        let manifest_path = find(&self.manifest_path)?;
        let registry_url = registry_url(&manifest_path, self.registry.as_deref())?;

        let mut dependency = if let Some(mut dependency) = crate_name.parse_as_version()? {
            // crate specifier includes a version (e.g. `docopt@0.8`)
            if let Some(ref url) = self.git {
                let url = url.clone();
                let version = dependency.version().unwrap().to_string();
                return Err(ErrorKind::GitUrlWithVersion(url, version).into());
            }

            if let Some(ref path) = self.path {
                let manifest = get_manifest_from_path(path)?;
                let dep_path = dunce::canonicalize(path)?;

                dependency = dependency.set_available_features(manifest.features()?);
                dependency = dependency.set_path(dep_path);
            } else {
                let features = get_features_from_registry(
                    &dependency.name,
                    dependency
                        .version()
                        .expect("version populated by `parse_as_version`"),
                    &registry_url,
                )?;
                dependency = dependency.set_available_features(features);
            }

            dependency
        } else if let Some(mut dependency) = crate_name.parse_crate_name_from_uri()? {
            // dev-dependencies do not need the version populated
            if !self.dev {
                let dep_path = dependency.path().map(ToOwned::to_owned);
                if let Some(dep_path) = dep_path {
                    if let Some(package) = workspace_members.iter().find(|p| {
                        p.manifest_path.parent().map(|p| p.as_std_path())
                            == Some(dep_path.as_path())
                    }) {
                        let v = format!(
                            "{prefix}{version}",
                            prefix = self.get_upgrade_prefix(),
                            version = package.version
                        );

                        dependency = dependency.set_version(&v);
                    }
                }
            }
            dependency
        } else {
            let mut dependency = Dependency::new(crate_name.name());

            if let Some(repo) = &self.git {
                assert!(self.vers.is_none());
                assert!(self.path.is_none());
                assert!(self.registry.is_none());
                let features = get_manifest_from_url(repo)?
                    .map(|m| m.features())
                    .transpose()?
                    .unwrap_or_else(Vec::new);

                dependency = dependency
                    .set_git(repo, self.branch.clone())
                    .set_available_features(features);
            } else {
                if let Some(path) = &self.path {
                    assert!(self.registry.is_none());
                    let manifest = get_manifest_from_path(path)?;

                    dependency = dependency.set_available_features(manifest.features()?);
                    dependency = dependency.set_path(dunce::canonicalize(path)?);
                }
                if let Some(version) = &self.vers {
                    dependency = dependency.set_version(parse_version_req(version)?);
                }

                if self.git.is_none() && self.path.is_none() && self.vers.is_none() {
                    // Only special-case workspaces when the user doesn't provide any extra
                    // information, otherwise, trust the user.
                    if let Some(package) = workspace_members
                        .iter()
                        .find(|p| p.name == crate_name.name())
                    {
                        dependency = dependency.set_path(
                            package
                                .manifest_path
                                .parent()
                                .expect("at least parent dir")
                                .as_std_path()
                                .to_owned(),
                        );
                        // dev-dependencies do not need the version populated
                        if !self.dev {
                            let v = format!(
                                "{prefix}{version}",
                                prefix = self.get_upgrade_prefix(),
                                version = package.version
                            );
                            dependency = dependency.set_version(&v);
                        }
                    } else {
                        dependency = get_latest_dependency(
                            crate_name.name(),
                            self.allow_prerelease,
                            &manifest_path,
                            Some(&registry_url),
                        )?;
                        let v = format!(
                            "{prefix}{version}",
                            prefix = self.get_upgrade_prefix(),
                            // If version is unavailable `get_latest_dependency` must have
                            // returned `Err(FetchVersionError::GetVersion)`
                            version = dependency.version().unwrap_or_else(|| unreachable!())
                        );
                        dependency = dependency.set_version(&v);
                    }
                }
            }

            dependency
        };

        if let Some(registry) = &self.registry {
            dependency = dependency.set_registry(registry);
        }

        Ok(dependency)
    }

    /// Build dependencies from arguments
    pub fn parse_dependencies(
        &self,
        requested_features: Option<Vec<String>>,
    ) -> Result<Vec<Dependency>> {
        let workspace_members = workspace_members(self.manifest_path.as_deref())?;

        if self.crates.len() > 1
            && (self.git.is_some() || self.path.is_some() || self.vers.is_some())
        {
            return Err(ErrorKind::MultipleCratesWithGitOrPathOrVers.into());
        }

        if self.crates.len() > 1 && self.rename.is_some() {
            return Err(ErrorKind::MultipleCratesWithRename.into());
        }

        if self.crates.len() > 1 && self.features.is_some() {
            return Err(ErrorKind::MultipleCratesWithFeatures.into());
        }

        self.crates
            .iter()
            .map(|crate_name| {
                self.parse_single_dependency(crate_name, &workspace_members)
                    .map(|x| {
                        let mut x = x
                            .set_optional(self.optional)
                            .set_features(requested_features.to_owned())
                            .set_default_features(!self.no_default_features);
                        if let Some(ref rename) = self.rename {
                            x = x.set_rename(rename);
                        }
                        x
                    })
            })
            .collect()
    }

    fn get_upgrade_prefix(&self) -> &'static str {
        match self.upgrade.as_ref() {
            "default" => "",
            "none" => "=",
            "patch" => "~",
            "minor" => "^",
            "all" => ">=",
            _ => unreachable!(),
        }
    }
}

#[cfg(test)]
impl Default for Args {
    fn default() -> Args {
        Args {
            crates: vec!["demo".to_owned()],
            rename: None,
            dev: false,
            build: false,
            vers: None,
            git: None,
            branch: None,
            path: None,
            target: None,
            optional: false,
            manifest_path: None,
            pkgid: None,
            upgrade: "minor".to_string(),
            allow_prerelease: false,
            features: None,
            no_default_features: false,
            quiet: false,
            offline: true,
            sort: false,
            registry: None,
        }
    }
}

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

    #[test]
    fn test_dependency_parsing() {
        let args = Args {
            vers: Some("0.4.2".to_owned()),
            ..Args::default()
        };

        assert_eq!(
            args.parse_dependencies(None).unwrap(),
            vec![Dependency::new("demo").set_version("0.4.2")]
        );
    }

    #[test]
    #[cfg(feature = "test-external-apis")]
    fn test_repo_as_arg_parsing() {
        let github_url = "https://github.com/killercup/cargo-edit/";
        let args_github = Args {
            crates: vec![github_url.to_owned()],
            ..Args::default()
        };
        assert_eq!(
            args_github.parse_dependencies(None).unwrap(),
            vec![Dependency::new("cargo-edit").set_git(github_url, None)]
        );

        let gitlab_url = "https://gitlab.com/Polly-lang/Polly.git";
        let args_gitlab = Args {
            crates: vec![gitlab_url.to_owned()],
            ..Args::default()
        };
        assert_eq!(
            args_gitlab.parse_dependencies(None).unwrap(),
            vec![Dependency::new("polly").set_git(gitlab_url, None)]
        );
    }

    #[test]
    fn test_path_as_arg_parsing() {
        let self_path = dunce::canonicalize(std::env::current_dir().unwrap()).unwrap();
        let args_path = Args {
            // Hacky to `display` but should generally work
            crates: vec![self_path.display().to_string()],
            ..Args::default()
        };
        assert_eq!(
            args_path.parse_dependencies(None).unwrap()[0]
                .path()
                .unwrap(),
            self_path
        );
    }

    #[test]
    fn verify_app() {
        use clap::IntoApp;
        Command::into_app().debug_assert()
    }
}
<|endoftext|>