Rust: Generalize shared type inference logic

Author: hvitved

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

@@ -2917,7 +2917,9 @@ private module Input3 implements InputSig3 {
     )
   }
 
-  class Closure extends Expr, Callable instanceof Rust::ClosureExpr { }
+  class Closure extends Expr, Callable instanceof Rust::ClosureExpr {
+    Expr getDefiningExpr() { result = this }
+  }
 
   class ClosureParameterPseudoType extends T::ClosureParameterPseudoType {
     Parameter getParameter() { result = this.getParam() }

shared/typeinference/codeql/typeinference/internal/TypeInference.qll

@@ -79,7 +79,14 @@ signature module InputSig1<LocationSig Location> {
    * For example `int` or ``IEnumerable`1``.
    */
   class Type {
-    /** Gets a type parameter of this type, if any. */
+    /** Gets the `i`th positional type parameter of this type, if any. */
+    TypeParameter getPositionalTypeParameter(int i);
+
+    /**
+     * Gets a type parameter of this type, if any.
+     *
+     * This may include non-positional type parameters as well.
+     */
     TypeParameter getATypeParameter();
 
     /** Gets a textual representation of this type. */
@@ -2041,7 +2048,7 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
        * item in Rust.
        */
       class Variable {
-        /** Gets the AST node that defines this variable. */
+        /** Gets the AST node that defines this variable, if any. */
         AstNode getDefiningNode();
 
         /** Gets an access to this variable. */
@@ -2055,7 +2062,7 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
       }
 
       /** A declaration. */
-      class Declaration extends AstNode {
+      class Declaration {
         /**
          * Gets the type mention of the entity that contains this declaration, if any.
          *
@@ -2088,6 +2095,12 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
          * a function.
          */
         TypeMention getType();
+
+        /** Gets a textual representation of this declaration. */
+        string toString();
+
+        /** Gets the location of this declaration. */
+        Location getLocation();
       }
 
       /**
@@ -2167,6 +2180,12 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
       /** A parameter. */
       class Parameter extends VariableDeclaration;
 
+      /**
+       * Holds if `p` is an implicit parameter declaration corresponding
+       * to variable `v`.
+       */
+      default predicate implicitParameterDecl(Parameter p, Variable v) { none() }
+
       /** A callable. This may include for example variant constructors. */
       class Callable extends Declaration {
         /**
@@ -2368,7 +2387,10 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
       }
 
       /** A closure/lambda expression. */
-      class Closure extends Callable, Expr;
+      class Closure extends Callable {
+        /** Gets the expression that defines this closure (typically the entity itself). */
+        Expr getDefiningExpr();
+      }
 
       /**
        * A special pseudo type representing a particular closure parameter without
@@ -2484,12 +2506,14 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
     module Make3<InputSig3 Input3> {
       private import Input3
 
-      private predicate closureStep(AstNode pattern, TypePath prefix1, Closure c, TypePath prefix2) {
-        exists(Parameter p |
+      pragma[nomagic]
+      private predicate closureStep(AstNode pattern, TypePath prefix1, Expr c, TypePath prefix2) {
+        exists(Closure c0, Parameter p |
           pattern = p.getPattern() and
-          p = c.getParameter(_) and
+          p = c0.getParameter(_) and
           prefix1.isEmpty() and
-          prefix2 = getClosureParameterTypePath(p)
+          prefix2 = getClosureParameterTypePath(p) and
+          c = c0.getDefiningExpr()
         )
       }
 
@@ -2513,10 +2537,16 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
               tm = decl.getType() and
               n = decl.getPattern()
             )
+            or
+            exists(Parameter p, Variable v |
+              implicitParameterDecl(p, v) and
+              result = p.getType().getTypeAt(path) and
+              n = v.getAnAccess()
+            )
           )
           or
           exists(Closure c, TypePath suffix |
-            n = c and
+            n = c.getDefiningExpr() and
             result = getCallableReturnType(c, suffix) and
             path = getClosureReturnTypePath(c).append(suffix)
           )
@@ -2615,8 +2645,11 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
             or
             result = inferLogicalOperationType(n, path)
             or
-            result = getClosureType(n) and
-            path.isEmpty()
+            exists(Closure c |
+              n = c.getDefiningExpr() and
+              result = getClosureType(c) and
+              path.isEmpty()
+            )
             or
             infersCertainTypeAt(n, path, result.getATypeParameter())
           ) and
@@ -2709,9 +2742,9 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
         or
         exists(Closure c |
           n1 = c.getBody() and
-          n2 = c and
+          n2 = c.getDefiningExpr() and
           prefix1.isEmpty() and
-          prefix2 = getClosureReturnTypePath(n2)
+          prefix2 = getClosureReturnTypePath(c)
         )
       }
 
@@ -2725,7 +2758,10 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
           or
           closureStep(n2, prefix2, n, prefix1) and
           // prevent closure parameter pseudo types from escaping the closure
-          not result.(ClosureParameterPseudoType).getParameter() = n.(Closure).getParameter(_)
+          not exists(Closure c |
+            n = c.getDefiningExpr() and
+            result.(ClosureParameterPseudoType).getParameter() = c.getParameter(_)
+          )
         )
       }
 
@@ -3076,7 +3112,7 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
         }
 
         pragma[nomagic]
-        private predicate hasUnknownTypeAt(AstNode n, TypePath path) {
+        predicate hasUnknownTypeAt(AstNode n, TypePath path) {
           inferType(n, path) instanceof UnknownType
         }
 
@@ -3209,8 +3245,8 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
             result.(ClosureParameterPseudoType).getParameter() = p
             or
             // step 3
-            hasClosureParameterPseudoType(c, path, p) and
-            n = c and
+            hasClosureParameterPseudoType(n, path, p) and
+            n = c.getDefiningExpr() and
             result instanceof UnknownType
           )
           or
@@ -3227,23 +3263,64 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
           or
           // The `step X` comments below refer to the steps for 'Case A' in the
           // QL doc for `ClosureParameterPseudoType`.
-          exists(Closure c, Parameter p |
+          exists(Closure c, Parameter p, Expr def |
             p = c.getParameter(_) and
-            not exists(p.getType())
+            not exists(p.getType()) and
+            def = c.getDefiningExpr()
           |
             // step 1
-            n = c and
+            n = def and
             path = getClosureParameterTypePath(p) and
             result instanceof UnknownType
             or
             // step 3
             n = p.getPattern() and
-            result = inferType(c, getClosureParameterTypePath(p).appendInverse(path)) and
+            result = inferType(def, getClosureParameterTypePath(p).appendInverse(path)) and
             not (path.isEmpty() and result instanceof UnknownType)
           )
         }
       }
 
+      /**
+       * Holds if `n` has unknown type at at `prefix`, but is still able to
+       * infer a known type at `suffix` for the `i`th type parameter of whatever
+       * the unknown type is.
+       *
+       * For example, in
+       *
+       * ```rust
+       * let mut x: Unresolvable<i32> = ...;
+       *
+       * x = resolvable(...);
+       * ```
+       *
+       * even though the root type is unresolvable at the declaration, we are still
+       * able to infer that the first type argument is `i32`. We can then combine this
+       * information with later inferred type information.
+       */
+      pragma[nomagic]
+      private predicate infersUnknownTypeArg(
+        AstNode n, TypePath prefix, int i, TypePath suffix, Type t
+      ) {
+        exists(TypeParameter tp, TypePath suffix0 |
+          ContextualTyping::hasUnknownTypeAt(n, prefix) and
+          suffix0.isCons(tp, suffix) and
+          tp = any(UnknownType ut).getPositionalTypeParameter(i) and
+          t = inferType(n, prefix.appendInverse(suffix0)) and
+          not t instanceof UnknownType
+        )
+      }
+
+      pragma[nomagic]
+      private predicate infersKnownAndUnknownType(AstNode n, TypePath path, int i, TypeParameter tp) {
+        ContextualTyping::hasUnknownTypeAt(n, path) and
+        exists(Type t |
+          t = inferType(n, path) and
+          not t instanceof UnknownType and
+          tp = t.getPositionalTypeParameter(i)
+        )
+      }
+
       /**
        * Gets an inferred candidate type of `n` at `path`.
        *
@@ -3275,6 +3352,12 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
         result instanceof UnknownType
         or
         result = ContextualTyping::inferTypeContextual(n, path)
+        or
+        exists(TypePath prefix, int i, TypePath suffix, TypeParameter tp |
+          infersUnknownTypeArg(n, prefix, i, suffix, result) and
+          infersKnownAndUnknownType(n, prefix, i, tp) and
+          path = prefix.append(TypePath::cons(tp, suffix))
+        )
       }
 
       /**
@@ -3298,7 +3381,10 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
           result instanceof ClosureParameterPseudoType
         ) and
         // prevent closure parameter pseudo types from escaping from the closure
-        not result.(ClosureParameterPseudoType).getParameter() = n.(Closure).getParameter(_)
+        not exists(Closure c |
+          n = c.getDefiningExpr() and
+          result.(ClosureParameterPseudoType).getParameter() = c.getParameter(_)
+        )
         or
         // If `n` has an explicitly unknown type at `prefix` and at the same time a certain
         // type at `prefix.suffix`, then extend the unknown type information to any path