From 55a4ce8646286646a45939588787ac016bd1aee7 Mon Sep 17 00:00:00 2001
From: Lucas Pedroza <lucas.pedroza@fauna.com>
Date: Tue, 12 Nov 2024 15:14:41 +0100
Subject: [PATCH] Handle unmaterialized sets in PageSerializer

---
 .../Serializers/PageSerializer.Tests.cs       | 12 ++++
 Fauna/Serialization/PageSerializer.cs         | 67 ++++++++++++-------
 2 files changed, 53 insertions(+), 26 deletions(-)

diff --git a/Fauna.Test/Serialization/Serializers/PageSerializer.Tests.cs b/Fauna.Test/Serialization/Serializers/PageSerializer.Tests.cs
index 6ab994bf..c3b63f16 100644
--- a/Fauna.Test/Serialization/Serializers/PageSerializer.Tests.cs
+++ b/Fauna.Test/Serialization/Serializers/PageSerializer.Tests.cs
@@ -64,4 +64,16 @@ public void DeserializeIntoPageWithUserDefinedClass()
         Assert.AreEqual(40, deserialized.Data[1].Age);
         Assert.AreEqual("next_page_cursor", deserialized.After);
     }
+
+    [Test]
+    public void DeserializeUnmaterializedSet()
+    {
+        const string token = "aftertoken";
+        const string wire = $$"""{"@set":"{{token}}"}""";
+
+        var serializer = Serializer.Generate<Page<object>>(s_ctx);
+        var deserialized = Helpers.Deserialize(serializer, s_ctx, wire)!;
+        Assert.IsEmpty(deserialized.Data);
+        Assert.AreEqual(token, deserialized.After);
+    }
 }
diff --git a/Fauna/Serialization/PageSerializer.cs b/Fauna/Serialization/PageSerializer.cs
index ffa9b7a9..3e4a3a71 100644
--- a/Fauna/Serialization/PageSerializer.cs
+++ b/Fauna/Serialization/PageSerializer.cs
@@ -1,4 +1,3 @@
-using Fauna.Exceptions;
 using Fauna.Mapping;
 using Fauna.Types;
 
@@ -13,11 +12,11 @@ public PageSerializer(ISerializer<T> elemSerializer)
         _dataSerializer = new ListSerializer<T>(elemSerializer);
     }
 
-    public override List<FaunaType> GetSupportedTypes() => new List<FaunaType> { FaunaType.Null, FaunaType.Set };
+    public override List<FaunaType> GetSupportedTypes() => [FaunaType.Null, FaunaType.Set];
 
     public override Page<T> Deserialize(MappingContext ctx, ref Utf8FaunaReader reader)
     {
-        var wrapInPage = false;
+        bool wrapInPage = false;
         var endToken = TokenType.None;
         switch (reader.CurrentTokenType)
         {
@@ -32,36 +31,52 @@ public override Page<T> Deserialize(MappingContext ctx, ref Utf8FaunaReader read
                 break;
         }
 
-        List<T>? data = null;
-        string? after = null;
-
         if (wrapInPage)
         {
-            data = _dataSerializer.Deserialize(ctx, ref reader);
+            var data = _dataSerializer.Deserialize(ctx, ref reader);
+            return new Page<T>(data, null);
         }
-        else
-        {
-            while (reader.Read() && reader.CurrentTokenType != endToken)
-            {
-                var fieldName = reader.GetString()!;
-                reader.Read();
 
-                switch (fieldName)
-                {
-                    case "data":
-                        data = _dataSerializer.Deserialize(ctx, ref reader);
-                        break;
-                    case "after":
-                        after = reader.GetString()!;
-                        break;
-                }
-            }
+        reader.Read();
+        return reader.CurrentTokenType == TokenType.String
+            ? HandleUnmaterialized(ctx, ref reader, endToken)
+            : HandleMaterialized(ctx, ref reader, endToken);
+    }
+
+    private Page<T> HandleUnmaterialized(MappingContext ctx, ref Utf8FaunaReader reader, TokenType endToken)
+    {
+        string after = reader.GetString()!;
+        reader.Read();
+        if (reader.CurrentTokenType != endToken)
+        {
+            throw UnexpectedToken(reader.CurrentTokenType);
         }
 
-        if (data is null)
-            throw new SerializationException($"No page data found while deserializing into {typeof(Page<T>)}");
+        return new Page<T>([], after);
+    }
+
+    private Page<T> HandleMaterialized(MappingContext ctx, ref Utf8FaunaReader reader, TokenType endToken)
+    {
+        List<T> data = [];
+        string? after = null;
+
+        do
+        {
+            string fieldName = reader.GetString()!;
+            reader.Read();
+
+            switch (fieldName)
+            {
+                case "data":
+                    data = _dataSerializer.Deserialize(ctx, ref reader);
+                    break;
+                case "after":
+                    after = reader.GetString()!;
+                    break;
+            }
+        } while (reader.Read() && reader.CurrentTokenType != endToken);
 
-        return new Page<T>(data!, after);
+        return new Page<T>(data, after);
     }
 
     public override void Serialize(MappingContext context, Utf8FaunaWriter writer, object? o)