Rust: Implement overloaded index expression in type inference

Author: paldepind

rust/ql/lib/codeql/rust/dataflow/internal/DataFlowImpl.qll

@@ -959,7 +959,11 @@ private module Cached {
 
   cached
   newtype TDataFlowCall =
-    TCall(CallCfgNode c) { Stages::DataFlowStage::ref() } or
+    TCall(CallCfgNode c) {
+      Stages::DataFlowStage::ref() and
+      // TODO: Handle index expressions as calls in data flow.
+      not c.getCall() instanceof IndexExpr
+    } or
     TSummaryCall(
       FlowSummaryImpl::Public::SummarizedCallable c, FlowSummaryImpl::Private::SummaryNode receiver
     ) {

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

@@ -160,4 +160,20 @@ module Impl {
       pos.asPosition() = 0 and result = super.getOperand(1)
     }
   }
+
+  private class IndexCall extends Call instanceof IndexExpr {
+    override string getMethodName() { result = "index" }
+
+    override Trait getTrait() { result.getCanonicalPath() = "core::ops::index::Index" }
+
+    override predicate implicitBorrowAt(ArgumentPosition pos, boolean certain) {
+      pos.isSelf() and certain = true
+    }
+
+    override Expr getArgument(ArgumentPosition pos) {
+      pos.isSelf() and result = super.getBase()
+      or
+      pos.asPosition() = 0 and result = super.getIndex()
+    }
+  }
 }

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

@@ -772,45 +772,48 @@ private Type inferCallExprBaseType(AstNode n, TypePath path) {
     n = a.getNodeAt(apos) and
     result = CallExprBaseMatching::inferAccessType(a, apos, path0)
   |
-    (
+    if
       apos.isBorrowed(true)
       or
-      // The desugaring of the unary `*e` is `*Deref::deref(&e)`. To handle the
-      // deref expression after the call we must strip a `&` from the type at
-      // the return position.
-      apos.isReturn() and a instanceof DerefExpr
-    ) and
-    path0.isCons(TRefTypeParameter(), path)
-    or
-    apos.isBorrowed(false) and
-    exists(Type argType | argType = inferType(n) |
-      if argType = TRefType()
+      // The desugaring of the unary `*e` is `*Deref::deref(&e)` and the
+      // desugaring of `a[b]` is `*Index::index(&a, b)`. To handle the deref
+      // expression after the call we must strip a `&` from the type at the
+      // return position.
+      apos.isReturn() and
+      (a instanceof DerefExpr or a instanceof IndexExpr)
+    then path0.isCons(TRefTypeParameter(), path)
+    else
+      if apos.isBorrowed(false)
       then
-        path = path0 and
-        path0.isCons(TRefTypeParameter(), _)
-        or
-        // adjust for implicit deref
-        not path0.isCons(TRefTypeParameter(), _) and
-        not (path0.isEmpty() and result = TRefType()) and
-        path = TypePath::cons(TRefTypeParameter(), path0)
-      else (
-        not (
-          argType.(StructType).asItemNode() instanceof StringStruct and
-          result.(StructType).asItemNode() instanceof Builtins::Str
-        ) and
-        (
-          not path0.isCons(TRefTypeParameter(), _) and
-          not (path0.isEmpty() and result = TRefType()) and
-          path = path0
-          or
-          // adjust for implicit borrow
-          path0.isCons(TRefTypeParameter(), path)
+        exists(Type argType | argType = inferType(n) |
+          if argType = TRefType()
+          then
+            path = path0 and
+            path0.isCons(TRefTypeParameter(), _)
+            or
+            // adjust for implicit deref
+            not path0.isCons(TRefTypeParameter(), _) and
+            not (path0.isEmpty() and result = TRefType()) and
+            path = TypePath::cons(TRefTypeParameter(), path0)
+          else (
+            not (
+              argType.(StructType).asItemNode() instanceof StringStruct and
+              result.(StructType).asItemNode() instanceof Builtins::Str
+            ) and
+            (
+              not path0.isCons(TRefTypeParameter(), _) and
+              not (path0.isEmpty() and result = TRefType()) and
+              path = path0
+              or
+              // adjust for implicit borrow
+              path0.isCons(TRefTypeParameter(), path)
+            )
+          )
         )
+      else (
+        not apos.isBorrowed(_) and
+        path = path0
       )
-    )
-    or
-    not apos.isBorrowed(_) and
-    path = path0
   )
 }
 
@@ -1116,8 +1119,8 @@ private class Vec extends Struct {
  */
 pragma[nomagic]
 private Type inferIndexExprType(IndexExpr ie, TypePath path) {
-  // TODO: Should be implemented as method resolution, using the special
-  // `std::ops::Index` trait.
+  // TODO: Method resolution to the `std::ops::Index` trait can handle the
+  // `Index` instances for slices and arrays.
   exists(TypePath exprPath, Builtins::BuiltinType t |
     TStruct(t) = inferType(ie.getIndex()) and
     (
@@ -1129,8 +1132,6 @@ private Type inferIndexExprType(IndexExpr ie, TypePath path) {
     ) and
     result = inferType(ie.getBase(), exprPath)
   |
-    exprPath.isCons(any(Vec v).getElementTypeParameter(), path)
-    or
     exprPath.isCons(any(ArrayTypeParameter tp), path)
     or
     exists(TypePath path0 |
@@ -1601,8 +1602,8 @@ private module Debug {
   private Locatable getRelevantLocatable() {
     exists(string filepath, int startline, int startcolumn, int endline, int endcolumn |
       result.getLocation().hasLocationInfo(filepath, startline, startcolumn, endline, endcolumn) and
-      filepath.matches("%/sqlx.rs") and
-      startline = [56 .. 60]
+      filepath.matches("%/main.rs") and
+      startline = [1854 .. 1880]
     )
   }
 

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

@@ -1856,21 +1856,21 @@ mod indexers {
 
         // MyVec::index
         fn index(&self, index: usize) -> &Self::Output {
-            &self.data[index] // $ fieldof=MyVec
+            &self.data[index] // $ fieldof=MyVec method=index
         }
     }
 
     fn analyze_slice(slice: &[S]) {
-        let x = slice[0].foo(); // $ method=foo type=x:S
+        let x = slice[0].foo(); // $ method=foo type=x:S method=index
     }
 
     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
+        vec[0].foo(); // $ method=MyVec::index method=foo
 
         let xs: [S; 1] = [S];
-        let x = xs[0].foo(); // $ method=foo type=x:S
+        let x = xs[0].foo(); // $ method=foo type=x:S method=index
 
         analyze_slice(&xs);
     }