<commit_before>//! Completion tests for pattern position.
use expect_test::{expect, Expect};

use crate::tests::{completion_list, BASE_ITEMS_FIXTURE};

fn check_empty(ra_fixture: &str, expect: Expect) {
    let actual = completion_list(ra_fixture);
    expect.assert_eq(&actual)
}

fn check(ra_fixture: &str, expect: Expect) {
    let actual = completion_list(&format!("{}\n{}", BASE_ITEMS_FIXTURE, ra_fixture));
    expect.assert_eq(&actual)
}

#[test]
fn ident_rebind_pat() {
    check_empty(
        r#"
fn quux() {
    let en$0 @ x
}
"#,
        expect![[r#"
            kw ref
            kw mut
        "#]],
    );
}

#[test]
fn ident_ref_pat() {
    check_empty(
        r#"
fn quux() {
    let ref en$0
}
"#,
        expect![[r#"
            kw mut
        "#]],
    );
    check_empty(
        r#"
fn quux() {
    let ref en$0 @ x
}
"#,
        expect![[r#"
            kw mut
        "#]],
    );
}

#[test]
fn ident_ref_mut_pat() {
    check_empty(
        r#"
fn quux() {
    let ref mut en$0
}
"#,
        expect![[r#""#]],
    );
    check_empty(
        r#"
fn quux() {
    let ref mut en$0 @ x
}
"#,
        expect![[r#""#]],
    );
}

#[test]
fn ref_pat() {
    check_empty(
        r#"
fn quux() {
    let &en$0
}
"#,
        expect![[r#"
            kw mut
        "#]],
    );
    check_empty(
        r#"
fn quux() {
    let &mut en$0
}
"#,
        expect![[r#""#]],
    );
    check_empty(
        r#"
fn foo() {
    for &$0 in () {}
}
"#,
        expect![[r#"
            kw mut
        "#]],
    );
}

#[test]
fn refutable() {
    check(
        r#"
fn foo() {
    if let a$0
}
"#,
        expect![[r#"
            kw ref
            kw mut
            en Enum
            bn Record    Record { field$1 }$0
            st Record
            bn Tuple     Tuple($1)$0
            st Tuple
            md module
            st Unit
            ma makro!(…) macro_rules! makro
            bn TupleV    TupleV($1)$0
            ev TupleV
            ct CONST
        "#]],
    );
}

#[test]
fn irrefutable() {
    check(
        r#"
enum SingleVariantEnum {
    Variant
}
use SingleVariantEnum::Variant;
fn foo() {
   let a$0
}
"#,
        expect![[r#"
            kw ref
            kw mut
            bn Record            Record { field$1 }$0
            st Record
            bn Tuple             Tuple($1)$0
            st Tuple
            ev Variant
            md module
            en SingleVariantEnum
            st Unit
            ma makro!(…)         macro_rules! makro
        "#]],
    );
}

#[test]
fn in_param() {
    check(
        r#"
fn foo(a$0) {
}
"#,
        expect![[r#"
            kw ref
            kw mut
            bn Record    Record { field$1 }: Record$0
            st Record
            bn Tuple     Tuple($1): Tuple$0
            st Tuple
            md module
            st Unit
            ma makro!(…) macro_rules! makro
        "#]],
    );
    check(
        r#"
fn foo(a$0: Tuple) {
}
"#,
        expect![[r#"
            kw ref
            kw mut
            bn Record    Record { field$1 }$0
            st Record
            bn Tuple     Tuple($1)$0
            st Tuple
            md module
            st Unit
            ma makro!(…) macro_rules! makro
        "#]],
    );
}

#[test]
fn only_fn_like_macros() {
    check_empty(
        r#"
macro_rules! m { ($e:expr) => { $e } }

#[rustc_builtin_macro]
macro Clone {}

fn foo() {
    let x$0
}
"#,
        expect![[r#"
            kw ref
            kw mut
            ma m!(…) macro_rules! m
        "#]],
    );
}

#[test]
fn in_simple_macro_call() {
    check_empty(
        r#"
macro_rules! m { ($e:expr) => { $e } }
enum E { X }

fn foo() {
   m!(match E::X { a$0 })
}
"#,
        expect![[r#"
            kw ref
            kw mut
            en E
            ma m!(…) macro_rules! m
        "#]],
    );
}

#[test]
fn omits_private_fields_pat() {
    check_empty(
        r#"
mod foo {
    pub struct Record { pub field: i32, _field: i32 }
    pub struct Tuple(pub u32, u32);
    pub struct Invisible(u32, u32);
}
use foo::*;

fn outer() {
    if let a$0
}
"#,
        expect![[r#"
            kw ref
            kw mut
            bn Record    Record { field$1, .. }$0
            st Record
            bn Tuple     Tuple($1, ..)$0
            st Tuple
            st Invisible
            md foo
        "#]],
    )
}

#[test]
fn completes_self_pats() {
    check_empty(
        r#"
struct Foo(i32);
impl Foo {
    fn foo() {
        match Foo(0) {
            a$0
        }
    }
}
    "#,
        expect![[r#"
            kw ref
            kw mut
            bn Self Self($1)$0
            sp Self
            bn Foo  Foo($1)$0
            st Foo
        "#]],
    )
}

#[test]
fn enum_qualified() {
    check(
        r#"
impl Enum {
    type AssocType = ();
    const ASSOC_CONST: () = ();
    fn assoc_fn() {}
}
fn func() {
    if let Enum::$0 = unknown {}
}
"#,
        expect![[r#"
            ev TupleV(…)   TupleV(u32)
            ev RecordV {…} RecordV { field: u32 }
            ev UnitV       UnitV
        "#]],
    );
}

#[test]
fn completes_in_record_field_pat() {
    check_empty(
        r#"
struct Foo { bar: Bar }
struct Bar(u32);
fn outer(Foo { bar: $0 }: Foo) {}
"#,
        expect![[r#"
            kw ref
            kw mut
            bn Foo Foo { bar$1 }$0
            st Foo
            bn Bar Bar($1)$0
            st Bar
        "#]],
    )
}

#[test]
fn skips_in_record_field_pat_name() {
    check_empty(
        r#"
struct Foo { bar: Bar }
struct Bar(u32);
fn outer(Foo { bar$0 }: Foo) {}
"#,
        expect![[r#""#]],
    )
}

#[test]
fn completes_in_fn_param() {
    check_empty(
        r#"
struct Foo { bar: Bar }
struct Bar(u32);
fn foo($0) {}
"#,
        expect![[r#"
            kw ref
            kw mut
            bn Foo Foo { bar$1 }: Foo$0
            st Foo
            bn Bar Bar($1): Bar$0
            st Bar
        "#]],
    )
}

#[test]
fn completes_in_closure_param() {
    check_empty(
        r#"
struct Foo { bar: Bar }
struct Bar(u32);
fn foo() {
    |$0| {};
}
"#,
        expect![[r#"
            kw ref
            kw mut
            bn Foo Foo { bar$1 }$0
            st Foo
            bn Bar Bar($1)$0
            st Bar
        "#]],
    )
}

#[test]
fn completes_no_delims_if_existing() {
    check_empty(
        r#"
struct Bar(u32);
fn foo() {
    match Bar(0) {
        B$0(b) => {}
    }
}
"#,
        expect![[r#"
            kw self::
            kw super::
            kw crate::
        "#]],
    );
    check_empty(
        r#"
struct Foo { bar: u32 }
fn foo() {
    match Foo { bar: 0 } {
        F$0 { bar } => {}
    }
}
"#,
        expect![[r#"
            kw return
            kw self
            kw super
            kw crate
            st Foo
            fn foo()  fn()
            bt u32
        "#]],
    );
    check_empty(
        r#"
enum Enum {
    TupleVariant(u32)
}
fn foo() {
    match Enum::TupleVariant(0) {
        Enum::T$0(b) => {}
    }
}
"#,
        expect![[r#""#]],
    );
    check_empty(
        r#"
enum Enum {
    RecordVariant { field: u32 }
}
fn foo() {
    match (Enum::RecordVariant { field: 0 }) {
        Enum::RecordV$0 { field } => {}
    }
}
"#,
        expect![[r#""#]],
    );
}
<commit_msg>Add completes_associated_const test<commit_after>//! Completion tests for pattern position.
use expect_test::{expect, Expect};

use crate::tests::{completion_list, BASE_ITEMS_FIXTURE};

fn check_empty(ra_fixture: &str, expect: Expect) {
    let actual = completion_list(ra_fixture);
    expect.assert_eq(&actual)
}

fn check(ra_fixture: &str, expect: Expect) {
    let actual = completion_list(&format!("{}\n{}", BASE_ITEMS_FIXTURE, ra_fixture));
    expect.assert_eq(&actual)
}

#[test]
fn ident_rebind_pat() {
    check_empty(
        r#"
fn quux() {
    let en$0 @ x
}
"#,
        expect![[r#"
            kw ref
            kw mut
        "#]],
    );
}

#[test]
fn ident_ref_pat() {
    check_empty(
        r#"
fn quux() {
    let ref en$0
}
"#,
        expect![[r#"
            kw mut
        "#]],
    );
    check_empty(
        r#"
fn quux() {
    let ref en$0 @ x
}
"#,
        expect![[r#"
            kw mut
        "#]],
    );
}

#[test]
fn ident_ref_mut_pat() {
    check_empty(
        r#"
fn quux() {
    let ref mut en$0
}
"#,
        expect![[r#""#]],
    );
    check_empty(
        r#"
fn quux() {
    let ref mut en$0 @ x
}
"#,
        expect![[r#""#]],
    );
}

#[test]
fn ref_pat() {
    check_empty(
        r#"
fn quux() {
    let &en$0
}
"#,
        expect![[r#"
            kw mut
        "#]],
    );
    check_empty(
        r#"
fn quux() {
    let &mut en$0
}
"#,
        expect![[r#""#]],
    );
    check_empty(
        r#"
fn foo() {
    for &$0 in () {}
}
"#,
        expect![[r#"
            kw mut
        "#]],
    );
}

#[test]
fn refutable() {
    check(
        r#"
fn foo() {
    if let a$0
}
"#,
        expect![[r#"
            kw ref
            kw mut
            en Enum
            bn Record    Record { field$1 }$0
            st Record
            bn Tuple     Tuple($1)$0
            st Tuple
            md module
            st Unit
            ma makro!(…) macro_rules! makro
            bn TupleV    TupleV($1)$0
            ev TupleV
            ct CONST
        "#]],
    );
}

#[test]
fn irrefutable() {
    check(
        r#"
enum SingleVariantEnum {
    Variant
}
use SingleVariantEnum::Variant;
fn foo() {
   let a$0
}
"#,
        expect![[r#"
            kw ref
            kw mut
            bn Record            Record { field$1 }$0
            st Record
            bn Tuple             Tuple($1)$0
            st Tuple
            ev Variant
            md module
            en SingleVariantEnum
            st Unit
            ma makro!(…)         macro_rules! makro
        "#]],
    );
}

#[test]
fn in_param() {
    check(
        r#"
fn foo(a$0) {
}
"#,
        expect![[r#"
            kw ref
            kw mut
            bn Record    Record { field$1 }: Record$0
            st Record
            bn Tuple     Tuple($1): Tuple$0
            st Tuple
            md module
            st Unit
            ma makro!(…) macro_rules! makro
        "#]],
    );
    check(
        r#"
fn foo(a$0: Tuple) {
}
"#,
        expect![[r#"
            kw ref
            kw mut
            bn Record    Record { field$1 }$0
            st Record
            bn Tuple     Tuple($1)$0
            st Tuple
            md module
            st Unit
            ma makro!(…) macro_rules! makro
        "#]],
    );
}

#[test]
fn only_fn_like_macros() {
    check_empty(
        r#"
macro_rules! m { ($e:expr) => { $e } }

#[rustc_builtin_macro]
macro Clone {}

fn foo() {
    let x$0
}
"#,
        expect![[r#"
            kw ref
            kw mut
            ma m!(…) macro_rules! m
        "#]],
    );
}

#[test]
fn in_simple_macro_call() {
    check_empty(
        r#"
macro_rules! m { ($e:expr) => { $e } }
enum E { X }

fn foo() {
   m!(match E::X { a$0 })
}
"#,
        expect![[r#"
            kw ref
            kw mut
            en E
            ma m!(…) macro_rules! m
        "#]],
    );
}

#[test]
fn omits_private_fields_pat() {
    check_empty(
        r#"
mod foo {
    pub struct Record { pub field: i32, _field: i32 }
    pub struct Tuple(pub u32, u32);
    pub struct Invisible(u32, u32);
}
use foo::*;

fn outer() {
    if let a$0
}
"#,
        expect![[r#"
            kw ref
            kw mut
            bn Record    Record { field$1, .. }$0
            st Record
            bn Tuple     Tuple($1, ..)$0
            st Tuple
            st Invisible
            md foo
        "#]],
    )
}

#[test]
fn completes_self_pats() {
    check_empty(
        r#"
struct Foo(i32);
impl Foo {
    fn foo() {
        match Foo(0) {
            a$0
        }
    }
}
    "#,
        expect![[r#"
            kw ref
            kw mut
            bn Self Self($1)$0
            sp Self
            bn Foo  Foo($1)$0
            st Foo
        "#]],
    )
}

#[test]
fn enum_qualified() {
    check(
        r#"
impl Enum {
    type AssocType = ();
    const ASSOC_CONST: () = ();
    fn assoc_fn() {}
}
fn func() {
    if let Enum::$0 = unknown {}
}
"#,
        expect![[r#"
            ev TupleV(…)   TupleV(u32)
            ev RecordV {…} RecordV { field: u32 }
            ev UnitV       UnitV
        "#]],
    );
}

#[test]
fn completes_in_record_field_pat() {
    check_empty(
        r#"
struct Foo { bar: Bar }
struct Bar(u32);
fn outer(Foo { bar: $0 }: Foo) {}
"#,
        expect![[r#"
            kw ref
            kw mut
            bn Foo Foo { bar$1 }$0
            st Foo
            bn Bar Bar($1)$0
            st Bar
        "#]],
    )
}

#[test]
fn skips_in_record_field_pat_name() {
    check_empty(
        r#"
struct Foo { bar: Bar }
struct Bar(u32);
fn outer(Foo { bar$0 }: Foo) {}
"#,
        expect![[r#""#]],
    )
}

#[test]
fn completes_in_fn_param() {
    check_empty(
        r#"
struct Foo { bar: Bar }
struct Bar(u32);
fn foo($0) {}
"#,
        expect![[r#"
            kw ref
            kw mut
            bn Foo Foo { bar$1 }: Foo$0
            st Foo
            bn Bar Bar($1): Bar$0
            st Bar
        "#]],
    )
}

#[test]
fn completes_in_closure_param() {
    check_empty(
        r#"
struct Foo { bar: Bar }
struct Bar(u32);
fn foo() {
    |$0| {};
}
"#,
        expect![[r#"
            kw ref
            kw mut
            bn Foo Foo { bar$1 }$0
            st Foo
            bn Bar Bar($1)$0
            st Bar
        "#]],
    )
}

#[test]
fn completes_no_delims_if_existing() {
    check_empty(
        r#"
struct Bar(u32);
fn foo() {
    match Bar(0) {
        B$0(b) => {}
    }
}
"#,
        expect![[r#"
            kw self::
            kw super::
            kw crate::
        "#]],
    );
    check_empty(
        r#"
struct Foo { bar: u32 }
fn foo() {
    match Foo { bar: 0 } {
        F$0 { bar } => {}
    }
}
"#,
        expect![[r#"
            kw return
            kw self
            kw super
            kw crate
            st Foo
            fn foo()  fn()
            bt u32
        "#]],
    );
    check_empty(
        r#"
enum Enum {
    TupleVariant(u32)
}
fn foo() {
    match Enum::TupleVariant(0) {
        Enum::T$0(b) => {}
    }
}
"#,
        expect![[r#""#]],
    );
    check_empty(
        r#"
enum Enum {
    RecordVariant { field: u32 }
}
fn foo() {
    match (Enum::RecordVariant { field: 0 }) {
        Enum::RecordV$0 { field } => {}
    }
}
"#,
        expect![[r#""#]],
    );
}

#[test]
fn completes_associated_const() {
    check_empty(
        r#"
#[derive(PartialEq, Eq)]
struct Ty(u8);

impl Ty {
    const ABC: Self = Self(0);
}

fn f(t: Ty) {
    match t {
        Ty::$0 => {}
        _ => {}
    }
}
"#,
        expect![[""]],
    )
}
<|endoftext|>


<commit_before>use super::namespaces;
use crate::{library::Library, nameutil::crate_name, version::Version};
use std::collections::btree_map::{BTreeMap, Iter};

/// Provides assistance in generating use declarations.
///
/// It takes into account that use declaration referring to names within the
/// same crate will look differently. It also avoids generating spurious
/// declarations referring to names from within the same module as the one we
/// are generating code for.
#[derive(Clone, Debug, Default)]
pub struct Imports {
    /// Name of the current crate.
    crate_name: String,
    /// Name defined within current module. It doesn't need use declaration.
    ///
    /// NOTE: Currently we don't need to support more than one such name.
    defined: Option<String>,
    defaults: ImportConditions,
    map: BTreeMap<String, ImportConditions>,
}

impl Imports {
    pub fn new(gir: &Library) -> Imports {
        Imports {
            crate_name: make_crate_name(gir),
            defined: None,
            defaults: ImportConditions::default(),
            map: BTreeMap::new(),
        }
    }

    pub fn with_defined(gir: &Library, name: &str) -> Imports {
        Imports {
            crate_name: make_crate_name(gir),
            defined: Some(name.to_owned()),
            defaults: ImportConditions::default(),
            map: BTreeMap::new(),
        }
    }

    pub fn set_defaults(&mut self, version: Option<Version>, constraint: &Option<String>) {
        let constraints = if let Some(constraint) = constraint {
            vec![constraint.clone()]
        } else {
            vec![]
        };
        self.defaults = ImportConditions {
            version,
            constraints,
        };
    }

    pub fn reset_defaults(&mut self) {
        self.defaults.clear();
    }

    /// Declares that name should be available through its last path component.
    ///
    /// For example, if name is `X::Y::Z` then it will be available as `Z`.
    /// Uses defaults.
    pub fn add(&mut self, name: &str) {
        if let Some(ref defined) = self.defined {
            if name == defined {
                return;
            }
        }
        if let Some(name) = self.strip_crate_name(name) {
            let defaults = &self.defaults;
            let entry = self
                .map
                .entry(name.to_owned())
                .or_insert_with(|| defaults.clone());
            entry.update_version(self.defaults.version);
            entry.update_constraints(self.defaults.constraints.clone());
        }
    }

    /// Declares that name should be available through its last path component.
    ///
    /// For example, if name is `X::Y::Z` then it will be available as `Z`.
    pub fn add_with_version(&mut self, name: &str, version: Option<Version>) {
        if let Some(ref defined) = self.defined {
            if name == defined {
                return;
            }
        }
        if let Some(name) = self.strip_crate_name(name) {
            let entry = self.map.entry(name.to_owned()).or_insert(ImportConditions {
                version,
                constraints: Vec::new(),
            });
            entry.update_version(version);
            // Since there is no constraint on this import, if any constraint
            // is present, we can just remove it.
            entry.constraints.clear();
        }
    }

    /// Declares that name should be available through its last path component and provides
    /// an optional feature constraint.
    ///
    /// For example, if name is `X::Y::Z` then it will be available as `Z`.
    pub fn add_with_constraint(
        &mut self,
        name: &str,
        version: Option<Version>,
        constraint: Option<&str>,
    ) {
        if let Some(ref defined) = self.defined {
            if name == defined {
                return;
            }
        }
        if let Some(name) = self.strip_crate_name(name) {
            let entry = if let Some(constraint) = constraint {
                let constraint = String::from(constraint);
                let entry = self.map.entry(name.to_owned()).or_insert(ImportConditions {
                    version,
                    constraints: vec![constraint.clone()],
                });
                entry.add_constraint(constraint);
                entry
            } else {
                let entry = self.map.entry(name.to_owned()).or_insert(ImportConditions {
                    version,
                    constraints: Vec::new(),
                });
                // Since there is no constraint on this import, if any constraint
                // is present, we can just remove it.
                entry.constraints.clear();
                entry
            };
            entry.update_version(version);
        }
    }

    /// Declares that name should be available through its full path.
    ///
    /// For example, if name is `X::Y` then it will be available as `X::Y`.
    pub fn add_used_type(&mut self, used_type: &str) {
        if let Some(i) = used_type.find("::") {
            if i == 0 {
                self.add(&used_type[2..]);
            } else {
                self.add(&used_type[..i]);
            }
        } else {
            self.add(used_type);
        }
    }

    pub fn add_used_types(&mut self, used_types: &[String]) {
        for s in used_types {
            self.add_used_type(s);
        }
    }

    /// Declares that name should be available through its full path.
    ///
    /// For example, if name is `X::Y` then it will be available as `X::Y`.
    pub fn add_used_type_with_version(&mut self, used_type: &str, version: Option<Version>) {
        if let Some(i) = used_type.find("::") {
            if i == 0 {
                self.add_with_version(&used_type[2..], version);
            } else {
                self.add_with_version(&used_type[..i], version);
            }
        } else {
            self.add_with_version(used_type, version);
        }
    }

    /// Tries to strip crate name prefix from given name.
    ///
    /// Returns `None` if name matches crate name exactly. Otherwise returns
    /// name with crate name prefix stripped or full name if there was no match.
    fn strip_crate_name<'a>(&self, name: &'a str) -> Option<&'a str> {
        let prefix = &self.crate_name;
        if !name.starts_with(prefix) {
            return Some(name);
        }
        let rest = &name[prefix.len()..];
        if rest.is_empty() {
            None
        } else if rest.starts_with("::") {
            Some(&rest["::".len()..])
        } else {
            // It was false positive, return the whole name.
            Some(name)
        }
    }

    pub fn iter(&self) -> Iter<'_, String, ImportConditions> {
        self.map.iter()
    }
}

#[derive(Clone, Debug, Default)]
pub struct ImportConditions {
    pub version: Option<Version>,
    pub constraints: Vec<String>,
}

impl ImportConditions {
    fn clear(&mut self) {
        *self = ImportConditions::default();
    }

    fn update_version(&mut self, version: Option<Version>) {
        if version < self.version {
            self.version = version;
        }
    }

    fn add_constraint(&mut self, constraint: String) {
        // If the import is already present but doesn't have any constraint,
        // we don't want to add one.
        if self.constraints.is_empty() {
            return;
        }
        // Otherwise, we just check if the constraint
        // is already present or not before adding it.
        if !self.constraints.iter().any(|x| x == &constraint) {
            self.constraints.push(constraint);
        }
    }

    fn update_constraints(&mut self, constraints: Vec<String>) {
        // If the import is already present but doesn't have any constraint,
        // we don't want to add one.
        if self.constraints.is_empty() {
            return;
        }
        if constraints.is_empty() {
            // Since there is no constraint on this import, if any constraint
            // is present, we can just remove it.
            self.constraints.clear();
        } else {
            // Otherwise, we just check if the constraint
            // is already present or not before adding it.
            for constraint in constraints {
                if !self.constraints.iter().any(|x| x == &constraint) {
                    self.constraints.push(constraint.clone());
                }
            }
        }
    }
}

fn make_crate_name(gir: &Library) -> String {
    let name = gir.namespace(namespaces::MAIN).name.as_str();
    if name == "GObject" {
        crate_name("GLib")
    } else {
        crate_name(name)
    }
}
<commit_msg>generate imports in the same order as rustfmt<commit_after>use super::namespaces;
use crate::{library::Library, nameutil::crate_name, version::Version};
use std::cmp::Ordering;
use std::collections::btree_map::BTreeMap;
use std::vec::IntoIter;

fn is_first_char_up(s: &str) -> bool {
    s.chars().next().unwrap().is_uppercase()
}

fn check_up_eq(a: &str, b: &str) -> Ordering {
    let is_a_up = is_first_char_up(a);
    let is_b_up = is_first_char_up(b);
    if is_a_up != is_b_up {
        if is_a_up {
            return Ordering::Greater;
        }
        return Ordering::Less;
    }
    Ordering::Equal
}

/// This function is used by the `Imports` type to generate output like `cargo fmt` would.
///
/// For example:
///
/// ```text
/// use gdk; // lowercases come first.
/// use Window;
///
/// use gdk::foo; // lowercases come first here as well.
/// use gdk::Foo;
/// ```
fn compare_imports(a: &(&String, &ImportConditions), b: &(&String, &ImportConditions)) -> Ordering {
    let s = check_up_eq(a.0, b.0);
    if s != Ordering::Equal {
        return s;
    }
    let mut a = a.0.split("::");
    let mut b = b.0.split("::");
    loop {
        match (a.next(), b.next()) {
            (Some(a), Some(b)) => {
                let s = check_up_eq(a, b);
                if s != Ordering::Equal {
                    break s;
                }
                let s = a.partial_cmp(b).unwrap();
                if s != Ordering::Equal {
                    break s;
                }
            }
            (Some(_), None) => break Ordering::Greater,
            (None, Some(_)) => break Ordering::Less,
            (None, None) => break Ordering::Equal,
        }
    }
}

/// Provides assistance in generating use declarations.
///
/// It takes into account that use declaration referring to names within the
/// same crate will look differently. It also avoids generating spurious
/// declarations referring to names from within the same module as the one we
/// are generating code for.
#[derive(Clone, Debug, Default)]
pub struct Imports {
    /// Name of the current crate.
    crate_name: String,
    /// Name defined within current module. It doesn't need use declaration.
    ///
    /// NOTE: Currently we don't need to support more than one such name.
    defined: Option<String>,
    defaults: ImportConditions,
    map: BTreeMap<String, ImportConditions>,
}

impl Imports {
    pub fn new(gir: &Library) -> Imports {
        Imports {
            crate_name: make_crate_name(gir),
            defined: None,
            defaults: ImportConditions::default(),
            map: BTreeMap::new(),
        }
    }

    pub fn with_defined(gir: &Library, name: &str) -> Imports {
        Imports {
            crate_name: make_crate_name(gir),
            defined: Some(name.to_owned()),
            defaults: ImportConditions::default(),
            map: BTreeMap::new(),
        }
    }

    pub fn set_defaults(&mut self, version: Option<Version>, constraint: &Option<String>) {
        let constraints = if let Some(constraint) = constraint {
            vec![constraint.clone()]
        } else {
            vec![]
        };
        self.defaults = ImportConditions {
            version,
            constraints,
        };
    }

    pub fn reset_defaults(&mut self) {
        self.defaults.clear();
    }

    /// Declares that name should be available through its last path component.
    ///
    /// For example, if name is `X::Y::Z` then it will be available as `Z`.
    /// Uses defaults.
    pub fn add(&mut self, name: &str) {
        if let Some(ref defined) = self.defined {
            if name == defined {
                return;
            }
        }
        if let Some(name) = self.strip_crate_name(name) {
            let defaults = &self.defaults;
            let entry = self
                .map
                .entry(name.to_owned())
                .or_insert_with(|| defaults.clone());
            entry.update_version(self.defaults.version);
            entry.update_constraints(self.defaults.constraints.clone());
        }
    }

    /// Declares that name should be available through its last path component.
    ///
    /// For example, if name is `X::Y::Z` then it will be available as `Z`.
    pub fn add_with_version(&mut self, name: &str, version: Option<Version>) {
        if let Some(ref defined) = self.defined {
            if name == defined {
                return;
            }
        }
        if let Some(name) = self.strip_crate_name(name) {
            let entry = self.map.entry(name.to_owned()).or_insert(ImportConditions {
                version,
                constraints: Vec::new(),
            });
            entry.update_version(version);
            // Since there is no constraint on this import, if any constraint
            // is present, we can just remove it.
            entry.constraints.clear();
        }
    }

    /// Declares that name should be available through its last path component and provides
    /// an optional feature constraint.
    ///
    /// For example, if name is `X::Y::Z` then it will be available as `Z`.
    pub fn add_with_constraint(
        &mut self,
        name: &str,
        version: Option<Version>,
        constraint: Option<&str>,
    ) {
        if let Some(ref defined) = self.defined {
            if name == defined {
                return;
            }
        }
        if let Some(name) = self.strip_crate_name(name) {
            let entry = if let Some(constraint) = constraint {
                let constraint = String::from(constraint);
                let entry = self.map.entry(name.to_owned()).or_insert(ImportConditions {
                    version,
                    constraints: vec![constraint.clone()],
                });
                entry.add_constraint(constraint);
                entry
            } else {
                let entry = self.map.entry(name.to_owned()).or_insert(ImportConditions {
                    version,
                    constraints: Vec::new(),
                });
                // Since there is no constraint on this import, if any constraint
                // is present, we can just remove it.
                entry.constraints.clear();
                entry
            };
            entry.update_version(version);
        }
    }

    /// Declares that name should be available through its full path.
    ///
    /// For example, if name is `X::Y` then it will be available as `X::Y`.
    pub fn add_used_type(&mut self, used_type: &str) {
        if let Some(i) = used_type.find("::") {
            if i == 0 {
                self.add(&used_type[2..]);
            } else {
                self.add(&used_type[..i]);
            }
        } else {
            self.add(used_type);
        }
    }

    pub fn add_used_types(&mut self, used_types: &[String]) {
        for s in used_types {
            self.add_used_type(s);
        }
    }

    /// Declares that name should be available through its full path.
    ///
    /// For example, if name is `X::Y` then it will be available as `X::Y`.
    pub fn add_used_type_with_version(&mut self, used_type: &str, version: Option<Version>) {
        if let Some(i) = used_type.find("::") {
            if i == 0 {
                self.add_with_version(&used_type[2..], version);
            } else {
                self.add_with_version(&used_type[..i], version);
            }
        } else {
            self.add_with_version(used_type, version);
        }
    }

    /// Tries to strip crate name prefix from given name.
    ///
    /// Returns `None` if name matches crate name exactly. Otherwise returns
    /// name with crate name prefix stripped or full name if there was no match.
    fn strip_crate_name<'a>(&self, name: &'a str) -> Option<&'a str> {
        let prefix = &self.crate_name;
        if !name.starts_with(prefix) {
            return Some(name);
        }
        let rest = &name[prefix.len()..];
        if rest.is_empty() {
            None
        } else if rest.starts_with("::") {
            Some(&rest["::".len()..])
        } else {
            // It was false positive, return the whole name.
            Some(name)
        }
    }

    pub fn iter(&self) -> IntoIter<(&String, &ImportConditions)> {
        let mut imports = self.map.iter().collect::<Vec<_>>();
        imports.sort_by(compare_imports);
        imports.into_iter()
    }
}

#[derive(Clone, Debug, Default)]
pub struct ImportConditions {
    pub version: Option<Version>,
    pub constraints: Vec<String>,
}

impl ImportConditions {
    fn clear(&mut self) {
        *self = ImportConditions::default();
    }

    fn update_version(&mut self, version: Option<Version>) {
        if version < self.version {
            self.version = version;
        }
    }

    fn add_constraint(&mut self, constraint: String) {
        // If the import is already present but doesn't have any constraint,
        // we don't want to add one.
        if self.constraints.is_empty() {
            return;
        }
        // Otherwise, we just check if the constraint
        // is already present or not before adding it.
        if !self.constraints.iter().any(|x| x == &constraint) {
            self.constraints.push(constraint);
        }
    }

    fn update_constraints(&mut self, constraints: Vec<String>) {
        // If the import is already present but doesn't have any constraint,
        // we don't want to add one.
        if self.constraints.is_empty() {
            return;
        }
        if constraints.is_empty() {
            // Since there is no constraint on this import, if any constraint
            // is present, we can just remove it.
            self.constraints.clear();
        } else {
            // Otherwise, we just check if the constraint
            // is already present or not before adding it.
            for constraint in constraints {
                if !self.constraints.iter().any(|x| x == &constraint) {
                    self.constraints.push(constraint.clone());
                }
            }
        }
    }
}

fn make_crate_name(gir: &Library) -> String {
    let name = gir.namespace(namespaces::MAIN).name.as_str();
    if name == "GObject" {
        crate_name("GLib")
    } else {
        crate_name(name)
    }
}
<|endoftext|>


<commit_before>use num::{FromPrimitive, ToPrimitive};
use std::marker::PhantomData;
use std::collections::HashSet;

use byteorder::{ByteOrder, LittleEndian};
use log::{debug, info};
use goblin::{Object};
use failure::{Error, Fail};

use super::super::super::arch::Arch;
use super::super::super::loader::{Permissions};
use super::super::super::workspace::Workspace;
use super::super::{Analyzer, AnalysisError};

#[derive(Debug, Fail)]
pub enum PtrAnalyzerError {
    #[fail(display = "Invalid relocation type")]
    InvalidRelocType,
}

pub struct PtrAnalyzer<A: Arch> {
    // This Analyzer must have a type parameter for it
    //  to implement Analyzer<A>.
    // however, it doesn't actually use this type itself.
    // so, we use a phantom data marker which has zero type,
    //  to ensure there is not an unused type parameter,
    //  which is a compile error.
    _phantom: PhantomData<A>,
}

impl<A: Arch> PtrAnalyzer<A> {
    pub fn new() -> PtrAnalyzer<A> {
        PtrAnalyzer {
            _phantom: PhantomData {},
        }
    }
}

struct Section<A: Arch> {
    start: A::RVA,
    end: A::RVA,
}

fn is_in_insn<A: Arch + 'static>(ws: &Workspace<A>, rva: A::RVA) -> bool {
    let start: usize = rva.to_usize().unwrap();
    // TODO: remove harded max insn length
    // TODO: underflow
    let end: usize = rva.to_usize().unwrap() - 0x10;

    for i in (start..end).rev() {
        let i = A::RVA::from_usize(i).unwrap();
        if let Some(meta) = ws.get_meta(i) {
            if !meta.is_insn() {
                continue;
            }

            if let Ok(len) = meta.get_insn_length() {
                if i + A::RVA::from_u8(len).unwrap() > rva {
                    return true;
                }
            }
        }
    }
    return false;
}

fn is_ptr<A: Arch + 'static>(ws: &Workspace<A>, rva: A::RVA) -> bool {
    if let Ok(ptr) = ws.read_va(rva) {
        if let Some(ptr) = ws.rva(ptr) {
            return ws.probe(ptr, 1);
        }
    }
    return false;
}

pub enum RelocationType {
    IMAGE_REL_BASED_ABSOLUTE,
    IMAGE_REL_BASED_HIGH,
    IMAGE_REL_BASED_LOW,
    IMAGE_REL_BASED_HIGHLOW,
    IMAGE_REL_BASED_HIGHADJ,

    // IMAGE_REL_BASED_MIPS_JMPADDR,
    // IMAGE_REL_BASED_ARM_MOV32,
    // IMAGE_REL_BASED_RISCV_HIGH20,
    IMAGE_REL_ARCH1,

    IMAGE_REL_RESERVED,

    // IMAGE_REL_BASED_THUMB_MOV32,
    // IMAGE_REL_BASED_RISCV_LOW12I,
    IMAGE_REL_ARCH2,

    IMAGE_REL_BASED_RISCV_LOW12S,
    IMAGE_REL_BASED_MIPS_JMPADDR16,
    IMAGE_REL_BASED_DIR64,
}

pub struct Reloc<A: Arch> {
    pub typ: RelocationType,
    pub offset: A::RVA,
}

fn parse_reloc<A: Arch>(base: A::RVA, entry: u16) -> Result<Reloc<A>, Error> {
    let reloc_type   = entry & 0b1111000000000000 >> 12;
    let reloc_offset = entry & 0b0000111111111111;

    // TODO: this should probably be on `RelocationType` as a `From` trait.
    let reloc_type = match reloc_type {
        0 => RelocationType::IMAGE_REL_BASED_ABSOLUTE,
        1 => RelocationType::IMAGE_REL_BASED_HIGH,
        2 => RelocationType::IMAGE_REL_BASED_LOW,
        3 => RelocationType::IMAGE_REL_BASED_HIGHLOW,
        4 => RelocationType::IMAGE_REL_BASED_HIGHADJ,
        5 => RelocationType::IMAGE_REL_ARCH1,
        6 => RelocationType::IMAGE_REL_RESERVED,
        7 => RelocationType::IMAGE_REL_ARCH2,
        8 => RelocationType::IMAGE_REL_BASED_RISCV_LOW12S,
        9 => RelocationType::IMAGE_REL_BASED_MIPS_JMPADDR16,
        10 => RelocationType::IMAGE_REL_BASED_DIR64,
        _ => return Err(PtrAnalyzerError::InvalidRelocType.into()),
    };

    Ok(Reloc {
        typ: reloc_type,
        offset: base + A::RVA::from_u16(reloc_offset).unwrap(),
    })
}

fn split_u32(e: u32) -> (u16, u16) {
    let m = ((e & 0x0000FFFF) >>  0) as u16;
    let n = ((e & 0xFFFF0000) >> 16) as u16;
    (m, n)
}

/// ```
/// use lancelot::rsrc::*;
/// use lancelot::arch::*;
/// use lancelot::workspace::Workspace;
/// use lancelot::analysis::pe::ptrs;
///
/// let mut ws = Workspace::<Arch64>::from_bytes("k32.dll", &get_buf(Rsrc::K32))
///    .disable_analysis()
///    .load().unwrap();
/// let relocs = ptrs::get_relocs(&ws).unwrap();
/// assert_eq!(relocs[0].offset,   0x76008);
/// assert_eq!(relocs[277].offset, 0xA8000);
/// ```
pub fn get_relocs<A: Arch + 'static>(ws: &Workspace<A>) -> Result<Vec<Reloc<A>>, Error> {
    let pe = match Object::parse(&ws.buf) {
        Ok(Object::PE(pe)) => pe,
        _ => return Ok(vec![]),
    };

    let opt_header = match pe.header.optional_header {
        Some(opt_header) => opt_header,
        _ => return Ok(vec![]),
    };

    let reloc_directory = match opt_header.data_directories.get_base_relocation_table() {
        Some(reloc_directory) => reloc_directory,
        _ => return Ok(vec![]),
    };

    let dir_start = A::RVA::from_u32(reloc_directory.virtual_address).unwrap();
    let buf = ws.read_bytes(dir_start, reloc_directory.size as usize)?;

    let entries: Vec<u32> = buf
        .chunks_exact(0x4)
        .map(|b| LittleEndian::read_u32(b))
        .collect();

    let mut ret = vec![];
    let mut index = 0;
    loop {
        // parse base relocation chunk
        // until there's no data left

        let page_rva = entries.get(index);
        let block_size = entries.get(index+1);
        let (page_rva, block_size) = match (page_rva, block_size) {
            (Some(0),        _               ) => break,
            (_,              Some(0)         ) => break,
            (Some(page_rva), Some(block_size)) => (A::RVA::from_u32(*page_rva).unwrap(), *block_size),
            _ => break,
        };

        // cast from u32 to usize
        let chunk_entries_count = ((block_size) / 4) as usize;

        for i in (2..chunk_entries_count) {
            if let Some(&entry) = entries.get(index + i) {
                let (m, n) = split_u32(entry);

                let reloc1 = parse_reloc::<A>(page_rva, m)?;
                if !ws.probe(reloc1.offset, 4) {
                    break
                }
                ret.push(reloc1);

                let reloc2 = parse_reloc::<A>(page_rva, n)?;
                if !ws.probe(reloc2.offset, 4) {
                    break
                }
                ret.push(reloc2);
            } else {
                break
            }
        };

        index += chunk_entries_count;
    }

    Ok(ret)
}

impl<A: Arch + 'static> Analyzer<A> for PtrAnalyzer<A> {
    fn get_name(&self) -> String {
        "pointer analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::PtrAnalyzer;
    ///
    /// let mut ws = Workspace::<Arch64>::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    /// let anal = PtrAnalyzer::<Arch64>::new();
    /// anal.analyze(&mut ws).unwrap();
    /// let meta = ws.get_meta(0xC7F0).unwrap();
    /// assert!(meta.is_insn());
    /// ```
    fn analyze(&self, ws: &mut Workspace<A>)-> Result<(), Error> {
        let pe = match Object::parse(&ws.buf) {
            Ok(Object::PE(pe)) => pe,
            _ => panic!("can't analyze unexpected format"),
        };

        let text_section = match ws.module.sections.iter()
            .filter(|&sec| sec.name == ".text")
            .next() {
                None => return Ok(()),
                Some(s) => s,
            };

        info!("found text section");

        let text_bounds = Section::<A> {
            start: text_section.addr,
            end: text_section.addr + A::RVA::from_usize(text_section.buf.len()).unwrap(),
        };

        let start: usize = text_bounds.start.to_usize().unwrap();
        let end: usize = text_bounds.end.to_usize().unwrap();

        // scan the relocations
        // looking for pointers into the .text section
        // to things that
        //   1. are not already in an instruction
        //   2. don't appear to be a pointer
        // and assume this is code.
        let o: Vec<A::RVA> = get_relocs(ws)?
            .iter()
            .map(|reloc| reloc.offset)
            .map(|rva| ws.read_va(rva))
            .filter_map(Result::ok)
            .filter_map(|va| ws.rva(va))
            .filter(|&rva| text_bounds.start <= rva)
            .filter(|&rva| rva < text_bounds.end)
            .filter(|&rva| !is_in_insn(ws, rva))
            .filter(|&rva| !is_ptr(ws, rva))
            // TODO: maybe ensure that the insn decodes.
            .collect();

        o.iter().for_each(|&rva| {
            debug!("found ptr from .text section to .text section at {:#x}", rva);

            // TODO: consume result
            ws.make_insn(rva).unwrap();
            // TODO: consume result
            ws.analyze().unwrap();
        });
        Ok(())
    }
}
<commit_msg>analysis/pe/ptrs: fix parsing of reloc entry<commit_after>use num::{FromPrimitive, ToPrimitive};
use std::marker::PhantomData;
use std::collections::HashSet;

use byteorder::{ByteOrder, LittleEndian};
use log::{debug, info};
use goblin::{Object};
use failure::{Error, Fail};

use super::super::super::arch::Arch;
use super::super::super::loader::{Permissions};
use super::super::super::workspace::Workspace;
use super::super::{Analyzer, AnalysisError};

#[derive(Debug, Fail)]
pub enum PtrAnalyzerError {
    #[fail(display = "Invalid relocation type")]
    InvalidRelocType,
}

pub struct PtrAnalyzer<A: Arch> {
    // This Analyzer must have a type parameter for it
    //  to implement Analyzer<A>.
    // however, it doesn't actually use this type itself.
    // so, we use a phantom data marker which has zero type,
    //  to ensure there is not an unused type parameter,
    //  which is a compile error.
    _phantom: PhantomData<A>,
}

impl<A: Arch> PtrAnalyzer<A> {
    pub fn new() -> PtrAnalyzer<A> {
        PtrAnalyzer {
            _phantom: PhantomData {},
        }
    }
}

struct Section<A: Arch> {
    start: A::RVA,
    end: A::RVA,
}

fn is_in_insn<A: Arch + 'static>(ws: &Workspace<A>, rva: A::RVA) -> bool {
    let start: usize = rva.to_usize().unwrap();
    // TODO: remove harded max insn length
    // TODO: underflow
    let end: usize = rva.to_usize().unwrap() - 0x10;

    for i in (start..end).rev() {
        let i = A::RVA::from_usize(i).unwrap();
        if let Some(meta) = ws.get_meta(i) {
            if !meta.is_insn() {
                continue;
            }

            if let Ok(len) = meta.get_insn_length() {
                if i + A::RVA::from_u8(len).unwrap() > rva {
                    return true;
                }
            }
        }
    }
    return false;
}

fn is_ptr<A: Arch + 'static>(ws: &Workspace<A>, rva: A::RVA) -> bool {
    if let Ok(ptr) = ws.read_va(rva) {
        if let Some(ptr) = ws.rva(ptr) {
            return ws.probe(ptr, 1);
        }
    }
    return false;
}

pub enum RelocationType {
    IMAGE_REL_BASED_ABSOLUTE,
    IMAGE_REL_BASED_HIGH,
    IMAGE_REL_BASED_LOW,
    IMAGE_REL_BASED_HIGHLOW,
    IMAGE_REL_BASED_HIGHADJ,

    // IMAGE_REL_BASED_MIPS_JMPADDR,
    // IMAGE_REL_BASED_ARM_MOV32,
    // IMAGE_REL_BASED_RISCV_HIGH20,
    IMAGE_REL_ARCH1,

    IMAGE_REL_RESERVED,

    // IMAGE_REL_BASED_THUMB_MOV32,
    // IMAGE_REL_BASED_RISCV_LOW12I,
    IMAGE_REL_ARCH2,

    IMAGE_REL_BASED_RISCV_LOW12S,
    IMAGE_REL_BASED_MIPS_JMPADDR16,
    IMAGE_REL_BASED_DIR64,
}

pub struct Reloc<A: Arch> {
    pub typ: RelocationType,
    pub offset: A::RVA,
}

fn parse_reloc<A: Arch>(base: A::RVA, entry: u16) -> Result<Reloc<A>, Error> {
    let reloc_type   = (entry & 0b1111000000000000) >> 12;
    let reloc_offset =  entry & 0b0000111111111111;

    // TODO: this should probably be on `RelocationType` as a `From` trait.
    let reloc_type = match reloc_type {
        0 => RelocationType::IMAGE_REL_BASED_ABSOLUTE,
        1 => RelocationType::IMAGE_REL_BASED_HIGH,
        2 => RelocationType::IMAGE_REL_BASED_LOW,
        3 => RelocationType::IMAGE_REL_BASED_HIGHLOW,
        4 => RelocationType::IMAGE_REL_BASED_HIGHADJ,
        5 => RelocationType::IMAGE_REL_ARCH1,
        6 => RelocationType::IMAGE_REL_RESERVED,
        7 => RelocationType::IMAGE_REL_ARCH2,
        8 => RelocationType::IMAGE_REL_BASED_RISCV_LOW12S,
        9 => RelocationType::IMAGE_REL_BASED_MIPS_JMPADDR16,
        10 => RelocationType::IMAGE_REL_BASED_DIR64,
        _ => return Err(PtrAnalyzerError::InvalidRelocType.into()),
    };

    Ok(Reloc {
        typ: reloc_type,
        offset: base + A::RVA::from_u16(reloc_offset).unwrap(),
    })
}

fn split_u32(e: u32) -> (u16, u16) {
    let m = ((e & 0x0000FFFF) >>  0) as u16;
    let n = ((e & 0xFFFF0000) >> 16) as u16;
    (m, n)
}

/// ```
/// use lancelot::rsrc::*;
/// use lancelot::arch::*;
/// use lancelot::workspace::Workspace;
/// use lancelot::analysis::pe::ptrs;
///
/// let mut ws = Workspace::<Arch64>::from_bytes("k32.dll", &get_buf(Rsrc::K32))
///    .disable_analysis()
///    .load().unwrap();
/// let relocs = ptrs::get_relocs(&ws).unwrap();
/// assert_eq!(relocs[0].offset,   0x76008);
/// assert_eq!(relocs[277].offset, 0xA8000);
/// ```
pub fn get_relocs<A: Arch + 'static>(ws: &Workspace<A>) -> Result<Vec<Reloc<A>>, Error> {
    let pe = match Object::parse(&ws.buf) {
        Ok(Object::PE(pe)) => pe,
        _ => return Ok(vec![]),
    };

    let opt_header = match pe.header.optional_header {
        Some(opt_header) => opt_header,
        _ => return Ok(vec![]),
    };

    let reloc_directory = match opt_header.data_directories.get_base_relocation_table() {
        Some(reloc_directory) => reloc_directory,
        _ => return Ok(vec![]),
    };

    let dir_start = A::RVA::from_u32(reloc_directory.virtual_address).unwrap();
    let buf = ws.read_bytes(dir_start, reloc_directory.size as usize)?;

    let entries: Vec<u32> = buf
        .chunks_exact(0x4)
        .map(|b| LittleEndian::read_u32(b))
        .collect();

    let mut ret = vec![];
    let mut index = 0;
    loop {
        // parse base relocation chunk
        // until there's no data left

        let page_rva = entries.get(index);
        let block_size = entries.get(index+1);
        let (page_rva, block_size) = match (page_rva, block_size) {
            (Some(0),        _               ) => break,
            (_,              Some(0)         ) => break,
            (Some(page_rva), Some(block_size)) => (A::RVA::from_u32(*page_rva).unwrap(), *block_size),
            _ => break,
        };

        // cast from u32 to usize
        let chunk_entries_count = ((block_size) / 4) as usize;

        for i in (2..chunk_entries_count) {
            if let Some(&entry) = entries.get(index + i) {
                let (m, n) = split_u32(entry);

                let reloc1 = parse_reloc::<A>(page_rva, m)?;
                if !ws.probe(reloc1.offset, 4) {
                    break
                }
                ret.push(reloc1);

                let reloc2 = parse_reloc::<A>(page_rva, n)?;
                if !ws.probe(reloc2.offset, 4) {
                    break
                }
                ret.push(reloc2);
            } else {
                break
            }
        };

        index += chunk_entries_count;
    }

    Ok(ret)
}

impl<A: Arch + 'static> Analyzer<A> for PtrAnalyzer<A> {
    fn get_name(&self) -> String {
        "pointer analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::PtrAnalyzer;
    ///
    /// let mut ws = Workspace::<Arch64>::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    /// let anal = PtrAnalyzer::<Arch64>::new();
    /// anal.analyze(&mut ws).unwrap();
    /// let meta = ws.get_meta(0xC7F0).unwrap();
    /// assert!(meta.is_insn());
    /// ```
    fn analyze(&self, ws: &mut Workspace<A>)-> Result<(), Error> {
        let pe = match Object::parse(&ws.buf) {
            Ok(Object::PE(pe)) => pe,
            _ => panic!("can't analyze unexpected format"),
        };

        let text_section = match ws.module.sections.iter()
            // currently limited to the text section.
            // TODO: accept any writable section.
            .filter(|&sec| sec.name == ".text")
            .next() {
                None => return Ok(()),
                Some(s) => s,
            };

        let text_bounds = Section::<A> {
            start: text_section.addr,
            end: text_section.addr + A::RVA::from_usize(text_section.buf.len()).unwrap(),
        };

        let start: usize = text_bounds.start.to_usize().unwrap();
        let end: usize = text_bounds.end.to_usize().unwrap();

        // scan the relocations
        // looking for pointers into the .text section
        // to things that
        //   1. are not already in an instruction
        //   2. don't appear to be a pointer
        // and assume this is code.
        let o: Vec<A::RVA> = get_relocs(ws)?
            .iter()
            .map(|reloc| reloc.offset)
            .map(|rva| ws.read_va(rva))
            .filter_map(Result::ok)
            .filter_map(|va| ws.rva(va))
            .filter(|&rva| text_bounds.start <= rva)
            .filter(|&rva| rva < text_bounds.end)
            .filter(|&rva| !is_in_insn(ws, rva))
            .filter(|&rva| !is_ptr(ws, rva))
            // TODO: maybe ensure that the insn decodes.
            .collect();

        o.iter().for_each(|&rva| {
            debug!("found ptr from .text section to .text section at {:#x}", rva);

            // TODO: consume result
            ws.make_insn(rva).unwrap();
            // TODO: consume result
            ws.analyze().unwrap();
        });
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Look up accessible paths for items.
use either::Either;
use hir::{
    AsAssocItem, AssocItem, Crate, ItemInNs, MacroDef, ModPath, Module, ModuleDef, PathResolution,
    PrefixKind, Semantics,
};
use rustc_hash::FxHashSet;
use syntax::{ast, AstNode};

use crate::{
    imports_locator::{self, AssocItemSearch, DEFAULT_QUERY_SEARCH_LIMIT},
    RootDatabase,
};

#[derive(Debug)]
pub enum ImportCandidate {
    // A path, qualified (`std::collections::HashMap`) or not (`HashMap`).
    Path(PathImportCandidate),
    /// A trait associated function (with no self parameter) or associated constant.
    /// For 'test_mod::TestEnum::test_function', `ty` is the `test_mod::TestEnum` expression type
    /// and `name` is the `test_function`
    TraitAssocItem(TraitImportCandidate),
    /// A trait method with self parameter.
    /// For 'test_enum.test_method()', `ty` is the `test_enum` expression type
    /// and `name` is the `test_method`
    TraitMethod(TraitImportCandidate),
}

#[derive(Debug)]
pub struct TraitImportCandidate {
    pub receiver_ty: hir::Type,
    pub name: NameToImport,
}

#[derive(Debug)]
pub struct PathImportCandidate {
    pub unresolved_qualifier: Option<ast::Path>,
    pub name: NameToImport,
}

#[derive(Debug)]
pub enum NameToImport {
    Exact(String),
    Fuzzy(String),
}

impl NameToImport {
    pub fn text(&self) -> &str {
        match self {
            NameToImport::Exact(text) => text.as_str(),
            NameToImport::Fuzzy(text) => text.as_str(),
        }
    }
}

#[derive(Debug)]
pub struct ImportAssets {
    import_candidate: ImportCandidate,
    module_with_candidate: hir::Module,
}

impl ImportAssets {
    pub fn for_method_call(
        method_call: &ast::MethodCallExpr,
        sema: &Semantics<RootDatabase>,
    ) -> Option<Self> {
        Some(Self {
            import_candidate: ImportCandidate::for_method_call(sema, method_call)?,
            module_with_candidate: sema.scope(method_call.syntax()).module()?,
        })
    }

    pub fn for_exact_path(
        fully_qualified_path: &ast::Path,
        sema: &Semantics<RootDatabase>,
    ) -> Option<Self> {
        let syntax_under_caret = fully_qualified_path.syntax();
        if syntax_under_caret.ancestors().find_map(ast::Use::cast).is_some() {
            return None;
        }
        Some(Self {
            import_candidate: ImportCandidate::for_regular_path(sema, fully_qualified_path)?,
            module_with_candidate: sema.scope(syntax_under_caret).module()?,
        })
    }

    pub fn for_fuzzy_path(
        module_with_candidate: Module,
        qualifier: Option<ast::Path>,
        fuzzy_name: String,
        sema: &Semantics<RootDatabase>,
    ) -> Option<Self> {
        Some(Self {
            import_candidate: ImportCandidate::for_fuzzy_path(qualifier, fuzzy_name, sema)?,
            module_with_candidate,
        })
    }

    pub fn for_fuzzy_method_call(
        module_with_method_call: Module,
        receiver_ty: hir::Type,
        fuzzy_method_name: String,
    ) -> Option<Self> {
        Some(Self {
            import_candidate: ImportCandidate::TraitMethod(TraitImportCandidate {
                receiver_ty,
                name: NameToImport::Fuzzy(fuzzy_method_name),
            }),
            module_with_candidate: module_with_method_call,
        })
    }
}

impl ImportAssets {
    pub fn import_candidate(&self) -> &ImportCandidate {
        &self.import_candidate
    }

    fn name_to_import(&self) -> &NameToImport {
        match &self.import_candidate {
            ImportCandidate::Path(candidate) => &candidate.name,
            ImportCandidate::TraitAssocItem(candidate)
            | ImportCandidate::TraitMethod(candidate) => &candidate.name,
        }
    }

    pub fn search_for_imports(
        &self,
        sema: &Semantics<RootDatabase>,
        prefix_kind: PrefixKind,
    ) -> Vec<(hir::ModPath, hir::ItemInNs)> {
        let _p = profile::span("import_assets::search_for_imports");
        self.search_for(sema, Some(prefix_kind))
    }

    /// This may return non-absolute paths if a part of the returned path is already imported into scope.
    pub fn search_for_relative_paths(
        &self,
        sema: &Semantics<RootDatabase>,
    ) -> Vec<(hir::ModPath, hir::ItemInNs)> {
        let _p = profile::span("import_assets::search_for_relative_paths");
        self.search_for(sema, None)
    }

    fn search_for(
        &self,
        sema: &Semantics<RootDatabase>,
        prefixed: Option<hir::PrefixKind>,
    ) -> Vec<(hir::ModPath, hir::ItemInNs)> {
        let current_crate = self.module_with_candidate.krate();

        let imports_for_candidate_name = match self.name_to_import() {
            NameToImport::Exact(exact_name) => {
                imports_locator::find_exact_imports(sema, current_crate, exact_name.clone())
            }
            // FIXME: ideally, we should avoid using `fst` for seacrhing trait imports for assoc items:
            // instead, we need to look up all trait impls for a certain struct and search through them only
            // see https://github.com/rust-analyzer/rust-analyzer/pull/7293#issuecomment-761585032
            // and https://rust-lang.zulipchat.com/#narrow/stream/185405-t-compiler.2Fwg-rls-2.2E0/topic/Blanket.20trait.20impls.20lookup
            // for the details
            NameToImport::Fuzzy(fuzzy_name) => {
                let (assoc_item_search, limit) = if self.import_candidate.is_trait_candidate() {
                    (AssocItemSearch::AssocItemsOnly, None)
                } else {
                    (AssocItemSearch::Exclude, Some(DEFAULT_QUERY_SEARCH_LIMIT))
                };
                imports_locator::find_similar_imports(
                    sema,
                    current_crate,
                    fuzzy_name.clone(),
                    assoc_item_search,
                    limit,
                )
            }
        };

        let mut res = self
            .applicable_defs(sema, prefixed, imports_for_candidate_name)
            .filter(|(use_path, _)| use_path.len() > 1)
            .collect::<Vec<_>>();
        res.sort_by_cached_key(|(path, _)| path.clone());
        res
    }

    fn applicable_defs<'a>(
        &self,
        sema: &'a Semantics<RootDatabase>,
        prefixed: Option<hir::PrefixKind>,
        unfiltered_imports: Box<dyn Iterator<Item = Either<ModuleDef, MacroDef>> + 'a>,
    ) -> Box<dyn Iterator<Item = (ModPath, ItemInNs)> + 'a> {
        let current_crate = self.module_with_candidate.krate();
        let db = sema.db;

        match &self.import_candidate {
            ImportCandidate::Path(path_candidate) => path_applicable_defs(
                sema,
                path_candidate,
                unfiltered_imports,
                self.module_with_candidate,
                prefixed,
            ),
            ImportCandidate::TraitAssocItem(trait_candidate) => trait_applicable_defs(
                db,
                current_crate,
                trait_candidate,
                true,
                unfiltered_imports,
                self.module_with_candidate,
                prefixed,
            ),
            ImportCandidate::TraitMethod(trait_candidate) => trait_applicable_defs(
                db,
                current_crate,
                trait_candidate,
                false,
                unfiltered_imports,
                self.module_with_candidate,
                prefixed,
            ),
        }
    }
}

fn path_applicable_defs<'a>(
    sema: &'a Semantics<RootDatabase>,
    path_candidate: &PathImportCandidate,
    unfiltered_defs: Box<dyn Iterator<Item = Either<ModuleDef, MacroDef>> + 'a>,
    module_with_candidate: Module,
    prefixed: Option<hir::PrefixKind>,
) -> Box<dyn Iterator<Item = (ModPath, ItemInNs)> + 'a> {
    let applicable_defs = unfiltered_defs
        .map(|candidate| candidate.either(ItemInNs::from, ItemInNs::from))
        .filter_map(move |item_to_search| {
            get_mod_path(sema.db, item_to_search, &module_with_candidate, prefixed)
                .zip(Some(item_to_search))
        });

    let unresolved_qualifier = match &path_candidate.unresolved_qualifier {
        Some(qualifier) => qualifier,
        None => {
            // TODO kb too many boxes tossed around
            return Box::new(applicable_defs);
        }
    };

    // TODO kb filter out items: found path should end with `qualifier::Name` or `qualifier::Something` for fuzzy search case.

    // TODO kb find a way to turn a qualifier into the corresponding ModuleDef. Maybe through the unfiltered data?
    if let Some(qualifier_start_resolution) = resolve_qualifier_start(sema, unresolved_qualifier) {
        // TODO kb ascend until an unresolved segment part appears
    } else {
        // first segment is already unresolved, need to turn it into ModuleDef somehow
    }

    return Box::new(applicable_defs);
}

fn resolve_qualifier_start(
    sema: &Semantics<RootDatabase>,
    qualifier: &ast::Path,
) -> Option<PathResolution> {
    let qualifier_start = qualifier.syntax().descendants().find_map(ast::NameRef::cast)?;
    let qualifier_start_path = qualifier_start.syntax().ancestors().find_map(ast::Path::cast)?;
    sema.resolve_path(&qualifier_start_path)
}

fn trait_applicable_defs<'a>(
    db: &'a RootDatabase,
    current_crate: Crate,
    trait_candidate: &TraitImportCandidate,
    trait_assoc_item: bool,
    unfiltered_defs: Box<dyn Iterator<Item = Either<ModuleDef, MacroDef>> + 'a>,
    module_with_candidate: Module,
    prefixed: Option<hir::PrefixKind>,
) -> Box<dyn Iterator<Item = (ModPath, ItemInNs)> + 'a> {
    let mut required_assoc_items = FxHashSet::default();

    let trait_candidates = unfiltered_defs
        .filter_map(|input| match input {
            Either::Left(module_def) => module_def.as_assoc_item(db),
            _ => None,
        })
        .filter_map(|assoc| {
            let assoc_item_trait = assoc.containing_trait(db)?;
            required_assoc_items.insert(assoc);
            Some(assoc_item_trait.into())
        })
        .collect();

    let mut applicable_defs = FxHashSet::default();

    if trait_assoc_item {
        trait_candidate.receiver_ty.iterate_path_candidates(
            db,
            current_crate,
            &trait_candidates,
            None,
            |_, assoc| {
                if required_assoc_items.contains(&assoc) {
                    if let AssocItem::Function(f) = assoc {
                        if f.self_param(db).is_some() {
                            return None;
                        }
                    }
                    applicable_defs.insert(assoc_to_module_def(assoc));
                }
                None::<()>
            },
        )
    } else {
        trait_candidate.receiver_ty.iterate_method_candidates(
            db,
            current_crate,
            &trait_candidates,
            None,
            |_, function| {
                let assoc = function.as_assoc_item(db)?;
                if required_assoc_items.contains(&assoc) {
                    applicable_defs.insert(assoc_to_module_def(assoc));
                }
                None::<()>
            },
        )
    };

    Box::new(
        applicable_defs
            .into_iter()
            .filter_map(move |candidate| {
                let canidate_trait = candidate.as_assoc_item(db)?.containing_trait(db)?;
                Some(ItemInNs::from(ModuleDef::from(canidate_trait)))
            })
            .filter_map(move |item_to_search| {
                get_mod_path(db, item_to_search, &module_with_candidate, prefixed)
                    .zip(Some(item_to_search))
            }),
    )
}

fn get_mod_path(
    db: &RootDatabase,
    item_to_search: ItemInNs,
    module_with_candidate: &Module,
    prefixed: Option<hir::PrefixKind>,
) -> Option<ModPath> {
    if let Some(prefix_kind) = prefixed {
        module_with_candidate.find_use_path_prefixed(db, item_to_search, prefix_kind)
    } else {
        module_with_candidate.find_use_path(db, item_to_search)
    }
}

fn assoc_to_module_def(assoc: AssocItem) -> ModuleDef {
    match assoc {
        AssocItem::Function(f) => f.into(),
        AssocItem::Const(c) => c.into(),
        AssocItem::TypeAlias(t) => t.into(),
    }
}

impl ImportCandidate {
    fn for_method_call(
        sema: &Semantics<RootDatabase>,
        method_call: &ast::MethodCallExpr,
    ) -> Option<Self> {
        match sema.resolve_method_call(method_call) {
            Some(_) => None,
            None => Some(Self::TraitMethod(TraitImportCandidate {
                receiver_ty: sema.type_of_expr(&method_call.receiver()?)?,
                name: NameToImport::Exact(method_call.name_ref()?.to_string()),
            })),
        }
    }

    fn for_regular_path(sema: &Semantics<RootDatabase>, path: &ast::Path) -> Option<Self> {
        if sema.resolve_path(path).is_some() {
            return None;
        }
        path_import_candidate(
            sema,
            path.qualifier(),
            NameToImport::Exact(path.segment()?.name_ref()?.to_string()),
        )
    }

    fn for_fuzzy_path(
        qualifier: Option<ast::Path>,
        fuzzy_name: String,
        sema: &Semantics<RootDatabase>,
    ) -> Option<Self> {
        path_import_candidate(sema, qualifier, NameToImport::Fuzzy(fuzzy_name))
    }

    fn is_trait_candidate(&self) -> bool {
        matches!(self, ImportCandidate::TraitAssocItem(_) | ImportCandidate::TraitMethod(_))
    }
}

fn path_import_candidate(
    sema: &Semantics<RootDatabase>,
    qualifier: Option<ast::Path>,
    name: NameToImport,
) -> Option<ImportCandidate> {
    Some(match qualifier {
        Some(qualifier) => match sema.resolve_path(&qualifier) {
            None => ImportCandidate::Path(PathImportCandidate {
                unresolved_qualifier: Some(qualifier),
                name,
            }),
            Some(hir::PathResolution::Def(hir::ModuleDef::Adt(assoc_item_path))) => {
                ImportCandidate::TraitAssocItem(TraitImportCandidate {
                    receiver_ty: assoc_item_path.ty(sema.db),
                    name,
                })
            }
            Some(_) => return None,
        },
        None => ImportCandidate::Path(PathImportCandidate { unresolved_qualifier: None, name }),
    })
}
<commit_msg>Simplify<commit_after>//! Look up accessible paths for items.
use either::Either;
use hir::{
    AsAssocItem, AssocItem, Crate, ItemInNs, MacroDef, ModPath, Module, ModuleDef, PathResolution,
    PrefixKind, Semantics,
};
use rustc_hash::FxHashSet;
use syntax::{ast, AstNode};

use crate::{
    imports_locator::{self, AssocItemSearch, DEFAULT_QUERY_SEARCH_LIMIT},
    RootDatabase,
};

#[derive(Debug)]
pub enum ImportCandidate {
    // A path, qualified (`std::collections::HashMap`) or not (`HashMap`).
    Path(PathImportCandidate),
    /// A trait associated function (with no self parameter) or associated constant.
    /// For 'test_mod::TestEnum::test_function', `ty` is the `test_mod::TestEnum` expression type
    /// and `name` is the `test_function`
    TraitAssocItem(TraitImportCandidate),
    /// A trait method with self parameter.
    /// For 'test_enum.test_method()', `ty` is the `test_enum` expression type
    /// and `name` is the `test_method`
    TraitMethod(TraitImportCandidate),
}

#[derive(Debug)]
pub struct TraitImportCandidate {
    pub receiver_ty: hir::Type,
    pub name: NameToImport,
}

#[derive(Debug)]
pub struct PathImportCandidate {
    pub unresolved_qualifier: Option<ast::Path>,
    pub name: NameToImport,
}

#[derive(Debug)]
pub enum NameToImport {
    Exact(String),
    Fuzzy(String),
}

impl NameToImport {
    pub fn text(&self) -> &str {
        match self {
            NameToImport::Exact(text) => text.as_str(),
            NameToImport::Fuzzy(text) => text.as_str(),
        }
    }
}

#[derive(Debug)]
pub struct ImportAssets {
    import_candidate: ImportCandidate,
    module_with_candidate: hir::Module,
}

impl ImportAssets {
    pub fn for_method_call(
        method_call: &ast::MethodCallExpr,
        sema: &Semantics<RootDatabase>,
    ) -> Option<Self> {
        Some(Self {
            import_candidate: ImportCandidate::for_method_call(sema, method_call)?,
            module_with_candidate: sema.scope(method_call.syntax()).module()?,
        })
    }

    pub fn for_exact_path(
        fully_qualified_path: &ast::Path,
        sema: &Semantics<RootDatabase>,
    ) -> Option<Self> {
        let syntax_under_caret = fully_qualified_path.syntax();
        if syntax_under_caret.ancestors().find_map(ast::Use::cast).is_some() {
            return None;
        }
        Some(Self {
            import_candidate: ImportCandidate::for_regular_path(sema, fully_qualified_path)?,
            module_with_candidate: sema.scope(syntax_under_caret).module()?,
        })
    }

    pub fn for_fuzzy_path(
        module_with_candidate: Module,
        qualifier: Option<ast::Path>,
        fuzzy_name: String,
        sema: &Semantics<RootDatabase>,
    ) -> Option<Self> {
        Some(Self {
            import_candidate: ImportCandidate::for_fuzzy_path(qualifier, fuzzy_name, sema)?,
            module_with_candidate,
        })
    }

    pub fn for_fuzzy_method_call(
        module_with_method_call: Module,
        receiver_ty: hir::Type,
        fuzzy_method_name: String,
    ) -> Option<Self> {
        Some(Self {
            import_candidate: ImportCandidate::TraitMethod(TraitImportCandidate {
                receiver_ty,
                name: NameToImport::Fuzzy(fuzzy_method_name),
            }),
            module_with_candidate: module_with_method_call,
        })
    }
}

impl ImportAssets {
    pub fn import_candidate(&self) -> &ImportCandidate {
        &self.import_candidate
    }

    fn name_to_import(&self) -> &NameToImport {
        match &self.import_candidate {
            ImportCandidate::Path(candidate) => &candidate.name,
            ImportCandidate::TraitAssocItem(candidate)
            | ImportCandidate::TraitMethod(candidate) => &candidate.name,
        }
    }

    pub fn search_for_imports(
        &self,
        sema: &Semantics<RootDatabase>,
        prefix_kind: PrefixKind,
    ) -> Vec<(hir::ModPath, hir::ItemInNs)> {
        let _p = profile::span("import_assets::search_for_imports");
        self.search_for(sema, Some(prefix_kind))
    }

    /// This may return non-absolute paths if a part of the returned path is already imported into scope.
    pub fn search_for_relative_paths(
        &self,
        sema: &Semantics<RootDatabase>,
    ) -> Vec<(hir::ModPath, hir::ItemInNs)> {
        let _p = profile::span("import_assets::search_for_relative_paths");
        self.search_for(sema, None)
    }

    fn search_for(
        &self,
        sema: &Semantics<RootDatabase>,
        prefixed: Option<hir::PrefixKind>,
    ) -> Vec<(hir::ModPath, hir::ItemInNs)> {
        let current_crate = self.module_with_candidate.krate();

        let imports_for_candidate_name = match self.name_to_import() {
            NameToImport::Exact(exact_name) => {
                imports_locator::find_exact_imports(sema, current_crate, exact_name.clone())
            }
            // FIXME: ideally, we should avoid using `fst` for seacrhing trait imports for assoc items:
            // instead, we need to look up all trait impls for a certain struct and search through them only
            // see https://github.com/rust-analyzer/rust-analyzer/pull/7293#issuecomment-761585032
            // and https://rust-lang.zulipchat.com/#narrow/stream/185405-t-compiler.2Fwg-rls-2.2E0/topic/Blanket.20trait.20impls.20lookup
            // for the details
            NameToImport::Fuzzy(fuzzy_name) => {
                let (assoc_item_search, limit) = if self.import_candidate.is_trait_candidate() {
                    (AssocItemSearch::AssocItemsOnly, None)
                } else {
                    (AssocItemSearch::Exclude, Some(DEFAULT_QUERY_SEARCH_LIMIT))
                };
                imports_locator::find_similar_imports(
                    sema,
                    current_crate,
                    fuzzy_name.clone(),
                    assoc_item_search,
                    limit,
                )
            }
        };

        let mut res = self
            .applicable_defs(sema, prefixed, imports_for_candidate_name)
            .filter(|(use_path, _)| use_path.len() > 1)
            .collect::<Vec<_>>();
        res.sort_by_cached_key(|(path, _)| path.clone());
        res
    }

    fn applicable_defs<'a>(
        &'a self,
        sema: &'a Semantics<RootDatabase>,
        prefixed: Option<hir::PrefixKind>,
        unfiltered_defs: impl Iterator<Item = Either<ModuleDef, MacroDef>> + 'a,
    ) -> Box<dyn Iterator<Item = (ModPath, ItemInNs)> + 'a> {
        let current_crate = self.module_with_candidate.krate();
        let db = sema.db;

        match &self.import_candidate {
            ImportCandidate::Path(path_candidate) => Box::new(path_applicable_defs(
                sema,
                path_candidate,
                unfiltered_defs
                    .into_iter()
                    .map(|def| def.either(ItemInNs::from, ItemInNs::from))
                    .filter_map(move |item_to_search| {
                        get_mod_path(db, item_to_search, &self.module_with_candidate, prefixed)
                            .zip(Some(item_to_search))
                    }),
            )),
            ImportCandidate::TraitAssocItem(trait_candidate) => Box::new(
                trait_applicable_defs(db, current_crate, trait_candidate, true, unfiltered_defs)
                    .into_iter()
                    .filter_map(move |item_to_search| {
                        get_mod_path(db, item_to_search, &self.module_with_candidate, prefixed)
                            .zip(Some(item_to_search))
                    }),
            ),
            ImportCandidate::TraitMethod(trait_candidate) => Box::new(
                trait_applicable_defs(db, current_crate, trait_candidate, false, unfiltered_defs)
                    .into_iter()
                    .filter_map(move |item_to_search| {
                        get_mod_path(db, item_to_search, &self.module_with_candidate, prefixed)
                            .zip(Some(item_to_search))
                    }),
            ),
        }
    }
}

fn path_applicable_defs<'a>(
    sema: &'a Semantics<RootDatabase>,
    path_candidate: &PathImportCandidate,
    unfiltered_defs: impl Iterator<Item = (ModPath, ItemInNs)> + 'a,
) -> impl Iterator<Item = (ModPath, ItemInNs)> + 'a {
    let unresolved_qualifier = match &path_candidate.unresolved_qualifier {
        Some(qualifier) => qualifier,
        None => {
            return unfiltered_defs;
        }
    };

    // TODO kb filter out items: found path should end with `qualifier::Name` or `qualifier::Something` for fuzzy search case.

    // TODO kb find a way to turn a qualifier into the corresponding ModuleDef. Maybe through the unfiltered data?
    if let Some(qualifier_start_resolution) = resolve_qualifier_start(sema, unresolved_qualifier) {
        // TODO kb ascend until an unresolved segment part appears
    } else {
        // first segment is already unresolved, need to turn it into ModuleDef somehow
    }

    return unfiltered_defs;
}

fn resolve_qualifier_start(
    sema: &Semantics<RootDatabase>,
    qualifier: &ast::Path,
) -> Option<PathResolution> {
    let qualifier_start = qualifier.syntax().descendants().find_map(ast::NameRef::cast)?;
    let qualifier_start_path = qualifier_start.syntax().ancestors().find_map(ast::Path::cast)?;
    sema.resolve_path(&qualifier_start_path)
}

fn trait_applicable_defs<'a>(
    db: &'a RootDatabase,
    current_crate: Crate,
    trait_candidate: &TraitImportCandidate,
    trait_assoc_item: bool,
    unfiltered_defs: impl Iterator<Item = Either<ModuleDef, MacroDef>> + 'a,
) -> FxHashSet<ItemInNs> {
    let mut required_assoc_items = FxHashSet::default();

    let trait_candidates = unfiltered_defs
        .filter_map(|input| match input {
            Either::Left(module_def) => module_def.as_assoc_item(db),
            _ => None,
        })
        .filter_map(|assoc| {
            let assoc_item_trait = assoc.containing_trait(db)?;
            required_assoc_items.insert(assoc);
            Some(assoc_item_trait.into())
        })
        .collect();

    let mut applicable_traits = FxHashSet::default();

    if trait_assoc_item {
        trait_candidate.receiver_ty.iterate_path_candidates(
            db,
            current_crate,
            &trait_candidates,
            None,
            |_, assoc| {
                if required_assoc_items.contains(&assoc) {
                    if let AssocItem::Function(f) = assoc {
                        if f.self_param(db).is_some() {
                            return None;
                        }
                    }
                    applicable_traits
                        .insert(ItemInNs::from(ModuleDef::from(assoc.containing_trait(db)?)));
                }
                None::<()>
            },
        )
    } else {
        trait_candidate.receiver_ty.iterate_method_candidates(
            db,
            current_crate,
            &trait_candidates,
            None,
            |_, function| {
                let assoc = function.as_assoc_item(db)?;
                if required_assoc_items.contains(&assoc) {
                    applicable_traits
                        .insert(ItemInNs::from(ModuleDef::from(assoc.containing_trait(db)?)));
                }
                None::<()>
            },
        )
    };

    applicable_traits
}

fn get_mod_path(
    db: &RootDatabase,
    item_to_search: ItemInNs,
    module_with_candidate: &Module,
    prefixed: Option<hir::PrefixKind>,
) -> Option<ModPath> {
    if let Some(prefix_kind) = prefixed {
        module_with_candidate.find_use_path_prefixed(db, item_to_search, prefix_kind)
    } else {
        module_with_candidate.find_use_path(db, item_to_search)
    }
}

impl ImportCandidate {
    fn for_method_call(
        sema: &Semantics<RootDatabase>,
        method_call: &ast::MethodCallExpr,
    ) -> Option<Self> {
        match sema.resolve_method_call(method_call) {
            Some(_) => None,
            None => Some(Self::TraitMethod(TraitImportCandidate {
                receiver_ty: sema.type_of_expr(&method_call.receiver()?)?,
                name: NameToImport::Exact(method_call.name_ref()?.to_string()),
            })),
        }
    }

    fn for_regular_path(sema: &Semantics<RootDatabase>, path: &ast::Path) -> Option<Self> {
        if sema.resolve_path(path).is_some() {
            return None;
        }
        path_import_candidate(
            sema,
            path.qualifier(),
            NameToImport::Exact(path.segment()?.name_ref()?.to_string()),
        )
    }

    fn for_fuzzy_path(
        qualifier: Option<ast::Path>,
        fuzzy_name: String,
        sema: &Semantics<RootDatabase>,
    ) -> Option<Self> {
        path_import_candidate(sema, qualifier, NameToImport::Fuzzy(fuzzy_name))
    }

    fn is_trait_candidate(&self) -> bool {
        matches!(self, ImportCandidate::TraitAssocItem(_) | ImportCandidate::TraitMethod(_))
    }
}

fn path_import_candidate(
    sema: &Semantics<RootDatabase>,
    qualifier: Option<ast::Path>,
    name: NameToImport,
) -> Option<ImportCandidate> {
    Some(match qualifier {
        Some(qualifier) => match sema.resolve_path(&qualifier) {
            None => ImportCandidate::Path(PathImportCandidate {
                unresolved_qualifier: Some(qualifier),
                name,
            }),
            Some(hir::PathResolution::Def(hir::ModuleDef::Adt(assoc_item_path))) => {
                ImportCandidate::TraitAssocItem(TraitImportCandidate {
                    receiver_ty: assoc_item_path.ty(sema.db),
                    name,
                })
            }
            Some(_) => return None,
        },
        None => ImportCandidate::Path(PathImportCandidate { unresolved_qualifier: None, name }),
    })
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::collections::hash_map::HashMap;
use std::collections::hash_set::HashSet;
use std::cmp;
use std::mem;

use js_types::js_type::{JsVar, JsType, JsPtrEnum};
use uuid::Uuid;

// Initial Arena size in bytes
const INITIAL_SIZE: usize = 1024;
// Minimum Arena capacity is at least 1 byte
const MIN_CAP: usize = 1;

pub struct Scope {
    parent: Option<Box<Scope>>,
    children: Vec<Box<Scope>>,
    black_set: HashMap<Uuid, RefCell<JsVar>>,
    grey_set: HashMap<Uuid, RefCell<JsVar>>,
    white_set: HashMap<Uuid, RefCell<JsVar>>,
    get_roots: Box<Fn() -> HashSet<Uuid>>,
}

impl Scope {
    pub fn new<F>(get_roots: F) -> Scope
        where F: Fn() -> HashSet<Uuid> + 'static {
        Scope {
            parent: None,
            children: Vec::new(),
            black_set: HashMap::new(),
            grey_set: HashMap::new(),
            white_set: HashMap::new(),
            get_roots: Box::new(get_roots),
        }
    }

    pub fn as_child<F>(parent: Box<Scope>, get_roots: F) -> Scope
        where F: Fn() -> HashSet<Uuid> + 'static {
        Scope {
            parent: Some(parent),
            children: Vec::new(),
            black_set: HashMap::new(),
            grey_set: HashMap::new(),
            white_set: HashMap::new(),
            get_roots: Box::new(get_roots),
        }
    }

    pub fn set_parent(&mut self, parent: Box<Scope>) {
        self.parent = Some(parent);
    }

    pub fn add_child(&mut self, child: Box<Scope>) {
        self.children.push(child);
    }

    pub fn alloc(&mut self, var: JsVar) -> Uuid {
        let uuid = var.uuid;
        self.white_set.insert(uuid, RefCell::new(var));
        uuid
    }

    pub fn dealloc(&mut self, uuid: &Uuid) -> bool {
        if let Some(_) = self.white_set.remove(uuid) { true } else { false }
    }

    pub fn get_var_copy(&self, uuid: &Uuid) -> Option<JsVar> {
        if let Some(var) = self.black_set.get(uuid) {
            Some(var.clone().into_inner())
        } else if let Some(var) = self.grey_set.get(uuid) {
            Some(var.clone().into_inner())
        } else if let Some(var) = self.white_set.get(uuid) {
            Some(var.clone().into_inner())
        } else { None }
    }

    pub fn update_var(&mut self, var: JsVar) -> bool {
        unimplemented!()
    }

    /// TODO Compute the roots of the current scope-- any variable that is
    /// directly referenced or declared within the scope. This might just be the
    /// key set of the uuid map(?) Not necessarily, I think. What if you do
    /// something like this:
    /// var x = {}
    /// var y = { 1: x }
    /// y = x
    /// y would be a root by the definition above, but is no longer reachable at
    /// the end of the scope because it now aliases x. A better definition would
    /// be "Any variable that is declared or referenced directly, but a direct
    /// reference (variable usage) supercedes a declaration." The above example
    /// demonstrates why this is necessary.
    /// This should come from the interpreter, so I shouldn't actually have to
    /// care about getting the root set myself.

    //pub fn compute_roots(&self) -> HashSet<Uuid> {
    //    self.get_roots();
    //}

    /// Roots always get marked as Black, since they're always reachable from
    /// the current scope. NB that this assumes all root references are actually
    /// valid reference types, i.e. they're not numbers, etc.
    pub fn mark_roots(&mut self, marks: HashSet<Uuid>) {
        for mark in marks.iter() {
            if let Some(var) = self.white_set.remove(mark) {
                let uuid = var.borrow().uuid;
                // Get all child references
                let child_ids = self.get_var_children(&var);
                self.black_set.insert(uuid, var);
                // Mark child references as grey
                self.grey_children(child_ids);
            }
        }
    }

    pub fn mark_phase(&mut self) {
        // Mark any grey object as black, and mark all white objs it refs as grey
        while let Some(&uuid) = self.grey_set.keys().take(1).next() {
            if let Some(var) = self.grey_set.remove(&uuid) {
                let child_ids = self.get_var_children(&var);
                self.black_set.insert(uuid, var);
                for child_id in child_ids {
                    if let Some(var) = self.white_set.remove(&child_id) {
                        self.grey_set.insert(child_id, var);
                    }
                }
            }
        }
    }

    fn grey_children(&mut self, child_ids: HashSet<Uuid>) {
        for child_id in child_ids {
            if let Some(var) = self.white_set.remove(&child_id) {
                self.grey_set.insert(child_id, var);
            }
        }
    }

    fn get_var_children(&self, var: &RefCell<JsVar>) -> HashSet<Uuid> {
        if let JsType::JsPtr(ref ptr) = (*var.borrow()).t {
            match ptr {
                &JsPtrEnum::JsObj(ref obj) => obj.get_children(),
                _ => HashSet::new(),
            }
        } else { HashSet::new() }
    }

}

#[cfg(test)]
mod tests {
    use super::*;
    use js_types::js_type::{JsVar, JsType, Marking};

    #[test]
    fn test_new_scope() {
        let test_jst0 = JsVar::new(JsType::JsNum(0.0));
        let test_jst1 = JsVar::new(JsType::JsNum(1.0));

        let mut test_scope = Scope::new("test");
        let uuid0 = test_scope.alloc(test_jst0);
        test_scope.mark_uuid(&uuid0, Marking::Black);
        let uuid1 = test_scope.alloc(test_jst1);
        test_scope.mark_uuid(&uuid1, Marking::White);
        assert_eq!(test_scope.get_jst_copy(&uuid0).unwrap().gc_flag, Marking::Black);
        assert_eq!(test_scope.get_jst_copy(&uuid1).unwrap().gc_flag, Marking::White);
    }
}
<commit_msg>Change Scope parent into an Rc to avoid owning it & improve Scope helpers<commit_after>use std::cell::RefCell;
use std::collections::hash_map::HashMap;
use std::collections::hash_set::HashSet;
use std::rc::Rc;
use std::cmp;
use std::mem;

use js_types::js_type::{JsVar, JsType, JsPtrEnum};
use uuid::Uuid;

// Initial Arena size in bytes
const INITIAL_SIZE: usize = 1024;
// Minimum Arena capacity is at least 1 byte
const MIN_CAP: usize = 1;

pub struct Scope {
    parent: Option<Rc<Scope>>,
    children: Vec<Box<Scope>>,
    black_set: HashMap<Uuid, RefCell<JsVar>>,
    grey_set: HashMap<Uuid, RefCell<JsVar>>,
    white_set: HashMap<Uuid, RefCell<JsVar>>,
    get_roots: Box<Fn() -> HashSet<Uuid>>,
}

impl Scope {
    pub fn new<F>(get_roots: F) -> Scope
        where F: Fn() -> HashSet<Uuid> + 'static {
        Scope {
            parent: None,
            children: Vec::new(),
            black_set: HashMap::new(),
            grey_set: HashMap::new(),
            white_set: HashMap::new(),
            get_roots: Box::new(get_roots),
        }
    }

    pub fn as_child<F>(parent: Rc<Scope>, get_roots: F) -> Scope
        where F: Fn() -> HashSet<Uuid> + 'static {
        Scope {
            parent: Some(parent),
            children: Vec::new(),
            black_set: HashMap::new(),
            grey_set: HashMap::new(),
            white_set: HashMap::new(),
            get_roots: Box::new(get_roots),
        }
    }

    pub fn set_parent(&mut self, parent: Rc<Scope>) {
        self.parent = Some(parent);
    }

    pub fn add_child(&mut self, child: Scope) {
        self.children.push(Box::new(child));
    }

    pub fn alloc(&mut self, var: JsVar) -> Uuid {
        let uuid = var.uuid;
        self.white_set.insert(uuid, RefCell::new(var));
        uuid
    }

    pub fn dealloc(&mut self, uuid: &Uuid) -> bool {
        if let Some(_) = self.white_set.remove(uuid) { true } else { false }
    }

    pub fn get_var_copy(&self, uuid: &Uuid) -> Option<JsVar> {
        self.find_id(uuid).map(|var| var.clone().into_inner())
    }

    pub fn update_var(&mut self, var: JsVar) -> bool {
        unimplemented!()
    }

    /// TODO Compute the roots of the current scope-- any variable that is
    /// directly referenced or declared within the scope. This might just be the
    /// key set of the uuid map(?) Not necessarily, I think. What if you do
    /// something like this:
    /// var x = {}
    /// var y = { 1: x }
    /// y = x
    /// y would be a root by the definition above, but is no longer reachable at
    /// the end of the scope because it now aliases x. A better definition would
    /// be "Any variable that is declared or referenced directly, but a direct
    /// reference (variable usage) supercedes a declaration." The above example
    /// demonstrates why this is necessary.
    /// This should come from the interpreter, so I shouldn't actually have to
    /// care about getting the root set myself.

    //pub fn compute_roots(&self) -> HashSet<Uuid> {
    //    self.get_roots();
    //}

    /// Roots always get marked as Black, since they're always reachable from
    /// the current scope. NB that this assumes all root references are actually
    /// valid reference types, i.e. they're not numbers, etc.
    pub fn mark_roots(&mut self, marks: HashSet<Uuid>) {
        for mark in marks.iter() {
            if let Some(var) = self.white_set.remove(mark) {
                let uuid = var.borrow().uuid;
                // Get all child references
                let child_ids = self.get_var_children(&var);
                self.black_set.insert(uuid, var);
                // Mark child references as grey
                self.grey_children(child_ids);
            }
        }
    }

    pub fn mark_phase(&mut self) {
        // Mark any grey object as black, and mark all white objs it refs as grey
        while let Some(&uuid) = self.grey_set.keys().take(1).next() {
            if let Some(var) = self.grey_set.remove(&uuid) {
                let child_ids = self.get_var_children(&var);
                self.black_set.insert(uuid, var);
                for child_id in child_ids {
                    if let Some(var) = self.white_set.remove(&child_id) {
                        self.grey_set.insert(child_id, var);
                    }
                }
            }
        }
    }

    pub fn sweep_phase(&mut self) {
        self.white_set.clear();
        self.white_set.shrink_to_fit();
    }

    fn find_id(&self, uuid: &Uuid) -> Option<&RefCell<JsVar>> {
        self.black_set.get(uuid).or_else(||
            self.grey_set.get(uuid).or_else(||
                self.white_set.get(uuid)))
    }

    fn grey_children(&mut self, child_ids: HashSet<Uuid>) {
        for child_id in child_ids {
            if let Some(var) = self.white_set.remove(&child_id) {
                self.grey_set.insert(child_id, var);
            }
        }
    }

    fn get_var_children(&self, var: &RefCell<JsVar>) -> HashSet<Uuid> {
        if let JsType::JsPtr(ref ptr) = (*var.borrow()).t {
            match ptr {
                &JsPtrEnum::JsObj(ref obj) => obj.get_children(),
                _ => HashSet::new(),
            }
        } else { HashSet::new() }
    }

}

#[cfg(test)]
mod tests {
    use super::*;
    use js_types::js_type::{JsVar, JsType, Marking};

    #[test]
    fn test_new_scope() {
        let test_jst0 = JsVar::new(JsType::JsNum(0.0));
        let test_jst1 = JsVar::new(JsType::JsNum(1.0));

        let mut test_scope = Scope::new("test");
        let uuid0 = test_scope.alloc(test_jst0);
        test_scope.mark_uuid(&uuid0, Marking::Black);
        let uuid1 = test_scope.alloc(test_jst1);
        test_scope.mark_uuid(&uuid1, Marking::White);
        assert_eq!(test_scope.get_jst_copy(&uuid0).unwrap().gc_flag, Marking::Black);
        assert_eq!(test_scope.get_jst_copy(&uuid1).unwrap().gc_flag, Marking::White);
    }
}
<|endoftext|>


<commit_before>use nom;
use nom::IResult::*;

use helper::*;
use super::{Date,Time,DateTime};

use macros::{take_2_digits,take_4_digits};

named!(sign <i32>, alt!(
        tag!("-") => { |_| -1 } |
        tag!("+") => { |_| 1 }
        )
    );

named!(positive_year  <i32>, map!(call!(take_4_digits), buf_to_i32));
named!(pub year <i32>, chain!(
        pref: opt!(sign) ~
        y:    positive_year
        ,
        || {
            pref.unwrap_or(1) * y
        }));

named!(pub month <u32>, map!(call!(take_2_digits), buf_to_u32));
named!(pub day   <u32>, map!(call!(take_2_digits), buf_to_u32));

named!(pub date <Date>, chain!(
        y: year ~ tag!("-") ~ m: month ~ tag!("-") ~ d: day
        , || { Date{ year: y, month: m, day: d }
        }
        ));

named!(pub hour   <u32>, map!(call!(take_2_digits), buf_to_u32));
named!(pub minute <u32>, map!(call!(take_2_digits), buf_to_u32));
named!(pub second <u32>, map!(call!(take_2_digits), buf_to_u32));

named!(pub time <Time>, chain!(
        h: hour      ~
           tag!(":") ~
        m: minute    ~
        s: empty_or!(chain!(
                tag!(":") ~ s:second , || { s }))
        ,
        || { Time{ hour: h,
                   minute: m,
                   second: s.unwrap_or(0),
                   tz_offset: 0 }
           }
        ));

named!(timezone_hour <i32>, chain!(
        s: sign ~
        h: hour ~
        m: empty_or!(
            chain!(
                tag!(":")? ~ m: minute , || { m }
            ))
        ,
        || { (s * (h as i32) * 3600) + (m.unwrap_or(0) * 60) as i32 }
        ));

named!(timezone_utc <i32>, map!(alt!(tag!("Z") | tag!("+")), |_| 0));
named!(pub timezone <i32>, alt!(timezone_utc | timezone_hour));

named!(pub datetime <DateTime>, chain!(
        d:   date      ~
             tag!("T") ~
        t:   time      ~
        tzo: empty_or!(call!(timezone))
        ,
        || {
            DateTime{
                date: d,
                time: t.set_tz(tzo.unwrap_or(0)),
            }
        }
        ));

named!(pub time_with_timezone <Time>, chain!(
        t:   time      ~
        tzo: empty_or!(call!(timezone))
        ,
        || {
            t.set_tz(tzo.unwrap_or(0)); t
        }
        ));
<commit_msg>finally removed easy.rs<commit_after><|endoftext|>