Fix crash in functions block (#550)

dotnet · May 2, 2019 · 3fe99e9 · 3fe99e9
1 parent 8168d8e
commit 3fe99e9
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 21 deletions.
diff --git a/src/Razor/src/Microsoft.AspNetCore.Razor.Language/Extensions/FunctionsDirectivePass.cs b/src/Razor/src/Microsoft.AspNetCore.Razor.Language/Extensions/FunctionsDirectivePass.cs
@@ -18,28 +18,27 @@ protected override void ExecuteCore(RazorCodeDocument codeDocument, DocumentInte
                 return;
             }
 
-            var directiveNodes = new List<IntermediateNode>();
-            foreach (var functions in documentNode.FindDirectiveReferences(FunctionsDirective.Directive))
-            {
-                directiveNodes.Add(functions.Node);
-            }
+            var directiveNodes = new List<IntermediateNodeReference>();
+            directiveNodes.AddRange(documentNode.FindDirectiveReferences(FunctionsDirective.Directive));
 
             if (FileKinds.IsComponent(codeDocument.GetFileKind()))
             {
-                foreach (var code in documentNode.FindDirectiveReferences(ComponentCodeDirective.Directive))
-                {
-                    directiveNodes.Add(code.Node);
-                }
+                directiveNodes.AddRange(documentNode.FindDirectiveReferences(ComponentCodeDirective.Directive));
             }
 
             // Now we have all the directive nodes, we want to add them to the end of the class node in document order.
-            var orderedDirectives = directiveNodes.OrderBy(n => n.Source?.AbsoluteIndex);
-            foreach (var node in orderedDirectives)
+            var orderedDirectives = directiveNodes.OrderBy(n => n.Node.Source?.AbsoluteIndex);
+            foreach (var directiveReference in orderedDirectives)
             {
+                var node = directiveReference.Node;
                 for (var i = 0; i < node.Children.Count; i++)
                 {
                     @class.Children.Add(node.Children[i]);
                 }
+
+                // We don't want to keep the original directive node around anymore.
+                // Otherwise this can cause unintended side effects in the subsequent passes.
+                directiveReference.Remove();
             }
         }
     }

diff --git a/src/Razor/test/RazorLanguage.Test/Extensions/FunctionsDirectivePassTest.cs b/src/Razor/test/RazorLanguage.Test/Extensions/FunctionsDirectivePassTest.cs
@@ -72,9 +72,7 @@ public void Execute_AddsStatementsToClassLevel()
                 node => CSharpCode(" var value = true; ", node));
 
             var method = @class.Children[0];
-            Children(
-                method,
-                node => Assert.IsType<DirectiveIntermediateNode>(node));
+            Assert.Empty(method.Children);
         }
 
         [Fact]
@@ -113,9 +111,7 @@ public void Execute_ComponentCodeDirective_AddsStatementsToClassLevel()
                 node => CSharpCode(" var value = true; ", node));
 
             var method = @class.Children[0];
-            Children(
-                method,
-                node => Assert.IsType<DirectiveIntermediateNode>(node));
+            Assert.Empty(method.Children);
         }
 
         [Fact]
@@ -161,10 +157,7 @@ public void Execute_FunctionsAndComponentCodeDirective_AddsStatementsToClassLeve
             var method = @class.Children[0];
             Children(
                 method,
-                node => Assert.IsType<HtmlContentIntermediateNode>(node),
-                node => Assert.IsType<DirectiveIntermediateNode>(node),
-                node => Assert.IsType<DirectiveIntermediateNode>(node),
-                node => Assert.IsType<DirectiveIntermediateNode>(node));
+                node => Assert.IsType<HtmlContentIntermediateNode>(node));
         }
 
         private static DocumentIntermediateNode Lower(RazorCodeDocument codeDocument, RazorProjectEngine projectEngine)