geographdb_core/

cfg_store.rs

//! CFG Store - Combines octree spatial indexing with persistent storage
//!
//! This module provides a high-level API for storing and querying
//! Control Flow Graph (CFG) blocks using 3D spatial indexing.

use anyhow::{Context, Result};
use std::collections::HashMap;
use std::path::Path;

use crate::spatial::{BoundingBox, Octree};
use crate::storage::{EdgeRec, MetadataRec, NodePoint, NodeRec, StorageManager};
use glam::Vec3;

/// CFG Block metadata stored separately from spatial index
#[derive(Debug, Clone)]
pub struct CfgBlockMetadata {
    pub block_kind: String,
    pub terminator: String,
    pub byte_start: u64,
    pub byte_end: u64,
    pub start_line: u64,
    pub start_col: u64,
    pub end_line: u64,
    pub end_col: u64,
}

/// CFG Block representation for spatial storage
#[derive(Debug, Clone)]
pub struct CfgBlock {
    pub id: u64,
    pub function_id: i64,
    pub block_kind: String,
    pub terminator: String,
    pub byte_start: u64,
    pub byte_end: u64,
    pub start_line: u64,
    pub start_col: u64,
    pub end_line: u64,
    pub end_col: u64,
    pub dominator_depth: u32,
    pub loop_nesting: u32,
    pub branch_count: u32,
}

impl CfgBlock {
    pub fn to_node_point(&self) -> NodePoint {
        NodePoint {
            id: self.id,
            x: self.dominator_depth as f32,
            y: self.loop_nesting as f32,
            z: self.function_id as f32, // Store function_id in z for spatial queries
        }
    }

    pub fn to_node_rec(&self, _assigned_id: u64) -> NodeRec {
        use crate::spatial::MortonEncode;
        // Debug: write to file
        if let Ok(mut f) = std::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open("/tmp/cfg_debug.log")
        {
            use std::io::Write;
            let _ = writeln!(
                f,
                "to_node_rec: id={}, function_id={}, z={}",
                self.id, self.function_id, self.function_id as f32
            );
        }
        NodeRec {
            id: self.id, // Use logical block ID for LSTS version tracking
            morton_code: <()>::encode_morton(
                self.dominator_depth,
                self.loop_nesting,
                self.function_id as u32, // Use function_id for spatial indexing (cast to u32)
            ),
            x: self.dominator_depth as f32,
            y: self.loop_nesting as f32,
            z: self.function_id as f32, // Store function_id in z for spatial queries
            edge_off: 0,
            edge_len: 0,
            // Store branch_count in flags field for later retrieval
            flags: self.branch_count,
            begin_ts: 0,
            end_ts: 0,
            tx_id: 0,
            visibility: 1,
            _padding: [0; 7],
        }
    }

    pub fn to_metadata(&self) -> CfgBlockMetadata {
        CfgBlockMetadata {
            block_kind: self.block_kind.clone(),
            terminator: self.terminator.clone(),
            byte_start: self.byte_start,
            byte_end: self.byte_end,
            start_line: self.start_line,
            start_col: self.start_col,
            end_line: self.end_line,
            end_col: self.end_col,
        }
    }

    pub fn from_metadata(
        id: u64,
        function_id: i64,
        meta: &CfgBlockMetadata,
        dominator_depth: u32,
        loop_nesting: u32,
        branch_count: u32,
    ) -> Self {
        CfgBlock {
            id,
            function_id,
            block_kind: meta.block_kind.clone(),
            terminator: meta.terminator.clone(),
            byte_start: meta.byte_start,
            byte_end: meta.byte_end,
            start_line: meta.start_line,
            start_col: meta.start_col,
            end_line: meta.end_line,
            end_col: meta.end_col,
            dominator_depth,
            loop_nesting,
            branch_count,
        }
    }
}

/// CFG Store - combines octree + storage for efficient CFG queries
pub struct CfgStore {
    /// Octree spatial index (public for persistence)
    pub octree: Octree,
    storage: StorageManager,
    function_bounds: HashMap<i64, BoundingBox>,
    /// Sidecar metadata storage indexed by storage-assigned block ID
    metadata: HashMap<u64, CfgBlockMetadata>,
    /// Map function_id to storage-assigned block IDs
    function_blocks: HashMap<i64, Vec<u64>>,
}

impl CfgStore {
    pub fn create(path: &Path) -> Result<Self> {
        let storage = StorageManager::create(path).context("Failed to create storage")?;
        let bounds = BoundingBox::new(Vec3::new(0.0, 0.0, 0.0), Vec3::new(100.0, 100.0, 100.0));
        let octree = Octree::new(bounds);
        Ok(Self {
            octree,
            storage,
            function_bounds: HashMap::new(),
            metadata: HashMap::new(),
            function_blocks: HashMap::new(),
        })
    }

    pub fn open(path: &Path) -> Result<Self> {
        let storage = StorageManager::open(path).context("Failed to open storage")?;
        let bounds = BoundingBox::new(Vec3::new(0.0, 0.0, 0.0), Vec3::new(100.0, 100.0, 100.0));
        let mut octree = Octree::new(bounds);

        // Rebuild in-memory data structures from persistent storage
        let mut metadata = HashMap::new();
        let mut function_blocks: HashMap<i64, Vec<u64>> = HashMap::new();
        let mut function_bounds: HashMap<i64, BoundingBox> = HashMap::new();

        // Iterate over all nodes to rebuild metadata and function_blocks
        for node_id in 0..storage.node_count() as u64 {
            if let Some(node) = storage.get_node(node_id) {
                // function_id stored in z coordinate during to_node_rec()
                let function_id = node.z as i64;

                // Load metadata from persistent storage
                let block_metadata = if let Some(meta_rec) = storage.get_metadata(node_id) {
                    CfgBlockMetadata {
                        block_kind: meta_rec.get_block_kind(),
                        terminator: meta_rec.get_terminator(),
                        byte_start: meta_rec.byte_start,
                        byte_end: meta_rec.byte_end,
                        start_line: meta_rec.start_line,
                        start_col: meta_rec.start_col,
                        end_line: meta_rec.end_line,
                        end_col: meta_rec.end_col,
                    }
                } else {
                    // Fallback to unknown for databases without metadata section
                    CfgBlockMetadata {
                        block_kind: "unknown".to_string(),
                        terminator: "unknown".to_string(),
                        byte_start: 0,
                        byte_end: 0,
                        start_line: 0,
                        start_col: 0,
                        end_line: 0,
                        end_col: 0,
                    }
                };
                metadata.insert(node_id, block_metadata);

                // Add to function_blocks mapping
                function_blocks
                    .entry(function_id)
                    .or_default()
                    .push(node_id);

                // Update function bounds
                let min = Vec3::new(node.x, node.y, node.z);
                let max = Vec3::new(node.x, node.y, node.z);
                let entry = function_bounds
                    .entry(function_id)
                    .or_insert(BoundingBox::new(min, max));
                entry.min = entry.min.min(min);
                entry.max = entry.max.max(max);

                // Rebuild octree spatial index
                let point = NodePoint {
                    id: node_id,
                    x: node.x,
                    y: node.y,
                    z: node.z,
                };
                octree.insert(point);
            }
        }

        Ok(Self {
            octree,
            storage,
            function_bounds,
            metadata,
            function_blocks,
        })
    }

    pub fn insert_block(&mut self, block: CfgBlock) -> Result<u64> {
        let function_id = block.function_id;

        // Insert into storage and get assigned ID
        let assigned_id = self
            .storage
            .insert_node(block.to_node_rec(self.storage.node_count() as u64))
            .context("Failed to insert into storage")?;

        // Create metadata record for persistent storage and store at the same ID as the node
        let meta_rec = MetadataRec::from_strings(
            &block.block_kind,
            &block.terminator,
            block.byte_start,
            block.byte_end,
            block.start_line,
            block.start_col,
            block.end_line,
            block.end_col,
        );

        // Persist metadata to disk at the same ID as the node
        self.storage
            .insert_metadata_at(assigned_id, meta_rec)
            .context("Failed to insert metadata")?;

        // Store metadata in memory HashMap for fast access
        let metadata = block.to_metadata();
        self.metadata.insert(assigned_id, metadata);

        // Track which blocks belong to which function
        self.function_blocks
            .entry(function_id)
            .or_default()
            .push(assigned_id);

        self.update_function_bounds(function_id, &block);
        self.octree.insert(block.to_node_point());

        Ok(assigned_id)
    }

    pub fn insert_blocks(&mut self, blocks: Vec<CfgBlock>) -> Result<()> {
        for block in blocks {
            self.insert_block(block)?;
        }
        Ok(())
    }

    pub fn query_nearby(&self, point: Vec3, radius: f32) -> Vec<NodePoint> {
        self.octree.query_sphere(point, radius)
    }

    /// Find the k nearest CFG blocks to a query point.
    ///
    /// Returns up to k nearest blocks sorted by ascending distance.
    /// Each result includes the distance squared for precision comparison.
    ///
    /// # Example
    /// ```ignore
    /// let nearest = store.query_knn(glam::Vec3::new(1.0, 2.0, 3.0), 5);
    /// for (node_point, dist_sq) in nearest {
    ///     println!("Block {} at distance {}", node_point.id, dist_sq.sqrt());
    /// }
    /// ```
    pub fn query_knn(&self, point: Vec3, k: usize) -> Vec<(NodePoint, f32)> {
        self.octree.query_knn(point, k)
    }

    pub fn block_count(&self) -> usize {
        self.storage.node_count()
    }

    pub fn metadata_count(&self) -> usize {
        self.storage.metadata_count()
    }

    pub fn get_stored_metadata(&self, id: u64) -> Option<&crate::storage::MetadataRec> {
        self.storage.get_metadata(id)
    }

    /// Flush any pending changes to disk
    pub fn flush(&mut self) -> Result<()> {
        self.storage.flush()
    }

    /// LSTS: Get a block at a specific timestamp (time-travel query)
    ///
    /// Returns the CfgBlock version that was visible at the given timestamp.
    /// This enables querying the CFG as it existed at any point in time.
    pub fn get_block_at_timestamp(
        &self,
        logical_block_id: u64,
        timestamp: u64,
    ) -> Option<CfgBlock> {
        // Use the storage manager's LSTS query to find the right version
        let _node = self
            .storage
            .get_node_at_timestamp(logical_block_id, timestamp)?;

        // Find the storage ID for this version
        // We need to scan the version index to find which physical storage ID
        // corresponds to this logical block at this timestamp
        let versions = self.storage.get_version_history(logical_block_id)?;

        for &version_id in versions {
            if let Some(node_version) = self.storage.get_node(version_id) {
                if node_version.begin_ts <= timestamp
                    && (node_version.end_ts == 0 || node_version.end_ts > timestamp)
                {
                    // Found the right version, now get the metadata
                    if let Some(meta) = self.metadata.get(&version_id) {
                        // Determine function_id from the metadata or node
                        // For now, we scan function_blocks to find which function owns this block
                        let function_id = self
                            .function_blocks
                            .iter()
                            .find(|(_, blocks)| blocks.contains(&version_id))
                            .map(|(fid, _)| *fid)
                            .unwrap_or(0);

                        return Some(CfgBlock::from_metadata(
                            logical_block_id,
                            function_id,
                            meta,
                            node_version.x as u32,
                            node_version.y as u32,
                            node_version.z as u32,
                        ));
                    }
                }
            }
        }

        None
    }

    /// Get all blocks for a function WITH FULL METADATA
    pub fn get_blocks_for_function(&self, function_id: i64) -> Vec<CfgBlock> {
        let mut blocks = Vec::new();

        // Get block IDs for this function
        if let Some(block_ids) = self.function_blocks.get(&function_id) {
            for &block_id in block_ids {
                // Get spatial data from storage
                if let Some(node) = self.storage.get_node(block_id) {
                    // Get metadata from sidecar table
                    if let Some(meta) = self.metadata.get(&block_id) {
                        // Use node.id (the logical block ID) instead of storage position
                        // Note: node.z is function_id, node.flags is branch_count
                        blocks.push(CfgBlock::from_metadata(
                            node.id,
                            function_id,
                            meta,
                            node.x as u32, // dominator_depth
                            node.y as u32, // loop_nesting
                            node.flags,    // branch_count (stored in flags)
                        ));
                    }
                }
            }
        }

        blocks
    }

    pub fn get_all_edges(&self) -> Vec<EdgeRec> {
        let mut edges = Vec::new();
        for i in 0..self.storage.edge_count() {
            if let Some(edge) = self.storage.get_edge(i as u64) {
                edges.push(*edge);
            }
        }
        edges
    }
    /// Insert an edge into storage
    pub fn insert_edge(&mut self, edge: EdgeRec) -> Result<u64> {
        self.storage
            .insert_edge(edge)
            .context("Failed to insert edge")
    }

    pub fn get_edges_for_node(&self, source_idx: u64) -> Vec<EdgeRec> {
        let mut edges = Vec::new();
        for i in 0..self.storage.edge_count() {
            if let Some(edge) = self.storage.get_edge(i as u64) {
                if edge.src == source_idx {
                    edges.push(*edge);
                }
            }
        }
        edges
    }

    /// Get the set of function IDs that have blocks in this store
    ///
    /// Returns all function IDs that have at least one CFG block.
    /// This is useful for vacuum operations to determine which functions are present.
    pub fn get_function_ids(&self) -> std::collections::HashSet<i64> {
        self.function_blocks.keys().copied().collect()
    }

    /// Get block IDs for a specific function
    ///
    /// Returns the block IDs that belong to the given function, or None if the
    /// function has no blocks in this store.
    pub fn get_block_ids_for_function(&self, function_id: i64) -> Option<&[u64]> {
        self.function_blocks.get(&function_id).map(|v| v.as_slice())
    }

    /// Get count of blocks for a specific function
    pub fn count_blocks_for_function(&self, function_id: i64) -> usize {
        self.function_blocks
            .get(&function_id)
            .map(|v| v.len())
            .unwrap_or(0)
    }

    /// Get the total edge count in storage
    pub fn edge_count(&self) -> usize {
        self.storage.edge_count()
    }

    fn update_function_bounds(&mut self, function_id: i64, block: &CfgBlock) {
        let pos = Vec3::new(
            block.dominator_depth as f32,
            block.loop_nesting as f32,
            block.branch_count as f32,
        );
        self.function_bounds
            .entry(function_id)
            .and_modify(|bounds| {
                bounds.min = bounds.min.min(pos);
                bounds.max = bounds.max.max(pos);
            })
            .or_insert_with(|| {
                BoundingBox::new(
                    pos - Vec3::new(5.0, 5.0, 5.0),
                    pos + Vec3::new(5.0, 5.0, 5.0),
                )
            });
    }
}
#[cfg(test)]
mod tests {
    use super::*;
    use crate::algorithms::astar::CfgGraphNode;
    use crate::algorithms::dominance::compute_dominance;
    use tempfile::tempdir;

    #[test]
    fn test_cfg_store_create() {
        let temp_dir = tempdir().unwrap();
        let db_path = temp_dir.path().join("test_cfg.db");
        let result = CfgStore::create(&db_path);
        assert!(result.is_ok(), "Should create CFG store");
    }

    /// Layer 1 Test: Verify dominator depths are computed correctly from CFG structure
    /// This test will FAIL initially because we need to integrate dominator analysis
    /// with CfgBlock coordinate computation
    #[test]
    fn test_dominator_depth_computation_from_cfg() {
        // Create a simple CFG: Entry -> A -> B -> C
        //                           \--> D ----/
        // Dominator depths should be:
        // Entry: 0 (dominates itself)
        // A: 1 (dominated by Entry)
        // B: 2 (dominated by Entry, A)
        // C: 3 (dominated by Entry, A, B) OR 2 (dominated by Entry, D) - depending on path
        // D: 1 (dominated by Entry)

        let nodes = vec![
            CfgGraphNode {
                id: 0,
                x: 0.0,
                y: 0.0,
                z: 0.0,
                successors: vec![1, 4],
            }, // Entry -> A, D
            CfgGraphNode {
                id: 1,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![2],
            }, // A -> B
            CfgGraphNode {
                id: 2,
                x: 2.0,
                y: 0.0,
                z: 0.0,
                successors: vec![3],
            }, // B -> C
            CfgGraphNode {
                id: 3,
                x: 3.0,
                y: 0.0,
                z: 0.0,
                successors: vec![],
            }, // C (exit)
            CfgGraphNode {
                id: 4,
                x: 1.0,
                y: 1.0,
                z: 0.0,
                successors: vec![3],
            }, // D -> C
        ];

        let dom_result = compute_dominance(&nodes, 0);

        // Layer 1: Basic dominator computation succeeded
        assert!(
            dom_result.dominators.contains_key(&0),
            "Layer 1: Entry should have dominators"
        );
        assert!(
            dom_result.dominators.contains_key(&3),
            "Layer 1: Exit block C should have dominators"
        );

        // Layer 2: Verify dominator depths are reasonable
        let entry_depth = dom_result.dominators.get(&0).unwrap().len() as u32;
        let a_depth = dom_result.dominators.get(&1).unwrap().len() as u32;
        let b_depth = dom_result.dominators.get(&2).unwrap().len() as u32;
        let c_depth = dom_result.dominators.get(&3).unwrap().len() as u32;
        let d_depth = dom_result.dominators.get(&4).unwrap().len() as u32;

        assert_eq!(
            entry_depth, 1,
            "Layer 2: Entry should have depth 1 (only itself)"
        );
        assert!(
            a_depth >= 2,
            "Layer 2: A should have depth >= 2 (Entry + A)"
        );
        assert!(b_depth >= a_depth, "Layer 2: B depth should be >= A depth");
        assert!(
            c_depth >= 2,
            "Layer 2: C should have at least entry + self dominators"
        );
        assert!(
            d_depth >= 2,
            "Layer 2: D should have depth >= 2 (Entry + D)"
        );

        // Layer 3: Verify CfgBlock coordinates would be computed correctly
        // This validates that the dominator depth maps to X coordinate
        let temp_dir = tempdir().unwrap();
        let db_path = temp_dir.path().join("test_dominator.db");
        let mut store = CfgStore::create(&db_path).unwrap();

        // Insert blocks with computed dominator depths
        let blocks = vec![
            CfgBlock {
                id: 0,
                function_id: 1,
                block_kind: "entry".to_string(),
                terminator: "fallthrough".to_string(),
                byte_start: 0,
                byte_end: 10,
                start_line: 1,
                start_col: 0,
                end_line: 1,
                end_col: 10,
                dominator_depth: entry_depth,
                loop_nesting: 0,
                branch_count: 0,
            },
            CfgBlock {
                id: 1,
                function_id: 1,
                block_kind: "block".to_string(),
                terminator: "fallthrough".to_string(),
                byte_start: 10,
                byte_end: 20,
                start_line: 2,
                start_col: 0,
                end_line: 2,
                end_col: 10,
                dominator_depth: a_depth,
                loop_nesting: 0,
                branch_count: 0,
            },
            CfgBlock {
                id: 2,
                function_id: 1,
                block_kind: "block".to_string(),
                terminator: "fallthrough".to_string(),
                byte_start: 20,
                byte_end: 30,
                start_line: 3,
                start_col: 0,
                end_line: 3,
                end_col: 10,
                dominator_depth: b_depth,
                loop_nesting: 0,
                branch_count: 0,
            },
        ];

        for block in &blocks {
            store.insert_block(block.clone()).unwrap();
        }

        // Retrieve and verify coordinates
        let retrieved = store.get_blocks_for_function(1);
        assert_eq!(retrieved.len(), 3, "Layer 3: Should retrieve all 3 blocks");

        // Verify dominator depths are preserved in spatial coordinates
        let entry_block = retrieved.iter().find(|b| b.block_kind == "entry").unwrap();
        assert_eq!(
            entry_block.dominator_depth, 1,
            "Layer 3: Entry dominator depth should be preserved"
        );

        // Layer 4: Spatial query should respect dominator structure
        // Blocks with similar dominator depths should be spatially close
        let nearby = store.query_nearby(glam::Vec3::new(entry_depth as f32, 0.0, 0.0), 5.0);
        assert!(
            !nearby.is_empty(),
            "Layer 4: Should find entry block by dominator depth coordinate"
        );
    }

    /// Phase D Layer 3 Test: Octree persistence integration with CfgStore
    ///
    /// Tests that CfgStore can save and restore its octree spatial index.
    /// This enables fast startup by avoiding octree rebuild from scratch.
    #[test]
    fn test_cfg_store_octree_persistence() {
        let temp_dir = tempdir().unwrap();
        let db_path = temp_dir.path().join("test_octree_persist.db");

        // Create CfgStore and insert blocks
        let mut store = CfgStore::create(&db_path).unwrap();

        let blocks = vec![
            CfgBlock {
                id: 1,
                function_id: 42,
                block_kind: "entry".to_string(),
                terminator: "fallthrough".to_string(),
                byte_start: 0,
                byte_end: 50,
                start_line: 1,
                start_col: 0,
                end_line: 5,
                end_col: 10,
                dominator_depth: 0,
                loop_nesting: 0,
                branch_count: 0,
            },
            CfgBlock {
                id: 2,
                function_id: 42,
                block_kind: "if".to_string(),
                terminator: "conditional".to_string(),
                byte_start: 50,
                byte_end: 100,
                start_line: 6,
                start_col: 4,
                end_line: 10,
                end_col: 15,
                dominator_depth: 1,
                loop_nesting: 0,
                branch_count: 1,
            },
            CfgBlock {
                id: 3,
                function_id: 42,
                block_kind: "loop".to_string(),
                terminator: "jump".to_string(),
                byte_start: 100,
                byte_end: 150,
                start_line: 11,
                start_col: 4,
                end_line: 15,
                end_col: 20,
                dominator_depth: 2,
                loop_nesting: 1,
                branch_count: 2,
            },
        ];

        for block in &blocks {
            store.insert_block(block.clone()).unwrap();
        }

        // Verify blocks are queryable
        let nearby_before = store.query_nearby(glam::Vec3::new(1.0, 0.0, 42.0), 5.0);
        assert!(
            !nearby_before.is_empty(),
            "Should find blocks before persistence"
        );

        // Serialize octree
        let octree_bytes = store.octree.to_bytes().expect("Should serialize octree");
        assert!(
            !octree_bytes.is_empty(),
            "Serialized octree should not be empty"
        );

        // Create new CfgStore and restore octree
        let mut store2 = CfgStore::create(&db_path).unwrap();
        store2.octree = Octree::from_bytes(&octree_bytes).expect("Should deserialize octree");

        // Verify blocks are still queryable after restore
        let nearby_after = store2.query_nearby(glam::Vec3::new(1.0, 0.0, 42.0), 5.0);
        assert!(
            !nearby_after.is_empty(),
            "Should find blocks after octree restore"
        );
        assert_eq!(
            nearby_before.len(),
            nearby_after.len(),
            "Should find same number of blocks"
        );
    }

    /// Phase D Layer 4 Test: Telemetry for octree operations
    ///
    /// Verifies loop guards prevent excessive recursion in octree operations.
    #[test]
    fn test_cfg_store_function_index_roundtrip() {
        let temp_dir = tempdir().unwrap();
        let db_path = temp_dir.path().join("test_function_roundtrip.db");

        // Create store and insert blocks for function 42
        let mut store = CfgStore::create(&db_path).unwrap();
        let block = CfgBlock {
            id: 1,
            function_id: 42,
            block_kind: "entry".to_string(),
            terminator: "fallthrough".to_string(),
            byte_start: 0,
            byte_end: 50,
            start_line: 1,
            start_col: 0,
            end_line: 5,
            end_col: 10,
            dominator_depth: 0,
            loop_nesting: 0,
            branch_count: 0,
        };
        store.insert_block(block.clone()).unwrap();
        assert_eq!(store.count_blocks_for_function(42), 1);

        // Open a fresh store and verify the function index is rebuilt
        let store2 = CfgStore::open(&db_path).unwrap();
        assert_eq!(
            store2.count_blocks_for_function(42),
            1,
            "function_blocks should be rebuilt from storage on open"
        );
    }
    #[test]
    #[cfg(feature = "telemetry")]
    fn test_octree_telemetry_loop_guard() {
        use crate::telemetry::LoopGuard;

        // Test that loop guard catches excessive iterations
        let guard = LoopGuard::new("octree_traversal", 100);

        // Simulate octree traversal
        for i in 0..150 {
            if i < 100 {
                assert!(guard.check().is_ok(), "Should allow iteration {}", i);
            } else {
                assert!(guard.check().is_err(), "Should fail at iteration {}", i);
                break;
            }
        }
    }
}