Rust: Make function trait syntax without return type default to unit

Author: paldepind

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

@@ -222,6 +222,33 @@ class NonAliasPathTypeMention extends PathTypeMention {
       result = this.getPositionalTypeArgument(pragma[only_bind_into](i), path) and
       tp = this.resolveRootType().getPositionalTypeParameter(pragma[only_bind_into](i))
     )
+    or
+    // Handle the special syntactic sugar for function traits. The syntactic
+    // form is detected by the presence of a parenthesized argument list which
+    // is a mandatory part of the syntax [1].
+    //
+    // For now we only support `FnOnce` as we can't support the "inherited"
+    // associated types of `Fn` and `FnMut` yet.
+    //
+    // [1]: https://doc.rust-lang.org/reference/paths.html#grammar-TypePathFn
+    exists(FnOnceTrait t, PathSegment s |
+      t = resolved and
+      s = this.getSegment() and
+      s.hasParenthesizedArgList()
+    |
+      tp = TTypeParamTypeParameter(t.getTypeParam()) and
+      result = s.getParenthesizedArgList().(TypeMention).resolveTypeAt(path)
+      or
+      tp = TAssociatedTypeTypeParameter(t.getOutputType()) and
+      (
+        result = s.getRetType().getTypeRepr().(TypeMention).resolveTypeAt(path)
+        or
+        // When the `-> ...` return type is omitted, it defaults to `()`.
+        not s.hasRetType() and
+        result instanceof UnitType and
+        path.isEmpty()
+      )
+    )
   }
 
   pragma[nomagic]
@@ -256,17 +283,6 @@ class NonAliasPathTypeMention extends PathTypeMention {
       result = alias.getTypeRepr() and
       tp = TAssociatedTypeTypeParameter(this.getResolvedAlias(pragma[only_bind_into](name)))
     )
-    or
-    // Handle the special syntactic sugar for function traits. For now we only
-    // support `FnOnce` as we can't support the "inherited" associated types of
-    // `Fn` and `FnMut` yet.
-    exists(FnOnceTrait t | t = resolved |
-      tp = TTypeParamTypeParameter(t.getTypeParam()) and
-      result = this.getSegment().getParenthesizedArgList()
-      or
-      tp = TAssociatedTypeTypeParameter(t.getOutputType()) and
-      result = this.getSegment().getRetType().getTypeRepr()
-    )
   }
 
   pragma[nomagic]

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

@@ -37,7 +37,7 @@ mod fn_once_trait {
     }
 
     fn return_type_omitted<F: FnOnce(bool)>(f: F) {
-        let _return = f(true); // $ MISSING: type=_return:()
+        let _return = f(true); // $ type=_return:()
     }
 
     fn argument_type<F: FnOnce(bool) -> i64>(f: F) {

rust/ql/test/library-tests/type-inference/type-inference.expected

@@ -4982,7 +4982,9 @@ inferType
 | closure.rs:36:25:36:28 | true |  | {EXTERNAL LOCATION} | bool |
 | closure.rs:39:45:39:45 | f |  | closure.rs:39:28:39:42 | F |
 | closure.rs:39:51:41:5 | { ... } |  | {EXTERNAL LOCATION} | () |
+| closure.rs:40:13:40:19 | _return |  | {EXTERNAL LOCATION} | () |
 | closure.rs:40:23:40:23 | f |  | closure.rs:39:28:39:42 | F |
+| closure.rs:40:23:40:29 | f(...) |  | {EXTERNAL LOCATION} | () |
 | closure.rs:40:25:40:28 | true |  | {EXTERNAL LOCATION} | bool |
 | closure.rs:43:46:43:46 | f |  | closure.rs:43:22:43:43 | F |
 | closure.rs:43:52:46:5 | { ... } |  | {EXTERNAL LOCATION} | () |