Skip to main content

trustformers_debug/
ai_code_analyzer.rs

1//! Rule-Based Static Code Analysis for Model Debugging
2//!
3//! This module identifies potential issues in neural-network training/inference
4//! source code, suggests optimizations, and flags common security anti-patterns.
5//!
6//! Despite the type name [`AICodeAnalyzer`] (kept for API stability), this is a
7//! **deterministic, rule-based static analyzer**: every finding comes from a
8//! fixed substring/pattern check against the literal source text (see
9//! `AICodeAnalyzer::perform_deep_analysis` and its sibling
10//! `detect_*`/`generate_*` methods), not from any trained model or live
11//! inference call. There is no network access, no model weights, and no
12//! non-deterministic behavior -- calling it twice on the same input always
13//! produces the same result. Each rule carries a fixed `confidence: f64`
14//! assigned when the rule was written, reflecting that rule's own
15//! specificity (how often that particular substring pattern is a true
16//! positive in practice) -- it is **not** a live-computed statistical
17//! probability from any model, and must not be read as one.
18// reason: debug/profiling scaffolding — structs are constructed and their fields/methods
19// are retained for the data model, serialization completeness, and future consumers that
20// do not yet read every member. Consolidated from many item-level #[allow(dead_code)].
21#![allow(dead_code)]
22
23use anyhow::Result;
24use serde::{Deserialize, Serialize};
25use std::collections::HashMap;
26use tokio::time::{Duration, Instant};
27use tracing::{debug, info};
28
29/// Rule-based static code analysis engine for model debugging.
30///
31/// See the module-level docs: this is a deterministic pattern matcher, not
32/// an AI model.
33#[derive(Debug)]
34pub struct AICodeAnalyzer {
35    config: AIAnalysisConfig,
36    analysis_cache: HashMap<String, CachedAnalysis>,
37    pattern_database: ModelPatternDatabase,
38    performance_monitor: AnalysisPerformanceMonitor,
39}
40
41/// Configuration for AI code analysis
42#[derive(Debug, Clone, Serialize, Deserialize)]
43pub struct AIAnalysisConfig {
44    /// Enable deep code analysis using AI models
45    pub enable_deep_analysis: bool,
46    /// Enable pattern recognition for common issues
47    pub enable_pattern_recognition: bool,
48    /// Enable optimization suggestions
49    pub enable_optimization_suggestions: bool,
50    /// Enable vulnerability detection
51    pub enable_vulnerability_detection: bool,
52    /// Enable performance prediction
53    pub enable_performance_prediction: bool,
54    /// Maximum analysis time per code segment (seconds)
55    pub max_analysis_time_secs: u64,
56    /// Confidence threshold for suggestions (0.0-1.0)
57    pub confidence_threshold: f64,
58    /// Enable caching of analysis results
59    pub enable_caching: bool,
60    /// Cache expiration time (hours)
61    pub cache_expiration_hours: u64,
62}
63
64impl Default for AIAnalysisConfig {
65    fn default() -> Self {
66        Self {
67            enable_deep_analysis: true,
68            enable_pattern_recognition: true,
69            enable_optimization_suggestions: true,
70            enable_vulnerability_detection: true,
71            enable_performance_prediction: true,
72            max_analysis_time_secs: 30,
73            confidence_threshold: 0.75,
74            enable_caching: true,
75            cache_expiration_hours: 24,
76        }
77    }
78}
79
80/// Cached analysis result
81#[derive(Debug, Clone, Serialize, Deserialize)]
82struct CachedAnalysis {
83    result: CodeAnalysisResult,
84    timestamp: std::time::SystemTime,
85    code_hash: String,
86}
87
88/// Performance monitor for analysis operations
89#[derive(Debug)]
90struct AnalysisPerformanceMonitor {
91    analysis_count: u64,
92    total_analysis_time: Duration,
93    cache_hits: u64,
94    cache_misses: u64,
95}
96
97impl AnalysisPerformanceMonitor {
98    fn new() -> Self {
99        Self {
100            analysis_count: 0,
101            total_analysis_time: Duration::from_secs(0),
102            cache_hits: 0,
103            cache_misses: 0,
104        }
105    }
106
107    fn record_analysis(&mut self, duration: Duration, cache_hit: bool) {
108        self.analysis_count += 1;
109        self.total_analysis_time += duration;
110        if cache_hit {
111            self.cache_hits += 1;
112        } else {
113            self.cache_misses += 1;
114        }
115    }
116
117    fn average_analysis_time(&self) -> Duration {
118        if self.analysis_count > 0 {
119            self.total_analysis_time / self.analysis_count as u32
120        } else {
121            Duration::from_secs(0)
122        }
123    }
124
125    fn cache_hit_rate(&self) -> f64 {
126        let total = self.cache_hits + self.cache_misses;
127        if total > 0 {
128            self.cache_hits as f64 / total as f64
129        } else {
130            0.0
131        }
132    }
133}
134
135impl AICodeAnalyzer {
136    /// Create a new AI code analyzer
137    pub fn new(config: AIAnalysisConfig) -> Self {
138        Self {
139            config,
140            analysis_cache: HashMap::new(),
141            pattern_database: ModelPatternDatabase::new(),
142            performance_monitor: AnalysisPerformanceMonitor::new(),
143        }
144    }
145
146    /// Analyze model code for potential issues and optimizations
147    pub async fn analyze_model_code(
148        &mut self,
149        code: &str,
150        context: ModelContext,
151    ) -> Result<CodeAnalysisResult> {
152        let start_time = Instant::now();
153        let code_hash = self.compute_code_hash(code);
154
155        // Check cache first
156        if self.config.enable_caching {
157            if let Some(cached) = self.get_cached_analysis(&code_hash) {
158                let result = cached.result.clone();
159                self.performance_monitor.record_analysis(start_time.elapsed(), true);
160                return Ok(result);
161            }
162        }
163
164        info!(
165            "Starting AI code analysis for {} lines of code",
166            code.lines().count()
167        );
168
169        let mut result = CodeAnalysisResult::new();
170
171        // Pattern recognition analysis
172        if self.config.enable_pattern_recognition {
173            let patterns = self.detect_code_patterns(code, &context).await?;
174            result.detected_patterns = patterns;
175        }
176
177        // Deep AI analysis
178        if self.config.enable_deep_analysis {
179            let issues = self.perform_deep_analysis(code, &context).await?;
180            result.identified_issues = issues;
181        }
182
183        // Optimization suggestions
184        if self.config.enable_optimization_suggestions {
185            let optimizations = self.generate_optimization_suggestions(code, &context).await?;
186            result.optimization_suggestions = optimizations;
187        }
188
189        // Vulnerability detection
190        if self.config.enable_vulnerability_detection {
191            let vulnerabilities = self.detect_vulnerabilities(code, &context).await?;
192            result.security_issues = vulnerabilities;
193        }
194
195        // Performance prediction
196        if self.config.enable_performance_prediction {
197            let predictions = self.predict_performance_characteristics(code, &context).await?;
198            result.performance_predictions = predictions;
199        }
200
201        // Calculate overall quality score
202        result.quality_score = self.calculate_quality_score(&result);
203        result.analysis_metadata = AnalysisMetadata {
204            analysis_duration: start_time.elapsed(),
205            confidence_score: self.calculate_confidence_score(&result),
206            analyzer_version: "1.0.0".to_string(),
207            timestamp: std::time::SystemTime::now(),
208        };
209
210        // Cache the result
211        if self.config.enable_caching {
212            self.cache_analysis(code_hash, &result);
213        }
214
215        self.performance_monitor.record_analysis(start_time.elapsed(), false);
216
217        info!(
218            "AI code analysis completed in {:?} with quality score: {:.2}",
219            start_time.elapsed(),
220            result.quality_score
221        );
222
223        Ok(result)
224    }
225
226    /// Analyze tensor operations for optimization opportunities
227    pub async fn analyze_tensor_operations(
228        &self,
229        operations: &[TensorOperation],
230    ) -> Result<TensorOptimizationReport> {
231        debug!("Analyzing {} tensor operations", operations.len());
232
233        let mut report = TensorOptimizationReport::new();
234
235        // Analyze operation patterns
236        report.fusion_opportunities = self.detect_fusion_opportunities(operations).await?;
237        report.memory_optimizations = self.detect_memory_optimizations(operations).await?;
238        report.parallelization_opportunities =
239            self.detect_parallelization_opportunities(operations).await?;
240        report.redundant_operations = self.detect_redundant_operations(operations).await?;
241
242        // Calculate potential speedup
243        report.estimated_speedup = self.estimate_optimization_speedup(&report);
244        report.estimated_memory_savings = self.estimate_memory_savings(&report);
245
246        Ok(report)
247    }
248
249    /// Perform automated debugging assistance
250    pub async fn automated_debugging_assistance(
251        &self,
252        error_context: &ErrorContext,
253    ) -> Result<DebuggingAssistance> {
254        info!(
255            "Providing automated debugging assistance for error: {}",
256            error_context.error_type
257        );
258
259        let mut assistance = DebuggingAssistance::new();
260
261        // Analyze error patterns
262        assistance.probable_causes = self.analyze_error_patterns(error_context).await?;
263        assistance.suggested_fixes = self.generate_suggested_fixes(error_context).await?;
264        assistance.debugging_steps = self.generate_debugging_steps(error_context).await?;
265        assistance.related_documentation = self.find_related_documentation(error_context).await?;
266
267        // Generate confidence score
268        assistance.confidence_score = self.calculate_debugging_confidence(&assistance);
269
270        Ok(assistance)
271    }
272
273    /// Get analysis performance metrics
274    pub fn get_performance_metrics(&self) -> AnalysisPerformanceMetrics {
275        AnalysisPerformanceMetrics {
276            total_analyses: self.performance_monitor.analysis_count,
277            average_analysis_time: self.performance_monitor.average_analysis_time(),
278            cache_hit_rate: self.performance_monitor.cache_hit_rate(),
279            cached_results: self.analysis_cache.len(),
280        }
281    }
282
283    // Private helper methods
284
285    async fn detect_code_patterns(
286        &self,
287        code: &str,
288        context: &ModelContext,
289    ) -> Result<Vec<DetectedPattern>> {
290        debug!("Detecting code patterns");
291
292        let mut patterns = Vec::new();
293
294        // Common anti-patterns in neural networks
295        if code.contains("torch.cuda.empty_cache()") && context.model_type == ModelType::Production
296        {
297            patterns.push(DetectedPattern {
298                pattern_type: PatternType::AntiPattern,
299                name: "Frequent CUDA Cache Clearing".to_string(),
300                description: "Frequent CUDA cache clearing can hurt performance".to_string(),
301                severity: Severity::Medium,
302                confidence: 0.85,
303                recommendations: vec![
304                    "Consider using gradient accumulation instead".to_string(),
305                    "Review memory management strategy".to_string(),
306                ],
307            });
308        }
309
310        // Gradient explosion patterns
311        if code.contains("grad_norm") && code.contains("clip") {
312            patterns.push(DetectedPattern {
313                pattern_type: PatternType::GoodPattern,
314                name: "Gradient Clipping".to_string(),
315                description: "Proper gradient clipping implementation detected".to_string(),
316                severity: Severity::Info,
317                confidence: 0.9,
318                recommendations: vec!["Consider adaptive gradient clipping".to_string()],
319            });
320        }
321
322        // Memory inefficient patterns
323        if code.contains("detach()") && code.contains("requires_grad") {
324            patterns.push(DetectedPattern {
325                pattern_type: PatternType::OptimizationOpportunity,
326                name: "Gradient Computation Inefficiency".to_string(),
327                description: "Potential inefficient gradient computation detected".to_string(),
328                severity: Severity::Medium,
329                confidence: 0.75,
330                recommendations: vec![
331                    "Consider using torch.no_grad() context".to_string(),
332                    "Review gradient requirements".to_string(),
333                ],
334            });
335        }
336
337        Ok(patterns)
338    }
339
340    /// Rule-based issue detection over the literal source text. See the
341    /// module docs: this is real pattern matching, not a live model call,
342    /// so there is no analysis latency to simulate here.
343    async fn perform_deep_analysis(
344        &self,
345        code: &str,
346        _context: &ModelContext,
347    ) -> Result<Vec<IdentifiedIssue>> {
348        debug!("Performing rule-based deep code analysis");
349
350        let mut issues = Vec::new();
351
352        // Check for numerical stability issues
353        if code.contains("log") && !code.contains("log1p") && code.contains("softmax") {
354            issues.push(IdentifiedIssue {
355                issue_type: IssueType::NumericalStability,
356                title: "Potential Numerical Instability in Log-Softmax".to_string(),
357                description: "Using log(softmax(x)) can cause numerical instability. Consider using log_softmax directly.".to_string(),
358                severity: Severity::High,
359                confidence: 0.88,
360                suggested_fix: "Replace log(softmax(x)) with log_softmax(x)".to_string(),
361                code_location: locate_in_code(code, "softmax"),
362            });
363        }
364
365        // Check for inefficient attention implementations
366        if code.contains("attention") && code.contains("matmul") && !code.contains("flash") {
367            issues.push(IdentifiedIssue {
368                issue_type: IssueType::Performance,
369                title: "Inefficient Attention Implementation".to_string(),
370                description:
371                    "Standard attention implementation may be inefficient for large sequences."
372                        .to_string(),
373                severity: Severity::Medium,
374                confidence: 0.75,
375                suggested_fix:
376                    "Consider using Flash Attention or other optimized attention mechanisms"
377                        .to_string(),
378                code_location: locate_in_code(code, "attention"),
379            });
380        }
381
382        // Check for memory leaks
383        if code.contains("accumulate") && !code.contains("zero_grad") {
384            issues.push(IdentifiedIssue {
385                issue_type: IssueType::MemoryLeak,
386                title: "Potential Gradient Accumulation Memory Leak".to_string(),
387                description: "Gradient accumulation without zero_grad() can cause memory leaks."
388                    .to_string(),
389                severity: Severity::High,
390                confidence: 0.82,
391                suggested_fix: "Ensure optimizer.zero_grad() is called appropriately".to_string(),
392                code_location: locate_in_code(code, "accumulate"),
393            });
394        }
395
396        Ok(issues)
397    }
398
399    async fn generate_optimization_suggestions(
400        &self,
401        code: &str,
402        context: &ModelContext,
403    ) -> Result<Vec<OptimizationSuggestion>> {
404        debug!("Generating optimization suggestions");
405
406        let mut suggestions = Vec::new();
407
408        // Suggest mixed precision training
409        if context.model_type == ModelType::Training && !code.contains("autocast") {
410            suggestions.push(OptimizationSuggestion {
411                optimization_type: OptimizationType::MixedPrecision,
412                title: "Enable Mixed Precision Training".to_string(),
413                description: "Mixed precision training can significantly speed up training and reduce memory usage.".to_string(),
414                potential_speedup: 1.5,
415                memory_savings: 0.4,
416                implementation_effort: ImplementationEffort::Low,
417                confidence: 0.9,
418                code_example: Some("with torch.autocast(device_type='cuda', dtype=torch.float16):".to_string()),
419            });
420        }
421
422        // Suggest model compilation
423        if context.model_type == ModelType::Production && !code.contains("compile") {
424            suggestions.push(OptimizationSuggestion {
425                optimization_type: OptimizationType::ModelCompilation,
426                title: "Enable Model Compilation".to_string(),
427                description: "Model compilation can provide significant inference speedups."
428                    .to_string(),
429                potential_speedup: 2.0,
430                memory_savings: 0.0,
431                implementation_effort: ImplementationEffort::Low,
432                confidence: 0.85,
433                code_example: Some("model = torch.compile(model)".to_string()),
434            });
435        }
436
437        // Suggest gradient checkpointing for large models
438        if context.model_size > 1_000_000_000 && !code.contains("checkpoint") {
439            suggestions.push(OptimizationSuggestion {
440                optimization_type: OptimizationType::MemoryOptimization,
441                title: "Enable Gradient Checkpointing".to_string(),
442                description:
443                    "Gradient checkpointing can significantly reduce memory usage for large models."
444                        .to_string(),
445                potential_speedup: 0.8, // Slight speed penalty
446                memory_savings: 0.6,
447                implementation_effort: ImplementationEffort::Medium,
448                confidence: 0.88,
449                code_example: Some("torch.utils.checkpoint.checkpoint(layer, x)".to_string()),
450            });
451        }
452
453        Ok(suggestions)
454    }
455
456    async fn detect_vulnerabilities(
457        &self,
458        code: &str,
459        context: &ModelContext,
460    ) -> Result<Vec<SecurityIssue>> {
461        debug!("Detecting security vulnerabilities");
462
463        let mut vulnerabilities = Vec::new();
464
465        // Check for unsafe pickle loading
466        if code.contains("pickle.load") && !code.contains("safe_load") {
467            vulnerabilities.push(SecurityIssue {
468                vulnerability_type: VulnerabilityType::CodeExecution,
469                title: "Unsafe Pickle Loading".to_string(),
470                description:
471                    "Loading pickle files can execute arbitrary code. Use safe alternatives."
472                        .to_string(),
473                severity: Severity::Critical,
474                confidence: 0.95,
475                mitigation: "Use torch.load with weights_only=True or safetensors".to_string(),
476                cve_references: vec!["CWE-502".to_string()],
477            });
478        }
479
480        // Check for model parameter exposure
481        if code.contains("state_dict")
482            && code.contains("save")
483            && context.model_type == ModelType::Production
484        {
485            vulnerabilities.push(SecurityIssue {
486                vulnerability_type: VulnerabilityType::DataExposure,
487                title: "Potential Model Parameter Exposure".to_string(),
488                description: "Saving full model state may expose sensitive parameters.".to_string(),
489                severity: Severity::Medium,
490                confidence: 0.7,
491                mitigation: "Consider differential privacy or parameter encryption".to_string(),
492                cve_references: vec![],
493            });
494        }
495
496        // Check for input validation
497        if code.contains("input") && !code.contains("validate") && !code.contains("sanitize") {
498            vulnerabilities.push(SecurityIssue {
499                vulnerability_type: VulnerabilityType::InputValidation,
500                title: "Missing Input Validation".to_string(),
501                description: "Input validation is important for preventing adversarial attacks."
502                    .to_string(),
503                severity: Severity::Medium,
504                confidence: 0.65,
505                mitigation: "Implement input validation and sanitization".to_string(),
506                cve_references: vec![],
507            });
508        }
509
510        Ok(vulnerabilities)
511    }
512
513    /// Rule-based performance prediction. Like [`Self::perform_deep_analysis`],
514    /// this is real (if heuristic) computation over `code`/`context`, not a
515    /// model call, so there is no latency to simulate.
516    async fn predict_performance_characteristics(
517        &self,
518        code: &str,
519        context: &ModelContext,
520    ) -> Result<PerformancePredictions> {
521        debug!("Predicting performance characteristics");
522
523        let mut predictions = PerformancePredictions::new();
524
525        // Predict memory usage based on model architecture
526        predictions.estimated_memory_usage = self.estimate_memory_usage(code, context);
527        predictions.estimated_training_time = self.estimate_training_time(code, context);
528        predictions.estimated_inference_latency = self.estimate_inference_latency(code, context);
529        predictions.scaling_characteristics = self.predict_scaling_behavior(code, context);
530
531        // Real, code/context-dependent bottleneck signals: each entry only
532        // appears when the pattern it describes was actually detected. The
533        // old implementation emitted both strings unconditionally,
534        // regardless of what `code`/`context` actually contained.
535        let mut signal_count: u32 = 0;
536        if code.contains("attention") && !code.contains("flash") {
537            predictions.predicted_bottlenecks.push(
538                "Attention computation may become a bottleneck for long sequences (no Flash \
539                 Attention detected in this code)"
540                    .to_string(),
541            );
542            signal_count += 1;
543        }
544        if context.model_size > 1_000_000_000 {
545            predictions.predicted_bottlenecks.push(
546                "Memory bandwidth may limit performance for large batch sizes (model exceeds \
547                 1B parameters)"
548                    .to_string(),
549            );
550            signal_count += 1;
551        }
552
553        // Confidence reflects how many independent real signals were
554        // found: an honest `0.0` ("no basis for a bottleneck prediction")
555        // when none matched, scaling up with corroborating signals. Never
556        // the old fixed `0.75` regardless of input.
557        predictions.confidence_score = match signal_count {
558            0 => 0.0,
559            1 => 0.6,
560            _ => 0.75,
561        };
562
563        Ok(predictions)
564    }
565
566    async fn detect_fusion_opportunities(
567        &self,
568        operations: &[TensorOperation],
569    ) -> Result<Vec<FusionOpportunity>> {
570        let mut opportunities = Vec::new();
571
572        // Detect MatMul + Add fusion (GEMM)
573        for window in operations.windows(2) {
574            if let [op1, op2] = window {
575                if matches!(op1.op_type, OperationType::MatMul)
576                    && matches!(op2.op_type, OperationType::Add)
577                {
578                    opportunities.push(FusionOpportunity {
579                        operations: vec![op1.clone(), op2.clone()],
580                        fusion_type: FusionType::GEMM,
581                        estimated_speedup: 1.3,
582                        description: "MatMul + Add can be fused into GEMM operation".to_string(),
583                    });
584                }
585            }
586        }
587
588        // Detect activation fusion opportunities
589        for window in operations.windows(2) {
590            if let [op1, op2] = window {
591                if matches!(op1.op_type, OperationType::Linear)
592                    && matches!(op2.op_type, OperationType::Activation)
593                {
594                    opportunities.push(FusionOpportunity {
595                        operations: vec![op1.clone(), op2.clone()],
596                        fusion_type: FusionType::LinearActivation,
597                        estimated_speedup: 1.2,
598                        description: "Linear + Activation can be fused".to_string(),
599                    });
600                }
601            }
602        }
603
604        Ok(opportunities)
605    }
606
607    async fn detect_memory_optimizations(
608        &self,
609        operations: &[TensorOperation],
610    ) -> Result<Vec<MemoryOptimization>> {
611        let mut optimizations = Vec::new();
612
613        // Detect in-place operation opportunities
614        for op in operations {
615            if op.can_be_inplace() && !op.is_inplace {
616                optimizations.push(MemoryOptimization {
617                    operation: op.clone(),
618                    optimization_type: MemoryOptimizationType::InPlace,
619                    memory_savings: op.output_size_bytes,
620                    description: format!("Operation {} can be performed in-place", op.name),
621                });
622            }
623        }
624
625        // Detect tensor reuse opportunities
626        let mut tensor_usage = HashMap::new();
627        for op in operations {
628            for input in &op.inputs {
629                *tensor_usage.entry(input.clone()).or_insert(0) += 1;
630            }
631        }
632
633        for (tensor, usage_count) in tensor_usage {
634            if usage_count == 1 {
635                optimizations.push(MemoryOptimization {
636                    operation: TensorOperation::default(),
637                    optimization_type: MemoryOptimizationType::TensorReuse,
638                    memory_savings: 0, // Would calculate based on tensor size
639                    description: format!("Tensor {} can be reused", tensor),
640                });
641            }
642        }
643
644        Ok(optimizations)
645    }
646
647    async fn detect_parallelization_opportunities(
648        &self,
649        operations: &[TensorOperation],
650    ) -> Result<Vec<ParallelizationOpportunity>> {
651        let mut opportunities = Vec::new();
652
653        // Detect independent operations that can run in parallel
654        for (i, op1) in operations.iter().enumerate() {
655            for op2 in operations.iter().skip(i + 1) {
656                if self.operations_are_independent(op1, op2) {
657                    opportunities.push(ParallelizationOpportunity {
658                        operations: vec![op1.clone(), op2.clone()],
659                        parallelization_type: ParallelizationType::DataParallel,
660                        estimated_speedup: 1.8,
661                        description: "Operations can run in parallel".to_string(),
662                    });
663                }
664            }
665        }
666
667        Ok(opportunities)
668    }
669
670    async fn detect_redundant_operations(
671        &self,
672        operations: &[TensorOperation],
673    ) -> Result<Vec<RedundantOperation>> {
674        let mut redundant = Vec::new();
675
676        // Detect duplicate operations
677        for (i, op1) in operations.iter().enumerate() {
678            for (_j, op2) in operations.iter().enumerate().skip(i + 1) {
679                if self.operations_are_equivalent(op1, op2) {
680                    redundant.push(RedundantOperation {
681                        original_operation: op1.clone(),
682                        redundant_operation: op2.clone(),
683                        redundancy_type: RedundancyType::Duplicate,
684                        description: "Operations produce identical results".to_string(),
685                    });
686                }
687            }
688        }
689
690        Ok(redundant)
691    }
692
693    // Analysis helper methods
694
695    fn operations_are_independent(&self, op1: &TensorOperation, op2: &TensorOperation) -> bool {
696        // Check if operations have no data dependencies
697        for input1 in &op1.inputs {
698            for output2 in &op2.outputs {
699                if input1 == output2 {
700                    return false;
701                }
702            }
703        }
704        for input2 in &op2.inputs {
705            for output1 in &op1.outputs {
706                if input2 == output1 {
707                    return false;
708                }
709            }
710        }
711        true
712    }
713
714    fn operations_are_equivalent(&self, op1: &TensorOperation, op2: &TensorOperation) -> bool {
715        op1.op_type == op2.op_type && op1.inputs == op2.inputs && op1.parameters == op2.parameters
716    }
717
718    fn compute_code_hash(&self, code: &str) -> String {
719        use std::collections::hash_map::DefaultHasher;
720        use std::hash::{Hash, Hasher};
721
722        let mut hasher = DefaultHasher::new();
723        code.hash(&mut hasher);
724        format!("{:x}", hasher.finish())
725    }
726
727    fn get_cached_analysis(&self, code_hash: &str) -> Option<&CachedAnalysis> {
728        self.analysis_cache.get(code_hash).and_then(|cached| {
729            let age = std::time::SystemTime::now()
730                .duration_since(cached.timestamp)
731                .unwrap_or_default();
732
733            if age.as_secs() < self.config.cache_expiration_hours * 3600 {
734                Some(cached)
735            } else {
736                None
737            }
738        })
739    }
740
741    fn cache_analysis(&mut self, code_hash: String, result: &CodeAnalysisResult) {
742        self.analysis_cache.insert(
743            code_hash.clone(),
744            CachedAnalysis {
745                result: result.clone(),
746                timestamp: std::time::SystemTime::now(),
747                code_hash,
748            },
749        );
750    }
751
752    fn calculate_quality_score(&self, result: &CodeAnalysisResult) -> f64 {
753        let mut score: f64 = 100.0;
754
755        // Deduct points for issues
756        for issue in &result.identified_issues {
757            match issue.severity {
758                Severity::Critical => score -= 20.0,
759                Severity::High => score -= 10.0,
760                Severity::Medium => score -= 5.0,
761                Severity::Low => score -= 2.0,
762                Severity::Info => score -= 0.0,
763            }
764        }
765
766        // Deduct points for security issues
767        for vulnerability in &result.security_issues {
768            match vulnerability.severity {
769                Severity::Critical => score -= 25.0,
770                Severity::High => score -= 15.0,
771                Severity::Medium => score -= 8.0,
772                Severity::Low => score -= 3.0,
773                Severity::Info => score -= 0.0,
774            }
775        }
776
777        // Add points for good patterns
778        for pattern in &result.detected_patterns {
779            if pattern.pattern_type == PatternType::GoodPattern {
780                score += 2.0;
781            }
782        }
783
784        score.max(0.0).min(100.0)
785    }
786
787    fn calculate_confidence_score(&self, result: &CodeAnalysisResult) -> f64 {
788        let mut total_confidence = 0.0;
789        let mut count = 0;
790
791        for issue in &result.identified_issues {
792            total_confidence += issue.confidence;
793            count += 1;
794        }
795
796        for pattern in &result.detected_patterns {
797            total_confidence += pattern.confidence;
798            count += 1;
799        }
800
801        if count > 0 {
802            total_confidence / count as f64
803        } else {
804            1.0
805        }
806    }
807
808    fn estimate_memory_usage(&self, code: &str, context: &ModelContext) -> f64 {
809        // Simplified estimation based on model size and code patterns
810        let base_memory = context.model_size as f64 * 4.0; // 4 bytes per parameter
811
812        let mut multiplier = 1.0;
813        if code.contains("gradient_accumulation") {
814            multiplier += 0.5;
815        }
816        if code.contains("mixed_precision") {
817            multiplier *= 0.7;
818        }
819
820        base_memory * multiplier / 1_000_000.0 // Convert to MB
821    }
822
823    fn estimate_training_time(&self, code: &str, context: &ModelContext) -> f64 {
824        // Simplified estimation in minutes per epoch
825        let base_time = (context.model_size as f64).log10() * 10.0;
826
827        let mut multiplier = 1.0;
828        if code.contains("mixed_precision") {
829            multiplier *= 0.6;
830        }
831        if code.contains("gradient_checkpointing") {
832            multiplier *= 1.3;
833        }
834
835        base_time * multiplier
836    }
837
838    fn estimate_inference_latency(&self, code: &str, context: &ModelContext) -> f64 {
839        // Simplified estimation in milliseconds
840        let base_latency = (context.model_size as f64).log10() * 5.0;
841
842        let mut multiplier = 1.0;
843        if code.contains("compile") {
844            multiplier *= 0.5;
845        }
846        if code.contains("quantization") {
847            multiplier *= 0.7;
848        }
849
850        base_latency * multiplier
851    }
852
853    fn predict_scaling_behavior(
854        &self,
855        _code: &str,
856        context: &ModelContext,
857    ) -> ScalingCharacteristics {
858        ScalingCharacteristics {
859            batch_size_scaling: if context.model_size > 1_000_000_000 {
860                ScalingBehavior::Sublinear
861            } else {
862                ScalingBehavior::Linear
863            },
864            sequence_length_scaling: ScalingBehavior::Quadratic, // Attention is O(n²)
865            model_size_scaling: ScalingBehavior::Linear,
866            memory_scaling: ScalingBehavior::Linear,
867        }
868    }
869
870    fn estimate_optimization_speedup(&self, report: &TensorOptimizationReport) -> f64 {
871        let mut speedup = 1.0;
872
873        for fusion in &report.fusion_opportunities {
874            speedup *= fusion.estimated_speedup;
875        }
876
877        for parallel in &report.parallelization_opportunities {
878            speedup *= parallel.estimated_speedup;
879        }
880
881        speedup.min(10.0) // Cap at 10x speedup
882    }
883
884    fn estimate_memory_savings(&self, report: &TensorOptimizationReport) -> f64 {
885        let total_savings: u64 =
886            report.memory_optimizations.iter().map(|opt| opt.memory_savings).sum();
887
888        total_savings as f64 / 1_000_000.0 // Convert to MB
889    }
890
891    async fn analyze_error_patterns(
892        &self,
893        error_context: &ErrorContext,
894    ) -> Result<Vec<ProbableCause>> {
895        let mut causes = Vec::new();
896
897        match error_context.error_type.as_str() {
898            "OutOfMemoryError" => {
899                causes.push(ProbableCause {
900                    cause: "Batch size too large".to_string(),
901                    probability: 0.8,
902                    evidence: vec!["GPU memory limit exceeded".to_string()],
903                });
904                causes.push(ProbableCause {
905                    cause: "Model too large for available memory".to_string(),
906                    probability: 0.6,
907                    evidence: vec!["Model parameter count".to_string()],
908                });
909            },
910            "GradientExplosion" => {
911                causes.push(ProbableCause {
912                    cause: "Learning rate too high".to_string(),
913                    probability: 0.7,
914                    evidence: vec!["Gradient norm increasing rapidly".to_string()],
915                });
916            },
917            _ => {
918                causes.push(ProbableCause {
919                    cause: "Unknown error pattern".to_string(),
920                    probability: 0.3,
921                    evidence: vec![],
922                });
923            },
924        }
925
926        Ok(causes)
927    }
928
929    async fn generate_suggested_fixes(
930        &self,
931        error_context: &ErrorContext,
932    ) -> Result<Vec<SuggestedFix>> {
933        let mut fixes = Vec::new();
934
935        match error_context.error_type.as_str() {
936            "OutOfMemoryError" => {
937                fixes.push(SuggestedFix {
938                    description: "Reduce batch size".to_string(),
939                    implementation: "batch_size = batch_size // 2".to_string(),
940                    confidence: 0.9,
941                    estimated_impact: "Should free ~50% of memory".to_string(),
942                });
943                fixes.push(SuggestedFix {
944                    description: "Enable gradient checkpointing".to_string(),
945                    implementation: "model.gradient_checkpointing_enable()".to_string(),
946                    confidence: 0.8,
947                    estimated_impact: "Reduces memory by ~40% with 10-20% speed penalty"
948                        .to_string(),
949                });
950            },
951            "GradientExplosion" => {
952                fixes.push(SuggestedFix {
953                    description: "Add gradient clipping".to_string(),
954                    implementation:
955                        "torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)"
956                            .to_string(),
957                    confidence: 0.95,
958                    estimated_impact: "Prevents gradient explosion".to_string(),
959                });
960            },
961            _ => {},
962        }
963
964        Ok(fixes)
965    }
966
967    async fn generate_debugging_steps(
968        &self,
969        error_context: &ErrorContext,
970    ) -> Result<Vec<DebuggingStep>> {
971        let mut steps = Vec::new();
972
973        steps.push(DebuggingStep {
974            step_number: 1,
975            description: "Check system resources".to_string(),
976            command: Some("nvidia-smi".to_string()),
977            expected_output: "GPU memory usage and availability".to_string(),
978        });
979
980        steps.push(DebuggingStep {
981            step_number: 2,
982            description: "Verify model configuration".to_string(),
983            command: Some("print(model)".to_string()),
984            expected_output: "Model architecture and parameter count".to_string(),
985        });
986
987        if error_context.error_type.as_str() == "OutOfMemoryError" {
988            steps.push(DebuggingStep {
989                step_number: 3,
990                description: "Check tensor shapes and batch size".to_string(),
991                command: Some(
992                    "print(f'Batch size: {batch_size}, Input shape: {input.shape}')".to_string(),
993                ),
994                expected_output: "Current batch size and input dimensions".to_string(),
995            });
996        }
997
998        Ok(steps)
999    }
1000
1001    async fn find_related_documentation(
1002        &self,
1003        error_context: &ErrorContext,
1004    ) -> Result<Vec<DocumentationReference>> {
1005        let mut references = Vec::new();
1006
1007        match error_context.error_type.as_str() {
1008            "OutOfMemoryError" => {
1009                references.push(DocumentationReference {
1010                    title: "Memory Management Best Practices".to_string(),
1011                    url: "https://docs.trustformers.ai/memory-management".to_string(),
1012                    relevance_score: 0.95,
1013                });
1014                references.push(DocumentationReference {
1015                    title: "Gradient Checkpointing Guide".to_string(),
1016                    url: "https://docs.trustformers.ai/gradient-checkpointing".to_string(),
1017                    relevance_score: 0.8,
1018                });
1019            },
1020            "GradientExplosion" => {
1021                references.push(DocumentationReference {
1022                    title: "Training Stability Guide".to_string(),
1023                    url: "https://docs.trustformers.ai/training-stability".to_string(),
1024                    relevance_score: 0.9,
1025                });
1026            },
1027            _ => {},
1028        }
1029
1030        Ok(references)
1031    }
1032
1033    fn calculate_debugging_confidence(&self, assistance: &DebuggingAssistance) -> f64 {
1034        let avg_cause_probability =
1035            assistance.probable_causes.iter().map(|cause| cause.probability).sum::<f64>()
1036                / assistance.probable_causes.len().max(1) as f64;
1037
1038        let avg_fix_confidence =
1039            assistance.suggested_fixes.iter().map(|fix| fix.confidence).sum::<f64>()
1040                / assistance.suggested_fixes.len().max(1) as f64;
1041
1042        (avg_cause_probability + avg_fix_confidence) / 2.0
1043    }
1044}
1045
1046// Supporting data structures and types
1047
1048/// Model pattern database for common patterns and anti-patterns
1049#[derive(Debug)]
1050struct ModelPatternDatabase {
1051    patterns: HashMap<String, PatternDefinition>,
1052}
1053
1054impl ModelPatternDatabase {
1055    fn new() -> Self {
1056        let mut patterns = HashMap::new();
1057
1058        // Add common patterns
1059        patterns.insert(
1060            "gradient_clipping".to_string(),
1061            PatternDefinition {
1062                name: "Gradient Clipping".to_string(),
1063                pattern_type: PatternType::GoodPattern,
1064                keywords: vec![
1065                    "clip_grad_norm".to_string(),
1066                    "gradient".to_string(),
1067                    "clip".to_string(),
1068                ],
1069                severity: Severity::Info,
1070                description: "Proper gradient clipping prevents gradient explosion".to_string(),
1071            },
1072        );
1073
1074        Self { patterns }
1075    }
1076}
1077
1078#[derive(Debug, Clone)]
1079struct PatternDefinition {
1080    name: String,
1081    pattern_type: PatternType,
1082    keywords: Vec<String>,
1083    severity: Severity,
1084    description: String,
1085}
1086
1087/// Model context for analysis
1088#[derive(Debug, Clone)]
1089pub struct ModelContext {
1090    pub model_type: ModelType,
1091    pub model_size: u64, // Number of parameters
1092    pub framework: String,
1093    pub target_hardware: String,
1094    pub training_stage: TrainingStage,
1095}
1096
1097#[derive(Debug, Clone, PartialEq)]
1098pub enum ModelType {
1099    Training,
1100    Inference,
1101    Production,
1102    Development,
1103}
1104
1105#[derive(Debug, Clone)]
1106pub enum TrainingStage {
1107    Training,
1108    Development,
1109    Pretraining,
1110    Finetuning,
1111    Evaluation,
1112    Inference,
1113}
1114
1115/// Comprehensive code analysis result
1116#[derive(Debug, Clone, Serialize, Deserialize)]
1117pub struct CodeAnalysisResult {
1118    pub quality_score: f64,
1119    pub detected_patterns: Vec<DetectedPattern>,
1120    pub identified_issues: Vec<IdentifiedIssue>,
1121    pub optimization_suggestions: Vec<OptimizationSuggestion>,
1122    pub security_issues: Vec<SecurityIssue>,
1123    pub performance_predictions: PerformancePredictions,
1124    pub analysis_metadata: AnalysisMetadata,
1125}
1126
1127impl CodeAnalysisResult {
1128    fn new() -> Self {
1129        Self {
1130            quality_score: 0.0,
1131            detected_patterns: Vec::new(),
1132            identified_issues: Vec::new(),
1133            optimization_suggestions: Vec::new(),
1134            security_issues: Vec::new(),
1135            performance_predictions: PerformancePredictions::new(),
1136            analysis_metadata: AnalysisMetadata::default(),
1137        }
1138    }
1139}
1140
1141#[derive(Debug, Clone, Serialize, Deserialize)]
1142pub struct DetectedPattern {
1143    pub pattern_type: PatternType,
1144    pub name: String,
1145    pub description: String,
1146    pub severity: Severity,
1147    pub confidence: f64,
1148    pub recommendations: Vec<String>,
1149}
1150
1151#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
1152pub enum PatternType {
1153    GoodPattern,
1154    AntiPattern,
1155    OptimizationOpportunity,
1156    SecurityConcern,
1157}
1158
1159#[derive(Debug, Clone, Serialize, Deserialize)]
1160pub struct IdentifiedIssue {
1161    pub issue_type: IssueType,
1162    pub title: String,
1163    pub description: String,
1164    pub severity: Severity,
1165    pub confidence: f64,
1166    pub suggested_fix: String,
1167    pub code_location: Option<CodeLocation>,
1168}
1169
1170#[derive(Debug, Clone, Serialize, Deserialize)]
1171pub enum IssueType {
1172    NumericalStability,
1173    Performance,
1174    MemoryLeak,
1175    LogicError,
1176    TypeMismatch,
1177    ResourceLeak,
1178}
1179
1180#[derive(Debug, Clone, Serialize, Deserialize)]
1181pub struct CodeLocation {
1182    pub file: String,
1183    pub line: u32,
1184    pub column: u32,
1185}
1186
1187/// Real 1-indexed `(line, column)` of the first occurrence of `needle` in
1188/// `code`, or `None` if `needle` does not occur. `file` is a fixed sentinel
1189/// (`analyze_model_code` receives only a code string, not a path) --
1190/// callers that need a real path should overwrite `CodeLocation::file` with
1191/// the source file they read `code` from.
1192///
1193/// Used so [`IdentifiedIssue::code_location`] points at the text that
1194/// actually triggered the rule, instead of always being `None`.
1195fn locate_in_code(code: &str, needle: &str) -> Option<CodeLocation> {
1196    let byte_offset = code.find(needle)?;
1197    let mut line: u32 = 1;
1198    let mut last_newline_offset: Option<usize> = None;
1199    for (idx, ch) in code[..byte_offset].char_indices() {
1200        if ch == '\n' {
1201            line += 1;
1202            last_newline_offset = Some(idx);
1203        }
1204    }
1205    let column = match last_newline_offset {
1206        Some(nl_idx) => (byte_offset - nl_idx) as u32,
1207        None => byte_offset as u32 + 1,
1208    };
1209
1210    Some(CodeLocation {
1211        file: "<analyzed source>".to_string(),
1212        line,
1213        column,
1214    })
1215}
1216
1217#[derive(Debug, Clone, Serialize, Deserialize)]
1218pub struct OptimizationSuggestion {
1219    pub optimization_type: OptimizationType,
1220    pub title: String,
1221    pub description: String,
1222    pub potential_speedup: f64,
1223    pub memory_savings: f64,
1224    pub implementation_effort: ImplementationEffort,
1225    pub confidence: f64,
1226    pub code_example: Option<String>,
1227}
1228
1229#[derive(Debug, Clone, Serialize, Deserialize)]
1230pub enum OptimizationType {
1231    MixedPrecision,
1232    ModelCompilation,
1233    MemoryOptimization,
1234    ComputationOptimization,
1235    IOOptimization,
1236    ParallelizationOptimization,
1237}
1238
1239#[derive(Debug, Clone, Serialize, Deserialize)]
1240pub enum ImplementationEffort {
1241    Low,
1242    Medium,
1243    High,
1244}
1245
1246#[derive(Debug, Clone, Serialize, Deserialize)]
1247pub struct SecurityIssue {
1248    pub vulnerability_type: VulnerabilityType,
1249    pub title: String,
1250    pub description: String,
1251    pub severity: Severity,
1252    pub confidence: f64,
1253    pub mitigation: String,
1254    pub cve_references: Vec<String>,
1255}
1256
1257#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
1258pub enum VulnerabilityType {
1259    CodeExecution,
1260    DataExposure,
1261    InputValidation,
1262    AuthenticationBypass,
1263    PrivilegeEscalation,
1264}
1265
1266#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
1267pub enum Severity {
1268    Critical,
1269    High,
1270    Medium,
1271    Low,
1272    Info,
1273}
1274
1275#[derive(Debug, Clone, Serialize, Deserialize)]
1276pub struct PerformancePredictions {
1277    pub estimated_memory_usage: f64,      // MB
1278    pub estimated_training_time: f64,     // minutes per epoch
1279    pub estimated_inference_latency: f64, // milliseconds
1280    pub scaling_characteristics: ScalingCharacteristics,
1281    pub predicted_bottlenecks: Vec<String>,
1282    pub confidence_score: f64,
1283}
1284
1285impl PerformancePredictions {
1286    fn new() -> Self {
1287        Self {
1288            estimated_memory_usage: 0.0,
1289            estimated_training_time: 0.0,
1290            estimated_inference_latency: 0.0,
1291            scaling_characteristics: ScalingCharacteristics::default(),
1292            predicted_bottlenecks: Vec::new(),
1293            confidence_score: 0.0,
1294        }
1295    }
1296}
1297
1298#[derive(Debug, Clone, Serialize, Deserialize)]
1299pub struct ScalingCharacteristics {
1300    pub batch_size_scaling: ScalingBehavior,
1301    pub sequence_length_scaling: ScalingBehavior,
1302    pub model_size_scaling: ScalingBehavior,
1303    pub memory_scaling: ScalingBehavior,
1304}
1305
1306impl Default for ScalingCharacteristics {
1307    fn default() -> Self {
1308        Self {
1309            batch_size_scaling: ScalingBehavior::Linear,
1310            sequence_length_scaling: ScalingBehavior::Linear,
1311            model_size_scaling: ScalingBehavior::Linear,
1312            memory_scaling: ScalingBehavior::Linear,
1313        }
1314    }
1315}
1316
1317#[derive(Debug, Clone, Serialize, Deserialize)]
1318pub enum ScalingBehavior {
1319    Constant,
1320    Linear,
1321    Quadratic,
1322    Exponential,
1323    Sublinear,
1324}
1325
1326#[derive(Debug, Clone, Serialize, Deserialize)]
1327pub struct AnalysisMetadata {
1328    pub analysis_duration: Duration,
1329    pub confidence_score: f64,
1330    pub analyzer_version: String,
1331    pub timestamp: std::time::SystemTime,
1332}
1333
1334impl Default for AnalysisMetadata {
1335    fn default() -> Self {
1336        Self {
1337            analysis_duration: Duration::from_secs(0),
1338            confidence_score: 0.0,
1339            analyzer_version: "1.0.0".to_string(),
1340            timestamp: std::time::SystemTime::now(),
1341        }
1342    }
1343}
1344
1345// Tensor operation analysis types
1346
1347#[derive(Debug, Clone)]
1348pub struct TensorOperation {
1349    pub name: String,
1350    pub op_type: OperationType,
1351    pub inputs: Vec<String>,
1352    pub outputs: Vec<String>,
1353    pub parameters: HashMap<String, String>,
1354    pub output_size_bytes: u64,
1355    pub is_inplace: bool,
1356}
1357
1358impl Default for TensorOperation {
1359    fn default() -> Self {
1360        Self {
1361            name: String::new(),
1362            op_type: OperationType::Unknown,
1363            inputs: Vec::new(),
1364            outputs: Vec::new(),
1365            parameters: HashMap::new(),
1366            output_size_bytes: 0,
1367            is_inplace: false,
1368        }
1369    }
1370}
1371
1372impl TensorOperation {
1373    fn can_be_inplace(&self) -> bool {
1374        matches!(
1375            self.op_type,
1376            OperationType::Add | OperationType::Mul | OperationType::Activation
1377        )
1378    }
1379}
1380
1381#[derive(Debug, Clone, PartialEq)]
1382pub enum OperationType {
1383    MatMul,
1384    Add,
1385    Mul,
1386    Conv2D,
1387    Linear,
1388    Activation,
1389    Pooling,
1390    BatchNorm,
1391    LayerNorm,
1392    Attention,
1393    Unknown,
1394}
1395
1396#[derive(Debug, Clone)]
1397pub struct TensorOptimizationReport {
1398    pub fusion_opportunities: Vec<FusionOpportunity>,
1399    pub memory_optimizations: Vec<MemoryOptimization>,
1400    pub parallelization_opportunities: Vec<ParallelizationOpportunity>,
1401    pub redundant_operations: Vec<RedundantOperation>,
1402    pub estimated_speedup: f64,
1403    pub estimated_memory_savings: f64,
1404}
1405
1406impl TensorOptimizationReport {
1407    fn new() -> Self {
1408        Self {
1409            fusion_opportunities: Vec::new(),
1410            memory_optimizations: Vec::new(),
1411            parallelization_opportunities: Vec::new(),
1412            redundant_operations: Vec::new(),
1413            estimated_speedup: 1.0,
1414            estimated_memory_savings: 0.0,
1415        }
1416    }
1417}
1418
1419#[derive(Debug, Clone)]
1420pub struct FusionOpportunity {
1421    pub operations: Vec<TensorOperation>,
1422    pub fusion_type: FusionType,
1423    pub estimated_speedup: f64,
1424    pub description: String,
1425}
1426
1427#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
1428pub enum FusionType {
1429    GEMM,
1430    LinearActivation,
1431    ConvBatchNorm,
1432    AttentionQKV,
1433}
1434
1435#[derive(Debug, Clone)]
1436pub struct MemoryOptimization {
1437    pub operation: TensorOperation,
1438    pub optimization_type: MemoryOptimizationType,
1439    pub memory_savings: u64,
1440    pub description: String,
1441}
1442
1443#[derive(Debug, Clone)]
1444pub enum MemoryOptimizationType {
1445    InPlace,
1446    TensorReuse,
1447    MemoryPool,
1448    GradientCheckpointing,
1449}
1450
1451#[derive(Debug, Clone)]
1452pub struct ParallelizationOpportunity {
1453    pub operations: Vec<TensorOperation>,
1454    pub parallelization_type: ParallelizationType,
1455    pub estimated_speedup: f64,
1456    pub description: String,
1457}
1458
1459#[derive(Debug, Clone)]
1460pub enum ParallelizationType {
1461    DataParallel,
1462    ModelParallel,
1463    PipelineParallel,
1464    TensorParallel,
1465}
1466
1467#[derive(Debug, Clone)]
1468pub struct RedundantOperation {
1469    pub original_operation: TensorOperation,
1470    pub redundant_operation: TensorOperation,
1471    pub redundancy_type: RedundancyType,
1472    pub description: String,
1473}
1474
1475#[derive(Debug, Clone)]
1476pub enum RedundancyType {
1477    Duplicate,
1478    Subsumed,
1479    Unnecessary,
1480}
1481
1482// Error context and debugging assistance types
1483
1484#[derive(Debug, Clone)]
1485pub struct ErrorContext {
1486    pub error_type: String,
1487    pub error_message: String,
1488    pub stack_trace: Option<String>,
1489    pub system_info: SystemInfo,
1490    pub model_info: Option<ModelContext>,
1491}
1492
1493#[derive(Debug, Clone)]
1494pub struct SystemInfo {
1495    pub gpu_memory_total: u64,
1496    pub gpu_memory_used: u64,
1497    pub cpu_count: u32,
1498    pub ram_total: u64,
1499    pub ram_used: u64,
1500}
1501
1502#[derive(Debug, Clone)]
1503pub struct DebuggingAssistance {
1504    pub probable_causes: Vec<ProbableCause>,
1505    pub suggested_fixes: Vec<SuggestedFix>,
1506    pub debugging_steps: Vec<DebuggingStep>,
1507    pub related_documentation: Vec<DocumentationReference>,
1508    pub confidence_score: f64,
1509}
1510
1511impl DebuggingAssistance {
1512    fn new() -> Self {
1513        Self {
1514            probable_causes: Vec::new(),
1515            suggested_fixes: Vec::new(),
1516            debugging_steps: Vec::new(),
1517            related_documentation: Vec::new(),
1518            confidence_score: 0.0,
1519        }
1520    }
1521}
1522
1523#[derive(Debug, Clone)]
1524pub struct ProbableCause {
1525    pub cause: String,
1526    pub probability: f64,
1527    pub evidence: Vec<String>,
1528}
1529
1530#[derive(Debug, Clone)]
1531pub struct SuggestedFix {
1532    pub description: String,
1533    pub implementation: String,
1534    pub confidence: f64,
1535    pub estimated_impact: String,
1536}
1537
1538#[derive(Debug, Clone)]
1539pub struct DebuggingStep {
1540    pub step_number: u32,
1541    pub description: String,
1542    pub command: Option<String>,
1543    pub expected_output: String,
1544}
1545
1546#[derive(Debug, Clone)]
1547pub struct DocumentationReference {
1548    pub title: String,
1549    pub url: String,
1550    pub relevance_score: f64,
1551}
1552
1553// Performance metrics
1554
1555#[derive(Debug, Serialize, Deserialize)]
1556pub struct AnalysisPerformanceMetrics {
1557    pub total_analyses: u64,
1558    pub average_analysis_time: Duration,
1559    pub cache_hit_rate: f64,
1560    pub cached_results: usize,
1561}
1562
1563/// Macro for quick AI code analysis
1564#[macro_export]
1565macro_rules! ai_analyze {
1566    ($code:expr, $context:expr) => {{
1567        let mut analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1568        analyzer.analyze_model_code($code, $context).await
1569    }};
1570}
1571
1572#[path = "ai_code_analyzer_tests.rs"]
1573mod ai_code_analyzer_tests;
1574
1575#[cfg(test)]
1576mod tests {
1577    use super::*;
1578
1579    #[tokio::test]
1580    async fn test_ai_code_analyzer_creation() {
1581        let analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1582        assert!(analyzer.config.enable_deep_analysis);
1583    }
1584
1585    #[tokio::test]
1586    async fn test_pattern_detection() {
1587        let mut analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1588
1589        let code = r#"
1590        import torch
1591
1592        def train_step(model, data):
1593            torch.cuda.empty_cache()  # Should trigger anti-pattern
1594            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # Good pattern
1595            return grad_norm
1596        "#;
1597
1598        let context = ModelContext {
1599            model_type: ModelType::Production,
1600            model_size: 1_000_000,
1601            framework: "PyTorch".to_string(),
1602            target_hardware: "CUDA".to_string(),
1603            training_stage: TrainingStage::Training,
1604        };
1605
1606        let result = analyzer
1607            .analyze_model_code(code, context)
1608            .await
1609            .expect("async operation failed");
1610        assert!(!result.detected_patterns.is_empty());
1611    }
1612
1613    #[tokio::test]
1614    async fn test_security_vulnerability_detection() {
1615        let mut analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1616
1617        let code = r#"
1618        import pickle
1619
1620        def load_model(path):
1621            with open(path, 'rb') as f:
1622                model = pickle.load(f)  # Should trigger security warning
1623            return model
1624        "#;
1625
1626        let context = ModelContext {
1627            model_type: ModelType::Production,
1628            model_size: 1_000_000,
1629            framework: "PyTorch".to_string(),
1630            target_hardware: "CUDA".to_string(),
1631            training_stage: TrainingStage::Inference,
1632        };
1633
1634        let result = analyzer
1635            .analyze_model_code(code, context)
1636            .await
1637            .expect("async operation failed");
1638        assert!(!result.security_issues.is_empty());
1639        assert_eq!(
1640            result.security_issues[0].vulnerability_type,
1641            VulnerabilityType::CodeExecution
1642        );
1643    }
1644
1645    #[tokio::test]
1646    async fn test_tensor_operation_analysis() {
1647        let analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1648
1649        let operations = vec![
1650            TensorOperation {
1651                name: "matmul1".to_string(),
1652                op_type: OperationType::MatMul,
1653                inputs: vec!["a".to_string(), "b".to_string()],
1654                outputs: vec!["c".to_string()],
1655                parameters: HashMap::new(),
1656                output_size_bytes: 1024,
1657                is_inplace: false,
1658            },
1659            TensorOperation {
1660                name: "add1".to_string(),
1661                op_type: OperationType::Add,
1662                inputs: vec!["c".to_string(), "bias".to_string()],
1663                outputs: vec!["d".to_string()],
1664                parameters: HashMap::new(),
1665                output_size_bytes: 1024,
1666                is_inplace: false,
1667            },
1668        ];
1669
1670        let report = analyzer
1671            .analyze_tensor_operations(&operations)
1672            .await
1673            .expect("tensor operation failed");
1674        assert!(!report.fusion_opportunities.is_empty());
1675        assert_eq!(report.fusion_opportunities[0].fusion_type, FusionType::GEMM);
1676    }
1677
1678    #[tokio::test]
1679    async fn test_performance_metrics() {
1680        let mut analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1681
1682        // Simulate some analyses
1683        let code = "print('hello')";
1684        let context = ModelContext {
1685            model_type: ModelType::Development,
1686            model_size: 1000,
1687            framework: "PyTorch".to_string(),
1688            target_hardware: "CPU".to_string(),
1689            training_stage: TrainingStage::Development,
1690        };
1691
1692        analyzer
1693            .analyze_model_code(code, context.clone())
1694            .await
1695            .expect("async operation failed");
1696        analyzer
1697            .analyze_model_code(code, context)
1698            .await
1699            .expect("async operation failed"); // Should hit cache
1700
1701        let metrics = analyzer.get_performance_metrics();
1702        assert_eq!(metrics.total_analyses, 2);
1703        assert!(metrics.cache_hit_rate > 0.0);
1704    }
1705
1706    #[tokio::test]
1707    async fn test_debugging_assistance() {
1708        let analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1709
1710        let error_context = ErrorContext {
1711            error_type: "OutOfMemoryError".to_string(),
1712            error_message: "CUDA out of memory".to_string(),
1713            stack_trace: None,
1714            system_info: SystemInfo {
1715                gpu_memory_total: 8_000_000_000,
1716                gpu_memory_used: 7_500_000_000,
1717                cpu_count: 8,
1718                ram_total: 32_000_000_000,
1719                ram_used: 16_000_000_000,
1720            },
1721            model_info: None,
1722        };
1723
1724        let assistance = analyzer
1725            .automated_debugging_assistance(&error_context)
1726            .await
1727            .expect("async operation failed");
1728        assert!(!assistance.probable_causes.is_empty());
1729        assert!(!assistance.suggested_fixes.is_empty());
1730        assert!(assistance.confidence_score > 0.0);
1731    }
1732
1733    fn make_context(model_size: u64) -> ModelContext {
1734        ModelContext {
1735            model_type: ModelType::Training,
1736            model_size,
1737            framework: "PyTorch".to_string(),
1738            target_hardware: "CUDA".to_string(),
1739            training_stage: TrainingStage::Training,
1740        }
1741    }
1742
1743    /// Regression test: `perform_deep_analysis` used to always set
1744    /// `code_location: None` ("Would be populated with actual line
1745    /// numbers"). It must now report the real line/column of the text that
1746    /// triggered the rule.
1747    #[tokio::test]
1748    async fn test_identified_issue_code_location_is_real_not_none() {
1749        let analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1750        let code = "def f(x):\n    return log(softmax(x))\n";
1751        let context = make_context(1_000);
1752
1753        let issues = analyzer
1754            .perform_deep_analysis(code, &context)
1755            .await
1756            .expect("deep analysis should succeed");
1757        let issue = issues
1758            .iter()
1759            .find(|i| matches!(i.issue_type, IssueType::NumericalStability))
1760            .expect("the log-softmax rule should have fired");
1761
1762        let location = issue
1763            .code_location
1764            .as_ref()
1765            .expect("code_location must be Some (a real position), not the old hardcoded None");
1766        // "softmax" occurs on line 2 (1-indexed), starting right after
1767        // "    return log(".
1768        assert_eq!(location.line, 2);
1769        assert_eq!(location.column, "    return log(".len() as u32 + 1);
1770    }
1771
1772    /// Regression test: `predict_performance_characteristics` used to
1773    /// unconditionally return the same two `predicted_bottlenecks` strings
1774    /// and a fixed `confidence_score: 0.75` for any input whatsoever. Code
1775    /// with none of the underlying signals must now honestly report no
1776    /// bottlenecks and zero confidence.
1777    #[tokio::test]
1778    async fn test_performance_predictions_are_honest_with_no_signals() {
1779        let analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1780        let code = "def f(x):\n    return x + 1\n";
1781        let context = make_context(1_000); // well under the 1B-parameter threshold
1782
1783        let predictions = analyzer
1784            .predict_performance_characteristics(code, &context)
1785            .await
1786            .expect("prediction should succeed");
1787
1788        assert!(
1789            predictions.predicted_bottlenecks.is_empty(),
1790            "must not report bottlenecks when no real signal was detected: {:?}",
1791            predictions.predicted_bottlenecks
1792        );
1793        assert_eq!(
1794            predictions.confidence_score, 0.0,
1795            "must not be the old fixed 0.75 when nothing was actually detected"
1796        );
1797    }
1798
1799    /// Companion to the above: with a real attention-without-flash signal
1800    /// present, the corresponding bottleneck must appear and confidence
1801    /// must be nonzero -- and the result must differ between inputs,
1802    /// proving these are no longer fixed constants.
1803    #[tokio::test]
1804    async fn test_performance_predictions_reflect_real_code_signals() {
1805        let analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1806        let attention_code = "out = attention(q, k, v, matmul_impl=True)";
1807        let plain_context = make_context(1_000);
1808
1809        let with_attention = analyzer
1810            .predict_performance_characteristics(attention_code, &plain_context)
1811            .await
1812            .expect("prediction should succeed");
1813        assert_eq!(with_attention.predicted_bottlenecks.len(), 1);
1814        assert!(with_attention.predicted_bottlenecks[0].contains("Attention"));
1815        assert_eq!(with_attention.confidence_score, 0.6);
1816
1817        let large_model_context = make_context(2_000_000_000);
1818        let with_both = analyzer
1819            .predict_performance_characteristics(attention_code, &large_model_context)
1820            .await
1821            .expect("prediction should succeed");
1822        assert_eq!(
1823            with_both.predicted_bottlenecks.len(),
1824            2,
1825            "a large model plus an attention signal must report both real bottlenecks"
1826        );
1827        assert_eq!(with_both.confidence_score, 0.75);
1828    }
1829
1830    /// Regression test: `perform_deep_analysis` and
1831    /// `predict_performance_characteristics` used to sleep for 100ms/50ms
1832    /// each ("Simulate AI analysis") regardless of input. A full
1833    /// `analyze_model_code` call (which invokes both, plus every other
1834    /// rule pass) must now complete essentially immediately.
1835    #[tokio::test]
1836    async fn test_analyze_model_code_has_no_artificial_latency() {
1837        let mut analyzer = AICodeAnalyzer::new(AIAnalysisConfig::default());
1838        let code = "out = attention(q, k, v, matmul_impl=True)\nlog(softmax(x))";
1839        let context = make_context(2_000_000_000);
1840
1841        let start = std::time::Instant::now();
1842        analyzer
1843            .analyze_model_code(code, context)
1844            .await
1845            .expect("analysis should succeed");
1846        let elapsed = start.elapsed();
1847
1848        // The old sleeps alone totalled 150ms; a real rule-based pass over
1849        // a two-line string should take microseconds. 50ms leaves generous
1850        // headroom for slow CI machines while still catching a reintroduced
1851        // sleep.
1852        assert!(
1853            elapsed < std::time::Duration::from_millis(50),
1854            "analysis took {elapsed:?}; the old implementation's artificial sleeps must be gone"
1855        );
1856    }
1857}