crates/syntax/src/ast/node_ext.rs

   1 //! Various extension methods to ast Nodes, which are hard to code-generate.
   2 //! Extensions for various expressions live in a sibling `expr_extensions` module.
   3
   4 use std::fmt;
   5
   6 use ast::AttrsOwner;
   7 use itertools::Itertools;
   8 use parser::SyntaxKind;
   9
  10 use crate::{
  11     ast::{self, support, AstNode, AstToken, NameOwner, SyntaxNode},
  12     SmolStr, SyntaxElement, SyntaxToken, T,
  13 };
  14
  15 impl ast::Lifetime {
  16     pub fn text(&self) -> &SmolStr {
  17         text_of_first_token(self.syntax())
  18     }
  19 }
  20
  21 impl ast::Name {
  22     pub fn text(&self) -> &SmolStr {
  23         text_of_first_token(self.syntax())
  24     }
  25 }
  26
  27 impl ast::NameRef {
  28     pub fn text(&self) -> &SmolStr {
  29         text_of_first_token(self.syntax())
  30     }
  31
  32     pub fn as_tuple_field(&self) -> Option<usize> {
  33         self.text().parse().ok()
  34     }
  35 }
  36
  37 fn text_of_first_token(node: &SyntaxNode) -> &SmolStr {
  38     node.green().children().next().and_then(|it| it.into_token()).unwrap().text()
  39 }
  40
  41 pub enum Macro {
  42     MacroRules(ast::MacroRules),
  43     MacroDef(ast::MacroDef),
  44 }
  45
  46 impl From<ast::MacroRules> for Macro {
  47     fn from(it: ast::MacroRules) -> Self {
  48         Macro::MacroRules(it)
  49     }
  50 }
  51
  52 impl From<ast::MacroDef> for Macro {
  53     fn from(it: ast::MacroDef) -> Self {
  54         Macro::MacroDef(it)
  55     }
  56 }
  57
  58 impl AstNode for Macro {
  59     fn can_cast(kind: SyntaxKind) -> bool {
  60         match kind {
  61             SyntaxKind::MACRO_RULES | SyntaxKind::MACRO_DEF => true,
  62             _ => false,
  63         }
  64     }
  65     fn cast(syntax: SyntaxNode) -> Option<Self> {
  66         let res = match syntax.kind() {
  67             SyntaxKind::MACRO_RULES => Macro::MacroRules(ast::MacroRules { syntax }),
  68             SyntaxKind::MACRO_DEF => Macro::MacroDef(ast::MacroDef { syntax }),
  69             _ => return None,
  70         };
  71         Some(res)
  72     }
  73     fn syntax(&self) -> &SyntaxNode {
  74         match self {
  75             Macro::MacroRules(it) => it.syntax(),
  76             Macro::MacroDef(it) => it.syntax(),
  77         }
  78     }
  79 }
  80
  81 impl NameOwner for Macro {
  82     fn name(&self) -> Option<ast::Name> {
  83         match self {
  84             Macro::MacroRules(mac) => mac.name(),
  85             Macro::MacroDef(mac) => mac.name(),
  86         }
  87     }
  88 }
  89
  90 impl AttrsOwner for Macro {}
  91
  92 #[derive(Debug, Clone, PartialEq, Eq)]
  93 pub enum AttrKind {
  94     Inner,
  95     Outer,
  96 }
  97
  98 impl ast::Attr {
  99     pub fn as_simple_atom(&self) -> Option<SmolStr> {
 100         if self.eq_token().is_some() || self.token_tree().is_some() {
 101             return None;
 102         }
 103         self.simple_name()
 104     }
 105
 106     pub fn as_simple_call(&self) -> Option<(SmolStr, ast::TokenTree)> {
 107         let tt = self.token_tree()?;
 108         Some((self.simple_name()?, tt))
 109     }
 110
 111     pub fn as_simple_key_value(&self) -> Option<(SmolStr, SmolStr)> {
 112         let lit = self.literal()?;
 113         let key = self.simple_name()?;
 114         let value_token = lit.syntax().first_token()?;
 115
 116         let value: SmolStr = ast::String::cast(value_token)?.value()?.into();
 117
 118         Some((key, value))
 119     }
 120
 121     pub fn simple_name(&self) -> Option<SmolStr> {
 122         let path = self.path()?;
 123         match (path.segment(), path.qualifier()) {
 124             (Some(segment), None) => Some(segment.syntax().first_token()?.text().clone()),
 125             _ => None,
 126         }
 127     }
 128
 129     pub fn kind(&self) -> AttrKind {
 130         let first_token = self.syntax().first_token();
 131         let first_token_kind = first_token.as_ref().map(SyntaxToken::kind);
 132         let second_token_kind =
 133             first_token.and_then(|token| token.next_token()).as_ref().map(SyntaxToken::kind);
 134
 135         match (first_token_kind, second_token_kind) {
 136             (Some(T![#]), Some(T![!])) => AttrKind::Inner,
 137             _ => AttrKind::Outer,
 138         }
 139     }
 140 }
 141
 142 #[derive(Debug, Clone, PartialEq, Eq)]
 143 pub enum PathSegmentKind {
 144     Name(ast::NameRef),
 145     Type { type_ref: Option<ast::Type>, trait_ref: Option<ast::PathType> },
 146     SelfKw,
 147     SuperKw,
 148     CrateKw,
 149 }
 150
 151 impl ast::PathSegment {
 152     pub fn parent_path(&self) -> ast::Path {
 153         self.syntax()
 154             .parent()
 155             .and_then(ast::Path::cast)
 156             .expect("segments are always nested in paths")
 157     }
 158
 159     pub fn crate_token(&self) -> Option<SyntaxToken> {
 160         self.name_ref().and_then(|it| it.crate_token())
 161     }
 162
 163     pub fn self_token(&self) -> Option<SyntaxToken> {
 164         self.name_ref().and_then(|it| it.self_token())
 165     }
 166
 167     pub fn super_token(&self) -> Option<SyntaxToken> {
 168         self.name_ref().and_then(|it| it.super_token())
 169     }
 170
 171     pub fn kind(&self) -> Option<PathSegmentKind> {
 172         let res = if let Some(name_ref) = self.name_ref() {
 173             match name_ref.syntax().first_token().map(|it| it.kind()) {
 174                 Some(T![self]) => PathSegmentKind::SelfKw,
 175                 Some(T![super]) => PathSegmentKind::SuperKw,
 176                 Some(T![crate]) => PathSegmentKind::CrateKw,
 177                 _ => PathSegmentKind::Name(name_ref),
 178             }
 179         } else {
 180             match self.syntax().first_child_or_token()?.kind() {
 181                 T![<] => {
 182                     // <T> or <T as Trait>
 183                     // T is any TypeRef, Trait has to be a PathType
 184                     let mut type_refs =
 185                         self.syntax().children().filter(|node| ast::Type::can_cast(node.kind()));
 186                     let type_ref = type_refs.next().and_then(ast::Type::cast);
 187                     let trait_ref = type_refs.next().and_then(ast::PathType::cast);
 188                     PathSegmentKind::Type { type_ref, trait_ref }
 189                 }
 190                 _ => return None,
 191             }
 192         };
 193         Some(res)
 194     }
 195 }
 196
 197 impl ast::Path {
 198     pub fn parent_path(&self) -> Option<ast::Path> {
 199         self.syntax().parent().and_then(ast::Path::cast)
 200     }
 201
 202     pub fn as_single_segment(&self) -> Option<ast::PathSegment> {
 203         match self.qualifier() {
 204             Some(_) => None,
 205             None => self.segment(),
 206         }
 207     }
 208 }
 209
 210 impl ast::UseTreeList {
 211     pub fn parent_use_tree(&self) -> ast::UseTree {
 212         self.syntax()
 213             .parent()
 214             .and_then(ast::UseTree::cast)
 215             .expect("UseTreeLists are always nested in UseTrees")
 216     }
 217
 218     pub fn has_inner_comment(&self) -> bool {
 219         self.syntax()
 220             .children_with_tokens()
 221             .filter_map(|it| it.into_token())
 222             .find_map(ast::Comment::cast)
 223             .is_some()
 224     }
 225 }
 226
 227 impl ast::Impl {
 228     pub fn self_ty(&self) -> Option<ast::Type> {
 229         match self.target() {
 230             (Some(t), None) | (_, Some(t)) => Some(t),
 231             _ => None,
 232         }
 233     }
 234
 235     pub fn trait_(&self) -> Option<ast::Type> {
 236         match self.target() {
 237             (Some(t), Some(_)) => Some(t),
 238             _ => None,
 239         }
 240     }
 241
 242     fn target(&self) -> (Option<ast::Type>, Option<ast::Type>) {
 243         let mut types = support::children(self.syntax());
 244         let first = types.next();
 245         let second = types.next();
 246         (first, second)
 247     }
 248 }
 249
 250 #[derive(Debug, Clone, PartialEq, Eq)]
 251 pub enum StructKind {
 252     Record(ast::RecordFieldList),
 253     Tuple(ast::TupleFieldList),
 254     Unit,
 255 }
 256
 257 impl StructKind {
 258     fn from_node<N: AstNode>(node: &N) -> StructKind {
 259         if let Some(nfdl) = support::child::<ast::RecordFieldList>(node.syntax()) {
 260             StructKind::Record(nfdl)
 261         } else if let Some(pfl) = support::child::<ast::TupleFieldList>(node.syntax()) {
 262             StructKind::Tuple(pfl)
 263         } else {
 264             StructKind::Unit
 265         }
 266     }
 267 }
 268
 269 impl ast::Struct {
 270     pub fn kind(&self) -> StructKind {
 271         StructKind::from_node(self)
 272     }
 273 }
 274
 275 impl ast::RecordExprField {
 276     pub fn for_field_name(field_name: &ast::NameRef) -> Option<ast::RecordExprField> {
 277         let candidate =
 278             field_name.syntax().parent().and_then(ast::RecordExprField::cast).or_else(|| {
 279                 field_name.syntax().ancestors().nth(4).and_then(ast::RecordExprField::cast)
 280             })?;
 281         if candidate.field_name().as_ref() == Some(field_name) {
 282             Some(candidate)
 283         } else {
 284             None
 285         }
 286     }
 287
 288     /// Deals with field init shorthand
 289     pub fn field_name(&self) -> Option<ast::NameRef> {
 290         if let Some(name_ref) = self.name_ref() {
 291             return Some(name_ref);
 292         }
 293         self.expr()?.name_ref()
 294     }
 295 }
 296
 297 pub enum NameOrNameRef {
 298     Name(ast::Name),
 299     NameRef(ast::NameRef),
 300 }
 301
 302 impl fmt::Display for NameOrNameRef {
 303     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 304         match self {
 305             NameOrNameRef::Name(it) => fmt::Display::fmt(it, f),
 306             NameOrNameRef::NameRef(it) => fmt::Display::fmt(it, f),
 307         }
 308     }
 309 }
 310
 311 impl ast::RecordPatField {
 312     /// Deals with field init shorthand
 313     pub fn field_name(&self) -> Option<NameOrNameRef> {
 314         if let Some(name_ref) = self.name_ref() {
 315             return Some(NameOrNameRef::NameRef(name_ref));
 316         }
 317         if let Some(ast::Pat::IdentPat(pat)) = self.pat() {
 318             let name = pat.name()?;
 319             return Some(NameOrNameRef::Name(name));
 320         }
 321         None
 322     }
 323 }
 324
 325 impl ast::Variant {
 326     pub fn parent_enum(&self) -> ast::Enum {
 327         self.syntax()
 328             .parent()
 329             .and_then(|it| it.parent())
 330             .and_then(ast::Enum::cast)
 331             .expect("EnumVariants are always nested in Enums")
 332     }
 333     pub fn kind(&self) -> StructKind {
 334         StructKind::from_node(self)
 335     }
 336 }
 337
 338 #[derive(Debug, Clone, PartialEq, Eq)]
 339 pub enum FieldKind {
 340     Name(ast::NameRef),
 341     Index(SyntaxToken),
 342 }
 343
 344 impl ast::FieldExpr {
 345     pub fn index_token(&self) -> Option<SyntaxToken> {
 346         self.syntax
 347             .children_with_tokens()
 348             // FIXME: Accepting floats here to reject them in validation later
 349             .find(|c| c.kind() == SyntaxKind::INT_NUMBER || c.kind() == SyntaxKind::FLOAT_NUMBER)
 350             .as_ref()
 351             .and_then(SyntaxElement::as_token)
 352             .cloned()
 353     }
 354
 355     pub fn field_access(&self) -> Option<FieldKind> {
 356         if let Some(nr) = self.name_ref() {
 357             Some(FieldKind::Name(nr))
 358         } else if let Some(tok) = self.index_token() {
 359             Some(FieldKind::Index(tok))
 360         } else {
 361             None
 362         }
 363     }
 364 }
 365
 366 pub struct SlicePatComponents {
 367     pub prefix: Vec<ast::Pat>,
 368     pub slice: Option<ast::Pat>,
 369     pub suffix: Vec<ast::Pat>,
 370 }
 371
 372 impl ast::SlicePat {
 373     pub fn components(&self) -> SlicePatComponents {
 374         let mut args = self.pats().peekable();
 375         let prefix = args
 376             .peeking_take_while(|p| match p {
 377                 ast::Pat::RestPat(_) => false,
 378                 ast::Pat::IdentPat(bp) => match bp.pat() {
 379                     Some(ast::Pat::RestPat(_)) => false,
 380                     _ => true,
 381                 },
 382                 ast::Pat::RefPat(rp) => match rp.pat() {
 383                     Some(ast::Pat::RestPat(_)) => false,
 384                     Some(ast::Pat::IdentPat(bp)) => match bp.pat() {
 385                         Some(ast::Pat::RestPat(_)) => false,
 386                         _ => true,
 387                     },
 388                     _ => true,
 389                 },
 390                 _ => true,
 391             })
 392             .collect();
 393         let slice = args.next();
 394         let suffix = args.collect();
 395
 396         SlicePatComponents { prefix, slice, suffix }
 397     }
 398 }
 399
 400 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
 401 pub enum SelfParamKind {
 402     /// self
 403     Owned,
 404     /// &self
 405     Ref,
 406     /// &mut self
 407     MutRef,
 408 }
 409
 410 impl ast::SelfParam {
 411     pub fn kind(&self) -> SelfParamKind {
 412         if self.amp_token().is_some() {
 413             if self.mut_token().is_some() {
 414                 SelfParamKind::MutRef
 415             } else {
 416                 SelfParamKind::Ref
 417             }
 418         } else {
 419             SelfParamKind::Owned
 420         }
 421     }
 422 }
 423
 424 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 425 pub enum TypeBoundKind {
 426     /// Trait
 427     PathType(ast::PathType),
 428     /// for<'a> ...
 429     ForType(ast::ForType),
 430     /// 'a
 431     Lifetime(ast::Lifetime),
 432 }
 433
 434 impl ast::TypeBound {
 435     pub fn kind(&self) -> TypeBoundKind {
 436         if let Some(path_type) = support::children(self.syntax()).next() {
 437             TypeBoundKind::PathType(path_type)
 438         } else if let Some(for_type) = support::children(self.syntax()).next() {
 439             TypeBoundKind::ForType(for_type)
 440         } else if let Some(lifetime) = self.lifetime() {
 441             TypeBoundKind::Lifetime(lifetime)
 442         } else {
 443             unreachable!()
 444         }
 445     }
 446 }
 447
 448 pub enum VisibilityKind {
 449     In(ast::Path),
 450     PubCrate,
 451     PubSuper,
 452     PubSelf,
 453     Pub,
 454 }
 455
 456 impl ast::Visibility {
 457     pub fn kind(&self) -> VisibilityKind {
 458         match self.path() {
 459             Some(path) => {
 460                 if let Some(segment) =
 461                     path.as_single_segment().filter(|it| it.coloncolon_token().is_none())
 462                 {
 463                     if segment.crate_token().is_some() {
 464                         return VisibilityKind::PubCrate;
 465                     } else if segment.super_token().is_some() {
 466                         return VisibilityKind::PubSuper;
 467                     } else if segment.self_token().is_some() {
 468                         return VisibilityKind::PubSelf;
 469                     }
 470                 }
 471                 VisibilityKind::In(path)
 472             }
 473             None => VisibilityKind::Pub,
 474         }
 475     }
 476 }
 477
 478 impl ast::LifetimeParam {
 479     pub fn lifetime_bounds(&self) -> impl Iterator<Item = SyntaxToken> {
 480         self.syntax()
 481             .children_with_tokens()
 482             .filter_map(|it| it.into_token())
 483             .skip_while(|x| x.kind() != T![:])
 484             .filter(|it| it.kind() == T![lifetime_ident])
 485     }
 486 }
 487
 488 impl ast::RangePat {
 489     pub fn start(&self) -> Option<ast::Pat> {
 490         self.syntax()
 491             .children_with_tokens()
 492             .take_while(|it| !(it.kind() == T![..] || it.kind() == T![..=]))
 493             .filter_map(|it| it.into_node())
 494             .find_map(ast::Pat::cast)
 495     }
 496
 497     pub fn end(&self) -> Option<ast::Pat> {
 498         self.syntax()
 499             .children_with_tokens()
 500             .skip_while(|it| !(it.kind() == T![..] || it.kind() == T![..=]))
 501             .filter_map(|it| it.into_node())
 502             .find_map(ast::Pat::cast)
 503     }
 504 }
 505
 506 impl ast::TokenTree {
 507     pub fn left_delimiter_token(&self) -> Option<SyntaxToken> {
 508         self.syntax()
 509             .first_child_or_token()?
 510             .into_token()
 511             .filter(|it| matches!(it.kind(), T!['{'] | T!['('] | T!['[']))
 512     }
 513
 514     pub fn right_delimiter_token(&self) -> Option<SyntaxToken> {
 515         self.syntax()
 516             .last_child_or_token()?
 517             .into_token()
 518             .filter(|it| matches!(it.kind(), T!['}'] | T![')'] | T![']']))
 519     }
 520 }
 521
 522 impl ast::GenericParamList {
 523     pub fn lifetime_params(&self) -> impl Iterator<Item = ast::LifetimeParam> {
 524         self.generic_params().filter_map(|param| match param {
 525             ast::GenericParam::LifetimeParam(it) => Some(it),
 526             ast::GenericParam::TypeParam(_) | ast::GenericParam::ConstParam(_) => None,
 527         })
 528     }
 529     pub fn type_params(&self) -> impl Iterator<Item = ast::TypeParam> {
 530         self.generic_params().filter_map(|param| match param {
 531             ast::GenericParam::TypeParam(it) => Some(it),
 532             ast::GenericParam::LifetimeParam(_) | ast::GenericParam::ConstParam(_) => None,
 533         })
 534     }
 535     pub fn const_params(&self) -> impl Iterator<Item = ast::ConstParam> {
 536         self.generic_params().filter_map(|param| match param {
 537             ast::GenericParam::ConstParam(it) => Some(it),
 538             ast::GenericParam::TypeParam(_) | ast::GenericParam::LifetimeParam(_) => None,
 539         })
 540     }
 541 }
 542
 543 impl ast::DocCommentsOwner for ast::SourceFile {}
 544 impl ast::DocCommentsOwner for ast::Fn {}
 545 impl ast::DocCommentsOwner for ast::Struct {}
 546 impl ast::DocCommentsOwner for ast::Union {}
 547 impl ast::DocCommentsOwner for ast::RecordField {}
 548 impl ast::DocCommentsOwner for ast::TupleField {}
 549 impl ast::DocCommentsOwner for ast::Enum {}
 550 impl ast::DocCommentsOwner for ast::Variant {}
 551 impl ast::DocCommentsOwner for ast::Trait {}
 552 impl ast::DocCommentsOwner for ast::Module {}
 553 impl ast::DocCommentsOwner for ast::Static {}
 554 impl ast::DocCommentsOwner for ast::Const {}
 555 impl ast::DocCommentsOwner for ast::TypeAlias {}
 556 impl ast::DocCommentsOwner for ast::Impl {}
 557 impl ast::DocCommentsOwner for ast::MacroRules {}
 558 impl ast::DocCommentsOwner for ast::MacroDef {}
 559 impl ast::DocCommentsOwner for ast::Macro {}
 560 impl ast::DocCommentsOwner for ast::Use {}