From dbe271949cc62515e3727c85aa34c5c4604f306f Mon Sep 17 00:00:00 2001
From: Aditya Sarpotdar <adityasa@microsoft.com>
Date: Thu, 27 Apr 2023 14:55:28 +0000
Subject: [PATCH] Initial commit

---
 .../src/Linq/TranslationContext.cs            |  6 +++
 .../BaselineTest/BaselineTests.cs             |  4 +-
 ...nqGeneralBaselineTests.TestLambdaReuse.xml | 38 +++++++++++++++++++
 .../LinqGeneralBaselineTests.cs               | 21 +++++++++-
 ...icrosoft.Azure.Cosmos.EmulatorTests.csproj |  6 +++
 5 files changed, 72 insertions(+), 3 deletions(-)
 create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/BaselineTest/TestBaseline/LinqGeneralBaselineTests.TestLambdaReuse.xml

diff --git a/Microsoft.Azure.Cosmos/src/Linq/TranslationContext.cs b/Microsoft.Azure.Cosmos/src/Linq/TranslationContext.cs
index 9205052604..8fc95d8701 100644
--- a/Microsoft.Azure.Cosmos/src/Linq/TranslationContext.cs
+++ b/Microsoft.Azure.Cosmos/src/Linq/TranslationContext.cs
@@ -119,6 +119,7 @@ public void PopParameter()
         {
             ParameterExpression last = this.lambdaParametersStack[this.lambdaParametersStack.Count - 1];
             this.lambdaParametersStack.RemoveAt(this.lambdaParametersStack.Count - 1);
+            this.substitutions.Remove(last);
         }
 
         /// <summary>
@@ -337,6 +338,11 @@ public Expression Lookup(ParameterExpression parameter)
             return null;
         }
 
+        internal void Remove(ParameterExpression parameter)
+        {
+            this.substitutionTable.Remove(parameter);
+        }
+
         public const string InputParameterName = "root";
     }
 
diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/BaselineTest/BaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/BaselineTest/BaselineTests.cs
index e9fbc54267..62f20f1fe4 100644
--- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/BaselineTest/BaselineTests.cs
+++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/BaselineTest/BaselineTests.cs
@@ -153,8 +153,8 @@ public void ExecuteTestSuite(IEnumerable<TInput> inputs, [CallerMemberName] stri
                     Please run the ..\azure-cosmos-dotnet-v3\UpdateContracts.ps1 script to update the baselines.
                     Expected: {baselineTextSuffix},
                     Actual:   {outputTextSuffix},
-                    OutputPath: {outputPath},
-                    BaselinePath: {baselinePath}");
+                    OutputPath: {Path.GetFullPath(outputPath)},
+                    BaselinePath: {Path.GetFullPath(baselinePath)}");
         }
 
         /// <summary>
diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/BaselineTest/TestBaseline/LinqGeneralBaselineTests.TestLambdaReuse.xml b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/BaselineTest/TestBaseline/LinqGeneralBaselineTests.TestLambdaReuse.xml
new file mode 100644
index 0000000000..bd2e168e76
--- /dev/null
+++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/BaselineTest/TestBaseline/LinqGeneralBaselineTests.TestLambdaReuse.xml
@@ -0,0 +1,38 @@
+<Results>
+  <Result>
+    <Input>
+      <Description><![CDATA[Where -> Where with same predicate instance]]></Description>
+      <Expression><![CDATA[query.Where(DisplayClass.predicate).Where(DisplayClass.predicate)]]></Expression>
+    </Input>
+    <Output>
+      <SqlQuery><![CDATA[
+SELECT VALUE root 
+FROM root 
+WHERE ((root["Int"] = 5) AND (root["Int"] = 5))]]></SqlQuery>
+    </Output>
+  </Result>
+  <Result>
+    <Input>
+      <Description><![CDATA[Where -> Select with same predicate instance]]></Description>
+      <Expression><![CDATA[query.Where(DisplayClass.predicate).Select(DisplayClass.predicate)]]></Expression>
+    </Input>
+    <Output>
+      <SqlQuery><![CDATA[
+SELECT VALUE (root["Int"] = 5) 
+FROM root 
+WHERE (root["Int"] = 5)]]></SqlQuery>
+    </Output>
+  </Result>
+  <Result>
+    <Input>
+      <Description><![CDATA[Where -> Select -> Where with same predicate instance]]></Description>
+      <Expression><![CDATA[query.Where(DisplayClass.predicate2).Select(c => new AnonymousType(Int = 10, Result = c)).Where(DisplayClass.predicate2)]]></Expression>
+    </Input>
+    <Output>
+      <SqlQuery><![CDATA[
+SELECT VALUE {"Int": 10, "Result": root} 
+FROM root 
+WHERE ((ToString(root) = "a") AND (ToString({"Int": 10, "Result": root}) = "a"))]]></SqlQuery>
+    </Output>
+  </Result>
+</Results>
\ No newline at end of file
diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/LinqGeneralBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/LinqGeneralBaselineTests.cs
index a9782bb26b..3fc3aff900 100644
--- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/LinqGeneralBaselineTests.cs
+++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/LinqGeneralBaselineTests.cs
@@ -1272,6 +1272,25 @@ orderby f.FamilyId descending
             this.ExecuteTestSuite(inputs);
         }
 
+        [TestMethod]
+        public void TestLambdaReuse()
+        {
+            List<LinqTestInput> inputs = new List<LinqTestInput>();
+
+            System.Linq.Expressions.Expression<Func<Family, bool>> predicate = f => f.Int == 5;
+            inputs.Add(new LinqTestInput("Where -> Where with same predicate instance", b => getQuery(b).Where(predicate).Where(predicate)));
+            inputs.Add(new LinqTestInput("Where -> Select with same predicate instance", b => getQuery(b).Where(predicate).Select(predicate)));
+
+            System.Linq.Expressions.Expression<Func<object, bool>> predicate2 = f => f.ToString() == "a";
+            inputs.Add(new LinqTestInput("Where -> Select -> Where with same predicate instance",
+                b => getQuery(b)
+                    .Where(predicate2)
+                    .Select(c => new { Int = 10, Result = c })
+                    .Where(predicate2)));
+
+            this.ExecuteTestSuite(inputs);
+        }
+
         [TestMethod]
         public void TestThenByTranslation()
         {
@@ -1922,7 +1941,7 @@ public async Task ValidateLinqQueries()
                 Gender = "female",
                 Grade = 1,
                 Pets = new List<Pet>() { pet, new Pet() { GivenName = "koko" } },
-                Things = new Dictionary<string, string>() { { "A", "B" }, { "C", "D" } },
+                Things = new Dictionary<string, string>() { { "A", "B" }, { "C", "D" } }
             };
 
             Address address = new Address { State = "NY", County = "Manhattan", City = "NY" };
diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Microsoft.Azure.Cosmos.EmulatorTests.csproj b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Microsoft.Azure.Cosmos.EmulatorTests.csproj
index a504adc855..51891d7e3f 100644
--- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Microsoft.Azure.Cosmos.EmulatorTests.csproj
+++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Microsoft.Azure.Cosmos.EmulatorTests.csproj
@@ -35,6 +35,7 @@
     <None Remove="BaselineTest\TestBaseline\EndToEndTraceWriterBaselineTests.StreamPointOperationsAsync.xml" />
     <None Remove="BaselineTest\TestBaseline\EndToEndTraceWriterBaselineTests.TypedPointOperationsAsync.xml" />
 	<None Remove="BaselineTest\TestBaseline\IndexMetricsParserBaselineTest.IndexUtilizationParse.xml" />
+	<None Remove="BaselineTest\TestBaseline\LinqGeneralBaselineTests.TestLambdaReuse.xml" />
     <None Remove="BaselineTest\TestBaseline\LinqTranslationBaselineTests.TestDateTimeJsonConverterTimezones.xml" />
     <None Remove="BaselineTest\TestBaseline\LinqTranslationBaselineTests.TestMemberAccessWithNullableTypes.xml" />
   </ItemGroup>
@@ -283,6 +284,11 @@
   <ItemGroup>
     <ProjectReference Include="..\..\src\Microsoft.Azure.Cosmos.csproj" />
   </ItemGroup>
+  <ItemGroup>
+    <Content Include="BaselineTest\TestBaseline\LinqGeneralBaselineTests.TestLambdaReuse.xml">
+      <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
+    </Content>
+  </ItemGroup>
   <ItemGroup>
     <Compile Update="Tracing\EndToEndTraceWriterBaselineTests.cs">
       <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>