Skip to content

Commit 84ae3af

Browse files
hvitvedhvitved
committed
Rust: Type inference for .await expressions
1 parent 24a4208 commit 84ae3af

File tree

6 files changed

+105
-9
lines changed

6 files changed

+105
-9
lines changed

‎rust/ql/lib/codeql/rust/frameworks/stdlib/Stdlib.qll

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,3 +49,19 @@ class ResultEnum extends Enum {
4949
/** Gets the `Err` variant. */
5050
Variant getErr() { result = this.getVariant("Err") }
5151
}
52+
53+
/**
54+
* The [`Future` trait][1].
55+
*
56+
* [1]: https://doc.rust-lang.org/std/future/trait.Future.html
57+
*/
58+
class FutureTrait extends Trait {
59+
FutureTrait() { this.getCanonicalPath() = "core::future::future::Future" }
60+
61+
/** Gets the `Output` associated type. */
62+
pragma[nomagic]
63+
TypeAlias getOutputType() {
64+
result = this.getAssocItemList().getAnAssocItem() and
65+
result.getName().getText() = "Output"
66+
}
67+
}

‎rust/ql/lib/codeql/rust/internal/Type.qll

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,9 @@ class TraitType extends Type, TTrait {
115115

116116
TraitType() { this = TTrait(trait) }
117117

118+
/** Gets the underlying trait. */
119+
Trait getTrait() { result = trait }
120+
118121
override StructField getStructField(string name) { none() }
119122

120123
override TupleField getTupleField(int i) { none() }

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

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -232,8 +232,12 @@ private predicate typeEquality(AstNode n1, TypePath path1, AstNode n2, TypePath
232232
n1 = n2.(ParenExpr).getExpr() and
233233
path1 = path2
234234
or
235-
n1 = n2.(BlockExpr).getStmtList().getTailExpr() and
236-
path1 = path2
235+
n2 =
236+
any(BlockExpr be |
237+
not be.isAsync() and
238+
n1 = be.getStmtList().getTailExpr() and
239+
path1 = path2
240+
)
237241
or
238242
n1 = n2.(IfExpr).getABranch() and
239243
path1 = path2
@@ -1000,6 +1004,22 @@ private StructType inferLiteralType(LiteralExpr le) {
10001004
)
10011005
}
10021006

1007+
pragma[nomagic]
1008+
private AssociatedTypeTypeParameter getFutureOutputTypeParameter() {
1009+
result.getTypeAlias() = any(FutureTrait ft).getOutputType()
1010+
}
1011+
1012+
pragma[nomagic]
1013+
private Type inferAwaitExprType(AwaitExpr ae, TypePath path) {
1014+
exists(TypePath exprPath | result = inferType(ae.getExpr(), exprPath) |
1015+
exprPath.isCons(getFutureOutputTypeParameter(), path)
1016+
or
1017+
path = exprPath and
1018+
not (exprPath.isEmpty() and result.(TraitType).getTrait() instanceof FutureTrait) and
1019+
not exprPath.isCons(getFutureOutputTypeParameter(), _)
1020+
)
1021+
}
1022+
10031023
cached
10041024
private module Cached {
10051025
private import codeql.rust.internal.CachedStages
@@ -1216,6 +1236,8 @@ private module Cached {
12161236
or
12171237
result = inferLiteralType(n) and
12181238
path.isEmpty()
1239+
or
1240+
result = inferAwaitExprType(n, path)
12191241
}
12201242
}
12211243

@@ -1232,7 +1254,7 @@ private module Debug {
12321254
exists(string filepath, int startline, int startcolumn, int endline, int endcolumn |
12331255
result.getLocation().hasLocationInfo(filepath, startline, startcolumn, endline, endcolumn) and
12341256
filepath.matches("%/main.rs") and
1235-
startline = 948
1257+
startline = [1276 .. 1278]
12361258
)
12371259
}
12381260

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

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ abstract class TypeMention extends AstNode {
1515

1616
/** Gets the sub mention at `path`. */
1717
pragma[nomagic]
18-
private TypeMention getMentionAt(TypePath path) {
18+
TypeMention getMentionAt(TypePath path) {
1919
path.isEmpty() and
2020
result = this
2121
or
@@ -44,6 +44,19 @@ class TypeReprMention extends TypeMention, TypeRepr {
4444
result = this.(PathTypeRepr).getPath().(PathMention).getTypeArgument(i)
4545
}
4646

47+
/**
48+
* Gets an [`impl Trait`][1] bound. For example, if this mention is
49+
* `impl A + B` then this predicate returns `A` and `B`.
50+
*
51+
* [1] https://doc.rust-lang.org/book/ch10-02-traits.html#traits-as-parameters
52+
*/
53+
private PathMention getAnImplTraitBound() {
54+
exists(TypeBoundList tbl |
55+
tbl = this.(ImplTraitTypeRepr).getTypeBoundList() and
56+
result = tbl.getABound().getTypeRepr().(PathTypeRepr).getPath()
57+
)
58+
}
59+
4760
override Type resolveType() {
4861
this instanceof ArrayTypeRepr and
4962
result = TArrayType()
@@ -52,12 +65,17 @@ class TypeReprMention extends TypeMention, TypeRepr {
5265
result = TRefType()
5366
or
5467
result = this.(PathTypeRepr).getPath().(PathMention).resolveType()
68+
or
69+
result = this.getAnImplTraitBound().resolveType()
5570
}
5671

5772
override Type resolveTypeAt(TypePath path) {
5873
result = this.(PathTypeRepr).getPath().(PathMention).resolveTypeAt(path)
5974
or
75+
result = this.getAnImplTraitBound().resolveTypeAt(path)
76+
or
6077
not exists(this.(PathTypeRepr).getPath()) and
78+
not exists(this.getAnImplTraitBound()) and
6179
result = super.resolveTypeAt(path)
6280
}
6381
}
@@ -71,6 +89,30 @@ abstract class PathMention extends TypeMention, Path {
7189
override TypeMention getTypeArgument(int i) {
7290
result = this.getSegment().getGenericArgList().getTypeArg(i)
7391
}
92+
93+
/** Gets the type argument for the associated type `alias`, if any. */
94+
pragma[nomagic]
95+
private TypeRepr getAnAssocTypeArgument(TypeAlias alias) {
96+
exists(TraitItemNode trait, string name, AssocTypeArg arg |
97+
alias = trait.getAnAssocItem() and
98+
name = alias.getName().getText() and
99+
trait = resolvePath(this) and
100+
arg = this.getSegment().getGenericArgList().getAGenericArg() and
101+
result = arg.getTypeRepr() and
102+
name = arg.getIdentifier().getText()
103+
)
104+
}
105+
106+
override TypeMention getMentionAt(TypePath path) {
107+
result = super.getMentionAt(path)
108+
or
109+
exists(TypeAlias alias, AssociatedTypeTypeParameter tp, TypeMention arg, TypePath suffix |
110+
arg = this.getAnAssocTypeArgument(alias) and
111+
result = arg.getMentionAt(suffix) and
112+
path = TypePath::cons(tp, suffix) and
113+
tp.getTypeAlias() = alias
114+
)
115+
}
74116
}
75117

76118
class NonAliasPathMention extends PathMention {

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1273,9 +1273,9 @@ mod async_ {
12731273
}
12741274

12751275
pub async fn f() {
1276-
f1().await.f(); // $ MISSING: method=S1f
1277-
f2().await.f(); // $ MISSING: method=S1f
1278-
f3().await.f(); // $ MISSING: method=S1f
1276+
f1().await.f(); // $ method=S1f
1277+
f2().await.f(); // $ method=S1f
1278+
f3().await.f(); // $ method=S1f
12791279
}
12801280
}
12811281

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

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1604,8 +1604,10 @@ inferType
16041604
| main.rs:1248:18:1248:21 | SelfParam | | main.rs:1245:5:1245:14 | S1 |
16051605
| main.rs:1251:25:1253:5 | { ... } | | main.rs:1245:5:1245:14 | S1 |
16061606
| main.rs:1252:9:1252:10 | S1 | | main.rs:1245:5:1245:14 | S1 |
1607-
| main.rs:1255:41:1259:5 | { ... } | | main.rs:1245:5:1245:14 | S1 |
1608-
| main.rs:1256:9:1258:9 | { ... } | | main.rs:1245:5:1245:14 | S1 |
1607+
| main.rs:1255:41:1259:5 | { ... } | | file:///RUSTUP_HOME/toolchain/lib/rustlib/src/rust/library/core/src/future/future.rs:7:1:105:1 | trait Future |
1608+
| main.rs:1255:41:1259:5 | { ... } | Output | main.rs:1245:5:1245:14 | S1 |
1609+
| main.rs:1256:9:1258:9 | { ... } | | file:///RUSTUP_HOME/toolchain/lib/rustlib/src/rust/library/core/src/future/future.rs:7:1:105:1 | trait Future |
1610+
| main.rs:1256:9:1258:9 | { ... } | Output | main.rs:1245:5:1245:14 | S1 |
16091611
| main.rs:1257:13:1257:14 | S1 | | main.rs:1245:5:1245:14 | S1 |
16101612
| main.rs:1266:17:1266:46 | SelfParam | | file:///RUSTUP_HOME/toolchain/lib/rustlib/src/rust/library/core/src/pin.rs:934:1:1104:1 | Pin |
16111613
| main.rs:1266:17:1266:46 | SelfParam | Ptr | file://:0:0:0:0 | & |
@@ -1617,9 +1619,20 @@ inferType
16171619
| main.rs:1267:13:1267:38 | ...::Ready(...) | | file:///RUSTUP_HOME/toolchain/lib/rustlib/src/rust/library/core/src/task/poll.rs:6:1:28:1 | Poll |
16181620
| main.rs:1267:13:1267:38 | ...::Ready(...) | T | main.rs:1245:5:1245:14 | S1 |
16191621
| main.rs:1267:36:1267:37 | S1 | | main.rs:1245:5:1245:14 | S1 |
1622+
| main.rs:1271:41:1273:5 | { ... } | | file:///RUSTUP_HOME/toolchain/lib/rustlib/src/rust/library/core/src/future/future.rs:7:1:105:1 | trait Future |
16201623
| main.rs:1271:41:1273:5 | { ... } | | main.rs:1261:5:1261:14 | S2 |
1624+
| main.rs:1271:41:1273:5 | { ... } | Output | main.rs:1245:5:1245:14 | S1 |
1625+
| main.rs:1272:9:1272:10 | S2 | | file:///RUSTUP_HOME/toolchain/lib/rustlib/src/rust/library/core/src/future/future.rs:7:1:105:1 | trait Future |
16211626
| main.rs:1272:9:1272:10 | S2 | | main.rs:1261:5:1261:14 | S2 |
1627+
| main.rs:1272:9:1272:10 | S2 | Output | main.rs:1245:5:1245:14 | S1 |
16221628
| main.rs:1276:9:1276:12 | f1(...) | | main.rs:1245:5:1245:14 | S1 |
1629+
| main.rs:1276:9:1276:18 | await ... | | main.rs:1245:5:1245:14 | S1 |
1630+
| main.rs:1277:9:1277:12 | f2(...) | | file:///RUSTUP_HOME/toolchain/lib/rustlib/src/rust/library/core/src/future/future.rs:7:1:105:1 | trait Future |
1631+
| main.rs:1277:9:1277:12 | f2(...) | Output | main.rs:1245:5:1245:14 | S1 |
1632+
| main.rs:1277:9:1277:18 | await ... | | main.rs:1245:5:1245:14 | S1 |
1633+
| main.rs:1278:9:1278:12 | f3(...) | | file:///RUSTUP_HOME/toolchain/lib/rustlib/src/rust/library/core/src/future/future.rs:7:1:105:1 | trait Future |
1634+
| main.rs:1278:9:1278:12 | f3(...) | Output | main.rs:1245:5:1245:14 | S1 |
1635+
| main.rs:1278:9:1278:18 | await ... | | main.rs:1245:5:1245:14 | S1 |
16231636
| main.rs:1284:5:1284:20 | ...::f(...) | | main.rs:67:5:67:21 | Foo |
16241637
| main.rs:1285:5:1285:60 | ...::g(...) | | main.rs:67:5:67:21 | Foo |
16251638
| main.rs:1285:20:1285:38 | ...::Foo {...} | | main.rs:67:5:67:21 | Foo |

0 commit comments

Comments
 (0)