crates/ide_assists/src/ast_transform.rs

   1 //! `AstTransformer`s are functions that replace nodes in an AST and can be easily combined.
   2 use hir::{HirDisplay, PathResolution, SemanticsScope};
   3 use ide_db::helpers::mod_path_to_ast;
   4 use rustc_hash::FxHashMap;
   5 use syntax::{
   6     algo::SyntaxRewriter,
   7     ast::{self, AstNode},
   8     SyntaxNode,
   9 };
  10
  11 pub fn apply<'a, N: AstNode>(transformer: &dyn AstTransform<'a>, node: N) -> N {
  12     SyntaxRewriter::from_fn(|element| match element {
  13         syntax::SyntaxElement::Node(n) => {
  14             let replacement = transformer.get_substitution(&n, transformer)?;
  15             Some(replacement.into())
  16         }
  17         _ => None,
  18     })
  19     .rewrite_ast(&node)
  20 }
  21
  22 /// `AstTransform` helps with applying bulk transformations to syntax nodes.
  23 ///
  24 /// This is mostly useful for IDE code generation. If you paste some existing
  25 /// code into a new context (for example, to add method overrides to an `impl`
  26 /// block), you generally want to appropriately qualify the names, and sometimes
  27 /// you might want to substitute generic parameters as well:
  28 ///
  29 /// ```
  30 /// mod x {
  31 ///   pub struct A;
  32 ///   pub trait T<U> { fn foo(&self, _: U) -> A; }
  33 /// }
  34 ///
  35 /// mod y {
  36 ///   use x::T;
  37 ///
  38 ///   impl T<()> for () {
  39 ///      // If we invoke **Add Missing Members** here, we want to copy-paste `foo`.
  40 ///      // But we want a slightly-modified version of it:
  41 ///      fn foo(&self, _: ()) -> x::A {}
  42 ///   }
  43 /// }
  44 /// ```
  45 ///
  46 /// So, a single `AstTransform` describes such function from `SyntaxNode` to
  47 /// `SyntaxNode`. Note that the API here is a bit too high-order and high-brow.
  48 /// We'd want to somehow express this concept simpler, but so far nobody got to
  49 /// simplifying this!
  50 pub trait AstTransform<'a> {
  51     fn get_substitution(
  52         &self,
  53         node: &SyntaxNode,
  54         recur: &dyn AstTransform<'a>,
  55     ) -> Option<SyntaxNode>;
  56
  57     fn or<T: AstTransform<'a> + 'a>(self, other: T) -> Box<dyn AstTransform<'a> + 'a>
  58     where
  59         Self: Sized + 'a,
  60     {
  61         Box::new(Or(Box::new(self), Box::new(other)))
  62     }
  63 }
  64
  65 struct Or<'a>(Box<dyn AstTransform<'a> + 'a>, Box<dyn AstTransform<'a> + 'a>);
  66
  67 impl<'a> AstTransform<'a> for Or<'a> {
  68     fn get_substitution(
  69         &self,
  70         node: &SyntaxNode,
  71         recur: &dyn AstTransform<'a>,
  72     ) -> Option<SyntaxNode> {
  73         self.0.get_substitution(node, recur).or_else(|| self.1.get_substitution(node, recur))
  74     }
  75 }
  76
  77 pub struct SubstituteTypeParams<'a> {
  78     source_scope: &'a SemanticsScope<'a>,
  79     substs: FxHashMap<hir::TypeParam, ast::Type>,
  80 }
  81
  82 impl<'a> SubstituteTypeParams<'a> {
  83     pub fn for_trait_impl(
  84         source_scope: &'a SemanticsScope<'a>,
  85         // FIXME: there's implicit invariant that `trait_` and  `source_scope` match...
  86         trait_: hir::Trait,
  87         impl_def: ast::Impl,
  88     ) -> SubstituteTypeParams<'a> {
  89         let substs = get_syntactic_substs(impl_def).unwrap_or_default();
  90         let generic_def: hir::GenericDef = trait_.into();
  91         let substs_by_param: FxHashMap<_, _> = generic_def
  92             .type_params(source_scope.db)
  93             .into_iter()
  94             // this is a trait impl, so we need to skip the first type parameter -- this is a bit hacky
  95             .skip(1)
  96             // The actual list of trait type parameters may be longer than the one
  97             // used in the `impl` block due to trailing default type parameters.
  98             // For that case we extend the `substs` with an empty iterator so we
  99             // can still hit those trailing values and check if they actually have
 100             // a default type. If they do, go for that type from `hir` to `ast` so
 101             // the resulting change can be applied correctly.
 102             .zip(substs.into_iter().map(Some).chain(std::iter::repeat(None)))
 103             .filter_map(|(k, v)| match v {
 104                 Some(v) => Some((k, v)),
 105                 None => {
 106                     let default = k.default(source_scope.db)?;
 107                     Some((
 108                         k,
 109                         ast::make::ty(
 110                             &default
 111                                 .display_source_code(source_scope.db, source_scope.module()?.into())
 112                                 .ok()?,
 113                         ),
 114                     ))
 115                 }
 116             })
 117             .collect();
 118         return SubstituteTypeParams { source_scope, substs: substs_by_param };
 119
 120         // FIXME: It would probably be nicer if we could get this via HIR (i.e. get the
 121         // trait ref, and then go from the types in the substs back to the syntax).
 122         fn get_syntactic_substs(impl_def: ast::Impl) -> Option<Vec<ast::Type>> {
 123             let target_trait = impl_def.trait_()?;
 124             let path_type = match target_trait {
 125                 ast::Type::PathType(path) => path,
 126                 _ => return None,
 127             };
 128             let generic_arg_list = path_type.path()?.segment()?.generic_arg_list()?;
 129
 130             let mut result = Vec::new();
 131             for generic_arg in generic_arg_list.generic_args() {
 132                 match generic_arg {
 133                     ast::GenericArg::TypeArg(type_arg) => result.push(type_arg.ty()?),
 134                     ast::GenericArg::AssocTypeArg(_)
 135                     | ast::GenericArg::LifetimeArg(_)
 136                     | ast::GenericArg::ConstArg(_) => (),
 137                 }
 138             }
 139
 140             Some(result)
 141         }
 142     }
 143 }
 144
 145 impl<'a> AstTransform<'a> for SubstituteTypeParams<'a> {
 146     fn get_substitution(
 147         &self,
 148         node: &SyntaxNode,
 149         _recur: &dyn AstTransform<'a>,
 150     ) -> Option<SyntaxNode> {
 151         let type_ref = ast::Type::cast(node.clone())?;
 152         let path = match &type_ref {
 153             ast::Type::PathType(path_type) => path_type.path()?,
 154             _ => return None,
 155         };
 156         let resolution = self.source_scope.speculative_resolve(&path)?;
 157         match resolution {
 158             hir::PathResolution::TypeParam(tp) => Some(self.substs.get(&tp)?.syntax().clone()),
 159             _ => None,
 160         }
 161     }
 162 }
 163
 164 pub struct QualifyPaths<'a> {
 165     target_scope: &'a SemanticsScope<'a>,
 166     source_scope: &'a SemanticsScope<'a>,
 167 }
 168
 169 impl<'a> QualifyPaths<'a> {
 170     pub fn new(target_scope: &'a SemanticsScope<'a>, source_scope: &'a SemanticsScope<'a>) -> Self {
 171         Self { target_scope, source_scope }
 172     }
 173 }
 174
 175 impl<'a> AstTransform<'a> for QualifyPaths<'a> {
 176     fn get_substitution(
 177         &self,
 178         node: &SyntaxNode,
 179         recur: &dyn AstTransform<'a>,
 180     ) -> Option<SyntaxNode> {
 181         // FIXME handle value ns?
 182         let from = self.target_scope.module()?;
 183         let p = ast::Path::cast(node.clone())?;
 184         if p.segment().and_then(|s| s.param_list()).is_some() {
 185             // don't try to qualify `Fn(Foo) -> Bar` paths, they are in prelude anyway
 186             return None;
 187         }
 188         let resolution = self.source_scope.speculative_resolve(&p)?;
 189         match resolution {
 190             PathResolution::Def(def) => {
 191                 let found_path = from.find_use_path(self.source_scope.db.upcast(), def)?;
 192                 let mut path = mod_path_to_ast(&found_path);
 193
 194                 let type_args = p
 195                     .segment()
 196                     .and_then(|s| s.generic_arg_list())
 197                     .map(|arg_list| apply(recur, arg_list));
 198                 if let Some(type_args) = type_args {
 199                     let last_segment = path.segment().unwrap();
 200                     path = path.with_segment(last_segment.with_generic_args(type_args))
 201                 }
 202
 203                 Some(path.syntax().clone())
 204             }
 205             PathResolution::Local(_)
 206             | PathResolution::TypeParam(_)
 207             | PathResolution::SelfType(_)
 208             | PathResolution::ConstParam(_) => None,
 209             PathResolution::Macro(_) => None,
 210             PathResolution::AssocItem(_) => None,
 211         }
 212     }
 213 }