From 41e338b475f8801517e09532956b7eb8d5664757 Mon Sep 17 00:00:00 2001
From: WindOctober <windocotber@gmail.com>
Date: Thu, 3 Sep 2026 16:45:11 +0800
Subject: [PATCH] Fix IfStmt single-statement branch accessors

---
 crates/oq3_syntax/src/ast/node_ext.rs       | 29 ++++++++++++++++-----
 crates/oq3_syntax/tests/if_stmt_branches.rs | 19 ++++++++++++++
 2 files changed, 41 insertions(+), 7 deletions(-)
 create mode 100644 crates/oq3_syntax/tests/if_stmt_branches.rs

diff --git a/crates/oq3_syntax/src/ast/node_ext.rs b/crates/oq3_syntax/src/ast/node_ext.rs
index fd035a8..b8161ce 100644
--- a/crates/oq3_syntax/src/ast/node_ext.rs
+++ b/crates/oq3_syntax/src/ast/node_ext.rs
@@ -185,6 +185,14 @@ impl ForStmt {
 // }
 
 impl ast::IfStmt {
+    fn branches(&self) -> impl Iterator<Item = BlockOrStmt> + '_ {
+        self.syntax().children().filter_map(|node| {
+            ast::BlockExpr::cast(node.clone())
+                .map(BlockOrStmt::BlockExpr)
+                .or_else(|| ast::Stmt::cast(node).map(BlockOrStmt::Stmt))
+        })
+    }
+
     pub fn condition(&self) -> Option<ast::Expr> {
         // If the condition is a BlockExpr, check if the then body is missing.
         // If it is, assume the condition is the expression that is missing instead.
@@ -196,16 +204,20 @@ impl ast::IfStmt {
         }
     }
 
+    // Return `Some` if the then branch is a curly-delimited block.
     pub fn then_branch_block(&self) -> Option<ast::BlockExpr> {
-        match support::children(self.syntax()).nth(1)? {
-            ast::Expr::BlockExpr(block) => Some(block),
+        match self.branches().next()? {
+            BlockOrStmt::BlockExpr(block) => Some(block),
             _ => None,
         }
     }
 
-    // Hmm. Not sure why this is not `nth(1)`. (It is equivalent to `nth(0)`.)
+    // Return `Some` if the then branch is a single statement.
     pub fn then_branch_stmt(&self) -> Option<ast::Stmt> {
-        support::child(&self.syntax)
+        match self.branches().next()? {
+            BlockOrStmt::Stmt(stmt) => Some(stmt),
+            _ => None,
+        }
     }
 
     // This is the `if` body, corresponding to the condition evaluating true.
@@ -221,15 +233,18 @@ impl ast::IfStmt {
 
     // Return `Some` if the else branch is present and is a curly-delimited block.
     pub fn else_branch_block(&self) -> Option<ast::BlockExpr> {
-        match support::children(self.syntax()).nth(2)? {
-            ast::Expr::BlockExpr(block) => Some(block),
+        match self.branches().nth(1)? {
+            BlockOrStmt::BlockExpr(block) => Some(block),
             _ => None,
         }
     }
 
     // Return `Some` if the else branch is present and is a single statement.
     pub fn else_branch_stmt(&self) -> Option<ast::Stmt> {
-        support::child(&self.syntax)
+        match self.branches().nth(1)? {
+            BlockOrStmt::Stmt(stmt) => Some(stmt),
+            _ => None,
+        }
     }
 
     // This is the `else` body, corresponding to the condition evaluating false.
diff --git a/crates/oq3_syntax/tests/if_stmt_branches.rs b/crates/oq3_syntax/tests/if_stmt_branches.rs
new file mode 100644
index 0000000..7f20d50
--- /dev/null
+++ b/crates/oq3_syntax/tests/if_stmt_branches.rs
@@ -0,0 +1,19 @@
+use oq3_syntax::{ast, AstNode, SourceFile};
+
+#[test]
+fn single_statement_if_branches_are_distinct() {
+    let parse = SourceFile::parse("if (true) x q; else z q;");
+    assert!(parse.errors().is_empty());
+
+    let if_stmt = match parse.tree().statements().next() {
+        Some(ast::Stmt::IfStmt(if_stmt)) => if_stmt,
+        _ => panic!("expected an if statement"),
+    };
+
+    let then_branch = if_stmt.then_branch_stmt().unwrap();
+    let else_branch = if_stmt.else_branch_stmt().unwrap();
+
+    // Regression test: the else accessor previously returned the then branch.
+    assert_eq!(then_branch.syntax().text().to_string(), "x q;");
+    assert_eq!(else_branch.syntax().text().to_string(), "z q;");
+}