Merge pull request #19026 from paldepind/rust-expr-type-eq

Rust: Handle type equality for a few more expression types

Author: paldepind

rust/ql/lib/codeql/rust/controlflow/internal/Completion.qll

@@ -90,7 +90,7 @@ class BooleanCompletion extends ConditionalCompletion, TBooleanCompletion {
       or
       e = parent.(BinaryLogicalOperation).getRhs()
       or
-      parent = any(IfExpr ie | e = [ie.getThen(), ie.getElse()])
+      e = parent.(IfExpr).getABranch()
       or
       e = parent.(MatchExpr).getAnArm().getExpr()
       or

rust/ql/lib/codeql/rust/controlflow/internal/ControlFlowGraphImpl.qll

@@ -336,7 +336,7 @@ module ExprTrees {
     override predicate first(AstNode node) { first(super.getCondition(), node) }
 
     override predicate propagatesAbnormal(AstNode child) {
-      child = [super.getCondition(), super.getThen(), super.getElse()]
+      child = [super.getCondition(), super.getABranch()]
     }
 
     private ConditionalCompletion conditionCompletion(Completion c) {

rust/ql/lib/codeql/rust/controlflow/internal/Splitting.qll

@@ -81,7 +81,7 @@ module ConditionalCompletionSplitting {
       (
         child = parent.(BinaryLogicalOperation).getAnOperand()
         or
-        parent = any(IfExpr ie | child = [ie.getThen(), ie.getElse()])
+        child = parent.(IfExpr).getABranch()
         or
         child = parent.(MatchExpr).getAnArm().getExpr()
         or

rust/ql/lib/codeql/rust/elements/internal/IfExprImpl.qll

@@ -39,5 +39,8 @@ module Impl {
       or
       index = 3 and this.hasElse() and result = "else {...}"
     }
+
+    /** Gets any of the branches of this if expression. */
+    Expr getABranch() { result = [this.getThen(), this.getElse()] }
   }
 }

rust/ql/lib/codeql/rust/internal/TypeInference.qll

@@ -158,6 +158,18 @@ private predicate typeEquality(AstNode n1, TypePath path1, AstNode n2, TypePath
   or
   n1 = n2.(BlockExpr).getStmtList().getTailExpr() and
   path1 = path2
+  or
+  n1 = n2.(IfExpr).getABranch() and
+  path1 = path2
+  or
+  n1 = n2.(MatchExpr).getAnArm().getExpr() and
+  path1 = path2
+  or
+  exists(BreakExpr break |
+    break.getExpr() = n1 and
+    break.getTarget() = n2.(LoopExpr) and
+    path1 = path2
+  )
 }
 
 pragma[nomagic]

rust/ql/test/library-tests/type-inference/main.rs

@@ -225,8 +225,8 @@ mod m6 {
     impl<T> MyEnum<T> {
         fn m1(self) -> T {
             match self {
-                MyEnum::C1(a) => a,    // missing
-                MyEnum::C2 { a } => a, // missing
+                MyEnum::C1(a) => a,
+                MyEnum::C2 { a } => a,
             }
         }
     }
@@ -402,8 +402,8 @@ mod m9 {
     impl<T> MyOption<MyOption<T>> {
         fn flatten(self) -> MyOption<T> {
             match self {
-                MyOption::MyNone() => MyOption::MyNone(), // missing inner type `Option<T>`
-                MyOption::MySome(x) => x,                 // missing
+                MyOption::MyNone() => MyOption::MyNone(),
+                MyOption::MySome(x) => x,
             }
         }
     }
@@ -432,6 +432,27 @@ mod m9 {
 
         let x6 = MyOption::MySome(MyOption::<S>::MyNone());
         println!("{:?}", MyOption::<MyOption<S>>::flatten(x6));
+
+        let from_if = if 1 + 1 > 2 {
+            MyOption::MyNone()
+        } else {
+            MyOption::MySome(S)
+        };
+        println!("{:?}", from_if);
+
+        let from_match = match 1 + 1 > 2 {
+            true => MyOption::MyNone(),
+            false => MyOption::MySome(S),
+        };
+        println!("{:?}", from_match);
+
+        let from_loop = loop {
+            if 1 + 1 > 2 {
+                break MyOption::MyNone();
+            }
+            break MyOption::MySome(S);
+        };
+        println!("{:?}", from_loop);
     }
 }