Rust: Simple type inference for index expressions

Author: hvitved

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 @@ private module Input1 implements InputSig1<Location> {
   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
@@ -569,7 +577,7 @@ private module CallExprBaseMatchingInput implements MatchingInputSig {
       exists(Param p, int i, boolean inMethod |
         paramPos(this.getParamList(), p, i, inMethod) and
         dpos = TPositionalDeclarationPosition(i, inMethod) and
-        result = inferAnnotatedType(p.getPat(), path)
+        result = p.getTypeRepr().(TypeMention).resolveTypeAt(path)
       )
       or
       exists(SelfParam self |
@@ -1010,6 +1018,36 @@ private StructType inferLiteralType(LiteralExpr le) {
   )
 }
 
+private class Vec extends Struct {
+  Vec() { this.getCanonicalPath() = "alloc::vec::Vec" }
+
+  TypeParamTypeParameter getElementTypeParameter() {
+    result.getTypeParam() = this.getGenericParamList().getTypeParam(0)
+  }
+}
+
+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, TypeParameter tp |
+    TStruct(t) = inferType(ie.getIndex()) and
+    (
+      t instanceof Builtins::I32
+      or
+      t instanceof Builtins::Usize
+    ) and
+    result = inferType(ie.getBase(), exprPath) and
+    exprPath.isCons(tp, path)
+  |
+    tp = any(Vec v).getElementTypeParameter()
+    or
+    tp instanceof ArrayTypeParameter
+    or
+    tp instanceof SliceTypeParameter
+  )
+}
+
 private module MethodCall {
   /** An expression that calls a method. */
   abstract private class MethodCallImpl extends Expr {
@@ -1347,6 +1385,8 @@ private module Cached {
     or
     result = inferLiteralType(n) and
     path.isEmpty()
+    or
+    result = inferIndexExprType(n, path)
   }
 }
 
@@ -1363,7 +1403,7 @@ private module Debug {
     exists(string filepath, int startline, int startcolumn, int endline, int endcolumn |
       result.getLocation().hasLocationInfo(filepath, startline, startcolumn, endline, endcolumn) and
       filepath.matches("%/main.rs") and
-      startline = 948
+      startline = 1718
     )
   }
 

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;