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

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

@@ -36,6 +36,10 @@ mod fn_once_trait {
         let _return = f(true); // $ type=_return:i64
     }
 
+    fn return_type_omitted<F: FnOnce(bool)>(f: F) {
+        let _return = f(true); // $ type=_return:()
+    }
+
     fn argument_type<F: FnOnce(bool) -> i64>(f: F) {
         let arg = Default::default(); // $ target=default type=arg:bool
         f(arg);