AyCode.Core/AyCode.Core/Serializers/ReferenceTracker.cs

using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;

namespace AyCode.Core.Serializers;

/// <summary>
/// Shared reference tracking logic for serialization.
/// Tracks object references to enable $id/$ref handling for circular references.
/// </summary>
public sealed class SerializationReferenceTracker
{
    private readonly Dictionary<object, int> _scanOccurrences;
    private readonly Dictionary<object, string> _writtenRefs;
    private readonly HashSet<object> _multiReferenced;
    private int _nextId;

    public SerializationReferenceTracker(int initialCapacity = 32)
    {
        _scanOccurrences = new(initialCapacity, ReferenceEqualityComparer.Instance);
        _writtenRefs = new(initialCapacity, ReferenceEqualityComparer.Instance);
        _multiReferenced = new(initialCapacity, ReferenceEqualityComparer.Instance);
        _nextId = 1;
    }

    /// <summary>
    /// Tracks an object during the scanning phase.
    /// Returns true if this is the first occurrence (should continue scanning children).
    /// Returns false if object was seen before (multi-referenced).
    /// </summary>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public bool TrackForScanning(object obj)
    {
        ref var count = ref CollectionsMarshal.GetValueRefOrAddDefault(_scanOccurrences, obj, out var exists);
        if (exists)
        {
            count++;
            _multiReferenced.Add(obj);
            return false;
        }
        count = 1;
        return true;
    }

    /// <summary>
    /// Checks if an object should have an $id written and returns the id.
    /// </summary>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public bool ShouldWriteId(object obj, out string id)
    {
        if (_multiReferenced.Contains(obj) && !_writtenRefs.ContainsKey(obj))
        {
            id = _nextId++.ToString();
            return true;
        }
        id = "";
        return false;
    }

    /// <summary>
    /// Marks an object as written with its assigned id.
    /// </summary>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public void MarkAsWritten(object obj, string id) => _writtenRefs[obj] = id;

    /// <summary>
    /// Tries to get an existing reference id for an object.
    /// If found, a $ref should be written instead of the full object.
    /// </summary>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public bool TryGetExistingRef(object obj, out string refId)
    {
        return _writtenRefs.TryGetValue(obj, out refId!);
    }

    /// <summary>
    /// Clears all tracking data for reuse.
    /// </summary>
    public void Clear()
    {
        _scanOccurrences.Clear();
        _writtenRefs.Clear();
        _multiReferenced.Clear();
        _nextId = 1;
    }
}

/// <summary>
/// Shared reference tracking logic for deserialization.
/// Resolves $id/$ref references during deserialization.
/// </summary>
public sealed class DeserializationReferenceTracker
{
    private Dictionary<string, object>? _idToObject;

    /// <summary>
    /// Registers an object with its $id.
    /// </summary>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public void RegisterObject(string id, object obj)
    {
        _idToObject ??= new Dictionary<string, object>(64, StringComparer.Ordinal);
        _idToObject[id] = obj;
    }

    /// <summary>
    /// Tries to get a referenced object by its $id.
    /// </summary>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public bool TryGetReferencedObject(string id, out object? obj)
    {
        if (_idToObject != null)
            return _idToObject.TryGetValue(id, out obj);
        obj = null;
        return false;
    }

    /// <summary>
    /// Clears all tracking data for reuse.
    /// </summary>
    public void Clear()
    {
        _idToObject?.Clear();
    }
}

/// <summary>
/// Reference equality comparer for object identity comparison.
/// Used for reference tracking dictionaries.
/// </summary>
public sealed class ReferenceEqualityComparer : IEqualityComparer<object>
{
    public static readonly ReferenceEqualityComparer Instance = new();
    
    private ReferenceEqualityComparer() { }
    
    public new bool Equals(object? x, object? y) => ReferenceEquals(x, y);
    public int GetHashCode(object obj) => RuntimeHelpers.GetHashCode(obj);
}