Rust: Improve type inference for closures

Author: hvitved

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

@@ -416,6 +416,14 @@ private predicate isPanicMacroCall(MacroExpr me) {
   me.getMacroCall().resolveMacro().(MacroRules).getName().getText() = "panic"
 }
 
+// Due to "binding modes" the type of the pattern is not necessarily the
+// same as the type of the initializer. The pattern being an identifier
+// pattern is sufficient to ensure that this is not the case.
+private predicate identLetStmt(LetStmt let, IdentPat lhs, Expr rhs) {
+  let.getPat() = lhs and
+  let.getInitializer() = rhs
+}
+
 /** Module for inferring certain type information. */
 module CertainTypeInference {
   pragma[nomagic]
@@ -493,11 +501,7 @@ module CertainTypeInference {
       // is not a certain type equality.
       exists(LetStmt let |
         not let.hasTypeRepr() and
-        // Due to "binding modes" the type of the pattern is not necessarily the
-        // same as the type of the initializer. The pattern being an identifier
-        // pattern is sufficient to ensure that this is not the case.
-        let.getPat().(IdentPat) = n1 and
-        let.getInitializer() = n2
+        identLetStmt(let, n1, n2)
       )
       or
       exists(LetExpr let |
@@ -521,6 +525,25 @@ module CertainTypeInference {
           )
         else prefix2.isEmpty()
       )
+    or
+    exists(CallExprImpl::DynamicCallExpr dce, TupleType tt, int i |
+      n1 = dce.getArgList() and
+      tt.getArity() = dce.getNumberOfSyntacticArguments() and
+      n2 = dce.getSyntacticPositionalArgument(i) and
+      prefix1 = TypePath::singleton(tt.getPositionalTypeParameter(i)) and
+      prefix2.isEmpty()
+    )
+    or
+    exists(ClosureExpr ce, int index |
+      n1 = ce and
+      n2 = ce.getParam(index).getPat() and
+      prefix1 = closureParameterPath(ce.getNumberOfParams(), index) and
+      prefix2.isEmpty()
+    )
+    or
+    n1 = any(ClosureExpr ce | not ce.hasRetType() and ce.getClosureBody() = n2) and
+    prefix1 = closureReturnPath() and
+    prefix2.isEmpty()
   }
 
   pragma[nomagic]
@@ -783,17 +806,6 @@ private predicate typeEquality(AstNode n1, TypePath prefix1, AstNode n2, TypePat
   n1.(ArrayRepeatExpr).getRepeatOperand() = n2 and
   prefix1 = TypePath::singleton(getArrayTypeParameter()) and
   prefix2.isEmpty()
-  or
-  exists(ClosureExpr ce, int index |
-    n1 = ce and
-    n2 = ce.getParam(index).getPat() and
-    prefix1 = closureParameterPath(ce.getNumberOfParams(), index) and
-    prefix2.isEmpty()
-  )
-  or
-  n1.(ClosureExpr).getClosureBody() = n2 and
-  prefix1 = closureReturnPath() and
-  prefix2.isEmpty()
 }
 
 /**
@@ -836,6 +848,19 @@ private predicate lubCoercion(AstNode parent, AstNode child, TypePath prefix) {
   )
 }
 
+private Type inferUnknownTypeFromAnnotation(AstNode n, TypePath path) {
+  inferType(n, path) = TUnknownType() and
+  // Normally, these are coercion sites, but in case a type is unknown we
+  // allow for type information to flow from the type annotation.
+  exists(TypeMention tm | result = tm.getTypeAt(path) |
+    tm = any(LetStmt let | identLetStmt(let, _, n)).getTypeRepr()
+    or
+    tm = any(ClosureExpr ce | n = ce.getBody()).getRetType().getTypeRepr()
+    or
+    tm = getReturnTypeMention(any(Function f | n = f.getBody()))
+  )
+}
+
 /**
  * Holds if the type tree of `n1` at `prefix1` should be equal to the type tree
  * of `n2` at `prefix2`, but type information should only propagate from `n1` to
@@ -1685,6 +1710,8 @@ private module AssocFunctionResolution {
    * 3. `AssocFunctionCallCallExpr`: a qualified function call, `Q::f(x)`; or
    * 4. `AssocFunctionCallOperation`: an operation expression, `x + y`, which is syntactic sugar
    *    for `Add::add(x, y)`.
+   * 5. `ClosureMethodCall`: a call to a closure, `c(x)`, which is syntactic sugar for
+   *    `c.call_once(x)`, `c.call_mut(x)`, or `c.call(x)`.
    *
    * Note that only in case 1 and 2 is auto-dereferencing and borrowing allowed.
    *
@@ -1701,7 +1728,7 @@ private module AssocFunctionResolution {
     pragma[nomagic]
     abstract predicate hasNameAndArity(string name, int arity);
 
-    abstract Expr getNonReturnNodeAt(FunctionPosition pos);
+    abstract AstNode getNonReturnNodeAt(FunctionPosition pos);
 
     AstNode getNodeAt(FunctionPosition pos) {
       result = this.getNonReturnNodeAt(pos)
@@ -2235,7 +2262,7 @@ private module AssocFunctionResolution {
     }
   }
 
-  private class AssocFunctionCallMethodCallExpr extends AssocFunctionCall instanceof MethodCallExpr {
+  private class MethodCallExprAssocFunctionCall extends AssocFunctionCall instanceof MethodCallExpr {
     override predicate hasNameAndArity(string name, int arity) {
       name = super.getIdentifier().getText() and
       arity = super.getNumberOfSyntacticArguments()
@@ -2255,7 +2282,7 @@ private module AssocFunctionResolution {
     override Trait getTrait() { none() }
   }
 
-  private class AssocFunctionCallIndexExpr extends AssocFunctionCall, IndexExpr {
+  private class IndexExprAssocFunctionCall extends AssocFunctionCall, IndexExpr {
     private predicate isInMutableContext() {
       // todo: does not handle all cases yet
       VariableImpl::assignmentOperationDescendant(_, this)
@@ -2285,8 +2312,8 @@ private module AssocFunctionResolution {
     }
   }
 
-  private class AssocFunctionCallCallExpr extends AssocFunctionCall, CallExpr {
-    AssocFunctionCallCallExpr() {
+  private class CallExprAssocFunctionCall extends AssocFunctionCall, CallExpr {
+    CallExprAssocFunctionCall() {
       exists(getCallExprPathQualifier(this)) and
       // even if a target cannot be resolved by path resolution, it may still
       // be possible to resolve a blanket implementation (so not `forex`)
@@ -2318,7 +2345,7 @@ private module AssocFunctionResolution {
     override Trait getTrait() { result = getCallExprTraitQualifier(this) }
   }
 
-  final class AssocFunctionCallOperation extends AssocFunctionCall, Operation {
+  final class OperationAssocFunctionCall extends AssocFunctionCall, Operation {
     override predicate hasNameAndArity(string name, int arity) {
       this.isOverloaded(_, name, _) and
       arity = this.getNumberOfOperands()
@@ -2376,6 +2403,29 @@ private module AssocFunctionResolution {
     override Trait getTrait() { this.isOverloaded(result, _, _) }
   }
 
+  private class DynamicAssocFunctionCall extends AssocFunctionCall instanceof CallExprImpl::DynamicCallExpr
+  {
+    pragma[nomagic]
+    override predicate hasNameAndArity(string name, int arity) {
+      name = "call_once" and // todo: handle call_mut and call
+      arity = 2 // args are passed in a tuple
+    }
+
+    override predicate hasReceiver() { any() }
+
+    override AstNode getNonReturnNodeAt(FunctionPosition pos) {
+      pos.asPosition() = 0 and
+      result = super.getFunction()
+      or
+      pos.asPosition() = 1 and
+      result = super.getArgList()
+    }
+
+    override predicate supportsAutoDerefAndBorrow() { any() }
+
+    override Trait getTrait() { result instanceof AnyFnTrait }
+  }
+
   pragma[nomagic]
   private AssocFunctionDeclaration getAssocFunctionSuccessor(
     ImplOrTraitItemNode i, string name, int arity
@@ -3239,7 +3289,7 @@ private module OperationMatchingInput implements MatchingInputSig {
     }
   }
 
-  class Access extends AssocFunctionResolution::AssocFunctionCallOperation {
+  class Access extends AssocFunctionResolution::OperationAssocFunctionCall {
     Type getTypeArgument(TypeArgumentPosition apos, TypePath path) { none() }
 
     pragma[nomagic]
@@ -3844,14 +3894,6 @@ private module InvokedClosureSatisfiesTypeInput implements SatisfiesTypeInputSig
   }
 }
 
-private module InvokedClosureSatisfiesType =
-  SatisfiesType<InvokedClosureExpr, InvokedClosureSatisfiesTypeInput>;
-
-/** Gets the type of `ce` when viewed as an implementation of `FnOnce`. */
-private Type invokedClosureFnTypeAt(InvokedClosureExpr ce, TypePath path) {
-  InvokedClosureSatisfiesType::satisfiesConstraintType(ce, _, path, result)
-}
-
 /**
  * Gets the root type of a closure.
  *
@@ -3878,73 +3920,39 @@ private TypePath closureParameterPath(int arity, int index) {
       TypePath::singleton(getTupleTypeParameter(arity, index)))
 }
 
-/** Gets the path to the return type of the `FnOnce` trait. */
-private TypePath fnReturnPath() {
-  result = TypePath::singleton(getAssociatedTypeTypeParameter(any(FnOnceTrait t).getOutputType()))
-}
-
-/**
- * Gets the path to the parameter type of the `FnOnce` trait with arity `arity`
- * and index `index`.
- */
 pragma[nomagic]
-private TypePath fnParameterPath(int arity, int index) {
-  result =
-    TypePath::cons(TTypeParamTypeParameter(any(FnOnceTrait t).getTypeParam()),
-      TypePath::singleton(getTupleTypeParameter(arity, index)))
-}
-
-pragma[nomagic]
-private Type inferDynamicCallExprType(Expr n, TypePath path) {
-  exists(InvokedClosureExpr ce |
-    // Propagate the function's return type to the call expression
-    exists(TypePath path0 | result = invokedClosureFnTypeAt(ce, path0) |
-      n = ce.getCall() and
-      path = path0.stripPrefix(fnReturnPath())
+private Type inferClosureExprType(AstNode n, TypePath path) {
+  exists(ClosureExpr ce |
+    n = ce and
+    (
+      path.isEmpty() and
+      result = closureRootType()
+      or
+      path = TypePath::singleton(TDynTraitTypeParameter(_, any(FnTrait t).getTypeParam())) and
+      result.(TupleType).getArity() = ce.getNumberOfParams()
       or
-      // Propagate the function's parameter type to the arguments
-      exists(int index |
-        n = ce.getCall().getSyntacticPositionalArgument(index) and
-        path =
-          path0.stripPrefix(fnParameterPath(ce.getCall().getArgList().getNumberOfArgs(), index))
+      exists(TypePath path0 |
+        result = ce.getRetType().getTypeRepr().(TypeMention).getTypeAt(path0) and
+        path = closureReturnPath().append(path0)
       )
     )
     or
-    // _If_ the invoked expression has the type of a closure, then we propagate
-    // the surrounding types into the closure.
-    exists(int arity, TypePath path0 | ce.getTypeAt(TypePath::nil()) = closureRootType() |
-      // Propagate the type of arguments to the parameter types of closure
-      exists(int index, ArgList args |
-        n = ce and
-        args = ce.getCall().getArgList() and
-        arity = args.getNumberOfArgs() and
-        result = inferType(args.getArg(index), path0) and
-        path = closureParameterPath(arity, index).append(path0)
-      )
-      or
-      // Propagate the type of the call expression to the return type of the closure
-      n = ce and
-      arity = ce.getCall().getArgList().getNumberOfArgs() and
-      result = inferType(ce.getCall(), path0) and
-      path = closureReturnPath().append(path0)
+    exists(Param p |
+      p = ce.getAParam() and
+      not p.hasTypeRepr() and
+      n = p.getPat() and
+      result = TUnknownType() and
+      path.isEmpty()
     )
   )
 }
 
 pragma[nomagic]
-private Type inferClosureExprType(AstNode n, TypePath path) {
-  exists(ClosureExpr ce |
-    n = ce and
-    path.isEmpty() and
-    result = closureRootType()
-    or
-    n = ce and
-    path = TypePath::singleton(TDynTraitTypeParameter(_, any(FnTrait t).getTypeParam())) and
-    result.(TupleType).getArity() = ce.getNumberOfParams()
-    or
-    // Propagate return type annotation to body
-    n = ce.getClosureBody() and
-    result = ce.getRetType().getTypeRepr().(TypeMention).getTypeAt(path)
+private TupleType inferArgList(ArgList args, TypePath path) {
+  exists(CallExprImpl::DynamicCallExpr dce |
+    args = dce.getArgList() and
+    result.getArity() = dce.getNumberOfSyntacticArguments() and
+    path.isEmpty()
   )
 }
 
@@ -3992,7 +4000,8 @@ private module Cached {
       or
       i instanceof ImplItemNode and dispatch = false
     |
-      result = call.(AssocFunctionResolution::AssocFunctionCall).resolveCallTarget(i, _, _, _)
+      result = call.(AssocFunctionResolution::AssocFunctionCall).resolveCallTarget(i, _, _, _) and
+      not call instanceof CallExprImpl::DynamicCallExpr // todo
     )
   }
 
@@ -4101,13 +4110,15 @@ private module Cached {
       or
       result = inferForLoopExprType(n, path)
       or
-      result = inferDynamicCallExprType(n, path)
-      or
       result = inferClosureExprType(n, path)
       or
+      result = inferArgList(n, path)
+      or
       result = inferStructPatType(n, path)
       or
       result = inferTupleStructPatType(n, path)
+      or
+      result = inferUnknownTypeFromAnnotation(n, path)
     )
   }
 }
@@ -4124,8 +4135,8 @@ private module Debug {
   Locatable getRelevantLocatable() {
     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 = 103
+      filepath.matches("%/regressions.rs") and
+      startline = 24
     )
   }