natsnet/dotnet/src/ZB.MOM.NatsNet.Server/Internal/DataStructures/GenericSublist.cs

// Copyright 2025 The NATS Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

namespace ZB.MOM.NatsNet.Server.Internal.DataStructures;

// Sublist is a routing mechanism to handle subject distribution and
// provides a facility to match subjects from published messages to
// interested subscribers. Subscribers can have wildcard subjects to
// match multiple published subjects.

/// <summary>
/// A value type used with <see cref="SimpleSublist"/> to track interest without
/// storing any associated data. Equivalent to Go's <c>struct{}</c>.
/// </summary>
public readonly struct EmptyStruct : IEquatable<EmptyStruct>
{
    public static readonly EmptyStruct Value = default;
    public bool Equals(EmptyStruct other) => true;
    public override bool Equals(object? obj) => obj is EmptyStruct;
    public override int GetHashCode() => 0;
    public static bool operator ==(EmptyStruct left, EmptyStruct right) => true;
    public static bool operator !=(EmptyStruct left, EmptyStruct right) => false;
}

/// <summary>
/// A thread-safe trie-based NATS subject routing list that efficiently stores and
/// retrieves subscriptions. Wildcards <c>*</c> (single-token) and <c>&gt;</c>
/// (full-wildcard) are supported.
/// </summary>
/// <typeparam name="T">The subscription value type. Must be non-null.</typeparam>
public class GenericSublist<T> where T : notnull
{
    // Token separator and wildcard constants (mirrors Go's const block).
    private const char Pwc = '*';
    private const char Fwc = '>';
    private const char Btsep = '.';

    // -------------------------------------------------------------------------
    // Public error singletons (mirrors Go's var block).
    // -------------------------------------------------------------------------

    /// <summary>Thrown when a subject is syntactically invalid.</summary>
    public static readonly ArgumentException ErrInvalidSubject =
        new("gsl: invalid subject");

    /// <summary>Thrown when a subscription is not found during removal.</summary>
    public static readonly KeyNotFoundException ErrNotFound =
        new("gsl: no matches found");

    /// <summary>Thrown when a value is already registered for the given subject.</summary>
    public static readonly InvalidOperationException ErrAlreadyRegistered =
        new("gsl: notification already registered");

    // -------------------------------------------------------------------------
    // Fields
    // -------------------------------------------------------------------------

    private readonly TrieLevel _root;
    private uint _count;
    private readonly ReaderWriterLockSlim _lock = new(LockRecursionPolicy.NoRecursion);

    // -------------------------------------------------------------------------
    // Construction
    // -------------------------------------------------------------------------

    internal GenericSublist()
    {
        _root = new TrieLevel();
    }

    /// <summary>Creates a new <see cref="GenericSublist{T}"/>.</summary>
    public static GenericSublist<T> NewSublist() => new();

    /// <summary>Creates a new <see cref="SimpleSublist"/>.</summary>
    public static SimpleSublist NewSimpleSublist() => new();

    // -------------------------------------------------------------------------
    // Public API
    // -------------------------------------------------------------------------

    /// <summary>Returns the total number of subscriptions stored.</summary>
    public uint Count
    {
        get
        {
            _lock.EnterReadLock();
            try { return _count; }
            finally { _lock.ExitReadLock(); }
        }
    }

    /// <summary>
    /// Inserts a subscription into the trie.
    /// Throws <see cref="ArgumentException"/> if <paramref name="subject"/> is invalid.
    /// </summary>
    public void Insert(string subject, T value)
    {
        _lock.EnterWriteLock();
        try
        {
            InsertCore(subject, value);
        }
        finally
        {
            _lock.ExitWriteLock();
        }
    }

    /// <summary>
    /// Removes a subscription from the trie.
    /// Throws <see cref="ArgumentException"/> if the subject is invalid, or
    /// <see cref="KeyNotFoundException"/> if not found.
    /// </summary>
    public void Remove(string subject, T value)
    {
        _lock.EnterWriteLock();
        try
        {
            RemoveCore(subject, value);
        }
        finally
        {
            _lock.ExitWriteLock();
        }
    }

    /// <summary>
    /// Calls <paramref name="action"/> for every value whose subscription matches
    /// the literal <paramref name="subject"/>.
    /// </summary>
    public void Match(string subject, Action<T> action)
    {
        _lock.EnterReadLock();
        try
        {
            var tokens = TokenizeForMatch(subject);
            if (tokens == null) return;
            MatchLevel(_root, tokens, 0, action);
        }
        finally
        {
            _lock.ExitReadLock();
        }
    }

    /// <summary>
    /// Calls <paramref name="action"/> for every value whose subscription matches
    /// <paramref name="subject"/> supplied as a UTF-8 byte span.
    /// </summary>
    public void MatchBytes(ReadOnlySpan<byte> subject, Action<T> action)
    {
        Match(System.Text.Encoding.UTF8.GetString(subject), action);
    }

    /// <summary>
    /// Returns <see langword="true"/> when at least one subscription matches
    /// <paramref name="subject"/>.
    /// </summary>
    public bool HasInterest(string subject)
    {
        _lock.EnterReadLock();
        try
        {
            var tokens = TokenizeForMatch(subject);
            if (tokens == null) return false;
            int dummy = 0;
            return MatchLevelForAny(_root, tokens, 0, ref dummy);
        }
        finally
        {
            _lock.ExitReadLock();
        }
    }

    /// <summary>
    /// Returns the number of subscriptions that match <paramref name="subject"/>.
    /// </summary>
    public int NumInterest(string subject)
    {
        _lock.EnterReadLock();
        try
        {
            var tokens = TokenizeForMatch(subject);
            if (tokens == null) return 0;
            int np = 0;
            MatchLevelForAny(_root, tokens, 0, ref np);
            return np;
        }
        finally
        {
            _lock.ExitReadLock();
        }
    }

    /// <summary>
    /// Returns <see langword="true"/> if the trie contains any subscription that
    /// could match a subject whose tokens begin with the tokens of
    /// <paramref name="subject"/>. Used for trie intersection checks.
    /// </summary>
    public bool HasInterestStartingIn(string subject)
    {
        _lock.EnterReadLock();
        try
        {
            var tokens = TokenizeSubjectIntoSlice(subject);
            return HasInterestStartingInLevel(_root, tokens, 0);
        }
        finally
        {
            _lock.ExitReadLock();
        }
    }

    // -------------------------------------------------------------------------
    // Internal helpers (accessible to tests in the same assembly).
    // -------------------------------------------------------------------------

    /// <summary>Returns the maximum depth of the trie. Used in tests.</summary>
    internal int NumLevels() => VisitLevel(_root, 0);

    // -------------------------------------------------------------------------
    // Private: Insert core (lock must be held by caller)
    // -------------------------------------------------------------------------

    private void InsertCore(string subject, T value)
    {
        var sfwc = false;   // seen full-wildcard token
        TrieNode? n = null;
        var l = _root;

        // Iterate tokens split by '.' using index arithmetic to avoid allocations.
        var start = 0;
        while (start <= subject.Length)
        {
            // Find end of this token.
            var end = subject.IndexOf(Btsep, start);
            var isLast = end < 0;
            if (isLast) end = subject.Length;

            var tokenLen = end - start;

            if (tokenLen == 0 || sfwc)
                throw new ArgumentException(ErrInvalidSubject.Message);

            if (tokenLen > 1)
            {
                var t = subject.Substring(start, tokenLen);
                if (!l.Nodes.TryGetValue(t, out n))
                {
                    n = new TrieNode();
                    l.Nodes[t] = n;
                }
            }
            else
            {
                switch (subject[start])
                {
                    case Pwc:
                        if (l.PwcNode == null) l.PwcNode = new TrieNode();
                        n = l.PwcNode;
                        break;
                    case Fwc:
                        if (l.FwcNode == null) l.FwcNode = new TrieNode();
                        n = l.FwcNode;
                        sfwc = true;
                        break;
                    default:
                        var t = subject.Substring(start, 1);
                        if (!l.Nodes.TryGetValue(t, out n))
                        {
                            n = new TrieNode();
                            l.Nodes[t] = n;
                        }
                        break;
                }
            }

            n.Next ??= new TrieLevel();
            l = n.Next;

            if (isLast) break;
            start = end + 1;
        }

        if (n == null)
            throw new ArgumentException(ErrInvalidSubject.Message);

        n.Subs[value] = subject;
        _count++;
    }

    // -------------------------------------------------------------------------
    // Private: Remove core (lock must be held by caller)
    // -------------------------------------------------------------------------

    private void RemoveCore(string subject, T value)
    {
        var sfwc = false;
        var l = _root;

        // We use a fixed-size stack-style array to track visited (level, node, token)
        // triples so we can prune upward after removal. 32 is the same as Go's [32]lnt.
        var levels = new LevelNodeToken[32];
        var levelCount = 0;
        TrieNode? n = null;

        var start = 0;
        while (start <= subject.Length)
        {
            var end = subject.IndexOf(Btsep, start);
            var isLast = end < 0;
            if (isLast) end = subject.Length;

            var tokenLen = end - start;

            if (tokenLen == 0 || sfwc)
                throw new ArgumentException(ErrInvalidSubject.Message);

            if (l == null!)
                throw new KeyNotFoundException(ErrNotFound.Message);

            var tokenStr = subject.Substring(start, tokenLen);

            if (tokenLen > 1)
            {
                l.Nodes.TryGetValue(tokenStr, out n);
            }
            else
            {
                switch (tokenStr[0])
                {
                    case Pwc:
                        n = l.PwcNode;
                        break;
                    case Fwc:
                        n = l.FwcNode;
                        sfwc = true;
                        break;
                    default:
                        l.Nodes.TryGetValue(tokenStr, out n);
                        break;
                }
            }

            if (n != null)
            {
                if (levelCount < levels.Length)
                    levels[levelCount++] = new LevelNodeToken(l, n, tokenStr);
                l = n.Next!;
            }
            else
            {
                l = null!;
            }

            if (isLast) break;
            start = end + 1;
        }

        // Remove from the final node's subscription map.
        if (!RemoveFromNode(n, value))
            throw new KeyNotFoundException(ErrNotFound.Message);

        _count--;

        // Prune empty nodes upward.
        for (var i = levelCount - 1; i >= 0; i--)
        {
            var (lv, nd, tk) = levels[i];
            if (nd.IsEmpty())
                lv.PruneNode(nd, tk);
        }
    }

    private static bool RemoveFromNode(TrieNode? n, T value)
    {
        if (n == null) return false;
        return n.Subs.Remove(value);
    }

    // -------------------------------------------------------------------------
    // Private: matchLevel - recursive trie descent with callback
    // Mirrors Go's matchLevel function exactly.
    // -------------------------------------------------------------------------

    private static void MatchLevel(TrieLevel? l, string[] tokens, int start, Action<T> action)
    {
        TrieNode? pwc = null;
        TrieNode? n = null;

        for (var i = start; i < tokens.Length; i++)
        {
            if (l == null) return;

            // Full-wildcard at this level matches everything at/below.
            if (l.FwcNode != null)
                CallbacksForResults(l.FwcNode, action);

            pwc = l.PwcNode;
            if (pwc != null)
                MatchLevel(pwc.Next, tokens, i + 1, action);

            l.Nodes.TryGetValue(tokens[i], out n);
            l = n?.Next;
        }

        // After consuming all tokens, emit subs from exact and pwc matches.
        if (n != null)
            CallbacksForResults(n, action);

        if (pwc != null)
            CallbacksForResults(pwc, action);
    }

    private static void CallbacksForResults(TrieNode n, Action<T> action)
    {
        foreach (var sub in n.Subs.Keys)
            action(sub);
    }

    // -------------------------------------------------------------------------
    // Private: matchLevelForAny - returns true on first match, counting via np
    // Mirrors Go's matchLevelForAny function exactly.
    // -------------------------------------------------------------------------

    private static bool MatchLevelForAny(TrieLevel? l, string[] tokens, int start, ref int np)
    {
        TrieNode? pwc = null;
        TrieNode? n = null;

        for (var i = start; i < tokens.Length; i++)
        {
            if (l == null) return false;

            if (l.FwcNode != null)
            {
                np += l.FwcNode.Subs.Count;
                return true;
            }

            pwc = l.PwcNode;
            if (pwc != null)
            {
                if (MatchLevelForAny(pwc.Next, tokens, i + 1, ref np))
                    return true;
            }

            l.Nodes.TryGetValue(tokens[i], out n);
            l = n?.Next;
        }

        if (n != null)
        {
            np += n.Subs.Count;
            if (n.Subs.Count > 0) return true;
        }

        if (pwc != null)
        {
            np += pwc.Subs.Count;
            return pwc.Subs.Count > 0;
        }

        return false;
    }

    // -------------------------------------------------------------------------
    // Private: hasInterestStartingIn - mirrors Go's hasInterestStartingIn
    // -------------------------------------------------------------------------

    private static bool HasInterestStartingInLevel(TrieLevel? l, string[] tokens, int start)
    {
        if (l == null) return false;
        if (start >= tokens.Length) return true;

        if (l.FwcNode != null) return true;

        var found = false;
        if (l.PwcNode != null)
            found = HasInterestStartingInLevel(l.PwcNode.Next, tokens, start + 1);

        if (!found && l.Nodes.TryGetValue(tokens[start], out var n))
            found = HasInterestStartingInLevel(n.Next, tokens, start + 1);

        return found;
    }

    // -------------------------------------------------------------------------
    // Private: numLevels helper - mirrors Go's visitLevel
    // -------------------------------------------------------------------------

    private static int VisitLevel(TrieLevel? l, int depth)
    {
        if (l == null || l.NumNodes() == 0) return depth;

        depth++;
        var maxDepth = depth;

        foreach (var n in l.Nodes.Values)
        {
            var d = VisitLevel(n.Next, depth);
            if (d > maxDepth) maxDepth = d;
        }

        if (l.PwcNode != null)
        {
            var d = VisitLevel(l.PwcNode.Next, depth);
            if (d > maxDepth) maxDepth = d;
        }

        if (l.FwcNode != null)
        {
            var d = VisitLevel(l.FwcNode.Next, depth);
            if (d > maxDepth) maxDepth = d;
        }

        return maxDepth;
    }

    // -------------------------------------------------------------------------
    // Private: tokenization helpers
    // -------------------------------------------------------------------------

    /// <summary>
    /// Tokenizes a subject for match/hasInterest operations.
    /// Returns <see langword="null"/> if the subject contains an empty token,
    /// because an empty token can never match any subscription in the trie.
    /// Mirrors Go's inline tokenization in <c>match()</c> and <c>hasInterest()</c>.
    /// </summary>
    private static string[]? TokenizeForMatch(string subject)
    {
        if (subject.Length == 0) return null;

        var tokens = new List<string>(8);
        var start = 0;

        for (var i = 0; i < subject.Length; i++)
        {
            if (subject[i] == Btsep)
            {
                if (i - start == 0) return null; // empty token
                tokens.Add(subject.Substring(start, i - start));
                start = i + 1;
            }
        }

        // Trailing separator produces empty last token.
        if (start >= subject.Length) return null;
        tokens.Add(subject.Substring(start));
        return tokens.ToArray();
    }

    /// <summary>
    /// Tokenizes a subject into a string array without validation.
    /// Mirrors Go's <c>tokenizeSubjectIntoSlice</c>.
    /// </summary>
    private static string[] TokenizeSubjectIntoSlice(string subject)
    {
        var tokens = new List<string>(8);
        var start = 0;
        for (var i = 0; i < subject.Length; i++)
        {
            if (subject[i] == Btsep)
            {
                tokens.Add(subject.Substring(start, i - start));
                start = i + 1;
            }
        }
        tokens.Add(subject.Substring(start));
        return tokens.ToArray();
    }

    // -------------------------------------------------------------------------
    // Private: Trie node and level types
    // -------------------------------------------------------------------------

    /// <summary>
    /// A trie node holding a subscription map and an optional link to the next level.
    /// Mirrors Go's <c>node[T]</c>.
    /// </summary>
    private sealed class TrieNode
    {
        /// <summary>Maps subscription value → original subject string.</summary>
        public readonly Dictionary<T, string> Subs = new();

        /// <summary>The next trie level below this node, or null if at a leaf.</summary>
        public TrieLevel? Next;

        /// <summary>
        /// Returns true when the node has no subscriptions and no live children.
        /// Used during removal to decide whether to prune this node.
        /// Mirrors Go's <c>node.isEmpty()</c>.
        /// </summary>
        public bool IsEmpty() => Subs.Count == 0 && (Next == null || Next.NumNodes() == 0);
    }

    /// <summary>
    /// A trie level containing named child nodes and special wildcard slots.
    /// Mirrors Go's <c>level[T]</c>.
    /// </summary>
    private sealed class TrieLevel
    {
        public readonly Dictionary<string, TrieNode> Nodes = new();
        public TrieNode? PwcNode;   // '*' single-token wildcard node
        public TrieNode? FwcNode;   // '>' full-wildcard node

        /// <summary>
        /// Returns the total count of live nodes at this level.
        /// Mirrors Go's <c>level.numNodes()</c>.
        /// </summary>
        public int NumNodes()
        {
            var num = Nodes.Count;
            if (PwcNode != null) num++;
            if (FwcNode != null) num++;
            return num;
        }

        /// <summary>
        /// Removes an empty node from this level, using reference equality to
        /// distinguish wildcard slots from named slots.
        /// Mirrors Go's <c>level.pruneNode()</c>.
        /// </summary>
        public void PruneNode(TrieNode n, string token)
        {
            if (ReferenceEquals(n, FwcNode))
                FwcNode = null;
            else if (ReferenceEquals(n, PwcNode))
                PwcNode = null;
            else
                Nodes.Remove(token);
        }
    }

    /// <summary>
    /// Tracks a (level, node, token) triple during removal for upward pruning.
    /// Mirrors Go's <c>lnt[T]</c>.
    /// </summary>
    private readonly struct LevelNodeToken
    {
        public readonly TrieLevel Level;
        public readonly TrieNode Node;
        public readonly string Token;

        public LevelNodeToken(TrieLevel level, TrieNode node, string token)
        {
            Level = level;
            Node = node;
            Token = token;
        }

        public void Deconstruct(out TrieLevel level, out TrieNode node, out string token)
        {
            level = Level;
            node = Node;
            token = Token;
        }
    }
}

/// <summary>
/// A lightweight sublist that tracks interest only, without storing any associated data.
/// Equivalent to Go's <c>SimpleSublist = GenericSublist[struct{}]</c>.
/// </summary>
public sealed class SimpleSublist : GenericSublist<EmptyStruct>
{
    internal SimpleSublist() { }
}