src/tools/rust-analyzer/crates/syntax/src/ast/expr_ext.rs

   1 //! Various extension methods to ast Expr Nodes, which are hard to code-generate.
   2 //!
   3 //! These methods should only do simple, shallow tasks related to the syntax of the node itself.
   4
   5 use crate::{
   6     ast::{
   7         self,
   8         operators::{ArithOp, BinaryOp, CmpOp, LogicOp, Ordering, RangeOp, UnaryOp},
   9         support, AstChildren, AstNode,
  10     },
  11     AstToken,
  12     SyntaxKind::*,
  13     SyntaxNode, SyntaxToken, T,
  14 };
  15
  16 impl ast::HasAttrs for ast::Expr {}
  17
  18 impl ast::Expr {
  19     pub fn is_block_like(&self) -> bool {
  20         matches!(
  21             self,
  22             ast::Expr::IfExpr(_)
  23                 | ast::Expr::LoopExpr(_)
  24                 | ast::Expr::ForExpr(_)
  25                 | ast::Expr::WhileExpr(_)
  26                 | ast::Expr::BlockExpr(_)
  27                 | ast::Expr::MatchExpr(_)
  28         )
  29     }
  30 }
  31
  32 #[derive(Debug, Clone, PartialEq, Eq)]
  33 pub enum ElseBranch {
  34     Block(ast::BlockExpr),
  35     IfExpr(ast::IfExpr),
  36 }
  37
  38 impl From<ast::BlockExpr> for ElseBranch {
  39     fn from(block_expr: ast::BlockExpr) -> Self {
  40         Self::Block(block_expr)
  41     }
  42 }
  43
  44 impl From<ast::IfExpr> for ElseBranch {
  45     fn from(if_expr: ast::IfExpr) -> Self {
  46         Self::IfExpr(if_expr)
  47     }
  48 }
  49
  50 impl ast::IfExpr {
  51     pub fn then_branch(&self) -> Option<ast::BlockExpr> {
  52         self.children_after_condition().next()
  53     }
  54
  55     pub fn else_branch(&self) -> Option<ElseBranch> {
  56         let res = match self.children_after_condition().nth(1) {
  57             Some(block) => ElseBranch::Block(block),
  58             None => {
  59                 let elif = self.children_after_condition().next()?;
  60                 ElseBranch::IfExpr(elif)
  61             }
  62         };
  63         Some(res)
  64     }
  65
  66     fn children_after_condition<N: AstNode>(&self) -> impl Iterator<Item = N> {
  67         self.syntax().children().skip(1).filter_map(N::cast)
  68     }
  69 }
  70
  71 #[test]
  72 fn if_block_condition() {
  73     let parse = ast::SourceFile::parse(
  74         r#"
  75         fn test() {
  76             if { true } { "if" }
  77             else if { false } { "first elif" }
  78             else if true { "second elif" }
  79             else if (true) { "third elif" }
  80             else { "else" }
  81         }
  82         "#,
  83     );
  84     let if_ = parse.tree().syntax().descendants().find_map(ast::IfExpr::cast).unwrap();
  85     assert_eq!(if_.then_branch().unwrap().syntax().text(), r#"{ "if" }"#);
  86     let elif = match if_.else_branch().unwrap() {
  87         ElseBranch::IfExpr(elif) => elif,
  88         ElseBranch::Block(_) => panic!("should be `else if`"),
  89     };
  90     assert_eq!(elif.then_branch().unwrap().syntax().text(), r#"{ "first elif" }"#);
  91     let elif = match elif.else_branch().unwrap() {
  92         ElseBranch::IfExpr(elif) => elif,
  93         ElseBranch::Block(_) => panic!("should be `else if`"),
  94     };
  95     assert_eq!(elif.then_branch().unwrap().syntax().text(), r#"{ "second elif" }"#);
  96     let elif = match elif.else_branch().unwrap() {
  97         ElseBranch::IfExpr(elif) => elif,
  98         ElseBranch::Block(_) => panic!("should be `else if`"),
  99     };
 100     assert_eq!(elif.then_branch().unwrap().syntax().text(), r#"{ "third elif" }"#);
 101     let else_ = match elif.else_branch().unwrap() {
 102         ElseBranch::Block(else_) => else_,
 103         ElseBranch::IfExpr(_) => panic!("should be `else`"),
 104     };
 105     assert_eq!(else_.syntax().text(), r#"{ "else" }"#);
 106 }
 107
 108 #[test]
 109 fn if_condition_with_if_inside() {
 110     let parse = ast::SourceFile::parse(
 111         r#"
 112         fn test() {
 113             if if true { true } else { false } { "if" }
 114             else { "else" }
 115         }
 116         "#,
 117     );
 118     let if_ = parse.tree().syntax().descendants().find_map(ast::IfExpr::cast).unwrap();
 119     assert_eq!(if_.then_branch().unwrap().syntax().text(), r#"{ "if" }"#);
 120     let else_ = match if_.else_branch().unwrap() {
 121         ElseBranch::Block(else_) => else_,
 122         ElseBranch::IfExpr(_) => panic!("should be `else`"),
 123     };
 124     assert_eq!(else_.syntax().text(), r#"{ "else" }"#);
 125 }
 126
 127 impl ast::PrefixExpr {
 128     pub fn op_kind(&self) -> Option<UnaryOp> {
 129         let res = match self.op_token()?.kind() {
 130             T![*] => UnaryOp::Deref,
 131             T![!] => UnaryOp::Not,
 132             T![-] => UnaryOp::Neg,
 133             _ => return None,
 134         };
 135         Some(res)
 136     }
 137
 138     pub fn op_token(&self) -> Option<SyntaxToken> {
 139         self.syntax().first_child_or_token()?.into_token()
 140     }
 141 }
 142
 143 impl ast::BinExpr {
 144     pub fn op_details(&self) -> Option<(SyntaxToken, BinaryOp)> {
 145         self.syntax().children_with_tokens().filter_map(|it| it.into_token()).find_map(|c| {
 146             #[rustfmt::skip]
 147             let bin_op = match c.kind() {
 148                 T![||] => BinaryOp::LogicOp(LogicOp::Or),
 149                 T![&&] => BinaryOp::LogicOp(LogicOp::And),
 150
 151                 T![==] => BinaryOp::CmpOp(CmpOp::Eq { negated: false }),
 152                 T![!=] => BinaryOp::CmpOp(CmpOp::Eq { negated: true }),
 153                 T![<=] => BinaryOp::CmpOp(CmpOp::Ord { ordering: Ordering::Less,    strict: false }),
 154                 T![>=] => BinaryOp::CmpOp(CmpOp::Ord { ordering: Ordering::Greater, strict: false }),
 155                 T![<]  => BinaryOp::CmpOp(CmpOp::Ord { ordering: Ordering::Less,    strict: true }),
 156                 T![>]  => BinaryOp::CmpOp(CmpOp::Ord { ordering: Ordering::Greater, strict: true }),
 157
 158                 T![+]  => BinaryOp::ArithOp(ArithOp::Add),
 159                 T![*]  => BinaryOp::ArithOp(ArithOp::Mul),
 160                 T![-]  => BinaryOp::ArithOp(ArithOp::Sub),
 161                 T![/]  => BinaryOp::ArithOp(ArithOp::Div),
 162                 T![%]  => BinaryOp::ArithOp(ArithOp::Rem),
 163                 T![<<] => BinaryOp::ArithOp(ArithOp::Shl),
 164                 T![>>] => BinaryOp::ArithOp(ArithOp::Shr),
 165                 T![^]  => BinaryOp::ArithOp(ArithOp::BitXor),
 166                 T![|]  => BinaryOp::ArithOp(ArithOp::BitOr),
 167                 T![&]  => BinaryOp::ArithOp(ArithOp::BitAnd),
 168
 169                 T![=]   => BinaryOp::Assignment { op: None },
 170                 T![+=]  => BinaryOp::Assignment { op: Some(ArithOp::Add) },
 171                 T![*=]  => BinaryOp::Assignment { op: Some(ArithOp::Mul) },
 172                 T![-=]  => BinaryOp::Assignment { op: Some(ArithOp::Sub) },
 173                 T![/=]  => BinaryOp::Assignment { op: Some(ArithOp::Div) },
 174                 T![%=]  => BinaryOp::Assignment { op: Some(ArithOp::Rem) },
 175                 T![<<=] => BinaryOp::Assignment { op: Some(ArithOp::Shl) },
 176                 T![>>=] => BinaryOp::Assignment { op: Some(ArithOp::Shr) },
 177                 T![^=]  => BinaryOp::Assignment { op: Some(ArithOp::BitXor) },
 178                 T![|=]  => BinaryOp::Assignment { op: Some(ArithOp::BitOr) },
 179                 T![&=]  => BinaryOp::Assignment { op: Some(ArithOp::BitAnd) },
 180
 181                 _ => return None,
 182             };
 183             Some((c, bin_op))
 184         })
 185     }
 186
 187     pub fn op_kind(&self) -> Option<BinaryOp> {
 188         self.op_details().map(|t| t.1)
 189     }
 190
 191     pub fn op_token(&self) -> Option<SyntaxToken> {
 192         self.op_details().map(|t| t.0)
 193     }
 194
 195     pub fn lhs(&self) -> Option<ast::Expr> {
 196         support::children(self.syntax()).next()
 197     }
 198
 199     pub fn rhs(&self) -> Option<ast::Expr> {
 200         support::children(self.syntax()).nth(1)
 201     }
 202
 203     pub fn sub_exprs(&self) -> (Option<ast::Expr>, Option<ast::Expr>) {
 204         let mut children = support::children(self.syntax());
 205         let first = children.next();
 206         let second = children.next();
 207         (first, second)
 208     }
 209 }
 210
 211 impl ast::RangeExpr {
 212     fn op_details(&self) -> Option<(usize, SyntaxToken, RangeOp)> {
 213         self.syntax().children_with_tokens().enumerate().find_map(|(ix, child)| {
 214             let token = child.into_token()?;
 215             let bin_op = match token.kind() {
 216                 T![..] => RangeOp::Exclusive,
 217                 T![..=] => RangeOp::Inclusive,
 218                 _ => return None,
 219             };
 220             Some((ix, token, bin_op))
 221         })
 222     }
 223
 224     pub fn op_kind(&self) -> Option<RangeOp> {
 225         self.op_details().map(|t| t.2)
 226     }
 227
 228     pub fn op_token(&self) -> Option<SyntaxToken> {
 229         self.op_details().map(|t| t.1)
 230     }
 231
 232     pub fn start(&self) -> Option<ast::Expr> {
 233         let op_ix = self.op_details()?.0;
 234         self.syntax()
 235             .children_with_tokens()
 236             .take(op_ix)
 237             .find_map(|it| ast::Expr::cast(it.into_node()?))
 238     }
 239
 240     pub fn end(&self) -> Option<ast::Expr> {
 241         let op_ix = self.op_details()?.0;
 242         self.syntax()
 243             .children_with_tokens()
 244             .skip(op_ix + 1)
 245             .find_map(|it| ast::Expr::cast(it.into_node()?))
 246     }
 247 }
 248
 249 impl ast::IndexExpr {
 250     pub fn base(&self) -> Option<ast::Expr> {
 251         support::children(self.syntax()).next()
 252     }
 253     pub fn index(&self) -> Option<ast::Expr> {
 254         support::children(self.syntax()).nth(1)
 255     }
 256 }
 257
 258 pub enum ArrayExprKind {
 259     Repeat { initializer: Option<ast::Expr>, repeat: Option<ast::Expr> },
 260     ElementList(AstChildren<ast::Expr>),
 261 }
 262
 263 impl ast::ArrayExpr {
 264     pub fn kind(&self) -> ArrayExprKind {
 265         if self.is_repeat() {
 266             ArrayExprKind::Repeat {
 267                 initializer: support::children(self.syntax()).next(),
 268                 repeat: support::children(self.syntax()).nth(1),
 269             }
 270         } else {
 271             ArrayExprKind::ElementList(support::children(self.syntax()))
 272         }
 273     }
 274
 275     fn is_repeat(&self) -> bool {
 276         self.semicolon_token().is_some()
 277     }
 278 }
 279
 280 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 281 pub enum LiteralKind {
 282     String(ast::String),
 283     ByteString(ast::ByteString),
 284     IntNumber(ast::IntNumber),
 285     FloatNumber(ast::FloatNumber),
 286     Char(ast::Char),
 287     Byte(ast::Byte),
 288     Bool(bool),
 289 }
 290
 291 impl ast::Literal {
 292     pub fn token(&self) -> SyntaxToken {
 293         self.syntax()
 294             .children_with_tokens()
 295             .find(|e| e.kind() != ATTR && !e.kind().is_trivia())
 296             .and_then(|e| e.into_token())
 297             .unwrap()
 298     }
 299
 300     pub fn kind(&self) -> LiteralKind {
 301         let token = self.token();
 302
 303         if let Some(t) = ast::IntNumber::cast(token.clone()) {
 304             return LiteralKind::IntNumber(t);
 305         }
 306         if let Some(t) = ast::FloatNumber::cast(token.clone()) {
 307             return LiteralKind::FloatNumber(t);
 308         }
 309         if let Some(t) = ast::String::cast(token.clone()) {
 310             return LiteralKind::String(t);
 311         }
 312         if let Some(t) = ast::ByteString::cast(token.clone()) {
 313             return LiteralKind::ByteString(t);
 314         }
 315         if let Some(t) = ast::Char::cast(token.clone()) {
 316             return LiteralKind::Char(t);
 317         }
 318         if let Some(t) = ast::Byte::cast(token.clone()) {
 319             return LiteralKind::Byte(t);
 320         }
 321
 322         match token.kind() {
 323             T![true] => LiteralKind::Bool(true),
 324             T![false] => LiteralKind::Bool(false),
 325             _ => unreachable!(),
 326         }
 327     }
 328 }
 329
 330 pub enum BlockModifier {
 331     Async(SyntaxToken),
 332     Unsafe(SyntaxToken),
 333     Try(SyntaxToken),
 334     Const(SyntaxToken),
 335     Label(ast::Label),
 336 }
 337
 338 impl ast::BlockExpr {
 339     pub fn modifier(&self) -> Option<BlockModifier> {
 340         self.async_token()
 341             .map(BlockModifier::Async)
 342             .or_else(|| self.unsafe_token().map(BlockModifier::Unsafe))
 343             .or_else(|| self.try_token().map(BlockModifier::Try))
 344             .or_else(|| self.const_token().map(BlockModifier::Const))
 345             .or_else(|| self.label().map(BlockModifier::Label))
 346     }
 347     /// false if the block is an intrinsic part of the syntax and can't be
 348     /// replaced with arbitrary expression.
 349     ///
 350     /// ```not_rust
 351     /// fn foo() { not_stand_alone }
 352     /// const FOO: () = { stand_alone };
 353     /// ```
 354     pub fn is_standalone(&self) -> bool {
 355         let parent = match self.syntax().parent() {
 356             Some(it) => it,
 357             None => return true,
 358         };
 359         !matches!(parent.kind(), FN | IF_EXPR | WHILE_EXPR | LOOP_EXPR)
 360     }
 361 }
 362
 363 #[test]
 364 fn test_literal_with_attr() {
 365     let parse = ast::SourceFile::parse(r#"const _: &str = { #[attr] "Hello" };"#);
 366     let lit = parse.tree().syntax().descendants().find_map(ast::Literal::cast).unwrap();
 367     assert_eq!(lit.token().text(), r#""Hello""#);
 368 }
 369
 370 impl ast::RecordExprField {
 371     pub fn parent_record_lit(&self) -> ast::RecordExpr {
 372         self.syntax().ancestors().find_map(ast::RecordExpr::cast).unwrap()
 373     }
 374 }
 375
 376 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 377 pub enum CallableExpr {
 378     Call(ast::CallExpr),
 379     MethodCall(ast::MethodCallExpr),
 380 }
 381
 382 impl ast::HasAttrs for CallableExpr {}
 383 impl ast::HasArgList for CallableExpr {}
 384
 385 impl AstNode for CallableExpr {
 386     fn can_cast(kind: parser::SyntaxKind) -> bool
 387     where
 388         Self: Sized,
 389     {
 390         ast::CallExpr::can_cast(kind) || ast::MethodCallExpr::can_cast(kind)
 391     }
 392
 393     fn cast(syntax: SyntaxNode) -> Option<Self>
 394     where
 395         Self: Sized,
 396     {
 397         if let Some(it) = ast::CallExpr::cast(syntax.clone()) {
 398             Some(Self::Call(it))
 399         } else {
 400             ast::MethodCallExpr::cast(syntax).map(Self::MethodCall)
 401         }
 402     }
 403
 404     fn syntax(&self) -> &SyntaxNode {
 405         match self {
 406             Self::Call(it) => it.syntax(),
 407             Self::MethodCall(it) => it.syntax(),
 408         }
 409     }
 410 }