public override Expression VisitSqlRoot()

in src/Microsoft.Health.Fhir.SqlServer/Features/Search/Expressions/Visitors/PartitionEliminationRewriter.cs [72:184]
84 lines of code
18 McCabe index (conditional complexity)

        public override Expression VisitSqlRoot(SqlRootExpression expression, object context)
        {
            if (_schemaInformation.Current < SchemaVersionConstants.PartitionedTables)
            {
                return expression;
            }

            // Look for primary key continuation token (PrimaryKeyParameter) or _type parameters

            int primaryKeyValueIndex = -1;
            bool hasTypeRestriction = false;
            for (var i = 0; i < expression.ResourceTableExpressions.Count; i++)
            {
                SearchParameterInfo parameter = expression.ResourceTableExpressions[i].Parameter;

                if (ReferenceEquals(parameter, SqlSearchParameters.PrimaryKeyParameter))
                {
                    primaryKeyValueIndex = i;
                }
                else if (ReferenceEquals(parameter, _resourceTypeSearchParameter))
                {
                    hasTypeRestriction = true;
                }
            }

            if (primaryKeyValueIndex < 0)
            {
                // no continuation token

                if (hasTypeRestriction)
                {
                    // This is already constrained to be one or more resource types.
                    return expression;
                }

                // Explicitly allow all resource types. SQL tends to create far better query plans than when there is no filter on ResourceTypeId.

                var updatedResourceTableExpressions = new List<SearchParameterExpressionBase>(expression.ResourceTableExpressions.Count + 1);
                updatedResourceTableExpressions.AddRange(expression.ResourceTableExpressions);
                updatedResourceTableExpressions.Add(AllTypesExpression);

                return new SqlRootExpression(expression.SearchParamTableExpressions, updatedResourceTableExpressions);
            }

            // There is a primary key continuation token.
            // Now look at the _type restrictions to construct a PrimaryKeyRange
            // that has only the allowed types.

            var primaryKeyParameter = (SearchParameterExpression)expression.ResourceTableExpressions[primaryKeyValueIndex];

            (short? singleAllowedResourceTypeId, BitArray allowedTypes) = TypeConstraintVisitor.Instance.Visit(expression, _model);

            var existingPrimaryKeyBinaryExpression = (BinaryExpression)primaryKeyParameter.Expression;
            var existingPrimaryKeyValue = (PrimaryKeyValue)existingPrimaryKeyBinaryExpression.Value;

            SearchParameterExpression newSearchParameterExpression;
            if (singleAllowedResourceTypeId != null || allowedTypes == null)
            {
                // we'll keep the existing _type parameter and just need to add a ResourceSurrogateId expression
                newSearchParameterExpression = Expression.SearchParameter(
                    SqlSearchParameters.ResourceSurrogateIdParameter,
                    new BinaryExpression(existingPrimaryKeyBinaryExpression.BinaryOperator, SqlFieldName.ResourceSurrogateId, null, existingPrimaryKeyValue.ResourceSurrogateId));
            }
            else
            {
                // Intersect allowed types with the direction of primary key parameter
                // e.g. if >, then eliminate all types that are <=
                switch (existingPrimaryKeyBinaryExpression.BinaryOperator)
                {
                    case BinaryOperator.GreaterThan:
                        for (int i = existingPrimaryKeyValue.ResourceTypeId; i >= 0; i--)
                        {
                            allowedTypes[i] = false;
                        }

                        break;
                    case BinaryOperator.LessThan:
                        for (int i = existingPrimaryKeyValue.ResourceTypeId; i < allowedTypes.Length; i++)
                        {
                            allowedTypes[i] = false;
                        }

                        break;
                    default:
                        throw new InvalidOperationException($"Unexpected operator {existingPrimaryKeyBinaryExpression.BinaryOperator}");
                }

                newSearchParameterExpression = Expression.SearchParameter(
                    primaryKeyParameter.Parameter,
                    new BinaryExpression(
                        existingPrimaryKeyBinaryExpression.BinaryOperator,
                        existingPrimaryKeyBinaryExpression.FieldName,
                        null,
                        new PrimaryKeyRange(existingPrimaryKeyValue, allowedTypes)));
            }

            var newResourceTableExpressions = new List<SearchParameterExpressionBase>();
            for (var i = 0; i < expression.ResourceTableExpressions.Count; i++)
            {
                if (i == primaryKeyValueIndex || // eliminate the existing primaryKey expression
                    (singleAllowedResourceTypeId == null && // if there are many possible types, the PrimaryKeyRange expression will already be constrained to those types
                     expression.ResourceTableExpressions[i] is SearchParameterExpression searchParameterExpression &&
                     searchParameterExpression.Parameter.Name == SearchParameterNames.ResourceType))
                {
                    continue;
                }

                newResourceTableExpressions.Add(expression.ResourceTableExpressions[i]);
            }

            newResourceTableExpressions.Add(newSearchParameterExpression);
            return new SqlRootExpression(expression.SearchParamTableExpressions, newResourceTableExpressions);
        }