Rust: Address review comments and handle ! types in type inference

Author: hvitved

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

@@ -291,6 +291,17 @@ private Type inferAnnotatedType(AstNode n, TypePath path) {
   result = n.(ShorthandSelfParameterMention).resolveTypeAt(path)
 }
 
+/**
+ * Holds if `me` is a call to the `panic!` macro.
+ *
+ * `panic!` needs special treatment, because it expands to a block expression
+ * that looks like it should have type `()` instead of the correct `!` type.
+ */
+pragma[nomagic]
+private predicate isPanicMacroCall(MacroExpr me) {
+  me.getMacroCall().resolveMacro().(MacroRules).getName().getText() = "panic"
+}
+
 /** Module for inferring certain type information. */
 module CertainTypeInference {
   pragma[nomagic]
@@ -443,6 +454,14 @@ module CertainTypeInference {
     or
     result = inferCastExprType(n, path)
     or
+    exprHasUnitType(n) and
+    path.isEmpty() and
+    result instanceof UnitType
+    or
+    isPanicMacroCall(n) and
+    path.isEmpty() and
+    result instanceof NeverType
+    or
     infersCertainTypeAt(n, path, result.getATypeParameter())
   }
 
@@ -579,7 +598,8 @@ private predicate typeEquality(AstNode n1, TypePath prefix1, AstNode n2, TypePat
       n2 = be.getRhs()
     )
     or
-    n1 = n2.(MacroExpr).getMacroCall().getMacroCallExpansion()
+    n1 = n2.(MacroExpr).getMacroCall().getMacroCallExpansion() and
+    not isPanicMacroCall(n2)
     or
     n1 = n2.(MacroPat).getMacroCall().getMacroCallExpansion()
     or
@@ -931,14 +951,17 @@ private predicate functionInfoBlanketLike(
  */
 bindingset[path, type]
 private predicate isComplexRootStripped(TypePath path, Type type) {
-  path.isEmpty() and
-  not validSelfType(type)
-  or
-  exists(TypeParameter tp |
-    complexSelfRoot(_, tp) and
-    path = TypePath::singleton(tp) and
-    exists(type)
-  )
+  (
+    path.isEmpty() and
+    not validSelfType(type)
+    or
+    exists(TypeParameter tp |
+      complexSelfRoot(_, tp) and
+      path = TypePath::singleton(tp) and
+      exists(type)
+    )
+  ) and
+  type != TNeverType()
 }
 
 /**
@@ -1540,7 +1563,8 @@ private module MethodResolution {
     MethodCall getMethodCall() { result = mc_ }
 
     Type getTypeAt(TypePath path) {
-      result = mc_.getACandidateReceiverTypeAtSubstituteLookupTraits(derefChain, borrow, path)
+      result = mc_.getACandidateReceiverTypeAtSubstituteLookupTraits(derefChain, borrow, path) and
+      not result = TNeverType()
     }
 
     pragma[nomagic]
@@ -2810,7 +2834,8 @@ private predicate isReturnExprCfgAncestor(AstNode n) {
 pragma[nomagic]
 predicate isUnitBlockExpr(BlockExpr be) {
   not be.getStmtList().hasTailExpr() and
-  not isReturnExprCfgAncestor(be)
+  not isReturnExprCfgAncestor(be) and
+  not be.hasLabel()
 }
 
 pragma[nomagic]
@@ -2831,6 +2856,15 @@ private Type inferBlockExprType(BlockExpr be, TypePath path) {
   )
 }
 
+pragma[nomagic]
+private predicate exprHasUnitType(Expr e) {
+  e = any(IfExpr ie | not ie.hasElse())
+  or
+  e instanceof WhileExpr
+  or
+  e instanceof ForExpr
+}
+
 final private class AwaitTarget extends Expr {
   AwaitTarget() { this = any(AwaitExpr ae).getExpr() }
 

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

@@ -444,7 +444,11 @@ TypeMention getSelfParamTypeMention(SelfParam self) {
 }
 
 /**
- * An element used to represent the implicit `()` return type of function.
+ * An element used to represent the implicit `()` return type of a function.
+ *
+ * Since the implicit type does not appear in the AST, we (somewhat arbitrarily)
+ * choose the name of the function as a type mention. This works because there
+ * is a one-to-one correspondence between a function and its name.
  */
 class ShorthandReturnTypeMention extends TypeMention instanceof Name {
   private Function f;

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

@@ -90,7 +90,10 @@ module SatisfiesBlanketConstraint<
 
     Location getLocation() { result = at.getLocation() }
 
-    Type getTypeAt(TypePath path) { result = at.getTypeAt(blanketPath.appendInverse(path)) }
+    Type getTypeAt(TypePath path) {
+      result = at.getTypeAt(blanketPath.appendInverse(path)) and
+      not result = TNeverType()
+    }
 
     string toString() { result = at.toString() + " [blanket at " + blanketPath.toString() + "]" }
   }

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

@@ -46,7 +46,7 @@ class FunctionPosition extends TFunctionPosition {
     result = f.getParam(this.asPosition()).getTypeRepr()
     or
     this.isReturn() and
-    result = f.getRetType().getTypeRepr()
+    result = getReturnTypeMention(f)
   }
 
   string toString() {
@@ -263,7 +263,10 @@ module ArgIsInstantiationOf<
   final private class ArgFinal = Arg;
 
   private class ArgSubst extends ArgFinal {
-    Type getTypeAt(TypePath path) { result = substituteLookupTraits(super.getTypeAt(path)) }
+    Type getTypeAt(TypePath path) {
+      result = substituteLookupTraits(super.getTypeAt(path)) and
+      not result = TNeverType()
+    }
   }
 
   private module IsInstantiationOfInput implements
@@ -368,10 +371,10 @@ module ArgsAreInstantiationsOf<ArgsAreInstantiationsOfInputSig Input> {
       CallAndPos cp, Input::Call call, FunctionPosition pos, int rnk, Function f,
       TypeAbstraction abs, AssocFunctionType constraint
     ) {
-      cp = MkCallAndPos(call, pos) and
+      cp = MkCallAndPos(call, pragma[only_bind_into](pos)) and
       call.hasTargetCand(abs, f) and
-      toCheckRanked(abs, f, pos, rnk) and
-      Input::toCheck(abs, f, pos, constraint)
+      toCheckRanked(abs, f, pragma[only_bind_into](pos), rnk) and
+      Input::toCheck(abs, f, pragma[only_bind_into](pos), constraint)
     }
 
     pragma[nomagic]

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

@@ -2841,7 +2841,7 @@ mod block_types {
     #[rustfmt::skip]
     fn f1(cond: bool) -> i32 {
         // Block that evaluates to unit
-        let a = { // $ MISSING: type=a:()
+        let a = { // $ type=a:()
             if cond {
                 return 12;
             }
@@ -2852,7 +2852,7 @@ mod block_types {
     #[rustfmt::skip]
     fn f2() -> i32 {
         // Block that does not evaluate to unit
-        let b = 'label: { // $ MISSING: b:i32 SPURIOUS: certainType=b:()
+        let b = 'label: { // $ MISSING: b:i32
             break 'label 12;
         };
         println!("b: {:?}", b);