wedb_embed 0.1.0

Embedded Kvrocks-compatible storage engine for WeDb
Documentation
use aok::Void;
use tempfile::tempdir;
use wedb_embed::WeDb;
use wedb_embed::key_composer::KeyComposer;
use wedb_embed::set::SetMeta;

#[ctor::ctor(unsafe)]
fn _log_init() {
    log_init::init();
}

#[test]
fn test_set_metadata_codec() -> Void {
    let meta = SetMeta::new(1700000000000, 101, 55);

    // 标准 26 字节
    let enc = meta.encode();
    assert_eq!(enc.len(), SetMeta::ENCODED_SIZE);
    let dec = SetMeta::decode(&enc).expect("decode failed");
    assert_eq!(dec.base.size, 55);
    assert_eq!(dec.base.version, 101);
    assert_eq!(dec.size(), 55);
    assert_eq!(dec.version(), 101);
    assert_eq!(dec.expire_at(), 1700000000000);
    assert!(!dec.is_empty());

    // Kvrocks 紧凑 25 字节
    let kv_enc = meta.encode_kvrocks();
    assert_eq!(kv_enc.len(), SetMeta::KVROCKS_ENCODED_SIZE);
    let kv_dec = SetMeta::decode(&kv_enc).expect("decode kvrocks failed");
    assert_eq!(kv_dec.base.size, 55);
    assert_eq!(kv_dec.base.version, 101);

    // new_with_version
    let v_meta = SetMeta::new_with_version(0, 10);
    assert_eq!(v_meta.size(), 10);
    assert!(v_meta.version() > 0);

    Ok(())
}

#[test]
fn test_set_basic_ops() -> Void {
    let dir = tempdir()?;
    let db = WeDb::open(dir.path())?;

    // 空操作测试
    assert_eq!(db.sadd("skey", &[] as &[&str])?, 0);
    assert_eq!(db.srem("skey", &[] as &[&str])?, 0);
    assert_eq!(db.scard("non_existent")?, 0);
    assert!(!db.sismember("non_existent", "m1")?);
    assert_eq!(
        db.smismember("non_existent", &["m1", "m2"])?,
        vec![false, false]
    );
    assert!(db.smembers("non_existent")?.is_empty());

    // 基础添加与基数
    assert_eq!(db.sadd("skey", &["m1", "m2", "m3"])?, 3);
    assert_eq!(db.sadd("skey", &["m2", "m4", "m4"])?, 1);
    assert_eq!(db.scard("skey")?, 4);

    assert!(db.sismember("skey", "m1")?);
    assert!(db.sismember("skey", "m2")?);
    assert!(db.sismember("skey", "m3")?);
    assert!(db.sismember("skey", "m4")?);
    assert!(!db.sismember("skey", "m5")?);

    let mism = db.smismember("skey", &["m1", "m5", "m3", "m99"])?;
    assert_eq!(mism, vec![true, false, true, false]);

    assert_eq!(db.srem("skey", &["m1", "m5"])?, 1);
    assert_eq!(db.scard("skey")?, 3);

    let members = db.smembers("skey")?;
    assert_eq!(members.len(), 3);

    // siter 流式遍历
    let mut collected = Vec::new();
    db.siter("skey", |m| {
        collected.push(m.to_vec());
        true
    })?;
    assert_eq!(collected.len(), 3);

    // SMOVE 测试
    assert!(db.smove("skey", "skey2", "m2")?);
    assert_eq!(db.scard("skey")?, 2);
    assert_eq!(db.scard("skey2")?, 1);
    assert!(!db.smove("skey", "skey2", "nonexistent")?);
    // 自身移动
    assert!(db.smove("skey2", "skey2", "m2")?);
    assert_eq!(db.scard("skey2")?, 1);

    // 清空集合并校验元数据删除
    assert_eq!(db.srem("skey", &["m3", "m4"])?, 2);
    assert_eq!(db.scard("skey")?, 0);
    assert_eq!(db.smembers("skey")?.len(), 0);

    Ok(())
}

#[test]
fn test_set_binary_safety() -> Void {
    let dir = tempdir()?;
    let db = WeDb::open(dir.path())?;

    let b1 = b"\x00\x01\x02\xff";
    let b2 = b"prefix:\x00:suffix";
    let b3 = b"\xfe\xed\xfa\xce";

    assert_eq!(
        db.sadd("bin_key", &[b1.as_slice(), b2.as_slice(), b3.as_slice()])?,
        3
    );
    assert_eq!(db.scard("bin_key")?, 3);
    assert!(db.sismember("bin_key", b1)?);
    assert!(db.sismember("bin_key", b2)?);
    assert!(db.sismember("bin_key", b3)?);

    let members = db.smembers("bin_key")?;
    assert_eq!(members.len(), 3);
    assert!(members.contains(&b1.to_vec()));
    assert!(members.contains(&b2.to_vec()));
    assert!(members.contains(&b3.to_vec()));

    assert_eq!(db.srem("bin_key", &[b2.as_slice()])?, 1);
    assert_eq!(db.scard("bin_key")?, 2);
    assert!(!db.sismember("bin_key", b2)?);

    Ok(())
}

#[test]
fn test_set_pop_and_randmember() -> Void {
    let dir = tempdir()?;
    let db = WeDb::open(dir.path())?;

    db.sadd("pop_set", &["a", "b", "c", "d", "e"])?;
    assert_eq!(db.scard("pop_set")?, 5);

    // SPOP 0
    assert!(db.spop("pop_set", 0)?.is_empty());
    assert_eq!(db.scard("pop_set")?, 5);

    // SRANDMEMBER 正数
    let r1 = db.srandmember("pop_set", 2)?;
    assert_eq!(r1.len(), 2);
    assert_ne!(r1[0], r1[1]); // 无重复

    let r_all = db.srandmember("pop_set", 10)?;
    assert_eq!(r_all.len(), 5);

    // SRANDMEMBER 负数(允许重复)
    let r_neg = db.srandmember("pop_set", -10)?;
    assert_eq!(r_neg.len(), 10);

    // SPOP 随机抽取 2 个
    let popped = db.spop("pop_set", 2)?;
    assert_eq!(popped.len(), 2);
    assert_eq!(db.scard("pop_set")?, 3);
    for p in &popped {
        assert!(!db.sismember("pop_set", p)?);
    }

    // SPOP 全部剩余
    let popped_rest = db.spop("pop_set", 10)?;
    assert_eq!(popped_rest.len(), 3);
    assert_eq!(db.scard("pop_set")?, 0);
    assert!(db.smembers("pop_set")?.is_empty());

    // 对空集合 pop
    assert!(db.spop("pop_set", 1)?.is_empty());

    Ok(())
}

#[test]
fn test_set_algebra_ops_and_scan() -> Void {
    let dir = tempdir()?;
    let db = WeDb::open(dir.path())?;

    db.sadd("set_a", &["1", "2", "3", "4"])?;
    db.sadd("set_b", &["3", "4", "5", "6"])?;
    db.sadd("set_c", &["4", "6", "7"])?;

    // SDIFF & SDIFFSTORE & SDIFFCARD
    let diff = db.sdiff(&["set_a", "set_b"])?;
    assert_eq!(diff.len(), 2);
    let diff_self = db.sdiff(&["set_a", "set_a"])?;
    assert!(diff_self.is_empty());
    assert_eq!(db.sdiffcard(&["set_a", "set_b"], 0)?, 2);
    assert_eq!(db.sdiffcard(&["set_a", "set_b"], 1)?, 1);

    assert_eq!(db.sdiffstore("diff_dst", &["set_a", "set_b"])?, 2);
    assert_eq!(db.scard("diff_dst")?, 2);

    // SUNION & SUNIONSTORE & SUNIONCARD
    let union_res = db.sunion(&["set_a", "set_b", "set_c"])?;
    assert_eq!(union_res.len(), 7);
    assert_eq!(db.sunioncard(&["set_a", "set_b", "set_c"], 0)?, 7);
    assert_eq!(db.sunioncard(&["set_a", "set_b", "set_c"], 4)?, 4);

    assert_eq!(
        db.sunionstore("union_dst", &["set_a", "set_b", "set_c"])?,
        7
    );
    assert_eq!(db.scard("union_dst")?, 7);

    // SINTER & SINTERSTORE & SINTERCARD
    let inter_res = db.sinter(&["set_a", "set_b", "set_c"])?;
    assert_eq!(inter_res, vec![b"4".to_vec()]);
    assert_eq!(
        db.sinterstore("inter_dst", &["set_a", "set_b", "set_c"])?,
        1
    );
    assert_eq!(db.sintercard(&["set_a", "set_b", "set_c"], 0)?, 1);
    assert_eq!(db.sintercard(&["set_a", "set_b", "set_c"], 5)?, 1);

    // SINTER 短路测试(与不存在的空 key 交集必为空)
    assert!(db.sinter(&["set_a", "empty_key"])?.is_empty());
    assert_eq!(db.sintercard(&["set_a", "empty_key"], 0)?, 0);

    // SSCAN
    let (cur, page) = db.sscan("union_dst", 0, None, Some(4))?;
    assert_eq!(page.len(), 4);
    assert_eq!(cur, 4);

    let (cur2, page2) = db.sscan("union_dst", cur, None, Some(10))?;
    assert_eq!(page2.len(), 3);
    assert_eq!(cur2, 0);

    // SSCAN Pattern 过滤
    let (_, pattern_res) = db.sscan("union_dst", 0, Some(b"[1-3]"), Some(10))?;
    assert_eq!(pattern_res.len(), 3);

    Ok(())
}

#[test]
fn test_set_expiration_behavior() -> Void {
    let dir = tempdir()?;
    let db = WeDb::open(dir.path())?;

    // 手动写入一个已过期的元数据
    let kc = KeyComposer::new("default");
    let meta_k = kc.set_meta("expired_set");
    let expired_meta = SetMeta::new(1000, 1, 5); // expire_at = 1000ms
    db.meta_ks
        .insert(meta_k.as_bytes(), expired_meta.encode())?;

    // 所有只读与删除操作对过期 key 应视为不存在
    assert_eq!(db.scard("expired_set")?, 0);
    assert!(!db.sismember("expired_set", "item")?);
    assert_eq!(
        db.smismember("expired_set", &["i1", "i2"])?,
        vec![false, false]
    );
    assert!(db.smembers("expired_set")?.is_empty());
    assert_eq!(db.srem("expired_set", &["item"])?, 0);
    assert!(db.spop("expired_set", 1)?.is_empty());
    assert!(db.srandmember("expired_set", 1)?.is_empty());
    let (cur, scan_items) = db.sscan("expired_set", 0, None, None)?;
    assert_eq!(cur, 0);
    assert!(scan_items.is_empty());

    // SADD 会重新创建并覆盖过期 key,并清理旧残留子键
    assert_eq!(db.sadd("expired_set", &["new_item"])?, 1);
    assert_eq!(db.scard("expired_set")?, 1);
    assert!(db.sismember("expired_set", "new_item")?);
    assert!(!db.sismember("expired_set", "item")?); // 确保旧键被彻底清理
    let sm = db.smembers("expired_set")?;
    assert_eq!(sm, vec![b"new_item".to_vec()]);

    Ok(())
}

#[test]
fn test_set_ttl_and_persist_helpers() -> Void {
    let dir = tempdir()?;
    let db = WeDb::open(dir.path())?;

    db.sadd("ttl_set", &["a", "b", "c"])?;
    assert_eq!(db.scard("ttl_set")?, 3);

    // 未设置过期时间时 TTL 为 -1
    assert_eq!(db.sttl("ttl_set")?, -1);

    // 设置绝对过期时间
    let future_ts = (ts_::sec() + 300) * 1000;
    assert!(db.sexpireat("ttl_set", future_ts)?);
    let ttl_val = db.sttl("ttl_set")?;
    assert!(ttl_val > 0 && ttl_val <= 300_000);

    // SPERSIST 移除过期时间
    assert!(db.spersist("ttl_set")?);
    assert_eq!(db.sttl("ttl_set")?, -1);

    // 不存在的 key
    assert!(!db.sexpireat("non_existent", future_ts)?);
    assert_eq!(db.sttl("non_existent")?, -2);
    assert!(!db.spersist("non_existent")?);

    Ok(())
}

#[test]
fn test_set_sinter_and_sdiff_adaptive_point_lookups() -> Void {
    let dir = tempdir()?;
    let db = WeDb::open(dir.path())?;

    // 创建一个小集合和一个大集合以触发自适应点查分支 (len * 4 < next_card)
    let mut small_items = Vec::new();
    for i in 0..5 {
        small_items.push(format!("item_{i}"));
    }
    let small_slices: Vec<&str> = small_items.iter().map(|s| s.as_str()).collect();
    db.sadd("small_set", &small_slices)?;

    let mut large_items = Vec::new();
    for i in 0..100 {
        large_items.push(format!("item_{i}"));
    }
    let large_slices: Vec<&str> = large_items.iter().map(|s| s.as_str()).collect();
    db.sadd("large_set", &large_slices)?;

    // SINTER:交集应为 small_set 的全部 5 个元素
    let inter = db.sinter(&["small_set", "large_set"])?;
    assert_eq!(inter.len(), 5);

    // SINTERCARD
    assert_eq!(db.sintercard(&["small_set", "large_set"], 0)?, 5);
    assert_eq!(db.sintercard(&["small_set", "large_set"], 3)?, 3);
    assert_eq!(db.sintercard(&["small_set", "large_set"], 10)?, 5);

    // 空集交集短路
    assert_eq!(db.sintercard(&["small_set", "non_existent"], 0)?, 0);
    assert!(db.sinter(&["small_set", "non_existent"])?.is_empty());

    // SDIFF 自适应点查:large_set - small_set (large_set 大,small_set 小 -> standard scan)
    let diff1 = db.sdiff(&["large_set", "small_set"])?;
    assert_eq!(diff1.len(), 95);

    // SDIFF 自适应点查:small_set - large_set (small_set 5 个,large_set 100 个 -> 点查)
    let diff2 = db.sdiff(&["small_set", "large_set"])?;
    assert_eq!(diff2.len(), 0);

    // SDIFFCARD
    assert_eq!(db.sdiffcard(&["large_set", "small_set"], 0)?, 95);
    assert_eq!(db.sdiffcard(&["large_set", "small_set"], 10)?, 10);

    Ok(())
}

#[test]
fn test_set_overwrite_and_smove_scenarios() -> Void {
    let dir = tempdir()?;
    let db = WeDb::open(dir.path())?;

    // 1. overwrite_set
    assert_eq!(db.overwrite_set("ow_key", &["x", "y", "z", "z"])?, 3);
    assert_eq!(db.scard("ow_key")?, 3);
    assert!(db.sismember("ow_key", "x")?);
    assert!(db.sismember("ow_key", "y")?);
    assert!(db.sismember("ow_key", "z")?);

    // 再次覆盖为空
    assert_eq!(db.overwrite_set("ow_key", &[] as &[&str])?, 0);
    assert_eq!(db.scard("ow_key")?, 0);
    assert!(db.smembers("ow_key")?.is_empty());

    // 2. SMOVE 各种场景
    db.sadd("src_set", &["alpha", "beta"])?;
    db.sadd("dst_set", &["gamma"])?;

    // 移动存在的元素
    assert!(db.smove("src_set", "dst_set", "alpha")?);
    assert_eq!(db.scard("src_set")?, 1);
    assert_eq!(db.scard("dst_set")?, 2);
    assert!(!db.sismember("src_set", "alpha")?);
    assert!(db.sismember("dst_set", "alpha")?);

    // 移动到已经包含该元素的目标集合
    assert!(db.smove("src_set", "dst_set", "beta")?);
    assert_eq!(db.scard("src_set")?, 0); // src 变为空并清除元数据
    assert_eq!(db.scard("dst_set")?, 3);

    // 对已空的源集合进行 move
    assert!(!db.smove("src_set", "dst_set", "beta")?);

    // 移动不存在的元素
    assert!(!db.smove("dst_set", "new_dst", "not_exist")?);

    // 移动到全新集合
    assert!(db.smove("dst_set", "fresh_dst", "gamma")?);
    assert_eq!(db.scard("fresh_dst")?, 1);
    assert_eq!(db.scard("dst_set")?, 2);

    Ok(())
}