srcgraph_metrics/

association_rules.rs

//! Apriori-style frequent itemset mining + association-rule generation over
//! per-class `callSequences` blobs.
//!
//! Each class node carries a `callSequences` JSON blob with shape
//!
//!   `{"sequences": [{"method": str, "calls": [str], …}, …]}`
//!
//! Every sequence with `len(calls) >= 2` becomes one transaction (the
//! `frozenset` of its `calls`). We then:
//!
//! 1. Mine frequent itemsets by Apriori (bottom-up by k), capped at `k = 5` to
//!    match the Python reference.
//! 2. Generate rules `A → C` for every non-empty proper subset `A` of each
//!    frequent itemset, scoring with `confidence = sup(A∪C)/sup(A)` and
//!    `lift = confidence / (sup(C)/n)`.
//! 3. Classify each rule as `invariant` (conf ≥ 0.99), `strong` (≥ 0.85), or
//!    `moderate` (≥ 0.5).
//!
//! Mirrors `analysis/association_rules.py` in the visiting tool.
//!
//! See `DESIGN.md` (Phase 3) at the workspace root.

use petgraph::Graph;
use serde::{Deserialize, Serialize};
use std::collections::{BTreeSet, HashMap};

use srcgraph_core::{ClassNode, EdgeKind};

/// One call-sequence record parsed from a class's `callSequences` blob.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CallSequence {
    pub method: String,
    pub calls: Vec<String>,
}

/// A frequent itemset and its support count across the transaction list.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FrequentItemset {
    /// Items in deterministic (sorted) order.
    pub items: Vec<String>,
    pub support: usize,
}

/// An association rule `antecedent → consequent` with support, confidence,
/// and lift.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AssociationRule {
    pub antecedent: Vec<String>,
    pub consequent: Vec<String>,
    pub support: usize,
    /// `support / |transactions|`, rounded to 4 decimals.
    pub support_pct: f64,
    /// `sup(A∪C) / sup(A)`, rounded to 4 decimals.
    pub confidence: f64,
    /// `confidence / (sup(C) / n)`, rounded to 4 decimals. `0.0` when `sup(C)`
    /// is unknown (consequent wasn't a frequent itemset).
    pub lift: f64,
    /// `"invariant"` (conf ≥ 0.99) / `"strong"` (≥ 0.85) / `"moderate"` (≥ 0.5).
    pub classification: String,
}

/// Whole-graph association-rule readout.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AssociationAnalysis {
    pub transactions: usize,
    pub rules: Vec<AssociationRule>,
    pub num_rules: usize,
    pub invariants: usize,
    pub strong: usize,
    pub moderate: usize,
    pub itemsets: usize,
}

/// Parse the `{"sequences": [...]}` blob attached to a single class.
///
/// Returns `None` if the blob isn't an object with a `sequences` array; missing
/// `method`/`calls` default to empty.
pub fn parse_call_sequences(blob: &serde_json::Value) -> Option<Vec<CallSequence>> {
    let seqs = blob.get("sequences")?.as_array()?;
    let mut out = Vec::with_capacity(seqs.len());
    for s in seqs {
        let Some(obj) = s.as_object() else { continue };
        let method = obj.get("method").and_then(|v| v.as_str()).unwrap_or("").to_owned();
        let calls = obj
            .get("calls")
            .and_then(|v| v.as_array())
            .map(|arr| {
                arr.iter()
                    .filter_map(|c| c.as_str().map(|s| s.to_owned()))
                    .collect()
            })
            .unwrap_or_default();
        out.push(CallSequence { method, calls });
    }
    Some(out)
}

/// Mine frequent itemsets by Apriori. Caps `k` at 5 (matches Python ref) to
/// keep candidate generation tractable on large unique-item alphabets.
///
/// Result is sorted by support descending, then by item-list ascending for
/// determinism among ties.
pub fn mine_frequent_itemsets(
    transactions: &[BTreeSet<String>],
    min_support: usize,
) -> Vec<FrequentItemset> {
    if transactions.is_empty() {
        return Vec::new();
    }

    // Frequent 1-itemsets via direct counting.
    let mut item_counts: HashMap<&str, usize> = HashMap::new();
    for txn in transactions {
        for item in txn {
            *item_counts.entry(item.as_str()).or_insert(0) += 1;
        }
    }
    let mut freq_items: Vec<String> = item_counts
        .iter()
        .filter(|(_, &c)| c >= min_support)
        .map(|(k, _)| (*k).to_owned())
        .collect();
    freq_items.sort();

    let mut frequent: HashMap<BTreeSet<String>, usize> = HashMap::new();
    for item in &freq_items {
        let mut s = BTreeSet::new();
        s.insert(item.clone());
        let c = item_counts[item.as_str()];
        frequent.insert(s, c);
    }

    // Iterate k = 2..=5; generate candidates from all k-combinations of
    // frequent 1-items, then keep those whose (k-1)-subsets are all frequent
    // (the Apriori property) and whose support meets the threshold.
    let mut k = 2usize;
    let mut prev_has_freq = !freq_items.is_empty();
    while prev_has_freq && k <= 5 {
        let mut new_count = 0usize;
        for combo in combinations(&freq_items, k) {
            let cand: BTreeSet<String> = combo.iter().cloned().collect();
            // All (k-1)-subsets frequent?
            let all_sub_freq = cand.iter().all(|item| {
                let mut sub = cand.clone();
                sub.remove(item);
                frequent.contains_key(&sub)
            });
            if !all_sub_freq {
                continue;
            }
            let support = transactions
                .iter()
                .filter(|txn| cand.iter().all(|x| txn.contains(x)))
                .count();
            if support >= min_support {
                frequent.insert(cand, support);
                new_count += 1;
            }
        }
        prev_has_freq = new_count > 0;
        k += 1;
    }

    let mut out: Vec<FrequentItemset> = frequent
        .into_iter()
        .map(|(set, support)| FrequentItemset {
            items: set.into_iter().collect(),
            support,
        })
        .collect();
    out.sort_by(|a, b| b.support.cmp(&a.support).then_with(|| a.items.cmp(&b.items)));
    out
}

/// Generate association rules from frequent itemsets.
///
/// Returns rules sorted by confidence desc, then support desc, then antecedent
/// ascending for deterministic ordering among ties.
pub fn generate_rules(
    itemsets: &[FrequentItemset],
    transactions: &[BTreeSet<String>],
    min_confidence: f64,
) -> Vec<AssociationRule> {
    let n_txns = transactions.len();
    if n_txns == 0 || itemsets.is_empty() {
        return Vec::new();
    }

    // Support lookup keyed on the itemset as BTreeSet.
    let support_map: HashMap<BTreeSet<String>, usize> = itemsets
        .iter()
        .map(|fi| (fi.items.iter().cloned().collect(), fi.support))
        .collect();

    let mut rules: Vec<AssociationRule> = Vec::new();

    for fi in itemsets {
        if fi.items.len() < 2 {
            continue;
        }
        let support = fi.support;
        // Enumerate every non-empty proper subset as antecedent.
        for i in 1..fi.items.len() {
            for ant_vec in combinations(&fi.items, i) {
                let antecedent: BTreeSet<String> = ant_vec.iter().cloned().collect();
                let full: BTreeSet<String> = fi.items.iter().cloned().collect();
                let consequent: BTreeSet<String> = full.difference(&antecedent).cloned().collect();

                let Some(&ant_support) = support_map.get(&antecedent) else {
                    continue;
                };
                if ant_support == 0 {
                    continue;
                }
                let confidence = support as f64 / ant_support as f64;
                if confidence < min_confidence {
                    continue;
                }
                let cons_support = support_map.get(&consequent).copied().unwrap_or(0);
                let lift = if cons_support > 0 {
                    confidence / (cons_support as f64 / n_txns as f64)
                } else {
                    0.0
                };
                let conf_r = round4(confidence);
                rules.push(AssociationRule {
                    antecedent: antecedent.into_iter().collect(),
                    consequent: consequent.into_iter().collect(),
                    support,
                    support_pct: round4(support as f64 / n_txns as f64),
                    confidence: conf_r,
                    lift: round4(lift),
                    classification: classify_rule(conf_r).to_owned(),
                });
            }
        }
    }

    rules.sort_by(|a, b| {
        b.confidence
            .partial_cmp(&a.confidence)
            .unwrap_or(std::cmp::Ordering::Equal)
            .then_with(|| b.support.cmp(&a.support))
            .then_with(|| a.antecedent.cmp(&b.antecedent))
            .then_with(|| a.consequent.cmp(&b.consequent))
    });
    rules
}

/// Classify by confidence threshold — `invariant` (≥0.99), `strong` (≥0.85),
/// `moderate` otherwise.
pub fn classify_rule(confidence: f64) -> &'static str {
    if confidence >= 0.99 {
        "invariant"
    } else if confidence >= 0.85 {
        "strong"
    } else {
        "moderate"
    }
}

/// Walk every node's `callSequences` blob and run end-to-end mining.
///
/// `min_support` is taken literally; the Python reference uses
/// `max(2, len(transactions) / 10)` adaptively — callers wanting that pattern
/// should compute it themselves (transactions count is reported back in the
/// result so a two-pass adaptive run is possible).
pub fn compute_association_analysis<N, E>(
    graph: &Graph<N, E>,
    min_support: usize,
    min_confidence: f64,
) -> AssociationAnalysis
where
    N: ClassNode,
    E: EdgeKind,
{
    let mut transactions: Vec<BTreeSet<String>> = Vec::new();

    for nx in graph.node_indices() {
        let node = &graph[nx];
        let Some(blob) = node.call_sequences() else {
            continue;
        };
        // The blob can be either the {sequences: [...]} object or a JSON string
        // when stored verbatim from GraphML — handle both, matching clone_detection.
        let parsed = if let Some(s) = blob.as_str() {
            serde_json::from_str::<serde_json::Value>(s)
                .ok()
                .and_then(|v| parse_call_sequences(&v))
        } else {
            parse_call_sequences(blob)
        };
        let Some(seqs) = parsed else {
            continue;
        };
        for seq in seqs {
            if seq.calls.len() >= 2 {
                transactions.push(seq.calls.into_iter().collect());
            }
        }
    }

    let itemsets = mine_frequent_itemsets(&transactions, min_support);
    let rules = generate_rules(&itemsets, &transactions, min_confidence);
    let invariants = rules.iter().filter(|r| r.classification == "invariant").count();
    let strong = rules.iter().filter(|r| r.classification == "strong").count();
    let moderate = rules.iter().filter(|r| r.classification == "moderate").count();

    AssociationAnalysis {
        transactions: transactions.len(),
        num_rules: rules.len(),
        invariants,
        strong,
        moderate,
        itemsets: itemsets.len(),
        rules,
    }
}

// ── helpers ──────────────────────────────────────────────────────────────────

fn round4(x: f64) -> f64 {
    (x * 10_000.0).round() / 10_000.0
}

/// All `k`-combinations of `items` in lexicographic index order. Empty when
/// `k > items.len()` or `k == 0`.
fn combinations<T: Clone>(items: &[T], k: usize) -> Vec<Vec<T>> {
    let n = items.len();
    if k == 0 || k > n {
        return Vec::new();
    }
    let mut out: Vec<Vec<T>> = Vec::new();
    let mut idx: Vec<usize> = (0..k).collect();
    loop {
        out.push(idx.iter().map(|&i| items[i].clone()).collect());
        // Find rightmost index we can still advance.
        let mut i = k;
        while i > 0 {
            i -= 1;
            if idx[i] < n - (k - i) {
                idx[i] += 1;
                for j in (i + 1)..k {
                    idx[j] = idx[j - 1] + 1;
                }
                break;
            }
            if i == 0 {
                return out;
            }
        }
        // Termination: when no index can advance, exit. The inner loop sets
        // `i = 0` and either advances [0] or returns above.
        if idx[0] > n - k {
            return out;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use srcgraph_core::{OwnedClassNode, OwnedGraph};
    use petgraph::Graph;
    use serde_json::json;

    fn class(id: &str, seqs: Option<serde_json::Value>) -> OwnedClassNode {
        OwnedClassNode {
            id: id.to_owned(),
            name: id.to_owned(),
            namespace: "test".to_owned(),
            line_count: 10,
            method_count: 1,
            halstead_eta1: 0,
            halstead_eta2: 0,
            halstead_n1: 0,
            halstead_n2: 0,
            method_connectivity: None,
            method_fingerprints: None,
            method_tokens: None,
            call_sequences: seqs,
            cyclomatic_complexity: None,
            path_conditions: None,
            invariants: None,
            error_messages: None,
            magic_numbers: None,
            dead_code: None,
            tenant_branches: None,
            state_transitions: None,
        }
    }

    fn simple_transactions() -> Vec<BTreeSet<String>> {
        vec![
            ["Validate", "Save", "Notify"].into_iter().map(String::from).collect(),
            ["Validate", "Save", "Notify"].into_iter().map(String::from).collect(),
            ["Validate", "Delete", "Notify"].into_iter().map(String::from).collect(),
            ["Validate", "Charge", "Notify"].into_iter().map(String::from).collect(),
        ]
    }

    // ── parse_call_sequences ─────────────────────────────────────────────

    #[test]
    fn parse_basic_sequences() {
        let blob = json!({"sequences": [
            {"method": "CreateOrder", "calls": ["Validate", "Save"]},
            {"method": "DeleteOrder", "calls": ["Validate", "Delete"]},
        ]});
        let seqs = parse_call_sequences(&blob).expect("parses");
        assert_eq!(seqs.len(), 2);
        assert_eq!(seqs[0].method, "CreateOrder");
        assert_eq!(seqs[0].calls, vec!["Validate", "Save"]);
    }

    #[test]
    fn parse_missing_key_yields_none() {
        let blob = json!({"other": []});
        assert!(parse_call_sequences(&blob).is_none());
    }

    #[test]
    fn parse_empty_sequences() {
        let blob = json!({"sequences": []});
        assert_eq!(parse_call_sequences(&blob).unwrap().len(), 0);
    }

    // ── combinations helper ──────────────────────────────────────────────

    #[test]
    fn combinations_basic() {
        let items: Vec<&str> = vec!["a", "b", "c", "d"];
        let cs = combinations(&items, 2);
        assert_eq!(cs.len(), 6);
        assert!(cs.contains(&vec!["a", "b"]));
        assert!(cs.contains(&vec!["c", "d"]));
    }

    #[test]
    fn combinations_k_too_big() {
        let items: Vec<&str> = vec!["a", "b"];
        assert!(combinations(&items, 3).is_empty());
    }

    #[test]
    fn combinations_k_zero() {
        let items: Vec<&str> = vec!["a"];
        assert!(combinations(&items, 0).is_empty());
    }

    // ── mine_frequent_itemsets ───────────────────────────────────────────

    #[test]
    fn mine_finds_singleton_in_all_txns() {
        let txns = simple_transactions();
        let isets = mine_frequent_itemsets(&txns, 2);
        let v = isets.iter().find(|fi| fi.items == vec!["Validate"]);
        assert!(v.is_some(), "expected Validate as frequent 1-itemset");
        assert_eq!(v.unwrap().support, 4);
    }

    #[test]
    fn mine_finds_pair() {
        let txns = simple_transactions();
        let isets = mine_frequent_itemsets(&txns, 2);
        let vn = isets
            .iter()
            .find(|fi| fi.items == vec!["Notify", "Validate"]);
        assert!(vn.is_some(), "expected {{Validate, Notify}} pair");
        assert_eq!(vn.unwrap().support, 4);
    }

    #[test]
    fn mine_respects_min_support() {
        let txns = simple_transactions();
        let high = mine_frequent_itemsets(&txns, 4);
        let low = mine_frequent_itemsets(&txns, 2);
        assert!(high.len() <= low.len());
        // At min_support=4 only items appearing in every txn survive: Validate, Notify, and their pair.
        for fi in &high {
            assert!(fi.support >= 4);
        }
    }

    #[test]
    fn mine_empty_transactions() {
        assert!(mine_frequent_itemsets(&[], 1).is_empty());
    }

    #[test]
    fn mine_sorted_by_support_desc() {
        let txns = simple_transactions();
        let isets = mine_frequent_itemsets(&txns, 2);
        let supports: Vec<usize> = isets.iter().map(|fi| fi.support).collect();
        let mut sorted = supports.clone();
        sorted.sort_by(|a, b| b.cmp(a));
        assert_eq!(supports, sorted);
    }

    // ── generate_rules ───────────────────────────────────────────────────

    #[test]
    fn generate_rules_produces_high_confidence_pair() {
        let txns = simple_transactions();
        let isets = mine_frequent_itemsets(&txns, 2);
        let rules = generate_rules(&isets, &txns, 0.5);
        assert!(!rules.is_empty());
        // Every txn contains both Validate and Notify, so Validate → Notify is conf 1.0.
        let vn = rules.iter().find(|r| {
            r.antecedent == vec!["Validate"] && r.consequent == vec!["Notify"]
        });
        assert!(vn.is_some());
        assert!(vn.unwrap().confidence >= 0.99);
    }

    #[test]
    fn generate_rules_sorted_by_confidence_desc() {
        let txns = simple_transactions();
        let isets = mine_frequent_itemsets(&txns, 2);
        let rules = generate_rules(&isets, &txns, 0.5);
        for w in rules.windows(2) {
            assert!(w[0].confidence >= w[1].confidence);
        }
    }

    #[test]
    fn generate_rules_classification_applied() {
        let txns = simple_transactions();
        let isets = mine_frequent_itemsets(&txns, 2);
        let rules = generate_rules(&isets, &txns, 0.5);
        for r in &rules {
            assert!(["invariant", "strong", "moderate"].contains(&r.classification.as_str()));
        }
    }

    #[test]
    fn generate_rules_empty_inputs() {
        assert!(generate_rules(&[], &[], 0.5).is_empty());
        let txns = simple_transactions();
        assert!(generate_rules(&[], &txns, 0.5).is_empty());
    }

    #[test]
    fn classify_thresholds() {
        assert_eq!(classify_rule(1.0), "invariant");
        assert_eq!(classify_rule(0.99), "invariant");
        assert_eq!(classify_rule(0.85), "strong");
        assert_eq!(classify_rule(0.5), "moderate");
        assert_eq!(classify_rule(0.49), "moderate");
    }

    // ── compute_association_analysis ─────────────────────────────────────

    #[test]
    fn compute_walks_graph_and_counts() {
        let blob = json!({"sequences": [
            {"method": "CreateOrder", "calls": ["Validate", "Save", "Notify"]},
            {"method": "UpdateOrder", "calls": ["Validate", "Save", "Notify"]},
            {"method": "DeleteOrder", "calls": ["Validate", "Delete", "Notify"]},
            {"method": "Payment", "calls": ["Validate", "Charge", "Notify", "Log"]},
        ]});
        let mut g: OwnedGraph = Graph::new();
        g.add_node(class("OrderSvc", Some(blob)));
        // Node without sequences — silently skipped.
        g.add_node(class("Plain", None));

        let r = compute_association_analysis(&g, 2, 0.5);
        assert_eq!(r.transactions, 4);
        assert!(r.num_rules > 0);
        assert!(r.itemsets > 0);
        // Validate → Notify is an invariant (every txn has both).
        assert!(r.invariants >= 1);
    }

    #[test]
    fn compute_accepts_string_encoded_blob() {
        let inner = json!({"sequences": [
            {"method": "A", "calls": ["X", "Y"]},
            {"method": "B", "calls": ["X", "Y"]},
        ]});
        let mut g: OwnedGraph = Graph::new();
        g.add_node(class("A", Some(serde_json::Value::String(inner.to_string()))));
        let r = compute_association_analysis(&g, 2, 0.5);
        assert_eq!(r.transactions, 2);
    }

    #[test]
    fn compute_empty_graph_zero_rules() {
        let g: OwnedGraph = Graph::new();
        let r = compute_association_analysis(&g, 2, 0.5);
        assert_eq!(r.transactions, 0);
        assert_eq!(r.num_rules, 0);
        assert_eq!(r.itemsets, 0);
    }

    #[test]
    fn compute_skips_short_sequences() {
        // calls of length < 2 are skipped (matches Python).
        let blob = json!({"sequences": [
            {"method": "A", "calls": ["only"]},
            {"method": "B", "calls": []},
        ]});
        let mut g: OwnedGraph = Graph::new();
        g.add_node(class("X", Some(blob)));
        let r = compute_association_analysis(&g, 1, 0.5);
        assert_eq!(r.transactions, 0);
    }
}