Save the full state

ltrzesniewski · Dec 2, 2023 · b48f38b · b48f38b
1 parent 3e8cbdc
commit b48f38b
Show file tree

Hide file tree

Showing 4 changed files with 62 additions and 70 deletions.
diff --git a/src/RazorBlade.Library/HtmlLayout.cs b/src/RazorBlade.Library/HtmlLayout.cs
@@ -1,5 +1,4 @@
 using System;
-using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -16,27 +15,14 @@ public abstract class HtmlLayout : HtmlTemplate, IRazorLayout
 
     async Task<IRazorLayout.IExecutionResult> IRazorLayout.ExecuteLayoutAsync(IRazorLayout.IExecutionResult input)
     {
-        input.CancellationToken.ThrowIfCancellationRequested();
-        var previousStatus = (Output, CancellationToken);
-
         try
         {
             _layoutInput = input;
-
-            var output = new StringWriter();
-
-            Output = output;
-            CancellationToken = input.CancellationToken;
-            // TODO fully reset/restore the state
-
-            await ExecuteAsync().ConfigureAwait(false);
-
-            return new ExecutionResult(this, output.GetStringBuilder());
+            return await ExecuteAsyncCore(input.CancellationToken);
         }
         finally
         {
             _layoutInput = null;
-            (Output, CancellationToken) = previousStatus;
         }
     }
 

diff --git a/src/RazorBlade.Library/RazorTemplate.cs b/src/RazorBlade.Library/RazorTemplate.cs
@@ -27,7 +27,7 @@ public abstract class RazorTemplate : IEncodedContent
     /// <summary>
     /// The cancellation token.
     /// </summary>
-    protected internal CancellationToken CancellationToken { get; set; }
+    protected internal CancellationToken CancellationToken { get; private set; }
 
     /// <summary>
     /// The layout to use.
@@ -116,41 +116,29 @@ private async Task<StringBuilder> RenderAsyncCore(CancellationToken cancellation
     {
         cancellationToken.ThrowIfCancellationRequested();
 
-        var previousState = (_sections, Output, CancellationToken, Layout);
+        var executionResult = await ExecuteAsyncCore(cancellationToken);
 
-        try
+        while (executionResult.Layout is { } layout)
         {
-            var output = new StringWriter();
-
-            _sections = null;
-            Output = output;
-            CancellationToken = cancellationToken;
-            Layout = null;
-
-            await ExecuteAsync().ConfigureAwait(false);
-
-            if (Layout is null)
-                return output.GetStringBuilder();
+            cancellationToken.ThrowIfCancellationRequested();
+            executionResult = await layout.ExecuteLayoutAsync(executionResult).ConfigureAwait(false);
+        }
 
-            IRazorLayout.IExecutionResult executionResult = new ExecutionResult(this, output.GetStringBuilder());
+        if (executionResult.Body is StringBuilderEncodedContent { StringBuilder: var outputStringBuilder })
+            return outputStringBuilder;
 
-            while (executionResult.Layout is { } layout)
-            {
-                cancellationToken.ThrowIfCancellationRequested();
-                executionResult = await layout.ExecuteLayoutAsync(executionResult).ConfigureAwait(false);
-            }
+        var outputStringWriter = new StringWriter();
+        executionResult.Body.WriteTo(outputStringWriter);
+        return outputStringWriter.GetStringBuilder();
+    }
 
-            if (executionResult.Body is StringBuilderEncodedContent { StringBuilder: var outputWithLayout })
-                return outputWithLayout;
+    private protected async Task<IRazorLayout.IExecutionResult> ExecuteAsyncCore(CancellationToken cancellationToken)
+    {
+        cancellationToken.ThrowIfCancellationRequested();
 
-            var outerBodyResult = new StringWriter();
-            executionResult.Body.WriteTo(outerBodyResult);
-            return outerBodyResult.GetStringBuilder();
-        }
-        finally
-        {
-            (_sections, Output, CancellationToken, Layout) = previousState;
-        }
+        using var executionScope = new ExecutionScope(this, cancellationToken);
+        await ExecuteAsync().ConfigureAwait(false);
+        return new ExecutionResult(this, executionScope.Output);
     }
 
     /// <summary>
@@ -238,9 +226,10 @@ protected internal void DefineSection(string name, Func<Task> action)
     void IEncodedContent.WriteTo(TextWriter textWriter)
         => Render(textWriter, CancellationToken.None);
 
-    private protected class ExecutionResult : IRazorLayout.IExecutionResult
+    private class ExecutionResult : IRazorLayout.IExecutionResult
     {
         private readonly RazorTemplate _page;
+        private readonly IReadOnlyDictionary<string, Func<Task>>? _sections;
 
         public IEncodedContent Body { get; }
         public IRazorLayout? Layout { get; }
@@ -249,31 +238,20 @@ private protected class ExecutionResult : IRazorLayout.IExecutionResult
         public ExecutionResult(RazorTemplate page, StringBuilder body)
         {
             _page = page;
+            _sections = page._sections;
             Body = new StringBuilderEncodedContent(body);
             Layout = page.Layout;
             CancellationToken = page.CancellationToken;
         }
 
         public async Task<IEncodedContent?> RenderSectionAsync(string name)
         {
-            if (!_page.Sections.TryGetValue(name, out var sectionAction))
+            if (_sections is null || !_sections.TryGetValue(name, out var sectionAction))
                 return null;
 
-            var previousOutput = _page.Output;
-
-            try
-            {
-                var output = new StringWriter();
-                _page.Output = output;
-
-                await sectionAction().ConfigureAwait(false);
-
-                return new StringBuilderEncodedContent(output.GetStringBuilder());
-            }
-            finally
-            {
-                _page.Output = previousOutput;
-            }
+            using var executionScope = new ExecutionScope(_page, CancellationToken);
+            await sectionAction().ConfigureAwait(false);
+            return new StringBuilderEncodedContent(executionScope.Output);
         }
     }
 
@@ -292,4 +270,38 @@ public void WriteTo(TextWriter textWriter)
         public override string ToString()
             => StringBuilder.ToString();
     }
+
+    private class ExecutionScope : IDisposable
+    {
+        private readonly RazorTemplate _page;
+
+        private readonly Dictionary<string, Func<Task>>? _sections;
+        private readonly TextWriter _output;
+        private readonly CancellationToken _cancellationToken;
+        private readonly IRazorLayout? _layout;
+
+        public StringBuilder Output { get; }
+
+        public ExecutionScope(RazorTemplate page, CancellationToken cancellationToken)
+        {
+            _page = page;
+
+            _sections = page._sections;
+            _output = page.Output;
+            _cancellationToken = page.CancellationToken;
+            _layout = page.Layout;
+
+            Output = new StringBuilder();
+            page.Output = new StringWriter(Output);
+            page.CancellationToken = cancellationToken;
+        }
+
+        public void Dispose()
+        {
+            _page._sections = _sections;
+            _page.Output = _output;
+            _page.CancellationToken = _cancellationToken;
+            _page.Layout = _layout;
+        }
+    }
 }
diff --git a/src/RazorBlade.Tests/HtmlTemplateTests.cs b/src/RazorBlade.Tests/HtmlTemplateTests.cs
@@ -108,10 +108,7 @@ private class Template : HtmlTemplate
         private readonly Action<Template> _executeAction;
 
         public Template(Action<Template> executeAction)
-        {
-            _executeAction = executeAction;
-            Output = new StringWriter();
-        }
+            => _executeAction = executeAction;
 
         protected internal override Task ExecuteAsync()
         {

diff --git a/src/RazorBlade.Tests/PlainTextTemplateTests.cs b/src/RazorBlade.Tests/PlainTextTemplateTests.cs
@@ -83,10 +83,7 @@ private class Template : PlainTextTemplate
         private readonly Action<Template> _executeAction;
 
         public Template(Action<Template> executeAction)
-        {
-            _executeAction = executeAction;
-            Output = new StringWriter();
-        }
+            => _executeAction = executeAction;
 
         protected internal override Task ExecuteAsync()
         {