sqry_core/graph/unified/storage/

metadata.rs

//! Sparse node metadata store for macro boundary analysis and classpath provenance.
//!
//! This module provides [`NodeMetadataStore`], a sparse metadata store keyed by
//! full [`NodeId`] (index + generation) to prevent stale metadata when the
//! generational arena reuses a slot index with a new generation.
//!
//! Only nodes with metadata get entries, keeping memory overhead
//! proportional to the number of annotated symbols rather than total node count.

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

use super::super::node::id::NodeId;

/// Optional metadata for nodes that participate in macro boundary analysis.
///
/// Stored separately from [`NodeEntry`] to avoid bloating the arena for the
/// majority of nodes that don't need macro metadata.
///
/// Each field is `Option` (or `Vec` with empty default) so only relevant
/// metadata consumes space in the serialized representation.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct MacroNodeMetadata {
    /// Whether this symbol was generated by macro expansion.
    pub macro_generated: Option<bool>,

    /// Qualified name of the macro that generated this symbol.
    pub macro_source: Option<String>,

    /// The cfg predicate string (e.g., `"test"`, `"feature = \"serde\""`)
    pub cfg_condition: Option<String>,

    /// Whether this cfg is active (`None` = unknown, requires external config).
    pub cfg_active: Option<bool>,

    /// Proc-macro kind for proc-macro function nodes.
    pub proc_macro_kind: Option<ProcMacroFunctionKind>,

    /// Whether expansion data came from cache vs live `cargo expand`.
    pub expansion_cached: Option<bool>,

    /// Unresolved attribute paths that could not be positively identified
    /// as proc-macro attributes. Stored for potential future resolution.
    pub unresolved_attributes: Vec<String>,
}

/// Classification of proc-macro function types.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ProcMacroFunctionKind {
    /// `#[proc_macro_derive(Name)]` — generates impls for structs/enums.
    Derive,
    /// `#[proc_macro_attribute]` — transforms annotated items.
    Attribute,
    /// `#[proc_macro]` — function-like `my_macro!(...)` invocation.
    FunctionLike,
}

/// Metadata for nodes originating from JVM classpath bytecode.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ClasspathNodeMetadata {
    /// Maven coordinates (e.g., `"com.google.guava:guava:33.0.0"`).
    pub coordinates: Option<String>,
    /// JAR file path this node was extracted from.
    pub jar_path: String,
    /// Fully qualified name in JVM format (e.g., `"java.util.HashMap"`).
    pub fqn: String,
    /// Whether this is a direct or transitive dependency.
    pub is_direct_dependency: bool,
}

/// Metadata that can be attached to graph nodes.
///
/// This enum is not directly `Serialize`/`Deserialize` because `postcard`
/// does not support Rust enum variants wrapping structs. Instead,
/// `NodeMetadataStore` handles serialization/deserialization through flat
/// entry structs with explicit discriminant bytes.
///
/// For JSON use (e.g., MCP export), use the `serde_json`-compatible
/// `to_json`/`from_json` convenience methods.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum NodeMetadata {
    /// Rust macro-related metadata.
    Macro(MacroNodeMetadata),
    /// JVM classpath provenance metadata.
    Classpath(ClasspathNodeMetadata),
}

/// Discriminant values for `NodeMetadata` wire format.
const NODE_METADATA_MACRO: u8 = 0;
const NODE_METADATA_CLASSPATH: u8 = 1;

/// Sparse metadata store keyed by full `NodeId` (index + generation).
///
/// Uses `(u32, u64)` tuple key to prevent stale metadata when the
/// generational arena reuses a slot index with a new generation.
/// A lookup with `NodeId { index: 5, generation: 3 }` will NOT match metadata
/// stored for `NodeId { index: 5, generation: 2 }`.
///
/// # Memory characteristics
///
/// For a typical large codebase (100K nodes), only ~5-10% of nodes have
/// metadata. A store with 10K entries at ~200 bytes each = ~2MB, which is
/// acceptable given snapshots are already 10-50MB.
///
/// # Serialization
///
/// The in-memory representation uses `HashMap` for O(1) lookups. For postcard
/// serialization (which doesn't support tuple keys natively), we serialize as
/// a `Vec` of `NodeMetadataEntry` structs with explicit `index` and `generation`
/// fields, then reconstruct the `HashMap` on deserialization.
///
/// For backwards compatibility, legacy entries serialized as bare
/// `MacroNodeMetadata` (without a wrapping `NodeMetadata` enum tag) are
/// transparently deserialized as `NodeMetadata::Macro`.
#[derive(Debug, Clone, Default)]
pub struct NodeMetadataStore {
    /// Metadata entries keyed by `(NodeId::index(), NodeId::generation())`.
    entries: HashMap<(u32, u64), NodeMetadata>,
}

/// Serialization wrapper for a V7 metadata entry with explicit discriminant.
#[derive(Debug, Clone, Serialize, Deserialize)]
struct NodeMetadataEntryV7 {
    index: u32,
    generation: u64,
    /// Discriminant: 0 = Macro, 1 = Classpath
    kind: u8,
    /// Macro metadata (present when kind == 0)
    macro_data: Option<MacroNodeMetadata>,
    /// Classpath metadata (present when kind == 1)
    classpath_data: Option<ClasspathNodeMetadata>,
}

impl Serialize for NodeMetadataStore {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        let entries: Vec<NodeMetadataEntryV7> = self
            .entries
            .iter()
            .map(|(&(index, generation), metadata)| match metadata {
                NodeMetadata::Macro(m) => NodeMetadataEntryV7 {
                    index,
                    generation,
                    kind: NODE_METADATA_MACRO,
                    macro_data: Some(m.clone()),
                    classpath_data: None,
                },
                NodeMetadata::Classpath(c) => NodeMetadataEntryV7 {
                    index,
                    generation,
                    kind: NODE_METADATA_CLASSPATH,
                    macro_data: None,
                    classpath_data: Some(c.clone()),
                },
            })
            .collect();
        entries.serialize(serializer)
    }
}

impl<'de> Deserialize<'de> for NodeMetadataStore {
    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        // V7 format: each entry has an explicit `kind` discriminant.
        // Legacy V6 entries had no `kind` field, so `kind` will default to 0
        // (which maps to Macro) via postcard's sequential field decoding.
        // However, since we bumped the snapshot version to V7, V6 snapshots
        // will be rejected at the magic byte check before reaching this code.
        let entries: Vec<NodeMetadataEntryV7> = Vec::deserialize(deserializer)?;
        let mut map = HashMap::with_capacity(entries.len());
        for e in entries {
            let metadata = match e.kind {
                NODE_METADATA_CLASSPATH => {
                    let data = e.classpath_data.ok_or_else(|| {
                        serde::de::Error::custom(
                            "missing classpath_data for Classpath metadata entry",
                        )
                    })?;
                    NodeMetadata::Classpath(data)
                }
                _ => {
                    // Default: treat as Macro (covers both explicit 0 and legacy)
                    let data = e.macro_data.unwrap_or_default();
                    NodeMetadata::Macro(data)
                }
            };
            map.insert((e.index, e.generation), metadata);
        }
        Ok(Self { entries: map })
    }
}

impl NodeMetadataStore {
    /// Create a new empty metadata store.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Get metadata for a node, if any exists.
    ///
    /// Returns `None` if no metadata is stored for this node, or if the
    /// generation doesn't match (indicating a stale reference).
    #[must_use]
    pub fn get(&self, node_id: NodeId) -> Option<&MacroNodeMetadata> {
        match self.entries.get(&(node_id.index(), node_id.generation()))? {
            NodeMetadata::Macro(m) => Some(m),
            NodeMetadata::Classpath(_) => None,
        }
    }

    /// Get the full `NodeMetadata` envelope for a node, if any exists.
    ///
    /// Returns `None` if no metadata is stored for this node, or if the
    /// generation doesn't match (indicating a stale reference).
    #[must_use]
    pub fn get_metadata(&self, node_id: NodeId) -> Option<&NodeMetadata> {
        self.entries.get(&(node_id.index(), node_id.generation()))
    }

    /// Get mutable metadata for a node, if any exists.
    #[must_use]
    pub fn get_mut(&mut self, node_id: NodeId) -> Option<&mut MacroNodeMetadata> {
        match self
            .entries
            .get_mut(&(node_id.index(), node_id.generation()))?
        {
            NodeMetadata::Macro(m) => Some(m),
            NodeMetadata::Classpath(_) => None,
        }
    }

    /// Insert macro metadata for a node, replacing any existing entry.
    ///
    /// Convenience method that wraps the metadata in `NodeMetadata::Macro`.
    pub fn insert(&mut self, node_id: NodeId, metadata: MacroNodeMetadata) {
        self.entries.insert(
            (node_id.index(), node_id.generation()),
            NodeMetadata::Macro(metadata),
        );
    }

    /// Insert typed metadata for a node, replacing any existing entry.
    pub fn insert_metadata(&mut self, node_id: NodeId, metadata: NodeMetadata) {
        self.entries
            .insert((node_id.index(), node_id.generation()), metadata);
    }

    /// Get or insert default macro metadata for a node.
    pub fn get_or_insert_default(&mut self, node_id: NodeId) -> &mut MacroNodeMetadata {
        let entry = self
            .entries
            .entry((node_id.index(), node_id.generation()))
            .or_insert_with(|| NodeMetadata::Macro(MacroNodeMetadata::default()));
        match entry {
            NodeMetadata::Macro(m) => m,
            // If a non-macro entry already exists at this key, this is a
            // programming error — callers should not mix metadata types for the
            // same node. Panic in debug, return a default in release.
            NodeMetadata::Classpath(_) => {
                panic!("get_or_insert_default called on a Classpath metadata entry")
            }
        }
    }

    /// Remove metadata for a node.
    pub fn remove(&mut self, node_id: NodeId) -> Option<MacroNodeMetadata> {
        match self
            .entries
            .remove(&(node_id.index(), node_id.generation()))?
        {
            NodeMetadata::Macro(m) => Some(m),
            NodeMetadata::Classpath(_) => None,
        }
    }

    /// Remove typed metadata for a node.
    pub fn remove_metadata(&mut self, node_id: NodeId) -> Option<NodeMetadata> {
        self.entries
            .remove(&(node_id.index(), node_id.generation()))
    }

    /// Returns the number of nodes with metadata.
    #[must_use]
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    /// Returns true if no nodes have metadata.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// Iterate over all metadata entries.
    pub fn iter(&self) -> impl Iterator<Item = ((u32, u64), &MacroNodeMetadata)> {
        self.entries.iter().filter_map(|(&k, v)| match v {
            NodeMetadata::Macro(m) => Some((k, m)),
            NodeMetadata::Classpath(_) => None,
        })
    }

    /// Iterate over all metadata entries as typed `NodeMetadata`.
    pub fn iter_all(&self) -> impl Iterator<Item = ((u32, u64), &NodeMetadata)> {
        self.entries.iter().map(|(&k, v)| (k, v))
    }

    /// Merge another metadata store into this one.
    ///
    /// Entries from `other` overwrite existing entries with the same key.
    pub fn merge(&mut self, other: &NodeMetadataStore) {
        for (&key, value) in &other.entries {
            self.entries.insert(key, value.clone());
        }
    }
}

impl PartialEq for NodeMetadataStore {
    fn eq(&self, other: &Self) -> bool {
        self.entries == other.entries
    }
}

impl Eq for NodeMetadataStore {}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_metadata_store_basic_operations() {
        let mut store = NodeMetadataStore::new();
        assert!(store.is_empty());
        assert_eq!(store.len(), 0);

        let node = NodeId::new(5, 1);
        let metadata = MacroNodeMetadata {
            macro_generated: Some(true),
            macro_source: Some("derive_Debug".to_string()),
            ..Default::default()
        };

        store.insert(node, metadata.clone());
        assert_eq!(store.len(), 1);
        assert!(!store.is_empty());

        let retrieved = store.get(node).unwrap();
        assert_eq!(retrieved.macro_generated, Some(true));
        assert_eq!(retrieved.macro_source.as_deref(), Some("derive_Debug"));
    }

    #[test]
    fn test_metadata_full_nodeid_key() {
        let mut store = NodeMetadataStore::new();

        let node_gen1 = NodeId::new(5, 1);
        let node_gen2 = NodeId::new(5, 2);

        store.insert(
            node_gen1,
            MacroNodeMetadata {
                macro_generated: Some(true),
                ..Default::default()
            },
        );

        // Same index, different generation → should NOT match
        assert!(store.get(node_gen2).is_none());

        // Same index, same generation → should match
        assert!(store.get(node_gen1).is_some());
    }

    #[test]
    fn test_metadata_slot_reuse_no_stale_data() {
        let mut store = NodeMetadataStore::new();

        // Simulate: node at index 5 gen 1 has metadata
        let old_node = NodeId::new(5, 1);
        store.insert(
            old_node,
            MacroNodeMetadata {
                cfg_condition: Some("test".to_string()),
                ..Default::default()
            },
        );

        // Simulate: slot 5 is reused with generation 2 (new node)
        let new_node = NodeId::new(5, 2);

        // New node should NOT see old metadata
        assert!(store.get(new_node).is_none());

        // Old node still accessible
        assert_eq!(
            store.get(old_node).unwrap().cfg_condition.as_deref(),
            Some("test")
        );
    }

    #[test]
    fn test_metadata_store_postcard_roundtrip() {
        let mut store = NodeMetadataStore::new();

        store.insert(
            NodeId::new(1, 0),
            MacroNodeMetadata {
                macro_generated: Some(true),
                macro_source: Some("derive_Debug".to_string()),
                cfg_condition: Some("test".to_string()),
                cfg_active: Some(true),
                proc_macro_kind: Some(ProcMacroFunctionKind::Derive),
                expansion_cached: Some(false),
                unresolved_attributes: vec!["my_attr".to_string()],
            },
        );

        store.insert(
            NodeId::new(42, 3),
            MacroNodeMetadata {
                cfg_condition: Some("feature = \"serde\"".to_string()),
                ..Default::default()
            },
        );

        let bytes = postcard::to_allocvec(&store).expect("serialize");
        let deserialized: NodeMetadataStore = postcard::from_bytes(&bytes).expect("deserialize");

        assert_eq!(store, deserialized);
    }

    #[test]
    fn test_empty_metadata_store_zero_overhead() {
        let store = NodeMetadataStore::new();
        let bytes = postcard::to_allocvec(&store).expect("serialize");

        // Empty HashMap serializes to a single varint length of 0
        assert!(
            bytes.len() <= 2,
            "Empty store should serialize to minimal bytes, got {} bytes",
            bytes.len()
        );
    }

    #[test]
    fn test_metadata_store_merge() {
        let mut store1 = NodeMetadataStore::new();
        let mut store2 = NodeMetadataStore::new();

        store1.insert(
            NodeId::new(1, 0),
            MacroNodeMetadata {
                macro_generated: Some(true),
                ..Default::default()
            },
        );

        store2.insert(
            NodeId::new(2, 0),
            MacroNodeMetadata {
                cfg_condition: Some("test".to_string()),
                ..Default::default()
            },
        );

        store1.merge(&store2);
        assert_eq!(store1.len(), 2);
        assert!(store1.get(NodeId::new(1, 0)).is_some());
        assert!(store1.get(NodeId::new(2, 0)).is_some());
    }

    #[test]
    fn test_proc_macro_function_kind_serde() {
        let kinds = [
            ProcMacroFunctionKind::Derive,
            ProcMacroFunctionKind::Attribute,
            ProcMacroFunctionKind::FunctionLike,
        ];

        for kind in kinds {
            let bytes = postcard::to_allocvec(&kind).expect("serialize");
            let deserialized: ProcMacroFunctionKind =
                postcard::from_bytes(&bytes).expect("deserialize");
            assert_eq!(kind, deserialized);
        }
    }

    #[test]
    fn test_metadata_get_or_insert_default() {
        let mut store = NodeMetadataStore::new();
        let node = NodeId::new(10, 0);

        // First access creates default
        let meta = store.get_or_insert_default(node);
        meta.cfg_condition = Some("test".to_string());

        // Second access returns existing
        let meta = store.get(node).unwrap();
        assert_eq!(meta.cfg_condition.as_deref(), Some("test"));
    }

    #[test]
    fn test_metadata_remove() {
        let mut store = NodeMetadataStore::new();
        let node = NodeId::new(1, 0);

        store.insert(
            node,
            MacroNodeMetadata {
                macro_generated: Some(true),
                ..Default::default()
            },
        );

        assert!(store.get(node).is_some());
        let removed = store.remove(node);
        assert!(removed.is_some());
        assert!(store.get(node).is_none());
        assert!(store.is_empty());
    }

    #[test]
    fn test_metadata_store_large_scale() {
        let mut store = NodeMetadataStore::new();

        // Insert 10K entries (simulating ~10% of a 100K-node codebase)
        for i in 0..10_000u32 {
            store.insert(
                NodeId::new(i, 0),
                MacroNodeMetadata {
                    cfg_condition: Some(format!("feature_{i}")),
                    ..Default::default()
                },
            );
        }

        assert_eq!(store.len(), 10_000);

        // Verify O(1) lookups
        assert!(store.get(NodeId::new(0, 0)).is_some());
        assert!(store.get(NodeId::new(5_000, 0)).is_some());
        assert!(store.get(NodeId::new(9_999, 0)).is_some());
        assert!(store.get(NodeId::new(10_000, 0)).is_none());

        // Verify round-trip
        let bytes = postcard::to_allocvec(&store).expect("serialize");
        let deserialized: NodeMetadataStore = postcard::from_bytes(&bytes).expect("deserialize");
        assert_eq!(store, deserialized);
    }

    #[test]
    fn test_classpath_metadata_insert_and_get() {
        let mut store = NodeMetadataStore::new();
        let node = NodeId::new(100, 0);

        let cp_meta = ClasspathNodeMetadata {
            coordinates: Some("com.google.guava:guava:33.0.0".to_string()),
            jar_path: "/home/user/.m2/repository/guava-33.0.0.jar".to_string(),
            fqn: "com.google.common.collect.ImmutableList".to_string(),
            is_direct_dependency: true,
        };

        store.insert_metadata(node, NodeMetadata::Classpath(cp_meta.clone()));
        assert_eq!(store.len(), 1);

        // get() should return None (only returns macro metadata)
        assert!(store.get(node).is_none());

        // get_metadata() should return the classpath metadata
        let retrieved = store.get_metadata(node).unwrap();
        match retrieved {
            NodeMetadata::Classpath(cp) => {
                assert_eq!(cp.fqn, "com.google.common.collect.ImmutableList");
                assert_eq!(
                    cp.coordinates.as_deref(),
                    Some("com.google.guava:guava:33.0.0")
                );
                assert!(cp.is_direct_dependency);
            }
            NodeMetadata::Macro(_) => panic!("expected Classpath variant"),
        }
    }

    #[test]
    fn test_classpath_metadata_postcard_roundtrip() {
        let mut store = NodeMetadataStore::new();

        // Mix of macro and classpath metadata
        store.insert(
            NodeId::new(1, 0),
            MacroNodeMetadata {
                macro_generated: Some(true),
                ..Default::default()
            },
        );

        store.insert_metadata(
            NodeId::new(2, 0),
            NodeMetadata::Classpath(ClasspathNodeMetadata {
                coordinates: Some("org.slf4j:slf4j-api:2.0.0".to_string()),
                jar_path: "slf4j-api-2.0.0.jar".to_string(),
                fqn: "org.slf4j.Logger".to_string(),
                is_direct_dependency: false,
            }),
        );

        let bytes = postcard::to_allocvec(&store).expect("serialize");
        let deserialized: NodeMetadataStore = postcard::from_bytes(&bytes).expect("deserialize");
        assert_eq!(store, deserialized);
        assert_eq!(deserialized.len(), 2);

        // Verify macro entry
        assert!(deserialized.get(NodeId::new(1, 0)).is_some());

        // Verify classpath entry via get_metadata
        let cp = deserialized.get_metadata(NodeId::new(2, 0)).unwrap();
        assert!(matches!(cp, NodeMetadata::Classpath(_)));
    }

    #[test]
    fn test_node_metadata_store_json_roundtrip() {
        let mut store = NodeMetadataStore::new();

        store.insert(
            NodeId::new(1, 0),
            MacroNodeMetadata {
                macro_generated: Some(true),
                macro_source: Some("serde_derive".to_string()),
                ..Default::default()
            },
        );

        store.insert_metadata(
            NodeId::new(2, 0),
            NodeMetadata::Classpath(ClasspathNodeMetadata {
                coordinates: None,
                jar_path: "rt.jar".to_string(),
                fqn: "java.lang.String".to_string(),
                is_direct_dependency: true,
            }),
        );

        let json = serde_json::to_string(&store).unwrap();
        let deserialized: NodeMetadataStore = serde_json::from_str(&json).unwrap();
        assert_eq!(store, deserialized);
    }

    #[test]
    fn test_iter_all_includes_both_types() {
        let mut store = NodeMetadataStore::new();

        store.insert(
            NodeId::new(1, 0),
            MacroNodeMetadata {
                macro_generated: Some(true),
                ..Default::default()
            },
        );

        store.insert_metadata(
            NodeId::new(2, 0),
            NodeMetadata::Classpath(ClasspathNodeMetadata {
                coordinates: None,
                jar_path: "test.jar".to_string(),
                fqn: "com.example.Test".to_string(),
                is_direct_dependency: true,
            }),
        );

        // iter() only yields macro entries
        let macro_entries: Vec<_> = store.iter().collect();
        assert_eq!(macro_entries.len(), 1);

        // iter_all() yields all entries
        let all_entries: Vec<_> = store.iter_all().collect();
        assert_eq!(all_entries.len(), 2);
    }

    #[test]
    fn test_remove_metadata_classpath() {
        let mut store = NodeMetadataStore::new();
        let node = NodeId::new(50, 0);

        store.insert_metadata(
            node,
            NodeMetadata::Classpath(ClasspathNodeMetadata {
                coordinates: None,
                jar_path: "test.jar".to_string(),
                fqn: "Test".to_string(),
                is_direct_dependency: true,
            }),
        );

        assert_eq!(store.len(), 1);

        // remove() returns None for non-macro entries
        let removed = store.remove(node);
        assert!(removed.is_none());
        // The entry is still gone from the store because remove() always removes
        assert!(store.is_empty());
    }

    #[test]
    fn test_remove_metadata_typed() {
        let mut store = NodeMetadataStore::new();
        let node = NodeId::new(50, 0);

        store.insert_metadata(
            node,
            NodeMetadata::Classpath(ClasspathNodeMetadata {
                coordinates: None,
                jar_path: "test.jar".to_string(),
                fqn: "Test".to_string(),
                is_direct_dependency: true,
            }),
        );

        // remove_metadata() returns the full NodeMetadata
        let removed = store.remove_metadata(node);
        assert!(matches!(removed, Some(NodeMetadata::Classpath(_))));
        assert!(store.is_empty());
    }
}