Handle await expressions in find references.

ICSharpCode.NRefactory.CSharp/Resolver/FindReferences.cs

@@ -244,6 +244,8 @@ namespace ICSharpCode.NRefactory.CSharp.Resolver
 					scope = FindMemberReferences(entity, m => new FindPropertyReferences((IProperty)m));
 					if (entity.Name == "Current")
 						additionalScope = FindEnumeratorCurrentReferences((IProperty)entity);
+					else if (entity.Name == "IsCompleted")
+						additionalScope = FindAwaiterIsCompletedReferences((IProperty)entity);
 					break;
 				case EntityType.Event:
 					scope = FindMemberReferences(entity, m => new FindEventReferences((IEvent)m));
@@ -662,6 +664,15 @@ namespace ICSharpCode.NRefactory.CSharp.Resolver
 				});
 		}
 
+		SearchScope FindAwaiterIsCompletedReferences(IProperty property)
+		{
+			return new SearchScope(
+				delegate(ICompilation compilation) {
+					IProperty imported = compilation.Import(property);
+					return imported != null ? new FindAwaiterIsCompletedReferencesNavigator(imported) : null;
+				});
+		}
+		
 		sealed class FindEnumeratorCurrentReferencesNavigator : FindReferenceNavigator
 		{
 			IProperty property;
@@ -682,6 +693,27 @@ namespace ICSharpCode.NRefactory.CSharp.Resolver
 				return ferr != null && ferr.CurrentProperty != null && findReferences.IsMemberMatch(property, ferr.CurrentProperty, true);
 			}
 		}
+
+		sealed class FindAwaiterIsCompletedReferencesNavigator : FindReferenceNavigator
+		{
+			IProperty property;
+			
+			public FindAwaiterIsCompletedReferencesNavigator(IProperty property)
+			{
+				this.property = property;
+			}
+			
+			internal override bool CanMatch(AstNode node)
+			{
+				return node is UnaryOperatorExpression;
+			}
+			
+			internal override bool IsMatch(ResolveResult rr)
+			{
+				AwaitResolveResult arr = rr as AwaitResolveResult;
+				return arr != null && arr.IsCompletedProperty != null && findReferences.IsMemberMatch(property, arr.IsCompletedProperty, true);
+			}
+		}
 		#endregion
 		
 		#region Find Method References
@@ -724,6 +756,11 @@ namespace ICSharpCode.NRefactory.CSharp.Resolver
 				case "MoveNext":
 					specialNodeType = typeof(ForeachStatement);
 					break;
+				case "GetAwaiter":
+				case "GetResult":
+				case "OnCompleted":
+					specialNodeType = typeof(UnaryOperatorExpression);
+					break;
 				default:
 					specialNodeType = null;
 					break;
@@ -794,6 +831,12 @@ namespace ICSharpCode.NRefactory.CSharp.Resolver
 						return IsMatch(ferr.GetEnumeratorCall)
 							|| (ferr.MoveNextMethod != null && findReferences.IsMemberMatch(method, ferr.MoveNextMethod, true));
 					}
+					var arr = rr as AwaitResolveResult;
+					if (arr != null) {
+						return IsMatch(arr.GetAwaiterInvocation)
+						    || (arr.GetResultMethod != null && findReferences.IsMemberMatch(method, arr.GetResultMethod, true))
+						    || (arr.OnCompletedMethod != null && findReferences.IsMemberMatch(method, arr.OnCompletedMethod, true));
+					}
 				}
 				var mrr = rr as MemberResolveResult;
 				return mrr != null && findReferences.IsMemberMatch(method, mrr.Member, mrr.IsVirtualCall);

ICSharpCode.NRefactory.Tests/CSharp/Resolver/FindReferencesTest.cs

@@ -211,5 +211,63 @@ class Calls {
 			Assert.IsTrue(actual.Any(r => r.StartLocation.Line == 9 && r is InvocationExpression));
 		}
 		#endregion
+
+		#region Await
+		const string awaitTest = @"using System;
+class MyAwaiter {
+	public bool IsCompleted { get { return false; } }
+	public void OnCompleted(Action continuation) {}
+	public int GetResult() { return 0; }
+}
+class MyAwaitable {
+	public MyAwaiter GetAwaiter() { return null; }
+}
+public class C {
+	public async void M() {
+		MyAwaitable x = null;
+		int i = await x;
+	}
+}";
+
+		[Test]
+		public void GetAwaiterReferenceInAwaitExpressionIsFound() {
+			Init(awaitTest);
+			var test = compilation.MainAssembly.TopLevelTypeDefinitions.Single(t => t.Name == "MyAwaitable");
+			var method = test.Methods.Single(m => m.Name == "GetAwaiter");
+			var actual = FindReferences(method).ToList();
+			Assert.IsTrue(actual.Any(r => r.StartLocation.Line == 8 && r is MethodDeclaration));
+			Assert.IsTrue(actual.Any(r => r.StartLocation.Line == 13 && r is UnaryOperatorExpression));
+		}
+
+		[Test]
+		public void GetResultReferenceInAwaitExpressionIsFound() {
+			Init(awaitTest);
+			var test = compilation.MainAssembly.TopLevelTypeDefinitions.Single(t => t.Name == "MyAwaiter");
+			var method = test.Methods.Single(m => m.Name == "GetResult");
+			var actual = FindReferences(method).ToList();
+			Assert.IsTrue(actual.Any(r => r.StartLocation.Line == 5 && r is MethodDeclaration));
+			Assert.IsTrue(actual.Any(r => r.StartLocation.Line == 13 && r is UnaryOperatorExpression));
+		}
+
+		[Test]
+		public void OnCompletedReferenceInAwaitExpressionIsFound() {
+			Init(awaitTest);
+			var test = compilation.MainAssembly.TopLevelTypeDefinitions.Single(t => t.Name == "MyAwaiter");
+			var method = test.Methods.Single(m => m.Name == "OnCompleted");
+			var actual = FindReferences(method).ToList();
+			Assert.IsTrue(actual.Any(r => r.StartLocation.Line == 4 && r is MethodDeclaration));
+			Assert.IsTrue(actual.Any(r => r.StartLocation.Line == 13 && r is UnaryOperatorExpression));
+		}
+
+		[Test]
+		public void IsCompletedReferenceInAwaitExpressionIsFound() {
+			Init(awaitTest);
+			var test = compilation.MainAssembly.TopLevelTypeDefinitions.Single(t => t.Name == "MyAwaiter");
+			var property = test.Properties.Single(m => m.Name == "IsCompleted");
+			var actual = FindReferences(property).ToList();
+			Assert.IsTrue(actual.Any(r => r.StartLocation.Line == 3 && r is PropertyDeclaration));
+			Assert.IsTrue(actual.Any(r => r.StartLocation.Line == 13 && r is UnaryOperatorExpression));
+		}
+		#endregion
 	}
 }