add chain to all expressions

poppingmoon · Jul 10, 2024 · 104e74b · 104e74b
1 parent 2fcca58
commit 104e74b
Show file tree

Hide file tree

Showing 4 changed files with 104 additions and 9 deletions.
diff --git a/src/parser/node.rs b/src/parser/node.rs
@@ -313,6 +313,7 @@ impl From<Assign> for ast::Assign {
 #[derive(Debug, PartialEq, Clone)]
 pub struct Not {
     pub expr: Box<Expression>,
+    pub chain: Option<Vec<ChainMember>>,
     pub loc: Option<Loc>,
 }
 
@@ -330,6 +331,7 @@ pub struct And {
     pub left: Box<Expression>,
     pub right: Box<Expression>,
     pub operator_loc: Loc,
+    pub chain: Option<Vec<ChainMember>>,
     pub loc: Option<Loc>,
 }
 
@@ -349,6 +351,7 @@ pub struct Or {
     pub left: Box<Expression>,
     pub right: Box<Expression>,
     pub operator_loc: Loc,
+    pub chain: Option<Vec<ChainMember>>,
     pub loc: Option<Loc>,
 }
 
@@ -369,6 +372,7 @@ pub struct If {
     pub then: Box<StatementOrExpression>,
     pub elseif: Vec<Elseif>,
     pub else_: Option<Box<StatementOrExpression>>,
+    pub chain: Option<Vec<ChainMember>>,
     pub loc: Option<Loc>,
 }
 

diff --git a/src/parser/parser.rs b/src/parser/parser.rs
@@ -83,6 +83,7 @@ peg::parser! {
                                 left: left.0.into(),
                                 right: right.0.into(),
                                 operator_loc: Loc{ start, end: end - 1 },
+                                chain: None,
                                 loc: None,
                             }
                         ),
@@ -96,6 +97,7 @@ peg::parser! {
                                 left: left.0.into(),
                                 right: right.0.into(),
                                 operator_loc: Loc{ start, end: end - 1 },
+                                chain: None,
                                 loc: None,
                             }
                         ),
@@ -603,6 +605,7 @@ peg::parser! {
             = start:position!() "!" expr:expr() end:position!() {
                 Not {
                     expr: expr.into(),
+                    chain: None,
                     loc: Some(Loc{ start, end: end - 1 }),
                 }
             }
@@ -612,10 +615,38 @@ peg::parser! {
         rule chain() -> Expression
             = e:expr3() chain:chain_member()+ {
                 match e {
-                    Expression::Not(_) => e,
-                    Expression::And(_) => e,
-                    Expression::Or(_) => e,
-                    Expression::If(_) => e,
+                    Expression::Not(not) => {
+                        let mut c = not.chain.unwrap_or_default();
+                        c.extend(chain);
+                        Expression::Not(Not {
+                            chain: Some(c),
+                            ..not
+                        })
+                    },
+                    Expression::And(and) => {
+                        let mut c = and.chain.unwrap_or_default();
+                        c.extend(chain);
+                        Expression::And(And {
+                            chain: Some(c),
+                            ..and
+                        })
+                    },
+                    Expression::Or(or) => {
+                        let mut c = or.chain.unwrap_or_default();
+                        c.extend(chain);
+                        Expression::Or(Or {
+                            chain: Some(c),
+                            ..or
+                        })
+                    },
+                    Expression::If(if_) => {
+                        let mut c = if_.chain.unwrap_or_default();
+                        c.extend(chain);
+                        Expression::If(If {
+                            chain: Some(c),
+                            ..if_
+                        })
+                    },
                     Expression::Fn(fn_) => {
                         let mut c = fn_.chain.unwrap_or_default();
                         c.extend(chain);
@@ -763,6 +794,7 @@ peg::parser! {
                     then: then.into(),
                     elseif: elseif.unwrap_or_default(),
                     else_: else_block.map(Into::into),
+                    chain: None,
                     loc: Some(Loc{ start, end: end - 1 }),
                 }
             }

diff --git a/src/parser/plugins/transform_chain.rs b/src/parser/plugins/transform_chain.rs
@@ -13,7 +13,19 @@ impl Visitor for ChainTransformer {
     ) -> Result<cst::Expression, AiScriptError> {
         // chain
         match &expression {
-            cst::Expression::Fn(cst::Fn_ {
+            cst::Expression::Not(cst::Not {
+                chain: Some(chain), ..
+            })
+            | cst::Expression::And(cst::And {
+                chain: Some(chain), ..
+            })
+            | cst::Expression::Or(cst::Or {
+                chain: Some(chain), ..
+            })
+            | cst::Expression::If(cst::If {
+                chain: Some(chain), ..
+            })
+            | cst::Expression::Fn(cst::Fn_ {
                 chain: Some(chain), ..
             })
             | cst::Expression::Match(cst::Match {
@@ -50,10 +62,22 @@ impl Visitor for ChainTransformer {
                 chain: Some(chain), ..
             }) => Ok(chain.iter().fold(
                 match &expression {
-                    cst::Expression::Not(not) => cst::Expression::Not(not.clone()),
-                    cst::Expression::And(and) => cst::Expression::And(and.clone()),
-                    cst::Expression::Or(or) => cst::Expression::Or(or.clone()),
-                    cst::Expression::If(if_) => cst::Expression::If(if_.clone()),
+                    cst::Expression::Not(not) => cst::Expression::Not(cst::Not {
+                        chain: None,
+                        ..not.clone()
+                    }),
+                    cst::Expression::And(and) => cst::Expression::And(cst::And {
+                        chain: None,
+                        ..and.clone()
+                    }),
+                    cst::Expression::Or(or) => cst::Expression::Or(cst::Or {
+                        chain: None,
+                        ..or.clone()
+                    }),
+                    cst::Expression::If(if_) => cst::Expression::If(cst::If {
+                        chain: None,
+                        ..if_.clone()
+                    }),
                     cst::Expression::Fn(fn_) => cst::Expression::Fn(cst::Fn_ {
                         chain: None,
                         ..fn_.clone()

diff --git a/tests/test.rs b/tests/test.rs
@@ -1596,6 +1596,41 @@ mod chain {
         }
         panic!();
     }
+
+    #[tokio::test]
+    async fn property_chain_with_if() {
+        let ast = Parser::default()
+            .parse(
+                r#"
+                (if a b else c).d
+                "#,
+            )
+            .unwrap();
+        let line = ast.first().unwrap().clone();
+        if let Node::Expression(Expression::Prop(Prop { target, name, .. })) = line.clone() {
+            assert_eq!(name, "d".to_string());
+            if let Expression::If(If {
+                cond,
+                then,
+                else_: Some(else_),
+                ..
+            }) = *target
+            {
+                if let (
+                    Expression::Identifier(cond),
+                    StatementOrExpression::Expression(Expression::Identifier(then)),
+                    StatementOrExpression::Expression(Expression::Identifier(else_)),
+                ) = (*cond, *then, *else_)
+                {
+                    assert_eq!(cond.name, "a".to_string());
+                    assert_eq!(then.name, "b".to_string());
+                    assert_eq!(else_.name, "c".to_string());
+                    return;
+                }
+            }
+        }
+        panic!();
+    }
 }
 
 mod template_syntax {