Add support for $count segments within $expand

src/Microsoft.AspNet.OData.Shared/ExpressionHelpers.cs

@@ -75,6 +75,22 @@ public static Expression Take(Expression source, int count, Type elementType, bo
             return takeQuery;
         }
 
+        public static Expression Count(Expression source, Type elementType)
+        {
+            MethodInfo countMethod;
+            if (typeof(IQueryable).IsAssignableFrom(source.Type))
+            {
+                countMethod = ExpressionHelperMethods.QueryableCountGeneric.MakeGenericMethod(elementType);
+            }
+            else
+            {
+                countMethod = ExpressionHelperMethods.EnumerableCountGeneric.MakeGenericMethod(elementType);
+            }
+
+            Expression countExpression = Expression.Call(null, countMethod, new[] { source });
+            return countExpression;
+        }
+
         public static Expression OrderByPropertyExpression(
             Expression source, 
             string propertyName, 

src/Microsoft.AspNet.OData.Shared/Formatter/Serialization/ODataResourceSerializer.cs

@@ -638,6 +638,7 @@ await writer.WriteEntityReferenceLinkAsync(new ODataEntityReferenceLink
                         await WriteComplexPropertiesAsync(selectExpandNode, resourceContext, writer);
                         await WriteDynamicComplexPropertiesAsync(resourceContext, writer);
                         await WriteNavigationLinksAsync(selectExpandNode, resourceContext, writer);
+                        await WriteExpandedCountPropertiesAsync(selectExpandNode, resourceContext, writer);
                         await WriteExpandedNavigationPropertiesAsync(selectExpandNode, resourceContext, writer);
                         await WriteReferencedNavigationPropertiesAsync(selectExpandNode, resourceContext, writer);
                         await writer.WriteEndAsync();
@@ -1396,6 +1397,32 @@ private async Task WriteExpandedNavigationPropertiesAsync(SelectExpandNode selec
             }
         }
 
+        private async Task WriteExpandedCountPropertiesAsync(SelectExpandNode selectExpandNode, ResourceContext resourceContext, ODataWriter writer)
+        {
+            Contract.Assert(resourceContext != null);
+            Contract.Assert(writer != null);
+
+            IDictionary<IEdmNavigationProperty, ExpandedCountSelectItem> navigationPropertiesToExpand = selectExpandNode.ExpandedCountProperties;
+            if (navigationPropertiesToExpand == null)
+            {
+                return;
+            }
+
+            foreach (KeyValuePair<IEdmNavigationProperty, ExpandedCountSelectItem> navPropertyToExpand in navigationPropertiesToExpand)
+            {
+                IEdmNavigationProperty navigationProperty = navPropertyToExpand.Key;
+
+                object propertyValue = resourceContext.GetPropertyValue(navigationProperty.Name);
+
+                // We should be able to write something like this:
+                // JsonWriter.WriteName(navigationProperty.Name + "@odata.count");
+                // JsonWriter.WriteValue(propertyValue)
+
+                // Or Even better have one method in ODL that we can call and write the odata.count
+                await Task.CompletedTask;
+            }
+        }
+
         private void WriteReferencedNavigationProperties(SelectExpandNode selectExpandNode, ResourceContext resourceContext, ODataWriter writer)
         {
             Contract.Assert(resourceContext != null);

src/Microsoft.AspNet.OData.Shared/Formatter/Serialization/SelectExpandNode.cs

@@ -148,6 +148,12 @@ public ISet<IEdmStructuralProperty> SelectedComplexProperties
         /// </summary>
         public IDictionary<IEdmNavigationProperty, ExpandedNavigationSelectItem> ExpandedProperties { get; internal set; }
 
+        /// <summary>
+        /// Gets the list of EDM navigation properties odata.count values to be added in the response.
+        /// It could be null if no navigation property with a $count segment.
+        /// </summary>
+        public IDictionary<IEdmNavigationProperty, ExpandedCountSelectItem> ExpandedCountProperties { get; internal set; }
+
         /// <summary>
         /// Gets the list of EDM navigation properties to be referenced in the response along with the nested query options embedded in the expand.
         /// It could be null if no navigation property to reference.
@@ -367,6 +373,19 @@ private void BuildExpandItem(ExpandedReferenceSelectItem expandReferenceItem,
 
                 if (structuralTypeInfo.IsNavigationPropertyDefined(firstNavigationSegment.NavigationProperty))
                 {
+                    // $expand=..../nav/$count
+                    ExpandedCountSelectItem expandedCount = expandReferenceItem as ExpandedCountSelectItem;
+                    if (expandedCount != null)
+                    {
+                        if (ExpandedCountProperties == null)
+                        {
+                            ExpandedCountProperties = new Dictionary<IEdmNavigationProperty, ExpandedCountSelectItem>();
+                        }
+
+                        ExpandedCountProperties[firstNavigationSegment.NavigationProperty] = expandedCount;
+                        return;
+                    }
+
                     // It's not allowed to have mulitple navigation expanded or referenced.
                     // for example: "$expand=nav($top=2),nav($skip=3)" is not allowed and will be merged (or throw exception) at ODL side.
                     ExpandedNavigationSelectItem expanded = expandReferenceItem as ExpandedNavigationSelectItem;

src/Microsoft.AspNet.OData.Shared/Query/Expressions/SelectExpandBinder.cs

@@ -596,18 +596,8 @@ private Expression CreateTotalCountExpression(Expression source, bool? countOpti
                 return countExpression;
             }
 
-            MethodInfo countMethod;
-            if (typeof(IQueryable).IsAssignableFrom(source.Type))
-            {
-                countMethod = ExpressionHelperMethods.QueryableCountGeneric.MakeGenericMethod(elementType);
-            }
-            else
-            {
-                countMethod = ExpressionHelperMethods.EnumerableCountGeneric.MakeGenericMethod(elementType);
-            }
-
             // call Count() method.
-            countExpression = Expression.Call(null, countMethod, new[] { source });
+            countExpression = ExpressionHelpers.Count(source, elementType);
 
             if (_settings.HandleNullPropagation == HandleNullPropagationOption.True)
             {
@@ -635,7 +625,7 @@ private Expression BuildPropertyContainer(Expression source, IEdmStructuredType
             {
                 foreach (var propertyToExpand in propertiesToExpand)
                 {
-                    // $expand=abc or $expand=abc/$ref
+                    // $expand=abc or $expand=abc/$ref or $expand=abc/$count
                     BuildExpandedProperty(source, structuredType, propertyToExpand.Key, propertyToExpand.Value, includedProperties);
                 }
             }
@@ -714,11 +704,23 @@ internal void BuildExpandedProperty(Expression source, IEdmStructuredType struct
             Expression countExpression = CreateTotalCountExpression(propertyValue, expandedItem.CountOption);
 
             int? modelBoundPageSize = querySettings == null ? null : querySettings.PageSize;
-            propertyValue = ProjectAsWrapper(propertyValue, subSelectExpandClause, edmEntityType, expandedItem.NavigationSource,
+
+            if(expandedItem is ExpandedCountSelectItem)
+            {
+                Type elementType;
+                if (TypeHelper.IsCollection(propertyValue.Type, out elementType))
+                {
+                    propertyValue = ExpressionHelpers.Count(propertyValue, elementType);
+                }
+            }
+            else
+            {
+                propertyValue = ProjectAsWrapper(propertyValue, subSelectExpandClause, edmEntityType, expandedItem.NavigationSource,
                 expandedItem.OrderByOption, // $orderby=...
                 expandedItem.TopOption, // $top=...
                 expandedItem.SkipOption, // $skip=...
                 modelBoundPageSize);
+            }
 
             NamedPropertyExpression propertyExpression = new NamedPropertyExpression(propertyName, propertyValue);
             if (subSelectExpandClause != null)
@@ -890,6 +892,13 @@ private static SelectExpandClause GetOrCreateSelectExpandClause(IEdmNavigationPr
                 return expandNavigationSelectItem.SelectAndExpand;
             }
 
+            // for $expand=.../$count, return null since we cannot have a select/expand after $count segment
+            ExpandedCountSelectItem expandedCountSelectItem = expandedItem as ExpandedCountSelectItem;
+            if (expandedCountSelectItem != null)
+            {
+                return null;
+            }
+
             // for $expand=..../$ref, just includes the keys properties
             IList<SelectItem> selectItems = new List<SelectItem>();
             foreach (IEdmStructuralProperty keyProperty in navigationProperty.ToEntityType().Key())

test/E2ETest/Microsoft.Test.E2E.AspNet.OData/QueryComposition/SelectExpandTests.cs

@@ -670,6 +670,52 @@ public async Task SelectWithByteArrayAndCharArrayAndIntArrayAndDoubleArrayWorks(
             JArray doubleData = (JArray)result["DoubleData"];
             Assert.Single(doubleData); // only one item
         }
+
+        [Fact]
+        public async Task DollarCountSegmentAfterDollarExpandWorks()
+        {
+            // Arrange
+            string queryUrl = string.Format("{0}/selectexpand/SelectCustomer?$expand=SelectOrders/$count", BaseAddress);
+            HttpRequestMessage request = new HttpRequestMessage(HttpMethod.Get, queryUrl);
+            request.Headers.Accept.Add(MediaTypeWithQualityHeaderValue.Parse("application/json"));
+            HttpClient client = new HttpClient();
+
+            // Act
+            HttpResponseMessage response = await client.SendAsync(request);
+
+            // Assert
+            Assert.NotNull(response);
+            Assert.Equal(HttpStatusCode.OK, response.StatusCode);
+
+            Assert.NotNull(response.Content);
+            JObject json = await response.Content.ReadAsObject<JObject>();
+            Assert.NotNull(json);
+
+            Assert.Equal("10", (string)json["[email protected]"]);
+        }
+
+        [Fact]
+        public async Task DollarCountSegmentAfterDollarExpandWithNestedFilterWorks()
+        {
+            // Arrange
+            string queryUrl = string.Format("{0}/selectexpand/SelectCustomer?$expand=SelectOrders/$count($filter=Id gt 5)", BaseAddress);
+            HttpRequestMessage request = new HttpRequestMessage(HttpMethod.Get, queryUrl);
+            request.Headers.Accept.Add(MediaTypeWithQualityHeaderValue.Parse("application/json"));
+            HttpClient client = new HttpClient();
+
+            // Act
+            HttpResponseMessage response = await client.SendAsync(request);
+
+            // Assert
+            Assert.NotNull(response);
+            Assert.Equal(HttpStatusCode.OK, response.StatusCode);
+
+            Assert.NotNull(response.Content);
+            JObject json = await response.Content.ReadAsObject<JObject>();
+            Assert.NotNull(json);
+
+            Assert.Equal("4", (string)json["[email protected]"]);
+        }
     }
 
     public class SelectCustomerController : TestODataController

test/UnitTest/Microsoft.AspNet.OData.Test.Shared/Query/Expressions/SelectExpandBinderTest.cs

@@ -759,6 +759,62 @@ public void ProjectAsWrapper_Element_ProjectedValueContains_SelectedTypeCastSubS
             Assert.Equal("201501", addressProperties["PostCode"]);
         }
 
+        [Fact]
+        public void ProjectAsWrapper_Element_ProjectedValueContainsCount_IfDollarCountInDollarExpand()
+        {
+            // Arrange
+            string expand = "Orders/$count";
+            QueryCustomer aCustomer = new QueryCustomer
+            {
+                Orders = new[]
+                {
+                    new QueryOrder { Id = 42 },
+                    new QueryVipOrder { Id = 38 }
+                }
+            };
+            Expression source = Expression.Constant(aCustomer);
+            SelectExpandClause selectExpandClause = ParseSelectExpand(null, expand, _model, _customer, _customers);
+            Assert.NotNull(selectExpandClause);
+
+            // Act
+            Expression projection = _binder.ProjectAsWrapper(source, selectExpandClause, _customer, _customers);
+
+            // Assert
+            Assert.Equal(ExpressionType.MemberInit, projection.NodeType);
+            Assert.NotEmpty((projection as MemberInitExpression).Bindings.Where(p => p.Member.Name == "Instance"));
+            SelectExpandWrapper<QueryCustomer> customerWrapper = Expression.Lambda(projection).Compile().DynamicInvoke() as SelectExpandWrapper<QueryCustomer>;
+            var orders = customerWrapper.Container.ToDictionary(PropertyMapper)["Orders"];
+            Assert.Equal((long)2, orders);
+        }
+
+        [Fact]
+        public void ProjectAsWrapper_Element_ProjectedValueContainsCount_IfDollarCountInDollarExpand_AndNestedFilterClause()
+        {
+            // Arrange
+            string expand = "Orders/$count($filter=Id eq 42)";
+            QueryCustomer aCustomer = new QueryCustomer
+            {
+                Orders = new[]
+                {
+                    new QueryOrder { Id = 42 },
+                    new QueryVipOrder { Id = 38 }
+                }
+            };
+            Expression source = Expression.Constant(aCustomer);
+            SelectExpandClause selectExpandClause = ParseSelectExpand(null, expand, _model, _customer, _customers);
+            Assert.NotNull(selectExpandClause);
+
+            // Act
+            Expression projection = _binder.ProjectAsWrapper(source, selectExpandClause, _customer, _customers);
+
+            // Assert
+            Assert.Equal(ExpressionType.MemberInit, projection.NodeType);
+            Assert.NotEmpty((projection as MemberInitExpression).Bindings.Where(p => p.Member.Name == "Instance"));
+            SelectExpandWrapper<QueryCustomer> customerWrapper = Expression.Lambda(projection).Compile().DynamicInvoke() as SelectExpandWrapper<QueryCustomer>;
+            var orders = customerWrapper.Container.ToDictionary(PropertyMapper)["Orders"];
+            Assert.Equal((long)1, orders);
+        }
+
         [Fact]
         public void ProjectAsWrapper_Element_ProjectedValueContainsSubKeys_IfDollarRefInDollarExpand()
         {