Add @code directive support for Blazor

src/Razor/src/Microsoft.AspNetCore.Razor.Language/Components/ComponentCodeDirective.cs

@@ -0,0 +1,28 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+
+namespace Microsoft.AspNetCore.Razor.Language.Components
+{
+    public static class ComponentCodeDirective
+    {
+        public static readonly DirectiveDescriptor Directive = DirectiveDescriptor.CreateDirective(
+            "code",
+            DirectiveKind.CodeBlock,
+            builder =>
+            {
+                builder.Description = Resources.FunctionsDirective_Description;
+            });
+
+        public static void Register(RazorProjectEngineBuilder builder)
+        {
+            if (builder == null)
+            {
+                throw new ArgumentNullException(nameof(builder));
+            }
+
+            builder.AddDirective(Directive, FileKinds.Component);
+        }
+    }
+}

src/Razor/src/Microsoft.AspNetCore.Razor.Language/Extensions/FunctionsDirectivePass.cs

@@ -1,6 +1,9 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
+using System.Collections.Generic;
+using System.Linq;
+using Microsoft.AspNetCore.Razor.Language.Components;
 using Microsoft.AspNetCore.Razor.Language.Intermediate;
 
 namespace Microsoft.AspNetCore.Razor.Language.Extensions
@@ -15,11 +18,27 @@ protected override void ExecuteCore(RazorCodeDocument codeDocument, DocumentInte
                 return;
             }
 
+            var directiveNodes = new List<IntermediateNode>();
             foreach (var functions in documentNode.FindDirectiveReferences(FunctionsDirective.Directive))
             {
-                for (var i = 0; i < functions.Node.Children.Count; i++)
+                directiveNodes.Add(functions.Node);
+            }
+
+            if (FileKinds.IsComponent(codeDocument.GetFileKind()))
+            {
+                foreach (var code in documentNode.FindDirectiveReferences(ComponentCodeDirective.Directive))
+                {
+                    directiveNodes.Add(code.Node);
+                }
+            }
+
+            // 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)
+            {
+                for (var i = 0; i < node.Children.Count; i++)
                 {
-                    @class.Children.Add(functions.Node.Children[i]);
+                    @class.Children.Add(node.Children[i]);
                 }
             }
         }

src/Razor/src/Microsoft.AspNetCore.Razor.Language/RazorProjectEngine.cs

@@ -211,6 +211,7 @@ private static void AddComponentFeatures(RazorProjectEngineBuilder builder)
             builder.Features.Add(new ComponentImportProjectFeature());
 
             // Directives (conditional on file kind)
+            ComponentCodeDirective.Register(builder);
             ComponentInjectDirective.Register(builder);
             ComponentLayoutDirective.Register(builder);
             ComponentPageDirective.Register(builder);

src/Razor/test/RazorLanguage.Test/Extensions/FunctionsDirectivePassTest.cs

@@ -1,6 +1,7 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
+using Microsoft.AspNetCore.Razor.Language.Components;
 using Microsoft.AspNetCore.Razor.Language.Intermediate;
 using Xunit;
 using static Microsoft.AspNetCore.Razor.Language.Intermediate.IntermediateNodeAssert;
@@ -76,6 +77,96 @@ public void Execute_AddsStatementsToClassLevel()
                 node => Assert.IsType<DirectiveIntermediateNode>(node));
         }
 
+        [Fact]
+        public void Execute_ComponentCodeDirective_AddsStatementsToClassLevel()
+        {
+            // Arrange
+            var projectEngine = CreateProjectEngine(b => b.AddDirective(ComponentCodeDirective.Directive));
+            var pass = new FunctionsDirectivePass()
+            {
+                Engine = projectEngine.Engine,
+            };
+
+            var sourceDocument = TestRazorSourceDocument.Create("@code { var value = true; }");
+            var codeDocument = RazorCodeDocument.Create(sourceDocument);
+            codeDocument.SetFileKind(FileKinds.Component);
+
+            var irDocument = Lower(codeDocument, projectEngine);
+
+            // Act
+            pass.Execute(codeDocument, irDocument);
+
+            // Assert
+            Children(
+                irDocument,
+                node => Assert.IsType<NamespaceDeclarationIntermediateNode>(node));
+
+            var @namespace = irDocument.Children[0];
+            Children(
+                @namespace,
+                node => Assert.IsType<ClassDeclarationIntermediateNode>(node));
+
+            var @class = @namespace.Children[0];
+            Children(
+                @class,
+                node => Assert.IsType<MethodDeclarationIntermediateNode>(node),
+                node => CSharpCode(" var value = true; ", node));
+
+            var method = @class.Children[0];
+            Children(
+                method,
+                node => Assert.IsType<DirectiveIntermediateNode>(node));
+        }
+
+        [Fact]
+        public void Execute_FunctionsAndComponentCodeDirective_AddsStatementsToClassLevel()
+        {
+            // Arrange
+            var projectEngine = CreateProjectEngine(b => b.AddDirective(ComponentCodeDirective.Directive));
+            var pass = new FunctionsDirectivePass()
+            {
+                Engine = projectEngine.Engine,
+            };
+
+            var sourceDocument = TestRazorSourceDocument.Create(@"
+@functions { var value1 = true; }
+@code { var value2 = true; }
+@functions { var value3 = true; }");
+            var codeDocument = RazorCodeDocument.Create(sourceDocument);
+            codeDocument.SetFileKind(FileKinds.Component);
+
+            var irDocument = Lower(codeDocument, projectEngine);
+
+            // Act
+            pass.Execute(codeDocument, irDocument);
+
+            // Assert
+            Children(
+                irDocument,
+                node => Assert.IsType<NamespaceDeclarationIntermediateNode>(node));
+
+            var @namespace = irDocument.Children[0];
+            Children(
+                @namespace,
+                node => Assert.IsType<ClassDeclarationIntermediateNode>(node));
+
+            var @class = @namespace.Children[0];
+            Children(
+                @class,
+                node => Assert.IsType<MethodDeclarationIntermediateNode>(node),
+                node => CSharpCode(" var value1 = true; ", node),
+                node => CSharpCode(" var value2 = true; ", node),
+                node => CSharpCode(" var value3 = true; ", node));
+
+            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));
+        }
+
         private static DocumentIntermediateNode Lower(RazorCodeDocument codeDocument, RazorProjectEngine projectEngine)
         {
             for (var i = 0; i < projectEngine.Phases.Count; i++)