natsdotnet/tests/NATS.Server.Tests/Internal/Avl/SequenceSetTests.cs

// Copyright 2024 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.

using System.Diagnostics;
using NATS.Server.Internal.Avl;

namespace NATS.Server.Tests.Internal.Avl;

/// <summary>
/// Tests for the AVL-backed SequenceSet, ported from Go server/avl/seqset_test.go
/// and server/avl/norace_test.go.
/// </summary>
public class SequenceSetTests
{
    private const int NumEntries = SequenceSet.NumEntries; // 2048
    private const int BitsPerBucket = SequenceSet.BitsPerBucket;
    private const int NumBuckets = SequenceSet.NumBuckets;

    // Go: TestSeqSetBasics server/avl/seqset_test.go:22
    [Fact]
    public void Basics_InsertExistsDelete()
    {
        var ss = new SequenceSet();

        ulong[] seqs = [22, 222, 2000, 2, 2, 4];
        foreach (var seq in seqs)
        {
            ss.Insert(seq);
            ss.Exists(seq).ShouldBeTrue();
        }

        ss.Nodes.ShouldBe(1);
        ss.Size.ShouldBe(seqs.Length - 1); // One dup (2 appears twice)
        var (lh, rh) = ss.Heights();
        lh.ShouldBe(0);
        rh.ShouldBe(0);
    }

    // Go: TestSeqSetLeftLean server/avl/seqset_test.go:38
    [Fact]
    public void LeftLean_TreeBalancesCorrectly()
    {
        var ss = new SequenceSet();

        // Insert from high to low to create a left-leaning tree.
        for (var i = (ulong)(4 * NumEntries); i > 0; i--)
        {
            ss.Insert(i);
        }

        ss.Nodes.ShouldBe(5);
        ss.Size.ShouldBe(4 * NumEntries);
        var (lh, rh) = ss.Heights();
        lh.ShouldBe(2);
        rh.ShouldBe(1);
    }

    // Go: TestSeqSetRightLean server/avl/seqset_test.go:52
    [Fact]
    public void RightLean_TreeBalancesCorrectly()
    {
        var ss = new SequenceSet();

        // Insert from low to high to create a right-leaning tree.
        for (var i = 0UL; i < (ulong)(4 * NumEntries); i++)
        {
            ss.Insert(i);
        }

        ss.Nodes.ShouldBe(4);
        ss.Size.ShouldBe(4 * NumEntries);
        var (lh, rh) = ss.Heights();
        lh.ShouldBe(1);
        rh.ShouldBe(2);
    }

    // Go: TestSeqSetCorrectness server/avl/seqset_test.go:66
    [Fact]
    public void Correctness_RandomInsertDelete()
    {
        // Generate 100k sequences across 500k range.
        const int num = 100_000;
        const int max = 500_000;

        var rng = new Random(42);
        var set = new HashSet<ulong>();
        var ss = new SequenceSet();

        for (var i = 0; i < num; i++)
        {
            var n = (ulong)rng.NextInt64(max + 1);
            ss.Insert(n);
            set.Add(n);
        }

        for (var i = 0UL; i <= max; i++)
        {
            ss.Exists(i).ShouldBe(set.Contains(i));
        }
    }

    // Go: TestSeqSetRange server/avl/seqset_test.go:85
    [Fact]
    public void Range_IteratesInOrder()
    {
        var num = 2 * NumEntries + 22;
        var nums = new List<ulong>(num);
        for (var i = 0; i < num; i++)
        {
            nums.Add((ulong)i);
        }

        // Shuffle and insert.
        var rng = new Random(42);
        Shuffle(nums, rng);

        var ss = new SequenceSet();
        foreach (var n in nums)
        {
            ss.Insert(n);
        }

        // Range should produce ascending order.
        var result = new List<ulong>();
        ss.Range(n =>
        {
            result.Add(n);
            return true;
        });

        result.Count.ShouldBe(num);
        for (var i = 0UL; i < (ulong)num; i++)
        {
            result[(int)i].ShouldBe(i);
        }

        // Test truncating the range call.
        result.Clear();
        ss.Range(n =>
        {
            if (n >= 10)
            {
                return false;
            }

            result.Add(n);
            return true;
        });

        result.Count.ShouldBe(10);
        for (var i = 0UL; i < 10; i++)
        {
            result[(int)i].ShouldBe(i);
        }
    }

    // Go: TestSeqSetDelete server/avl/seqset_test.go:123
    [Fact]
    public void Delete_VariousPatterns()
    {
        var ss = new SequenceSet();

        ulong[] seqs = [22, 222, 2222, 2, 2, 4];
        foreach (var seq in seqs)
        {
            ss.Insert(seq);
        }

        foreach (var seq in seqs)
        {
            ss.Delete(seq);
            ss.Exists(seq).ShouldBeFalse();
        }

        ss.Root.ShouldBeNull();
    }

    // Go: TestSeqSetInsertAndDeletePedantic server/avl/seqset_test.go:139
    [Fact]
    public void InsertAndDelete_PedanticVerification()
    {
        var ss = new SequenceSet();

        var num = 50 * NumEntries + 22;
        var nums = new List<ulong>(num);
        for (var i = 0; i < num; i++)
        {
            nums.Add((ulong)i);
        }

        var rng = new Random(42);
        Shuffle(nums, rng);

        // Insert all, verify balanced after each insert.
        foreach (var n in nums)
        {
            ss.Insert(n);
            VerifyBalanced(ss);
        }

        ss.Root.ShouldNotBeNull();

        // Delete all, verify balanced after each delete.
        foreach (var n in nums)
        {
            ss.Delete(n);
            VerifyBalanced(ss);
            ss.Exists(n).ShouldBeFalse();
            if (ss.Size > 0)
            {
                ss.Root.ShouldNotBeNull();
            }
        }

        ss.Root.ShouldBeNull();
    }

    // Go: TestSeqSetMinMax server/avl/seqset_test.go:181
    [Fact]
    public void MinMax_TracksCorrectly()
    {
        var ss = new SequenceSet();

        // Simple single node.
        ulong[] seqs = [22, 222, 2222, 2, 2, 4];
        foreach (var seq in seqs)
        {
            ss.Insert(seq);
        }

        var (min, max) = ss.MinMax();
        min.ShouldBe(2UL);
        max.ShouldBe(2222UL);

        // Multi-node
        ss.Empty();

        var num = 22 * NumEntries + 22;
        var nums = new List<ulong>(num);
        for (var i = 0; i < num; i++)
        {
            nums.Add((ulong)i);
        }

        var rng = new Random(42);
        Shuffle(nums, rng);
        foreach (var n in nums)
        {
            ss.Insert(n);
        }

        (min, max) = ss.MinMax();
        min.ShouldBe(0UL);
        max.ShouldBe((ulong)(num - 1));
    }

    // Go: TestSeqSetClone server/avl/seqset_test.go:210
    [Fact]
    public void Clone_IndependentCopy()
    {
        // Generate 100k sequences across 500k range.
        const int num = 100_000;
        const int max = 500_000;

        var rng = new Random(42);
        var ss = new SequenceSet();
        for (var i = 0; i < num; i++)
        {
            ss.Insert((ulong)rng.NextInt64(max + 1));
        }

        var ssc = ss.Clone();
        ssc.Size.ShouldBe(ss.Size);
        ssc.Nodes.ShouldBe(ss.Nodes);
    }

    // Go: TestSeqSetUnion server/avl/seqset_test.go:225
    [Fact]
    public void Union_MergesSets()
    {
        var ss1 = new SequenceSet();
        var ss2 = new SequenceSet();

        ulong[] seqs1 = [22, 222, 2222, 2, 2, 4];
        foreach (var seq in seqs1)
        {
            ss1.Insert(seq);
        }

        ulong[] seqs2 = [33, 333, 3333, 3, 33_333, 333_333];
        foreach (var seq in seqs2)
        {
            ss2.Insert(seq);
        }

        var ss = SequenceSet.CreateUnion(ss1, ss2);
        ss.Size.ShouldBe(11);

        ulong[] allSeqs = [.. seqs1, .. seqs2];
        foreach (var n in allSeqs)
        {
            ss.Exists(n).ShouldBeTrue();
        }
    }

    // Go: TestSeqSetFirst server/avl/seqset_test.go:247
    [Fact]
    public void First_ReturnsMinimum()
    {
        var ss = new SequenceSet();

        ulong[] seqs = [22, 222, 2222, 222_222];
        foreach (var seq in seqs)
        {
            // Normal case where we pick first/base.
            ss.Insert(seq);
            ss.Root!.Base.ShouldBe((seq / (ulong)NumEntries) * (ulong)NumEntries);
            ss.Empty();

            // Where we set the minimum start value.
            ss.SetInitialMin(seq);
            ss.Insert(seq);
            ss.Root!.Base.ShouldBe(seq);
            ss.Empty();
        }
    }

    // Go: TestSeqSetDistinctUnion server/avl/seqset_test.go:265
    [Fact]
    public void DistinctUnion_NoOverlap()
    {
        var ss1 = new SequenceSet();
        ulong[] seqs1 = [1, 10, 100, 200];
        foreach (var seq in seqs1)
        {
            ss1.Insert(seq);
        }

        var ss2 = new SequenceSet();
        ulong[] seqs2 = [5000, 6100, 6200, 6222];
        foreach (var seq in seqs2)
        {
            ss2.Insert(seq);
        }

        var ss = ss1.Clone();
        ulong[] allSeqs = [.. seqs1, .. seqs2];

        ss.Union(ss2);
        ss.Size.ShouldBe(allSeqs.Length);
        foreach (var seq in allSeqs)
        {
            ss.Exists(seq).ShouldBeTrue();
        }
    }

    // Go: TestSeqSetDecodeV1 server/avl/seqset_test.go:289
    [Fact]
    public void DecodeV1_BackwardsCompatible()
    {
        // Encoding from v1 which was 64 buckets.
        ulong[] seqs = [22, 222, 2222, 222_222, 2_222_222];
        var encStr =
            "FgEDAAAABQAAAABgAwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABAADgIQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABAA==";

        var enc = Convert.FromBase64String(encStr);
        var (ss, _) = SequenceSet.Decode(enc);

        ss.Size.ShouldBe(seqs.Length);
        foreach (var seq in seqs)
        {
            ss.Exists(seq).ShouldBeTrue();
        }
    }

    // Go: TestNoRaceSeqSetSizeComparison server/avl/norace_test.go:33
    [Fact]
    public void SizeComparison_LargeSet()
    {
        // Create 5M random entries out of 7M range.
        const int num = 5_000_000;
        const int max = 7_000_000;

        var rng = new Random(42);
        var seqs = new ulong[num];
        for (var i = 0; i < num; i++)
        {
            seqs[i] = (ulong)rng.NextInt64(max + 1);
        }

        // Insert into a dictionary to compare.
        var dmap = new HashSet<ulong>(num);
        foreach (var n in seqs)
        {
            dmap.Add(n);
        }

        // Insert into SequenceSet.
        var ss = new SequenceSet();
        foreach (var n in seqs)
        {
            ss.Insert(n);
        }

        // Verify sizes match.
        ss.Size.ShouldBe(dmap.Count);

        // Verify SequenceSet uses very few nodes relative to its element count.
        // With 2048 entries per node and 7M range, we expect ~ceil(7M/2048) = ~3419 nodes at most.
        ss.Nodes.ShouldBeLessThan(5000);
    }

    // Go: TestNoRaceSeqSetEncodeLarge server/avl/norace_test.go:81
    [Fact]
    public void EncodeLarge_RoundTrips()
    {
        const int num = 2_500_000;
        const int max = 5_000_000;

        var rng = new Random(42);
        var ss = new SequenceSet();
        for (var i = 0; i < num; i++)
        {
            ss.Insert((ulong)rng.NextInt64(max + 1));
        }

        var sw = Stopwatch.StartNew();
        var buf = ss.Encode();
        sw.Stop();

        // Encode should be fast (the Go test uses 1ms, we allow more for .NET JIT).
        sw.Elapsed.ShouldBeLessThan(TimeSpan.FromSeconds(1));

        sw.Restart();
        var (ss2, bytesRead) = SequenceSet.Decode(buf);
        sw.Stop();

        sw.Elapsed.ShouldBeLessThan(TimeSpan.FromSeconds(1));
        bytesRead.ShouldBe(buf.Length);
        ss2.Nodes.ShouldBe(ss.Nodes);
        ss2.Size.ShouldBe(ss.Size);
    }

    // Go: TestNoRaceSeqSetRelativeSpeed server/avl/norace_test.go:123
    [Fact]
    public void RelativeSpeed_Performance()
    {
        const int num = 1_000_000;
        const int max = 3_000_000;

        var rng = new Random(42);
        var seqs = new ulong[num];
        for (var i = 0; i < num; i++)
        {
            seqs[i] = (ulong)rng.NextInt64(max + 1);
        }

        // SequenceSet insert.
        var sw = Stopwatch.StartNew();
        var ss = new SequenceSet();
        foreach (var n in seqs)
        {
            ss.Insert(n);
        }

        var ssInsert = sw.Elapsed;

        // SequenceSet lookup.
        sw.Restart();
        foreach (var n in seqs)
        {
            ss.Exists(n).ShouldBeTrue();
        }

        var ssLookup = sw.Elapsed;

        // Dictionary insert.
        sw.Restart();
        var dmap = new HashSet<ulong>();
        foreach (var n in seqs)
        {
            dmap.Add(n);
        }

        var mapInsert = sw.Elapsed;

        // Dictionary lookup.
        sw.Restart();
        foreach (var n in seqs)
        {
            dmap.Contains(n).ShouldBeTrue();
        }

        var mapLookup = sw.Elapsed;

        // Relaxed bounds: SequenceSet insert should be no more than 10x slower.
        // (.NET JIT and test host overhead can be significant vs Go's simpler runtime.)
        ssInsert.ShouldBeLessThan(mapInsert * 10);
        ssLookup.ShouldBeLessThan(mapLookup * 10);
    }

    /// <summary>Verifies the AVL tree is balanced at every node.</summary>
    private static void VerifyBalanced(SequenceSet ss)
    {
        if (ss.Root == null)
        {
            return;
        }

        // Check all node heights and balance factors.
        SequenceSet.Node.NodeIter(ss.Root, n =>
        {
            var expectedHeight = SequenceSet.Node.MaxHeight(n) + 1;
            n.Height.ShouldBe(expectedHeight, $"Node height is wrong for node with base {n.Base}");
        });

        var bf = SequenceSet.Node.BalanceFactor(ss.Root);
        bf.ShouldBeInRange(-1, 1, "Tree is unbalanced at root");
    }

    /// <summary>Fisher-Yates shuffle.</summary>
    private static void Shuffle(List<ulong> list, Random rng)
    {
        for (var i = list.Count - 1; i > 0; i--)
        {
            var j = rng.Next(i + 1);
            (list[i], list[j]) = (list[j], list[i]);
        }
    }
}