/*
 * Copyright (c) 2006-Present, Redis Ltd.
 * All rights reserved.
 *
 * Licensed under your choice of the Redis Source Available License 2.0
 * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the
 * GNU Affero General Public License v3 (AGPLv3).
*/


#include "gtest/gtest.h"
#include "trie/trie_node.h"
#include "trie/trie_node_internal.h"  // whitebox: subtreeMaxScore invariant checks
#include "trie/trie.h"
#include "redismock/redismock.h"

#include <set>
#include <string>
#include <memory>
#include <functional>

typedef std::set<std::string> ElemSet;

class TrieTest : public ::testing::Test {};

static bool trieInsert(Trie *t, const char *s, size_t n) {
  return Trie_InsertStringBuffer(t, s, n, 1, 1, NULL, 0);
}
static bool trieInsert(Trie *t, const char *s) {
  return trieInsert(t, s, strlen(s));
}
static bool trieInsert(Trie *t, const std::string &s) {
  return trieInsert(t, s.c_str(), s.size());
}

static void *triePayload(Trie *t, const char *s, size_t len, bool exact) {
  if (len > TRIE_INITIAL_STRING_LEN * sizeof(rune)) {
    return nullptr;
  }
  runeBuf buf;
  rune *runes = runeBufFill(s, len, &buf, &len);
  TrieNode *node = Trie_GetNode(t, runes, len, exact, NULL);
  runeBufFree(&buf);
  return TrieNode_GetPayloadData(node);
}

static int rangeFunc(const rune *u16, size_t nrune, void *ctx, void *payload, size_t numDocsInTerm) {
  size_t n;
  char *s = runesToStr(u16, nrune, &n);
  std::string xs(s, n);
  free(s);
  ElemSet *e = (ElemSet *)ctx;
  assert(e->end() == e->find(xs));
  e->insert(xs);
  return REDISEARCH_OK;
}

static ElemSet trieIterRange(Trie *t, const char *begin, size_t nbegin, const char *end,
                             size_t nend) {
  rune r1[256] = {0};
  rune r2[256] = {0};
  size_t nr1, nr2;

  rune *r1Ptr = r1;
  rune *r2Ptr = r2;

  nr1 = strToRunesN(begin, nbegin, r1);
  nr2 = strToRunesN(end, nend, r2);

  if (!begin) {
    r1Ptr = NULL;
    nr1 = -1;
  }

  if (!end) {
    r2Ptr = NULL;
    nr2 = -1;
  }

  ElemSet foundElements;
  Trie_IterateRange(t, r1Ptr, nr1, true, r2Ptr, nr2, false,
                        rangeFunc, &foundElements);
  return foundElements;
}

static ElemSet trieIterRange(Trie *t, const char *begin, const char *end) {
  return trieIterRange(t, begin, begin ? strlen(begin) : 0, end, end ? strlen(end) : 0);
}

TEST_F(TrieTest, testBasicRange) {
  Trie *t = NewTrie(NULL, Trie_Sort_Lex);
  rune rbuf[TRIE_INITIAL_STRING_LEN + 1];
  for (size_t ii = 0; ii < 1000; ++ii) {
    char buf[64];
    snprintf(buf, sizeof(buf), "%lu", (unsigned long)ii);
    auto n = trieInsert(t, buf);
    ASSERT_TRUE(n);
  }

  //TrieNode_Print(t->root, 0, 0);

  // Get all numbers within the lexical range of 1 and 1Z
  auto ret = trieIterRange(t, "1", "1Z");
  ASSERT_EQ(111, ret.size());

  // What does a NULL range return? the entire trie
  ret = trieIterRange(t, NULL, NULL);
  ASSERT_EQ(Trie_Size(t), ret.size());

  // Min and max the same- should return only one value
  ret = trieIterRange(t, "1", "1");
  ASSERT_EQ(1, ret.size());

  ret = trieIterRange(t, "10", 2, "11", 2);
  ASSERT_EQ(11, ret.size());

  // Min and Min+1
  ret = trieIterRange(t, "10", 2, "10\x01", 3);
  ASSERT_EQ(1, ret.size());

  // No min, but has a max
  ret = trieIterRange(t, NULL, "5");
  ASSERT_EQ(445, ret.size());

  TrieType_Free(t);
}

TEST_F(TrieTest, testBasicRangeWithScore) {
  Trie *t = NewTrie(NULL, Trie_Sort_Score);
  rune rbuf[TRIE_INITIAL_STRING_LEN + 1];
  for (size_t ii = 0; ii < 1000; ++ii) {
    char buf[64];
    snprintf(buf, sizeof(buf), "%lu", (unsigned long)ii);
    auto n = trieInsert(t, buf);
    ASSERT_TRUE(n);
  }

  //TrieNode_Print(t->root, 0, 0);

  // Get all numbers within the lexical range of 1 and 1Z
  auto ret = trieIterRange(t, "1", "1Z");
  ASSERT_EQ(111, ret.size());

  // What does a NULL range return? the entire trie
  ret = trieIterRange(t, NULL, NULL);
  ASSERT_EQ(Trie_Size(t), ret.size());

  // Min and max the same- should return only one value
  ret = trieIterRange(t, "1", "1");
  ASSERT_EQ(1, ret.size());

  ret = trieIterRange(t, "10", 2, "11", 2);
  ASSERT_EQ(11, ret.size());

  // Min and Min+1
  ret = trieIterRange(t, "10", 2, "10\x01", 3);
  ASSERT_EQ(1, ret.size());

  // No min, but has a max
  ret = trieIterRange(t, NULL, "5");
  ASSERT_EQ(445, ret.size());

  TrieType_Free(t);
}

// Regression test for `rangeIterate` double-emission when the boundary value
// is a proper prefix of a child's collapsed label. The fixture uses textual
// terms (rather than testBasicRange's numeric ones) so that root children
// like "ban" have multi-character labels — only those expose the bug, because
// `rsb_gt`/`rsb_lt` treat e.g. "b" < "ban" and thus include the boundary
// child in the for-loop alongside the dedicated boundary recursion above it.
// `rangeFunc`'s `assert(e->end() == e->find(xs))` aborts on any duplicate
// emission, so a clean run is the regression signal.
TEST_F(TrieTest, testRangeBoundaryPrefix) {
  Trie *t = NewTrie(NULL, Trie_Sort_Lex);
  for (const char *term : {"apple", "banana", "band", "bandana", "cherry", "date"}) {
    ASSERT_TRUE(trieInsert(t, term));
  }

  // Min-only: [b, +inf). "b" is a proper prefix of the collapsed "ban" label.
  // Pre-fix: ban-subtree fires once via the boundary recursion and again
  // via the for-loop (`rsb_gt("b")` returns the same index as `beginEqIdx`).
  auto retMin = trieIterRange(t, "b", NULL);
  ElemSet expectedMin{"banana", "band", "bandana", "cherry", "date"};
  EXPECT_EQ(expectedMin, retMin);

  // Max-only: (-inf, banb). "ban" is a proper prefix of "banb", so
  // `rsb_lt("banb")` includes the "ban" subtree alongside the boundary
  // recursion. After the fix only entries strictly less than "banb" surface
  // ("band"/"bandana" are lex-greater than "banb").
  auto retMax = trieIterRange(t, NULL, "banb");
  ElemSet expectedMax{"apple", "banana"};
  EXPECT_EQ(expectedMax, retMax);

  TrieType_Free(t);
}

/**
 * This test ensures that the stack isn't overflown from all the frames.
 * The maximum trie depth cannot be greater than the maximum length of the
 * string.
 */
TEST_F(TrieTest, testDeepEntry) {
  Trie *t = NewTrie(NULL, Trie_Sort_Score);
  const size_t maxbuf = TRIE_INITIAL_STRING_LEN - 1;
  char manyOnes[maxbuf + 1];
  for (size_t ii = 0; ii < maxbuf; ++ii) {
    manyOnes[ii] = '1';
  }

  manyOnes[maxbuf] = 0;
  size_t manyLen = strlen(manyOnes);

  for (size_t ii = 0; ii < manyLen; ++ii) {
    size_t curlen = ii + 1;
    int rc = trieInsert(t, manyOnes, curlen);
    ASSERT_TRUE(rc);
    // printf("Inserting with len=%u: %d\n", curlen, rc);
  }

  auto ret = trieIterRange(t, "1", "1Z");
  ASSERT_EQ(maxbuf, ret.size());
  TrieType_Free(t);
}

/**
 * This test ensures payload isn't corrupted when the trie changes.
 */
TEST_F(TrieTest, testPayload) {
  char buf1[] = "world";

  Trie *t = NewTrie(NULL, Trie_Sort_Score);

  RSPayload payload = { .data = buf1, .len = 2 };
  Trie_InsertStringBuffer(t, buf1, 2, 1, 1, &payload, 0);
  payload.len = 4;
  Trie_InsertStringBuffer(t, buf1, 4, 1, 1, &payload, 0);
  payload.len = 5;
  Trie_InsertStringBuffer(t, buf1, 5, 1, 1, &payload, 0);
  payload.len = 3;
  Trie_InsertStringBuffer(t, buf1, 3, 1, 1, &payload, 0);

  char buf2[] = "work";
  payload = { .data = buf2, .len = 4 };
  Trie_InsertStringBuffer(t, buf2, 4, 1, 1, &payload, 0);


  // check for prefix of existing term
  // with exact returns null, w/o return load of next term
  ASSERT_EQ(strncmp((char*)triePayload(t, buf1, 1, 0), "wo", 2), 0);
  ASSERT_TRUE((char*)triePayload(t, buf1, 1, 1) == NULL);

  ASSERT_EQ(strncmp((char*)triePayload(t, buf1, 2, 1), "wo", 2), 0);
  ASSERT_EQ(strncmp((char*)triePayload(t, buf1, 3, 1), "wor", 3), 0);
  ASSERT_EQ(strncmp((char*)triePayload(t, buf1, 4, 1), "worl", 4), 0);
  ASSERT_EQ(strncmp((char*)triePayload(t, buf1, 5, 1), "world", 5), 0);
  ASSERT_EQ(strncmp((char*)triePayload(t, buf2, 4, 1), "work", 4), 0);

  ASSERT_EQ(Trie_Delete(t, buf1, 3), 1);
  ASSERT_EQ(strncmp((char*)triePayload(t, buf1, 2, 1), "wo", 2), 0);
  ASSERT_TRUE((char*)triePayload(t, buf1, 3, 1) == NULL);
  ASSERT_EQ(strncmp((char*)triePayload(t, buf1, 4, 1), "worl", 4), 0);
  ASSERT_EQ(strncmp((char*)triePayload(t, buf1, 5, 1), "world", 5), 0);
  ASSERT_EQ(strncmp((char*)triePayload(t, buf2, 4, 1), "work", 4), 0);

  ASSERT_EQ(Trie_Delete(t, buf1, 4), 1);
  ASSERT_EQ(strncmp((char*)triePayload(t, buf1, 2, 1), "wo", 2), 0);
  ASSERT_TRUE((char*)triePayload(t, buf1, 3, 1) == NULL);
  ASSERT_TRUE((char*)triePayload(t, buf1, 4, 1) == NULL);
  ASSERT_EQ(strncmp((char*)triePayload(t, buf1, 5, 1), "world", 5), 0);
  ASSERT_EQ(strncmp((char*)triePayload(t, buf2, 4, 1), "work", 4), 0);

  // testing with exact = 0
  // "wor" node exists with NULL payload.
  ASSERT_TRUE((char*)triePayload(t, buf1, 3, 0) == NULL);
  // "worl" does not exist but is partial offset of =>`wor`+`ld`.
  // payload of `ld` is returned.
  ASSERT_EQ(strncmp((char*)triePayload(t, buf1, 4, 0), "world", 5), 0);

  TrieType_Free(t);
}

/**
 * This test check free callback.
 */
void trieFreeCb(void *val) {
  char **str = (char **)val;
  rm_free(*str);
}

TEST_F(TrieTest, testFreeCallback) {
  Trie *t = NewTrie(trieFreeCb, Trie_Sort_Score);

  char buf[] = "world";
  char *str = rm_strdup("hello");

  RSPayload payload = { .data = (char *)&str, .len = sizeof(str) };
  Trie_InsertStringBuffer(t, buf, 5, 1, 1, &payload, 0);

  TrieType_Free(t);
}

void checkNext(TrieIterator *iter, const char *str) {
  char buf[16];
  rune *rstr = (rune *)&buf;
  t_len rlen;
  float score;
  RSPayload payload;

  TrieIterator_Next(iter, &rstr, &rlen, &payload, &score, NULL, NULL);
  size_t len;
  char *res_str = runesToStr(rstr, rlen, &len);
  ASSERT_STREQ(res_str, str);
  rm_free(res_str);
}

TEST_F(TrieTest, testLexOrder) {
  Trie *t = NewTrie(trieFreeCb, Trie_Sort_Lex);

  trieInsert(t, "hello");
  trieInsert(t, "world");
  trieInsert(t, "helen");
  trieInsert(t, "foo");
  trieInsert(t, "bar");
  trieInsert(t, "help");

  TrieIterator *iter = Trie_IterateAll(t);
  checkNext(iter, "bar");
  checkNext(iter, "foo");
  checkNext(iter, "helen");
  checkNext(iter, "hello");
  checkNext(iter, "help");
  checkNext(iter, "world");
  TrieIterator_Free(iter);

  Trie_Delete(t, "bar", 3);
  Trie_Delete(t, "hello", 5);
  Trie_Delete(t, "world", 5);

  iter = Trie_IterateAll(t);
  checkNext(iter, "foo");
  checkNext(iter, "helen");
  checkNext(iter, "help");
  TrieIterator_Free(iter);

  TrieType_Free(t);
}

bool trieInsertByScore(Trie *t, const char *s, float score) {
  return Trie_InsertStringBuffer(t, s, strlen(s), score, 1, NULL, 0);
}

bool trieContains(Trie *t, const char *s) {
  runeBuf buf;
  size_t len = strlen(s);
  rune *runes = runeBufFill(s, len, &buf, &len);
  if (!runes) {
    return false;
  }
  TrieNode *node = Trie_GetNode(t, runes, len, 0, NULL);
  runeBufFree(&buf);
  return node != NULL;
}

TEST_F(TrieTest, testScoreOrder) {
  Trie *t = NewTrie(trieFreeCb, Trie_Sort_Score);

  trieInsertByScore(t, "hello", 4);
  trieInsertByScore(t, "world", 2);
  trieInsertByScore(t, "foo", 6);
  trieInsertByScore(t, "bar", 1);
  trieInsertByScore(t, "help", 3);
  trieInsertByScore(t, "helen", 5);

  TrieIterator *iter = Trie_IterateAll(t);
  checkNext(iter, "foo");
  checkNext(iter, "helen");
  checkNext(iter, "hello");
  checkNext(iter, "help");
  checkNext(iter, "world");
  checkNext(iter, "bar");
  TrieIterator_Free(iter);

  Trie_Delete(t, "hello", 5);
  Trie_Delete(t, "world", 5);
  Trie_Delete(t, "bar", 3);

  iter = Trie_IterateAll(t);
  checkNext(iter, "foo");
  checkNext(iter, "helen");
  checkNext(iter, "help");
  TrieIterator_Free(iter);

  TrieType_Free(t);
}

// Fetch a node by its UTF-8 key through the public wrapper, so the test reads
// subtreeMaxScore without reaching into the opaque Trie struct for its root.
static TrieNode *getNode(Trie *t, const char *s) {
  runeBuf buf;
  size_t len = strlen(s);
  rune *runes = runeBufFill(s, len, &buf, &len);
  TrieNode *n = Trie_GetNode(t, runes, len, true, NULL);
  runeBufFree(&buf);
  return n;
}

// Regression test for the subtreeMaxScore staleness bug. subtreeMaxScore is the
// branch-and-bound upper bound Trie_CollectFuzzy (FT.SUGGET) prunes on, so it
// must always cover a node's own score and every descendant's. The buggy code
// folded only the score *delta* into the bound on two insert paths, leaving it
// under-estimated and causing valid suggestions to be silently pruned.
TEST_F(TrieTest, testSubtreeMaxScoreCoversIncrAndSplit) {
  Trie *t = NewTrie(NULL, Trie_Sort_Score);

  // ADD_INCR path: "beer"/"beet" share the internal node "bee". INCR "beer" by 3
  // (5 -> 8); the buggy code folded only the delta (3), leaving the bounds at 5.
  Trie_InsertStringBuffer(t, "beer", 4, 5.0, 0, NULL, 1);
  Trie_InsertStringBuffer(t, "beet", 4, 5.0, 0, NULL, 1);
  Trie_InsertStringBuffer(t, "beer", 4, 3.0, 1, NULL, 1);

  // Split-exact path: "zoom" is one compressed node; inserting its proper prefix
  // "zoo" splits it and makes "zoo" terminal with score 9. The buggy code never
  // folded that score into the split node's bound, leaving it at "zoom"'s 5.
  Trie_InsertStringBuffer(t, "zoom", 4, 5.0, 0, NULL, 1);
  Trie_InsertStringBuffer(t, "zoo", 3, 9.0, 0, NULL, 1);

  EXPECT_FLOAT_EQ(getNode(t, "beer")->score, 8.0f);      // INCR applied the delta
  EXPECT_GE(getNode(t, "beer")->subtreeMaxScore, 8.0f);  // bound covers the total
  EXPECT_GE(getNode(t, "bee")->subtreeMaxScore, 8.0f);   // ancestor covers descendant
  EXPECT_GE(getNode(t, "zoo")->subtreeMaxScore, 9.0f);   // split node folded its score

  TrieType_Free(t);
}

// Add a UTF-8 key directly to a raw root node, bypassing the Trie wrapper so
// the test owns the root pointer and can walk the structure afterwards.
static void addRaw(TrieNode **root, const char *s, float score, TrieAddOp op) {
  runeBuf buf;
  size_t len = strlen(s);
  rune *runes = runeBufFill(s, len, &buf, &len);
  TrieNode_Add(root, runes, len, NULL, score, op, NULL, 1);
  runeBufFree(&buf);
}

// Recursively assert that every node keeps its children ordered by descending
// subtreeMaxScore — the order the score-mode iterator relies on to visit
// high-scoring branches first.
static void assertChildrenScoreOrdered(const TrieNode *n) {
  TrieNode **children = TrieNode_Children(n);
  for (t_len i = 0; i < TrieNode_NumChildren(n); i++) {
    if (i + 1 < TrieNode_NumChildren(n)) {
      EXPECT_GE(children[i]->subtreeMaxScore, children[i + 1]->subtreeMaxScore)
          << "children out of descending subtreeMaxScore order";
    }
    assertChildrenScoreOrdered(children[i]);
  }
}

// The insert path restores child order with a single-element rotation instead of
// a full sort. Storm a small trie with rank-crossing INCRs, then verify both the
// order invariant and (via exact lookups) that the rotation kept the parallel
// child-key array consistent with the children.
TEST_F(TrieTest, testScoreOrderMaintainedAfterIncrStorm) {
  rune emptyRoot[1] = {0};
  TrieNode *root = __newTrieNode(emptyRoot, 0, 0, NULL, 0, 0, 0.0f, 0, Trie_Sort_Score, 0);

  const char *keys[] = {"alpha", "alps", "beer", "beet", "bee", "gamma", "gap", "delta"};
  const size_t numKeys = sizeof(keys) / sizeof(keys[0]);
  float expected[numKeys];

  for (size_t i = 0; i < numKeys; i++) {
    addRaw(&root, keys[i], 1.0f, ADD_REPLACE);
    expected[i] = 1.0f;
  }

  // Uneven, shifting deltas so sibling ranks keep crossing at every level and
  // the rotation has to move children by more than one slot.
  for (int round = 0; round < 20; round++) {
    for (size_t i = 0; i < numKeys; i++) {
      float delta = (float)((i + round) % 4 + 1);
      addRaw(&root, keys[i], delta, ADD_INCR);
      expected[i] += delta;
    }
    assertChildrenScoreOrdered(root);
  }

  for (size_t i = 0; i < numKeys; i++) {
    runeBuf buf;
    size_t len = strlen(keys[i]);
    rune *runes = runeBufFill(keys[i], len, &buf, &len);
    TrieNode *node = TrieNode_Get(root, runes, len, true, NULL);
    runeBufFree(&buf);
    ASSERT_NE(node, nullptr) << keys[i];
    EXPECT_FLOAT_EQ(node->score, expected[i]) << keys[i];
  }

  TrieNode_Free(root, NULL);
}

// Assert the first rune of each child of root, in child-array order.
static void assertChildOrder(TrieNode *root, const char *firstRunes) {
  size_t expected = strlen(firstRunes);
  ASSERT_EQ(TrieNode_NumChildren(root), expected);
  TrieNode **children = TrieNode_Children(root);
  for (size_t i = 0; i < expected; i++) {
    EXPECT_EQ(children[i]->str[0], (rune)firstRunes[i]) << "child " << i;
  }
}

// Whitebox tests for __trieNode_rotateChildIntoPlace: raise one child's bound,
// rotate, check order, tie stability, and key-cache consistency.
TEST_F(TrieTest, testRotateChildIntoPlace) {
  rune emptyRoot[1] = {0};
  TrieNode *root = __newTrieNode(emptyRoot, 0, 0, NULL, 0, 0, 0.0f, 0, Trie_Sort_Score, 0);

  // descending scores append in order: children are [delta, charlie, bravo, alpha]
  addRaw(&root, "delta", 9.0f, ADD_REPLACE);
  addRaw(&root, "charlie", 7.0f, ADD_REPLACE);
  addRaw(&root, "bravo", 5.0f, ADD_REPLACE);
  addRaw(&root, "alpha", 3.0f, ADD_REPLACE);
  assertChildOrder(root, "dcba");
  TrieNode **children = TrieNode_Children(root);

  // no-move: bravo's bound rises but stays below its left neighbor
  children[2]->subtreeMaxScore = 6.0f;
  __trieNode_rotateChildIntoPlace(root, 2);
  assertChildOrder(root, "dcba");

  // tie stability: bound rises to exactly charlie's; no move
  children[2]->subtreeMaxScore = 7.0f;
  __trieNode_rotateChildIntoPlace(root, 2);
  assertChildOrder(root, "dcba");

  // multi-slot move: alpha's bound rises past bravo and charlie but not delta
  children[3]->subtreeMaxScore = 8.0f;
  __trieNode_rotateChildIntoPlace(root, 3);
  assertChildOrder(root, "dacb");

  // move to front: bravo's bound rises past everything
  children[3]->subtreeMaxScore = 10.0f;
  __trieNode_rotateChildIntoPlace(root, 3);
  assertChildOrder(root, "bdac");

  // key cache stayed in sync: every key still reachable by exact lookup
  const char *keys[] = {"alpha", "bravo", "charlie", "delta"};
  const float scores[] = {3.0f, 5.0f, 7.0f, 9.0f};
  for (size_t i = 0; i < 4; i++) {
    runeBuf buf;
    size_t len = strlen(keys[i]);
    rune *runes = runeBufFill(keys[i], len, &buf, &len);
    TrieNode *node = TrieNode_Get(root, runes, len, true, NULL);
    runeBufFree(&buf);
    ASSERT_NE(node, nullptr) << keys[i];
    EXPECT_FLOAT_EQ(node->score, scores[i]) << keys[i];
  }

  TrieNode_Free(root, NULL);
}

/* leave for future benchmarks if needed
TEST_F(TrieTest, testbenchmark) {
  Trie *t = NewTrie(trieFreeCb, Trie_Sort_Lex);
  char buf[128];
  int count = 1024 * 1024 * 8;
  for (size_t i = 0; i < count; ++i) {
    int random = rand() % (count / 5);
    snprintf(buf, sizeof(buf), "%x", random);
    Trie_InsertStringBuffer(t, buf, strlen(buf), 1, 0,10 NULL);
  }

  TrieType_Free(t);
}*/

// Helper function to compare two tries for equality
static bool compareTrieContents(Trie *original, Trie *loaded) {
  if (Trie_Size(original) != Trie_Size(loaded)) {
    return false;
  }

  // Compare all entries using iterators
  TrieIterator *origIter = Trie_IterateAll(original);
  TrieIterator *loadedIter = Trie_IterateAll(loaded);

  std::unique_ptr<TrieIterator, std::function<void(TrieIterator *)>> origIterPtr(origIter, [](TrieIterator *iter) {
    TrieIterator_Free(iter);
  });
  std::unique_ptr<TrieIterator, std::function<void(TrieIterator *)>> loadedIterPtr(loadedIter, [](TrieIterator *iter) {
    TrieIterator_Free(iter);
  });

  rune *origRstr, *loadedRstr;
  t_len origLen, loadedLen;
  float origScore, loadedScore;
  RSPayload origPayload, loadedPayload;

  while (true) {
    int origHasNext = TrieIterator_Next(origIter, &origRstr, &origLen, &origPayload, &origScore, NULL, NULL);
    int loadedHasNext = TrieIterator_Next(loadedIter, &loadedRstr, &loadedLen, &loadedPayload, &loadedScore, NULL, NULL);

    if (origHasNext != loadedHasNext) {
      return false;
    }

    if (!origHasNext) {
      break; // Both iterators finished
    }

    // Compare strings
    if (origLen != loadedLen) {
      return false;
    }

    size_t origStrLen, loadedStrLen;
    char *origStr = runesToStr(origRstr, origLen, &origStrLen);
    char *loadedStr = runesToStr(loadedRstr, loadedLen, &loadedStrLen);

    std::unique_ptr<char, std::function<void(char *)>> origStrPtr(origStr, [](char *str) { rm_free(str); });
    std::unique_ptr<char, std::function<void(char *)>> loadedStrPtr(loadedStr, [](char *str) { rm_free(str); });

    if (origStrLen != loadedStrLen || strncmp(origStr, loadedStr, origStrLen) != 0) {
      return false;
    }

    // Compare scores
    if (origScore != loadedScore) {
      return false;
    }

    // Compare payloads
    if (origPayload.len != loadedPayload.len) {
      return false;
    }

    if (origPayload.len > 0 && loadedPayload.len > 0) {
      if (memcmp(origPayload.data, loadedPayload.data, origPayload.len) != 0) {
        return false;
      }
    } else if ((origPayload.data == NULL) != (loadedPayload.data == NULL)) {
      return false;
    }
  }

  return true;
}

TEST_F(TrieTest, testBasicRdbSaveLoad) {
  // Create a trie with some test data
  Trie *originalTrie = NewTrie(NULL, Trie_Sort_Score);
  std::unique_ptr<Trie, std::function<void(Trie *)>> originalTriePtr(originalTrie, [](Trie *trie) {
    TrieType_Free(trie);
  });

  // Insert complex test data with prefixes and extensions to stress the trie
  trieInsertByScore(originalTrie, "app", 5.0);         // Base word
  trieInsertByScore(originalTrie, "apple", 3.0);       // Extension of "app"
  trieInsertByScore(originalTrie, "application", 7.0); // Extension of "app"
  trieInsertByScore(originalTrie, "apply", 1.0);       // Extension of "app"
  trieInsertByScore(originalTrie, "applied", 4.0);     // Extension of "apply"
  trieInsertByScore(originalTrie, "book", 6.0);        // Base word
  trieInsertByScore(originalTrie, "books", 8.0);       // Extension of "book"
  trieInsertByScore(originalTrie, "booking", 2.0);     // Extension of "book"

  ASSERT_EQ(8, Trie_Size(originalTrie));

  // Create RDB IO context
  RedisModuleIO *io = RMCK_CreateRdbIO();
  std::unique_ptr<RedisModuleIO, std::function<void(RedisModuleIO *)>> ioPtr(io, [](RedisModuleIO *io) {
    RMCK_FreeRdbIO(io);
  });
  ASSERT_TRUE(io != nullptr);

  // Save the trie to RDB
  TrieType_RdbSave(io, originalTrie);
  EXPECT_EQ(0, RMCK_IsIOError(io));

  // Reset read position to load it back
  io->read_pos = 0;

  // Load the trie from RDB
  Trie *loadedTrie = (Trie *)TrieType_RdbLoad(io, TRIE_ENCVER_CURRENT);
  std::unique_ptr<Trie, std::function<void(Trie *)>> loadedTriePtr(loadedTrie, [](Trie *trie) {
    TrieType_Free(trie);
  });
  ASSERT_TRUE(loadedTrie != nullptr);
  EXPECT_EQ(0, RMCK_IsIOError(io));

  // Compare the original and loaded tries
  EXPECT_EQ(Trie_Size(originalTrie), Trie_Size(loadedTrie));

  // Verify all entries are present in the loaded trie
  EXPECT_TRUE(trieContains(loadedTrie, "app"));
  EXPECT_TRUE(trieContains(loadedTrie, "apple"));
  EXPECT_TRUE(trieContains(loadedTrie, "application"));
  EXPECT_TRUE(trieContains(loadedTrie, "apply"));
  EXPECT_TRUE(trieContains(loadedTrie, "applied"));
  EXPECT_TRUE(trieContains(loadedTrie, "book"));
  EXPECT_TRUE(trieContains(loadedTrie, "books"));
  EXPECT_TRUE(trieContains(loadedTrie, "booking"));
}

TEST_F(TrieTest, testRdbSaveLoadWithPayloads) {
  // Create a trie with payloads
  Trie *originalTrie = NewTrie(NULL, Trie_Sort_Score);
  std::unique_ptr<Trie, std::function<void(Trie *)>> originalTriePtr(originalTrie, [](Trie *trie) {
    TrieType_Free(trie);
  });

  // Insert complex test data with payloads - includes prefixes and extensions
  char payload1[] = "payload_run";
  char payload2[] = "payload_running";
  char payload3[] = "payload_runner";

  RSPayload p1 = {.data = payload1, .len = strlen(payload1)};
  RSPayload p2 = {.data = payload2, .len = strlen(payload2)};
  RSPayload p3 = {.data = payload3, .len = strlen(payload3)};

  bool r1 = Trie_InsertStringBuffer(originalTrie, "run", 3, 5.0, 0, &p1, 0);        // Base word with payload
  bool r2 = Trie_InsertStringBuffer(originalTrie, "running", 7, 3.0, 0, &p2, 0);    // Extension with payload
  bool r3 = Trie_InsertStringBuffer(originalTrie, "runner", 6, 4.0, 0, &p3, 0);     // Extension with payload

  EXPECT_EQ(3, Trie_Size(originalTrie));

  // Create RDB IO context
  RedisModuleIO *io = RMCK_CreateRdbIO();
  std::unique_ptr<RedisModuleIO, std::function<void(RedisModuleIO *)>> ioPtr(io, [](RedisModuleIO *io) {
    RMCK_FreeRdbIO(io);
  });
  ASSERT_TRUE(io != nullptr);

  // Save the trie to RDB
  TrieType_RdbSave(io, originalTrie);
  EXPECT_EQ(0, RMCK_IsIOError(io));

  // Reset read position to load it back
  io->read_pos = 0;

  // Load the trie from RDB (with payloads)
  Trie *loadedTrie = (Trie *)TrieType_GenericLoad(io, true, true, Trie_Sort_Score);
  std::unique_ptr<Trie, std::function<void(Trie *)>> loadedTriePtr(loadedTrie, [](Trie *trie) {
    TrieType_Free(trie);
  });
  ASSERT_TRUE(loadedTrie != nullptr);
  EXPECT_EQ(0, RMCK_IsIOError(io));

  // Compare the original and loaded tries
  EXPECT_EQ(Trie_Size(originalTrie), Trie_Size(loadedTrie));

  // Verify all entries are present in the loaded trie
  EXPECT_TRUE(trieContains(loadedTrie, "run"));
  EXPECT_TRUE(trieContains(loadedTrie, "running"));
  EXPECT_TRUE(trieContains(loadedTrie, "runner"));

  // Verify specific payloads are preserved
  void *loadedPayload1 = triePayload(loadedTrie, "run", 3, true);
  void *loadedPayload2 = triePayload(loadedTrie, "running", 7, true);
  void *loadedPayload3 = triePayload(loadedTrie, "runner", 6, true);

  ASSERT_TRUE(loadedPayload1 != nullptr);
  ASSERT_TRUE(loadedPayload2 != nullptr);
  ASSERT_TRUE(loadedPayload3 != nullptr);

  EXPECT_EQ(0, strncmp(payload1, (char *)loadedPayload1, strlen(payload1)));
  EXPECT_EQ(0, strncmp(payload2, (char *)loadedPayload2, strlen(payload2)));
  EXPECT_EQ(0, strncmp(payload3, (char *)loadedPayload3, strlen(payload3)));
}

TEST_F(TrieTest, testRdbSaveLoadPayloadsNotSerialized) {
  // Create a trie with payloads but save without serializing them
  Trie *originalTrie = NewTrie(NULL, Trie_Sort_Score);
  std::unique_ptr<Trie, std::function<void(Trie *)>> originalTriePtr(originalTrie, [](Trie *trie) {
    TrieType_Free(trie);
  });

  // Insert complex test data with payloads - includes prefixes and extensions
  char payload1[] = "payload_car";
  char payload2[] = "payload_care";
  char payload3[] = "payload_careful";

  RSPayload p1 = {.data = payload1, .len = strlen(payload1)};
  RSPayload p2 = {.data = payload2, .len = strlen(payload2)};
  RSPayload p3 = {.data = payload3, .len = strlen(payload3)};

  Trie_InsertStringBuffer(originalTrie, "car", 3, 8.0, 0, &p1, 0);        // Base word with payload
  Trie_InsertStringBuffer(originalTrie, "care", 4, 6.0, 0, &p2, 0);       // Extension with payload
  Trie_InsertStringBuffer(originalTrie, "careful", 7, 4.0, 0, &p3, 0);    // Extension with payload

  EXPECT_EQ(3, Trie_Size(originalTrie));

  // Create RDB IO context
  RedisModuleIO *io = RMCK_CreateRdbIO();
  std::unique_ptr<RedisModuleIO, std::function<void(RedisModuleIO *)>> ioPtr(io, [](RedisModuleIO *io) {
    RMCK_FreeRdbIO(io);
  });
  ASSERT_TRUE(io != nullptr);

  // Save the trie to RDB WITHOUT payloads (savePayloads = false) and numDocs (saveNumDocs = false)
  TrieType_GenericSave(io, originalTrie, false, false);
  EXPECT_EQ(0, RMCK_IsIOError(io));

  // Reset read position to load it back
  io->read_pos = 0;

  // Load the trie from RDB WITHOUT payloads (loadPayloads = false) and numDocs (loadNumDocs = false)
  Trie *loadedTrie = (Trie *)TrieType_GenericLoad(io, false, false, Trie_Sort_Score);
  std::unique_ptr<Trie, std::function<void(Trie *)>> loadedTriePtr(loadedTrie, [](Trie *trie) {
    TrieType_Free(trie);
  });
  ASSERT_TRUE(loadedTrie != nullptr);
  EXPECT_EQ(0, RMCK_IsIOError(io));

  // Compare the original and loaded tries - sizes should match
  EXPECT_EQ(Trie_Size(originalTrie), Trie_Size(loadedTrie));

  // Verify all entries are present in the loaded trie
  EXPECT_TRUE(trieContains(loadedTrie, "car"));
  EXPECT_TRUE(trieContains(loadedTrie, "care"));
  EXPECT_TRUE(trieContains(loadedTrie, "careful"));

  // Verify that payloads are NOT preserved (should be null)
  void *loadedPayload1 = triePayload(loadedTrie, "car", 3, true);
  void *loadedPayload2 = triePayload(loadedTrie, "care", 4, true);
  void *loadedPayload3 = triePayload(loadedTrie, "careful", 7, true);

  EXPECT_TRUE(loadedPayload1 == nullptr);  // Payload should not be preserved
  EXPECT_TRUE(loadedPayload2 == nullptr);  // Payload should not be preserved
  EXPECT_TRUE(loadedPayload3 == nullptr);  // Payload should not be preserved
}

TEST_F(TrieTest, testRdbSaveLoadWithoutPayloads) {
  // Create a trie and insert entries WITHOUT payloads
  Trie *originalTrie = NewTrie(NULL, Trie_Sort_Score);
  std::unique_ptr<Trie, std::function<void(Trie *)>> originalTriePtr(originalTrie, [](Trie *trie) {
    TrieType_Free(trie);
  });

  char payload1[] = "payload_1";
  char payload2[] = "payload_2";

  RSPayload p1 = {.data = payload1, .len = strlen(payload1)};
  RSPayload p2 = {.data = payload2, .len = strlen(payload2)};

  // Insert complex test data WITHOUT payloads - includes prefixes and extensions
  Trie_InsertStringBuffer(originalTrie, "hello", 5, 8.0, 0, NULL, 0);     // Base word without payload
  Trie_InsertStringBuffer(originalTrie, "hell", 4, 6.0, 0, &p1, 0);      // Prefix with payload
  Trie_InsertStringBuffer(originalTrie, "help", 4, 7.0, 0, NULL, 0);      // Related word without payload
  Trie_InsertStringBuffer(originalTrie, "helper", 6, 5.0, 0, &p2, 0);    // Extension with payload

  EXPECT_EQ(4, Trie_Size(originalTrie));

  // Create RDB IO context
  RedisModuleIO *io = RMCK_CreateRdbIO();
  std::unique_ptr<RedisModuleIO, std::function<void(RedisModuleIO *)>> ioPtr(io, [](RedisModuleIO *io) {
    RMCK_FreeRdbIO(io);
  });
  ASSERT_TRUE(io != nullptr);

  // Save the trie to RDB
  TrieType_GenericSave(io, originalTrie, false, false);
  EXPECT_EQ(0, RMCK_IsIOError(io));

  // Reset read position to load it back
  io->read_pos = 0;

  // Load the trie from RDB WITHOUT payloads (loadPayloads = false) and numDocs (loadNumDocs = false) to match the save operation
  Trie *loadedTrie = (Trie *)TrieType_GenericLoad(io, false, false, Trie_Sort_Score);
  std::unique_ptr<Trie, std::function<void(Trie *)>> loadedTriePtr(loadedTrie, [](Trie *trie) {
    TrieType_Free(trie);
  });
  ASSERT_TRUE(loadedTrie != nullptr);
  EXPECT_EQ(0, RMCK_IsIOError(io));

  // Compare sizes - entries should be preserved
  EXPECT_EQ(Trie_Size(originalTrie), Trie_Size(loadedTrie));

  // Verify all entries are present in the loaded trie
  EXPECT_TRUE(trieContains(loadedTrie, "hello"));
  EXPECT_TRUE(trieContains(loadedTrie, "hell"));
  EXPECT_TRUE(trieContains(loadedTrie, "help"));
  EXPECT_TRUE(trieContains(loadedTrie, "helper"));

  // Verify that payloads remain NULL (since none were inserted)
  void *loadedPayload1 = triePayload(loadedTrie, "hello", 5, true);
  void *loadedPayload2 = triePayload(loadedTrie, "hell", 4, true);
  void *loadedPayload3 = triePayload(loadedTrie, "help", 4, true);
  void *loadedPayload4 = triePayload(loadedTrie, "helper", 6, true);

  EXPECT_TRUE(loadedPayload1 == nullptr);  // No payload was inserted
  EXPECT_TRUE(loadedPayload2 == nullptr);  // No payload was inserted
  EXPECT_TRUE(loadedPayload3 == nullptr);  // No payload was inserted
  EXPECT_TRUE(loadedPayload4 == nullptr);  // No payload was inserted
}

TEST_F(TrieTest, testRdbSaveLoadEmptyTrie) {
  // Create an empty trie
  Trie *originalTrie = NewTrie(NULL, Trie_Sort_Score);
  std::unique_ptr<Trie, std::function<void(Trie *)>> originalTriePtr(originalTrie, [](Trie *trie) {
    TrieType_Free(trie);
  });

  ASSERT_EQ(0, Trie_Size(originalTrie));

  // Create RDB IO context
  RedisModuleIO *io = RMCK_CreateRdbIO();
  std::unique_ptr<RedisModuleIO, std::function<void(RedisModuleIO *)>> ioPtr(io, [](RedisModuleIO *io) {
    RMCK_FreeRdbIO(io);
  });
  ASSERT_TRUE(io != nullptr);

  // Save the empty trie to RDB
  TrieType_RdbSave(io, originalTrie);
  EXPECT_EQ(0, RMCK_IsIOError(io));

  // Reset read position to load it back
  io->read_pos = 0;

  // Load the trie from RDB
  Trie *loadedTrie = (Trie *)TrieType_RdbLoad(io, TRIE_ENCVER_CURRENT);
  std::unique_ptr<Trie, std::function<void(Trie *)>> loadedTriePtr(loadedTrie, [](Trie *trie) {
    TrieType_Free(trie);
  });
  ASSERT_TRUE(loadedTrie != nullptr);
  EXPECT_EQ(0, RMCK_IsIOError(io));

  // Compare the original and loaded tries
  EXPECT_EQ(0, Trie_Size(loadedTrie));
  EXPECT_EQ(Trie_Size(originalTrie), Trie_Size(loadedTrie));
}

TEST_F(TrieTest, testRdbSaveLoadLexSortedTrie) {
  // Create a trie with lexical sorting - this is the only difference from testBasicRdbSaveLoad
  Trie *originalTrie = NewTrie(NULL, Trie_Sort_Lex);
  std::unique_ptr<Trie, std::function<void(Trie *)>> originalTriePtr(originalTrie, [](Trie *trie) {
    TrieType_Free(trie);
  });

  // Insert complex test data with prefixes, extensions, and overlapping words
  // This stresses the trie implementation with hierarchical relationships
  trieInsertByScore(originalTrie, "test", 5.0);        // Base word
  trieInsertByScore(originalTrie, "testing", 4.0);     // Extension of "test"
  trieInsertByScore(originalTrie, "tester", 3.0);      // Another extension of "test"
  trieInsertByScore(originalTrie, "tests", 6.0);       // Plural of "test"
  trieInsertByScore(originalTrie, "te", 2.0);          // Prefix of "test"
  trieInsertByScore(originalTrie, "hello", 8.0);       // Base word
  trieInsertByScore(originalTrie, "hell", 7.0);        // Prefix of "hello"
  trieInsertByScore(originalTrie, "help", 9.0);        // Shares prefix "hel" with "hello"
  trieInsertByScore(originalTrie, "helper", 1.0);      // Extension of "help"
  trieInsertByScore(originalTrie, "helping", 10.0);    // Another extension of "help"
  trieInsertByScore(originalTrie, "car", 11.0);        // Base word
  trieInsertByScore(originalTrie, "care", 12.0);       // Extension of "car"
  trieInsertByScore(originalTrie, "careful", 13.0);    // Extension of "care"
  trieInsertByScore(originalTrie, "carefully", 14.0);  // Extension of "careful"

  ASSERT_EQ(14, Trie_Size(originalTrie));

  // Verify all entries exist in the original trie
  EXPECT_TRUE(trieContains(originalTrie, "test"));
  EXPECT_TRUE(trieContains(originalTrie, "testing"));
  EXPECT_TRUE(trieContains(originalTrie, "tester"));
  EXPECT_TRUE(trieContains(originalTrie, "tests"));
  EXPECT_TRUE(trieContains(originalTrie, "te"));
  EXPECT_TRUE(trieContains(originalTrie, "hello"));
  EXPECT_TRUE(trieContains(originalTrie, "hell"));
  EXPECT_TRUE(trieContains(originalTrie, "help"));
  EXPECT_TRUE(trieContains(originalTrie, "helper"));
  EXPECT_TRUE(trieContains(originalTrie, "helping"));
  EXPECT_TRUE(trieContains(originalTrie, "car"));
  EXPECT_TRUE(trieContains(originalTrie, "care"));
  EXPECT_TRUE(trieContains(originalTrie, "careful"));
  EXPECT_TRUE(trieContains(originalTrie, "carefully"));

  // Create RDB IO context
  RedisModuleIO *io = RMCK_CreateRdbIO();
  std::unique_ptr<RedisModuleIO, std::function<void(RedisModuleIO *)>> ioPtr(io, [](RedisModuleIO *io) {
    RMCK_FreeRdbIO(io);
  });
  ASSERT_TRUE(io != nullptr);

  // Save the trie to RDB
  TrieType_RdbSave(io, originalTrie);
  EXPECT_EQ(0, RMCK_IsIOError(io));

  // Reset read position to load it back
  io->read_pos = 0;

  // Load the trie from RDB
  Trie *loadedTrie = (Trie *)TrieType_RdbLoad(io, TRIE_ENCVER_CURRENT);
  std::unique_ptr<Trie, std::function<void(Trie *)>> loadedTriePtr(loadedTrie, [](Trie *trie) {
    TrieType_Free(trie);
  });
  ASSERT_TRUE(loadedTrie != nullptr);
  EXPECT_EQ(0, RMCK_IsIOError(io));

  // Compare the original and loaded tries
  EXPECT_EQ(Trie_Size(originalTrie), Trie_Size(loadedTrie));

  // Note: TrieType_RdbLoad (the registered TrieType callback) always reconstructs
  // with Trie_Sort_Score because the only producer of the registered type is FT.SUGADD.
  // All entries should still be present, even though the sorting mode changed.

  // Verify all entries are present in the loaded trie
  EXPECT_TRUE(trieContains(loadedTrie, "test"));
  EXPECT_TRUE(trieContains(loadedTrie, "testing"));
  EXPECT_TRUE(trieContains(loadedTrie, "tester"));
  EXPECT_TRUE(trieContains(loadedTrie, "tests"));
  EXPECT_TRUE(trieContains(loadedTrie, "te"));
  EXPECT_TRUE(trieContains(loadedTrie, "hello"));
  EXPECT_TRUE(trieContains(loadedTrie, "hell"));
  EXPECT_TRUE(trieContains(loadedTrie, "help"));
  EXPECT_TRUE(trieContains(loadedTrie, "helper"));
  EXPECT_TRUE(trieContains(loadedTrie, "helping"));
  EXPECT_TRUE(trieContains(loadedTrie, "car"));
  EXPECT_TRUE(trieContains(loadedTrie, "care"));
  EXPECT_TRUE(trieContains(loadedTrie, "careful"));
  EXPECT_TRUE(trieContains(loadedTrie, "carefully"));

  // Since the sorting mode changes during RDB load, we can't use compareTrieContents
  // which expects the same iteration order. Instead, we verify that all entries exist
  // and the size matches.
}

// Helper function to insert with numDocs
static bool trieInsertWithNumDocs(Trie *t, const char *s, float score, size_t numDocs) {
  return Trie_InsertStringBuffer(t, s, strlen(s), score, 0, NULL, numDocs);
}

// Helper function to get numDocs from a trie node
static size_t trieGetNumDocs(Trie *t, const char *s) {
  runeBuf buf;
  size_t runeLen = strlen(s);
  rune *runes = runeBufFill(s, runeLen, &buf, &runeLen);
  TrieNode *node = Trie_GetNode(t, runes, runeLen, true, NULL);
  runeBufFree(&buf);
  if (node == NULL) {
    return 0;
  }
  return TrieNode_NumDocs(node);
}

TEST_F(TrieTest, testDecrementNumDocsRepresentable) {
  Trie *t = NewTrie(NULL, Trie_Sort_Score);
  std::unique_ptr<Trie, std::function<void(Trie *)>> tp(t, [](Trie *p) { TrieType_Free(p); });

  trieInsertWithNumDocs(t, "hello", 1.0, 3);

  EXPECT_EQ(TRIE_DECR_UPDATED, Trie_DecrementNumDocs(t, "hello", strlen("hello"), 1));
  EXPECT_EQ(2, trieGetNumDocs(t, "hello"));

  EXPECT_EQ(TRIE_DECR_DELETED, Trie_DecrementNumDocs(t, "hello", strlen("hello"), 2));
  EXPECT_EQ(0, trieGetNumDocs(t, "hello"));
}

// A representable term never inserted must stay distinct (NOT_FOUND) from the
// unrepresentable case, so the disk-compaction caller can still assert on it.
TEST_F(TrieTest, testDecrementNumDocsMissingRepresentable) {
  Trie *t = NewTrie(NULL, Trie_Sort_Score);
  std::unique_ptr<Trie, std::function<void(Trie *)>> tp(t, [](Trie *p) { TrieType_Free(p); });

  trieInsertWithNumDocs(t, "hello", 1.0, 3);
  EXPECT_EQ(TRIE_DECR_NOT_FOUND, Trie_DecrementNumDocs(t, "absent", strlen("absent"), 1));
}

// A term at/over the trie's TRIE_INITIAL_STRING_LEN rune cap is never inserted,
// so decrementing it reports UNSUPPORTED, not NOT_FOUND.
TEST_F(TrieTest, testDecrementNumDocsUnrepresentableLongTerm) {
  Trie *t = NewTrie(NULL, Trie_Sort_Score);
  std::unique_ptr<Trie, std::function<void(Trie *)>> tp(t, [](Trie *p) { TrieType_Free(p); });

  // 255 runes: representable, inserts and decrements normally.
  std::string ascii255(TRIE_INITIAL_STRING_LEN - 1, 'a');
  ASSERT_TRUE(trieInsertWithNumDocs(t, ascii255.c_str(), 1.0, 1));
  EXPECT_EQ(TRIE_DECR_DELETED, Trie_DecrementNumDocs(t, ascii255.c_str(), ascii255.size(), 1));

  // 256 / 257 runes: trip the rune-length guard.
  std::string ascii256(TRIE_INITIAL_STRING_LEN, 'a');
  std::string ascii257(TRIE_INITIAL_STRING_LEN + 1, 'a');
  EXPECT_EQ(TRIE_DECR_UNSUPPORTED, Trie_DecrementNumDocs(t, ascii256.c_str(), ascii256.size(), 1));
  EXPECT_EQ(TRIE_DECR_UNSUPPORTED, Trie_DecrementNumDocs(t, ascii257.c_str(), ascii257.size(), 1));

  // 256 copies of U+754C (界), 3 bytes each: trips the earlier byte-length guard.
  std::string cjk256;
  for (int i = 0; i < TRIE_INITIAL_STRING_LEN; i++) cjk256 += "\xE7\x95\x8C";
  EXPECT_EQ(TRIE_DECR_UNSUPPORTED, Trie_DecrementNumDocs(t, cjk256.c_str(), cjk256.size(), 1));
}

// Regression: TrieType_GenericLoad must preserve sort mode across reload,
// or Trie_IterateRange's binary search breaks on Lex-sourced tries.
TEST_F(TrieTest, testRdbSaveLoadLexRangePreservesQueries) {
  Trie *original = NewTrie(NULL, Trie_Sort_Lex);
  std::unique_ptr<Trie, std::function<void(Trie *)>> originalPtr(
      original, [](Trie *t) { TrieType_Free(t); });

  // Single-char top-level entries so they become direct children of root.
  // Scores scrambled so score-descending order != lex-ascending order.
  struct E { const char *term; float score; };
  E entries[] = {
      {"a", 5},  {"b", 11}, {"c", 2},  {"d", 8},  {"e", 13},
      {"f", 1},  {"g", 7},  {"h", 4},  {"i", 10}, {"j", 6},
      {"k", 12}, {"l", 3},  {"m", 9},
  };
  for (auto &e : entries) {
    ASSERT_TRUE(Trie_InsertStringBuffer(original, e.term, 1, e.score, 1, NULL, 0));
  }
  ASSERT_EQ(13, Trie_Size(original));

  // Ground truth: lex range [d, j) on the original Lex trie.
  ElemSet expected = trieIterRange(original, "d", "j");
  ASSERT_EQ((ElemSet{"d", "e", "f", "g", "h", "i"}), expected);

  RedisModuleIO *io = RMCK_CreateRdbIO();
  std::unique_ptr<RedisModuleIO, std::function<void(RedisModuleIO *)>> ioPtr(
      io, [](RedisModuleIO *p) { RMCK_FreeRdbIO(p); });
  ASSERT_TRUE(io != nullptr);

  // Round-trip via the generic loader (same path as disk-spec / legacy RDB load).
  TrieType_GenericSave(io, original, false, false);
  ASSERT_EQ(0, RMCK_IsIOError(io));
  io->read_pos = 0;
  Trie *loaded = (Trie *)TrieType_GenericLoad(io, false, false, Trie_Sort_Lex);
  std::unique_ptr<Trie, std::function<void(Trie *)>> loadedPtr(
      loaded, [](Trie *t) { TrieType_Free(t); });
  ASSERT_TRUE(loaded != nullptr);
  ASSERT_EQ(Trie_Size(original), Trie_Size(loaded));

  // Same range query on the same data must produce the same result set.
  ElemSet actual = trieIterRange(loaded, "d", "j");
  EXPECT_EQ(expected, actual);
}

TEST_F(TrieTest, testRdbSaveLoadWithNumDocs) {
  // Create a trie with numDocs values
  Trie *originalTrie = NewTrie(NULL, Trie_Sort_Score);
  std::unique_ptr<Trie, std::function<void(Trie *)>> originalTriePtr(originalTrie, [](Trie *trie) {
    TrieType_Free(trie);
  });

  // Insert words with common prefixes and various numDocs values
  trieInsertWithNumDocs(originalTrie, "help", 1.0, 10);     // numDocs = 10
  trieInsertWithNumDocs(originalTrie, "helping", 2.0, 20);  // numDocs = 20
  trieInsertWithNumDocs(originalTrie, "helper", 3.0, 30);   // numDocs = 30
  trieInsertWithNumDocs(originalTrie, "A", 4.0, 100);       // numDocs = 100
  trieInsertWithNumDocs(originalTrie, "AB", 5.0, 200);      // numDocs = 200
  trieInsertWithNumDocs(originalTrie, "ABC", 6.0, 300);     // numDocs = 300

  ASSERT_EQ(6, Trie_Size(originalTrie));

  // Verify original numDocs values
  EXPECT_EQ(10, trieGetNumDocs(originalTrie, "help"));
  EXPECT_EQ(20, trieGetNumDocs(originalTrie, "helping"));
  EXPECT_EQ(30, trieGetNumDocs(originalTrie, "helper"));
  EXPECT_EQ(100, trieGetNumDocs(originalTrie, "A"));
  EXPECT_EQ(200, trieGetNumDocs(originalTrie, "AB"));
  EXPECT_EQ(300, trieGetNumDocs(originalTrie, "ABC"));

  // Create RDB IO context
  RedisModuleIO *io = RMCK_CreateRdbIO();
  std::unique_ptr<RedisModuleIO, std::function<void(RedisModuleIO *)>> ioPtr(io, [](RedisModuleIO *io) {
    RMCK_FreeRdbIO(io);
  });
  ASSERT_TRUE(io != nullptr);

  // Save the trie to RDB
  TrieType_RdbSave(io, originalTrie);
  EXPECT_EQ(0, RMCK_IsIOError(io));

  // Reset read position to load it back
  io->read_pos = 0;

  // Load the trie from RDB
  Trie *loadedTrie = (Trie *)TrieType_RdbLoad(io, TRIE_ENCVER_CURRENT);
  std::unique_ptr<Trie, std::function<void(Trie *)>> loadedTriePtr(loadedTrie, [](Trie *trie) {
    TrieType_Free(trie);
  });
  ASSERT_TRUE(loadedTrie != nullptr);
  EXPECT_EQ(0, RMCK_IsIOError(io));

  // Verify the loaded trie has the same size
  EXPECT_EQ(Trie_Size(originalTrie), Trie_Size(loadedTrie));

  // Verify all entries are present
  EXPECT_TRUE(trieContains(loadedTrie, "help"));
  EXPECT_TRUE(trieContains(loadedTrie, "helping"));
  EXPECT_TRUE(trieContains(loadedTrie, "helper"));
  EXPECT_TRUE(trieContains(loadedTrie, "A"));
  EXPECT_TRUE(trieContains(loadedTrie, "AB"));
  EXPECT_TRUE(trieContains(loadedTrie, "ABC"));

  // Verify numDocs values are preserved after RDB load
  EXPECT_EQ(10, trieGetNumDocs(loadedTrie, "help"));
  EXPECT_EQ(20, trieGetNumDocs(loadedTrie, "helping"));
  EXPECT_EQ(30, trieGetNumDocs(loadedTrie, "helper"));
  EXPECT_EQ(100, trieGetNumDocs(loadedTrie, "A"));
  EXPECT_EQ(200, trieGetNumDocs(loadedTrie, "AB"));
  EXPECT_EQ(300, trieGetNumDocs(loadedTrie, "ABC"));

  // Verify numDocs via iterator as well
  TrieIterator *it = Trie_IterateAll(loadedTrie);
  rune *rstr;
  t_len len;
  float score;
  size_t numDocs;
  RSPayload payload = {.data = NULL, .len = 0};
  int count = 0;

  while (TrieIterator_Next(it, &rstr, &len, &payload, &score, &numDocs, NULL)) {
    count++;
    size_t slen;
    char *s = runesToStr(rstr, len, &slen);
    std::string term(s, slen);
    free(s);

    if (term == "help") {
      EXPECT_EQ(10, numDocs);
    } else if (term == "helping") {
      EXPECT_EQ(20, numDocs);
    } else if (term == "helper") {
      EXPECT_EQ(30, numDocs);
    } else if (term == "A") {
      EXPECT_EQ(100, numDocs);
    } else if (term == "AB") {
      EXPECT_EQ(200, numDocs);
    } else if (term == "ABC") {
      EXPECT_EQ(300, numDocs);
    }
  }
  EXPECT_EQ(6, count);
  TrieIterator_Free(it);
}

// Regression: Trie_CollectFuzzy(trim=1) used to mutate ret->top before reading the
// tail entries it was about to free. Vector_Get then returned 0 without
// touching the out pointer, so TrieSearchResult_Free(h) ran on stack garbage
// and the allocator aborted. This test forces the trim drop branch (one entry
// whose score dominates the rest by > SCORE_TRIM_FACTOR) and asserts both
// non-abort and the surviving count.
TEST_F(TrieTest, testSearchTrimDropsTail) {
  Trie *t = NewTrie(NULL, Trie_Sort_Score);

  trieInsertByScore(t, "ha", 100.0f);
  trieInsertByScore(t, "hb", 1.0f);
  trieInsertByScore(t, "hc", 1.0f);
  trieInsertByScore(t, "hd", 1.0f);
  trieInsertByScore(t, "he", 1.0f);

  // num=10, maxDist=0, mode=TRIE_MATCH_PREFIX, trim=1, optimize=0
  Vector *res = Trie_CollectFuzzy(t, "h", 1, 10, 0, TRIE_MATCH_PREFIX, 1, 0);
  ASSERT_TRUE(res != NULL);

  // Only the dominant entry should survive the trim: 1.0 < 100.0 / 10.0.
  ASSERT_EQ(1, Vector_Size(res));

  TrieSearchResult *e;
  ASSERT_EQ(1, Vector_Get(res, 0, &e));
  ASSERT_EQ(2, e->len);
  ASSERT_EQ(0, memcmp(e->str, "ha", 2));
  TrieSearchResult_Free(e);
  Vector_Free(res);

  TrieType_Free(t);
}