compliance-scanner-agent/compliance-agent/src/controls/index.rs

//! In-memory embedding index over the control corpus, for region → control
//! retrieval.
//!
//! At master-control scale (~13.6k) findings can't be mapped by CWE (the master
//! controls carry none), so we map by *similarity*: embed each control's
//! requirement text once, then for a code region pull the top-K nearest controls
//! to hand to the grounded judge. This is the retrieval half of the semantic path.

use std::path::Path;

use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

use compliance_core::control_check::ControlCheckSpec;
use compliance_core::error::CoreError;

use crate::llm::LlmClient;

/// A control spec paired with its requirement-text embedding.
pub struct ControlIndex {
    entries: Vec<(ControlCheckSpec, Vec<f64>)>,
}

/// On-disk form of the index: the corpus identity hash plus every spec+embedding.
/// The hash lets a later scan reuse the embeddings only if the corpus is unchanged.
#[derive(Serialize, Deserialize)]
struct PersistedIndex {
    corpus_hash: String,
    entries: Vec<PersistedEntry>,
}

#[derive(Serialize, Deserialize)]
struct PersistedEntry {
    spec: ControlCheckSpec,
    embedding: Vec<f64>,
}

/// Stable hash of the corpus identity (each control's id + requirement text, in
/// order). Same catalog → same hash → the cached embeddings are reused instead of
/// re-embedding the whole corpus.
fn corpus_hash(specs: &[ControlCheckSpec]) -> String {
    let mut hasher = Sha256::new();
    for s in specs {
        hasher.update(s.control_id.as_bytes());
        hasher.update([0u8]);
        hasher.update(s.requirement.as_bytes());
        hasher.update([0u8]);
    }
    format!("{:x}", hasher.finalize())
}

impl ControlIndex {
    /// Build directly from precomputed embeddings (used by tests + callers that
    /// already embedded the corpus).
    pub fn from_embeddings(entries: Vec<(ControlCheckSpec, Vec<f64>)>) -> Self {
        Self { entries }
    }

    /// Load the index from `cache_path` if it still matches the current corpus,
    /// otherwise embed the corpus and persist it there. This turns the per-scan
    /// re-embed of the whole (~13.6k) master-control corpus into a one-time cost
    /// that survives across scans; the cache self-invalidates when the catalog
    /// changes (its [`corpus_hash`] no longer matches).
    pub async fn load_or_build(
        llm: &LlmClient,
        specs: Vec<ControlCheckSpec>,
        cache_path: &Path,
    ) -> Result<Self, CoreError> {
        let hash = corpus_hash(&specs);
        if let Some(index) = Self::load_cache(cache_path, &hash).await {
            tracing::debug!(
                controls = index.len(),
                "reusing cached control embedding index"
            );
            return Ok(index);
        }
        let index = Self::build(llm, specs).await?;
        if let Err(e) = index.write_cache(cache_path, &hash).await {
            tracing::warn!(error = %e, "failed to persist control embedding index");
        }
        Ok(index)
    }

    /// Read a persisted index, returning it only if its corpus hash matches.
    async fn load_cache(path: &Path, hash: &str) -> Option<Self> {
        let raw = tokio::fs::read(path).await.ok()?;
        let persisted: PersistedIndex = serde_json::from_slice(&raw).ok()?;
        if persisted.corpus_hash != hash {
            return None;
        }
        Some(Self {
            entries: persisted
                .entries
                .into_iter()
                .map(|e| (e.spec, e.embedding))
                .collect(),
        })
    }

    /// Persist the index atomically (temp file + rename) keyed by corpus hash.
    async fn write_cache(&self, path: &Path, hash: &str) -> Result<(), CoreError> {
        if let Some(parent) = path.parent() {
            tokio::fs::create_dir_all(parent).await?;
        }
        let persisted = PersistedIndex {
            corpus_hash: hash.to_string(),
            entries: self
                .entries
                .iter()
                .map(|(spec, emb)| PersistedEntry {
                    spec: spec.clone(),
                    embedding: emb.clone(),
                })
                .collect(),
        };
        let raw = serde_json::to_vec(&persisted)?;
        let tmp = path.with_extension("json.tmp");
        tokio::fs::write(&tmp, &raw).await?;
        tokio::fs::rename(&tmp, path).await?;
        Ok(())
    }

    /// Build by embedding each control's requirement text.
    pub async fn build(llm: &LlmClient, specs: Vec<ControlCheckSpec>) -> Result<Self, CoreError> {
        if specs.is_empty() {
            return Ok(Self {
                entries: Vec::new(),
            });
        }
        let texts: Vec<String> = specs.iter().map(|s| s.requirement.clone()).collect();
        let embeddings = llm
            .embed(texts)
            .await
            .map_err(|e| CoreError::Llm(e.to_string()))?;
        Ok(Self {
            entries: specs.into_iter().zip(embeddings).collect(),
        })
    }

    pub fn len(&self) -> usize {
        self.entries.len()
    }

    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// The top-`k` control specs whose embedding is nearest (cosine) to `query`.
    pub fn nearest(&self, query: &[f64], k: usize) -> Vec<ControlCheckSpec> {
        let mut scored: Vec<(f64, &ControlCheckSpec)> = self
            .entries
            .iter()
            .map(|(spec, emb)| (cosine(query, emb), spec))
            .collect();
        scored.sort_by(|a, b| b.0.total_cmp(&a.0));
        scored.into_iter().take(k).map(|(_, s)| s.clone()).collect()
    }
}

/// Cosine similarity; 0.0 for length-mismatched, empty, or zero vectors.
fn cosine(a: &[f64], b: &[f64]) -> f64 {
    if a.len() != b.len() || a.is_empty() {
        return 0.0;
    }
    let dot: f64 = a.iter().zip(b).map(|(x, y)| x * y).sum();
    let na: f64 = a.iter().map(|x| x * x).sum();
    let nb: f64 = b.iter().map(|x| x * x).sum();
    if na == 0.0 || nb == 0.0 {
        return 0.0;
    }
    dot / (na.sqrt() * nb.sqrt())
}

#[cfg(test)]
mod tests {
    use super::*;
    use compliance_core::models::finding::Severity;

    fn spec(id: &str) -> ControlCheckSpec {
        ControlCheckSpec {
            control_id: id.into(),
            title: id.into(),
            requirement: id.into(),
            default_cwe: None,
            severity: Severity::Medium,
        }
    }

    #[test]
    fn nearest_ranks_by_cosine() {
        let index = ControlIndex::from_embeddings(vec![
            (spec("a"), vec![1.0, 0.0]),
            (spec("b"), vec![0.0, 1.0]),
            (spec("c"), vec![0.7, 0.7]),
        ]);
        let hits = index.nearest(&[0.9, 0.1], 2);
        assert_eq!(hits.len(), 2);
        assert_eq!(hits[0].control_id, "a"); // closest to [0.9,0.1]
    }

    #[test]
    fn cosine_edges_are_zero() {
        assert_eq!(cosine(&[1.0], &[1.0, 2.0]), 0.0); // length mismatch
        assert_eq!(cosine(&[0.0, 0.0], &[1.0, 1.0]), 0.0); // zero vector
        assert!((cosine(&[1.0, 0.0], &[1.0, 0.0]) - 1.0).abs() < 1e-9); // identical
    }

    #[test]
    fn corpus_hash_is_stable_and_identity_sensitive() {
        let a = corpus_hash(&[spec("x"), spec("y")]);
        assert_eq!(a, corpus_hash(&[spec("x"), spec("y")])); // same corpus → same hash
        assert_ne!(a, corpus_hash(&[spec("y"), spec("x")])); // reorder → different
        assert_ne!(a, corpus_hash(&[spec("x")])); // fewer controls → different
    }

    #[tokio::test]
    #[allow(clippy::unwrap_used)]
    async fn cache_round_trips_and_misses_on_corpus_change() {
        let dir = std::env::temp_dir().join(format!("cidx-{}", uuid::Uuid::new_v4()));
        let path = dir.join("control-index.json");
        let specs = [spec("a"), spec("b")];
        let hash = corpus_hash(&specs);
        let index = ControlIndex::from_embeddings(vec![
            (spec("a"), vec![1.0, 0.0]),
            (spec("b"), vec![0.0, 1.0]),
        ]);
        index.write_cache(&path, &hash).await.unwrap();

        // matching corpus hash → hit
        let loaded = ControlIndex::load_cache(&path, &hash).await.unwrap();
        assert_eq!(loaded.len(), 2);
        assert_eq!(loaded.nearest(&[0.9, 0.1], 1)[0].control_id, "a");
        // corpus changed → miss (forces a rebuild)
        assert!(ControlIndex::load_cache(&path, "differenthash")
            .await
            .is_none());
        // absent file → miss, not an error
        assert!(
            ControlIndex::load_cache(dir.join("nope.json").as_path(), &hash)
                .await
                .is_none()
        );
        let _ = std::fs::remove_dir_all(&dir);
    }
}