Rust: Simple type inference for index expressions

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

@@ -15,10 +15,13 @@ newtype TType =
   TTrait(Trait t) or
   TArrayType() or // todo: add size?
   TRefType() or // todo: add mut?
+  TSliceType() or
   TTypeParamTypeParameter(TypeParam t) or
   TAssociatedTypeTypeParameter(TypeAlias t) { any(TraitItemNode trait).getAnAssocItem() = t } or
+  TArrayTypeParameter() or
   TRefTypeParameter() or
-  TSelfTypeParameter(Trait t)
+  TSelfTypeParameter(Trait t) or
+  TSliceTypeParameter()
 
 /**
  * A type without type arguments.
@@ -145,7 +148,8 @@ class ArrayType extends Type, TArrayType {
   override TupleField getTupleField(int i) { none() }
 
   override TypeParameter getTypeParameter(int i) {
-    none() // todo
+    result = TArrayTypeParameter() and
+    i = 0
   }
 
   override string toString() { result = "[]" }
@@ -176,6 +180,29 @@ class RefType extends Type, TRefType {
   override Location getLocation() { result instanceof EmptyLocation }
 }
 
+/**
+ * A slice type.
+ *
+ * Slice types like `[i64]` are modeled as normal generic types
+ * with a single type argument.
+ */
+class SliceType extends Type, TSliceType {
+  SliceType() { this = TSliceType() }
+
+  override StructField getStructField(string name) { none() }
+
+  override TupleField getTupleField(int i) { none() }
+
+  override TypeParameter getTypeParameter(int i) {
+    result = TSliceTypeParameter() and
+    i = 0
+  }
+
+  override string toString() { result = "[]" }
+
+  override Location getLocation() { result instanceof EmptyLocation }
+}
+
 /** A type parameter. */
 abstract class TypeParameter extends Type {
   override StructField getStructField(string name) { none() }
@@ -255,13 +282,27 @@ class AssociatedTypeTypeParameter extends TypeParameter, TAssociatedTypeTypePara
   override Location getLocation() { result = typeAlias.getLocation() }
 }
 
+/** An implicit array type parameter. */
+class ArrayTypeParameter extends TypeParameter, TArrayTypeParameter {
+  override string toString() { result = "[T;...]" }
+
+  override Location getLocation() { result instanceof EmptyLocation }
+}
+
 /** An implicit reference type parameter. */
 class RefTypeParameter extends TypeParameter, TRefTypeParameter {
   override string toString() { result = "&T" }
 
   override Location getLocation() { result instanceof EmptyLocation }
 }
 
+/** An implicit slice type parameter. */
+class SliceTypeParameter extends TypeParameter, TSliceTypeParameter {
+  override string toString() { result = "[T]" }
+
+  override Location getLocation() { result instanceof EmptyLocation }
+}
+
 /**
  * The implicit `Self` type parameter of a trait, that refers to the
  * implementing type of the trait.

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

@@ -80,10 +80,18 @@
   int getTypeParameterId(TypeParameter tp) {
     tp =
       rank[result](TypeParameter tp0, int kind, int id |
-        tp0 instanceof RefTypeParameter and
+        tp0 instanceof ArrayTypeParameter and
         kind = 0 and
         id = 0
         or
+        tp0 instanceof RefTypeParameter and
+        kind = 0 and
+        id = 1
+        or
+        tp0 instanceof SliceTypeParameter and
+        kind = 0 and
+        id = 2
+        or
         kind = 1 and
         exists(AstNode node | id = idOfTypeParameterAstNode(node) |
           node = tp0.(TypeParamTypeParameter).getTypeParam() or
@@ -1010,6 +1018,50 @@
   )
 }
 
+private class Vec extends Struct {
+  Vec() { this.getCanonicalPath() = "alloc::vec::Vec" }
+
+  TypeParamTypeParameter getElementTypeParameter() {
+    result.getTypeParam() = this.getGenericParamList().getTypeParam(0)
+  }
+}
+
+/**
+ * According to [the Rust reference][1]: _"array and slice-typed expressions
+ * can be indexed with a `usize` index ... For other types an index expression
+ * `a[b]` is equivalent to *std::ops::Index::index(&a, b)"_.
+ *
+ * The logic below handles array and slice indexing, but for other types it is
+ * currently limited to `Vec`.
+ *
+ * [1]: https://doc.rust-lang.org/reference/expressions/array-expr.html#r-expr.array.index
+ */
+pragma[nomagic]
+private Type inferIndexExprType(IndexExpr ie, TypePath path) {
+  // TODO: Should be implemented as method resolution, using the special
+  // `std::ops::Index` trait.
+  exists(TypePath exprPath, Builtins::BuiltinType t |
+    TStruct(t) = inferType(ie.getIndex()) and
+    (
+      // also allow `i32`, since that is currently the type that we infer for
+      // integer literals like `0`
+      t instanceof Builtins::I32
+      or
+      t instanceof Builtins::Usize
+    ) and
+    result = inferType(ie.getBase(), exprPath)
+  |
+    exprPath.isCons(any(Vec v).getElementTypeParameter(), path)
+    or
+    exprPath.isCons(any(ArrayTypeParameter tp), path)
+    or
+    exists(TypePath path0 |
+      exprPath.isCons(any(RefTypeParameter tp), path0) and
+      path0.isCons(any(SliceTypeParameter tp), path)
+    )
+  )
+}
+
 private module MethodCall {
   /** An expression that calls a method. */
   abstract private class MethodCallImpl extends Expr {
@@ -1347,6 +1399,8 @@
     or
     result = inferLiteralType(n) and
     path.isEmpty()
+    or
+    result = inferIndexExprType(n, path)
   }
 }
 

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

@@ -43,6 +43,12 @@ class RefTypeReprMention extends TypeMention instanceof RefTypeRepr {
   override Type resolveType() { result = TRefType() }
 }
 
+class SliceTypeReprMention extends TypeMention instanceof SliceTypeRepr {
+  override TypeMention getTypeArgument(int i) { result = super.getTypeRepr() and i = 0 }
+
+  override Type resolveType() { result = TSliceType() }
+}
+
 class PathTypeReprMention extends TypeMention instanceof PathTypeRepr {
   Path path;
   ItemNode resolved;

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

@@ -1630,6 +1630,58 @@ mod overloadable_operators {
     }
 }
 
+mod indexers {
+    use std::ops::Index;
+
+    #[derive(Debug)]
+    struct S;
+
+    impl S {
+        fn foo(&self) -> Self {
+            S
+        }
+    }
+
+    #[derive(Debug)]
+    struct MyVec<T> {
+        data: Vec<T>,
+    }
+
+    impl<T> MyVec<T> {
+        fn new() -> Self {
+            MyVec { data: Vec::new() }
+        }
+
+        fn push(&mut self, value: T) {
+            self.data.push(value); // $ fieldof=MyVec
+        }
+    }
+
+    impl<T> Index<usize> for MyVec<T> {
+        type Output = T;
+
+        // MyVec::index
+        fn index(&self, index: usize) -> &Self::Output {
+            &self.data[index] // $ fieldof=MyVec
+        }
+    }
+
+    fn analyze_slice(slice: &[S]) {
+        let x = slice[0].foo(); // $ method=foo type=x:S
+    }
+
+    pub fn f() {
+        let mut vec = MyVec::new(); // $ type=vec:T.S
+        vec.push(S); // $ method=push
+        vec[0].foo(); // $ MISSING: method=foo -- type inference does not support the `Index` trait yet
+
+        let xs: [S; 1] = [S];
+        let x = xs[0].foo(); // $ method=foo type=x:S
+
+        analyze_slice(&xs);
+    }
+}
+
 fn main() {
     field_access::f();
     method_impl::f();
@@ -1649,4 +1701,5 @@ fn main() {
     try_expressions::f();
     builtins::f();
     operators::f();
+    indexers::f();
 }

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

@@ -2375,7 +2375,82 @@ inferType
 | main.rs:1629:13:1629:20 | vec2_not |  | main.rs:1278:5:1283:5 | Vec2 |
 | main.rs:1629:24:1629:26 | ! ... |  | main.rs:1278:5:1283:5 | Vec2 |
 | main.rs:1629:25:1629:26 | v1 |  | main.rs:1278:5:1283:5 | Vec2 |
-| main.rs:1635:5:1635:20 | ...::f(...) |  | main.rs:67:5:67:21 | Foo |
-| main.rs:1636:5:1636:60 | ...::g(...) |  | main.rs:67:5:67:21 | Foo |
-| main.rs:1636:20:1636:38 | ...::Foo {...} |  | main.rs:67:5:67:21 | Foo |
-| main.rs:1636:41:1636:59 | ...::Foo {...} |  | main.rs:67:5:67:21 | Foo |
+| main.rs:1640:16:1640:20 | SelfParam |  | file://:0:0:0:0 | & |
+| main.rs:1640:16:1640:20 | SelfParam | &T | main.rs:1636:5:1637:13 | S |
+| main.rs:1640:31:1642:9 | { ... } |  | main.rs:1636:5:1637:13 | S |
+| main.rs:1641:13:1641:13 | S |  | main.rs:1636:5:1637:13 | S |
+| main.rs:1651:26:1653:9 | { ... } |  | main.rs:1645:5:1648:5 | MyVec |
+| main.rs:1651:26:1653:9 | { ... } | T | main.rs:1650:10:1650:10 | T |
+| main.rs:1652:13:1652:38 | MyVec {...} |  | main.rs:1645:5:1648:5 | MyVec |
+| main.rs:1652:13:1652:38 | MyVec {...} | T | main.rs:1650:10:1650:10 | T |
+| main.rs:1655:17:1655:25 | SelfParam |  | file://:0:0:0:0 | & |
+| main.rs:1655:17:1655:25 | SelfParam | &T | main.rs:1645:5:1648:5 | MyVec |
+| main.rs:1655:17:1655:25 | SelfParam | &T.T | main.rs:1650:10:1650:10 | T |
+| main.rs:1655:28:1655:32 | value |  | main.rs:1650:10:1650:10 | T |
+| main.rs:1656:13:1656:16 | self |  | file://:0:0:0:0 | & |
+| main.rs:1656:13:1656:16 | self | &T | main.rs:1645:5:1648:5 | MyVec |
+| main.rs:1656:13:1656:16 | self | &T.T | main.rs:1650:10:1650:10 | T |
+| main.rs:1656:28:1656:32 | value |  | main.rs:1650:10:1650:10 | T |
+| main.rs:1664:18:1664:22 | SelfParam |  | file://:0:0:0:0 | & |
+| main.rs:1664:18:1664:22 | SelfParam | &T | main.rs:1645:5:1648:5 | MyVec |
+| main.rs:1664:18:1664:22 | SelfParam | &T.T | main.rs:1660:10:1660:10 | T |
+| main.rs:1664:25:1664:29 | index |  | file:///BUILTINS/types.rs:20:1:21:17 | usize |
+| main.rs:1664:56:1666:9 | { ... } |  | file://:0:0:0:0 | & |
+| main.rs:1664:56:1666:9 | { ... } | &T | main.rs:1660:10:1660:10 | T |
+| main.rs:1665:13:1665:29 | &... |  | file://:0:0:0:0 | & |
+| main.rs:1665:13:1665:29 | &... | &T | main.rs:1660:10:1660:10 | T |
+| main.rs:1665:14:1665:17 | self |  | file://:0:0:0:0 | & |
+| main.rs:1665:14:1665:17 | self | &T | main.rs:1645:5:1648:5 | MyVec |
+| main.rs:1665:14:1665:17 | self | &T.T | main.rs:1660:10:1660:10 | T |
+| main.rs:1665:24:1665:28 | index |  | file:///BUILTINS/types.rs:20:1:21:17 | usize |
+| main.rs:1669:22:1669:26 | slice |  | file://:0:0:0:0 | & |
+| main.rs:1669:22:1669:26 | slice | &T | file://:0:0:0:0 | [] |
+| main.rs:1669:22:1669:26 | slice | &T.[T] | main.rs:1636:5:1637:13 | S |
+| main.rs:1670:13:1670:13 | x |  | main.rs:1636:5:1637:13 | S |
+| main.rs:1670:17:1670:21 | slice |  | file://:0:0:0:0 | & |
+| main.rs:1670:17:1670:21 | slice | &T | file://:0:0:0:0 | [] |
+| main.rs:1670:17:1670:21 | slice | &T.[T] | main.rs:1636:5:1637:13 | S |
+| main.rs:1670:17:1670:24 | slice[0] |  | main.rs:1636:5:1637:13 | S |
+| main.rs:1670:17:1670:30 | ... .foo() |  | main.rs:1636:5:1637:13 | S |
+| main.rs:1670:23:1670:23 | 0 |  | file:///BUILTINS/types.rs:12:1:12:15 | i32 |
+| main.rs:1674:13:1674:19 | mut vec |  | main.rs:1645:5:1648:5 | MyVec |
+| main.rs:1674:13:1674:19 | mut vec | T | main.rs:1636:5:1637:13 | S |
+| main.rs:1674:23:1674:34 | ...::new(...) |  | main.rs:1645:5:1648:5 | MyVec |
+| main.rs:1674:23:1674:34 | ...::new(...) | T | main.rs:1636:5:1637:13 | S |
+| main.rs:1675:9:1675:11 | vec |  | main.rs:1645:5:1648:5 | MyVec |
+| main.rs:1675:9:1675:11 | vec | T | main.rs:1636:5:1637:13 | S |
+| main.rs:1675:18:1675:18 | S |  | main.rs:1636:5:1637:13 | S |
+| main.rs:1676:9:1676:11 | vec |  | main.rs:1645:5:1648:5 | MyVec |
+| main.rs:1676:9:1676:11 | vec | T | main.rs:1636:5:1637:13 | S |
+| main.rs:1676:13:1676:13 | 0 |  | file:///BUILTINS/types.rs:12:1:12:15 | i32 |
+| main.rs:1678:13:1678:14 | xs |  | file://:0:0:0:0 | [] |
+| main.rs:1678:13:1678:14 | xs |  | file://:0:0:0:0 | [] |
+| main.rs:1678:13:1678:14 | xs | [T;...] | main.rs:1636:5:1637:13 | S |
+| main.rs:1678:13:1678:14 | xs | [T] | main.rs:1636:5:1637:13 | S |
+| main.rs:1678:21:1678:21 | 1 |  | file:///BUILTINS/types.rs:12:1:12:15 | i32 |
+| main.rs:1678:26:1678:28 | [...] |  | file://:0:0:0:0 | [] |
+| main.rs:1678:26:1678:28 | [...] |  | file://:0:0:0:0 | [] |
+| main.rs:1678:26:1678:28 | [...] | [T;...] | main.rs:1636:5:1637:13 | S |
+| main.rs:1678:26:1678:28 | [...] | [T] | main.rs:1636:5:1637:13 | S |
+| main.rs:1678:27:1678:27 | S |  | main.rs:1636:5:1637:13 | S |
+| main.rs:1679:13:1679:13 | x |  | main.rs:1636:5:1637:13 | S |
+| main.rs:1679:17:1679:18 | xs |  | file://:0:0:0:0 | [] |
+| main.rs:1679:17:1679:18 | xs |  | file://:0:0:0:0 | [] |
+| main.rs:1679:17:1679:18 | xs | [T;...] | main.rs:1636:5:1637:13 | S |
+| main.rs:1679:17:1679:18 | xs | [T] | main.rs:1636:5:1637:13 | S |
+| main.rs:1679:17:1679:21 | xs[0] |  | main.rs:1636:5:1637:13 | S |
+| main.rs:1679:17:1679:27 | ... .foo() |  | main.rs:1636:5:1637:13 | S |
+| main.rs:1679:20:1679:20 | 0 |  | file:///BUILTINS/types.rs:12:1:12:15 | i32 |
+| main.rs:1681:23:1681:25 | &xs |  | file://:0:0:0:0 | & |
+| main.rs:1681:23:1681:25 | &xs | &T | file://:0:0:0:0 | [] |
+| main.rs:1681:23:1681:25 | &xs | &T | file://:0:0:0:0 | [] |
+| main.rs:1681:23:1681:25 | &xs | &T.[T;...] | main.rs:1636:5:1637:13 | S |
+| main.rs:1681:23:1681:25 | &xs | &T.[T] | main.rs:1636:5:1637:13 | S |
+| main.rs:1681:24:1681:25 | xs |  | file://:0:0:0:0 | [] |
+| main.rs:1681:24:1681:25 | xs |  | file://:0:0:0:0 | [] |
+| main.rs:1681:24:1681:25 | xs | [T;...] | main.rs:1636:5:1637:13 | S |
+| main.rs:1681:24:1681:25 | xs | [T] | main.rs:1636:5:1637:13 | S |
+| main.rs:1687:5:1687:20 | ...::f(...) |  | main.rs:67:5:67:21 | Foo |
+| main.rs:1688:5:1688:60 | ...::g(...) |  | main.rs:67:5:67:21 | Foo |
+| main.rs:1688:20:1688:38 | ...::Foo {...} |  | main.rs:67:5:67:21 | Foo |
+| main.rs:1688:41:1688:59 | ...::Foo {...} |  | main.rs:67:5:67:21 | Foo |

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

@@ -70,10 +70,13 @@ module TypeTest implements TestSig {
     exists(AstNode n, TypePath path, Type t |
       t = TypeInference::inferType(n, path) and
       location = n.getLocation() and
-      element = n.toString() and
       if path.isEmpty()
       then value = element + ":" + t
       else value = element + ":" + path.toString() + "." + t.toString()
+    |
+      element = n.toString()
+      or
+      element = n.(IdentPat).getName().getText()
     )
   }
 }