Merge pull request #19754 from geoffw0/typeinfer

Rust: Type inference for `for` loops and array expressions

Author: geoffw0

rust/ql/lib/change-notes/2025-06-24-type-inference.md

@@ -0,0 +1,4 @@
+---
+category: minorAnalysis
+---
+* Added type inference for `for` loops and array expressions.

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

@@ -285,6 +285,16 @@ private predicate typeEquality(AstNode n1, TypePath prefix1, AstNode n2, TypePat
       prefix2.isEmpty()
     )
   )
+  or
+  // an array list expression (`[1, 2, 3]`) has the type of the first (any) element
+  n1.(ArrayListExpr).getExpr(_) = n2 and
+  prefix1 = TypePath::singleton(TArrayTypeParameter()) and
+  prefix2.isEmpty()
+  or
+  // an array repeat expression (`[1; 3]`) has the type of the repeat operand
+  n1.(ArrayRepeatExpr).getRepeatOperand() = n2 and
+  prefix1 = TypePath::singleton(TArrayTypeParameter()) and
+  prefix2.isEmpty()
 }
 
 pragma[nomagic]
@@ -1037,6 +1047,12 @@ private class Vec extends Struct {
   }
 }
 
+/**
+ * Gets the root type of the array expression `ae`.
+ */
+pragma[nomagic]
+private Type inferArrayExprType(ArrayExpr ae) { exists(ae) and result = TArrayType() }
+
 /**
  * According to [the Rust reference][1]: _"array and slice-typed expressions
  * can be indexed with a `usize` index ... For other types an index expression
@@ -1073,6 +1089,26 @@ private Type inferIndexExprType(IndexExpr ie, TypePath path) {
   )
 }
 
+pragma[nomagic]
+private Type inferForLoopExprType(AstNode n, TypePath path) {
+  // type of iterable -> type of pattern (loop variable)
+  exists(ForExpr fe, Type iterableType, TypePath iterablePath |
+    n = fe.getPat() and
+    iterableType = inferType(fe.getIterable(), iterablePath) and
+    result = iterableType and
+    (
+      iterablePath.isCons(any(Vec v).getElementTypeParameter(), path)
+      or
+      iterablePath.isCons(any(ArrayTypeParameter tp), path)
+      or
+      iterablePath
+          .stripPrefix(TypePath::cons(TRefTypeParameter(),
+              TypePath::singleton(any(SliceTypeParameter tp)))) = path
+      // TODO: iterables (general case for containers, ranges etc)
+    )
+  )
+}
+
 final class MethodCall extends Call {
   MethodCall() {
     exists(this.getReceiver()) and
@@ -1518,7 +1554,12 @@ private module Cached {
     or
     result = inferAwaitExprType(n, path)
     or
+    result = inferArrayExprType(n) and
+    path.isEmpty()
+    or
     result = inferIndexExprType(n, path)
+    or
+    result = inferForLoopExprType(n, path)
   }
 }
 

rust/ql/test/library-tests/type-inference/CONSISTENCY/PathResolutionConsistency.expected

@@ -1,2 +1,9 @@
 multipleCallTargets
 | dereference.rs:61:15:61:24 | e1.deref() |
+| main.rs:1963:13:1963:31 | ...::from(...) |
+| main.rs:1964:13:1964:31 | ...::from(...) |
+| main.rs:1965:13:1965:31 | ...::from(...) |
+| main.rs:1970:13:1970:31 | ...::from(...) |
+| main.rs:1971:13:1971:31 | ...::from(...) |
+| main.rs:1972:13:1972:31 | ...::from(...) |
+| main.rs:2006:21:2006:43 | ...::from(...) |

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

@@ -1910,11 +1910,7 @@ mod method_determined_by_argument_type {
     impl MyAdd<bool> for i64 {
         // MyAdd<bool>::my_add
         fn my_add(&self, value: bool) -> Self {
-            if value {
-                1
-            } else {
-                0
-            }
+            if value { 1 } else { 0 }
         }
     }
 
@@ -1926,6 +1922,122 @@ mod method_determined_by_argument_type {
     }
 }
 
+mod loops {
+    struct MyCallable {}
+
+    impl MyCallable {
+        fn new() -> Self {
+            MyCallable {}
+        }
+
+        fn call(&self) -> i64 {
+            1
+        }
+    }
+
+    pub fn f() {
+        // for loops with arrays
+
+        for i in [1, 2, 3] {} // $ type=i:i32
+        for i in [1, 2, 3].map(|x| x + 1) {} // $ method=map MISSING: type=i:i32
+        for i in [1, 2, 3].into_iter() {} // $ method=into_iter MISSING: type=i:i32
+
+        let vals1 = [1u8, 2, 3]; // $ type=vals1:[T;...].u8
+        for u in vals1 {} // $ type=u:u8
+
+        let vals2 = [1u16; 3]; // $ type=vals2:[T;...].u16
+        for u in vals2 {} // $ type=u:u16
+
+        let vals3: [u32; 3] = [1, 2, 3]; // $ type=vals3:[T;...].u32
+        for u in vals3 {} // $ type=u:u32
+
+        let vals4: [u64; 3] = [1; 3]; // $ type=vals4:[T;...].u64
+        for u in vals4 {} // $ type=u:u64
+
+        let mut strings1 = ["foo", "bar", "baz"]; // $ type=strings1:[T;...].str
+        for s in &strings1 {} // $ MISSING: type=s:&T.str
+        for s in &mut strings1 {} // $ MISSING: type=s:&T.str
+        for s in strings1 {} // $ type=s:str
+
+        let strings2 = [ // $ type=strings2:[T;...].String
+            String::from("foo"),
+            String::from("bar"),
+            String::from("baz"),
+        ];
+        for s in strings2 {} // $ type=s:String
+
+        let strings3 = &[ // $ type=strings3:&T.[T;...].String
+            String::from("foo"),
+            String::from("bar"),
+            String::from("baz"),
+        ];
+        for s in strings3 {} // $ MISSING: type=s:String
+
+        let callables = [MyCallable::new(), MyCallable::new(), MyCallable::new()]; // $ MISSING: type=callables:[T;...].MyCallable; 3
+        for c in callables // $ type=c:MyCallable
+        {
+            let result = c.call(); // $ type=result:i64 method=call
+        }
+
+        // for loops with ranges
+
+        for i in 0..10 {} // $ MISSING: type=i:i32
+        for u in [0u8..10] {} // $ MISSING: type=u:u8
+        let range = 0..10; // $ MISSING: type=range:Range type=range:Idx.i32
+        for i in range {} // $ MISSING: type=i:i32
+
+        let range1 = std::ops::Range { // $ type=range1:Range type=range1:Idx.u16
+            start: 0u16,
+            end: 10u16,
+        };
+        for u in range1 {} // $ MISSING: type=u:u16
+
+        // for loops with containers
+
+        let vals3 = vec![1, 2, 3]; // $ MISSING: type=vals3:Vec type=vals3:T.i32
+        for i in vals3 {} // $ MISSING: type=i:i32
+
+        let vals4a: Vec<u16> = [1u16, 2, 3].to_vec(); // $ type=vals4a:Vec type=vals4a:T.u16
+        for u in vals4a {} // $ type=u:u16
+
+        let vals4b = [1u16, 2, 3].to_vec(); // $ MISSING: type=vals4b:Vec type=vals4b:T.u16
+        for u in vals4b {} // $ MISSING: type=u:u16
+
+        let vals5 = Vec::from([1u32, 2, 3]); // $ type=vals5:Vec MISSING: type=vals5:T.u32
+        for u in vals5 {} // $ MISSING: type=u:u32
+
+        let vals6: Vec<&u64> = [1u64, 2, 3].iter().collect(); // $ type=vals6:Vec type=vals6:T.&T.u64
+        for u in vals6 {} // $ type=u:&T.u64
+
+        let mut vals7 = Vec::new(); // $ type=vals7:Vec MISSING: type=vals7:T.u8
+        vals7.push(1u8); // $ method=push
+        for u in vals7 {} // $ MISSING: type=u:u8
+
+        let matrix1 = vec![vec![1, 2], vec![3, 4]]; // $ MISSING: type=matrix1:Vec type=matrix1:T.Vec type=matrix1:T.T.i32
+        for row in matrix1 {
+            // $ MISSING: type=row:Vec type=row:T.i32
+            for cell in row { // $ MISSING: type=cell:i32
+            }
+        }
+
+        let mut map1 = std::collections::HashMap::new(); // $ MISSING: type=map1:Hashmap type=map1:K.i32 type=map1:V.Box type1=map1:V.T.&T.str
+        map1.insert(1, Box::new("one")); // $ method=insert
+        map1.insert(2, Box::new("two")); // $ method=insert
+        for key in map1.keys() {} // $ method=keys MISSING: type=key:i32
+        for value in map1.values() {} // $ method=values MISSING: type=value:Box type=value:T.&T.str
+        for (key, value) in map1.iter() {} // $ method=iter MISSING: type=key:i32 type=value:Box type=value:T.&T.str
+        for (key, value) in &map1 {} // $ MISSING: type=key:i32 type=value:Box type=value:T.&T.str
+
+        // while loops
+
+        let mut a: i64 = 0; // $ type=a:i64
+        while a < 10 // $ method=lt type=a:i64
+        {
+            a += 1; // $ type=a:i64 method=add_assign
+        }
+    }
+}
+
 mod dereference;
 
 fn main() {
@@ -1950,6 +2062,7 @@ fn main() {
     async_::f();
     impl_trait::f();
     indexers::f();
+    loops::f();
     macros::f();
     method_determined_by_argument_type::f();
     dereference::test();