Hierarchical State Restoring (#6)

* Add StateManager

* Rename method

* Rename variable

* Add state for counters

* Remove ViewModelStateKey

* Add EncodeState method

* Cleanup

* Fix summary

* Minor fixes

* Fix GetNestedState

* Add summary

* Add tests

* Add State directory

* Extract IStateManager

* Refactor EncodeState

* Add RestoreStateAsync

* Cleanup

* Add tests

* Cleanup

* Cleanup

* PR fixes

* Make InjectionMaps private

* Fix formatting

* Fix PR

sample/BitzArt.Blazor.MVVM.SampleApp/BitzArt.Blazor.MVVM.SampleApp.Client/Pages/CounterPage.razor

@@ -10,7 +10,7 @@
 
 <p>@ViewModel.State?.Text</p>
 
-<Counter ViewModel="ViewModel.Counter1ViewModel" />
-<Counter ViewModel="ViewModel.Counter2ViewModel" />
+<Counter ViewModel="ViewModel.Counter1" />
+<Counter ViewModel="ViewModel.Counter2" />
 
 <ComponentStateContainer ViewModel="ViewModel" />

sample/BitzArt.Blazor.MVVM.SampleApp/BitzArt.Blazor.MVVM.SampleApp.Client/ViewModels/CounterPageViewModel.cs

@@ -7,29 +7,19 @@ public class CounterPageViewModel(
     : ViewModel<CounterPageViewModelState>
 {
     [Inject]
-    public CounterViewModel Counter1ViewModel { get; set; } = null!;
+    public CounterViewModel Counter1 { get; set; } = null!;
 
     [Inject]
-    public CounterViewModel Counter2ViewModel { get; set; } = null!;
+    public CounterViewModel Counter2 { get; set; } = null!;
 
     public override void InitializeState()
     {
         State.Text = $"ViewModel State initialized on: {renderingEnvironment}";
-
-        OnStateRestored();
-        Counter2ViewModel.State!.Count += 100;
-    }
-
-    public override void OnStateRestored()
-    {
-        Counter1ViewModel.State = State.Counter1State;
-        Counter2ViewModel.State = State.Counter2State;
+        Counter2.State!.Count = 100;
     }
 }
 
 public class CounterPageViewModelState
 {
     public string Text { get; set; } = "State not initialized";
-    public CounterState Counter1State { get; set; } = new();
-    public CounterState Counter2State { get; set; } = new();
 }

sample/BitzArt.Blazor.MVVM.SampleApp/BitzArt.Blazor.MVVM.SampleApp.Client/ViewModels/CounterViewModel.cs

@@ -1,16 +1,19 @@
 namespace BitzArt.Blazor.MVVM.SampleApp;
 
-public class CounterViewModel : ViewModel
+public class CounterViewModel : ViewModel<CounterState>
 {
-    public CounterState? State { get; set; }
-
     private readonly Timer _timer;
 
     public CounterViewModel()
     {
         _timer = new Timer(TimerIncrementCount, null, 1000, 1000);
     }
 
+    public override void InitializeState()
+    {
+        State.Count = 0;
+    }
+
     private void TimerIncrementCount(object? state)
     {
         if (State is null) return;
@@ -30,5 +33,5 @@ public void IncrementCount()
 
 public class CounterState
 {
-    public int Count { get; set; } = 0;
+    public int? Count { get; set; };
 }

src/BitzArt.Blazor.MVVM/BitzArt.Blazor.MVVM.csproj

@@ -15,14 +15,9 @@
 		<PackageProjectUrl>https://bitzart.github.io/Blazor.MVVM/</PackageProjectUrl>
 		<PackageReadmeFile>README.md</PackageReadmeFile>
 	</PropertyGroup>
-
-	<ItemGroup>
-	  <Compile Remove="Models\ComponentStateContainer.razor.cs" />
-	</ItemGroup>
-
+
 	<ItemGroup>
 		<None Include="..\..\README.md" Pack="True" Visible="False" PackagePath="\" />
-		<None Include="ComponentStateContainer.razor.cs" />
 	</ItemGroup>
 
 	<ItemGroup>

src/BitzArt.Blazor.MVVM/Builder/BlazorMvvmBuilder.cs

@@ -15,7 +15,9 @@ internal class BlazorMvvmBuilder : IBlazorMvvmBuilder
 
     public IViewModelFactory Factory { get; set; }
 
-    public BlazorMvvmBuilder(IServiceCollection serviceCollection, IViewModelFactory factory)
+    public BlazorMvvmBuilder(
+        IServiceCollection serviceCollection,
+        IViewModelFactory factory)
     {
         ServiceCollection = serviceCollection;
         Factory = factory;
@@ -24,5 +26,6 @@ public BlazorMvvmBuilder(IServiceCollection serviceCollection, IViewModelFactory
             throw new InvalidOperationException("IViewModelFactory is already registered in this ServiceCollection.");
 
         ServiceCollection.AddSingleton(Factory);
+        ServiceCollection.AddSingleton<BlazorViewModelStateManager>();
     }
 }

src/BitzArt.Blazor.MVVM/Factory/ViewModelFactory.cs

@@ -12,7 +12,7 @@ public interface IViewModelFactory
 
 internal class ViewModelFactory : IViewModelFactory
 {
-    public ICollection<ViewModelInjectionMap> InjectionMaps { get; set; }
+    private Dictionary<Type, ViewModelInjectionMap> InjectionMaps { get; set; }
 
     public ViewModelFactory()
     {
@@ -21,14 +21,13 @@ public ViewModelFactory()
 
     public void AddViewModel(Type viewModelType, string registrationKey)
     {
-        if (!typeof(ViewModel).IsAssignableFrom(viewModelType)) throw new InvalidOperationException(
-            $"Type {viewModelType.Name} is not a ViewModel.");
+        if (!typeof(ViewModel).IsAssignableFrom(viewModelType)) 
+            throw new InvalidOperationException($"Type {viewModelType.Name} is not a ViewModel.");
 
-        if (InjectionMaps.Any(x => x.ViewModelType == viewModelType))
-            throw new InvalidOperationException(
-                $"ViewModel {viewModelType.Name} is already registered in the factory.");
+        if (InjectionMaps.ContainsKey(viewModelType))
+            throw new InvalidOperationException($"ViewModel {viewModelType.Name} is already registered in the factory.");
 
-        InjectionMaps.Add(new(viewModelType, registrationKey));
+        InjectionMaps.Add(viewModelType, new(viewModelType, registrationKey));
     }
 
     public TViewModel Create<TViewModel>(IServiceProvider serviceProvider)
@@ -37,9 +36,8 @@ public TViewModel Create<TViewModel>(IServiceProvider serviceProvider)
 
     public ViewModel Create(IServiceProvider serviceProvider, Type viewModelType)
     {
-        var viewModelMap = InjectionMaps.FirstOrDefault(x => x.ViewModelType == viewModelType)
-            ?? throw new InvalidOperationException(
-                $"ViewModel {viewModelType.Name} is not registered in the factory.");
+        var viewModelMap = InjectionMaps.GetValueOrDefault(viewModelType)
+            ?? throw new InvalidOperationException($"ViewModel {viewModelType.Name} is not registered in the factory.");
 
         var viewModel = (ViewModel)serviceProvider.GetRequiredKeyedService(typeof(ViewModel), viewModelMap.RegistrationKey);
         foreach (var injection in viewModelMap.Injections)
@@ -50,4 +48,12 @@ public ViewModel Create(IServiceProvider serviceProvider, Type viewModelType)
         viewModel.OnDependenciesInjected();
         return viewModel;
     }
+
+    public ViewModelInjectionMap GetInjectionMap(Type viewModelType)
+    {
+        var injectionMap = InjectionMaps.GetValueOrDefault(viewModelType)
+            ?? throw new InvalidOperationException($"ViewModel {viewModelType.Name} is not registered in the factory.");
+
+        return injectionMap;
+    }
 }

src/BitzArt.Blazor.MVVM/Models/PageBase.cs

@@ -19,6 +19,8 @@ public abstract class PageBase<TViewModel> : ComponentBase<TViewModel>, IStateCo
         set => base.ViewModel = value;
     }
 
+    [Inject] private BlazorViewModelStateManager StateManager { get; set; } = null!;
+
     private const string StateKey = "state";
 
     protected override async Task RestoreStateAsync()
@@ -28,8 +30,6 @@ protected override async Task RestoreStateAsync()
 
     private async Task RestoreComponentStateAsync(ViewModel viewModel)
     {
-        if (viewModel is not IStatefulViewModel statefulViewModel) return;
-
         var isPrerender = RenderingEnvironment.IsPrerender;
         var state = isPrerender
             ? null
@@ -39,24 +39,13 @@ private async Task RestoreComponentStateAsync(ViewModel viewModel)
         {
             var buffer = Convert.FromBase64String(state);
             var json = Encoding.UTF8.GetString(buffer);
-            statefulViewModel.State = JsonSerializer.Deserialize(json, statefulViewModel.StateType, StateJsonOptionsProvider.Options)!;
-
-            statefulViewModel.OnStateRestored();
-            await statefulViewModel.OnStateRestoredAsync();
+
+            await StateManager.RestoreStateAsync(viewModel, json);
         }
         else
         {
-            statefulViewModel.State = Activator.CreateInstance(statefulViewModel.StateType)!;
-
-            statefulViewModel.OnStateChanged();
-            await statefulViewModel.OnStateChangedAsync();
-
-            statefulViewModel.InitializeState();
-            await statefulViewModel.InitializeStateAsync();
+            await StateManager.InitializeStateAsync(viewModel);
         }
-
-        statefulViewModel.OnStateChanged();
-        await statefulViewModel.OnStateChangedAsync();
     }
 
     /// <summary>
@@ -75,6 +64,7 @@ internal static class StateJsonOptionsProvider
     public static readonly JsonSerializerOptions Options = new()
     {
         PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
+        DictionaryKeyPolicy = JsonNamingPolicy.CamelCase,
         PropertyNameCaseInsensitive = true,
         DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull
     };

src/BitzArt.Blazor.MVVM/State/BlazorViewModelStateManager.cs

@@ -0,0 +1,123 @@
+using System.Text.Json;
+using System.Text.Json.Nodes;
+
+namespace BitzArt.Blazor.MVVM;
+
+internal class BlazorViewModelStateManager(IViewModelFactory viewModelFactory)
+{
+    private ViewModelFactory _viewModelFactory { get; } = (ViewModelFactory)viewModelFactory;
+
+    private const string _nestedStatePrefix = "n__";
+
+    /// <summary>
+    /// Serializes states in <see cref="ViewModel"/>s hierarchy to JSON encoded as UTF-8 bytes.
+    /// </summary>
+    public byte[]? SerializeState(ViewModel viewModel)
+    {
+        var injectionMap = _viewModelFactory.GetInjectionMap(viewModel.GetType());
+        var state = GetState(viewModel, injectionMap);
+        if (state is null) return null;
+
+        return JsonSerializer.SerializeToUtf8Bytes(state, StateJsonOptionsProvider.Options);
+    }
+
+    private Dictionary<string, object?>? GetState(ViewModel viewModel, ViewModelInjectionMap injectionMap)
+    {
+        var state = new Dictionary<string, object?>();
+
+        if (viewModel is IStatefulViewModel statefulViewModel)
+        {
+            foreach (var property in statefulViewModel.StateType.GetProperties())
+                state.Add(property.Name, property.GetValue(statefulViewModel.State));
+        }
+
+        foreach (var injection in injectionMap.Injections)
+        {
+            var property = injection.Property;
+            var nestedViewModel = property.GetValue(viewModel) as ViewModel;
+            var nestedMap = _viewModelFactory.GetInjectionMap(injection.ViewModelType);
+            var nestedState = GetState(nestedViewModel!, nestedMap);
+
+            if (nestedState is not null)
+                state.Add($"{_nestedStatePrefix}{property.Name}", nestedState);
+        }
+
+        return state.Values.Count != 0 ? state : null;
+    }
+
+    /// <summary>
+    /// Restores states in <see cref="ViewModel"/>s hierarchy from JSON string.
+    /// </summary>
+    public async Task RestoreStateAsync(ViewModel viewModel, string json)
+    {
+        var node = JsonNode.Parse(json);
+        if (node is null) return;
+
+        var injectionMap = _viewModelFactory.GetInjectionMap(viewModel.GetType());
+        await RestoreStateAsync(viewModel, node, injectionMap);
+    }
+
+    private async Task RestoreStateAsync(ViewModel viewModel, JsonNode node, ViewModelInjectionMap injectionMap)
+    {
+        foreach (var injection in injectionMap.Injections)
+        {
+            var property = injection.Property;
+            var jsonKey = $"{_nestedStatePrefix}{property.Name}";
+            var nestedNode = node[jsonKey];
+
+            if (nestedNode is null) continue;
+
+            var nestedViewModel = property.GetValue(viewModel) as ViewModel;
+            var nestedMap = _viewModelFactory.GetInjectionMap(injection.ViewModelType);
+
+            await RestoreStateAsync(nestedViewModel!, nestedNode, nestedMap);
+
+            (node as JsonObject)!.Remove(jsonKey);
+        }
+
+        if (viewModel is not IStatefulViewModel statefulViewModel) return;
+
+        var state = JsonSerializer.Deserialize(node, statefulViewModel.StateType, StateJsonOptionsProvider.Options);
+        statefulViewModel.State = state!;
+
+        statefulViewModel.OnStateRestored();
+        await statefulViewModel.OnStateRestoredAsync();
+
+        statefulViewModel.OnStateChanged();
+        await statefulViewModel.OnStateChangedAsync();
+    }
+
+    /// <summary>
+    /// Initializes states in <see cref="ViewModel"/>s hierarchy.
+    /// </summary>
+    public async Task InitializeStateAsync(ViewModel viewModel)
+    {
+        var injectionMap = _viewModelFactory.GetInjectionMap(viewModel.GetType());
+        await InitializeStateAsync(viewModel, injectionMap);
+    }
+
+    private async Task InitializeStateAsync(ViewModel viewModel, ViewModelInjectionMap injectionMap)
+    {
+        foreach (var injection in injectionMap.Injections)
+        {
+            var property = injection.Property;
+            var nestedViewModel = property.GetValue(viewModel) as ViewModel;
+            var nestedMap = _viewModelFactory.GetInjectionMap(injection.ViewModelType);
+
+            await InitializeStateAsync(nestedViewModel!, nestedMap);
+        }
+
+        if (viewModel is not IStatefulViewModel statefulViewModel) return;
+
+        statefulViewModel.State = Activator.CreateInstance(statefulViewModel.StateType)!;
+
+        statefulViewModel.OnStateChanged();
+        await statefulViewModel.OnStateChangedAsync();
+
+        statefulViewModel.InitializeState();
+        await statefulViewModel.InitializeStateAsync();
+
+        statefulViewModel.OnStateChanged();
+        await statefulViewModel.OnStateChangedAsync();
+    }
+}