Rust: Improvements to tuple type inference based on PR feedback

Author: paldepind

rust/ql/.generated.list

@@ -1,4 +1,3 @@
-// generated by codegen, remove this comment if you wish to edit this file
 /**
  * This module provides a hand-modifiable wrapper around the generated class `TuplePat`.
  *
@@ -12,12 +11,25 @@ private import codeql.rust.elements.internal.generated.TuplePat
  * be referenced directly.
  */
 module Impl {
+  private import rust
+
+  // the following QLdoc is generated: if you need to edit it, do it in the schema file
   /**
    * A tuple pattern. For example:
    * ```rust
    * let (x, y) = (1, 2);
    * let (a, b, ..,  z) = (1, 2, 3, 4, 5);
    * ```
    */
-  class TuplePat extends Generated::TuplePat { }
+  class TuplePat extends Generated::TuplePat {
+    /**
+     * Gets the arity of the tuple matched by this pattern, if any.
+     *
+     * This is the number of fields in the tuple pattern if and only if the
+     * pattern does not contain a `..` pattern.
+     */
+    int getTupleArity() {
+      result = this.getNumberOfFields() and not this.getAField() instanceof RestPat
+    }
+  }
 }

rust/ql/.gitattributes

@@ -10,7 +10,12 @@ private import codeql.rust.elements.internal.generated.Synth
 cached
 newtype TType =
   TTuple(int arity) {
-    arity = any(TupleTypeRepr t).getNumberOfFields() and
+    arity =
+      [
+        any(TupleTypeRepr t).getNumberOfFields(),
+        any(TupleExpr e).getNumberOfFields(),
+        any(TuplePat p).getNumberOfFields()
+      ] and
     Stages::TypeInferenceStage::ref()
   } or
   TStruct(Struct s) or
@@ -59,40 +64,33 @@ abstract class Type extends TType {
   abstract Location getLocation();
 }
 
-/** The unit type `()`. */
-class UnitType extends Type, TTuple {
-  UnitType() { this = TTuple(0) }
-
-  override StructField getStructField(string name) { none() }
-
-  override TupleField getTupleField(int i) { none() }
-
-  override TypeParameter getTypeParameter(int i) { none() }
-
-  override string toString() { result = "()" }
-
-  override Location getLocation() { result instanceof EmptyLocation }
-}
-
 /** A tuple type `(T, ...)`. */
 class TupleType extends Type, TTuple {
   private int arity;
 
-  TupleType() { this = TTuple(arity) and arity > 0 }
+  TupleType() { this = TTuple(arity) }
 
   override StructField getStructField(string name) { none() }
 
   override TupleField getTupleField(int i) { none() }
 
   override TypeParameter getTypeParameter(int i) { result = TTupleTypeParameter(arity, i) }
 
+  /** Gets the arity of this tuple type. */
   int getArity() { result = arity }
 
   override string toString() { result = "(T_" + arity + ")" }
 
   override Location getLocation() { result instanceof EmptyLocation }
 }
 
+/** The unit type `()`. */
+class UnitType extends TupleType, TTuple {
+  UnitType() { this = TTuple(0) }
+
+  override string toString() { result = "()" }
+}
+
 abstract private class StructOrEnumType extends Type {
   abstract ItemNode asItemNode();
 }
@@ -355,8 +353,9 @@ class AssociatedTypeTypeParameter extends TypeParameter, TAssociatedTypeTypePara
 /**
  * A tuple type parameter. For instance the `T` in `(T, U)`.
  *
- *  Since tuples are structural their parameters can be represented simply as
- *  their positional index.
+ * Since tuples are structural their type parameters can be represented as their
+ * positional index. The type inference library requires that type parameters
+ * belong to a single type, so we also include the arity of the tuple type.
  */
 class TupleTypeParameter extends TypeParameter, TTupleTypeParameter {
   private int arity;
@@ -371,8 +370,8 @@ class TupleTypeParameter extends TypeParameter, TTupleTypeParameter {
   /** Gets the index of this tuple type parameter. */
   int getIndex() { result = index }
 
-  /** Gets the arity of this tuple type parameter. */
-  int getArity() { result = arity }
+  /** Gets the tuple type that corresponds to this tuple type parameter. */
+  TupleType getTupleType() { result = TTuple(arity) }
 }
 
 /** An implicit array type parameter. */

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

@@ -108,7 +108,7 @@ private module Input1 implements InputSig1<Location> {
           maxArity = max(int i | i = any(TupleType tt).getArity()) and
           tp0 = ttp and
           kind = 2 and
-          id = ttp.getArity() * maxArity + ttp.getIndex()
+          id = ttp.getTupleType().getArity() * maxArity + ttp.getIndex()
         )
       |
         tp0 order by kind, id
@@ -335,7 +335,7 @@ private predicate typeEquality(AstNode n1, TypePath prefix1, AstNode n2, TypePat
     arity = n2.(TupleExpr).getNumberOfFields() and
     n1 = n2.(TupleExpr).getField(i)
     or
-    arity = n2.(TuplePat).getNumberOfFields() and
+    arity = n2.(TuplePat).getTupleArity() and
     n1 = n2.(TuplePat).getField(i)
   )
   or
@@ -553,9 +553,9 @@ private Type inferStructExprType(AstNode n, TypePath path) {
 }
 
 pragma[nomagic]
-private Type inferTupleExprRootType(TupleExpr te) {
-  // `typeEquality` handles the non-root case
-  result = TTuple(te.getNumberOfFields())
+private Type inferTupleRootType(AstNode n) {
+  // `typeEquality` handles the non-root cases
+  result = TTuple([n.(TupleExpr).getNumberOfFields(), n.(TuplePat).getTupleArity()])
 }
 
 pragma[nomagic]
@@ -1091,16 +1091,27 @@ private Type inferTupleIndexExprType(FieldExpr fe, TypePath path) {
 
 /** Infers the type of `t` in `t.n` when `t` is a tuple. */
 private Type inferTupleContainerExprType(Expr e, TypePath path) {
-  // NOTE: For a field expression `t.n` where `n` is a number `t` might both be
-  // a tuple struct or a tuple. It is only correct to let type information flow
-  // from `t.n` to tuple type parameters of `t` in the latter case. Hence we
-  // include the condition that the root type of `t` must be a tuple type.
+  // NOTE: For a field expression `t.n` where `n` is a number `t` might be a
+  // tuple as in:
+  // ```rust
+  // let t = (Default::default(), 2);
+  // let s: String = t.0;
+  // ```
+  // But it could also be a tuple struct as in:
+  // ```rust
+  // struct T(String, u32);
+  // let t = T(Default::default(), 2);
+  // let s: String = t.0;
+  // ```
+  // We need type information to flow from `t.n` to tuple type parameters of `t`
+  // in the former case but not the latter case. Hence we include the condition
+  // that the root type of `t` must be a tuple type.
   exists(int i, TypePath path0, FieldExpr fe, int arity |
     e = fe.getContainer() and
     fe.getIdentifier().getText() = i.toString() and
     arity = inferType(fe.getContainer()).(TupleType).getArity() and
     result = inferType(fe, path0) and
-    path = TypePath::cons(TTupleTypeParameter(arity, i), path0) // FIXME:
+    path = TypePath::cons(TTupleTypeParameter(arity, i), path0)
   )
 }
 
@@ -1992,7 +2003,7 @@ private module Cached {
     or
     result = inferStructExprType(n, path)
     or
-    result = inferTupleExprRootType(n) and
+    result = inferTupleRootType(n) and
     path.isEmpty()
     or
     result = inferPathExprType(n, path)

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

@@ -2357,6 +2357,13 @@ mod tuples {
         let pair = (a, b); // $ type=pair:0(2).i64 type=pair:1(2).bool
         let i: i64 = pair.0;
         let j: bool = pair.1;
+
+        let pair = [1, 1].into(); // $ type=pair:0(2).i32 MISSING: target=into
+        match pair {
+            (0,0) => print!("unexpected"),
+            _ => print!("expected"),
+        }
+        let x = pair.0; // $ type=x:i32
     }
 }
 

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

@@ -704,7 +704,7 @@ pub fn complex_nested_patterns() {
         }
         // Catch-all with identifier pattern
         other => {
-            let other_complex = other; // $ MISSING: type=other_complex:?
+            let other_complex = other; // $ type=other_complex:0(2).Point type=other_complex:1(2).MyOption
             println!("Other complex data: {:?}", other_complex);
         }
     }
@@ -766,7 +766,7 @@ pub fn patterns_in_function_parameters() {
 
     // Call the functions to use them
     let point = Point { x: 5, y: 10 };
-    let extracted = extract_point(point); // $ target=extract_point MISSING: type=extracted:?
+    let extracted = extract_point(point); // $ target=extract_point type=extracted:0(2).i32 type=extracted:1(2).i32
 
     let color = Color(200, 100, 50);
     let red = extract_color(color); // $ target=extract_color type=red:u8