erpc_analysis/algorithms/
classification.rs

Help
use log::info;
use std::sync::Arc;

use crate::db_trait::{AnalysisDatabase, AnalysisError};
use crate::models::partitions::{
    ConnectedComponent, PartitionClassificationResult,
};

/// Analyzer for classifying network partitions by various attributes
pub struct PartitionClassifier {
    db_client: Arc<dyn AnalysisDatabase>,
}

impl PartitionClassifier {
    /// Create a new PartitionClassifier with the given database client
    pub fn new(db_client: Arc<dyn AnalysisDatabase>) -> Self {
        Self { db_client }
    }

    /// Classify connected components by geographic location (country)
    pub async fn classify_by_geography(
        &self,
        components: &[ConnectedComponent],
    ) -> Result<PartitionClassificationResult, AnalysisError> {
        info!("=== Starting Geographic Classification Analysis ===");

        let result = self
            .db_client
            .classify_components_by_geography(components)
            .await?;

        info!("=== Geographic Classification Complete ===");

        info!(
            "Found {} geographic groups with {:.1}% coverage",
            result.metrics.total_groups,
            result.metrics.classification_coverage
        );

        Ok(result)
    }

    /// Classify connected components by ASN
    pub async fn classify_by_asn(
        &self,
        components: &[ConnectedComponent],
    ) -> Result<PartitionClassificationResult, AnalysisError> {
        info!("=== Starting ASN Classification Analysis ===");

        let result = self
            .db_client
            .classify_components_by_asn(components)
            .await?;

        info!("=== ASN Classification Complete ===");

        info!(
            "Found {} ASN groups with {:.1}% coverage",
            result.metrics.total_groups,
            result.metrics.classification_coverage
        );

        Ok(result)
    }

    /// Classify connected components by relay family relationships
    pub async fn classify_by_family(
        &self,
        components: &[ConnectedComponent],
    ) -> Result<PartitionClassificationResult, AnalysisError> {
        info!("=== Starting Family Classification Analysis ===");

        let result = self
            .db_client
            .classify_components_by_family(components)
            .await?;

        info!("=== Family Classification Complete ===");

        info!(
            "Found {} family groups with {:.1}% coverage",
            result.metrics.total_groups,
            result.metrics.classification_coverage
        );

        Ok(result)
    }

    /// Display detailed geographic classification results
    pub fn display_geographic_classification(
        &self,
        result: &PartitionClassificationResult,
        config: &crate::config::AnalysisSettings,
    ) -> Result<(), Box<dyn std::error::Error>> {
        info!("Geographic Classification Analysis:");
        info!("Total Countries: {}", result.metrics.total_groups);
        info!(
            "Countries with Partitions: {}",
            result.metrics.groups_with_partitions
        );
        info!(
            "Classification Coverage: {:.1}%",
            result.metrics.classification_coverage
        );
        info!(
            "Largest Country Group: {} relays",
            result.metrics.largest_group_size
        );
        info!(
            "Average Country Group Size: {:.1} relays",
            result.metrics.average_group_size
        );
        info!("Diversity Score: {:.3}", result.metrics.diversity_score);
        info!(
            "Partition Correlation: {:.2}",
            result.metrics.partition_correlation
        );

        // Show top countries by size
        let mut sorted_groups = result.groups.clone();
        sorted_groups.sort_by(|a, b| {
            b.relay_fingerprints.len().cmp(&a.relay_fingerprints.len())
        });

        info!("Top Countries by Relay Count:");
        for (i, group) in sorted_groups
            .iter()
            .take(config.max_display_components)
            .enumerate()
        {
            let partition_info = if group.component_mapping.len() > 1 {
                format!(
                    " (across {} partitions)",
                    group.component_mapping.len()
                )
            } else {
                " (single partition)".to_string()
            };

            info!(
                "{}. {}: {} relays, isolation: {:.1}%{}",
                i + 1,
                group.identifier,
                group.relay_fingerprints.len(),
                group.isolation_score,
                partition_info
            );
        }

        if !result.unclassified_relays.is_empty() {
            // Calculate total relays from classified + unclassified
            let total_classified = result
                .groups
                .iter()
                .map(|g| g.relay_fingerprints.len())
                .sum::<usize>();
            let total_relays =
                total_classified + result.unclassified_relays.len();
            let unclassified_percent = if total_relays > 0 {
                (result.unclassified_relays.len() as f64 / total_relays as f64)
                    * 100.0
            } else {
                0.0
            };

            info!(
                "Unclassified relays: {} ({:.3}% of total)",
                result.unclassified_relays.len(),
                unclassified_percent
            );
        }

        Ok(())
    }

    /// Display detailed ASN classification results
    pub fn display_asn_classification(
        &self,
        result: &PartitionClassificationResult,
        config: &crate::config::AnalysisSettings,
    ) -> Result<(), Box<dyn std::error::Error>> {
        info!("ASN Classification Analysis:");
        info!("Total ASNs: {}", result.metrics.total_groups);
        info!(
            "ASNs with Partitions: {}",
            result.metrics.groups_with_partitions
        );
        info!(
            "Classification Coverage: {:.1}%",
            result.metrics.classification_coverage
        );
        info!(
            "Largest ASN Group: {} relays",
            result.metrics.largest_group_size
        );
        info!(
            "Average ASN Group Size: {:.1} relays",
            result.metrics.average_group_size
        );
        info!("Diversity Score: {:.3}", result.metrics.diversity_score);
        info!(
            "Partition Correlation: {:.2}",
            result.metrics.partition_correlation
        );

        // Show top ASNs by size
        let mut sorted_groups = result.groups.clone();
        sorted_groups.sort_by(|a, b| {
            b.relay_fingerprints.len().cmp(&a.relay_fingerprints.len())
        });

        info!("Top ASNs by Relay Count:");
        for (i, group) in sorted_groups
            .iter()
            .take(config.max_display_components)
            .enumerate()
        {
            let partition_info = if group.component_mapping.len() > 1 {
                format!(
                    " (across {} partitions)",
                    group.component_mapping.len()
                )
            } else {
                " (single partition)".to_string()
            };

            info!(
                "{}. AS{}: {} relays, isolation: {:.1}%{}",
                i + 1,
                group.identifier,
                group.relay_fingerprints.len(),
                group.isolation_score,
                partition_info
            );
        }

        if !result.unclassified_relays.is_empty() {
            // Calculate total relays from classified + unclassified
            let total_classified = result
                .groups
                .iter()
                .map(|g| g.relay_fingerprints.len())
                .sum::<usize>();
            let total_relays =
                total_classified + result.unclassified_relays.len();
            let unclassified_percent = if total_relays > 0 {
                (result.unclassified_relays.len() as f64 / total_relays as f64)
                    * 100.0
            } else {
                0.0
            };

            info!(
                "Unclassified relays: {} ({:.3}% of total)",
                result.unclassified_relays.len(),
                unclassified_percent
            );
        }

        Ok(())
    }

    /// Display detailed family classification results
    pub fn display_family_classification(
        &self,
        result: &PartitionClassificationResult,
        config: &crate::config::AnalysisSettings,
    ) -> Result<(), Box<dyn std::error::Error>> {
        info!("Family Classification Analysis:");
        info!("Total Families: {}", result.metrics.total_groups);
        info!(
            "Families with Partitions: {}",
            result.metrics.groups_with_partitions
        );
        info!(
            "Classification Coverage: {:.1}%",
            result.metrics.classification_coverage
        );
        info!(
            "Largest Family: {} relays",
            result.metrics.largest_group_size
        );
        info!(
            "Average Family Size: {:.1} relays",
            result.metrics.average_group_size
        );
        info!("Diversity Score: {:.3}", result.metrics.diversity_score);
        info!(
            "Partition Correlation: {:.2}",
            result.metrics.partition_correlation
        );

        // Show top families by size
        let mut sorted_groups = result.groups.clone();
        sorted_groups.sort_by(|a, b| {
            b.relay_fingerprints.len().cmp(&a.relay_fingerprints.len())
        });

        info!("Top Families by Relay Count:");
        for (i, group) in sorted_groups
            .iter()
            .take(config.max_display_components)
            .enumerate()
        {
            let partition_info = if group.component_mapping.len() > 1 {
                format!(
                    " (across {} partitions)",
                    group.component_mapping.len()
                )
            } else {
                " (single partition)".to_string()
            };

            info!(
                "{}. {}: {} relays, isolation: {:.1}%{}",
                i + 1,
                if group.identifier.len() > 12 {
                    format!("{}...", &group.identifier[0..12])
                } else {
                    group.identifier.clone()
                },
                group.relay_fingerprints.len(),
                group.isolation_score,
                partition_info
            );
        }

        if !result.unclassified_relays.is_empty() {
            // Calculate total relays from classified + unclassified
            let total_classified = result
                .groups
                .iter()
                .map(|g| g.relay_fingerprints.len())
                .sum::<usize>();
            let total_relays =
                total_classified + result.unclassified_relays.len();
            let unclassified_percent = if total_relays > 0 {
                (result.unclassified_relays.len() as f64 / total_relays as f64)
                    * 100.0
            } else {
                0.0
            };

            info!(
                "Non-family relays: {} ({:.3}% of total)",
                result.unclassified_relays.len(),
                unclassified_percent
            );
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::db_trait::mock::MockDatabase;
    use crate::models::partitions::ConnectedComponent;

    /// Test geographic classification
    #[tokio::test]
    async fn test_geographic_classification() {
        let db = Arc::new(MockDatabase::new());
        let classifier = PartitionClassifier::new(db);

        let components = vec![ConnectedComponent {
            component_id: 1,
            relay_fingerprints: vec![
                "RELAY_A".to_string(), // Hash % 4 = 0 -> US
                "RELAY_B".to_string(), // Hash % 4 = 1 -> DE
                "RELAY_C".to_string(), // Hash % 4 = 2 -> FR
                "RELAY_D".to_string(), // Hash % 4 = 3 -> NL
            ],
            size: 4,
        }];

        let result =
            classifier.classify_by_geography(&components).await.unwrap();

        assert!(!result.groups.is_empty());
        assert_eq!(result.metrics.classification_coverage, 100.0);

        // Verify component mappings use actual component IDs
        for group in &result.groups {
            assert!(group.component_mapping.contains_key(&1));
        }
    }

    /// Test ASN classification
    #[tokio::test]
    async fn test_asn_classification() {
        let db = Arc::new(MockDatabase::new());
        let classifier = PartitionClassifier::new(db);

        let components = vec![ConnectedComponent {
            component_id: 42,
            relay_fingerprints: vec![
                "RELAY_X".to_string(), // Hash % 3 = 0 -> 13335
                "RELAY_Y".to_string(), // Hash % 3 = 1 -> 210558
                "RELAY_Z".to_string(), // Hash % 3 = 2 -> 64512
            ],
            size: 3,
        }];

        let result = classifier.classify_by_asn(&components).await.unwrap();

        assert!(!result.groups.is_empty());
        assert_eq!(result.metrics.classification_coverage, 100.0);

        // Verify component mappings use actual component IDs (42)
        for group in &result.groups {
            assert!(group.component_mapping.contains_key(&42));
        }
    }

    /// Test isolation score calculation
    #[tokio::test]
    async fn test_isolation_score() {
        let db = Arc::new(MockDatabase::new());
        let classifier = PartitionClassifier::new(db);

        // Empty case
        let result = classifier.classify_by_geography(&[]).await.unwrap();
        assert_eq!(result.groups.len(), 0);
        assert_eq!(result.metrics.classification_coverage, 0.0);

        // Test with relays that create both fragmented and cohesive groups
        let components = vec![
            // Component 0: All relays hash to same country (perfect cohesion)
            ConnectedComponent {
                component_id: 10,
                relay_fingerprints: vec!["AAAA".to_string()], // Single relay
                size: 1,
            },
            // Component 1: Multiple relays
            ConnectedComponent {
                component_id: 20,
                relay_fingerprints: vec![
                    "BBBB".to_string(),
                    "CCCC".to_string(),
                ],
                size: 2,
            },
        ];

        let result =
            classifier.classify_by_geography(&components).await.unwrap();

        // Verify that isolation scores are calculated correctly
        // Groups with all relays in one component should have 0% isolation
        // Groups split across components should have >0% isolation
        assert!(!result.groups.is_empty());

        for group in &result.groups {
            // Isolation score should be between 0 and 100
            assert!(group.isolation_score >= 0.0);
            assert!(group.isolation_score <= 100.0);

            // Component mappings should use actual component IDs
            let component_ids: Vec<i64> =
                group.component_mapping.keys().copied().collect();
            assert!(component_ids.iter().any(|&id| id == 10 || id == 20));
        }
    }

    /// Test family classification
    #[tokio::test]
    async fn test_family_classification() {
        let db = Arc::new(MockDatabase::new());
        let classifier = PartitionClassifier::new(db);

        let components = vec![ConnectedComponent {
            component_id: 5,
            relay_fingerprints: vec![
                "RELAY_FAM1".to_string(), // Some get families, some don't
                "RELAY_FAM2".to_string(),
                "RELAY_FAM3".to_string(),
                "RELAY_LONE".to_string(),
            ],
            size: 4,
        }];

        let result = classifier.classify_by_family(&components).await.unwrap();

        // Family classification may have some unclassified relays
        let total_classified: usize = result
            .groups
            .iter()
            .map(|g| g.relay_fingerprints.len())
            .sum();
        let total_unclassified = result.unclassified_relays.len();

        assert_eq!(total_classified + total_unclassified, 4);

        // Verify component mappings use actual component ID (5)
        for group in &result.groups {
            assert!(group.component_mapping.contains_key(&5));
        }

        let expected_coverage = (total_classified as f64 / 4.0) * 100.0;
        assert!(
            (result.metrics.classification_coverage - expected_coverage).abs()
                < 0.1
        );
    }

    /// Test error handling for classification failures
    #[tokio::test]
    async fn test_classification_error_handling() {
        let components = vec![ConnectedComponent {
            component_id: 0,
            relay_fingerprints: vec!["FAIL_RELAY".to_string()],
            size: 1,
        }];

        let db = Arc::new(
            MockDatabase::new().fail_on("classify_components_by_geography"),
        );
        let classifier = PartitionClassifier::new(db);

        let result = classifier.classify_by_geography(&components).await;
        assert!(result.is_err(), "Should fail when database operation fails");
    }

    /// Test display methods don't panic with valid data
    #[tokio::test]
    async fn test_display_methods() {
        let components = vec![ConnectedComponent {
            component_id: 0,
            relay_fingerprints: vec!["TEST_RELAY".to_string()],
            size: 1,
        }];

        let db = Arc::new(MockDatabase::new());
        let classifier = PartitionClassifier::new(db);

        let geo_result =
            classifier.classify_by_geography(&components).await.unwrap();
        let asn_result =
            classifier.classify_by_asn(&components).await.unwrap();
        let family_result =
            classifier.classify_by_family(&components).await.unwrap();

        // These should not panic
        let config = crate::config::AnalysisSettings::default();
        classifier
            .display_geographic_classification(&geo_result, &config)
            .unwrap();
        classifier
            .display_asn_classification(&asn_result, &config)
            .unwrap();
        classifier
            .display_family_classification(&family_result, &config)
            .unwrap();
    }
}
erpc_analysis/algorithms/classification.rs

erpc_analysis/algorithms/
classification.rs
