<commit_before>//! This modules implements "expand macro" functionality in the IDE

use crate::{db::RootDatabase, FilePosition};
use hir::db::AstDatabase;
use ra_db::SourceDatabase;
use rustc_hash::FxHashMap;

use ra_syntax::{
    algo::{find_node_at_offset, replace_descendants},
    ast::{self},
    AstNode, NodeOrToken, SyntaxKind, SyntaxNode, WalkEvent, T,
};

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

pub(crate) fn expand_macro(db: &RootDatabase, position: FilePosition) -> Option<ExpandedMacro> {
    let parse = db.parse(position.file_id);
    let file = parse.tree();
    let name_ref = find_node_at_offset::<ast::NameRef>(file.syntax(), position.offset)?;
    let mac = name_ref.syntax().ancestors().find_map(ast::MacroCall::cast)?;

    let source = hir::Source::new(position.file_id.into(), mac.syntax());
    let expanded = expand_macro_recur(db, source, source.with_value(&mac))?;

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

fn expand_macro_recur(
    db: &RootDatabase,
    source: hir::Source<&SyntaxNode>,
    macro_call: hir::Source<&ast::MacroCall>,
) -> Option<SyntaxNode> {
    let analyzer = hir::SourceAnalyzer::new(db, source, None);
    let expansion = analyzer.expand(db, macro_call)?;
    let macro_file_id = expansion.file_id();
    let mut expanded: SyntaxNode = db.parse_or_expand(macro_file_id)?;

    let children = expanded.descendants().filter_map(ast::MacroCall::cast);
    let mut replaces = FxHashMap::default();

    for child in children.into_iter() {
        let node = hir::Source::new(macro_file_id, &child);
        let new_node = expand_macro_recur(db, source, node)?;

        // Replace the whole node if it is root
        // `replace_descendants` will not replace the parent node
        // but `SyntaxNode::descendants include itself
        if expanded == *child.syntax() {
            expanded = new_node;
        } else {
            replaces.insert(child.syntax().clone().into(), new_node.into());
        }
    }

    Some(replace_descendants(&expanded, &replaces))
}

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

    let mut res = String::new();
    let mut token_iter = syn
        .preorder_with_tokens()
        .filter_map(|event| {
            if let WalkEvent::Enter(NodeOrToken::Token(token)) = event {
                Some(token)
            } else {
                None
            }
        })
        .peekable();

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

    while let Some(token) = token_iter.next() {
        let mut is_next = |f: fn(SyntaxKind) -> bool, default| -> bool {
            token_iter.peek().map(|it| f(it.kind())).unwrap_or(default)
        };
        let is_last = |f: fn(SyntaxKind) -> bool, default| -> bool {
            last.map(|it| f(it)).unwrap_or(default)
        };

        res += &match token.kind() {
            k @ _ if is_text(k) && is_next(|it| !it.is_punct(), true) => {
                token.text().to_string() + " "
            }
            L_CURLY if is_next(|it| it != R_CURLY, true) => {
                indent += 1;
                let leading_space = if is_last(|it| is_text(it), false) { " " } else { "" };
                format!("{}{{\n{}", leading_space, "  ".repeat(indent))
            }
            R_CURLY if is_last(|it| it != L_CURLY, true) => {
                indent = indent.checked_sub(1).unwrap_or(0);
                format!("\n{}}}", "  ".repeat(indent))
            }
            R_CURLY => format!("}}\n{}", "  ".repeat(indent)),
            T![;] => format!(";\n{}", "  ".repeat(indent)),
            T![->] => " -> ".to_string(),
            T![=] => " = ".to_string(),
            T![=>] => " => ".to_string(),
            _ => token.text().to_string(),
        };

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

    return res;

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mock_analysis::analysis_and_position;
    use insta::assert_snapshot;

    fn check_expand_macro(fixture: &str) -> ExpandedMacro {
        let (analysis, pos) = analysis_and_position(fixture);
        analysis.expand_macro(pos).unwrap().unwrap()
    }

    #[test]
    fn macro_expand_recursive_expansion() {
        let res = check_expand_macro(
            r#"
        //- /lib.rs
        macro_rules! bar {
            () => { fn  b() {} }
        }
        macro_rules! foo {
            () => { bar!(); }
        }
        macro_rules! baz {
            () => { foo!(); }
        }
        f<|>oo!();
        "#,
        );

        assert_eq!(res.name, "foo");
        assert_snapshot!(res.expansion, @r###"
fn b(){}
"###);
    }

    #[test]
    fn macro_expand_multiple_lines() {
        let res = check_expand_macro(
            r#"
        //- /lib.rs
        macro_rules! foo {
            () => {
                fn some_thing() -> u32 {
                    let a = 0;
                    a + 10
                }
            }
        }
        f<|>oo!();
        "#,
        );

        assert_eq!(res.name, "foo");
        assert_snapshot!(res.expansion, @r###"
fn some_thing() -> u32 {
  let a = 0;
  a+10
}
"###);
    }

    #[test]
    fn macro_expand_match_ast() {
        let res = check_expand_macro(
            r#"
        //- /lib.rs
        macro_rules! match_ast {
            (match $node:ident { $($tt:tt)* }) => { match_ast!(match ($node) { $($tt)* }) };

            (match ($node:expr) {
                $( ast::$ast:ident($it:ident) => $res:block, )*
                _ => $catch_all:expr $(,)?
            }) => {{
                $( if let Some($it) = ast::$ast::cast($node.clone()) $res else )*
                { $catch_all }
            }};
        }

        fn main() {
            mat<|>ch_ast! {
                match container {
                    ast::TraitDef(it) => {},
                    ast::ImplBlock(it) => {},
                    _ => { continue },
                }
            }
        }
        "#,
        );

        assert_eq!(res.name, "match_ast");
        assert_snapshot!(res.expansion, @r###"
{
  if let Some(it) = ast::TraitDef::cast(container.clone()){}
  else if let Some(it) = ast::ImplBlock::cast(container.clone()){}
  else {
    {
      continue
    }
  }
}
"###);
    }

    #[test]
    fn macro_expand_match_ast_inside_let_statement() {
        let res = check_expand_macro(
            r#"
        //- /lib.rs
        macro_rules! match_ast {
            (match $node:ident { $($tt:tt)* }) => { match_ast!(match ($node) { $($tt)* }) };
            (match ($node:expr) {}) => {{}};
        }

        fn main() {
            let p = f(|it| {
                let res = mat<|>ch_ast! { match c {}};
                Some(res)
            })?;
        }
        "#,
        );

        assert_eq!(res.name, "match_ast");
        assert_snapshot!(res.expansion, @r###"{}"###);
    }
}
<commit_msg>Improve fail case in expand_macro<commit_after>//! This modules implements "expand macro" functionality in the IDE

use crate::{db::RootDatabase, FilePosition};
use hir::db::AstDatabase;
use ra_db::SourceDatabase;
use rustc_hash::FxHashMap;

use ra_syntax::{
    algo::{find_node_at_offset, replace_descendants},
    ast::{self},
    AstNode, NodeOrToken, SyntaxKind, SyntaxNode, WalkEvent, T,
};

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

pub(crate) fn expand_macro(db: &RootDatabase, position: FilePosition) -> Option<ExpandedMacro> {
    let parse = db.parse(position.file_id);
    let file = parse.tree();
    let name_ref = find_node_at_offset::<ast::NameRef>(file.syntax(), position.offset)?;
    let mac = name_ref.syntax().ancestors().find_map(ast::MacroCall::cast)?;

    let source = hir::Source::new(position.file_id.into(), mac.syntax());
    let expanded = expand_macro_recur(db, source, source.with_value(&mac))?;

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

fn expand_macro_recur(
    db: &RootDatabase,
    source: hir::Source<&SyntaxNode>,
    macro_call: hir::Source<&ast::MacroCall>,
) -> Option<SyntaxNode> {
    let analyzer = hir::SourceAnalyzer::new(db, source, None);
    let expansion = analyzer.expand(db, macro_call)?;
    let macro_file_id = expansion.file_id();
    let mut expanded: SyntaxNode = db.parse_or_expand(macro_file_id)?;

    let children = expanded.descendants().filter_map(ast::MacroCall::cast);
    let mut replaces = FxHashMap::default();

    for child in children.into_iter() {
        let node = hir::Source::new(macro_file_id, &child);
        if let Some(new_node) = expand_macro_recur(db, source, node) {
            // Replace the whole node if it is root
            // `replace_descendants` will not replace the parent node
            // but `SyntaxNode::descendants include itself
            if expanded == *child.syntax() {
                expanded = new_node;
            } else {
                replaces.insert(child.syntax().clone().into(), new_node.into());
            }
        }
    }

    Some(replace_descendants(&expanded, &replaces))
}

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

    let mut res = String::new();
    let mut token_iter = syn
        .preorder_with_tokens()
        .filter_map(|event| {
            if let WalkEvent::Enter(NodeOrToken::Token(token)) = event {
                Some(token)
            } else {
                None
            }
        })
        .peekable();

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

    while let Some(token) = token_iter.next() {
        let mut is_next = |f: fn(SyntaxKind) -> bool, default| -> bool {
            token_iter.peek().map(|it| f(it.kind())).unwrap_or(default)
        };
        let is_last = |f: fn(SyntaxKind) -> bool, default| -> bool {
            last.map(|it| f(it)).unwrap_or(default)
        };

        res += &match token.kind() {
            k @ _ if is_text(k) && is_next(|it| !it.is_punct(), true) => {
                token.text().to_string() + " "
            }
            L_CURLY if is_next(|it| it != R_CURLY, true) => {
                indent += 1;
                let leading_space = if is_last(|it| is_text(it), false) { " " } else { "" };
                format!("{}{{\n{}", leading_space, "  ".repeat(indent))
            }
            R_CURLY if is_last(|it| it != L_CURLY, true) => {
                indent = indent.checked_sub(1).unwrap_or(0);
                format!("\n{}}}", "  ".repeat(indent))
            }
            R_CURLY => format!("}}\n{}", "  ".repeat(indent)),
            T![;] => format!(";\n{}", "  ".repeat(indent)),
            T![->] => " -> ".to_string(),
            T![=] => " = ".to_string(),
            T![=>] => " => ".to_string(),
            _ => token.text().to_string(),
        };

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

    return res;

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mock_analysis::analysis_and_position;
    use insta::assert_snapshot;

    fn check_expand_macro(fixture: &str) -> ExpandedMacro {
        let (analysis, pos) = analysis_and_position(fixture);
        analysis.expand_macro(pos).unwrap().unwrap()
    }

    #[test]
    fn macro_expand_recursive_expansion() {
        let res = check_expand_macro(
            r#"
        //- /lib.rs
        macro_rules! bar {
            () => { fn  b() {} }
        }
        macro_rules! foo {
            () => { bar!(); }
        }
        macro_rules! baz {
            () => { foo!(); }
        }
        f<|>oo!();
        "#,
        );

        assert_eq!(res.name, "foo");
        assert_snapshot!(res.expansion, @r###"
fn b(){}
"###);
    }

    #[test]
    fn macro_expand_multiple_lines() {
        let res = check_expand_macro(
            r#"
        //- /lib.rs
        macro_rules! foo {
            () => {
                fn some_thing() -> u32 {
                    let a = 0;
                    a + 10
                }
            }
        }
        f<|>oo!();
        "#,
        );

        assert_eq!(res.name, "foo");
        assert_snapshot!(res.expansion, @r###"
fn some_thing() -> u32 {
  let a = 0;
  a+10
}
"###);
    }

    #[test]
    fn macro_expand_match_ast() {
        let res = check_expand_macro(
            r#"
        //- /lib.rs
        macro_rules! match_ast {
            (match $node:ident { $($tt:tt)* }) => { match_ast!(match ($node) { $($tt)* }) };

            (match ($node:expr) {
                $( ast::$ast:ident($it:ident) => $res:block, )*
                _ => $catch_all:expr $(,)?
            }) => {{
                $( if let Some($it) = ast::$ast::cast($node.clone()) $res else )*
                { $catch_all }
            }};
        }

        fn main() {
            mat<|>ch_ast! {
                match container {
                    ast::TraitDef(it) => {},
                    ast::ImplBlock(it) => {},
                    _ => { continue },
                }
            }
        }
        "#,
        );

        assert_eq!(res.name, "match_ast");
        assert_snapshot!(res.expansion, @r###"
{
  if let Some(it) = ast::TraitDef::cast(container.clone()){}
  else if let Some(it) = ast::ImplBlock::cast(container.clone()){}
  else {
    {
      continue
    }
  }
}
"###);
    }

    #[test]
    fn macro_expand_match_ast_inside_let_statement() {
        let res = check_expand_macro(
            r#"
        //- /lib.rs
        macro_rules! match_ast {
            (match $node:ident { $($tt:tt)* }) => { match_ast!(match ($node) { $($tt)* }) };
            (match ($node:expr) {}) => {{}};
        }

        fn main() {
            let p = f(|it| {
                let res = mat<|>ch_ast! { match c {}};
                Some(res)
            })?;
        }
        "#,
        );

        assert_eq!(res.name, "match_ast");
        assert_snapshot!(res.expansion, @r###"{}"###);
    }

    #[test]
    fn macro_expand_inner_macro_fail_to_expand() {
        let res = check_expand_macro(
            r#"
        //- /lib.rs
        macro_rules! bar {
            (BAD) => {};
        }
        macro_rules! foo {
            () => {bar!()};
        }

        fn main() {
            let res = fo<|>o!();
        }
        "#,
        );

        assert_eq!(res.name, "foo");
        assert_snapshot!(res.expansion, @r###"bar!()"###);
    }
}
<|endoftext|>


<commit_before>/// Declares an item with a doc attribute computed by some macro expression.
/// This allows documentation to be dynamically generated based on input.
/// Necessary to work around <https://github.com/rust-lang/rust/issues/52607>.
macro_rules! doc_concat {
    ( $( #[doc = $doc:expr] $( $thing:tt )* )* ) => ( $( #[doc = $doc] $( $thing )* )* );
}

macro_rules! partial_eq_string {
    ($id:ty $([$( $g:ident ),*])?) => {
        partial_eq_string!(@imp $(<$($g),*>)?, $id, str);
        partial_eq_string!(@imp $(<$($g),*>)?, $id, &str);
        partial_eq_string!(@imp $(<$($g),*>)?, $id, String);
        partial_eq_string!(@imp $(<$($g),*>)?, str, $id);
        partial_eq_string!(@imp $(<$($g),*>)?, &str, $id);
        partial_eq_string!(@imp $(<$($g),*>)?, String, $id);
    };
    (@imp $(<$( $g:ident ),*>)?, $l:ty, $r:ty) => {
        impl $(<$($g),*>)? ::std::cmp::PartialEq<$r> for $l {
            fn eq(&self, other: &$r) -> bool {
                ::std::convert::AsRef::<str>::as_ref(self)
                    == ::std::convert::AsRef::<str>::as_ref(other)
            }
        }
    }
}

macro_rules! as_str_based_impls {
    ($id:ty) => {
        impl ::std::convert::AsRef<str> for $id {
            fn as_ref(&self) -> &str {
                self.as_str()
            }
        }

        impl ::std::fmt::Display for $id {
            fn fmt(&self, f: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result {
                write!(f, "{}", self.as_str())
            }
        }

        impl ::std::cmp::PartialEq for $id {
            fn eq(&self, other: &Self) -> bool {
                self.as_str() == other.as_str()
            }
        }

        impl ::std::cmp::Eq for $id {}

        impl ::std::cmp::PartialOrd for $id {
            fn partial_cmp(&self, other: &Self) -> Option<::std::cmp::Ordering> {
                ::std::cmp::PartialOrd::partial_cmp(self.as_str(), other.as_str())
            }
        }

        impl ::std::cmp::Ord for $id {
            fn cmp(&self, other: &Self) -> ::std::cmp::Ordering {
                ::std::cmp::Ord::cmp(self.as_str(), other.as_str())
            }
        }

        impl ::std::hash::Hash for $id {
            fn hash<H: ::std::hash::Hasher>(&self, state: &mut H) {
                self.as_str().hash(state);
            }
        }

        #[cfg(feature = "serde")]
        impl ::serde::Serialize for $id {
            fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
            where
                S: ::serde::Serializer,
            {
                serializer.serialize_str(self.as_str())
            }
        }
    };
}

macro_rules! common_impls {
    ($id:ty, $try_from:ident, $desc:literal) => {
        impl $id {
            doc_concat! {
                #[doc = concat!("Creates a string slice from this `", stringify!($id), "`")]
                pub fn as_str(&self) -> &str {
                    &self.full_id
                }
            }

            doc_concat! {
                #[doc = concat!("Creates a byte slice from this `", stringify!($id), "`")]
                pub fn as_bytes(&self) -> &[u8] {
                    self.full_id.as_bytes()
                }
            }

            doc_concat! {
                #[doc = concat!("Converts this `", stringify!($id), "` into a `String`")]
                pub fn into_string(self) -> String {
                    self.full_id.into()
                }
            }

            doc_concat! {
                #[doc = concat!("Converts this `", stringify!($id), "` into a `Vec<u8>`")]
                pub fn into_bytes(self) -> Vec<u8> {
                    Box::<[u8]>::from(self.full_id).into()
                }
            }
        }

        impl ::std::convert::From<$id> for ::std::string::String {
            fn from(id: $id) -> Self {
                id.into_string()
            }
        }

        impl ::std::convert::From<$id> for ::std::vec::Vec<u8> {
            fn from(id: $id) -> Self {
                id.into_bytes()
            }
        }

        impl ::std::str::FromStr for $id {
            type Err = crate::Error;

            fn from_str(s: &str) -> Result<Self, Self::Err> {
                $try_from(s)
            }
        }

        impl ::std::convert::TryFrom<&str> for $id {
            type Error = crate::Error;

            fn try_from(s: &str) -> Result<Self, Self::Error> {
                $try_from(s)
            }
        }

        impl ::std::convert::TryFrom<String> for $id {
            type Error = crate::Error;

            fn try_from(s: String) -> Result<Self, Self::Error> {
                $try_from(s)
            }
        }

        #[cfg(feature = "serde")]
        impl<'de> ::serde::Deserialize<'de> for $id {
            fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
            where
                D: ::serde::Deserializer<'de>,
            {
                crate::deserialize_id(deserializer, $desc)
            }
        }

        as_str_based_impls!($id);
        partial_eq_string!($id);
    };
}

macro_rules! opaque_identifier {
    (
        $( #[doc = $docs:literal] )*
        $vis:vis type $id:ident;
    ) => {
        $( #[doc = $docs] )*
        #[repr(transparent)]
        pub struct $id(str);

        ::paste::paste! {
            doc_concat! {
                #[doc = concat!("An owned [", stringify!($id), "].")]
                pub type [<$id Box>] = ::std::boxed::Box<$id>;
            }
        }

        impl $id {
            #[allow(clippy::transmute_ptr_to_ptr)]
            fn from_borrowed(s: &str) -> &Self {
                unsafe { ::std::mem::transmute(s) }
            }

            pub(super) fn from_owned(s: ::std::boxed::Box<str>) -> ::std::boxed::Box<Self> {
                unsafe { ::std::mem::transmute(s) }
            }

            fn into_owned(self: ::std::boxed::Box<Self>) -> ::std::boxed::Box<str> {
                unsafe { ::std::mem::transmute(self) }
            }

            doc_concat! {
                #[doc = concat!("Creates a string slice from this `", stringify!($id), "`.")]
                pub fn as_str(&self) -> &str {
                    &self.0
                }
            }

            doc_concat! {
                #[doc = concat!("Creates a byte slice from this `", stringify!($id), "`.")]
                pub fn as_bytes(&self) -> &[u8] {
                    self.0.as_bytes()
                }
            }
        }

        impl ::std::fmt::Debug for $id {
            fn fmt(&self, f: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result {
                f.write_str(self.as_ref())
            }
        }

        impl Clone for ::std::boxed::Box<$id> {
            fn clone(&self) -> Self {
                (**self).to_owned()
            }
        }

        impl ToOwned for $id {
            type Owned = ::std::boxed::Box<$id>;

            fn to_owned(&self) -> Self::Owned {
                Self::from_owned(self.0.to_owned().into_boxed_str())
            }
        }

        impl From<&$id> for ::std::boxed::Box<$id> {
            fn from(id: &$id) -> Self {
                id.to_owned()
            }
        }

        impl AsRef<str> for ::std::boxed::Box<$id> {
            fn as_ref(&self) -> &str {
                self.as_str()
            }
        }

        impl<'a> From<&'a str> for &'a $id {
            fn from(s: &'a str) -> Self {
                $id::from_borrowed(s)
            }
        }

        impl From<&str> for ::std::boxed::Box<$id> {
            fn from(s: &str) -> Self {
                $id::from_owned(s.into())
            }
        }

        impl From<String> for ::std::boxed::Box<$id> {
            fn from(s: String) -> Self {
                $id::from_owned(s.into())
            }
        }

        impl From<::std::boxed::Box<$id>> for String {
            fn from(id: ::std::boxed::Box<$id>) -> Self {
                id.into_owned().into()
            }
        }

        #[cfg(feature = "serde")]
        impl<'de> ::serde::Deserialize<'de> for ::std::boxed::Box<$id> {
            fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
            where
                D: ::serde::Deserializer<'de>,
            {
                ::std::boxed::Box::<str>::deserialize(deserializer).map($id::from_owned)
            }
        }

        as_str_based_impls!($id);
        partial_eq_string!($id);
        partial_eq_string!(::std::boxed::Box<$id>);
    };
}
<commit_msg>identifiers: Don't avoid prelude in hygienic crate-internal macros<commit_after>/// Declares an item with a doc attribute computed by some macro expression.
/// This allows documentation to be dynamically generated based on input.
/// Necessary to work around <https://github.com/rust-lang/rust/issues/52607>.
macro_rules! doc_concat {
    ( $( #[doc = $doc:expr] $( $thing:tt )* )* ) => ( $( #[doc = $doc] $( $thing )* )* );
}

macro_rules! partial_eq_string {
    ($id:ty $([$( $g:ident ),*])?) => {
        partial_eq_string!(@imp $(<$($g),*>)?, $id, str);
        partial_eq_string!(@imp $(<$($g),*>)?, $id, &str);
        partial_eq_string!(@imp $(<$($g),*>)?, $id, String);
        partial_eq_string!(@imp $(<$($g),*>)?, str, $id);
        partial_eq_string!(@imp $(<$($g),*>)?, &str, $id);
        partial_eq_string!(@imp $(<$($g),*>)?, String, $id);
    };
    (@imp $(<$( $g:ident ),*>)?, $l:ty, $r:ty) => {
        impl $(<$($g),*>)? PartialEq<$r> for $l {
            fn eq(&self, other: &$r) -> bool {
                AsRef::<str>::as_ref(self)
                    == AsRef::<str>::as_ref(other)
            }
        }
    }
}

macro_rules! as_str_based_impls {
    ($id:ty) => {
        impl AsRef<str> for $id {
            fn as_ref(&self) -> &str {
                self.as_str()
            }
        }

        impl std::fmt::Display for $id {
            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                write!(f, "{}", self.as_str())
            }
        }

        impl PartialEq for $id {
            fn eq(&self, other: &Self) -> bool {
                self.as_str() == other.as_str()
            }
        }

        impl std::cmp::Eq for $id {}

        impl std::cmp::PartialOrd for $id {
            fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
                std::cmp::PartialOrd::partial_cmp(self.as_str(), other.as_str())
            }
        }

        impl std::cmp::Ord for $id {
            fn cmp(&self, other: &Self) -> std::cmp::Ordering {
                std::cmp::Ord::cmp(self.as_str(), other.as_str())
            }
        }

        impl std::hash::Hash for $id {
            fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
                self.as_str().hash(state);
            }
        }

        #[cfg(feature = "serde")]
        impl serde::Serialize for $id {
            fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
            where
                S: serde::Serializer,
            {
                serializer.serialize_str(self.as_str())
            }
        }
    };
}

macro_rules! common_impls {
    ($id:ty, $try_from:ident, $desc:literal) => {
        impl $id {
            doc_concat! {
                #[doc = concat!("Creates a string slice from this `", stringify!($id), "`")]
                pub fn as_str(&self) -> &str {
                    &self.full_id
                }
            }

            doc_concat! {
                #[doc = concat!("Creates a byte slice from this `", stringify!($id), "`")]
                pub fn as_bytes(&self) -> &[u8] {
                    self.full_id.as_bytes()
                }
            }

            doc_concat! {
                #[doc = concat!("Converts this `", stringify!($id), "` into a `String`")]
                pub fn into_string(self) -> String {
                    self.full_id.into()
                }
            }

            doc_concat! {
                #[doc = concat!("Converts this `", stringify!($id), "` into a `Vec<u8>`")]
                pub fn into_bytes(self) -> Vec<u8> {
                    Box::<[u8]>::from(self.full_id).into()
                }
            }
        }

        impl From<$id> for String {
            fn from(id: $id) -> Self {
                id.into_string()
            }
        }

        impl From<$id> for Vec<u8> {
            fn from(id: $id) -> Self {
                id.into_bytes()
            }
        }

        impl std::str::FromStr for $id {
            type Err = crate::Error;

            fn from_str(s: &str) -> Result<Self, Self::Err> {
                $try_from(s)
            }
        }

        impl std::convert::TryFrom<&str> for $id {
            type Error = crate::Error;

            fn try_from(s: &str) -> Result<Self, Self::Error> {
                $try_from(s)
            }
        }

        impl std::convert::TryFrom<String> for $id {
            type Error = crate::Error;

            fn try_from(s: String) -> Result<Self, Self::Error> {
                $try_from(s)
            }
        }

        #[cfg(feature = "serde")]
        impl<'de> serde::Deserialize<'de> for $id {
            fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
            where
                D: serde::Deserializer<'de>,
            {
                crate::deserialize_id(deserializer, $desc)
            }
        }

        as_str_based_impls!($id);
        partial_eq_string!($id);
    };
}

macro_rules! opaque_identifier {
    (
        $( #[doc = $docs:literal] )*
        $vis:vis type $id:ident;
    ) => {
        $( #[doc = $docs] )*
        #[repr(transparent)]
        pub struct $id(str);

        paste::paste! {
            doc_concat! {
                #[doc = concat!("An owned [", stringify!($id), "].")]
                pub type [<$id Box>] = Box<$id>;
            }
        }

        impl $id {
            #[allow(clippy::transmute_ptr_to_ptr)]
            fn from_borrowed(s: &str) -> &Self {
                unsafe { std::mem::transmute(s) }
            }

            pub(super) fn from_owned(s: Box<str>) -> Box<Self> {
                unsafe { std::mem::transmute(s) }
            }

            fn into_owned(self: Box<Self>) -> Box<str> {
                unsafe { std::mem::transmute(self) }
            }

            doc_concat! {
                #[doc = concat!("Creates a string slice from this `", stringify!($id), "`.")]
                pub fn as_str(&self) -> &str {
                    &self.0
                }
            }

            doc_concat! {
                #[doc = concat!("Creates a byte slice from this `", stringify!($id), "`.")]
                pub fn as_bytes(&self) -> &[u8] {
                    self.0.as_bytes()
                }
            }
        }

        impl std::fmt::Debug for $id {
            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                f.write_str(self.as_ref())
            }
        }

        impl Clone for Box<$id> {
            fn clone(&self) -> Self {
                (**self).to_owned()
            }
        }

        impl ToOwned for $id {
            type Owned = Box<$id>;

            fn to_owned(&self) -> Self::Owned {
                Self::from_owned(self.0.to_owned().into_boxed_str())
            }
        }

        impl From<&$id> for Box<$id> {
            fn from(id: &$id) -> Self {
                id.to_owned()
            }
        }

        impl AsRef<str> for Box<$id> {
            fn as_ref(&self) -> &str {
                self.as_str()
            }
        }

        impl<'a> From<&'a str> for &'a $id {
            fn from(s: &'a str) -> Self {
                $id::from_borrowed(s)
            }
        }

        impl From<&str> for Box<$id> {
            fn from(s: &str) -> Self {
                $id::from_owned(s.into())
            }
        }

        impl From<String> for Box<$id> {
            fn from(s: String) -> Self {
                $id::from_owned(s.into())
            }
        }

        impl From<Box<$id>> for String {
            fn from(id: Box<$id>) -> Self {
                id.into_owned().into()
            }
        }

        #[cfg(feature = "serde")]
        impl<'de> serde::Deserialize<'de> for Box<$id> {
            fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
            where
                D: serde::Deserializer<'de>,
            {
                Box::<str>::deserialize(deserializer).map($id::from_owned)
            }
        }

        as_str_based_impls!($id);
        partial_eq_string!($id);
        partial_eq_string!(Box<$id>);
    };
}
<|endoftext|>


<commit_before>// Copyright 2019 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::collections::BTreeMap;

use super::constants::*;
use super::virtio_input_absinfo;
use super::virtio_input_bitmap;
use super::virtio_input_device_ids;
use super::VirtioInputConfig;

/// Instantiates a VirtioInputConfig object with the default configuration for a trackpad. It
/// supports touch, left button and right button events, as well as X and Y axis.
pub fn new_trackpad_config(width: u32, height: u32) -> VirtioInputConfig {
    VirtioInputConfig::new(
        virtio_input_device_ids::new(0, 0, 0, 0),
        b"Crosvm Virtio Trackpad".to_vec(),
        b"virtio-trackpad".to_vec(),
        virtio_input_bitmap::new([0u8; 128]),
        default_trackpad_events(),
        default_trackpad_absinfo(width, height),
    )
}

/// Instantiates a VirtioInputConfig object with the default configuration for a mouse.
/// It supports left, right and middle buttons, as wel as X, Y and wheel relative axes.
pub fn new_mouse_config() -> VirtioInputConfig {
    VirtioInputConfig::new(
        virtio_input_device_ids::new(0, 0, 0, 0),
        b"Crosvm Virtio Mouse".to_vec(),
        b"virtio-mouse".to_vec(),
        virtio_input_bitmap::new([0u8; 128]),
        default_mouse_events(),
        BTreeMap::new(),
    )
}

/// Instantiates a VirtioInputConfig object with the default configuration for a keyboard.
/// It supports the same keys as a en-us keyboard and the CAPSLOCK, NUMLOCK and SCROLLLOCK leds.
pub fn new_keyboard_config() -> VirtioInputConfig {
    VirtioInputConfig::new(
        virtio_input_device_ids::new(0, 0, 0, 0),
        b"Crosvm Virtio Keyboard".to_vec(),
        b"virtio-keyboard".to_vec(),
        virtio_input_bitmap::new([0u8; 128]),
        default_keyboard_events(),
        BTreeMap::new(),
    )
}

/// Instantiates a VirtioInputConfig object with the default configuration for a touchscreen (no
/// multitouch support).
pub fn new_single_touch_config(width: u32, height: u32) -> VirtioInputConfig {
    VirtioInputConfig::new(
        virtio_input_device_ids::new(0, 0, 0, 0),
        b"Crosvm Virtio Touchscreen".to_vec(),
        b"virtio-touchscreen".to_vec(),
        virtio_input_bitmap::from_bits(&[INPUT_PROP_DIRECT]),
        default_touchscreen_events(),
        default_touchscreen_absinfo(width, height),
    )
}

/// Instantiates a VirtioInputConfig object with the default configuration for a multitouch
/// touchscreen.
pub fn new_multi_touch_config(width: u32, height: u32) -> VirtioInputConfig {
    VirtioInputConfig::new(
        virtio_input_device_ids::new(0, 0, 0, 0),
        b"Crosvm Virtio Multitouch Touchscreen".to_vec(),
        b"virtio-touchscreen".to_vec(),
        virtio_input_bitmap::from_bits(&[INPUT_PROP_DIRECT]),
        default_multitouchscreen_events(),
        default_multitouchscreen_absinfo(width, height, 0, 0),
    )
}

fn default_touchscreen_absinfo(width: u32, height: u32) -> BTreeMap<u16, virtio_input_absinfo> {
    let mut absinfo: BTreeMap<u16, virtio_input_absinfo> = BTreeMap::new();
    absinfo.insert(ABS_X, virtio_input_absinfo::new(0, width, 0, 0));
    absinfo.insert(ABS_Y, virtio_input_absinfo::new(0, height, 0, 0));
    absinfo
}

fn default_touchscreen_events() -> BTreeMap<u16, virtio_input_bitmap> {
    let mut supported_events: BTreeMap<u16, virtio_input_bitmap> = BTreeMap::new();
    supported_events.insert(EV_KEY, virtio_input_bitmap::from_bits(&[BTN_TOUCH]));
    supported_events.insert(EV_ABS, virtio_input_bitmap::from_bits(&[ABS_X, ABS_Y]));
    supported_events
}

fn default_multitouchscreen_absinfo(
    width: u32,
    height: u32,
    slot: u32,
    id: u32,
) -> BTreeMap<u16, virtio_input_absinfo> {
    let mut absinfo: BTreeMap<u16, virtio_input_absinfo> = BTreeMap::new();
    absinfo.insert(ABS_MT_SLOT, virtio_input_absinfo::new(0, slot, 0, 0));
    absinfo.insert(ABS_MT_TRACKING_ID, virtio_input_absinfo::new(0, id, 0, 0));
    absinfo.insert(ABS_MT_POSITION_X, virtio_input_absinfo::new(0, width, 0, 0));
    absinfo.insert(
        ABS_MT_POSITION_Y,
        virtio_input_absinfo::new(0, height, 0, 0),
    );
    absinfo
}

fn default_multitouchscreen_events() -> BTreeMap<u16, virtio_input_bitmap> {
    let mut supported_events: BTreeMap<u16, virtio_input_bitmap> = BTreeMap::new();
    supported_events.insert(EV_KEY, virtio_input_bitmap::from_bits(&[BTN_TOUCH]));
    supported_events.insert(
        EV_ABS,
        virtio_input_bitmap::from_bits(&[
            ABS_MT_SLOT,
            ABS_MT_TRACKING_ID,
            ABS_MT_POSITION_X,
            ABS_MT_POSITION_Y,
        ]),
    );
    supported_events
}

fn default_trackpad_absinfo(width: u32, height: u32) -> BTreeMap<u16, virtio_input_absinfo> {
    let mut absinfo: BTreeMap<u16, virtio_input_absinfo> = BTreeMap::new();
    absinfo.insert(ABS_X, virtio_input_absinfo::new(0, width, 0, 0));
    absinfo.insert(ABS_Y, virtio_input_absinfo::new(0, height, 0, 0));
    absinfo
}

fn default_trackpad_events() -> BTreeMap<u16, virtio_input_bitmap> {
    let mut supported_events: BTreeMap<u16, virtio_input_bitmap> = BTreeMap::new();
    supported_events.insert(
        EV_KEY,
        virtio_input_bitmap::from_bits(&[BTN_TOOL_FINGER, BTN_TOUCH, BTN_LEFT, BTN_RIGHT]),
    );
    supported_events.insert(EV_ABS, virtio_input_bitmap::from_bits(&[ABS_X, ABS_Y]));
    supported_events
}

fn default_mouse_events() -> BTreeMap<u16, virtio_input_bitmap> {
    let mut supported_events: BTreeMap<u16, virtio_input_bitmap> = BTreeMap::new();
    supported_events.insert(
        EV_KEY,
        virtio_input_bitmap::from_bits(&[BTN_LEFT, BTN_RIGHT, BTN_MIDDLE]),
    );
    supported_events.insert(
        EV_REL,
        virtio_input_bitmap::from_bits(&[REL_X, REL_Y, REL_WHEEL]),
    );
    supported_events
}

fn default_keyboard_events() -> BTreeMap<u16, virtio_input_bitmap> {
    let mut supported_events: BTreeMap<u16, virtio_input_bitmap> = BTreeMap::new();
    supported_events.insert(
        EV_KEY,
        virtio_input_bitmap::from_bits(&[
            KEY_ESC,
            KEY_1,
            KEY_2,
            KEY_3,
            KEY_4,
            KEY_5,
            KEY_6,
            KEY_7,
            KEY_8,
            KEY_9,
            KEY_0,
            KEY_MINUS,
            KEY_EQUAL,
            KEY_BACKSPACE,
            KEY_TAB,
            KEY_Q,
            KEY_W,
            KEY_E,
            KEY_R,
            KEY_T,
            KEY_Y,
            KEY_U,
            KEY_I,
            KEY_O,
            KEY_P,
            KEY_LEFTBRACE,
            KEY_RIGHTBRACE,
            KEY_ENTER,
            KEY_LEFTCTRL,
            KEY_A,
            KEY_S,
            KEY_D,
            KEY_F,
            KEY_G,
            KEY_H,
            KEY_J,
            KEY_K,
            KEY_L,
            KEY_SEMICOLON,
            KEY_APOSTROPHE,
            KEY_GRAVE,
            KEY_LEFTSHIFT,
            KEY_BACKSLASH,
            KEY_Z,
            KEY_X,
            KEY_C,
            KEY_V,
            KEY_B,
            KEY_N,
            KEY_M,
            KEY_COMMA,
            KEY_DOT,
            KEY_SLASH,
            KEY_RIGHTSHIFT,
            KEY_KPASTERISK,
            KEY_LEFTALT,
            KEY_SPACE,
            KEY_CAPSLOCK,
            KEY_F1,
            KEY_F2,
            KEY_F3,
            KEY_F4,
            KEY_F5,
            KEY_F6,
            KEY_F7,
            KEY_F8,
            KEY_F9,
            KEY_F10,
            KEY_NUMLOCK,
            KEY_SCROLLLOCK,
            KEY_KP7,
            KEY_KP8,
            KEY_KP9,
            KEY_KPMINUS,
            KEY_KP4,
            KEY_KP5,
            KEY_KP6,
            KEY_KPPLUS,
            KEY_KP1,
            KEY_KP2,
            KEY_KP3,
            KEY_KP0,
            KEY_KPDOT,
            KEY_F11,
            KEY_F12,
            KEY_KPENTER,
            KEY_RIGHTCTRL,
            KEY_KPSLASH,
            KEY_SYSRQ,
            KEY_RIGHTALT,
            KEY_HOME,
            KEY_UP,
            KEY_PAGEUP,
            KEY_LEFT,
            KEY_RIGHT,
            KEY_END,
            KEY_DOWN,
            KEY_PAGEDOWN,
            KEY_INSERT,
            KEY_DELETE,
            KEY_PAUSE,
            KEY_MENU,
            KEY_PRINT,
            KEY_POWER,
            KEY_HOMEPAGE,
            KEY_MUTE,
            KEY_VOLUMEDOWN,
            KEY_VOLUMEUP,
        ]),
    );
    supported_events.insert(
        EV_REP,
        virtio_input_bitmap::from_bits(&[REP_DELAY, REP_PERIOD]),
    );
    supported_events.insert(
        EV_LED,
        virtio_input_bitmap::from_bits(&[LED_CAPSL, LED_NUML, LED_SCROLLL]),
    );
    supported_events
}
<commit_msg>Fix max multitouch slot & tracking ID values.<commit_after>// Copyright 2019 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::collections::BTreeMap;

use super::constants::*;
use super::virtio_input_absinfo;
use super::virtio_input_bitmap;
use super::virtio_input_device_ids;
use super::VirtioInputConfig;

/// Instantiates a VirtioInputConfig object with the default configuration for a trackpad. It
/// supports touch, left button and right button events, as well as X and Y axis.
pub fn new_trackpad_config(width: u32, height: u32) -> VirtioInputConfig {
    VirtioInputConfig::new(
        virtio_input_device_ids::new(0, 0, 0, 0),
        b"Crosvm Virtio Trackpad".to_vec(),
        b"virtio-trackpad".to_vec(),
        virtio_input_bitmap::new([0u8; 128]),
        default_trackpad_events(),
        default_trackpad_absinfo(width, height),
    )
}

/// Instantiates a VirtioInputConfig object with the default configuration for a mouse.
/// It supports left, right and middle buttons, as wel as X, Y and wheel relative axes.
pub fn new_mouse_config() -> VirtioInputConfig {
    VirtioInputConfig::new(
        virtio_input_device_ids::new(0, 0, 0, 0),
        b"Crosvm Virtio Mouse".to_vec(),
        b"virtio-mouse".to_vec(),
        virtio_input_bitmap::new([0u8; 128]),
        default_mouse_events(),
        BTreeMap::new(),
    )
}

/// Instantiates a VirtioInputConfig object with the default configuration for a keyboard.
/// It supports the same keys as a en-us keyboard and the CAPSLOCK, NUMLOCK and SCROLLLOCK leds.
pub fn new_keyboard_config() -> VirtioInputConfig {
    VirtioInputConfig::new(
        virtio_input_device_ids::new(0, 0, 0, 0),
        b"Crosvm Virtio Keyboard".to_vec(),
        b"virtio-keyboard".to_vec(),
        virtio_input_bitmap::new([0u8; 128]),
        default_keyboard_events(),
        BTreeMap::new(),
    )
}

/// Instantiates a VirtioInputConfig object with the default configuration for a touchscreen (no
/// multitouch support).
pub fn new_single_touch_config(width: u32, height: u32) -> VirtioInputConfig {
    VirtioInputConfig::new(
        virtio_input_device_ids::new(0, 0, 0, 0),
        b"Crosvm Virtio Touchscreen".to_vec(),
        b"virtio-touchscreen".to_vec(),
        virtio_input_bitmap::from_bits(&[INPUT_PROP_DIRECT]),
        default_touchscreen_events(),
        default_touchscreen_absinfo(width, height),
    )
}

/// Instantiates a VirtioInputConfig object with the default configuration for a multitouch
/// touchscreen.
pub fn new_multi_touch_config(width: u32, height: u32) -> VirtioInputConfig {
    VirtioInputConfig::new(
        virtio_input_device_ids::new(0, 0, 0, 0),
        b"Crosvm Virtio Multitouch Touchscreen".to_vec(),
        b"virtio-touchscreen".to_vec(),
        virtio_input_bitmap::from_bits(&[INPUT_PROP_DIRECT]),
        default_multitouchscreen_events(),
        default_multitouchscreen_absinfo(width, height, 10, 10),
    )
}

fn default_touchscreen_absinfo(width: u32, height: u32) -> BTreeMap<u16, virtio_input_absinfo> {
    let mut absinfo: BTreeMap<u16, virtio_input_absinfo> = BTreeMap::new();
    absinfo.insert(ABS_X, virtio_input_absinfo::new(0, width, 0, 0));
    absinfo.insert(ABS_Y, virtio_input_absinfo::new(0, height, 0, 0));
    absinfo
}

fn default_touchscreen_events() -> BTreeMap<u16, virtio_input_bitmap> {
    let mut supported_events: BTreeMap<u16, virtio_input_bitmap> = BTreeMap::new();
    supported_events.insert(EV_KEY, virtio_input_bitmap::from_bits(&[BTN_TOUCH]));
    supported_events.insert(EV_ABS, virtio_input_bitmap::from_bits(&[ABS_X, ABS_Y]));
    supported_events
}

fn default_multitouchscreen_absinfo(
    width: u32,
    height: u32,
    slot: u32,
    id: u32,
) -> BTreeMap<u16, virtio_input_absinfo> {
    let mut absinfo: BTreeMap<u16, virtio_input_absinfo> = BTreeMap::new();
    absinfo.insert(ABS_MT_SLOT, virtio_input_absinfo::new(0, slot, 0, 0));
    absinfo.insert(ABS_MT_TRACKING_ID, virtio_input_absinfo::new(0, id, 0, 0));
    absinfo.insert(ABS_MT_POSITION_X, virtio_input_absinfo::new(0, width, 0, 0));
    absinfo.insert(
        ABS_MT_POSITION_Y,
        virtio_input_absinfo::new(0, height, 0, 0),
    );
    absinfo
}

fn default_multitouchscreen_events() -> BTreeMap<u16, virtio_input_bitmap> {
    let mut supported_events: BTreeMap<u16, virtio_input_bitmap> = BTreeMap::new();
    supported_events.insert(EV_KEY, virtio_input_bitmap::from_bits(&[BTN_TOUCH]));
    supported_events.insert(
        EV_ABS,
        virtio_input_bitmap::from_bits(&[
            ABS_MT_SLOT,
            ABS_MT_TRACKING_ID,
            ABS_MT_POSITION_X,
            ABS_MT_POSITION_Y,
        ]),
    );
    supported_events
}

fn default_trackpad_absinfo(width: u32, height: u32) -> BTreeMap<u16, virtio_input_absinfo> {
    let mut absinfo: BTreeMap<u16, virtio_input_absinfo> = BTreeMap::new();
    absinfo.insert(ABS_X, virtio_input_absinfo::new(0, width, 0, 0));
    absinfo.insert(ABS_Y, virtio_input_absinfo::new(0, height, 0, 0));
    absinfo
}

fn default_trackpad_events() -> BTreeMap<u16, virtio_input_bitmap> {
    let mut supported_events: BTreeMap<u16, virtio_input_bitmap> = BTreeMap::new();
    supported_events.insert(
        EV_KEY,
        virtio_input_bitmap::from_bits(&[BTN_TOOL_FINGER, BTN_TOUCH, BTN_LEFT, BTN_RIGHT]),
    );
    supported_events.insert(EV_ABS, virtio_input_bitmap::from_bits(&[ABS_X, ABS_Y]));
    supported_events
}

fn default_mouse_events() -> BTreeMap<u16, virtio_input_bitmap> {
    let mut supported_events: BTreeMap<u16, virtio_input_bitmap> = BTreeMap::new();
    supported_events.insert(
        EV_KEY,
        virtio_input_bitmap::from_bits(&[BTN_LEFT, BTN_RIGHT, BTN_MIDDLE]),
    );
    supported_events.insert(
        EV_REL,
        virtio_input_bitmap::from_bits(&[REL_X, REL_Y, REL_WHEEL]),
    );
    supported_events
}

fn default_keyboard_events() -> BTreeMap<u16, virtio_input_bitmap> {
    let mut supported_events: BTreeMap<u16, virtio_input_bitmap> = BTreeMap::new();
    supported_events.insert(
        EV_KEY,
        virtio_input_bitmap::from_bits(&[
            KEY_ESC,
            KEY_1,
            KEY_2,
            KEY_3,
            KEY_4,
            KEY_5,
            KEY_6,
            KEY_7,
            KEY_8,
            KEY_9,
            KEY_0,
            KEY_MINUS,
            KEY_EQUAL,
            KEY_BACKSPACE,
            KEY_TAB,
            KEY_Q,
            KEY_W,
            KEY_E,
            KEY_R,
            KEY_T,
            KEY_Y,
            KEY_U,
            KEY_I,
            KEY_O,
            KEY_P,
            KEY_LEFTBRACE,
            KEY_RIGHTBRACE,
            KEY_ENTER,
            KEY_LEFTCTRL,
            KEY_A,
            KEY_S,
            KEY_D,
            KEY_F,
            KEY_G,
            KEY_H,
            KEY_J,
            KEY_K,
            KEY_L,
            KEY_SEMICOLON,
            KEY_APOSTROPHE,
            KEY_GRAVE,
            KEY_LEFTSHIFT,
            KEY_BACKSLASH,
            KEY_Z,
            KEY_X,
            KEY_C,
            KEY_V,
            KEY_B,
            KEY_N,
            KEY_M,
            KEY_COMMA,
            KEY_DOT,
            KEY_SLASH,
            KEY_RIGHTSHIFT,
            KEY_KPASTERISK,
            KEY_LEFTALT,
            KEY_SPACE,
            KEY_CAPSLOCK,
            KEY_F1,
            KEY_F2,
            KEY_F3,
            KEY_F4,
            KEY_F5,
            KEY_F6,
            KEY_F7,
            KEY_F8,
            KEY_F9,
            KEY_F10,
            KEY_NUMLOCK,
            KEY_SCROLLLOCK,
            KEY_KP7,
            KEY_KP8,
            KEY_KP9,
            KEY_KPMINUS,
            KEY_KP4,
            KEY_KP5,
            KEY_KP6,
            KEY_KPPLUS,
            KEY_KP1,
            KEY_KP2,
            KEY_KP3,
            KEY_KP0,
            KEY_KPDOT,
            KEY_F11,
            KEY_F12,
            KEY_KPENTER,
            KEY_RIGHTCTRL,
            KEY_KPSLASH,
            KEY_SYSRQ,
            KEY_RIGHTALT,
            KEY_HOME,
            KEY_UP,
            KEY_PAGEUP,
            KEY_LEFT,
            KEY_RIGHT,
            KEY_END,
            KEY_DOWN,
            KEY_PAGEDOWN,
            KEY_INSERT,
            KEY_DELETE,
            KEY_PAUSE,
            KEY_MENU,
            KEY_PRINT,
            KEY_POWER,
            KEY_HOMEPAGE,
            KEY_MUTE,
            KEY_VOLUMEDOWN,
            KEY_VOLUMEUP,
        ]),
    );
    supported_events.insert(
        EV_REP,
        virtio_input_bitmap::from_bits(&[REP_DELAY, REP_PERIOD]),
    );
    supported_events.insert(
        EV_LED,
        virtio_input_bitmap::from_bits(&[LED_CAPSL, LED_NUML, LED_SCROLLL]),
    );
    supported_events
}
<|endoftext|>


<commit_before>use hir::ModuleDef;
use ide_db::helpers::{import_assets::NameToImport, mod_path_to_ast};
use ide_db::items_locator;
use itertools::Itertools;
use syntax::{
    ast::{self, make, AstNode, NameOwner},
    SyntaxKind::{IDENT, WHITESPACE},
    TextSize,
};

use crate::{
    assist_context::{AssistBuilder, AssistContext, Assists},
    utils::{
        add_trait_assoc_items_to_impl, filter_assoc_items, generate_trait_impl_text,
        render_snippet, Cursor, DefaultMethods,
    },
    AssistId, AssistKind,
};

// Assist: replace_derive_with_manual_impl
//
// Converts a `derive` impl into a manual one.
//
// ```
// # trait Debug { fn fmt(&self, f: &mut Formatter) -> Result<()>; }
// #[derive(Deb$0ug, Display)]
// struct S;
// ```
// ->
// ```
// # trait Debug { fn fmt(&self, f: &mut Formatter) -> Result<()>; }
// #[derive(Display)]
// struct S;
//
// impl Debug for S {
//     fn fmt(&self, f: &mut Formatter) -> Result<()> {
//         ${0:todo!()}
//     }
// }
// ```
pub(crate) fn replace_derive_with_manual_impl(
    acc: &mut Assists,
    ctx: &AssistContext,
) -> Option<()> {
    let attr = ctx.find_node_at_offset::<ast::Attr>()?;

    let has_derive = attr
        .syntax()
        .descendants_with_tokens()
        .filter(|t| t.kind() == IDENT)
        .find_map(syntax::NodeOrToken::into_token)
        .filter(|t| t.text() == "derive")
        .is_some();
    if !has_derive {
        return None;
    }

    let trait_token = ctx.token_at_offset().find(|t| t.kind() == IDENT && t.text() != "derive")?;
    let trait_path = make::path_unqualified(make::path_segment(make::name_ref(trait_token.text())));

    let adt = attr.syntax().parent().and_then(ast::Adt::cast)?;
    let annotated_name = adt.name()?;
    let insert_pos = adt.syntax().text_range().end();

    let current_module = ctx.sema.scope(annotated_name.syntax()).module()?;
    let current_crate = current_module.krate();

    let found_traits = items_locator::items_with_name(
        &ctx.sema,
        current_crate,
        NameToImport::Exact(trait_token.text().to_string()),
        items_locator::AssocItemSearch::Exclude,
        Some(items_locator::DEFAULT_QUERY_SEARCH_LIMIT),
    )
    .filter_map(|item| match ModuleDef::from(item.as_module_def_id()?) {
        ModuleDef::Trait(trait_) => Some(trait_),
        _ => None,
    })
    .flat_map(|trait_| {
        current_module
            .find_use_path(ctx.sema.db, hir::ModuleDef::Trait(trait_))
            .as_ref()
            .map(mod_path_to_ast)
            .zip(Some(trait_))
    });

    let mut no_traits_found = true;
    for (trait_path, trait_) in found_traits.inspect(|_| no_traits_found = false) {
        add_assist(acc, ctx, &attr, &trait_path, Some(trait_), &adt, &annotated_name, insert_pos)?;
    }
    if no_traits_found {
        add_assist(acc, ctx, &attr, &trait_path, None, &adt, &annotated_name, insert_pos)?;
    }
    Some(())
}

fn add_assist(
    acc: &mut Assists,
    ctx: &AssistContext,
    attr: &ast::Attr,
    trait_path: &ast::Path,
    trait_: Option<hir::Trait>,
    adt: &ast::Adt,
    annotated_name: &ast::Name,
    insert_pos: TextSize,
) -> Option<()> {
    let target = attr.syntax().text_range();
    let input = attr.token_tree()?;
    let label = format!("Convert to manual  `impl {} for {}`", trait_path, annotated_name);
    let trait_name = trait_path.segment().and_then(|seg| seg.name_ref())?;

    acc.add(
        AssistId("replace_derive_with_manual_impl", AssistKind::Refactor),
        label,
        target,
        |builder| {
            let impl_def_with_items =
                impl_def_from_trait(&ctx.sema, annotated_name, trait_, trait_path);
            update_attribute(builder, &input, &trait_name, &attr);
            let trait_path = format!("{}", trait_path);
            match (ctx.config.snippet_cap, impl_def_with_items) {
                (None, _) => {
                    builder.insert(insert_pos, generate_trait_impl_text(adt, &trait_path, ""))
                }
                (Some(cap), None) => builder.insert_snippet(
                    cap,
                    insert_pos,
                    generate_trait_impl_text(adt, &trait_path, "    $0"),
                ),
                (Some(cap), Some((impl_def, first_assoc_item))) => {
                    let mut cursor = Cursor::Before(first_assoc_item.syntax());
                    let placeholder;
                    if let ast::AssocItem::Fn(ref func) = first_assoc_item {
                        if let Some(m) = func.syntax().descendants().find_map(ast::MacroCall::cast)
                        {
                            if m.syntax().text() == "todo!()" {
                                placeholder = m;
                                cursor = Cursor::Replace(placeholder.syntax());
                            }
                        }
                    }

                    builder.insert_snippet(
                        cap,
                        insert_pos,
                        format!("\n\n{}", render_snippet(cap, impl_def.syntax(), cursor)),
                    )
                }
            };
        },
    )
}

fn impl_def_from_trait(
    sema: &hir::Semantics<ide_db::RootDatabase>,
    annotated_name: &ast::Name,
    trait_: Option<hir::Trait>,
    trait_path: &ast::Path,
) -> Option<(ast::Impl, ast::AssocItem)> {
    let trait_ = trait_?;
    let target_scope = sema.scope(annotated_name.syntax());
    let trait_items = filter_assoc_items(sema.db, &trait_.items(sema.db), DefaultMethods::No);
    if trait_items.is_empty() {
        return None;
    }
    let impl_def = make::impl_trait(
        trait_path.clone(),
        make::path_unqualified(make::path_segment(make::name_ref(&annotated_name.text()))),
    );
    let (impl_def, first_assoc_item) =
        add_trait_assoc_items_to_impl(sema, trait_items, trait_, impl_def, target_scope);
    Some((impl_def, first_assoc_item))
}

fn update_attribute(
    builder: &mut AssistBuilder,
    input: &ast::TokenTree,
    trait_name: &ast::NameRef,
    attr: &ast::Attr,
) {
    let trait_name = trait_name.text();
    let new_attr_input = input
        .syntax()
        .descendants_with_tokens()
        .filter(|t| t.kind() == IDENT)
        .filter_map(|t| t.into_token().map(|t| t.text().to_string()))
        .filter(|t| t != &trait_name)
        .collect::<Vec<_>>();
    let has_more_derives = !new_attr_input.is_empty();

    if has_more_derives {
        let new_attr_input = format!("({})", new_attr_input.iter().format(", "));
        builder.replace(input.syntax().text_range(), new_attr_input);
    } else {
        let attr_range = attr.syntax().text_range();
        builder.delete(attr_range);

        if let Some(line_break_range) = attr
            .syntax()
            .next_sibling_or_token()
            .filter(|t| t.kind() == WHITESPACE)
            .map(|t| t.text_range())
        {
            builder.delete(line_break_range);
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn add_custom_impl_debug() {
        check_assist(
            replace_derive_with_manual_impl,
            "
mod fmt {
    pub struct Error;
    pub type Result = Result<(), Error>;
    pub struct Formatter<'a>;
    pub trait Debug {
        fn fmt(&self, f: &mut Formatter<'_>) -> Result;
    }
}

#[derive(Debu$0g)]
struct Foo {
    bar: String,
}
",
            "
mod fmt {
    pub struct Error;
    pub type Result = Result<(), Error>;
    pub struct Formatter<'a>;
    pub trait Debug {
        fn fmt(&self, f: &mut Formatter<'_>) -> Result;
    }
}

struct Foo {
    bar: String,
}

impl fmt::Debug for Foo {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        ${0:todo!()}
    }
}
",
        )
    }
    #[test]
    fn add_custom_impl_all() {
        check_assist(
            replace_derive_with_manual_impl,
            "
mod foo {
    pub trait Bar {
        type Qux;
        const Baz: usize = 42;
        const Fez: usize;
        fn foo();
        fn bar() {}
    }
}

#[derive($0Bar)]
struct Foo {
    bar: String,
}
",
            "
mod foo {
    pub trait Bar {
        type Qux;
        const Baz: usize = 42;
        const Fez: usize;
        fn foo();
        fn bar() {}
    }
}

struct Foo {
    bar: String,
}

impl foo::Bar for Foo {
    $0type Qux;

    const Baz: usize = 42;

    const Fez: usize;

    fn foo() {
        todo!()
    }
}
",
        )
    }
    #[test]
    fn add_custom_impl_for_unique_input() {
        check_assist(
            replace_derive_with_manual_impl,
            "
#[derive(Debu$0g)]
struct Foo {
    bar: String,
}
            ",
            "
struct Foo {
    bar: String,
}

impl Debug for Foo {
    $0
}
            ",
        )
    }

    #[test]
    fn add_custom_impl_for_with_visibility_modifier() {
        check_assist(
            replace_derive_with_manual_impl,
            "
#[derive(Debug$0)]
pub struct Foo {
    bar: String,
}
            ",
            "
pub struct Foo {
    bar: String,
}

impl Debug for Foo {
    $0
}
            ",
        )
    }

    #[test]
    fn add_custom_impl_when_multiple_inputs() {
        check_assist(
            replace_derive_with_manual_impl,
            "
#[derive(Display, Debug$0, Serialize)]
struct Foo {}
            ",
            "
#[derive(Display, Serialize)]
struct Foo {}

impl Debug for Foo {
    $0
}
            ",
        )
    }

    #[test]
    fn test_ignore_derive_macro_without_input() {
        check_assist_not_applicable(
            replace_derive_with_manual_impl,
            "
#[derive($0)]
struct Foo {}
            ",
        )
    }

    #[test]
    fn test_ignore_if_cursor_on_param() {
        check_assist_not_applicable(
            replace_derive_with_manual_impl,
            "
#[derive$0(Debug)]
struct Foo {}
            ",
        );

        check_assist_not_applicable(
            replace_derive_with_manual_impl,
            "
#[derive(Debug)$0]
struct Foo {}
            ",
        )
    }

    #[test]
    fn test_ignore_if_not_derive() {
        check_assist_not_applicable(
            replace_derive_with_manual_impl,
            "
#[allow(non_camel_$0case_types)]
struct Foo {}
            ",
        )
    }
}
<commit_msg>Simplify<commit_after>use hir::ModuleDef;
use ide_db::helpers::{import_assets::NameToImport, mod_path_to_ast};
use ide_db::items_locator;
use itertools::Itertools;
use syntax::{
    ast::{self, make, AstNode, NameOwner},
    SyntaxKind::{IDENT, WHITESPACE},
};

use crate::{
    assist_context::{AssistBuilder, AssistContext, Assists},
    utils::{
        add_trait_assoc_items_to_impl, filter_assoc_items, generate_trait_impl_text,
        render_snippet, Cursor, DefaultMethods,
    },
    AssistId, AssistKind,
};

// Assist: replace_derive_with_manual_impl
//
// Converts a `derive` impl into a manual one.
//
// ```
// # trait Debug { fn fmt(&self, f: &mut Formatter) -> Result<()>; }
// #[derive(Deb$0ug, Display)]
// struct S;
// ```
// ->
// ```
// # trait Debug { fn fmt(&self, f: &mut Formatter) -> Result<()>; }
// #[derive(Display)]
// struct S;
//
// impl Debug for S {
//     fn fmt(&self, f: &mut Formatter) -> Result<()> {
//         ${0:todo!()}
//     }
// }
// ```
pub(crate) fn replace_derive_with_manual_impl(
    acc: &mut Assists,
    ctx: &AssistContext,
) -> Option<()> {
    let attr = ctx.find_node_at_offset::<ast::Attr>()?;
    let (name, args) = attr.as_simple_call()?;
    if name != "derive" {
        return None;
    }

    let trait_token = args.syntax().token_at_offset(ctx.offset()).find(|t| t.kind() == IDENT)?;
    let trait_name = trait_token.text();

    let adt = attr.syntax().parent().and_then(ast::Adt::cast)?;

    let current_module = ctx.sema.scope(adt.syntax()).module()?;
    let current_crate = current_module.krate();

    let found_traits = items_locator::items_with_name(
        &ctx.sema,
        current_crate,
        NameToImport::Exact(trait_name.to_string()),
        items_locator::AssocItemSearch::Exclude,
        Some(items_locator::DEFAULT_QUERY_SEARCH_LIMIT),
    )
    .filter_map(|item| match ModuleDef::from(item.as_module_def_id()?) {
        ModuleDef::Trait(trait_) => Some(trait_),
        _ => None,
    })
    .flat_map(|trait_| {
        current_module
            .find_use_path(ctx.sema.db, hir::ModuleDef::Trait(trait_))
            .as_ref()
            .map(mod_path_to_ast)
            .zip(Some(trait_))
    });

    let mut no_traits_found = true;
    for (trait_path, trait_) in found_traits.inspect(|_| no_traits_found = false) {
        add_assist(acc, ctx, &attr, &args, &trait_path, Some(trait_), &adt)?;
    }
    if no_traits_found {
        let trait_path = make::path_unqualified(make::path_segment(make::name_ref(trait_name)));
        add_assist(acc, ctx, &attr, &args, &trait_path, None, &adt)?;
    }
    Some(())
}

fn add_assist(
    acc: &mut Assists,
    ctx: &AssistContext,
    attr: &ast::Attr,
    input: &ast::TokenTree,
    trait_path: &ast::Path,
    trait_: Option<hir::Trait>,
    adt: &ast::Adt,
) -> Option<()> {
    let target = attr.syntax().text_range();
    let annotated_name = adt.name()?;
    let label = format!("Convert to manual `impl {} for {}`", trait_path, annotated_name);
    let trait_name = trait_path.segment().and_then(|seg| seg.name_ref())?;

    acc.add(
        AssistId("replace_derive_with_manual_impl", AssistKind::Refactor),
        label,
        target,
        |builder| {
            let insert_pos = adt.syntax().text_range().end();
            let impl_def_with_items =
                impl_def_from_trait(&ctx.sema, &annotated_name, trait_, trait_path);
            update_attribute(builder, &input, &trait_name, &attr);
            let trait_path = format!("{}", trait_path);
            match (ctx.config.snippet_cap, impl_def_with_items) {
                (None, _) => {
                    builder.insert(insert_pos, generate_trait_impl_text(adt, &trait_path, ""))
                }
                (Some(cap), None) => builder.insert_snippet(
                    cap,
                    insert_pos,
                    generate_trait_impl_text(adt, &trait_path, "    $0"),
                ),
                (Some(cap), Some((impl_def, first_assoc_item))) => {
                    let mut cursor = Cursor::Before(first_assoc_item.syntax());
                    let placeholder;
                    if let ast::AssocItem::Fn(ref func) = first_assoc_item {
                        if let Some(m) = func.syntax().descendants().find_map(ast::MacroCall::cast)
                        {
                            if m.syntax().text() == "todo!()" {
                                placeholder = m;
                                cursor = Cursor::Replace(placeholder.syntax());
                            }
                        }
                    }

                    builder.insert_snippet(
                        cap,
                        insert_pos,
                        format!("\n\n{}", render_snippet(cap, impl_def.syntax(), cursor)),
                    )
                }
            };
        },
    )
}

fn impl_def_from_trait(
    sema: &hir::Semantics<ide_db::RootDatabase>,
    annotated_name: &ast::Name,
    trait_: Option<hir::Trait>,
    trait_path: &ast::Path,
) -> Option<(ast::Impl, ast::AssocItem)> {
    let trait_ = trait_?;
    let target_scope = sema.scope(annotated_name.syntax());
    let trait_items = filter_assoc_items(sema.db, &trait_.items(sema.db), DefaultMethods::No);
    if trait_items.is_empty() {
        return None;
    }
    let impl_def = make::impl_trait(
        trait_path.clone(),
        make::path_unqualified(make::path_segment(make::name_ref(&annotated_name.text()))),
    );
    let (impl_def, first_assoc_item) =
        add_trait_assoc_items_to_impl(sema, trait_items, trait_, impl_def, target_scope);
    Some((impl_def, first_assoc_item))
}

fn update_attribute(
    builder: &mut AssistBuilder,
    input: &ast::TokenTree,
    trait_name: &ast::NameRef,
    attr: &ast::Attr,
) {
    let trait_name = trait_name.text();
    let new_attr_input = input
        .syntax()
        .descendants_with_tokens()
        .filter(|t| t.kind() == IDENT)
        .filter_map(|t| t.into_token().map(|t| t.text().to_string()))
        .filter(|t| t != &trait_name)
        .collect::<Vec<_>>();
    let has_more_derives = !new_attr_input.is_empty();

    if has_more_derives {
        let new_attr_input = format!("({})", new_attr_input.iter().format(", "));
        builder.replace(input.syntax().text_range(), new_attr_input);
    } else {
        let attr_range = attr.syntax().text_range();
        builder.delete(attr_range);

        if let Some(line_break_range) = attr
            .syntax()
            .next_sibling_or_token()
            .filter(|t| t.kind() == WHITESPACE)
            .map(|t| t.text_range())
        {
            builder.delete(line_break_range);
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn add_custom_impl_debug() {
        check_assist(
            replace_derive_with_manual_impl,
            "
mod fmt {
    pub struct Error;
    pub type Result = Result<(), Error>;
    pub struct Formatter<'a>;
    pub trait Debug {
        fn fmt(&self, f: &mut Formatter<'_>) -> Result;
    }
}

#[derive(Debu$0g)]
struct Foo {
    bar: String,
}
",
            "
mod fmt {
    pub struct Error;
    pub type Result = Result<(), Error>;
    pub struct Formatter<'a>;
    pub trait Debug {
        fn fmt(&self, f: &mut Formatter<'_>) -> Result;
    }
}

struct Foo {
    bar: String,
}

impl fmt::Debug for Foo {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        ${0:todo!()}
    }
}
",
        )
    }
    #[test]
    fn add_custom_impl_all() {
        check_assist(
            replace_derive_with_manual_impl,
            "
mod foo {
    pub trait Bar {
        type Qux;
        const Baz: usize = 42;
        const Fez: usize;
        fn foo();
        fn bar() {}
    }
}

#[derive($0Bar)]
struct Foo {
    bar: String,
}
",
            "
mod foo {
    pub trait Bar {
        type Qux;
        const Baz: usize = 42;
        const Fez: usize;
        fn foo();
        fn bar() {}
    }
}

struct Foo {
    bar: String,
}

impl foo::Bar for Foo {
    $0type Qux;

    const Baz: usize = 42;

    const Fez: usize;

    fn foo() {
        todo!()
    }
}
",
        )
    }
    #[test]
    fn add_custom_impl_for_unique_input() {
        check_assist(
            replace_derive_with_manual_impl,
            "
#[derive(Debu$0g)]
struct Foo {
    bar: String,
}
            ",
            "
struct Foo {
    bar: String,
}

impl Debug for Foo {
    $0
}
            ",
        )
    }

    #[test]
    fn add_custom_impl_for_with_visibility_modifier() {
        check_assist(
            replace_derive_with_manual_impl,
            "
#[derive(Debug$0)]
pub struct Foo {
    bar: String,
}
            ",
            "
pub struct Foo {
    bar: String,
}

impl Debug for Foo {
    $0
}
            ",
        )
    }

    #[test]
    fn add_custom_impl_when_multiple_inputs() {
        check_assist(
            replace_derive_with_manual_impl,
            "
#[derive(Display, Debug$0, Serialize)]
struct Foo {}
            ",
            "
#[derive(Display, Serialize)]
struct Foo {}

impl Debug for Foo {
    $0
}
            ",
        )
    }

    #[test]
    fn test_ignore_derive_macro_without_input() {
        check_assist_not_applicable(
            replace_derive_with_manual_impl,
            "
#[derive($0)]
struct Foo {}
            ",
        )
    }

    #[test]
    fn test_ignore_if_cursor_on_param() {
        check_assist_not_applicable(
            replace_derive_with_manual_impl,
            "
#[derive$0(Debug)]
struct Foo {}
            ",
        );

        check_assist_not_applicable(
            replace_derive_with_manual_impl,
            "
#[derive(Debug)$0]
struct Foo {}
            ",
        )
    }

    #[test]
    fn test_ignore_if_not_derive() {
        check_assist_not_applicable(
            replace_derive_with_manual_impl,
            "
#[allow(non_camel_$0case_types)]
struct Foo {}
            ",
        )
    }
}
<|endoftext|>


<commit_before>use cipherstate::CipherState;
use error::NoiseError;
use handshakepattern::{Token, HandshakePattern};
use symmetricstate::SymmetricState;
use traits::{DH, Cipher, Hash, U8Array};

/// Noise handshake state.
///
/// # Panics
///
/// `HandshakeState` must be used correctly, or its methods will likely panic:
///
/// Keys required by the handshake pattern must be set;
///
/// `write_message` and `read_message` must be called in right turns;
///
/// `write_message` and `read_message` must not be called after `completed`.
pub struct HandshakeState<D: DH, C: Cipher, H: Hash> {
    symmetric: SymmetricState<C, H>,
    s: Option<D::Key>,
    e: Option<D::Key>,
    rs: Option<D::Pubkey>,
    re: Option<D::Pubkey>,
    is_initiator: bool,
    pattern: HandshakePattern,
    message_index: usize,
}

impl<D, C, H> HandshakeState<D, C, H>
    where D: DH,
          C: Cipher,
          H: Hash
{
    /// Get protocol name, e.g. Noise_IK_25519_ChaChaPoly_BLAKE2s.
    pub fn get_name(has_psk: bool, pattern_name: &str) -> String {
        format!("Noise{}_{}_{}_{}_{}",
                if has_psk { "PSK" } else { "" },
                pattern_name,
                D::name(),
                C::name(),
                H::name())
    }

    /// Initialize a handshake state.
    ///
    /// `HandshakeState` does not generate a new ephemeral key when seeing a `E` toekn (for now?).
    pub fn new(pattern: HandshakePattern,
               is_initiator: bool,
               prologue: &[u8],
               psk: Option<&[u8]>,
               s: Option<D::Key>,
               e: Option<D::Key>,
               rs: Option<D::Pubkey>,
               re: Option<D::Pubkey>)
               -> Self {
        let mut symmetric = SymmetricState::new(Self::get_name(psk.is_some(), pattern.get_name())
            .as_bytes());

        // Mix in prologue.
        symmetric.mix_hash(prologue);

        // Mix in pre-shared key.
        if let Some(psk) = psk {
            assert_eq!(psk.len(), 32);
            symmetric.mix_preshared_key(psk);
        }

        // Mix in static keys known ahead of time.
        for t in pattern.get_pre_i() {
            match *t {
                Token::S => {
                    if is_initiator {
                        symmetric.mix_hash(D::pubkey(s.as_ref().unwrap()).as_slice());
                    } else {
                        symmetric.mix_hash(rs.as_ref().unwrap().as_slice());
                    }
                }
                Token::E => {
                    if is_initiator {
                        let e = D::pubkey(e.as_ref().unwrap());
                        symmetric.mix_hash(e.as_slice());
                        if symmetric.has_preshared_key() {
                            symmetric.mix_key(e.as_slice());
                        }
                    } else {
                        let re = re.as_ref().unwrap().as_slice();
                        symmetric.mix_hash(re);
                        if symmetric.has_preshared_key() {
                            symmetric.mix_key(re);
                        }
                    }
                }
                _ => panic!("Unexpected token in pre message"),
            }
        }
        for t in pattern.get_pre_r() {
            match *t {
                Token::S => {
                    if is_initiator {
                        symmetric.mix_hash(rs.as_ref().unwrap().as_slice());
                    } else {
                        symmetric.mix_hash(D::pubkey(s.as_ref().unwrap()).as_slice());
                    }
                }
                Token::E => {
                    if is_initiator {
                        let re = re.as_ref().unwrap().as_slice();
                        symmetric.mix_hash(re);
                        if symmetric.has_preshared_key() {
                            symmetric.mix_key(re);
                        }
                    } else {
                        let e = D::pubkey(e.as_ref().unwrap());
                        symmetric.mix_hash(e.as_slice());
                        if symmetric.has_preshared_key() {
                            symmetric.mix_key(e.as_slice());
                        }
                    }
                }
                _ => panic!("Unexpected token in pre message"),
            }
        }

        HandshakeState {
            symmetric: symmetric,
            s: s,
            e: e,
            rs: rs,
            re: re,
            is_initiator: is_initiator,
            pattern: pattern,
            message_index: 0,
        }
    }

    /// Takes a payload and return a packet that you should send to the peer.
    pub fn write_message(&mut self, payload: &[u8]) -> Vec<u8> {
        // Check that it is our turn to send.
        assert!(self.message_index % 2 == if self.is_initiator { 0 } else { 1 });

        // Get the message pattern.
        // Clone to make the borrow check happy.
        let m = self.pattern.get_message_patterns()[self.message_index].clone();
        self.message_index += 1;

        let mut out = Vec::new();

        // Process tokens.
        for t in m {
            match t {
                Token::E => {
                    // Spec says that we should generate new ephemeral key, but that would make
                    // testing very difficult.
                    let e_pk = D::pubkey(self.e.as_ref().unwrap());
                    self.symmetric.mix_hash(e_pk.as_slice());
                    if self.symmetric.has_preshared_key() {
                        self.symmetric.mix_key(e_pk.as_slice());
                    }
                    out.extend_from_slice(e_pk.as_slice());
                }
                Token::S => {
                    let encrypted_s = self.symmetric
                        .encrypt_and_hash_vec(D::pubkey(self.s.as_ref().unwrap()).as_slice());
                    out.extend_from_slice(&encrypted_s);
                }
                t => self.perform_dh(t),
            }
        }

        let encrypted_payload = self.symmetric.encrypt_and_hash_vec(payload);
        out.extend_from_slice(&encrypted_payload);

        out
    }

    /// Update handshake state and get payload, given a packet.
    ///
    /// If the packet fails to decrypt, the whole HandshakeState may be in invalid state, and
    /// should not be used any more. Expect to `get_re` before falling back to `XXfallback`.
    pub fn read_message(&mut self, data: &[u8]) -> Result<Vec<u8>, NoiseError> {
        // Check that it is our turn to recv.
        assert!(self.message_index % 2 == if self.is_initiator { 1 } else { 0 });

        // Get the message pattern.
        let m = self.pattern.get_message_patterns()[self.message_index].clone();
        self.message_index += 1;

        let mut data = data;
        // Consume the next `n` bytes of data.
        let mut get = |n| if data.len() >= n {
            let ret = &data[..n];
            data = &data[n..];
            Ok(ret)
        } else {
            Err(NoiseError::TooShort)
        };

        // Process tokens.
        for t in m {
            match t {
                Token::E => {
                    let mut re = D::Pubkey::new();
                    re.as_mut().copy_from_slice(get(D::pub_len())?);
                    self.symmetric.mix_hash(re.as_slice());
                    if self.symmetric.has_preshared_key() {
                        self.symmetric.mix_key(re.as_slice());
                    }
                    self.re = Some(re);
                }
                Token::S => {
                    let temp = get(if self.symmetric.has_key() {
                        D::pub_len() + 16
                    } else {
                        D::pub_len()
                    })?;
                    let mut rs = D::Pubkey::new();
                    self.symmetric.decrypt_and_hash(temp, rs.as_mut())?;
                    self.rs = Some(rs);
                }
                t => self.perform_dh(t),
            }
        }

        Ok(self.symmetric.decrypt_and_hash_vec(data)?)
    }

    /// Whether handshake has completed.
    pub fn completed(&self) -> bool {
        self.message_index == self.pattern.get_message_patterns().len()
    }

    /// Get handshake hash. Useful for e.g., channel binding.
    ///
    /// Should be called after handshake is `completed()`.
    pub fn get_hash(&self) -> &[u8] {
        self.symmetric.get_hash()
    }

    /// Get ciphers that can be used to encrypt/decrypt furthur messages.
    /// The first `CiperState` is for initiator to responder, and the second for responder
    /// to initiator.
    ///
    /// Should be called after handshake is `completed()`.
    pub fn get_ciphers(&self) -> (CipherState<C>, CipherState<C>) {
        self.symmetric.split()
    }

    /// Get remote semi-ephemeral pubkey.
    ///
    /// Returns `None` if we do not know.
    ///
    /// Useful for noise-pipes.
    pub fn get_re(&self) -> Option<D::Pubkey> {
        self.re
    }

    fn perform_dh(&mut self, t: Token) {
        let dh = |a: Option<&D::Key>, b: Option<&D::Pubkey>| D::dh(a.unwrap(), b.unwrap());

        let k = match t {
            Token::EE => dh(self.e.as_ref(), self.re.as_ref()),
            Token::ES => {
                if self.is_initiator {
                    dh(self.e.as_ref(), self.rs.as_ref())
                } else {
                    dh(self.s.as_ref(), self.re.as_ref())
                }
            }
            Token::SE => {
                if self.is_initiator {
                    dh(self.s.as_ref(), self.re.as_ref())
                } else {
                    dh(self.e.as_ref(), self.rs.as_ref())
                }
            }
            Token::SS => dh(self.s.as_ref(), self.rs.as_ref()),
            _ => unreachable!(),
        };

        self.symmetric.mix_key(k.as_slice());
    }
}

/// Builder for `HandshakeState`.
pub struct HandshakeStateBuilder<'a, D: DH> {
    pattern: Option<HandshakePattern>,
    is_initiator: Option<bool>,
    prologue: Option<&'a [u8]>,
    psk: Option<&'a [u8]>,
    s: Option<D::Key>,
    e: Option<D::Key>,
    rs: Option<D::Pubkey>,
    re: Option<D::Pubkey>,
}

impl<'a, D> HandshakeStateBuilder<'a, D>
    where D: DH
{
    pub fn new() -> Self {
        HandshakeStateBuilder {
            pattern: None,
            is_initiator: None,
            prologue: None,
            psk: None,
            s: None,
            e: None,
            rs: None,
            re: None,
        }
    }

    pub fn set_pattern(&mut self, p: HandshakePattern) -> &mut Self {
        self.pattern = Some(p);
        self
    }

    pub fn set_is_initiator(&mut self, is: bool) -> &mut Self {
        self.is_initiator = Some(is);
        self
    }

    pub fn set_prologue(&mut self, prologue: &'a [u8]) -> &mut Self {
        self.prologue = Some(prologue);
        self
    }

    pub fn set_psk(&mut self, psk: &'a [u8]) -> &mut Self {
        self.psk = Some(psk);
        self
    }

    pub fn set_e(&mut self, e: &D::Key) -> &mut Self {
        self.e = Some(*e);
        self
    }

    pub fn set_s(&mut self, s: &D::Key) -> &mut Self {
        self.s = Some(*s);
        self
    }

    pub fn set_re(&mut self, re: &D::Pubkey) -> &mut Self {
        self.re = Some(*re);
        self
    }

    pub fn set_rs(&mut self, rs: &D::Pubkey) -> &mut Self {
        self.rs = Some(*rs);
        self
    }

    /// Build `HandshakeState`.
    ///
    /// # Panics
    ///
    /// `pattern`, `prologue` and `is_initiator` must be set.
    pub fn build_handshake_state<C, H>(&self) -> HandshakeState<D, C, H>
        where C: Cipher,
              H: Hash
    {
        HandshakeState::new(self.pattern.as_ref().unwrap().clone(),
                            self.is_initiator.unwrap(),
                            self.prologue.unwrap(),
                            self.psk,
                            self.s,
                            self.e,
                            self.rs,
                            self.re)
    }
}
<commit_msg>HandshakeState: add method to get remote static pubkey<commit_after>use cipherstate::CipherState;
use error::NoiseError;
use handshakepattern::{Token, HandshakePattern};
use symmetricstate::SymmetricState;
use traits::{DH, Cipher, Hash, U8Array};

/// Noise handshake state.
///
/// # Panics
///
/// `HandshakeState` must be used correctly, or its methods will likely panic:
///
/// Keys required by the handshake pattern must be set;
///
/// `write_message` and `read_message` must be called in right turns;
///
/// `write_message` and `read_message` must not be called after `completed`.
pub struct HandshakeState<D: DH, C: Cipher, H: Hash> {
    symmetric: SymmetricState<C, H>,
    s: Option<D::Key>,
    e: Option<D::Key>,
    rs: Option<D::Pubkey>,
    re: Option<D::Pubkey>,
    is_initiator: bool,
    pattern: HandshakePattern,
    message_index: usize,
}

impl<D, C, H> HandshakeState<D, C, H>
    where D: DH,
          C: Cipher,
          H: Hash
{
    /// Get protocol name, e.g. Noise_IK_25519_ChaChaPoly_BLAKE2s.
    pub fn get_name(has_psk: bool, pattern_name: &str) -> String {
        format!("Noise{}_{}_{}_{}_{}",
                if has_psk { "PSK" } else { "" },
                pattern_name,
                D::name(),
                C::name(),
                H::name())
    }

    /// Initialize a handshake state.
    ///
    /// `HandshakeState` does not generate a new ephemeral key when seeing a `E` toekn (for now?).
    pub fn new(pattern: HandshakePattern,
               is_initiator: bool,
               prologue: &[u8],
               psk: Option<&[u8]>,
               s: Option<D::Key>,
               e: Option<D::Key>,
               rs: Option<D::Pubkey>,
               re: Option<D::Pubkey>)
               -> Self {
        let mut symmetric = SymmetricState::new(Self::get_name(psk.is_some(), pattern.get_name())
            .as_bytes());

        // Mix in prologue.
        symmetric.mix_hash(prologue);

        // Mix in pre-shared key.
        if let Some(psk) = psk {
            assert_eq!(psk.len(), 32);
            symmetric.mix_preshared_key(psk);
        }

        // Mix in static keys known ahead of time.
        for t in pattern.get_pre_i() {
            match *t {
                Token::S => {
                    if is_initiator {
                        symmetric.mix_hash(D::pubkey(s.as_ref().unwrap()).as_slice());
                    } else {
                        symmetric.mix_hash(rs.as_ref().unwrap().as_slice());
                    }
                }
                Token::E => {
                    if is_initiator {
                        let e = D::pubkey(e.as_ref().unwrap());
                        symmetric.mix_hash(e.as_slice());
                        if symmetric.has_preshared_key() {
                            symmetric.mix_key(e.as_slice());
                        }
                    } else {
                        let re = re.as_ref().unwrap().as_slice();
                        symmetric.mix_hash(re);
                        if symmetric.has_preshared_key() {
                            symmetric.mix_key(re);
                        }
                    }
                }
                _ => panic!("Unexpected token in pre message"),
            }
        }
        for t in pattern.get_pre_r() {
            match *t {
                Token::S => {
                    if is_initiator {
                        symmetric.mix_hash(rs.as_ref().unwrap().as_slice());
                    } else {
                        symmetric.mix_hash(D::pubkey(s.as_ref().unwrap()).as_slice());
                    }
                }
                Token::E => {
                    if is_initiator {
                        let re = re.as_ref().unwrap().as_slice();
                        symmetric.mix_hash(re);
                        if symmetric.has_preshared_key() {
                            symmetric.mix_key(re);
                        }
                    } else {
                        let e = D::pubkey(e.as_ref().unwrap());
                        symmetric.mix_hash(e.as_slice());
                        if symmetric.has_preshared_key() {
                            symmetric.mix_key(e.as_slice());
                        }
                    }
                }
                _ => panic!("Unexpected token in pre message"),
            }
        }

        HandshakeState {
            symmetric: symmetric,
            s: s,
            e: e,
            rs: rs,
            re: re,
            is_initiator: is_initiator,
            pattern: pattern,
            message_index: 0,
        }
    }

    /// Takes a payload and return a packet that you should send to the peer.
    pub fn write_message(&mut self, payload: &[u8]) -> Vec<u8> {
        // Check that it is our turn to send.
        assert!(self.message_index % 2 == if self.is_initiator { 0 } else { 1 });

        // Get the message pattern.
        // Clone to make the borrow check happy.
        let m = self.pattern.get_message_patterns()[self.message_index].clone();
        self.message_index += 1;

        let mut out = Vec::new();

        // Process tokens.
        for t in m {
            match t {
                Token::E => {
                    // Spec says that we should generate new ephemeral key, but that would make
                    // testing very difficult.
                    let e_pk = D::pubkey(self.e.as_ref().unwrap());
                    self.symmetric.mix_hash(e_pk.as_slice());
                    if self.symmetric.has_preshared_key() {
                        self.symmetric.mix_key(e_pk.as_slice());
                    }
                    out.extend_from_slice(e_pk.as_slice());
                }
                Token::S => {
                    let encrypted_s = self.symmetric
                        .encrypt_and_hash_vec(D::pubkey(self.s.as_ref().unwrap()).as_slice());
                    out.extend_from_slice(&encrypted_s);
                }
                t => self.perform_dh(t),
            }
        }

        let encrypted_payload = self.symmetric.encrypt_and_hash_vec(payload);
        out.extend_from_slice(&encrypted_payload);

        out
    }

    /// Update handshake state and get payload, given a packet.
    ///
    /// If the packet fails to decrypt, the whole HandshakeState may be in invalid state, and
    /// should not be used any more. Expect to `get_re` before falling back to `XXfallback`.
    pub fn read_message(&mut self, data: &[u8]) -> Result<Vec<u8>, NoiseError> {
        // Check that it is our turn to recv.
        assert!(self.message_index % 2 == if self.is_initiator { 1 } else { 0 });

        // Get the message pattern.
        let m = self.pattern.get_message_patterns()[self.message_index].clone();
        self.message_index += 1;

        let mut data = data;
        // Consume the next `n` bytes of data.
        let mut get = |n| if data.len() >= n {
            let ret = &data[..n];
            data = &data[n..];
            Ok(ret)
        } else {
            Err(NoiseError::TooShort)
        };

        // Process tokens.
        for t in m {
            match t {
                Token::E => {
                    let mut re = D::Pubkey::new();
                    re.as_mut().copy_from_slice(get(D::pub_len())?);
                    self.symmetric.mix_hash(re.as_slice());
                    if self.symmetric.has_preshared_key() {
                        self.symmetric.mix_key(re.as_slice());
                    }
                    self.re = Some(re);
                }
                Token::S => {
                    let temp = get(if self.symmetric.has_key() {
                        D::pub_len() + 16
                    } else {
                        D::pub_len()
                    })?;
                    let mut rs = D::Pubkey::new();
                    self.symmetric.decrypt_and_hash(temp, rs.as_mut())?;
                    self.rs = Some(rs);
                }
                t => self.perform_dh(t),
            }
        }

        Ok(self.symmetric.decrypt_and_hash_vec(data)?)
    }

    /// Whether handshake has completed.
    pub fn completed(&self) -> bool {
        self.message_index == self.pattern.get_message_patterns().len()
    }

    /// Get handshake hash. Useful for e.g., channel binding.
    ///
    /// Should be called after handshake is `completed()`.
    pub fn get_hash(&self) -> &[u8] {
        self.symmetric.get_hash()
    }

    /// Get ciphers that can be used to encrypt/decrypt furthur messages.
    /// The first `CiperState` is for initiator to responder, and the second for responder
    /// to initiator.
    ///
    /// Should be called after handshake is `completed()`.
    pub fn get_ciphers(&self) -> (CipherState<C>, CipherState<C>) {
        self.symmetric.split()
    }

    /// Get remote static pubkey, if available.
    pub fn get_rs(&self) -> Option<D::Pubkey> {
        self.rs
    }

    /// Get remote semi-ephemeral pubkey.
    ///
    /// Returns `None` if we do not know.
    ///
    /// Useful for noise-pipes.
    pub fn get_re(&self) -> Option<D::Pubkey> {
        self.re
    }

    fn perform_dh(&mut self, t: Token) {
        let dh = |a: Option<&D::Key>, b: Option<&D::Pubkey>| D::dh(a.unwrap(), b.unwrap());

        let k = match t {
            Token::EE => dh(self.e.as_ref(), self.re.as_ref()),
            Token::ES => {
                if self.is_initiator {
                    dh(self.e.as_ref(), self.rs.as_ref())
                } else {
                    dh(self.s.as_ref(), self.re.as_ref())
                }
            }
            Token::SE => {
                if self.is_initiator {
                    dh(self.s.as_ref(), self.re.as_ref())
                } else {
                    dh(self.e.as_ref(), self.rs.as_ref())
                }
            }
            Token::SS => dh(self.s.as_ref(), self.rs.as_ref()),
            _ => unreachable!(),
        };

        self.symmetric.mix_key(k.as_slice());
    }
}

/// Builder for `HandshakeState`.
pub struct HandshakeStateBuilder<'a, D: DH> {
    pattern: Option<HandshakePattern>,
    is_initiator: Option<bool>,
    prologue: Option<&'a [u8]>,
    psk: Option<&'a [u8]>,
    s: Option<D::Key>,
    e: Option<D::Key>,
    rs: Option<D::Pubkey>,
    re: Option<D::Pubkey>,
}

impl<'a, D> HandshakeStateBuilder<'a, D>
    where D: DH
{
    pub fn new() -> Self {
        HandshakeStateBuilder {
            pattern: None,
            is_initiator: None,
            prologue: None,
            psk: None,
            s: None,
            e: None,
            rs: None,
            re: None,
        }
    }

    pub fn set_pattern(&mut self, p: HandshakePattern) -> &mut Self {
        self.pattern = Some(p);
        self
    }

    pub fn set_is_initiator(&mut self, is: bool) -> &mut Self {
        self.is_initiator = Some(is);
        self
    }

    pub fn set_prologue(&mut self, prologue: &'a [u8]) -> &mut Self {
        self.prologue = Some(prologue);
        self
    }

    pub fn set_psk(&mut self, psk: &'a [u8]) -> &mut Self {
        self.psk = Some(psk);
        self
    }

    pub fn set_e(&mut self, e: &D::Key) -> &mut Self {
        self.e = Some(*e);
        self
    }

    pub fn set_s(&mut self, s: &D::Key) -> &mut Self {
        self.s = Some(*s);
        self
    }

    pub fn set_re(&mut self, re: &D::Pubkey) -> &mut Self {
        self.re = Some(*re);
        self
    }

    pub fn set_rs(&mut self, rs: &D::Pubkey) -> &mut Self {
        self.rs = Some(*rs);
        self
    }

    /// Build `HandshakeState`.
    ///
    /// # Panics
    ///
    /// `pattern`, `prologue` and `is_initiator` must be set.
    pub fn build_handshake_state<C, H>(&self) -> HandshakeState<D, C, H>
        where C: Cipher,
              H: Hash
    {
        HandshakeState::new(self.pattern.as_ref().unwrap().clone(),
                            self.is_initiator.unwrap(),
                            self.prologue.unwrap(),
                            self.psk,
                            self.s,
                            self.e,
                            self.rs,
                            self.re)
    }
}
<|endoftext|>