jmt-blake3 0.1.2

// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use rand::{rngs::OsRng, Rng};

use crate::{
    mock::MockTreeStore,
    node_type::{Node, NodeKey},
    tree_cache::TreeCache,
    types::{nibble::nibble_path::NibblePath, Version, PRE_GENESIS_VERSION},
    KeyHash,
};

fn random_leaf_with_key(next_version: Version) -> (Node, NodeKey) {
    let key: [u8; 32] = OsRng.gen();
    let value: [u8; 32] = OsRng.gen();
    let key_hash: KeyHash = key.as_ref().into();
    let node = Node::new_leaf(key_hash, value.to_vec());
    let node_key = NodeKey::new(next_version, NibblePath::new(key_hash.0.to_vec()));
    (node, node_key)
}

#[tokio::test]
async fn test_get_node() {
    let next_version = 0;
    let db = MockTreeStore::default();
    let cache = TreeCache::new(&db, next_version).await.unwrap();

    let (node, node_key) = random_leaf_with_key(next_version);
    db.put_node(node_key.clone(), node.clone()).await.unwrap();

    assert_eq!(cache.get_node(&node_key).await.unwrap(), node);
}

#[tokio::test]
async fn test_root_node() {
    let next_version = 0;
    let db = MockTreeStore::default();
    let mut cache = TreeCache::new(&db, next_version).await.unwrap();
    assert_eq!(*cache.get_root_node_key(), NodeKey::new_empty_path(0));

    let (node, node_key) = random_leaf_with_key(next_version);
    db.put_node(node_key.clone(), node).await.unwrap();
    cache.set_root_node_key(node_key.clone());

    assert_eq!(*cache.get_root_node_key(), node_key);
}

#[tokio::test]
async fn test_pre_genesis() {
    let next_version = 0;
    let db = MockTreeStore::default();
    let pre_genesis_root_key = NodeKey::new_empty_path(PRE_GENESIS_VERSION);
    let (pre_genesis_only_node, _) = random_leaf_with_key(PRE_GENESIS_VERSION);
    db.put_node(pre_genesis_root_key.clone(), pre_genesis_only_node)
        .await
        .unwrap();

    let cache = TreeCache::new(&db, next_version).await.unwrap();
    assert_eq!(*cache.get_root_node_key(), pre_genesis_root_key);
}

#[tokio::test]
async fn test_freeze_with_delete() {
    let next_version = 0;
    let db = MockTreeStore::default();
    let mut cache = TreeCache::new(&db, next_version).await.unwrap();

    assert_eq!(*cache.get_root_node_key(), NodeKey::new_empty_path(0));

    let (node1, node1_key) = random_leaf_with_key(next_version);
    cache.put_node(node1_key.clone(), node1.clone()).unwrap();
    let (node2, node2_key) = random_leaf_with_key(next_version);
    cache.put_node(node2_key.clone(), node2.clone()).unwrap();
    assert_eq!(cache.get_node(&node1_key).await.unwrap(), node1);
    assert_eq!(cache.get_node(&node2_key).await.unwrap(), node2);
    cache.freeze().await;
    assert_eq!(cache.get_node(&node1_key).await.unwrap(), node1);
    assert_eq!(cache.get_node(&node2_key).await.unwrap(), node2);

    cache.delete_node(&node1_key, true /* is_leaf */);
    cache.freeze().await;
    let (_, update_batch) = cache.into();
    assert_eq!(update_batch.node_batch.len(), 3);
    assert_eq!(update_batch.stale_node_index_batch.len(), 1);
}