<commit_before>use std::env;
use std::fs;
use std::path::{Path, PathBuf};

use cfg::CargonautsConfig;
use heck::KebabCase;
use quote::{ToTokens, Tokens, Ident};

use ast::*;

pub struct Routes {
    routes: Vec<Route>,
}

impl Routes {
    pub fn new(resources: &[Resource], cfg: Option<&CargonautsConfig>) -> Routes {
        Routes {
            routes: Route::build(resources, cfg),
        }
    }
}

impl ToTokens for Routes {
    fn to_tokens(&self, tokens: &mut Tokens) {
        let routes = &self.routes;
        tokens.append(quote! ({
            use ::cargonauts::server::NewService;
            let routes = vec!#routes.into_iter().collect();
            ::cargonauts::routing::RoutingTable::new(routes, assets, asset_handler, env)
        }));
    }
}

#[derive(Clone)]
struct Route {
    resource: String,
    endpoint: String,
    method: String,
    format: String,
    rel: Option<String>,
    middleware: Option<String>,
    template_root: PathBuf,
}

impl Route {
    fn build(resources: &[Resource], cfg: Option<&CargonautsConfig>) -> Vec<Route> {
        let mut vec = vec![];

        let root = {
            let root: PathBuf = env::var("CARGO_MANIFEST_DIR").unwrap().into();
            let sub = cfg.and_then(|cfg| cfg.templates()).unwrap_or(Path::new("src/templates"));
            root.join(sub)
        };

        for resource in resources {
            let endpoint = resource.header.endpoint.clone().unwrap_or_else(|| {
                resource.header.ty.to_kebab_case()
            });

            let resource_ty = &resource.header.ty;

            for (method, attrs) in resource.members.iter().filter_map(|m| m.as_method()) {
                let middleware = attrs.iter().filter_map(|attr| attr.as_middleware()).next();
                vec.push(Route {
                    resource: resource_ty.clone(),
                    endpoint: endpoint.clone(),
                    method: method.method.clone(),
                    format: method.format.clone(),
                    rel: None,
                    middleware,
                    template_root: root.clone(),
                })
            }

            for relation in resource.members.iter().filter_map(|m| m.as_relation()) {
                for (method, attrs) in relation.members.iter().filter_map(|m| m.as_method()) {
                    let middleware = attrs.iter().filter_map(|attr| attr.as_middleware()).next();
                    vec.push(Route {
                        resource: resource_ty.clone(),
                        endpoint: endpoint.clone(),
                        method: method.method.clone(),
                        format: method.format.clone(),
                        rel: Some(relation.rel.clone()),
                        middleware,
                        template_root: root.clone(),
                    })
                }
            }
        }

        vec
    }
}

impl ToTokens for Route {
    fn to_tokens(&self, tokens: &mut Tokens) {
        let endpoint = &self.endpoint;
        let template = load_template(self.template_root.clone(),
                                     &self.resource,
                                     &self.method,
                                     self.rel.as_ref());
        let resource = Ident::new(&self.resource[..]);
        let format = Ident::new(&self.format[..]);
        let method = method_for(&self.method, self.rel.as_ref(), &resource);

        let route_key = if let Some(ref rel) = self.rel {
            quote!(::cargonauts::routing::RouteKey::new(#endpoint, <#method as ::cargonauts::method::Method<#resource>>::ROUTE, Some(#rel)))
        } else {
            quote!(::cargonauts::routing::RouteKey::new(#endpoint, <#method as ::cargonauts::method::Method<#resource>>::ROUTE, None))
        };

        if let Some(ref middleware) = self.middleware {
            let middleware = Ident::new(&middleware[..]);
            tokens.append(quote!({
                let route_key = #route_key;
                let service = ::cargonauts::routing::EndpointService::<_, _, (#resource, #format, #method)>::new(#template);
                let middleware = <#middleware as Default>::default():
                let service = ::cargonauts::middleware::Middleware::wrap(middleware, service)
                (route_key, Box::new(serice) as ::cargonatus::routing::Handler)
            }));
        } else {
            tokens.append(quote!({
                let route_key = #route_key;
                let service = ::cargonauts::routing::EndpointService::<_, _, (#resource, #format, #method)>::new(#template);
                (route_key, Box::new(service) as ::cargonauts::routing::Handler)
            }));
        }

    }
}

fn method_for(method: &str, rel: Option<&String>, resource: &Ident) -> Tokens {
    // Special cases to allow mainsail methods to have associated types
    match method {
        "Post"  => {
            quote!(Post<Identifier = <#resource as ::cargonauts::api::Resource>::Identifier, Post = <#resource as ::cargonauts::api::Post>::Post>)
        }
        "Patch" => {
            quote!(Patch<Identifier = <#resource as ::cargonauts::api::Resource>::Identifier, Patch = <#resource as ::cargonauts::api::Patch>::Patch>)
        }
        "PostRelated" => {
            let rel = Ident::new(&rel.unwrap()[..]);
            quote!(PostRelated<#rel, Identifier = <#resource as ::cargonauts::api::Resource>::Identifier, Post = <#resource as ::cargonauts::api::PostRelated<#rel>::Post>)
        }
        "UpdateRelated" => {
            let rel = Ident::new(&rel.unwrap()[..]);
            quote!(UpdateRelated<#rel, Identifier = <#resource as ::cargonauts::api::Resource>::Identifier, Update = <#resource as ::cargonauts::api::UpdateRelated<#rel>::Update>)
        }
        _       => {
            let method = Ident::new(method);
            match rel {
                Some(rel)   => {
                    let rel = Ident::new(&rel[..]);
                    quote!(#method<#rel, Identifier = <#resource as ::cargonauts::api::Resource>::Identifier>)
                }
                None        => {
                    quote!(#method<Identifier = <#resource as ::cargonauts::api::Resource>::Identifier>)
                }
            }
        }
    }
}

fn load_template(mut path: PathBuf, resource: &str, method: &str, rel: Option<&String>) -> Tokens {
    path.push(resource);
    if let Some(rel) = rel { path.push(rel); }

    let match_method = |entry: &fs::DirEntry| {
        entry.file_name().into_string().unwrap().starts_with(&method.to_kebab_case())
    };

    let mut walk_dir = fs::read_dir(path).into_iter().flat_map(|ls| ls.into_iter().map(|e| e.unwrap()));

    if let Some(entry) = walk_dir.find(match_method) {
        let path = entry.path();
        let path = path.to_string_lossy();
        quote!(Some(::cargonauts::format::Template::static_prepare(include_str!(#path))))
    } else {
        quote!(None)
    }
}
<commit_msg>Fix types in middleware codegen.<commit_after>use std::env;
use std::fs;
use std::path::{Path, PathBuf};

use cfg::CargonautsConfig;
use heck::KebabCase;
use quote::{ToTokens, Tokens, Ident};

use ast::*;

pub struct Routes {
    routes: Vec<Route>,
}

impl Routes {
    pub fn new(resources: &[Resource], cfg: Option<&CargonautsConfig>) -> Routes {
        Routes {
            routes: Route::build(resources, cfg),
        }
    }
}

impl ToTokens for Routes {
    fn to_tokens(&self, tokens: &mut Tokens) {
        let routes = &self.routes;
        tokens.append(quote! ({
            use ::cargonauts::server::NewService;
            let routes = vec!#routes.into_iter().collect();
            ::cargonauts::routing::RoutingTable::new(routes, assets, asset_handler, env)
        }));
    }
}

#[derive(Clone)]
struct Route {
    resource: String,
    endpoint: String,
    method: String,
    format: String,
    rel: Option<String>,
    middleware: Option<String>,
    template_root: PathBuf,
}

impl Route {
    fn build(resources: &[Resource], cfg: Option<&CargonautsConfig>) -> Vec<Route> {
        let mut vec = vec![];

        let root = {
            let root: PathBuf = env::var("CARGO_MANIFEST_DIR").unwrap().into();
            let sub = cfg.and_then(|cfg| cfg.templates()).unwrap_or(Path::new("src/templates"));
            root.join(sub)
        };

        for resource in resources {
            let endpoint = resource.header.endpoint.clone().unwrap_or_else(|| {
                resource.header.ty.to_kebab_case()
            });

            let resource_ty = &resource.header.ty;

            for (method, attrs) in resource.members.iter().filter_map(|m| m.as_method()) {
                let middleware = attrs.iter().filter_map(|attr| attr.as_middleware()).next();
                vec.push(Route {
                    resource: resource_ty.clone(),
                    endpoint: endpoint.clone(),
                    method: method.method.clone(),
                    format: method.format.clone(),
                    rel: None,
                    middleware,
                    template_root: root.clone(),
                })
            }

            for relation in resource.members.iter().filter_map(|m| m.as_relation()) {
                for (method, attrs) in relation.members.iter().filter_map(|m| m.as_method()) {
                    let middleware = attrs.iter().filter_map(|attr| attr.as_middleware()).next();
                    vec.push(Route {
                        resource: resource_ty.clone(),
                        endpoint: endpoint.clone(),
                        method: method.method.clone(),
                        format: method.format.clone(),
                        rel: Some(relation.rel.clone()),
                        middleware,
                        template_root: root.clone(),
                    })
                }
            }
        }

        vec
    }
}

impl ToTokens for Route {
    fn to_tokens(&self, tokens: &mut Tokens) {
        let endpoint = &self.endpoint;
        let template = load_template(self.template_root.clone(),
                                     &self.resource,
                                     &self.method,
                                     self.rel.as_ref());
        let resource = Ident::new(&self.resource[..]);
        let format = Ident::new(&self.format[..]);
        let method = method_for(&self.method, self.rel.as_ref(), &resource);

        let route_key = if let Some(ref rel) = self.rel {
            quote!(::cargonauts::routing::RouteKey::new(#endpoint, <#method as ::cargonauts::method::Method<#resource>>::ROUTE, Some(#rel)))
        } else {
            quote!(::cargonauts::routing::RouteKey::new(#endpoint, <#method as ::cargonauts::method::Method<#resource>>::ROUTE, None))
        };

        if let Some(ref middleware) = self.middleware {
            let middleware = Ident::new(&middleware[..]);
            tokens.append(quote!({
                let route_key = #route_key;
                let service = ::cargonauts::routing::EndpointService::<_, _, (#resource, #format, #method)>::new(#template);
                let middleware = <#middleware as Default>::default();
                let service = ::cargonauts::middleware::Middleware::wrap(middleware, service);
                (route_key, Box::new(service) as ::cargonauts::routing::Handler)
            }));
        } else {
            tokens.append(quote!({
                let route_key = #route_key;
                let service = ::cargonauts::routing::EndpointService::<_, _, (#resource, #format, #method)>::new(#template);
                (route_key, Box::new(service) as ::cargonauts::routing::Handler)
            }));
        }

    }
}

fn method_for(method: &str, rel: Option<&String>, resource: &Ident) -> Tokens {
    // Special cases to allow mainsail methods to have associated types
    match method {
        "Post"  => {
            quote!(Post<Identifier = <#resource as ::cargonauts::api::Resource>::Identifier, Post = <#resource as ::cargonauts::api::Post>::Post>)
        }
        "Patch" => {
            quote!(Patch<Identifier = <#resource as ::cargonauts::api::Resource>::Identifier, Patch = <#resource as ::cargonauts::api::Patch>::Patch>)
        }
        "PostRelated" => {
            let rel = Ident::new(&rel.unwrap()[..]);
            quote!(PostRelated<#rel, Identifier = <#resource as ::cargonauts::api::Resource>::Identifier, Post = <#resource as ::cargonauts::api::PostRelated<#rel>::Post>)
        }
        "UpdateRelated" => {
            let rel = Ident::new(&rel.unwrap()[..]);
            quote!(UpdateRelated<#rel, Identifier = <#resource as ::cargonauts::api::Resource>::Identifier, Update = <#resource as ::cargonauts::api::UpdateRelated<#rel>::Update>)
        }
        _       => {
            let method = Ident::new(method);
            match rel {
                Some(rel)   => {
                    let rel = Ident::new(&rel[..]);
                    quote!(#method<#rel, Identifier = <#resource as ::cargonauts::api::Resource>::Identifier>)
                }
                None        => {
                    quote!(#method<Identifier = <#resource as ::cargonauts::api::Resource>::Identifier>)
                }
            }
        }
    }
}

fn load_template(mut path: PathBuf, resource: &str, method: &str, rel: Option<&String>) -> Tokens {
    path.push(resource);
    if let Some(rel) = rel { path.push(rel); }

    let match_method = |entry: &fs::DirEntry| {
        entry.file_name().into_string().unwrap().starts_with(&method.to_kebab_case())
    };

    let mut walk_dir = fs::read_dir(path).into_iter().flat_map(|ls| ls.into_iter().map(|e| e.unwrap()));

    if let Some(entry) = walk_dir.find(match_method) {
        let path = entry.path();
        let path = path.to_string_lossy();
        quote!(Some(::cargonauts::format::Template::static_prepare(include_str!(#path))))
    } else {
        quote!(None)
    }
}
<|endoftext|>


<commit_before>//! This module handles fuzzy-searching of functions, structs and other symbols
//! by name across the whole workspace and dependencies.
//!
//! It works by building an incrementally-updated text-search index of all
//! symbols. The backbone of the index is the **awesome** `fst` crate by
//! @BurntSushi.
//!
//! In a nutshell, you give a set of strings to `fst`, and it builds a
//! finite state machine describing this set of strings. The strings which
//! could fuzzy-match a pattern can also be described by a finite state machine.
//! What is freaking cool is that you can now traverse both state machines in
//! lock-step to enumerate the strings which are both in the input set and
//! fuzz-match the query. Or, more formally, given two languages described by
//! FSTs, one can build a product FST which describes the intersection of the
//! languages.
//!
//! `fst` does not support cheap updating of the index, but it supports unioning
//! of state machines. So, to account for changing source code, we build an FST
//! for each library (which is assumed to never change) and an FST for each Rust
//! file in the current workspace, and run a query against the union of all
//! those FSTs.
use std::{
    cmp::Ordering,
    hash::{Hash, Hasher},
    sync::Arc,
    mem,
    fmt,
};

use fst::{self, Streamer};
use ra_syntax::{
    SyntaxNode, SyntaxNodePtr, SourceFile, SmolStr, TreeArc, AstNode,
    algo::{visit::{visitor, Visitor}},
    SyntaxKind::{self, *},
    ast::{self, NameOwner},
    WalkEvent,
    TextRange,
};
use ra_db::{
    SourceRootId, SourceDatabase,
    salsa::{self, ParallelDatabase},
};
use rayon::prelude::*;

use crate::{
    FileId, Query,
    db::RootDatabase,
};

#[salsa::query_group(SymbolsDatabaseStorage)]
pub(crate) trait SymbolsDatabase: hir::db::HirDatabase {
    fn file_symbols(&self, file_id: FileId) -> Arc<SymbolIndex>;
    #[salsa::input]
    fn library_symbols(&self, id: SourceRootId) -> Arc<SymbolIndex>;
    /// The set of "local" (that is, from the current workspace) roots.
    /// Files in local roots are assumed to change frequently.
    #[salsa::input]
    fn local_roots(&self) -> Arc<Vec<SourceRootId>>;
    /// The set of roots for crates.io libraries.
    /// Files in libraries are assumed to never change.
    #[salsa::input]
    fn library_roots(&self) -> Arc<Vec<SourceRootId>>;
}

fn file_symbols(db: &impl SymbolsDatabase, file_id: FileId) -> Arc<SymbolIndex> {
    db.check_canceled();
    let source_file = db.parse(file_id);

    let symbols = source_file_to_file_symbols(&source_file, file_id);

    // FIXME: add macros here

    Arc::new(SymbolIndex::new(symbols))
}

pub(crate) fn world_symbols(db: &RootDatabase, query: Query) -> Vec<FileSymbol> {
    /// Need to wrap Snapshot to provide `Clone` impl for `map_with`
    struct Snap(salsa::Snapshot<RootDatabase>);
    impl Clone for Snap {
        fn clone(&self) -> Snap {
            Snap(self.0.snapshot())
        }
    }

    let buf: Vec<Arc<SymbolIndex>> = if query.libs {
        let snap = Snap(db.snapshot());
        db.library_roots()
            .par_iter()
            .map_with(snap, |db, &lib_id| db.0.library_symbols(lib_id))
            .collect()
    } else {
        let mut files = Vec::new();
        for &root in db.local_roots().iter() {
            let sr = db.source_root(root);
            files.extend(sr.files.values().map(|&it| it))
        }

        let snap = Snap(db.snapshot());
        files.par_iter().map_with(snap, |db, &file_id| db.0.file_symbols(file_id)).collect()
    };
    query.search(&buf)
}

pub(crate) fn index_resolve(db: &RootDatabase, name_ref: &ast::NameRef) -> Vec<FileSymbol> {
    let name = name_ref.text();
    let mut query = Query::new(name.to_string());
    query.exact();
    query.limit(4);
    crate::symbol_index::world_symbols(db, query)
}

#[derive(Default)]
pub(crate) struct SymbolIndex {
    symbols: Vec<FileSymbol>,
    map: fst::Map,
}

impl fmt::Debug for SymbolIndex {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("SymbolIndex").field("n_symbols", &self.symbols.len()).finish()
    }
}

impl PartialEq for SymbolIndex {
    fn eq(&self, other: &SymbolIndex) -> bool {
        self.symbols == other.symbols
    }
}

impl Eq for SymbolIndex {}

impl Hash for SymbolIndex {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.symbols.hash(hasher)
    }
}

impl SymbolIndex {
    fn new(mut symbols: Vec<FileSymbol>) -> SymbolIndex {
        fn cmp(s1: &FileSymbol, s2: &FileSymbol) -> Ordering {
            unicase::Ascii::new(s1.name.as_str()).cmp(&unicase::Ascii::new(s2.name.as_str()))
        }
        symbols.par_sort_by(cmp);
        symbols.dedup_by(|s1, s2| cmp(s1, s2) == Ordering::Equal);
        let names = symbols.iter().map(|it| it.name.as_str().to_lowercase());
        let map = fst::Map::from_iter(names.zip(0u64..)).unwrap();
        SymbolIndex { symbols, map }
    }

    pub(crate) fn len(&self) -> usize {
        self.symbols.len()
    }

    pub(crate) fn memory_size(&self) -> usize {
        self.map.as_fst().size() + self.symbols.len() * mem::size_of::<FileSymbol>()
    }

    pub(crate) fn for_files(
        files: impl ParallelIterator<Item = (FileId, TreeArc<SourceFile>)>,
    ) -> SymbolIndex {
        let symbols = files
            .flat_map(|(file_id, file)| source_file_to_file_symbols(&file, file_id))
            .collect::<Vec<_>>();
        SymbolIndex::new(symbols)
    }
}

impl Query {
    pub(crate) fn search(self, indices: &[Arc<SymbolIndex>]) -> Vec<FileSymbol> {
        let mut op = fst::map::OpBuilder::new();
        for file_symbols in indices.iter() {
            let automaton = fst::automaton::Subsequence::new(&self.lowercased);
            op = op.add(file_symbols.map.search(automaton))
        }
        let mut stream = op.union();
        let mut res = Vec::new();
        while let Some((_, indexed_values)) = stream.next() {
            if res.len() >= self.limit {
                break;
            }
            for indexed_value in indexed_values {
                let file_symbols = &indices[indexed_value.index];
                let idx = indexed_value.value as usize;

                let symbol = &file_symbols.symbols[idx];
                if self.only_types && !is_type(symbol.ptr.kind()) {
                    continue;
                }
                if self.exact && symbol.name != self.query {
                    continue;
                }
                res.push(symbol.clone());
            }
        }
        res
    }
}

fn is_type(kind: SyntaxKind) -> bool {
    match kind {
        STRUCT_DEF | ENUM_DEF | TRAIT_DEF | TYPE_ALIAS_DEF => true,
        _ => false,
    }
}

/// The actual data that is stored in the index. It should be as compact as
/// possible.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub(crate) struct FileSymbol {
    pub(crate) file_id: FileId,
    pub(crate) name: SmolStr,
    pub(crate) ptr: SyntaxNodePtr,
    pub(crate) name_range: Option<TextRange>,
    pub(crate) container_name: Option<SmolStr>,
}

fn source_file_to_file_symbols(source_file: &SourceFile, file_id: FileId) -> Vec<FileSymbol> {
    let mut symbols = Vec::new();
    let mut stack = Vec::new();

    for event in source_file.syntax().preorder() {
        match event {
            WalkEvent::Enter(node) => {
                if let Some(mut symbol) = to_file_symbol(node, file_id) {
                    symbol.container_name = stack.last().cloned();

                    stack.push(symbol.name.clone());
                    symbols.push(symbol);
                }
            }

            WalkEvent::Leave(node) => {
                if to_symbol(node).is_some() {
                    stack.pop();
                }
            }
        }
    }

    symbols
}

fn to_symbol(node: &SyntaxNode) -> Option<(SmolStr, SyntaxNodePtr, TextRange)> {
    fn decl<N: NameOwner>(node: &N) -> Option<(SmolStr, SyntaxNodePtr, TextRange)> {
        let name = node.name()?;
        let name_range = name.syntax().range();
        let name = name.text().clone();
        let ptr = SyntaxNodePtr::new(node.syntax());

        Some((name, ptr, name_range))
    }
    visitor()
        .visit(decl::<ast::FnDef>)
        .visit(decl::<ast::StructDef>)
        .visit(decl::<ast::EnumDef>)
        .visit(decl::<ast::TraitDef>)
        .visit(decl::<ast::Module>)
        .visit(decl::<ast::TypeAliasDef>)
        .visit(decl::<ast::ConstDef>)
        .visit(decl::<ast::StaticDef>)
        .accept(node)?
}

fn to_file_symbol(node: &SyntaxNode, file_id: FileId) -> Option<FileSymbol> {
    to_symbol(node).map(move |(name, ptr, name_range)| FileSymbol {
        name,
        ptr,
        file_id,
        name_range: Some(name_range),
        container_name: None,
    })
}

#[cfg(test)]
mod tests {
    use ra_syntax::SmolStr;
    use crate::{
        display::NavigationTarget,
        mock_analysis::single_file,
        Query,
};

    #[test]
    fn test_world_symbols_with_no_container() {
        let code = r#"
    enum FooInner { }
    "#;

        let mut symbols = get_symbols_matching(code, "FooInner");

        let s = symbols.pop().unwrap();

        assert_eq!(s.name(), "FooInner");
        assert!(s.container_name().is_none());
    }

    #[test]
    fn test_world_symbols_include_container_name() {
        let code = r#"
fn foo() {
    enum FooInner { }
}
    "#;

        let mut symbols = get_symbols_matching(code, "FooInner");

        let s = symbols.pop().unwrap();

        assert_eq!(s.name(), "FooInner");
        assert_eq!(s.container_name(), Some(&SmolStr::new("foo")));

        let code = r#"
mod foo {
    struct FooInner;
}
    "#;

        let mut symbols = get_symbols_matching(code, "FooInner");

        let s = symbols.pop().unwrap();

        assert_eq!(s.name(), "FooInner");
        assert_eq!(s.container_name(), Some(&SmolStr::new("foo")));
    }

    fn get_symbols_matching(text: &str, query: &str) -> Vec<NavigationTarget> {
        let (analysis, _) = single_file(text);
        analysis.symbol_search(Query::new(query.into())).unwrap()
    }
}
<commit_msg>Allows searching for case-equivalent symbols (fixes #1151)<commit_after>//! This module handles fuzzy-searching of functions, structs and other symbols
//! by name across the whole workspace and dependencies.
//!
//! It works by building an incrementally-updated text-search index of all
//! symbols. The backbone of the index is the **awesome** `fst` crate by
//! @BurntSushi.
//!
//! In a nutshell, you give a set of strings to `fst`, and it builds a
//! finite state machine describing this set of strings. The strings which
//! could fuzzy-match a pattern can also be described by a finite state machine.
//! What is freaking cool is that you can now traverse both state machines in
//! lock-step to enumerate the strings which are both in the input set and
//! fuzz-match the query. Or, more formally, given two languages described by
//! FSTs, one can build a product FST which describes the intersection of the
//! languages.
//!
//! `fst` does not support cheap updating of the index, but it supports unioning
//! of state machines. So, to account for changing source code, we build an FST
//! for each library (which is assumed to never change) and an FST for each Rust
//! file in the current workspace, and run a query against the union of all
//! those FSTs.
use std::{
    hash::{Hash, Hasher},
    sync::Arc,
    mem,
    fmt,
};

use fst::{self, Streamer};
use ra_syntax::{
    SyntaxNode, SyntaxNodePtr, SourceFile, SmolStr, TreeArc, AstNode,
    algo::{visit::{visitor, Visitor}},
    SyntaxKind::{self, *},
    ast::{self, NameOwner},
    WalkEvent,
    TextRange,
};
use ra_db::{
    SourceRootId, SourceDatabase,
    salsa::{self, ParallelDatabase},
};
use rayon::prelude::*;

use crate::{
    FileId, Query,
    db::RootDatabase,
};

#[salsa::query_group(SymbolsDatabaseStorage)]
pub(crate) trait SymbolsDatabase: hir::db::HirDatabase {
    fn file_symbols(&self, file_id: FileId) -> Arc<SymbolIndex>;
    #[salsa::input]
    fn library_symbols(&self, id: SourceRootId) -> Arc<SymbolIndex>;
    /// The set of "local" (that is, from the current workspace) roots.
    /// Files in local roots are assumed to change frequently.
    #[salsa::input]
    fn local_roots(&self) -> Arc<Vec<SourceRootId>>;
    /// The set of roots for crates.io libraries.
    /// Files in libraries are assumed to never change.
    #[salsa::input]
    fn library_roots(&self) -> Arc<Vec<SourceRootId>>;
}

fn file_symbols(db: &impl SymbolsDatabase, file_id: FileId) -> Arc<SymbolIndex> {
    db.check_canceled();
    let source_file = db.parse(file_id);

    let symbols = source_file_to_file_symbols(&source_file, file_id);

    // FIXME: add macros here

    Arc::new(SymbolIndex::new(symbols))
}

pub(crate) fn world_symbols(db: &RootDatabase, query: Query) -> Vec<FileSymbol> {
    /// Need to wrap Snapshot to provide `Clone` impl for `map_with`
    struct Snap(salsa::Snapshot<RootDatabase>);
    impl Clone for Snap {
        fn clone(&self) -> Snap {
            Snap(self.0.snapshot())
        }
    }

    let buf: Vec<Arc<SymbolIndex>> = if query.libs {
        let snap = Snap(db.snapshot());
        db.library_roots()
            .par_iter()
            .map_with(snap, |db, &lib_id| db.0.library_symbols(lib_id))
            .collect()
    } else {
        let mut files = Vec::new();
        for &root in db.local_roots().iter() {
            let sr = db.source_root(root);
            files.extend(sr.files.values().map(|&it| it))
        }

        let snap = Snap(db.snapshot());
        files.par_iter().map_with(snap, |db, &file_id| db.0.file_symbols(file_id)).collect()
    };
    query.search(&buf)
}

pub(crate) fn index_resolve(db: &RootDatabase, name_ref: &ast::NameRef) -> Vec<FileSymbol> {
    let name = name_ref.text();
    let mut query = Query::new(name.to_string());
    query.exact();
    query.limit(4);
    crate::symbol_index::world_symbols(db, query)
}

#[derive(Default)]
pub(crate) struct SymbolIndex {
    symbols: Vec<FileSymbol>,
    map: fst::Map,
}

impl fmt::Debug for SymbolIndex {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("SymbolIndex").field("n_symbols", &self.symbols.len()).finish()
    }
}

impl PartialEq for SymbolIndex {
    fn eq(&self, other: &SymbolIndex) -> bool {
        self.symbols == other.symbols
    }
}

impl Eq for SymbolIndex {}

impl Hash for SymbolIndex {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.symbols.hash(hasher)
    }
}

impl SymbolIndex {
    fn new(mut symbols: Vec<FileSymbol>) -> SymbolIndex {
        fn cmp_key<'a>(s1: &'a FileSymbol) -> impl Ord + Eq + 'a {
            unicase::Ascii::new(s1.name.as_str())
        }

        symbols.par_sort_by(|s1, s2| cmp_key(s1).cmp(&cmp_key(s2)));

        let mut builder = fst::MapBuilder::memory();

        let mut last_batch_start = 0;

        for idx in 0..symbols.len() {
            if symbols.get(last_batch_start).map(cmp_key) == symbols.get(idx + 1).map(cmp_key) {
                continue;
            }

            let start = last_batch_start;
            let end = idx + 1;
            last_batch_start = end;

            let key = symbols[start].name.as_str().to_lowercase();
            let value = SymbolIndex::range_to_map_value(start, end);

            builder.insert(key, value).unwrap();
        }

        let map = fst::Map::from_bytes(builder.into_inner().unwrap()).unwrap();
        SymbolIndex { symbols, map }
    }

    pub(crate) fn len(&self) -> usize {
        self.symbols.len()
    }

    pub(crate) fn memory_size(&self) -> usize {
        self.map.as_fst().size() + self.symbols.len() * mem::size_of::<FileSymbol>()
    }

    pub(crate) fn for_files(
        files: impl ParallelIterator<Item = (FileId, TreeArc<SourceFile>)>,
    ) -> SymbolIndex {
        let symbols = files
            .flat_map(|(file_id, file)| source_file_to_file_symbols(&file, file_id))
            .collect::<Vec<_>>();
        SymbolIndex::new(symbols)
    }

    fn range_to_map_value(start: usize, end: usize) -> u64 {
        debug_assert![start <= (std::u32::MAX as usize)];
        debug_assert![end <= (std::u32::MAX as usize)];

        ((start as u64) << 32) + end as u64
    }

    fn map_value_to_range(value: u64) -> (usize, usize) {
        let end = value as u32 as usize;
        let start = (value >> 32) as usize;
        (start, end)
    }
}

impl Query {
    pub(crate) fn search(self, indices: &[Arc<SymbolIndex>]) -> Vec<FileSymbol> {
        let mut op = fst::map::OpBuilder::new();
        for file_symbols in indices.iter() {
            let automaton = fst::automaton::Subsequence::new(&self.lowercased);
            op = op.add(file_symbols.map.search(automaton))
        }
        let mut stream = op.union();
        let mut res = Vec::new();
        while let Some((_, indexed_values)) = stream.next() {
            if res.len() >= self.limit {
                break;
            }
            for indexed_value in indexed_values {
                let symbol_index = &indices[indexed_value.index];
                let (start, end) = SymbolIndex::map_value_to_range(indexed_value.value);

                for symbol in &symbol_index.symbols[start..end] {
                    if self.only_types && !is_type(symbol.ptr.kind()) {
                        continue;
                    }
                    if self.exact && symbol.name != self.query {
                        continue;
                    }
                    res.push(symbol.clone());
                }
            }
        }
        res
    }
}

fn is_type(kind: SyntaxKind) -> bool {
    match kind {
        STRUCT_DEF | ENUM_DEF | TRAIT_DEF | TYPE_ALIAS_DEF => true,
        _ => false,
    }
}

/// The actual data that is stored in the index. It should be as compact as
/// possible.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub(crate) struct FileSymbol {
    pub(crate) file_id: FileId,
    pub(crate) name: SmolStr,
    pub(crate) ptr: SyntaxNodePtr,
    pub(crate) name_range: Option<TextRange>,
    pub(crate) container_name: Option<SmolStr>,
}

fn source_file_to_file_symbols(source_file: &SourceFile, file_id: FileId) -> Vec<FileSymbol> {
    let mut symbols = Vec::new();
    let mut stack = Vec::new();

    for event in source_file.syntax().preorder() {
        match event {
            WalkEvent::Enter(node) => {
                if let Some(mut symbol) = to_file_symbol(node, file_id) {
                    symbol.container_name = stack.last().cloned();

                    stack.push(symbol.name.clone());
                    symbols.push(symbol);
                }
            }

            WalkEvent::Leave(node) => {
                if to_symbol(node).is_some() {
                    stack.pop();
                }
            }
        }
    }

    symbols
}

fn to_symbol(node: &SyntaxNode) -> Option<(SmolStr, SyntaxNodePtr, TextRange)> {
    fn decl<N: NameOwner>(node: &N) -> Option<(SmolStr, SyntaxNodePtr, TextRange)> {
        let name = node.name()?;
        let name_range = name.syntax().range();
        let name = name.text().clone();
        let ptr = SyntaxNodePtr::new(node.syntax());

        Some((name, ptr, name_range))
    }
    visitor()
        .visit(decl::<ast::FnDef>)
        .visit(decl::<ast::StructDef>)
        .visit(decl::<ast::EnumDef>)
        .visit(decl::<ast::TraitDef>)
        .visit(decl::<ast::Module>)
        .visit(decl::<ast::TypeAliasDef>)
        .visit(decl::<ast::ConstDef>)
        .visit(decl::<ast::StaticDef>)
        .accept(node)?
}

fn to_file_symbol(node: &SyntaxNode, file_id: FileId) -> Option<FileSymbol> {
    to_symbol(node).map(move |(name, ptr, name_range)| FileSymbol {
        name,
        ptr,
        file_id,
        name_range: Some(name_range),
        container_name: None,
    })
}

#[cfg(test)]
mod tests {
    use ra_syntax::{
        SmolStr,
        SyntaxKind::{FN_DEF, STRUCT_DEF}
};
    use crate::{
        display::NavigationTarget,
        mock_analysis::single_file,
        Query,
};

    #[test]
    fn test_world_symbols_with_no_container() {
        let code = r#"
    enum FooInner { }
    "#;

        let mut symbols = get_symbols_matching(code, "FooInner");

        let s = symbols.pop().unwrap();

        assert_eq!(s.name(), "FooInner");
        assert!(s.container_name().is_none());
    }

    #[test]
    fn test_world_symbols_include_container_name() {
        let code = r#"
fn foo() {
    enum FooInner { }
}
    "#;

        let mut symbols = get_symbols_matching(code, "FooInner");

        let s = symbols.pop().unwrap();

        assert_eq!(s.name(), "FooInner");
        assert_eq!(s.container_name(), Some(&SmolStr::new("foo")));

        let code = r#"
mod foo {
    struct FooInner;
}
    "#;

        let mut symbols = get_symbols_matching(code, "FooInner");

        let s = symbols.pop().unwrap();

        assert_eq!(s.name(), "FooInner");
        assert_eq!(s.container_name(), Some(&SmolStr::new("foo")));
    }

    #[test]
    fn test_world_symbols_are_case_sensitive() {
        let code = r#"
fn foo() {}

struct Foo;
        "#;

        let symbols = get_symbols_matching(code, "Foo");

        let fn_match = symbols.iter().find(|s| s.name() == "foo").map(|s| s.kind());
        let struct_match = symbols.iter().find(|s| s.name() == "Foo").map(|s| s.kind());

        assert_eq!(fn_match, Some(FN_DEF));
        assert_eq!(struct_match, Some(STRUCT_DEF));
    }

    fn get_symbols_matching(text: &str, query: &str) -> Vec<NavigationTarget> {
        let (analysis, _) = single_file(text);
        analysis.symbol_search(Query::new(query.into())).unwrap()
    }
}
<|endoftext|>


<commit_before>//! This module handles fuzzy-searching of functions, structs and other symbols
//! by name across the whole workspace and dependencies.
//!
//! It works by building an incrementally-updated text-search index of all
//! symbols. The backbone of the index is the **awesome** `fst` crate by
//! @BurntSushi.
//!
//! In a nutshell, you give a set of strings to `fst`, and it builds a
//! finite state machine describing this set of strings. The strings which
//! could fuzzy-match a pattern can also be described by a finite state machine.
//! What is freaking cool is that you can now traverse both state machines in
//! lock-step to enumerate the strings which are both in the input set and
//! fuzz-match the query. Or, more formally, given two languages described by
//! FSTs, one can build a product FST which describes the intersection of the
//! languages.
//!
//! `fst` does not support cheap updating of the index, but it supports unioning
//! of state machines. So, to account for changing source code, we build an FST
//! for each library (which is assumed to never change) and an FST for each Rust
//! file in the current workspace, and run a query against the union of all
//! those FSTs.
use std::{
    hash::{Hash, Hasher},
    sync::Arc,
    mem,
    fmt,
};

use fst::{self, Streamer};
use ra_syntax::{
    SyntaxNode, SyntaxNodePtr, SourceFile, SmolStr, TreeArc, AstNode,
    algo::{visit::{visitor, Visitor}},
    SyntaxKind::{self, *},
    ast::{self, NameOwner},
    WalkEvent,
    TextRange,
};
use ra_db::{
    SourceRootId, SourceDatabase,
    salsa::{self, ParallelDatabase},
};
use rayon::prelude::*;

use crate::{
    FileId, Query,
    db::RootDatabase,
};

#[salsa::query_group(SymbolsDatabaseStorage)]
pub(crate) trait SymbolsDatabase: hir::db::HirDatabase {
    fn file_symbols(&self, file_id: FileId) -> Arc<SymbolIndex>;
    #[salsa::input]
    fn library_symbols(&self, id: SourceRootId) -> Arc<SymbolIndex>;
    /// The set of "local" (that is, from the current workspace) roots.
    /// Files in local roots are assumed to change frequently.
    #[salsa::input]
    fn local_roots(&self) -> Arc<Vec<SourceRootId>>;
    /// The set of roots for crates.io libraries.
    /// Files in libraries are assumed to never change.
    #[salsa::input]
    fn library_roots(&self) -> Arc<Vec<SourceRootId>>;
}

fn file_symbols(db: &impl SymbolsDatabase, file_id: FileId) -> Arc<SymbolIndex> {
    db.check_canceled();
    let source_file = db.parse(file_id);

    let symbols = source_file_to_file_symbols(&source_file, file_id);

    // FIXME: add macros here

    Arc::new(SymbolIndex::new(symbols))
}

pub(crate) fn world_symbols(db: &RootDatabase, query: Query) -> Vec<FileSymbol> {
    /// Need to wrap Snapshot to provide `Clone` impl for `map_with`
    struct Snap(salsa::Snapshot<RootDatabase>);
    impl Clone for Snap {
        fn clone(&self) -> Snap {
            Snap(self.0.snapshot())
        }
    }

    let buf: Vec<Arc<SymbolIndex>> = if query.libs {
        let snap = Snap(db.snapshot());
        db.library_roots()
            .par_iter()
            .map_with(snap, |db, &lib_id| db.0.library_symbols(lib_id))
            .collect()
    } else {
        let mut files = Vec::new();
        for &root in db.local_roots().iter() {
            let sr = db.source_root(root);
            files.extend(sr.files.values().map(|&it| it))
        }

        let snap = Snap(db.snapshot());
        files.par_iter().map_with(snap, |db, &file_id| db.0.file_symbols(file_id)).collect()
    };
    query.search(&buf)
}

pub(crate) fn index_resolve(db: &RootDatabase, name_ref: &ast::NameRef) -> Vec<FileSymbol> {
    let name = name_ref.text();
    let mut query = Query::new(name.to_string());
    query.exact();
    query.limit(4);
    crate::symbol_index::world_symbols(db, query)
}

#[derive(Default)]
pub(crate) struct SymbolIndex {
    symbols: Vec<FileSymbol>,
    map: fst::Map,
}

impl fmt::Debug for SymbolIndex {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("SymbolIndex").field("n_symbols", &self.symbols.len()).finish()
    }
}

impl PartialEq for SymbolIndex {
    fn eq(&self, other: &SymbolIndex) -> bool {
        self.symbols == other.symbols
    }
}

impl Eq for SymbolIndex {}

impl Hash for SymbolIndex {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.symbols.hash(hasher)
    }
}

impl SymbolIndex {
    fn new(mut symbols: Vec<FileSymbol>) -> SymbolIndex {
        fn cmp_key<'a>(s1: &'a FileSymbol) -> impl Ord + Eq + 'a {
            unicase::Ascii::new(s1.name.as_str())
        }

        symbols.par_sort_by(|s1, s2| cmp_key(s1).cmp(&cmp_key(s2)));

        let mut builder = fst::MapBuilder::memory();

        let mut last_batch_start = 0;

        for idx in 0..symbols.len() {
            if symbols.get(last_batch_start).map(cmp_key) == symbols.get(idx + 1).map(cmp_key) {
                continue;
            }

            let start = last_batch_start;
            let end = idx + 1;
            last_batch_start = end;

            let key = symbols[start].name.as_str().to_lowercase();
            let value = SymbolIndex::range_to_map_value(start, end);

            builder.insert(key, value).unwrap();
        }

        let map = fst::Map::from_bytes(builder.into_inner().unwrap()).unwrap();
        SymbolIndex { symbols, map }
    }

    pub(crate) fn len(&self) -> usize {
        self.symbols.len()
    }

    pub(crate) fn memory_size(&self) -> usize {
        self.map.as_fst().size() + self.symbols.len() * mem::size_of::<FileSymbol>()
    }

    pub(crate) fn for_files(
        files: impl ParallelIterator<Item = (FileId, TreeArc<SourceFile>)>,
    ) -> SymbolIndex {
        let symbols = files
            .flat_map(|(file_id, file)| source_file_to_file_symbols(&file, file_id))
            .collect::<Vec<_>>();
        SymbolIndex::new(symbols)
    }

    fn range_to_map_value(start: usize, end: usize) -> u64 {
        debug_assert![start <= (std::u32::MAX as usize)];
        debug_assert![end <= (std::u32::MAX as usize)];

        ((start as u64) << 32) + end as u64
    }

    fn map_value_to_range(value: u64) -> (usize, usize) {
        let end = value as u32 as usize;
        let start = (value >> 32) as usize;
        (start, end)
    }
}

impl Query {
    pub(crate) fn search(self, indices: &[Arc<SymbolIndex>]) -> Vec<FileSymbol> {
        let mut op = fst::map::OpBuilder::new();
        for file_symbols in indices.iter() {
            let automaton = fst::automaton::Subsequence::new(&self.lowercased);
            op = op.add(file_symbols.map.search(automaton))
        }
        let mut stream = op.union();
        let mut res = Vec::new();
        while let Some((_, indexed_values)) = stream.next() {
            if res.len() >= self.limit {
                break;
            }
            for indexed_value in indexed_values {
                let symbol_index = &indices[indexed_value.index];
                let (start, end) = SymbolIndex::map_value_to_range(indexed_value.value);

                for symbol in &symbol_index.symbols[start..end] {
                    if self.only_types && !is_type(symbol.ptr.kind()) {
                        continue;
                    }
                    if self.exact && symbol.name != self.query {
                        continue;
                    }
                    res.push(symbol.clone());
                }
            }
        }
        res
    }
}

fn is_type(kind: SyntaxKind) -> bool {
    match kind {
        STRUCT_DEF | ENUM_DEF | TRAIT_DEF | TYPE_ALIAS_DEF => true,
        _ => false,
    }
}

/// The actual data that is stored in the index. It should be as compact as
/// possible.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub(crate) struct FileSymbol {
    pub(crate) file_id: FileId,
    pub(crate) name: SmolStr,
    pub(crate) ptr: SyntaxNodePtr,
    pub(crate) name_range: Option<TextRange>,
    pub(crate) container_name: Option<SmolStr>,
}

fn source_file_to_file_symbols(source_file: &SourceFile, file_id: FileId) -> Vec<FileSymbol> {
    let mut symbols = Vec::new();
    let mut stack = Vec::new();

    for event in source_file.syntax().preorder() {
        match event {
            WalkEvent::Enter(node) => {
                if let Some(mut symbol) = to_file_symbol(node, file_id) {
                    symbol.container_name = stack.last().cloned();

                    stack.push(symbol.name.clone());
                    symbols.push(symbol);
                }
            }

            WalkEvent::Leave(node) => {
                if to_symbol(node).is_some() {
                    stack.pop();
                }
            }
        }
    }

    symbols
}

fn to_symbol(node: &SyntaxNode) -> Option<(SmolStr, SyntaxNodePtr, TextRange)> {
    fn decl<N: NameOwner>(node: &N) -> Option<(SmolStr, SyntaxNodePtr, TextRange)> {
        let name = node.name()?;
        let name_range = name.syntax().range();
        let name = name.text().clone();
        let ptr = SyntaxNodePtr::new(node.syntax());

        Some((name, ptr, name_range))
    }
    visitor()
        .visit(decl::<ast::FnDef>)
        .visit(decl::<ast::StructDef>)
        .visit(decl::<ast::EnumDef>)
        .visit(decl::<ast::TraitDef>)
        .visit(decl::<ast::Module>)
        .visit(decl::<ast::TypeAliasDef>)
        .visit(decl::<ast::ConstDef>)
        .visit(decl::<ast::StaticDef>)
        .accept(node)?
}

fn to_file_symbol(node: &SyntaxNode, file_id: FileId) -> Option<FileSymbol> {
    to_symbol(node).map(move |(name, ptr, name_range)| FileSymbol {
        name,
        ptr,
        file_id,
        name_range: Some(name_range),
        container_name: None,
    })
}

#[cfg(test)]
mod tests {
    use ra_syntax::{
        SmolStr,
        SyntaxKind::{FN_DEF, STRUCT_DEF}
};
    use crate::{
        display::NavigationTarget,
        mock_analysis::single_file,
        Query,
};

    #[test]
    fn test_world_symbols_with_no_container() {
        let code = r#"
    enum FooInner { }
    "#;

        let mut symbols = get_symbols_matching(code, "FooInner");

        let s = symbols.pop().unwrap();

        assert_eq!(s.name(), "FooInner");
        assert!(s.container_name().is_none());
    }

    #[test]
    fn test_world_symbols_include_container_name() {
        let code = r#"
fn foo() {
    enum FooInner { }
}
    "#;

        let mut symbols = get_symbols_matching(code, "FooInner");

        let s = symbols.pop().unwrap();

        assert_eq!(s.name(), "FooInner");
        assert_eq!(s.container_name(), Some(&SmolStr::new("foo")));

        let code = r#"
mod foo {
    struct FooInner;
}
    "#;

        let mut symbols = get_symbols_matching(code, "FooInner");

        let s = symbols.pop().unwrap();

        assert_eq!(s.name(), "FooInner");
        assert_eq!(s.container_name(), Some(&SmolStr::new("foo")));
    }

    #[test]
    fn test_world_symbols_are_case_sensitive() {
        let code = r#"
fn foo() {}

struct Foo;
        "#;

        let symbols = get_symbols_matching(code, "Foo");

        let fn_match = symbols.iter().find(|s| s.name() == "foo").map(|s| s.kind());
        let struct_match = symbols.iter().find(|s| s.name() == "Foo").map(|s| s.kind());

        assert_eq!(fn_match, Some(FN_DEF));
        assert_eq!(struct_match, Some(STRUCT_DEF));
    }

    fn get_symbols_matching(text: &str, query: &str) -> Vec<NavigationTarget> {
        let (analysis, _) = single_file(text);
        analysis.symbol_search(Query::new(query.into())).unwrap()
    }
}
<commit_msg>CR corrections<commit_after>//! This module handles fuzzy-searching of functions, structs and other symbols
//! by name across the whole workspace and dependencies.
//!
//! It works by building an incrementally-updated text-search index of all
//! symbols. The backbone of the index is the **awesome** `fst` crate by
//! @BurntSushi.
//!
//! In a nutshell, you give a set of strings to `fst`, and it builds a
//! finite state machine describing this set of strings. The strings which
//! could fuzzy-match a pattern can also be described by a finite state machine.
//! What is freaking cool is that you can now traverse both state machines in
//! lock-step to enumerate the strings which are both in the input set and
//! fuzz-match the query. Or, more formally, given two languages described by
//! FSTs, one can build a product FST which describes the intersection of the
//! languages.
//!
//! `fst` does not support cheap updating of the index, but it supports unioning
//! of state machines. So, to account for changing source code, we build an FST
//! for each library (which is assumed to never change) and an FST for each Rust
//! file in the current workspace, and run a query against the union of all
//! those FSTs.
use std::{
    hash::{Hash, Hasher},
    sync::Arc,
    mem,
    fmt,
};

use fst::{self, Streamer};
use ra_syntax::{
    SyntaxNode, SyntaxNodePtr, SourceFile, SmolStr, TreeArc, AstNode,
    algo::{visit::{visitor, Visitor}},
    SyntaxKind::{self, *},
    ast::{self, NameOwner},
    WalkEvent,
    TextRange,
};
use ra_db::{
    SourceRootId, SourceDatabase,
    salsa::{self, ParallelDatabase},
};
use rayon::prelude::*;

use crate::{
    FileId, Query,
    db::RootDatabase,
};

#[salsa::query_group(SymbolsDatabaseStorage)]
pub(crate) trait SymbolsDatabase: hir::db::HirDatabase {
    fn file_symbols(&self, file_id: FileId) -> Arc<SymbolIndex>;
    #[salsa::input]
    fn library_symbols(&self, id: SourceRootId) -> Arc<SymbolIndex>;
    /// The set of "local" (that is, from the current workspace) roots.
    /// Files in local roots are assumed to change frequently.
    #[salsa::input]
    fn local_roots(&self) -> Arc<Vec<SourceRootId>>;
    /// The set of roots for crates.io libraries.
    /// Files in libraries are assumed to never change.
    #[salsa::input]
    fn library_roots(&self) -> Arc<Vec<SourceRootId>>;
}

fn file_symbols(db: &impl SymbolsDatabase, file_id: FileId) -> Arc<SymbolIndex> {
    db.check_canceled();
    let source_file = db.parse(file_id);

    let symbols = source_file_to_file_symbols(&source_file, file_id);

    // FIXME: add macros here

    Arc::new(SymbolIndex::new(symbols))
}

pub(crate) fn world_symbols(db: &RootDatabase, query: Query) -> Vec<FileSymbol> {
    /// Need to wrap Snapshot to provide `Clone` impl for `map_with`
    struct Snap(salsa::Snapshot<RootDatabase>);
    impl Clone for Snap {
        fn clone(&self) -> Snap {
            Snap(self.0.snapshot())
        }
    }

    let buf: Vec<Arc<SymbolIndex>> = if query.libs {
        let snap = Snap(db.snapshot());
        db.library_roots()
            .par_iter()
            .map_with(snap, |db, &lib_id| db.0.library_symbols(lib_id))
            .collect()
    } else {
        let mut files = Vec::new();
        for &root in db.local_roots().iter() {
            let sr = db.source_root(root);
            files.extend(sr.files.values().map(|&it| it))
        }

        let snap = Snap(db.snapshot());
        files.par_iter().map_with(snap, |db, &file_id| db.0.file_symbols(file_id)).collect()
    };
    query.search(&buf)
}

pub(crate) fn index_resolve(db: &RootDatabase, name_ref: &ast::NameRef) -> Vec<FileSymbol> {
    let name = name_ref.text();
    let mut query = Query::new(name.to_string());
    query.exact();
    query.limit(4);
    crate::symbol_index::world_symbols(db, query)
}

#[derive(Default)]
pub(crate) struct SymbolIndex {
    symbols: Vec<FileSymbol>,
    map: fst::Map,
}

impl fmt::Debug for SymbolIndex {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("SymbolIndex").field("n_symbols", &self.symbols.len()).finish()
    }
}

impl PartialEq for SymbolIndex {
    fn eq(&self, other: &SymbolIndex) -> bool {
        self.symbols == other.symbols
    }
}

impl Eq for SymbolIndex {}

impl Hash for SymbolIndex {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.symbols.hash(hasher)
    }
}

impl SymbolIndex {
    fn new(mut symbols: Vec<FileSymbol>) -> SymbolIndex {
        fn cmp_key<'a>(s1: &'a FileSymbol) -> impl Ord + 'a {
            unicase::Ascii::new(s1.name.as_str())
        }

        symbols.par_sort_by(|s1, s2| cmp_key(s1).cmp(&cmp_key(s2)));

        let mut builder = fst::MapBuilder::memory();

        let mut last_batch_start = 0;

        for idx in 0..symbols.len() {
            if symbols.get(last_batch_start).map(cmp_key) == symbols.get(idx + 1).map(cmp_key) {
                continue;
            }

            let start = last_batch_start;
            let end = idx + 1;
            last_batch_start = end;

            let key = symbols[start].name.as_str().to_lowercase();
            let value = SymbolIndex::range_to_map_value(start, end);

            builder.insert(key, value).unwrap();
        }

        let map = fst::Map::from_bytes(builder.into_inner().unwrap()).unwrap();
        SymbolIndex { symbols, map }
    }

    pub(crate) fn len(&self) -> usize {
        self.symbols.len()
    }

    pub(crate) fn memory_size(&self) -> usize {
        self.map.as_fst().size() + self.symbols.len() * mem::size_of::<FileSymbol>()
    }

    pub(crate) fn for_files(
        files: impl ParallelIterator<Item = (FileId, TreeArc<SourceFile>)>,
    ) -> SymbolIndex {
        let symbols = files
            .flat_map(|(file_id, file)| source_file_to_file_symbols(&file, file_id))
            .collect::<Vec<_>>();
        SymbolIndex::new(symbols)
    }

    fn range_to_map_value(start: usize, end: usize) -> u64 {
        debug_assert![start <= (std::u32::MAX as usize)];
        debug_assert![end <= (std::u32::MAX as usize)];

        ((start as u64) << 32) | end as u64
    }

    fn map_value_to_range(value: u64) -> (usize, usize) {
        let end = value as u32 as usize;
        let start = (value >> 32) as usize;
        (start, end)
    }
}

impl Query {
    pub(crate) fn search(self, indices: &[Arc<SymbolIndex>]) -> Vec<FileSymbol> {
        let mut op = fst::map::OpBuilder::new();
        for file_symbols in indices.iter() {
            let automaton = fst::automaton::Subsequence::new(&self.lowercased);
            op = op.add(file_symbols.map.search(automaton))
        }
        let mut stream = op.union();
        let mut res = Vec::new();
        while let Some((_, indexed_values)) = stream.next() {
            if res.len() >= self.limit {
                break;
            }
            for indexed_value in indexed_values {
                let symbol_index = &indices[indexed_value.index];
                let (start, end) = SymbolIndex::map_value_to_range(indexed_value.value);

                for symbol in &symbol_index.symbols[start..end] {
                    if self.only_types && !is_type(symbol.ptr.kind()) {
                        continue;
                    }
                    if self.exact && symbol.name != self.query {
                        continue;
                    }
                    res.push(symbol.clone());
                }
            }
        }
        res
    }
}

fn is_type(kind: SyntaxKind) -> bool {
    match kind {
        STRUCT_DEF | ENUM_DEF | TRAIT_DEF | TYPE_ALIAS_DEF => true,
        _ => false,
    }
}

/// The actual data that is stored in the index. It should be as compact as
/// possible.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub(crate) struct FileSymbol {
    pub(crate) file_id: FileId,
    pub(crate) name: SmolStr,
    pub(crate) ptr: SyntaxNodePtr,
    pub(crate) name_range: Option<TextRange>,
    pub(crate) container_name: Option<SmolStr>,
}

fn source_file_to_file_symbols(source_file: &SourceFile, file_id: FileId) -> Vec<FileSymbol> {
    let mut symbols = Vec::new();
    let mut stack = Vec::new();

    for event in source_file.syntax().preorder() {
        match event {
            WalkEvent::Enter(node) => {
                if let Some(mut symbol) = to_file_symbol(node, file_id) {
                    symbol.container_name = stack.last().cloned();

                    stack.push(symbol.name.clone());
                    symbols.push(symbol);
                }
            }

            WalkEvent::Leave(node) => {
                if to_symbol(node).is_some() {
                    stack.pop();
                }
            }
        }
    }

    symbols
}

fn to_symbol(node: &SyntaxNode) -> Option<(SmolStr, SyntaxNodePtr, TextRange)> {
    fn decl<N: NameOwner>(node: &N) -> Option<(SmolStr, SyntaxNodePtr, TextRange)> {
        let name = node.name()?;
        let name_range = name.syntax().range();
        let name = name.text().clone();
        let ptr = SyntaxNodePtr::new(node.syntax());

        Some((name, ptr, name_range))
    }
    visitor()
        .visit(decl::<ast::FnDef>)
        .visit(decl::<ast::StructDef>)
        .visit(decl::<ast::EnumDef>)
        .visit(decl::<ast::TraitDef>)
        .visit(decl::<ast::Module>)
        .visit(decl::<ast::TypeAliasDef>)
        .visit(decl::<ast::ConstDef>)
        .visit(decl::<ast::StaticDef>)
        .accept(node)?
}

fn to_file_symbol(node: &SyntaxNode, file_id: FileId) -> Option<FileSymbol> {
    to_symbol(node).map(move |(name, ptr, name_range)| FileSymbol {
        name,
        ptr,
        file_id,
        name_range: Some(name_range),
        container_name: None,
    })
}

#[cfg(test)]
mod tests {
    use ra_syntax::{
        SmolStr,
        SyntaxKind::{FN_DEF, STRUCT_DEF}
};
    use crate::{
        display::NavigationTarget,
        mock_analysis::single_file,
        Query,
};

    #[test]
    fn test_world_symbols_with_no_container() {
        let code = r#"
    enum FooInner { }
    "#;

        let mut symbols = get_symbols_matching(code, "FooInner");

        let s = symbols.pop().unwrap();

        assert_eq!(s.name(), "FooInner");
        assert!(s.container_name().is_none());
    }

    #[test]
    fn test_world_symbols_include_container_name() {
        let code = r#"
fn foo() {
    enum FooInner { }
}
    "#;

        let mut symbols = get_symbols_matching(code, "FooInner");

        let s = symbols.pop().unwrap();

        assert_eq!(s.name(), "FooInner");
        assert_eq!(s.container_name(), Some(&SmolStr::new("foo")));

        let code = r#"
mod foo {
    struct FooInner;
}
    "#;

        let mut symbols = get_symbols_matching(code, "FooInner");

        let s = symbols.pop().unwrap();

        assert_eq!(s.name(), "FooInner");
        assert_eq!(s.container_name(), Some(&SmolStr::new("foo")));
    }

    #[test]
    fn test_world_symbols_are_case_sensitive() {
        let code = r#"
fn foo() {}

struct Foo;
        "#;

        let symbols = get_symbols_matching(code, "Foo");

        let fn_match = symbols.iter().find(|s| s.name() == "foo").map(|s| s.kind());
        let struct_match = symbols.iter().find(|s| s.name() == "Foo").map(|s| s.kind());

        assert_eq!(fn_match, Some(FN_DEF));
        assert_eq!(struct_match, Some(STRUCT_DEF));
    }

    fn get_symbols_matching(text: &str, query: &str) -> Vec<NavigationTarget> {
        let (analysis, _) = single_file(text);
        analysis.symbol_search(Query::new(query.into())).unwrap()
    }
}
<|endoftext|>


<commit_before>//! Handles dynamic library loading for proc macro

use crate::{proc_macro::bridge, rustc_server::TokenStream};
use std::fs::File;
use std::io::Read;
use std::path::Path;

use goblin::{mach::Mach, Object};
use libloading::Library;
use ra_proc_macro::ProcMacroKind;

static NEW_REGISTRAR_SYMBOL: &str = "__rustc_proc_macro_decls_";
static _OLD_REGISTRAR_SYMBOL: &str = "__rustc_derive_registrar_";

fn read_bytes(file: &Path) -> Option<Vec<u8>> {
    let mut fd = File::open(file).ok()?;
    let mut buffer = Vec::new();
    fd.read_to_end(&mut buffer).ok()?;

    Some(buffer)
}

fn get_symbols_from_lib(file: &Path) -> Option<Vec<String>> {
    let buffer = read_bytes(file)?;
    let object = Object::parse(&buffer).ok()?;

    return match object {
        Object::Elf(elf) => {
            let symbols = elf.dynstrtab.to_vec().ok()?;
            let names = symbols.iter().map(|s| s.to_string()).collect();

            Some(names)
        }

        Object::PE(pe) => {
            let symbol_names =
                pe.exports.iter().flat_map(|s| s.name).map(|n| n.to_string()).collect();
            Some(symbol_names)
        }

        Object::Mach(mach) => match mach {
            Mach::Binary(binary) => {
                let exports = binary.exports().ok()?;
                let names = exports.iter().map(|s| s.name.clone()).collect();

                Some(names)
            }

            Mach::Fat(_) => None,
        },

        Object::Archive(_) | Object::Unknown(_) => None,
    };
}

fn is_derive_registrar_symbol(symbol: &str) -> bool {
    symbol.contains(NEW_REGISTRAR_SYMBOL)
}

#[cfg(not(target_os = "macos"))]
fn adjust_symbol_name(name: &str) -> String {
    name.to_string()
}

#[cfg(target_os = "macos")]
fn adjust_symbol_name(s: &str) -> String {
    // In macos doc:
    // https://developer.apple.com/library/archive/documentation/System/Conceptual/ManPages_iPhoneOS/man3/dlsym.3.html
    // Unlike other dyld API's, the symbol name passed to dlsym() must NOT be
    // prepended with an underscore.
    if s.starts_with("_") {
        s[1..s.len()].to_string()
    } else {
        s.to_string()
    }
}

fn find_registrar_symbol(file: &Path) -> Option<String> {
    let symbols = get_symbols_from_lib(file)?;

    symbols.iter().find(|s| is_derive_registrar_symbol(s)).map(|s| adjust_symbol_name(&s))
}

/// Loads dynamic library in platform dependent manner.
///
/// For unix, you have to use RTLD_DEEPBIND flag to escape problems described
/// [here](https://github.com/fedochet/rust-proc-macro-panic-inside-panic-expample)
/// and [here](https://github.com/rust-lang/rust/issues/60593).
///
/// Usage of RTLD_DEEPBIND
/// [here](https://github.com/fedochet/rust-proc-macro-panic-inside-panic-expample/issues/1)
///
/// It seems that on Windows that behaviour is default, so we do nothing in that case.
#[cfg(windows)]
fn load_library(file: &Path) -> Result<Library, libloading::Error> {
    Library::new(file)
}

#[cfg(unix)]
fn load_library(file: &Path) -> Result<Library, libloading::Error> {
    use libloading::os::unix::Library as UnixLibrary;
    use std::os::raw::c_int;

    const RTLD_NOW: c_int = 0x00002;
    const RTLD_DEEPBIND: c_int = 0x00008;

    UnixLibrary::open(Some(file), RTLD_NOW | RTLD_DEEPBIND).map(|lib| lib.into())
}

struct ProcMacroLibraryLibloading {
    // Hold the dylib to prevent it for unloadeding
    _lib: Library,
    exported_macros: Vec<bridge::client::ProcMacro>,
}

impl ProcMacroLibraryLibloading {
    fn open(file: &Path) -> Result<Self, String> {
        let symbol_name = find_registrar_symbol(file)
            .ok_or(format!("Cannot find registrar symbol in file {:?}", file))?;

        let lib = load_library(file).map_err(|e| e.to_string())?;

        let exported_macros = {
            let macros: libloading::Symbol<&&[bridge::client::ProcMacro]> =
                unsafe { lib.get(symbol_name.as_bytes()) }.map_err(|e| e.to_string())?;

            macros.to_vec()
        };

        Ok(ProcMacroLibraryLibloading { _lib: lib, exported_macros })
    }
}

type ProcMacroLibraryImpl = ProcMacroLibraryLibloading;

pub struct Expander {
    libs: Vec<ProcMacroLibraryImpl>,
}

impl Expander {
    pub fn new<P: AsRef<Path>>(lib: &P) -> Result<Expander, String> {
        let mut libs = vec![];

        /* Some libraries for dynamic loading require canonicalized path (even when it is
        already absolute
        */
        let lib =
            lib.as_ref().canonicalize().expect(&format!("Cannot canonicalize {:?}", lib.as_ref()));

        let library = ProcMacroLibraryImpl::open(&lib)?;
        libs.push(library);

        Ok(Expander { libs })
    }

    pub fn expand(
        &self,
        macro_name: &str,
        macro_body: &ra_tt::Subtree,
        attributes: Option<&ra_tt::Subtree>,
    ) -> Result<ra_tt::Subtree, bridge::PanicMessage> {
        let parsed_body = TokenStream::with_subtree(macro_body.clone());

        let parsed_attributes = attributes
            .map_or(crate::rustc_server::TokenStream::new(), |attr| {
                TokenStream::with_subtree(attr.clone())
            });

        for lib in &self.libs {
            for proc_macro in &lib.exported_macros {
                match proc_macro {
                    bridge::client::ProcMacro::CustomDerive { trait_name, client, .. }
                        if *trait_name == macro_name =>
                    {
                        let res = client.run(
                            &crate::proc_macro::bridge::server::SameThread,
                            crate::rustc_server::Rustc::default(),
                            parsed_body,
                        );

                        return res.map(|it| it.subtree);
                    }

                    bridge::client::ProcMacro::Bang { name, client } if *name == macro_name => {
                        let res = client.run(
                            &crate::proc_macro::bridge::server::SameThread,
                            crate::rustc_server::Rustc::default(),
                            parsed_body,
                        );

                        return res.map(|it| it.subtree);
                    }

                    bridge::client::ProcMacro::Attr { name, client } if *name == macro_name => {
                        let res = client.run(
                            &crate::proc_macro::bridge::server::SameThread,
                            crate::rustc_server::Rustc::default(),
                            parsed_attributes,
                            parsed_body,
                        );

                        return res.map(|it| it.subtree);
                    }

                    _ => {
                        continue;
                    }
                }
            }
        }

        Err(bridge::PanicMessage::String("Nothing to expand".to_string()))
    }

    pub fn list_macros(&self) -> Result<Vec<(String, ProcMacroKind)>, bridge::PanicMessage> {
        let mut result = vec![];

        for lib in &self.libs {
            for proc_macro in &lib.exported_macros {
                let res = match proc_macro {
                    bridge::client::ProcMacro::CustomDerive { trait_name, .. } => {
                        (trait_name.to_string(), ProcMacroKind::CustomDerive)
                    }
                    bridge::client::ProcMacro::Bang { name, .. } => {
                        (name.to_string(), ProcMacroKind::FuncLike)
                    }
                    bridge::client::ProcMacro::Attr { name, .. } => {
                        (name.to_string(), ProcMacroKind::Attr)
                    }
                };
                result.push(res);
            }
        }

        Ok(result)
    }
}
<commit_msg>tidy up a bit<commit_after>//! Handles dynamic library loading for proc macro

use crate::{proc_macro::bridge, rustc_server::TokenStream};
use std::path::Path;

use goblin::{mach::Mach, Object};
use libloading::Library;
use ra_proc_macro::ProcMacroKind;

static NEW_REGISTRAR_SYMBOL: &str = "_rustc_proc_macro_decls_";

fn get_symbols_from_lib(file: &Path) -> Option<Vec<String>> {
    let buffer = std::fs::read(file).ok()?;
    let object = Object::parse(&buffer).ok()?;

    return match object {
        Object::Elf(elf) => {
            let symbols = elf.dynstrtab.to_vec().ok()?;
            let names = symbols.iter().map(|s| s.to_string()).collect();
            Some(names)
        }
        Object::PE(pe) => {
            let symbol_names =
                pe.exports.iter().flat_map(|s| s.name).map(|n| n.to_string()).collect();
            Some(symbol_names)
        }
        Object::Mach(mach) => match mach {
            Mach::Binary(binary) => {
                let exports = binary.exports().ok()?;
                let names = exports
                    .iter()
                    .map(|s| {
                        // In macos doc:
                        // https://developer.apple.com/library/archive/documentation/System/Conceptual/ManPages_iPhoneOS/man3/dlsym.3.html
                        // Unlike other dyld API's, the symbol name passed to dlsym() must NOT be
                        // prepended with an underscore.
                        if s.name.starts_with("_") {
                            s.name[1..].to_string()
                        } else {
                            s.name.to_string()
                        }
                    })
                    .collect();
                Some(names)
            }
            Mach::Fat(_) => None,
        },
        Object::Archive(_) | Object::Unknown(_) => None,
    };
}

fn is_derive_registrar_symbol(symbol: &str) -> bool {
    symbol.contains(NEW_REGISTRAR_SYMBOL)
}

fn find_registrar_symbol(file: &Path) -> Option<String> {
    let symbols = get_symbols_from_lib(file)?;

    symbols.iter().find(|s| is_derive_registrar_symbol(s)).map(|s| s.clone())
}

/// Loads dynamic library in platform dependent manner.
///
/// For unix, you have to use RTLD_DEEPBIND flag to escape problems described
/// [here](https://github.com/fedochet/rust-proc-macro-panic-inside-panic-expample)
/// and [here](https://github.com/rust-lang/rust/issues/60593).
///
/// Usage of RTLD_DEEPBIND
/// [here](https://github.com/fedochet/rust-proc-macro-panic-inside-panic-expample/issues/1)
///
/// It seems that on Windows that behaviour is default, so we do nothing in that case.
#[cfg(windows)]
fn load_library(file: &Path) -> Result<Library, libloading::Error> {
    Library::new(file)
}

#[cfg(unix)]
fn load_library(file: &Path) -> Result<Library, libloading::Error> {
    use libloading::os::unix::Library as UnixLibrary;
    use std::os::raw::c_int;

    const RTLD_NOW: c_int = 0x00002;
    const RTLD_DEEPBIND: c_int = 0x00008;

    UnixLibrary::open(Some(file), RTLD_NOW | RTLD_DEEPBIND).map(|lib| lib.into())
}

struct ProcMacroLibraryLibloading {
    // Hold the dylib to prevent it for unloadeding
    _lib: Library,
    exported_macros: Vec<bridge::client::ProcMacro>,
}

impl ProcMacroLibraryLibloading {
    fn open(file: &Path) -> Result<Self, String> {
        let symbol_name = find_registrar_symbol(file)
            .ok_or(format!("Cannot find registrar symbol in file {:?}", file))?;

        let lib = load_library(file).map_err(|e| e.to_string())?;
        let exported_macros = {
            let macros: libloading::Symbol<&&[bridge::client::ProcMacro]> =
                unsafe { lib.get(symbol_name.as_bytes()) }.map_err(|e| e.to_string())?;
            macros.to_vec()
        };

        Ok(ProcMacroLibraryLibloading { _lib: lib, exported_macros })
    }
}

type ProcMacroLibraryImpl = ProcMacroLibraryLibloading;

pub struct Expander {
    libs: Vec<ProcMacroLibraryImpl>,
}

impl Expander {
    pub fn new<P: AsRef<Path>>(lib: &P) -> Result<Expander, String> {
        let mut libs = vec![];
        /* Some libraries for dynamic loading require canonicalized path (even when it is
        already absolute
        */
        let lib =
            lib.as_ref().canonicalize().expect(&format!("Cannot canonicalize {:?}", lib.as_ref()));

        let library = ProcMacroLibraryImpl::open(&lib)?;
        libs.push(library);

        Ok(Expander { libs })
    }

    pub fn expand(
        &self,
        macro_name: &str,
        macro_body: &ra_tt::Subtree,
        attributes: Option<&ra_tt::Subtree>,
    ) -> Result<ra_tt::Subtree, bridge::PanicMessage> {
        let parsed_body = TokenStream::with_subtree(macro_body.clone());

        let parsed_attributes = attributes
            .map_or(crate::rustc_server::TokenStream::new(), |attr| {
                TokenStream::with_subtree(attr.clone())
            });

        for lib in &self.libs {
            for proc_macro in &lib.exported_macros {
                match proc_macro {
                    bridge::client::ProcMacro::CustomDerive { trait_name, client, .. }
                        if *trait_name == macro_name =>
                    {
                        let res = client.run(
                            &crate::proc_macro::bridge::server::SameThread,
                            crate::rustc_server::Rustc::default(),
                            parsed_body,
                        );
                        return res.map(|it| it.subtree);
                    }
                    bridge::client::ProcMacro::Bang { name, client } if *name == macro_name => {
                        let res = client.run(
                            &crate::proc_macro::bridge::server::SameThread,
                            crate::rustc_server::Rustc::default(),
                            parsed_body,
                        );
                        return res.map(|it| it.subtree);
                    }
                    bridge::client::ProcMacro::Attr { name, client } if *name == macro_name => {
                        let res = client.run(
                            &crate::proc_macro::bridge::server::SameThread,
                            crate::rustc_server::Rustc::default(),
                            parsed_attributes,
                            parsed_body,
                        );

                        return res.map(|it| it.subtree);
                    }
                    _ => continue,
                }
            }
        }

        Err(bridge::PanicMessage::String("Nothing to expand".to_string()))
    }

    pub fn list_macros(&self) -> Result<Vec<(String, ProcMacroKind)>, bridge::PanicMessage> {
        let mut result = vec![];

        for lib in &self.libs {
            for proc_macro in &lib.exported_macros {
                let res = match proc_macro {
                    bridge::client::ProcMacro::CustomDerive { trait_name, .. } => {
                        (trait_name.to_string(), ProcMacroKind::CustomDerive)
                    }
                    bridge::client::ProcMacro::Bang { name, .. } => {
                        (name.to_string(), ProcMacroKind::FuncLike)
                    }
                    bridge::client::ProcMacro::Attr { name, .. } => {
                        (name.to_string(), ProcMacroKind::Attr)
                    }
                };
                result.push(res);
            }
        }

        Ok(result)
    }
}
<|endoftext|>


<commit_before>// Spans are encoded using 1-bit tag and 2 different encoding formats (one for each tag value).
// One format is used for keeping span data inline,
// another contains index into an out-of-line span interner.
// The encoding format for inline spans were obtained by optimizing over crates in rustc/libstd.
// See https://internals.rust-lang.org/t/rfc-compiler-refactoring-spans/1357/28

use crate::hygiene::SyntaxContext;
use crate::SESSION_GLOBALS;
use crate::{BytePos, SpanData};

use rustc_data_structures::fx::FxHashMap;

/// A compressed span.
///
/// `SpanData` is 12 bytes, which is a bit too big to stick everywhere. `Span`
/// is a form that only takes up 8 bytes, with less space for the length and
/// context. The vast majority (99.9%+) of `SpanData` instances will fit within
/// those 8 bytes; any `SpanData` whose fields don't fit into a `Span` are
/// stored in a separate interner table, and the `Span` will index into that
/// table. Interning is rare enough that the cost is low, but common enough
/// that the code is exercised regularly.
///
/// An earlier version of this code used only 4 bytes for `Span`, but that was
/// slower because only 80--90% of spans could be stored inline (even less in
/// very large crates) and so the interner was used a lot more.
///
/// Inline (compressed) format:
/// - `span.base_or_index == span_data.lo`
/// - `span.len_or_tag == len == span_data.hi - span_data.lo` (must be `<= MAX_LEN`)
/// - `span.ctxt == span_data.ctxt` (must be `<= MAX_CTXT`)
///
/// Interned format:
/// - `span.base_or_index == index` (indexes into the interner table)
/// - `span.len_or_tag == LEN_TAG` (high bit set, all other bits are zero)
/// - `span.ctxt == 0`
///
/// The inline form uses 0 for the tag value (rather than 1) so that we don't
/// need to mask out the tag bit when getting the length, and so that the
/// dummy span can be all zeroes.
///
/// Notes about the choice of field sizes:
/// - `base` is 32 bits in both `Span` and `SpanData`, which means that `base`
///   values never cause interning. The number of bits needed for `base`
///   depends on the crate size. 32 bits allows up to 4 GiB of code in a crate.
///   `script-servo` is the largest crate in `rustc-perf`, requiring 26 bits
///   for some spans.
/// - `len` is 15 bits in `Span` (a u16, minus 1 bit for the tag) and 32 bits
///   in `SpanData`, which means that large `len` values will cause interning.
///   The number of bits needed for `len` does not depend on the crate size.
///   The most common number of bits for `len` are 0--7, with a peak usually at
///   3 or 4, and then it drops off quickly from 8 onwards. 15 bits is enough
///   for 99.99%+ of cases, but larger values (sometimes 20+ bits) might occur
///   dozens of times in a typical crate.
/// - `ctxt` is 16 bits in `Span` and 32 bits in `SpanData`, which means that
///   large `ctxt` values will cause interning. The number of bits needed for
///   `ctxt` values depend partly on the crate size and partly on the form of
///   the code. No crates in `rustc-perf` need more than 15 bits for `ctxt`,
///   but larger crates might need more than 16 bits.
///
#[derive(Clone, Copy, Eq, PartialEq, Hash)]
pub struct Span {
    base_or_index: u32,
    len_or_tag: u16,
    ctxt_or_zero: u16,
}

const LEN_TAG: u16 = 0b1000_0000_0000_0000;
const MAX_LEN: u32 = 0b0111_1111_1111_1111;
const MAX_CTXT: u32 = 0b1111_1111_1111_1111;

/// Dummy span, both position and length are zero, syntax context is zero as well.
pub const DUMMY_SP: Span = Span { base_or_index: 0, len_or_tag: 0, ctxt_or_zero: 0 };

impl Span {
    #[inline]
    pub fn new(mut lo: BytePos, mut hi: BytePos, ctxt: SyntaxContext) -> Self {
        if lo > hi {
            std::mem::swap(&mut lo, &mut hi);
        }

        let (base, len, ctxt2) = (lo.0, hi.0 - lo.0, ctxt.as_u32());

        if len <= MAX_LEN && ctxt2 <= MAX_CTXT {
            // Inline format.
            Span { base_or_index: base, len_or_tag: len as u16, ctxt_or_zero: ctxt2 as u16 }
        } else {
            // Interned format.
            let index = with_span_interner(|interner| interner.intern(&SpanData { lo, hi, ctxt }));
            Span { base_or_index: index, len_or_tag: LEN_TAG, ctxt_or_zero: 0 }
        }
    }

    #[inline]
    pub fn data(self) -> SpanData {
        if self.len_or_tag != LEN_TAG {
            // Inline format.
            debug_assert!(self.len_or_tag as u32 <= MAX_LEN);
            SpanData {
                lo: BytePos(self.base_or_index),
                hi: BytePos(self.base_or_index + self.len_or_tag as u32),
                ctxt: SyntaxContext::from_u32(self.ctxt_or_zero as u32),
            }
        } else {
            // Interned format.
            debug_assert!(self.ctxt_or_zero == 0);
            let index = self.base_or_index;
            with_span_interner(|interner| *interner.get(index))
        }
    }
}

#[derive(Default)]
pub struct SpanInterner {
    spans: FxHashMap<SpanData, u32>,
    span_data: Vec<SpanData>,
}

impl SpanInterner {
    fn intern(&mut self, span_data: &SpanData) -> u32 {
        if let Some(index) = self.spans.get(span_data) {
            return *index;
        }

        let index = self.spans.len() as u32;
        self.span_data.push(*span_data);
        self.spans.insert(*span_data, index);
        index
    }

    #[inline]
    fn get(&self, index: u32) -> &SpanData {
        &self.span_data[index as usize]
    }
}

// If an interner exists, return it. Otherwise, prepare a fresh one.
#[inline]
fn with_span_interner<T, F: FnOnce(&mut SpanInterner) -> T>(f: F) -> T {
    SESSION_GLOBALS.with(|session_globals| f(&mut *session_globals.span_interner.lock()))
}
<commit_msg>rustc_span: use IndexSet in SpanInterner<commit_after>// Spans are encoded using 1-bit tag and 2 different encoding formats (one for each tag value).
// One format is used for keeping span data inline,
// another contains index into an out-of-line span interner.
// The encoding format for inline spans were obtained by optimizing over crates in rustc/libstd.
// See https://internals.rust-lang.org/t/rfc-compiler-refactoring-spans/1357/28

use crate::hygiene::SyntaxContext;
use crate::SESSION_GLOBALS;
use crate::{BytePos, SpanData};

use rustc_data_structures::fx::FxIndexSet;

/// A compressed span.
///
/// `SpanData` is 12 bytes, which is a bit too big to stick everywhere. `Span`
/// is a form that only takes up 8 bytes, with less space for the length and
/// context. The vast majority (99.9%+) of `SpanData` instances will fit within
/// those 8 bytes; any `SpanData` whose fields don't fit into a `Span` are
/// stored in a separate interner table, and the `Span` will index into that
/// table. Interning is rare enough that the cost is low, but common enough
/// that the code is exercised regularly.
///
/// An earlier version of this code used only 4 bytes for `Span`, but that was
/// slower because only 80--90% of spans could be stored inline (even less in
/// very large crates) and so the interner was used a lot more.
///
/// Inline (compressed) format:
/// - `span.base_or_index == span_data.lo`
/// - `span.len_or_tag == len == span_data.hi - span_data.lo` (must be `<= MAX_LEN`)
/// - `span.ctxt == span_data.ctxt` (must be `<= MAX_CTXT`)
///
/// Interned format:
/// - `span.base_or_index == index` (indexes into the interner table)
/// - `span.len_or_tag == LEN_TAG` (high bit set, all other bits are zero)
/// - `span.ctxt == 0`
///
/// The inline form uses 0 for the tag value (rather than 1) so that we don't
/// need to mask out the tag bit when getting the length, and so that the
/// dummy span can be all zeroes.
///
/// Notes about the choice of field sizes:
/// - `base` is 32 bits in both `Span` and `SpanData`, which means that `base`
///   values never cause interning. The number of bits needed for `base`
///   depends on the crate size. 32 bits allows up to 4 GiB of code in a crate.
///   `script-servo` is the largest crate in `rustc-perf`, requiring 26 bits
///   for some spans.
/// - `len` is 15 bits in `Span` (a u16, minus 1 bit for the tag) and 32 bits
///   in `SpanData`, which means that large `len` values will cause interning.
///   The number of bits needed for `len` does not depend on the crate size.
///   The most common number of bits for `len` are 0--7, with a peak usually at
///   3 or 4, and then it drops off quickly from 8 onwards. 15 bits is enough
///   for 99.99%+ of cases, but larger values (sometimes 20+ bits) might occur
///   dozens of times in a typical crate.
/// - `ctxt` is 16 bits in `Span` and 32 bits in `SpanData`, which means that
///   large `ctxt` values will cause interning. The number of bits needed for
///   `ctxt` values depend partly on the crate size and partly on the form of
///   the code. No crates in `rustc-perf` need more than 15 bits for `ctxt`,
///   but larger crates might need more than 16 bits.
///
#[derive(Clone, Copy, Eq, PartialEq, Hash)]
pub struct Span {
    base_or_index: u32,
    len_or_tag: u16,
    ctxt_or_zero: u16,
}

const LEN_TAG: u16 = 0b1000_0000_0000_0000;
const MAX_LEN: u32 = 0b0111_1111_1111_1111;
const MAX_CTXT: u32 = 0b1111_1111_1111_1111;

/// Dummy span, both position and length are zero, syntax context is zero as well.
pub const DUMMY_SP: Span = Span { base_or_index: 0, len_or_tag: 0, ctxt_or_zero: 0 };

impl Span {
    #[inline]
    pub fn new(mut lo: BytePos, mut hi: BytePos, ctxt: SyntaxContext) -> Self {
        if lo > hi {
            std::mem::swap(&mut lo, &mut hi);
        }

        let (base, len, ctxt2) = (lo.0, hi.0 - lo.0, ctxt.as_u32());

        if len <= MAX_LEN && ctxt2 <= MAX_CTXT {
            // Inline format.
            Span { base_or_index: base, len_or_tag: len as u16, ctxt_or_zero: ctxt2 as u16 }
        } else {
            // Interned format.
            let index = with_span_interner(|interner| interner.intern(&SpanData { lo, hi, ctxt }));
            Span { base_or_index: index, len_or_tag: LEN_TAG, ctxt_or_zero: 0 }
        }
    }

    #[inline]
    pub fn data(self) -> SpanData {
        if self.len_or_tag != LEN_TAG {
            // Inline format.
            debug_assert!(self.len_or_tag as u32 <= MAX_LEN);
            SpanData {
                lo: BytePos(self.base_or_index),
                hi: BytePos(self.base_or_index + self.len_or_tag as u32),
                ctxt: SyntaxContext::from_u32(self.ctxt_or_zero as u32),
            }
        } else {
            // Interned format.
            debug_assert!(self.ctxt_or_zero == 0);
            let index = self.base_or_index;
            with_span_interner(|interner| *interner.get(index))
        }
    }
}

#[derive(Default)]
pub struct SpanInterner {
    spans: FxIndexSet<SpanData>,
}

impl SpanInterner {
    fn intern(&mut self, span_data: &SpanData) -> u32 {
        let (index, _) = self.spans.insert_full(*span_data);
        index as u32
    }

    #[inline]
    fn get(&self, index: u32) -> &SpanData {
        &self.spans[index as usize]
    }
}

// If an interner exists, return it. Otherwise, prepare a fresh one.
#[inline]
fn with_span_interner<T, F: FnOnce(&mut SpanInterner) -> T>(f: F) -> T {
    SESSION_GLOBALS.with(|session_globals| f(&mut *session_globals.span_interner.lock()))
}
<|endoftext|>