Rust: Store arity in tuple type parameters

Type parameters are required to belong to a single type only. Since we store the arity for tuple types, we need to store the arity in tuple type parameters as well such that we can associate them to the tuple type of the same arity.

Author: paldepind

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

@@ -10,7 +10,8 @@ private import codeql.rust.elements.internal.generated.Synth
 cached
 newtype TType =
   TTuple(int arity) {
-    exists(any(TupleTypeRepr t).getField(arity)) and Stages::TypeInferenceStage::ref()
+    arity = any(TupleTypeRepr t).getNumberOfFields() and
+    Stages::TypeInferenceStage::ref()
   } or
   TStruct(Struct s) or
   TEnum(Enum e) or
@@ -19,7 +20,7 @@ newtype TType =
   TRefType() or // todo: add mut?
   TImplTraitType(ImplTraitTypeRepr impl) or
   TSliceType() or
-  TTupleTypeParameter(int i) { exists(TTuple(i)) } or
+  TTupleTypeParameter(int arity, int i) { exists(TTuple(arity)) and i in [0 .. arity - 1] } or
   TTypeParamTypeParameter(TypeParam t) or
   TAssociatedTypeTypeParameter(TypeAlias t) { any(TraitItemNode trait).getAnAssocItem() = t } or
   TArrayTypeParameter() or
@@ -83,7 +84,7 @@ class TupleType extends Type, TTuple {
 
   override TupleField getTupleField(int i) { none() }
 
-  override TypeParameter getTypeParameter(int i) { result = TTupleTypeParameter(i) and i < arity }
+  override TypeParameter getTypeParameter(int i) { result = TTupleTypeParameter(arity, i) }
 
   int getArity() { result = arity }
 
@@ -358,12 +359,20 @@ class AssociatedTypeTypeParameter extends TypeParameter, TAssociatedTypeTypePara
  *  their positional index.
  */
 class TupleTypeParameter extends TypeParameter, TTupleTypeParameter {
-  override string toString() { result = this.getIndex().toString() }
+  private int arity;
+  private int index;
+
+  TupleTypeParameter() { this = TTupleTypeParameter(arity, index) }
+
+  override string toString() { result = index.toString() + "(" + arity + ")" }
 
   override Location getLocation() { result instanceof EmptyLocation }
 
   /** Gets the index of this tuple type parameter. */
-  int getIndex() { this = TTupleTypeParameter(result) }
+  int getIndex() { result = index }
+
+  /** Gets the arity of this tuple type parameter. */
+  int getArity() { result = arity }
 }
 
 /** An implicit array type parameter. */

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

@@ -104,8 +104,12 @@ private module Input1 implements InputSig1<Location> {
           node = tp0.(ImplTraitTypeTypeParameter).getImplTraitTypeRepr()
         )
         or
-        kind = 2 and
-        id = tp0.(TupleTypeParameter).getIndex()
+        exists(TupleTypeParameter ttp, int maxArity |
+          maxArity = max(int i | i = any(TupleType tt).getArity()) and
+          tp0 = ttp and
+          kind = 2 and
+          id = ttp.getArity() * maxArity + ttp.getIndex()
+        )
       |
         tp0 order by kind, id
       )
@@ -324,11 +328,14 @@ private predicate typeEquality(AstNode n1, TypePath prefix1, AstNode n2, TypePat
   prefix1.isEmpty() and
   prefix2 = TypePath::singleton(TRefTypeParameter())
   or
-  exists(int i |
+  exists(int i, int arity |
     prefix1.isEmpty() and
-    prefix2 = TypePath::singleton(TTupleTypeParameter(i))
+    prefix2 = TypePath::singleton(TTupleTypeParameter(arity, i))
   |
-    n1 = n2.(TupleExpr).getField(i) or
+    arity = n2.(TupleExpr).getNumberOfFields() and
+    n1 = n2.(TupleExpr).getField(i)
+    or
+    arity = n2.(TuplePat).getNumberOfFields() and
     n1 = n2.(TuplePat).getField(i)
   )
   or
@@ -1077,7 +1084,7 @@ private Type inferTupleIndexExprType(FieldExpr fe, TypePath path) {
   exists(int i, TypePath path0 |
     fe.getIdentifier().getText() = i.toString() and
     result = inferType(fe.getContainer(), path0) and
-    path0.isCons(TTupleTypeParameter(i), path) and
+    path0.isCons(TTupleTypeParameter(_, i), path) and
     fe.getIdentifier().getText() = i.toString()
   )
 }
@@ -1088,12 +1095,12 @@ private Type inferTupleContainerExprType(Expr e, TypePath path) {
   // 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.
-  exists(int i, TypePath path0, FieldExpr fe |
+  exists(int i, TypePath path0, FieldExpr fe, int arity |
     e = fe.getContainer() and
     fe.getIdentifier().getText() = i.toString() and
-    inferType(fe.getContainer()) instanceof TupleType and
+    arity = inferType(fe.getContainer()).(TupleType).getArity() and
     result = inferType(fe, path0) and
-    path = TypePath::cons(TTupleTypeParameter(i), path0)
+    path = TypePath::cons(TTupleTypeParameter(arity, i), path0) // FIXME:
   )
 }
 

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

@@ -21,7 +21,7 @@ class TupleTypeReprMention extends TypeMention instanceof TupleTypeRepr {
     or
     exists(TypePath suffix, int i |
       result = super.getField(i).(TypeMention).resolveTypeAt(suffix) and
-      path = TypePath::cons(TTupleTypeParameter(i), suffix)
+      path = TypePath::cons(TTupleTypeParameter(super.getNumberOfFields(), i), suffix)
     )
   }
 }

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

@@ -2354,7 +2354,7 @@ mod tuples {
         // `a` and `b` to be inferred.
         let a = Default::default(); // $ target=default type=a:i64
         let b = Default::default(); // $ target=default type=b:bool
-        let pair = (a, b); // $ type=pair:0.i64 type=pair:1.bool
+        let pair = (a, b); // $ type=pair:0(2).i64 type=pair:1(2).bool
         let i: i64 = pair.0;
         let j: bool = pair.1;
     }

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

@@ -460,7 +460,7 @@ pub fn tuple_patterns() {
     // With rest pattern
     match tuple {
         (first, ..) => {
-            let tuple_first = first; // $ type=tuple_first:i32
+            let tuple_first = first; // $ MISSING: type=tuple_first:i32
             println!("First element: {}", tuple_first);
         }
     }
@@ -630,7 +630,7 @@ pub fn rest_patterns() {
     // RestPat - Rest patterns (..)
     match tuple {
         (first, ..) => {
-            let rest_first = first; // $ type=rest_first:i32
+            let rest_first = first; // $ MISSING: type=rest_first:i32
             println!("First with rest: {}", rest_first);
         }
     }
@@ -644,7 +644,7 @@ pub fn rest_patterns() {
 
     match tuple {
         (first, .., last) => {
-            let rest_start = first; // $ type=rest_start:i32
+            let rest_start = first; // $ MISSING: type=rest_start:i32
             let rest_end = last; // $ MISSING: type=rest_end:u8
             println!("First and last: {}, {}", rest_start, rest_end);
         }
@@ -772,7 +772,7 @@ pub fn patterns_in_function_parameters() {
     let red = extract_color(color); // $ target=extract_color type=red:u8
 
     let tuple = (42i32, 3.14f64, true);
-    let tuple_extracted = extract_tuple(tuple); // $ target=extract_tuple type=tuple_extracted:0.i32 type=tuple_extracted:1.bool
+    let tuple_extracted = extract_tuple(tuple); // $ target=extract_tuple type=tuple_extracted:0(2).i32 type=tuple_extracted:1(2).bool
 }
 
 #[rustfmt::skip]