Rust: Implement type inference for ref expression as type equality

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

@@ -257,6 +257,10 @@ private predicate typeEquality(AstNode n1, TypePath prefix1, AstNode n2, TypePat
     )
   )
   or
+  n1 = n2.(RefExpr).getExpr() and
+  prefix1.isEmpty() and
+  prefix2 = TypePath::singleton(TRefTypeParameter())
+  or
   n1 = n2.(DerefExpr).getExpr() and
   prefix1 = TypePath::singleton(TRefTypeParameter()) and
   prefix2.isEmpty()
@@ -973,43 +977,9 @@ private Type inferFieldExprType(AstNode n, TypePath path) {
   )
 }
 
-/**
- * Gets the type of `n` at `path`, where `n` is either a reference expression
- * `& x` or an expression `x` inside a reference expression `& x`.
- */
+/** Gets the root type of the reference expression `re`. */
 pragma[nomagic]
-private Type inferRefExprType(Expr e, TypePath path) {
-  exists(RefExpr re |
-    e = re and
-    path.isEmpty() and
-    result = TRefType()
-    or
-    e = re and
-    exists(TypePath exprPath | result = inferType(re.getExpr(), exprPath) |
-      if exprPath.isCons(TRefTypeParameter(), _)
-      then
-        // `&x` simply means `x` when `x` already has reference type
-        path = exprPath
-      else (
-        path = TypePath::cons(TRefTypeParameter(), exprPath) and
-        not (exprPath.isEmpty() and result = TRefType())
-      )
-    )
-    or
-    e = re.getExpr() and
-    exists(TypePath exprPath, TypePath refPath, Type exprType |
-      result = inferType(re, exprPath) and
-      exprPath.isCons(TRefTypeParameter(), refPath) and
-      exprType = inferType(e)
-    |
-      if exprType = TRefType()
-      then
-        // `&x` simply means `x` when `x` already has reference type
-        path = exprPath
-      else path = refPath
-    )
-  )
-}
+private Type inferRefExprType(RefExpr re) { exists(re) and result = TRefType() }
 
 pragma[nomagic]
 private Type inferTryExprType(TryExpr te, TypePath path) {
@@ -1505,7 +1475,8 @@ private module Cached {
     or
     result = inferFieldExprType(n, path)
     or
-    result = inferRefExprType(n, path)
+    result = inferRefExprType(n) and
+    path.isEmpty()
     or
     result = inferTryExprType(n, path)
     or

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

@@ -1154,6 +1154,17 @@ mod implicit_self_borrow {
 }
 
 mod borrowed_typed {
+    #[derive(Debug, Copy, Clone, Default)]
+    struct MyFlag {
+        bool: bool,
+    }
+
+    impl MyFlag {
+        fn flip(&mut self) {
+            self.bool = !self.bool; // $ fieldof=MyFlag method=not
+        }
+    }
+
     struct S;
 
     impl S {
@@ -1179,6 +1190,14 @@ mod borrowed_typed {
         x.f1(); // $ method=f1
         x.f2(); // $ method=f2
         S::f3(&x);
+
+        let n = **&&true; // $ type=n:bool
+
+        // In this example the type of `flag` must be inferred at the call to
+        // `flip` and flow through the borrow in the argument.
+        let mut flag = Default::default();
+        MyFlag::flip(&mut flag);
+        println!("{:?}", flag); // $ type=flag:MyFlag
     }
 }