Rust: Implement basic type inference in QL

Author: hvitved

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

@@ -17,6 +17,7 @@ module Impl {
   private import codeql.rust.elements.internal.CallExprImpl::Impl
   private import codeql.rust.elements.internal.PathExprImpl::Impl
   private import codeql.rust.elements.internal.PathResolution
+  private import codeql.rust.elements.internal.TypeInference
 
   pragma[nomagic]
   Resolvable getCallResolvable(CallExprBase call) {
@@ -35,9 +36,11 @@ module Impl {
      * be statically resolved.
      */
     Callable getStaticTarget() {
-      getCallResolvable(this).resolvesAsItem(result)
-      or
+      // getCallResolvable(this).resolvesAsItem(result)
+      // or
       result = resolvePath(this.(CallExpr).getFunction().(PathExpr).getPath())
+      or
+      result = resolveMethodCallExpr(this)
     }
   }
 }

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

@@ -184,10 +184,10 @@ abstract private class ImplOrTraitItemNode extends ItemNode {
   }
 }
 
-private class ImplItemNode extends ImplOrTraitItemNode instanceof Impl {
+class ImplItemNode extends ImplOrTraitItemNode instanceof Impl {
   override string getName() { result = "(impl)" }
 
-  override Visibility getVisibility() { none() }
+  override Visibility getVisibility() { result = Impl.super.getVisibility() }
 }
 
 private class MacroCallItemNode extends ItemNode instanceof MacroCall {
@@ -232,6 +232,12 @@ private class BlockExprItemNode extends ItemNode instanceof BlockExpr {
   override Visibility getVisibility() { none() }
 }
 
+private class TypeParamItemNode extends ItemNode instanceof TypeParam {
+  override string getName() { result = TypeParam.super.getName().getText() }
+
+  override Visibility getVisibility() { none() }
+}
+
 /** Holds if `item` has the name `name` and is a top-level item inside `f`. */
 private predicate sourceFileEdge(SourceFile f, string name, ItemNode item) {
   item = f.getAnItem() and

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

@@ -0,0 +1,319 @@
+/** Provides functionality for inferring types. */
+
+private import rust
+private import PathResolution
+
+/** Gets the singleton type path `i`. */
+bindingset[i]
+private TypePath typePath(int i) { result = i.toString() }
+
+bindingset[s]
+private predicate decodeTypePathComponent(string s, int i) { i = s.toInt() }
+
+/**
+ * A path into a (constructed) type.
+ *
+ * Paths are represented in left-to-right order, for example, a path `0.1` into the
+ * type `C1<C2<A,B>,C3<C,D>>` points at the type `B`.
+ *
+ * Type paths are used to represent constructed types without using a `newtype`, which
+ * makes it practically feasible to do type inference in mutual recursion with call
+ * resolution.
+ *
+ * As an example, the type above can be represented by the following set of tuples
+ *
+ * `TypePath` | `Type`
+ * ---------- | ------
+ * `""`       | ``C1``
+ * `"0"`      | ``C2``
+ * `"0.0"`    | `A`
+ * `"0.1"`    | `B`
+ * `"1"`      | ``C3``
+ * `"1.0"`    | `C`
+ * `"1.1"`    | `D`
+ */
+class TypePath extends string {
+  bindingset[this]
+  TypePath() { exists(this) }
+
+  predicate isEmpty() { this = "" }
+
+  /** Gets the path obtained by appending `suffix` onto this path. */
+  bindingset[suffix, result]
+  bindingset[this, result]
+  bindingset[this, suffix]
+  TypePath append(TypePath suffix) {
+    if this.isEmpty()
+    then result = suffix
+    else
+      if suffix.isEmpty()
+      then result = this
+      else result = this + "." + suffix
+  }
+
+  /** Holds if this path starts with `prefix`, followed by `i`. */
+  bindingset[this]
+  predicate endsWith(TypePath prefix, int i) {
+    decodeTypePathComponent(this, i) and
+    prefix.isEmpty()
+    or
+    exists(int last |
+      last = max(this.indexOf(".")) and
+      prefix = this.prefix(last) and
+      decodeTypePathComponent(this.suffix(last + 1), i)
+    )
+  }
+
+  /** Holds if this path starts with `i`, followed by `suffix`. */
+  bindingset[this]
+  predicate startsWith(int i, TypePath suffix) {
+    decodeTypePathComponent(this, i) and
+    suffix.isEmpty()
+    or
+    exists(int first |
+      first = min(this.indexOf(".")) and
+      suffix = this.suffix(first + 1) and
+      decodeTypePathComponent(this.prefix(first), i)
+    )
+  }
+}
+
+private predicate letStmtTyped(LetStmt let, Pat pat, TypeRepr t) {
+  pat = let.getPat() and
+  t = let.getTypeRepr() and
+  not t instanceof InferTypeRepr
+}
+
+private predicate paramTyped(Param p, Pat pat, TypeRepr t) {
+  pat = p.getPat() and
+  t = p.getTypeRepr() and
+  not t instanceof InferTypeRepr
+}
+
+private predicate isTargetTyped(AstNode n) {
+  exists(Variable v |
+    n = v.getPat() and
+    not letStmtTyped(_, n, _) and
+    not paramTyped(_, n, _)
+  )
+}
+
+// todo: add more cases
+private newtype TType =
+  TStruct(Struct s) or
+  TEnum(Enum e) or
+  TArrayType() or // todo: add size?
+  TRefType() or // todo: add mut, lifetime?
+  TTypeParameter(TypeParam t)
+
+abstract private class Type extends TType {
+  pragma[nomagic]
+  abstract Function getMethod(string name);
+
+  abstract string toString();
+
+  abstract Location getLocation();
+}
+
+class StructType extends Type {
+  private Struct struct;
+
+  StructType() { this = TStruct(struct) }
+
+  override Function getMethod(string name) {
+    exists(ImplItemNode i |
+      struct = resolvePath(i.(Impl).getSelfTy().(PathTypeRepr).getPath()) and
+      result = i.getASuccessor(name)
+    )
+  }
+
+  override string toString() { result = struct.toString() }
+
+  override Location getLocation() { result = struct.getLocation() }
+}
+
+class EnumType extends Type {
+  private Enum enum;
+
+  EnumType() { this = TEnum(enum) }
+
+  override Function getMethod(string name) {
+    exists(ImplItemNode i |
+      enum = resolvePath(i.(Impl).getSelfTy().(PathTypeRepr).getPath()) and
+      result = i.getASuccessor(name)
+    )
+  }
+
+  override string toString() { result = enum.toString() }
+
+  override Location getLocation() { result = enum.getLocation() }
+}
+
+class ArrayType extends Type {
+  ArrayType() { this = TArrayType() }
+
+  override Function getMethod(string name) { none() }
+
+  override string toString() { result = "[]" }
+
+  override Location getLocation() { result instanceof EmptyLocation }
+}
+
+class RefType extends Type {
+  RefType() { this = TRefType() }
+
+  override Function getMethod(string name) { none() }
+
+  override string toString() { result = "&" }
+
+  override Location getLocation() { result instanceof EmptyLocation }
+}
+
+class TypeParameter extends Type {
+  private TypeParam typeParam;
+
+  TypeParameter() { this = TTypeParameter(typeParam) }
+
+  override Function getMethod(string name) { none() }
+
+  override string toString() { result = typeParam.toString() }
+
+  override Location getLocation() { result = typeParam.getLocation() }
+}
+
+abstract private class TypeReprOrPath extends AstNode {
+  TypeRepr getTypeReprArgument(int i) {
+    result = this.(ArrayTypeRepr).getElementTypeRepr() and
+    i = 0
+    or
+    result = this.(RefTypeRepr).getTypeRepr() and
+    i = 0
+    or
+    result = this.(Path).getPart().getGenericArgList().getGenericArg(i).(TypeArg).getTypeRepr()
+    or
+    result = this.(PathTypeRepr).getPath().(PathTypeReprOrPath).getTypeReprArgument(i)
+  }
+}
+
+private class TypeReprTypeReprOrPath extends TypeReprOrPath, TypeRepr { }
+
+private class PathTypeReprOrPath extends TypeReprOrPath, Path { }
+
+private TypeReprOrPath getTypeReprAt(TypeReprOrPath t, TypePath path) {
+  path.isEmpty() and
+  result = t
+  or
+  exists(int i, TypeReprOrPath arg, TypePath suffix |
+    arg = t.getTypeReprArgument(i) and
+    result = getTypeReprAt(arg, suffix) and
+    path = typePath(i).append(suffix)
+  )
+}
+
+private Type resolveTypeRepr(TypeReprOrPath t) {
+  t instanceof ArrayTypeRepr and
+  result = TArrayType()
+  or
+  t instanceof RefTypeRepr and
+  result = TRefType()
+  or
+  t =
+    any(Path path |
+      result = TStruct(resolvePath(path))
+      or
+      result = TEnum(resolvePath(path))
+      or
+      result = TTypeParameter(resolvePath(path))
+    )
+  or
+  result = resolveTypeRepr(t.(PathTypeRepr).getPath())
+}
+
+private Type resolveTypeAt(TypeReprOrPath t, TypePath path) {
+  result = resolveTypeRepr(getTypeReprAt(t, path))
+}
+
+private Type resolveVariableType(AstNode n, TypePath path) {
+  exists(TypeRepr t | letStmtTyped(_, n, t) or paramTyped(_, n, t) |
+    result = resolveTypeAt(t, path)
+  )
+  or
+  exists(Variable v |
+    result = resolveType(v.getPat(), path) and
+    n = v.getAnAccess()
+  )
+}
+
+pragma[nomagic]
+private Type resolveTargetTyped(AstNode n, TypePath path) {
+  isTargetTyped(n) and
+  exists(LetStmt let |
+    let.getPat() = n and
+    result = resolveType(let.getInitializer(), path)
+  )
+}
+
+private Type resolveRecordExprType(RecordExpr re, TypePath path) {
+  result = resolveTypeAt(re.getPath(), path)
+}
+
+pragma[nomagic]
+private Type resolveCallExprType(CallExpr ce, TypePath path) {
+  exists(ItemNode i | i = resolvePath(ce.getFunction().(PathExpr).getPath()) |
+    result = resolveTypeAt(i.(Function).getRetType().getTypeRepr(), path)
+    or
+    exists(Enum e |
+      i = e.getVariantList().getAVariant() and
+      result = TEnum(e) and
+      path.isEmpty()
+    )
+  )
+}
+
+pragma[nomagic]
+Function resolveMethodCallExpr(MethodCallExpr mce) {
+  exists(Type t |
+    t = resolveType(mce.getReceiver()) and
+    if t = TRefType()
+    then
+      // for reference types, lookup the method in the type being referenced
+      result = resolveType(mce.getReceiver(), "0").getMethod(mce.getNameRef().getText())
+    else result = t.getMethod(mce.getNameRef().getText())
+  )
+}
+
+pragma[nomagic]
+private Type resolveMethodCallExprType(MethodCallExpr mce, TypePath path) {
+  exists(Function f |
+    f = resolveMethodCallExpr(mce) and
+    result = resolveTypeAt(f.getRetType().getTypeRepr(), path)
+  )
+}
+
+pragma[nomagic]
+private Type resolveRefExprType(RefExpr re, TypePath path) {
+  exists(re) and
+  path.isEmpty() and
+  result = TRefType()
+  or
+  exists(TypePath suffix |
+    result = resolveType(re.getExpr(), suffix) and
+    path = typePath(0).append(suffix)
+  )
+}
+
+Type resolveType(AstNode n, TypePath path) {
+  result = resolveVariableType(n, path)
+  or
+  result = resolveTargetTyped(n, path)
+  or
+  result = resolveRecordExprType(n, path)
+  or
+  result = resolveCallExprType(n, path)
+  or
+  result = resolveMethodCallExprType(n, path)
+  or
+  result = resolveRefExprType(n, path)
+}
+
+Type resolveType(AstNode n) { result = resolveType(n, "") }