armdb 0.2.0

sharded bitcask key-value storage optimized for NVMe
Documentation
use armdb::{Config, ConstTree, DbError};
use tempfile::tempdir;

#[cfg(feature = "var-collections")]
use armdb::VarTree;

// ---------------------------------------------------------------------------
// db.meta format mismatch tests
// ---------------------------------------------------------------------------

#[test]
fn test_shard_count_mismatch() {
    let dir = tempdir().unwrap();

    // Create database with 3 shards
    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), Config::test()).unwrap();
        tree.put(&1u64.to_be_bytes(), &100u64.to_be_bytes())
            .unwrap();
        tree.close().unwrap();
    }

    // Reopen with different shard_count → FormatMismatch
    {
        let mut config = Config::test();
        config.shard_count = 5; // was 3
        let err = unwrap_err(ConstTree::<[u8; 8], 8>::open(dir.path(), config));
        assert!(
            matches!(err, DbError::FormatMismatch(_)),
            "expected FormatMismatch, got: {err}"
        );
        let msg = err.to_string();
        assert!(
            msg.contains("shard_count"),
            "error should mention shard_count: {msg}"
        );
    }
}

#[test]
fn test_shard_prefix_bits_mismatch() {
    let dir = tempdir().unwrap();

    // Create database with shard_prefix_bits = 0
    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), Config::test()).unwrap();
        tree.put(&1u64.to_be_bytes(), &100u64.to_be_bytes())
            .unwrap();
        tree.close().unwrap();
    }

    // Reopen with different shard_prefix_bits → FormatMismatch
    {
        let mut config = Config::test();
        config.shard_prefix_bits = 32; // was 0
        let err = unwrap_err(ConstTree::<[u8; 8], 8>::open(dir.path(), config));
        assert!(
            matches!(err, DbError::FormatMismatch(_)),
            "expected FormatMismatch, got: {err}"
        );
        let msg = err.to_string();
        assert!(
            msg.contains("shard_prefix_bits"),
            "error should mention shard_prefix_bits: {msg}"
        );
    }
}

/// F-3: non-encryption build must reject a db.meta with the encrypted bit set.
#[cfg(not(feature = "encryption"))]
#[test]
fn test_rejects_encrypted_meta_on_non_encryption_build() {
    let dir = tempdir().unwrap();
    std::fs::create_dir_all(dir.path()).unwrap();

    std::fs::write(dir.path().join("db.meta"), [2u8, 0, 1, 0]).unwrap();

    std::fs::create_dir_all(dir.path().join("shard_000")).unwrap();
    std::fs::create_dir_all(dir.path().join("shard_001")).unwrap();

    let err = unwrap_err(ConstTree::<[u8; 8], 8>::open(dir.path(), Config::test()));
    assert!(
        matches!(err, DbError::FormatMismatch(_)),
        "expected FormatMismatch for encrypted db on non-encryption build, got: {err}"
    );
    let msg = err.to_string();
    assert!(
        msg.contains("encrypted"),
        "error should mention encryption: {msg}"
    );
}

#[test]
fn test_db_meta_invalid_size() {
    let dir = tempdir().unwrap();

    // Create db directory with a malformed db.meta
    std::fs::create_dir_all(dir.path()).unwrap();
    std::fs::write(dir.path().join("db.meta"), [0u8; 2]).unwrap(); // wrong size

    let err = unwrap_err(ConstTree::<[u8; 8], 8>::open(dir.path(), Config::test()));
    assert!(
        matches!(err, DbError::FormatMismatch(_)),
        "expected FormatMismatch, got: {err}"
    );
}

#[test]
fn test_same_config_reopens_fine() {
    let dir = tempdir().unwrap();
    let config = Config::test();

    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), config.clone()).unwrap();
        tree.put(&1u64.to_be_bytes(), &42u64.to_be_bytes()).unwrap();
        tree.close().unwrap();
    }

    // Same config → should succeed
    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), config).unwrap();
        assert_eq!(tree.get(&1u64.to_be_bytes()), Some(42u64.to_be_bytes()));
    }
}

/// F-5: opening an existing DB with a missing shard directory must fail.
#[test]
fn test_rejects_missing_shard_dir_on_existing_db() {
    let dir = tempdir().unwrap();
    let config = Config::test(); // shard_count = 2

    // Create and close a valid database.
    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), config.clone()).unwrap();
        tree.put(&1u64.to_be_bytes(), &42u64.to_be_bytes()).unwrap();
        tree.close().unwrap();
    }

    // Delete one shard directory.
    std::fs::remove_dir_all(dir.path().join("shard_001")).unwrap();

    // Reopen must fail — not silently recreate the empty shard.
    let err = unwrap_err(ConstTree::<[u8; 8], 8>::open(dir.path(), config));
    assert!(
        matches!(err, DbError::FormatMismatch(_)),
        "expected FormatMismatch for missing shard dir, got: {err}"
    );
    let msg = err.to_string();
    assert!(
        msg.contains("shard") && msg.contains("missing"),
        "error should mention missing shard: {msg}"
    );
}

/// F-8B case 1: db.meta deleted but all shard dirs remain → must reject.
#[test]
fn test_rejects_orphaned_shard_dirs_meta_deleted() {
    let dir = tempdir().unwrap();
    let config = Config::test();

    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), config.clone()).unwrap();
        tree.put(&1u64.to_be_bytes(), &42u64.to_be_bytes()).unwrap();
        tree.close().unwrap();
    }

    // Delete only db.meta — shard dirs still present.
    std::fs::remove_file(dir.path().join("db.meta")).unwrap();

    let err = unwrap_err(ConstTree::<[u8; 8], 8>::open(dir.path(), config));
    assert!(
        matches!(err, DbError::FormatMismatch(_)),
        "expected FormatMismatch for orphaned shard dirs, got: {err}"
    );
    let msg = err.to_string();
    assert!(
        msg.contains("db.meta") && msg.contains("shard"),
        "error should mention db.meta and shard dirs: {msg}"
    );
}

/// F-8B case 2: db.meta AND shard_000 deleted, but shard_001 remains → must reject.
#[test]
fn test_rejects_partial_orphaned_shard_dirs() {
    let dir = tempdir().unwrap();
    let config = Config::test();

    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), config.clone()).unwrap();
        tree.put(&1u64.to_be_bytes(), &42u64.to_be_bytes()).unwrap();
        tree.close().unwrap();
    }

    // Delete db.meta and shard_000, leave shard_001.
    std::fs::remove_file(dir.path().join("db.meta")).unwrap();
    std::fs::remove_dir_all(dir.path().join("shard_000")).unwrap();

    let err = unwrap_err(ConstTree::<[u8; 8], 8>::open(dir.path(), config));
    assert!(
        matches!(err, DbError::FormatMismatch(_)),
        "expected FormatMismatch for partial orphaned shard dirs, got: {err}"
    );
}

// ---------------------------------------------------------------------------
// hints=false recovery tests
// ---------------------------------------------------------------------------

#[test]
fn test_const_recovery_no_hints() {
    let dir = tempdir().unwrap();
    let mut config = Config::test();
    config.hints = false;

    // Write entries with hints disabled
    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), config.clone()).unwrap();

        for i in 0u64..100 {
            tree.put(&i.to_be_bytes(), &(i * 10).to_be_bytes()).unwrap();
        }

        tree.close().unwrap();
    }

    // Verify no hint files were created
    let has_hints = walkdir(dir.path(), ".hint");
    assert!(!has_hints, "no .hint files should exist when hints=false");

    // Reopen and verify all entries recovered via full scan
    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), config).unwrap();
        assert_eq!(tree.len(), 100);

        for i in 0u64..100 {
            let actual = tree
                .get(&i.to_be_bytes())
                .unwrap_or_else(|| panic!("key {i} not found"));
            assert_eq!(actual, (i * 10).to_be_bytes());
        }
    }
}

#[test]
fn test_const_recovery_no_hints_with_deletes() {
    let dir = tempdir().unwrap();
    let mut config = Config::test();
    config.hints = false;

    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), config.clone()).unwrap();

        for i in 0u64..100 {
            tree.put(&i.to_be_bytes(), &(i * 10).to_be_bytes()).unwrap();
        }
        for i in 0u64..50 {
            tree.delete(&i.to_be_bytes()).unwrap();
        }

        tree.close().unwrap();
    }

    {
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), config).unwrap();
        assert_eq!(tree.len(), 50);

        for i in 0u64..50 {
            assert!(tree.get(&i.to_be_bytes()).is_none());
        }
        for i in 50u64..100 {
            let actual = tree
                .get(&i.to_be_bytes())
                .unwrap_or_else(|| panic!("key {i} not found"));
            assert_eq!(actual, (i * 10).to_be_bytes());
        }
    }
}

#[cfg(feature = "var-collections")]
#[test]
fn test_var_recovery_no_hints() {
    let dir = tempdir().unwrap();
    let mut config = Config::test();
    config.hints = false;

    {
        let tree = VarTree::<[u8; 8]>::open(dir.path(), config.clone()).unwrap();

        for i in 0u64..100 {
            let value = format!("val_{i}").into_bytes();
            tree.put(&i.to_be_bytes(), &value).unwrap();
        }

        tree.close().unwrap();
    }

    let has_hints = walkdir(dir.path(), ".hint");
    assert!(!has_hints, "no .hint files should exist when hints=false");

    {
        let tree = VarTree::<[u8; 8]>::open(dir.path(), config).unwrap();
        assert_eq!(tree.len(), 100);

        for i in 0u64..100 {
            let expected = format!("val_{i}").into_bytes();
            let actual = tree
                .get(&i.to_be_bytes())
                .unwrap_or_else(|| panic!("key {i} not found"));
            assert_eq!(actual.as_bytes(), expected.as_slice());
        }
    }
}

/// Switch from hints=true to hints=false between restarts.
/// Old hint files should be ignored, recovery via full scan.
#[test]
fn test_switch_hints_true_to_false() {
    let dir = tempdir().unwrap();

    // Open with hints=true, write, close (generates hint files)
    {
        let mut config = Config::test();
        config.hints = true;
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), config).unwrap();

        for i in 0u64..50 {
            tree.put(&i.to_be_bytes(), &(i * 10).to_be_bytes()).unwrap();
        }

        tree.close().unwrap();
    }

    let had_hints = walkdir(dir.path(), ".hint");
    assert!(had_hints, "hint files should exist after hints=true close");

    // Reopen with hints=false — should recover fine via full scan
    {
        let mut config = Config::test();
        config.hints = false;
        let tree = ConstTree::<[u8; 8], 8>::open(dir.path(), config).unwrap();
        assert_eq!(tree.len(), 50);

        for i in 0u64..50 {
            let actual = tree
                .get(&i.to_be_bytes())
                .unwrap_or_else(|| panic!("key {i} not found"));
            assert_eq!(actual, (i * 10).to_be_bytes());
        }
    }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Extract error from Result, panicking if Ok.
fn unwrap_err<T>(result: Result<T, DbError>) -> DbError {
    match result {
        Err(e) => e,
        Ok(_) => panic!("expected error, got Ok"),
    }
}

/// Check if any file with the given extension exists recursively under `dir`.
fn walkdir(dir: &std::path::Path, ext: &str) -> bool {
    fn walk(dir: &std::path::Path, ext: &str) -> bool {
        let Ok(entries) = std::fs::read_dir(dir) else {
            return false;
        };
        for entry in entries.flatten() {
            let path = entry.path();
            if path.is_dir() {
                if walk(&path, ext) {
                    return true;
                }
            } else if path.to_string_lossy().ends_with(ext) {
                return true;
            }
        }
        false
    }
    walk(dir, ext)
}