fix(ast)!: add missing ChainExpression from TSNonNullExpression (#7377)

closes #7375

* `foo?.bar!`
* `foo?.[bar]!`

`TSNonNullExpression` was not wrapped inside `ChainExpression`.

crates/oxc_ast/src/ast/js.rs

@@ -950,6 +950,8 @@ inherit_variants! {
 #[generate_derive(CloneIn, GetSpan, GetSpanMut, ContentEq, ContentHash, ESTree)]
 pub enum ChainElement<'a> {
     CallExpression(Box<'a, CallExpression<'a>>) = 0,
+    /// `foo?.baz!` or `foo?.[bar]!`
+    TSNonNullExpression(Box<'a, TSNonNullExpression<'a>>) = 1,
     // `MemberExpression` variants added here by `inherit_variants!` macro
     @inherit MemberExpression
 }

crates/oxc_ast/src/ast_impl/js.rs

@@ -462,10 +462,9 @@ impl<'a> MemberExpression<'a> {
     #[allow(missing_docs)]
     pub fn through_optional_is_specific_member_access(&self, object: &str, property: &str) -> bool {
         let object_matches = match self.object().without_parentheses() {
-            Expression::ChainExpression(x) => match &x.expression {
-                ChainElement::CallExpression(_) => false,
-                match_member_expression!(ChainElement) => {
-                    let member_expr = x.expression.to_member_expression();
+            Expression::ChainExpression(x) => match x.expression.member_expression() {
+                None => false,
+                Some(member_expr) => {
                     member_expr.object().without_parentheses().is_specific_id(object)
                 }
             },
@@ -523,6 +522,19 @@ impl<'a> StaticMemberExpression<'a> {
     }
 }
 
+impl<'a> ChainElement<'a> {
+    /// Returns the member expression.
+    pub fn member_expression(&self) -> Option<&MemberExpression<'a>> {
+        match self {
+            ChainElement::TSNonNullExpression(e) => match &e.expression {
+                match_member_expression!(Expression) => e.expression.as_member_expression(),
+                _ => None,
+            },
+            _ => self.as_member_expression(),
+        }
+    }
+}
+
 impl<'a> CallExpression<'a> {
     #[allow(missing_docs)]
     pub fn callee_name(&self) -> Option<&str> {

crates/oxc_ast/src/generated/ast_builder.rs

@@ -3092,6 +3092,22 @@ impl<'a> AstBuilder<'a> {
         )))
     }
 
+    /// Build a [`ChainElement::TSNonNullExpression`]
+    ///
+    /// This node contains a [`TSNonNullExpression`] that will be stored in the memory arena.
+    ///
+    /// ## Parameters
+    /// - span: The [`Span`] covering this node
+    /// - expression
+    #[inline]
+    pub fn chain_element_ts_non_null_expression(
+        self,
+        span: Span,
+        expression: Expression<'a>,
+    ) -> ChainElement<'a> {
+        ChainElement::TSNonNullExpression(self.alloc(self.ts_non_null_expression(span, expression)))
+    }
+
     /// Build a [`ParenthesizedExpression`].
     ///
     /// If you want the built node to be allocated in the memory arena, use [`AstBuilder::alloc_parenthesized_expression`] instead.

crates/oxc_ast/src/generated/derive_clone_in.rs

@@ -1210,6 +1210,9 @@ impl<'old_alloc, 'new_alloc> CloneIn<'new_alloc> for ChainElement<'old_alloc> {
             Self::CallExpression(it) => {
                 ChainElement::CallExpression(CloneIn::clone_in(it, allocator))
             }
+            Self::TSNonNullExpression(it) => {
+                ChainElement::TSNonNullExpression(CloneIn::clone_in(it, allocator))
+            }
             Self::ComputedMemberExpression(it) => {
                 ChainElement::ComputedMemberExpression(CloneIn::clone_in(it, allocator))
             }

crates/oxc_ast/src/generated/derive_content_eq.rs

@@ -1310,6 +1310,10 @@ impl<'a> ContentEq for ChainElement<'a> {
                 Self::CallExpression(other) if ContentEq::content_eq(it, other) => true,
                 _ => false,
             },
+            Self::TSNonNullExpression(it) => match other {
+                Self::TSNonNullExpression(other) if ContentEq::content_eq(it, other) => true,
+                _ => false,
+            },
             Self::ComputedMemberExpression(it) => match other {
                 Self::ComputedMemberExpression(other) if ContentEq::content_eq(it, other) => true,
                 _ => false,

crates/oxc_ast/src/generated/derive_content_hash.rs

@@ -641,6 +641,7 @@ impl<'a> ContentHash for ChainElement<'a> {
         ContentHash::content_hash(&discriminant(self), state);
         match self {
             Self::CallExpression(it) => ContentHash::content_hash(it, state),
+            Self::TSNonNullExpression(it) => ContentHash::content_hash(it, state),
             Self::ComputedMemberExpression(it) => ContentHash::content_hash(it, state),
             Self::StaticMemberExpression(it) => ContentHash::content_hash(it, state),
             Self::PrivateFieldExpression(it) => ContentHash::content_hash(it, state),

crates/oxc_ast/src/generated/derive_estree.rs

@@ -838,6 +838,7 @@ impl<'a> Serialize for ChainElement<'a> {
     fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
         match self {
             ChainElement::CallExpression(x) => Serialize::serialize(x, serializer),
+            ChainElement::TSNonNullExpression(x) => Serialize::serialize(x, serializer),
             ChainElement::ComputedMemberExpression(x) => Serialize::serialize(x, serializer),
             ChainElement::StaticMemberExpression(x) => Serialize::serialize(x, serializer),
             ChainElement::PrivateFieldExpression(x) => Serialize::serialize(x, serializer),

crates/oxc_ast/src/generated/derive_get_span.rs

@@ -609,6 +609,7 @@ impl<'a> GetSpan for ChainElement<'a> {
     fn span(&self) -> Span {
         match self {
             Self::CallExpression(it) => GetSpan::span(it.as_ref()),
+            Self::TSNonNullExpression(it) => GetSpan::span(it.as_ref()),
             Self::ComputedMemberExpression(it) => GetSpan::span(it.as_ref()),
             Self::StaticMemberExpression(it) => GetSpan::span(it.as_ref()),
             Self::PrivateFieldExpression(it) => GetSpan::span(it.as_ref()),

crates/oxc_ast/src/generated/derive_get_span_mut.rs

@@ -609,6 +609,7 @@ impl<'a> GetSpanMut for ChainElement<'a> {
     fn span_mut(&mut self) -> &mut Span {
         match self {
             Self::CallExpression(it) => GetSpanMut::span_mut(&mut **it),
+            Self::TSNonNullExpression(it) => GetSpanMut::span_mut(&mut **it),
             Self::ComputedMemberExpression(it) => GetSpanMut::span_mut(&mut **it),
             Self::StaticMemberExpression(it) => GetSpanMut::span_mut(&mut **it),
             Self::PrivateFieldExpression(it) => GetSpanMut::span_mut(&mut **it),

crates/oxc_ast/src/generated/visit.rs

@@ -2909,6 +2909,7 @@ pub mod walk {
     pub fn walk_chain_element<'a, V: Visit<'a>>(visitor: &mut V, it: &ChainElement<'a>) {
         match it {
             ChainElement::CallExpression(it) => visitor.visit_call_expression(it),
+            ChainElement::TSNonNullExpression(it) => visitor.visit_ts_non_null_expression(it),
             match_member_expression!(ChainElement) => {
                 visitor.visit_member_expression(it.to_member_expression())
             }

crates/oxc_ast/src/generated/visit_mut.rs

@@ -3033,6 +3033,7 @@ pub mod walk_mut {
     pub fn walk_chain_element<'a, V: VisitMut<'a>>(visitor: &mut V, it: &mut ChainElement<'a>) {
         match it {
             ChainElement::CallExpression(it) => visitor.visit_call_expression(it),
+            ChainElement::TSNonNullExpression(it) => visitor.visit_ts_non_null_expression(it),
             match_member_expression!(ChainElement) => {
                 visitor.visit_member_expression(it.to_member_expression_mut())
             }

crates/oxc_codegen/src/gen.rs

@@ -2049,6 +2049,7 @@ impl<'a> GenExpr for ChainExpression<'a> {
     fn gen_expr(&self, p: &mut Codegen, precedence: Precedence, ctx: Context) {
         p.wrap(precedence >= Precedence::Postfix, |p| match &self.expression {
             ChainElement::CallExpression(expr) => expr.print_expr(p, precedence, ctx),
+            ChainElement::TSNonNullExpression(expr) => expr.print_expr(p, precedence, ctx),
             match_member_expression!(ChainElement) => {
                 self.expression.to_member_expression().print_expr(p, precedence, ctx);
             }

crates/oxc_linter/src/ast_util.rs

@@ -344,9 +344,9 @@ pub fn is_method_call<'a>(
     let callee_without_parentheses = call_expr.callee.without_parentheses();
     let member_expr = match callee_without_parentheses {
         match_member_expression!(Expression) => callee_without_parentheses.to_member_expression(),
-        Expression::ChainExpression(chain) => match chain.expression {
-            match_member_expression!(ChainElement) => chain.expression.to_member_expression(),
-            ChainElement::CallExpression(_) => return false,
+        Expression::ChainExpression(chain) => match chain.expression.member_expression() {
+            Some(e) => e,
+            None => return false,
         },
         _ => return false,
     };

crates/oxc_linter/src/rules/eslint/getter_return.rs

@@ -129,7 +129,7 @@ impl GetterReturn {
                 match_member_expression!(ChainElement) => {
                     Self::handle_member_expression(ce.expression.to_member_expression())
                 }
-                ChainElement::CallExpression(_) => {
+                ChainElement::CallExpression(_) | ChainElement::TSNonNullExpression(_) => {
                     false // todo: make a test for this
                 }
             },

crates/oxc_linter/src/rules/oxc/no_optional_chaining.rs

@@ -99,6 +99,8 @@ fn test() {
         ("var x = '?.'?.['?.']", None),
         ("var x = '?.'?.['?.']", None),
         ("a?.c?.b<c>", None),
+        ("foo?.bar!", None),
+        ("foo?.[bar]!", None),
         (
             "var x = a?.b",
             Some(serde_json::json!([{

crates/oxc_linter/src/snapshots/no_optional_chaining.snap

@@ -62,6 +62,18 @@ snapshot_kind: text
    · ───────
    ╰────
 
+  ⚠ oxc(no-optional-chaining): Optional chaining is not allowed.
+   ╭─[no_optional_chaining.tsx:1:1]
+ 1 │ foo?.bar!
+   · ─────────
+   ╰────
+
+  ⚠ oxc(no-optional-chaining): Optional chaining is not allowed.
+   ╭─[no_optional_chaining.tsx:1:1]
+ 1 │ foo?.[bar]!
+   · ───────────
+   ╰────
+
   ⚠ oxc(no-optional-chaining): Optional chaining is not allowed.
    ╭─[no_optional_chaining.tsx:1:9]
  1 │ var x = a?.b

crates/oxc_parser/src/js/expression.rs

@@ -587,12 +587,15 @@ impl<'a> ParserImpl<'a> {
     fn map_to_chain_expression(&mut self, span: Span, expr: Expression<'a>) -> Expression<'a> {
         match expr {
             match_member_expression!(Expression) => {
-                let member_expr = MemberExpression::try_from(expr).unwrap();
+                let member_expr = expr.into_member_expression();
                 self.ast.expression_chain(span, ChainElement::from(member_expr))
             }
             Expression::CallExpression(result) => {
                 self.ast.expression_chain(span, ChainElement::CallExpression(result))
             }
+            Expression::TSNonNullExpression(result) => {
+                self.ast.expression_chain(span, ChainElement::TSNonNullExpression(result))
+            }
             expr => expr,
         }
     }

crates/oxc_prettier/src/format/mod.rs

@@ -2438,6 +2438,7 @@ impl<'a> Format<'a> for ChainElement<'a> {
     fn format(&self, p: &mut Prettier<'a>) -> Doc<'a> {
         match self {
             Self::CallExpression(expr) => expr.format(p),
+            Self::TSNonNullExpression(expr) => expr.format(p),
             match_member_expression!(Self) => self.to_member_expression().format(p),
         }
     }

crates/oxc_prettier/src/needs_parens.rs

@@ -669,6 +669,9 @@ impl<'a> Prettier<'a> {
                 ChainElement::CallExpression(e) => {
                     Self::starts_with_no_lookahead_token(&e.callee, span)
                 }
+                ChainElement::TSNonNullExpression(e) => {
+                    Self::starts_with_no_lookahead_token(&e.expression, span)
+                }
                 ChainElement::ComputedMemberExpression(e) => {
                     Self::starts_with_no_lookahead_token(&e.object, span)
                 }

crates/oxc_transformer/src/lib.rs

@@ -220,6 +220,12 @@ impl<'a, 'ctx> Traverse<'a> for TransformerImpl<'a, 'ctx> {
         self.x1_jsx.enter_call_expression(expr, ctx);
     }
 
+    fn enter_chain_element(&mut self, element: &mut ChainElement<'a>, ctx: &mut TraverseCtx<'a>) {
+        if let Some(typescript) = self.x0_typescript.as_mut() {
+            typescript.enter_chain_element(element, ctx);
+        }
+    }
+
     fn enter_class(&mut self, class: &mut Class<'a>, ctx: &mut TraverseCtx<'a>) {
         if let Some(typescript) = self.x0_typescript.as_mut() {
             typescript.enter_class(class, ctx);

crates/oxc_transformer/src/typescript/annotations.rs

@@ -181,6 +181,21 @@ impl<'a, 'ctx> Traverse<'a> for TypeScriptAnnotations<'a, 'ctx> {
         expr.type_parameters = None;
     }
 
+    fn enter_chain_element(&mut self, element: &mut ChainElement<'a>, ctx: &mut TraverseCtx<'a>) {
+        if let ChainElement::TSNonNullExpression(e) = element {
+            *element = match ctx.ast.move_expression(e.expression.get_inner_expression_mut()) {
+                Expression::CallExpression(call_expr) => ChainElement::CallExpression(call_expr),
+                expr @ match_member_expression!(Expression) => {
+                    ChainElement::from(expr.into_member_expression())
+                }
+                _ => {
+                    /* syntax error */
+                    return;
+                }
+            }
+        }
+    }
+
     fn enter_class(&mut self, class: &mut Class<'a>, _ctx: &mut TraverseCtx<'a>) {
         class.type_parameters = None;
         class.super_type_parameters = None;

crates/oxc_transformer/src/typescript/mod.rs

@@ -106,6 +106,10 @@ impl<'a, 'ctx> Traverse<'a> for TypeScript<'a, 'ctx> {
         self.annotations.enter_call_expression(expr, ctx);
     }
 
+    fn enter_chain_element(&mut self, element: &mut ChainElement<'a>, ctx: &mut TraverseCtx<'a>) {
+        self.annotations.enter_chain_element(element, ctx);
+    }
+
     fn enter_class(&mut self, class: &mut Class<'a>, ctx: &mut TraverseCtx<'a>) {
         self.annotations.enter_class(class, ctx);
     }

crates/oxc_traverse/src/generated/walk.rs

@@ -1271,6 +1271,9 @@ pub(crate) unsafe fn walk_chain_element<'a, Tr: Traverse<'a>>(
         ChainElement::CallExpression(node) => {
             walk_call_expression(traverser, (&mut **node) as *mut _, ctx)
         }
+        ChainElement::TSNonNullExpression(node) => {
+            walk_ts_non_null_expression(traverser, (&mut **node) as *mut _, ctx)
+        }
         ChainElement::ComputedMemberExpression(_)
         | ChainElement::StaticMemberExpression(_)
         | ChainElement::PrivateFieldExpression(_) => {

npm/oxc-types/types.d.ts

@@ -452,7 +452,12 @@ export interface ChainExpression extends Span {
   expression: ChainElement;
 }
 
-export type ChainElement = CallExpression | ComputedMemberExpression | StaticMemberExpression | PrivateFieldExpression;
+export type ChainElement =
+  | CallExpression
+  | TSNonNullExpression
+  | ComputedMemberExpression
+  | StaticMemberExpression
+  | PrivateFieldExpression;
 
 export interface ParenthesizedExpression extends Span {
   type: 'ParenthesizedExpression';