Skip to main content

verificar/generator/
coverage.rs

1//! Coverage-guided generation (NAUTILUS-style)
2//!
3//! Implements grammar-aware fuzzing with coverage feedback to prioritize
4//! unexplored AST paths. Based on Aschermann et al. (2019) NAUTILUS.
5//!
6//! # Key Concepts
7//!
8//! - **Coverage Map**: Tracks which AST paths/features have been explored
9//! - **Corpus**: Collection of interesting inputs that increased coverage
10//! - **Energy**: Selection probability based on coverage potential
11//! - **Grammar-aware mutation**: Mutate AST nodes while maintaining validity
12
13#![allow(clippy::self_only_used_in_recursion)]
14#![allow(clippy::match_same_arms)]
15
16use std::collections::HashSet;
17
18use super::{BinaryOp, CompareOp, GeneratedCode, PythonEnumerator, PythonNode, UnaryOp};
19use crate::Language;
20
21/// Coverage information for a generated program
22#[derive(Debug, Clone, Default)]
23pub struct CoverageMap {
24    /// AST node types seen
25    node_types: HashSet<String>,
26    /// AST paths (parent->child relationships) seen
27    ast_paths: HashSet<(String, String)>,
28    /// Feature combinations seen
29    feature_combos: HashSet<String>,
30}
31
32impl CoverageMap {
33    /// Create a new empty coverage map
34    #[must_use]
35    pub fn new() -> Self {
36        Self::default()
37    }
38
39    /// Record a node type as covered
40    pub fn record_node(&mut self, node_type: &str) {
41        self.node_types.insert(node_type.to_string());
42    }
43
44    /// Record an AST path (parent->child) as covered
45    pub fn record_path(&mut self, parent: &str, child: &str) {
46        self.ast_paths
47            .insert((parent.to_string(), child.to_string()));
48    }
49
50    /// Record a feature combination
51    pub fn record_feature(&mut self, feature: &str) {
52        self.feature_combos.insert(feature.to_string());
53    }
54
55    /// Check if this coverage has any new information compared to another
56    #[must_use]
57    pub fn has_new_coverage(&self, existing: &Self) -> bool {
58        // Check for new node types
59        for node in &self.node_types {
60            if !existing.node_types.contains(node) {
61                return true;
62            }
63        }
64
65        // Check for new AST paths
66        for path in &self.ast_paths {
67            if !existing.ast_paths.contains(path) {
68                return true;
69            }
70        }
71
72        // Check for new feature combos
73        for feature in &self.feature_combos {
74            if !existing.feature_combos.contains(feature) {
75                return true;
76            }
77        }
78
79        false
80    }
81
82    /// Merge another coverage map into this one
83    pub fn merge(&mut self, other: &Self) {
84        self.node_types.extend(other.node_types.iter().cloned());
85        self.ast_paths.extend(other.ast_paths.iter().cloned());
86        self.feature_combos
87            .extend(other.feature_combos.iter().cloned());
88    }
89
90    /// Get the total number of covered items
91    #[must_use]
92    pub fn coverage_count(&self) -> usize {
93        self.node_types.len() + self.ast_paths.len() + self.feature_combos.len()
94    }
95
96    /// Get covered node types
97    #[must_use]
98    pub fn node_types(&self) -> &HashSet<String> {
99        &self.node_types
100    }
101
102    /// Get covered AST paths
103    #[must_use]
104    pub fn ast_paths(&self) -> &HashSet<(String, String)> {
105        &self.ast_paths
106    }
107}
108
109/// Entry in the corpus of interesting inputs
110#[derive(Debug, Clone)]
111pub struct CorpusEntry {
112    /// The generated code
113    pub code: GeneratedCode,
114    /// Coverage achieved by this input
115    pub coverage: CoverageMap,
116    /// Energy score (selection probability weight)
117    pub energy: f64,
118    /// Number of times this entry has been selected
119    pub selection_count: usize,
120    /// AST representation for mutation
121    pub ast: Option<PythonNode>,
122}
123
124impl CorpusEntry {
125    /// Create a new corpus entry
126    #[must_use]
127    pub fn new(code: GeneratedCode, coverage: CoverageMap) -> Self {
128        Self {
129            code,
130            coverage,
131            energy: 1.0,
132            selection_count: 0,
133            ast: None,
134        }
135    }
136
137    /// Create a corpus entry with AST
138    #[must_use]
139    pub fn with_ast(code: GeneratedCode, coverage: CoverageMap, ast: PythonNode) -> Self {
140        Self {
141            code,
142            coverage,
143            energy: 1.0,
144            selection_count: 0,
145            ast: Some(ast),
146        }
147    }
148
149    /// Update energy based on coverage potential
150    pub fn update_energy(&mut self, global_coverage: &CoverageMap) {
151        // Higher energy for entries with unique coverage
152        let unique_nodes = self
153            .coverage
154            .node_types
155            .difference(&global_coverage.node_types)
156            .count();
157        let unique_paths = self
158            .coverage
159            .ast_paths
160            .difference(&global_coverage.ast_paths)
161            .count();
162
163        // Energy decays with selection count but boosted by unique coverage
164        let decay = 1.0 / (1.0 + self.selection_count as f64 * 0.1);
165        let uniqueness_boost = 1.0 + (unique_nodes + unique_paths) as f64 * 0.5;
166
167        self.energy = decay * uniqueness_boost;
168    }
169}
170
171/// NAUTILUS-style coverage-guided generator
172#[derive(Debug)]
173pub struct NautilusGenerator {
174    /// Corpus of interesting inputs
175    corpus: Vec<CorpusEntry>,
176    /// Global coverage map
177    global_coverage: CoverageMap,
178    /// Maximum corpus size
179    max_corpus_size: usize,
180    /// Maximum AST depth for generation
181    max_depth: usize,
182    /// Target language
183    language: Language,
184    /// Random seed for reproducibility
185    seed: u64,
186    /// Simple RNG state
187    rng_state: u64,
188}
189
190impl NautilusGenerator {
191    /// Create a new NAUTILUS generator
192    #[must_use]
193    pub fn new(language: Language, max_depth: usize) -> Self {
194        Self {
195            corpus: Vec::new(),
196            global_coverage: CoverageMap::new(),
197            max_corpus_size: 1000,
198            max_depth,
199            language,
200            seed: 42,
201            rng_state: 42,
202        }
203    }
204
205    /// Set the random seed
206    #[must_use]
207    pub fn with_seed(mut self, seed: u64) -> Self {
208        self.seed = seed;
209        self.rng_state = seed;
210        self
211    }
212
213    /// Set maximum corpus size
214    #[must_use]
215    pub fn with_max_corpus(mut self, size: usize) -> Self {
216        self.max_corpus_size = size;
217        self
218    }
219
220    /// Simple xorshift64 PRNG
221    fn next_random(&mut self) -> u64 {
222        let mut x = self.rng_state;
223        x ^= x << 13;
224        x ^= x >> 7;
225        x ^= x << 17;
226        self.rng_state = x;
227        x
228    }
229
230    /// Random float in [0, 1)
231    fn random_float(&mut self) -> f64 {
232        (self.next_random() as f64) / (u64::MAX as f64)
233    }
234
235    /// Initialize corpus with seed inputs from exhaustive enumeration
236    pub fn initialize_corpus(&mut self) {
237        let enumerator = PythonEnumerator::new(self.max_depth.min(2));
238        let seeds = enumerator.enumerate_programs();
239
240        for program in seeds.into_iter().take(self.max_corpus_size / 10) {
241            let coverage = self.compute_coverage(&program.code);
242            if coverage.has_new_coverage(&self.global_coverage) {
243                self.global_coverage.merge(&coverage);
244                self.corpus.push(CorpusEntry::new(program, coverage));
245            }
246        }
247    }
248
249    /// Initialize corpus with ASTs for mutation
250    pub fn initialize_corpus_with_ast(&mut self) {
251        let enumerator = PythonEnumerator::new(self.max_depth.min(2));
252
253        // Get statements with their ASTs
254        for stmt in enumerator.enumerate_statements(self.max_depth.min(2)) {
255            let code = stmt.to_code(0);
256            let ast_depth = stmt.depth();
257
258            let program = GeneratedCode {
259                code: code.clone(),
260                language: self.language,
261                ast_depth,
262                features: Self::extract_features(&stmt),
263            };
264
265            let coverage = self.compute_coverage_from_ast(&stmt);
266
267            if coverage.has_new_coverage(&self.global_coverage) {
268                self.global_coverage.merge(&coverage);
269                self.corpus
270                    .push(CorpusEntry::with_ast(program, coverage, stmt));
271            }
272
273            if self.corpus.len() >= self.max_corpus_size / 10 {
274                break;
275            }
276        }
277    }
278
279    /// Extract features from an AST node
280    fn extract_features(node: &PythonNode) -> Vec<String> {
281        let mut features = Vec::new();
282
283        match node {
284            PythonNode::IntLit(_) => features.push("literal".to_string()),
285            PythonNode::FloatLit(_) => {
286                features.push("literal".to_string());
287                features.push("float".to_string());
288            }
289            PythonNode::StrLit(_) => {
290                features.push("literal".to_string());
291                features.push("string".to_string());
292            }
293            PythonNode::BoolLit(_) => {
294                features.push("literal".to_string());
295                features.push("boolean".to_string());
296            }
297            PythonNode::NoneLit => {
298                features.push("literal".to_string());
299                features.push("none".to_string());
300            }
301            PythonNode::Name(_) => features.push("variable".to_string()),
302            PythonNode::BinOp { op, .. } => {
303                features.push("binary_op".to_string());
304                features.push(format!("op_{}", op.to_str()));
305            }
306            PythonNode::UnaryOp { op, .. } => {
307                features.push("unary_op".to_string());
308                features.push(format!("op_{}", op.to_str()));
309            }
310            PythonNode::Compare { op, .. } => {
311                features.push("comparison".to_string());
312                features.push(format!("cmp_{}", op.to_str()));
313            }
314            PythonNode::Assign { .. } => features.push("assignment".to_string()),
315            PythonNode::Return(_) => features.push("return".to_string()),
316            PythonNode::If { orelse, .. } => {
317                features.push("conditional".to_string());
318                if !orelse.is_empty() {
319                    features.push("else_branch".to_string());
320                }
321            }
322            PythonNode::While { .. } => {
323                features.push("loop".to_string());
324                features.push("while_loop".to_string());
325            }
326            PythonNode::For { .. } => {
327                features.push("loop".to_string());
328                features.push("for_loop".to_string());
329            }
330            PythonNode::FuncDef { .. } => features.push("function_def".to_string()),
331            PythonNode::Call { .. } => features.push("function_call".to_string()),
332            PythonNode::List(_) => {
333                features.push("collection".to_string());
334                features.push("list".to_string());
335            }
336            PythonNode::Module(_) => features.push("module".to_string()),
337            PythonNode::Pass => features.push("pass".to_string()),
338            PythonNode::Break => {
339                features.push("control_flow".to_string());
340                features.push("break".to_string());
341            }
342            PythonNode::Continue => {
343                features.push("control_flow".to_string());
344                features.push("continue".to_string());
345            }
346        }
347
348        features
349    }
350
351    /// Compute coverage from source code (static analysis)
352    fn compute_coverage(&self, code: &str) -> CoverageMap {
353        let mut coverage = CoverageMap::new();
354
355        // Simple lexical coverage analysis
356        if code.contains("def ") {
357            coverage.record_node("function_def");
358        }
359        if code.contains("if ") {
360            coverage.record_node("if_stmt");
361        }
362        if code.contains("while ") {
363            coverage.record_node("while_stmt");
364        }
365        if code.contains("for ") {
366            coverage.record_node("for_stmt");
367        }
368        if code.contains("return ") || code.contains("return\n") {
369            coverage.record_node("return_stmt");
370        }
371        if code.contains(" = ") {
372            coverage.record_node("assignment");
373        }
374        if code.contains('+') || code.contains('-') || code.contains('*') || code.contains('/') {
375            coverage.record_node("binary_op");
376        }
377        if code.contains('[') {
378            coverage.record_node("list");
379        }
380
381        coverage
382    }
383
384    /// Compute coverage from AST (more precise)
385    fn compute_coverage_from_ast(&self, node: &PythonNode) -> CoverageMap {
386        let mut coverage = CoverageMap::new();
387        self.visit_ast_for_coverage(node, None, &mut coverage);
388        coverage
389    }
390
391    /// Recursively visit AST nodes to compute coverage
392    fn visit_ast_for_coverage(
393        &self,
394        node: &PythonNode,
395        parent: Option<&str>,
396        coverage: &mut CoverageMap,
397    ) {
398        let node_type = Self::node_type_name(node);
399        coverage.record_node(&node_type);
400
401        if let Some(p) = parent {
402            coverage.record_path(p, &node_type);
403        }
404
405        for feature in Self::extract_features(node) {
406            coverage.record_feature(&feature);
407        }
408
409        self.visit_children(node, &node_type, coverage);
410    }
411
412    fn visit_children(&self, node: &PythonNode, node_type: &str, coverage: &mut CoverageMap) {
413        match node {
414            PythonNode::Module(stmts) => {
415                for stmt in stmts {
416                    self.visit_ast_for_coverage(stmt, Some(node_type), coverage);
417                }
418            }
419            PythonNode::BinOp { left, right, .. } => {
420                self.visit_ast_for_coverage(left, Some(node_type), coverage);
421                self.visit_ast_for_coverage(right, Some(node_type), coverage);
422            }
423            PythonNode::UnaryOp { operand, .. } => {
424                self.visit_ast_for_coverage(operand, Some(node_type), coverage);
425            }
426            PythonNode::Compare { left, right, .. } => {
427                self.visit_ast_for_coverage(left, Some(node_type), coverage);
428                self.visit_ast_for_coverage(right, Some(node_type), coverage);
429            }
430            PythonNode::Assign { value, .. } => {
431                self.visit_ast_for_coverage(value, Some(node_type), coverage);
432            }
433            PythonNode::Return(Some(expr)) => {
434                self.visit_ast_for_coverage(expr, Some(node_type), coverage);
435            }
436            PythonNode::If { test, body, orelse } => {
437                self.visit_ast_for_coverage(test, Some(node_type), coverage);
438                for stmt in body {
439                    self.visit_ast_for_coverage(stmt, Some(node_type), coverage);
440                }
441                for stmt in orelse {
442                    self.visit_ast_for_coverage(stmt, Some(node_type), coverage);
443                }
444            }
445            PythonNode::While { test, body } => {
446                self.visit_ast_for_coverage(test, Some(node_type), coverage);
447                for stmt in body {
448                    self.visit_ast_for_coverage(stmt, Some(node_type), coverage);
449                }
450            }
451            PythonNode::For { iter, body, .. } => {
452                self.visit_ast_for_coverage(iter, Some(node_type), coverage);
453                for stmt in body {
454                    self.visit_ast_for_coverage(stmt, Some(node_type), coverage);
455                }
456            }
457            PythonNode::FuncDef { body, .. } => {
458                for stmt in body {
459                    self.visit_ast_for_coverage(stmt, Some(node_type), coverage);
460                }
461            }
462            PythonNode::Call { args, .. } => {
463                for arg in args {
464                    self.visit_ast_for_coverage(arg, Some(node_type), coverage);
465                }
466            }
467            PythonNode::List(items) => {
468                for item in items {
469                    self.visit_ast_for_coverage(item, Some(node_type), coverage);
470                }
471            }
472            PythonNode::IntLit(_)
473            | PythonNode::FloatLit(_)
474            | PythonNode::StrLit(_)
475            | PythonNode::BoolLit(_)
476            | PythonNode::NoneLit
477            | PythonNode::Name(_)
478            | PythonNode::Return(None)
479            | PythonNode::Pass
480            | PythonNode::Break
481            | PythonNode::Continue => {}
482        }
483    }
484
485    /// Get the type name of a node
486    fn node_type_name(node: &PythonNode) -> String {
487        match node {
488            PythonNode::Module(_) => "Module".to_string(),
489            PythonNode::IntLit(_) => "IntLit".to_string(),
490            PythonNode::FloatLit(_) => "FloatLit".to_string(),
491            PythonNode::StrLit(_) => "StrLit".to_string(),
492            PythonNode::BoolLit(_) => "BoolLit".to_string(),
493            PythonNode::NoneLit => "NoneLit".to_string(),
494            PythonNode::Name(_) => "Name".to_string(),
495            PythonNode::BinOp { op, .. } => format!("BinOp_{}", op.to_str()),
496            PythonNode::UnaryOp { op, .. } => format!("UnaryOp_{}", op.to_str()),
497            PythonNode::Compare { op, .. } => format!("Compare_{}", op.to_str()),
498            PythonNode::Assign { .. } => "Assign".to_string(),
499            PythonNode::Return(_) => "Return".to_string(),
500            PythonNode::If { .. } => "If".to_string(),
501            PythonNode::While { .. } => "While".to_string(),
502            PythonNode::For { .. } => "For".to_string(),
503            PythonNode::FuncDef { .. } => "FuncDef".to_string(),
504            PythonNode::Call { .. } => "Call".to_string(),
505            PythonNode::List(_) => "List".to_string(),
506            PythonNode::Pass => "Pass".to_string(),
507            PythonNode::Break => "Break".to_string(),
508            PythonNode::Continue => "Continue".to_string(),
509        }
510    }
511
512    /// Select and mark an entry (mutable version for updating selection count)
513    fn select_entry_mut(&mut self) -> Option<usize> {
514        if self.corpus.is_empty() {
515            return None;
516        }
517
518        let total_energy: f64 = self.corpus.iter().map(|e| e.energy).sum();
519        if total_energy <= 0.0 {
520            let idx = (self.next_random() as usize) % self.corpus.len();
521            self.corpus[idx].selection_count += 1;
522            return Some(idx);
523        }
524
525        let mut threshold = self.random_float() * total_energy;
526        for (i, entry) in self.corpus.iter_mut().enumerate() {
527            threshold -= entry.energy;
528            if threshold <= 0.0 {
529                entry.selection_count += 1;
530                return Some(i);
531            }
532        }
533
534        let last_idx = self.corpus.len() - 1;
535        self.corpus[last_idx].selection_count += 1;
536        Some(last_idx)
537    }
538
539    /// Add a new entry to the corpus if it has new coverage
540    pub fn add_to_corpus(&mut self, code: GeneratedCode, coverage: CoverageMap) -> bool {
541        if !coverage.has_new_coverage(&self.global_coverage) {
542            return false;
543        }
544
545        self.global_coverage.merge(&coverage);
546
547        // If corpus is full, replace lowest energy entry
548        if self.corpus.len() >= self.max_corpus_size {
549            if let Some(min_idx) = self
550                .corpus
551                .iter()
552                .enumerate()
553                .min_by(|(_, a), (_, b)| {
554                    a.energy
555                        .partial_cmp(&b.energy)
556                        .unwrap_or(std::cmp::Ordering::Equal)
557                })
558                .map(|(i, _)| i)
559            {
560                self.corpus[min_idx] = CorpusEntry::new(code, coverage);
561            }
562        } else {
563            self.corpus.push(CorpusEntry::new(code, coverage));
564        }
565
566        // Update all energies
567        let global = self.global_coverage.clone();
568        for entry in &mut self.corpus {
569            entry.update_energy(&global);
570        }
571
572        true
573    }
574
575    /// Generate programs using coverage-guided fuzzing
576    pub fn generate(&mut self, count: usize) -> Vec<GeneratedCode> {
577        // Initialize corpus if empty
578        if self.corpus.is_empty() {
579            self.initialize_corpus_with_ast();
580        }
581
582        let mut results = Vec::with_capacity(count);
583        let mut iterations = 0;
584        let max_iterations = count * 10; // Prevent infinite loops
585
586        while results.len() < count && iterations < max_iterations {
587            iterations += 1;
588
589            // Select a corpus entry
590            if let Some(idx) = self.select_entry_mut() {
591                // Get necessary info from entry before mutable operations
592                let has_ast = self.corpus[idx].ast.is_some();
593                let ast_clone = self.corpus[idx].ast.clone();
594                let code_clone = self.corpus[idx].code.clone();
595
596                // Decide whether to mutate
597                let should_mutate = self.random_float() < 0.7 && has_ast;
598
599                if should_mutate {
600                    // Mutate the cloned AST (safe: has_ast check ensures Some)
601                    if let Some(ast) = ast_clone {
602                        if let Some(mutated) = self.mutate_ast(&ast) {
603                            let code = mutated.to_code(0);
604                            let coverage = self.compute_coverage_from_ast(&mutated);
605
606                            let program = GeneratedCode {
607                                code,
608                                language: self.language,
609                                ast_depth: mutated.depth(),
610                                features: Self::extract_features(&mutated),
611                            };
612
613                            // Add to corpus if new coverage
614                            self.add_to_corpus(program.clone(), coverage);
615                            results.push(program);
616                        }
617                    }
618                } else {
619                    // Use existing entry
620                    results.push(code_clone);
621                }
622            } else {
623                // No corpus - generate fresh
624                let enumerator = PythonEnumerator::new(self.max_depth);
625                let programs = enumerator.enumerate_programs();
626                if let Some(program) = programs.into_iter().next() {
627                    results.push(program);
628                }
629            }
630        }
631
632        results
633    }
634
635    /// Mutate an AST node
636    fn mutate_ast(&mut self, node: &PythonNode) -> Option<PythonNode> {
637        let mutation_type = self.next_random() % 4;
638
639        match mutation_type {
640            0 => self.mutate_operator(node),
641            1 => self.mutate_literal(node),
642            2 => self.insert_wrapper(node),
643            _ => self.delete_subtree(node),
644        }
645    }
646
647    /// Mutate operators in the AST
648    fn mutate_operator(&mut self, node: &PythonNode) -> Option<PythonNode> {
649        match node {
650            PythonNode::BinOp { left, right, .. } => {
651                let ops = BinaryOp::all();
652                let new_op = ops[(self.next_random() as usize) % ops.len()];
653                Some(PythonNode::BinOp {
654                    left: left.clone(),
655                    op: new_op,
656                    right: right.clone(),
657                })
658            }
659            PythonNode::Compare { left, right, .. } => {
660                let ops = CompareOp::all();
661                let new_op = ops[(self.next_random() as usize) % ops.len()];
662                Some(PythonNode::Compare {
663                    left: left.clone(),
664                    op: new_op,
665                    right: right.clone(),
666                })
667            }
668            PythonNode::UnaryOp { operand, .. } => {
669                let ops = UnaryOp::all();
670                let new_op = ops[(self.next_random() as usize) % ops.len()];
671                Some(PythonNode::UnaryOp {
672                    op: new_op,
673                    operand: operand.clone(),
674                })
675            }
676            _ => None,
677        }
678    }
679
680    /// Mutate literals in the AST
681    fn mutate_literal(&mut self, node: &PythonNode) -> Option<PythonNode> {
682        match node {
683            PythonNode::IntLit(n) => {
684                let mutations = [0, 1, -1, i64::MAX, i64::MIN, *n + 1, n.saturating_sub(1)];
685                let new_val = mutations[(self.next_random() as usize) % mutations.len()];
686                Some(PythonNode::IntLit(new_val))
687            }
688            PythonNode::BoolLit(b) => Some(PythonNode::BoolLit(!b)),
689            PythonNode::StrLit(s) => {
690                let mutations = ["", " ", "\\n", "\\t", &format!("{s}x")];
691                let new_val = mutations[(self.next_random() as usize) % mutations.len()];
692                Some(PythonNode::StrLit(new_val.to_string()))
693            }
694            _ => None,
695        }
696    }
697
698    /// Insert a wrapper around an expression
699    fn insert_wrapper(&mut self, node: &PythonNode) -> Option<PythonNode> {
700        // Only wrap expressions
701        match node {
702            PythonNode::IntLit(_)
703            | PythonNode::FloatLit(_)
704            | PythonNode::Name(_)
705            | PythonNode::BinOp { .. } => {
706                let ops = UnaryOp::all();
707                let op = ops[(self.next_random() as usize) % ops.len()];
708                Some(PythonNode::UnaryOp {
709                    op,
710                    operand: Box::new(node.clone()),
711                })
712            }
713            _ => None,
714        }
715    }
716
717    /// Delete/simplify a subtree
718    fn delete_subtree(&mut self, node: &PythonNode) -> Option<PythonNode> {
719        match node {
720            PythonNode::BinOp { left, right, .. } => {
721                // Replace binary op with one of its operands
722                if self.random_float() < 0.5 {
723                    Some((**left).clone())
724                } else {
725                    Some((**right).clone())
726                }
727            }
728            PythonNode::UnaryOp { operand, .. } => Some((**operand).clone()),
729            PythonNode::If { body, .. } => {
730                // Simplify to just the first statement in body
731                body.first().cloned()
732            }
733            _ => None,
734        }
735    }
736
737    /// Get current corpus size
738    #[must_use]
739    pub fn corpus_size(&self) -> usize {
740        self.corpus.len()
741    }
742
743    /// Get global coverage statistics
744    #[must_use]
745    pub fn coverage_stats(&self) -> CoverageStats {
746        CoverageStats {
747            total_coverage: self.global_coverage.coverage_count(),
748            node_types_covered: self.global_coverage.node_types.len(),
749            ast_paths_covered: self.global_coverage.ast_paths.len(),
750            features_covered: self.global_coverage.feature_combos.len(),
751            corpus_size: self.corpus.len(),
752        }
753    }
754}
755
756/// Statistics about coverage
757#[derive(Debug, Clone)]
758pub struct CoverageStats {
759    /// Total coverage count
760    pub total_coverage: usize,
761    /// Number of unique node types covered
762    pub node_types_covered: usize,
763    /// Number of unique AST paths covered
764    pub ast_paths_covered: usize,
765    /// Number of unique features covered
766    pub features_covered: usize,
767    /// Current corpus size
768    pub corpus_size: usize,
769}
770
771#[cfg(test)]
772mod tests {
773    use super::*;
774
775    #[test]
776    fn test_coverage_map_new() {
777        let coverage = CoverageMap::new();
778        assert_eq!(coverage.coverage_count(), 0);
779    }
780
781    #[test]
782    fn test_coverage_map_record_node() {
783        let mut coverage = CoverageMap::new();
784        coverage.record_node("function_def");
785        coverage.record_node("if_stmt");
786
787        assert!(coverage.node_types().contains("function_def"));
788        assert!(coverage.node_types().contains("if_stmt"));
789        assert_eq!(coverage.node_types().len(), 2);
790    }
791
792    #[test]
793    fn test_coverage_map_record_path() {
794        let mut coverage = CoverageMap::new();
795        coverage.record_path("function_def", "return_stmt");
796
797        assert!(coverage
798            .ast_paths()
799            .contains(&("function_def".to_string(), "return_stmt".to_string())));
800    }
801
802    #[test]
803    fn test_coverage_map_record_feature() {
804        let mut coverage = CoverageMap::new();
805        coverage.record_feature("loops");
806        coverage.record_feature("conditionals");
807        assert_eq!(coverage.coverage_count(), 2);
808    }
809
810    #[test]
811    fn test_coverage_map_has_new_coverage() {
812        let mut existing = CoverageMap::new();
813        existing.record_node("function_def");
814
815        let mut new_coverage = CoverageMap::new();
816        new_coverage.record_node("function_def");
817        assert!(!new_coverage.has_new_coverage(&existing));
818
819        new_coverage.record_node("while_stmt");
820        assert!(new_coverage.has_new_coverage(&existing));
821    }
822
823    #[test]
824    fn test_coverage_map_has_new_path_coverage() {
825        let mut existing = CoverageMap::new();
826        existing.record_path("a", "b");
827
828        let mut new_coverage = CoverageMap::new();
829        new_coverage.record_path("a", "b");
830        assert!(!new_coverage.has_new_coverage(&existing));
831
832        new_coverage.record_path("a", "c");
833        assert!(new_coverage.has_new_coverage(&existing));
834    }
835
836    #[test]
837    fn test_coverage_map_has_new_feature_coverage() {
838        let mut existing = CoverageMap::new();
839        existing.record_feature("feat1");
840
841        let mut new_coverage = CoverageMap::new();
842        new_coverage.record_feature("feat1");
843        assert!(!new_coverage.has_new_coverage(&existing));
844
845        new_coverage.record_feature("feat2");
846        assert!(new_coverage.has_new_coverage(&existing));
847    }
848
849    #[test]
850    fn test_coverage_map_merge() {
851        let mut map1 = CoverageMap::new();
852        map1.record_node("a");
853        map1.record_node("b");
854
855        let mut map2 = CoverageMap::new();
856        map2.record_node("b");
857        map2.record_node("c");
858
859        map1.merge(&map2);
860        assert_eq!(map1.node_types().len(), 3);
861        assert!(map1.node_types().contains("a"));
862        assert!(map1.node_types().contains("b"));
863        assert!(map1.node_types().contains("c"));
864    }
865
866    #[test]
867    fn test_coverage_map_default() {
868        let coverage = CoverageMap::default();
869        assert_eq!(coverage.coverage_count(), 0);
870    }
871
872    #[test]
873    fn test_coverage_map_debug() {
874        let mut coverage = CoverageMap::new();
875        coverage.record_node("test");
876        let debug = format!("{:?}", coverage);
877        assert!(debug.contains("CoverageMap"));
878    }
879
880    #[test]
881    fn test_coverage_map_clone() {
882        let mut coverage = CoverageMap::new();
883        coverage.record_node("test");
884        let cloned = coverage.clone();
885        assert_eq!(cloned.coverage_count(), coverage.coverage_count());
886    }
887
888    #[test]
889    fn test_corpus_entry_new() {
890        let code = GeneratedCode {
891            code: "x = 1".to_string(),
892            language: Language::Python,
893            ast_depth: 1,
894            features: vec!["assignment".to_string()],
895        };
896        let coverage = CoverageMap::new();
897        let entry = CorpusEntry::new(code, coverage);
898
899        assert_eq!(entry.energy, 1.0);
900        assert_eq!(entry.selection_count, 0);
901        assert!(entry.ast.is_none());
902    }
903
904    #[test]
905    fn test_corpus_entry_with_ast() {
906        let code = GeneratedCode {
907            code: "x = 1".to_string(),
908            language: Language::Python,
909            ast_depth: 1,
910            features: vec!["assignment".to_string()],
911        };
912        let coverage = CoverageMap::new();
913        let ast = PythonNode::IntLit(1);
914        let entry = CorpusEntry::with_ast(code, coverage, ast);
915
916        assert!(entry.ast.is_some());
917    }
918
919    #[test]
920    fn test_corpus_entry_update_energy() {
921        let code = GeneratedCode {
922            code: "x = 1".to_string(),
923            language: Language::Python,
924            ast_depth: 1,
925            features: vec![],
926        };
927        let mut coverage = CoverageMap::new();
928        coverage.record_node("unique_node");
929        let mut entry = CorpusEntry::new(code, coverage);
930
931        let global = CoverageMap::new();
932        entry.update_energy(&global);
933
934        // Energy should be boosted because of unique coverage
935        assert!(entry.energy > 1.0);
936    }
937
938    #[test]
939    fn test_corpus_entry_debug() {
940        let code = GeneratedCode {
941            code: "x = 1".to_string(),
942            language: Language::Python,
943            ast_depth: 1,
944            features: vec![],
945        };
946        let coverage = CoverageMap::new();
947        let entry = CorpusEntry::new(code, coverage);
948        let debug = format!("{:?}", entry);
949        assert!(debug.contains("CorpusEntry"));
950    }
951
952    #[test]
953    fn test_corpus_entry_clone() {
954        let code = GeneratedCode {
955            code: "x = 1".to_string(),
956            language: Language::Python,
957            ast_depth: 1,
958            features: vec![],
959        };
960        let coverage = CoverageMap::new();
961        let entry = CorpusEntry::new(code, coverage);
962        let cloned = entry.clone();
963        assert_eq!(cloned.energy, entry.energy);
964    }
965
966    #[test]
967    fn test_nautilus_generator_new() {
968        let gen = NautilusGenerator::new(Language::Python, 3);
969        assert_eq!(gen.corpus_size(), 0);
970        assert_eq!(gen.language, Language::Python);
971    }
972
973    #[test]
974    fn test_nautilus_generator_with_max_corpus() {
975        let gen = NautilusGenerator::new(Language::Python, 2).with_max_corpus(500);
976        assert_eq!(gen.max_corpus_size, 500);
977    }
978
979    #[test]
980    fn test_nautilus_generator_initialize_corpus() {
981        let mut gen = NautilusGenerator::new(Language::Python, 2);
982        gen.initialize_corpus();
983
984        assert!(gen.corpus_size() > 0, "Corpus should be initialized");
985    }
986
987    #[test]
988    fn test_nautilus_generator_generate() {
989        let mut gen = NautilusGenerator::new(Language::Python, 2).with_seed(42);
990        let programs = gen.generate(5);
991
992        assert!(!programs.is_empty(), "Should generate programs");
993        for prog in &programs {
994            assert_eq!(prog.language, Language::Python);
995        }
996    }
997
998    #[test]
999    fn test_nautilus_generator_coverage_stats() {
1000        let mut gen = NautilusGenerator::new(Language::Python, 2);
1001        gen.initialize_corpus_with_ast();
1002
1003        let stats = gen.coverage_stats();
1004        assert!(stats.node_types_covered > 0, "Should cover some node types");
1005        assert!(stats.corpus_size > 0, "Corpus should have entries");
1006    }
1007
1008    #[test]
1009    fn test_nautilus_generator_with_seed() {
1010        let mut gen1 = NautilusGenerator::new(Language::Python, 2).with_seed(123);
1011        let mut gen2 = NautilusGenerator::new(Language::Python, 2).with_seed(123);
1012
1013        gen1.initialize_corpus_with_ast();
1014        gen2.initialize_corpus_with_ast();
1015
1016        // Both should have same corpus size with same seed
1017        assert_eq!(gen1.corpus_size(), gen2.corpus_size());
1018    }
1019
1020    #[test]
1021    fn test_add_to_corpus_new_coverage() {
1022        let mut gen = NautilusGenerator::new(Language::Python, 2);
1023
1024        let code = GeneratedCode {
1025            code: "def foo(): pass".to_string(),
1026            language: Language::Python,
1027            ast_depth: 2,
1028            features: vec!["function_def".to_string()],
1029        };
1030
1031        let mut coverage = CoverageMap::new();
1032        coverage.record_node("unique_node_type");
1033
1034        let added = gen.add_to_corpus(code, coverage);
1035        assert!(added, "Should add entry with new coverage");
1036        assert_eq!(gen.corpus_size(), 1);
1037    }
1038
1039    #[test]
1040    fn test_add_to_corpus_duplicate_coverage() {
1041        let mut gen = NautilusGenerator::new(Language::Python, 2);
1042
1043        // Add first entry
1044        let code1 = GeneratedCode {
1045            code: "x = 1".to_string(),
1046            language: Language::Python,
1047            ast_depth: 1,
1048            features: vec![],
1049        };
1050        let mut coverage1 = CoverageMap::new();
1051        coverage1.record_node("assignment");
1052        gen.add_to_corpus(code1, coverage1);
1053
1054        // Try to add entry with same coverage
1055        let code2 = GeneratedCode {
1056            code: "y = 2".to_string(),
1057            language: Language::Python,
1058            ast_depth: 1,
1059            features: vec![],
1060        };
1061        let mut coverage2 = CoverageMap::new();
1062        coverage2.record_node("assignment");
1063
1064        let added = gen.add_to_corpus(code2, coverage2);
1065        assert!(!added, "Should not add entry with duplicate coverage");
1066        assert_eq!(gen.corpus_size(), 1);
1067    }
1068
1069    #[test]
1070    fn test_extract_features_literals() {
1071        let int_features = NautilusGenerator::extract_features(&PythonNode::IntLit(42));
1072        assert!(int_features.contains(&"literal".to_string()));
1073
1074        let float_features = NautilusGenerator::extract_features(&PythonNode::FloatLit(3.14));
1075        assert!(float_features.contains(&"float".to_string()));
1076
1077        let str_features =
1078            NautilusGenerator::extract_features(&PythonNode::StrLit("hello".to_string()));
1079        assert!(str_features.contains(&"string".to_string()));
1080
1081        let bool_features = NautilusGenerator::extract_features(&PythonNode::BoolLit(true));
1082        assert!(bool_features.contains(&"boolean".to_string()));
1083
1084        let none_features = NautilusGenerator::extract_features(&PythonNode::NoneLit);
1085        assert!(none_features.contains(&"none".to_string()));
1086    }
1087
1088    #[test]
1089    fn test_extract_features_control_flow() {
1090        let break_features = NautilusGenerator::extract_features(&PythonNode::Break);
1091        assert!(break_features.contains(&"control_flow".to_string()));
1092        assert!(break_features.contains(&"break".to_string()));
1093
1094        let continue_features = NautilusGenerator::extract_features(&PythonNode::Continue);
1095        assert!(continue_features.contains(&"continue".to_string()));
1096    }
1097
1098    #[test]
1099    fn test_extract_features_loops() {
1100        let while_node = PythonNode::While {
1101            test: Box::new(PythonNode::BoolLit(true)),
1102            body: vec![PythonNode::Pass],
1103        };
1104        let while_features = NautilusGenerator::extract_features(&while_node);
1105        assert!(while_features.contains(&"loop".to_string()));
1106        assert!(while_features.contains(&"while_loop".to_string()));
1107
1108        let for_node = PythonNode::For {
1109            target: "i".to_string(),
1110            iter: Box::new(PythonNode::List(vec![])),
1111            body: vec![PythonNode::Pass],
1112        };
1113        let for_features = NautilusGenerator::extract_features(&for_node);
1114        assert!(for_features.contains(&"for_loop".to_string()));
1115    }
1116
1117    #[test]
1118    fn test_extract_features_collections() {
1119        let list_features = NautilusGenerator::extract_features(&PythonNode::List(vec![]));
1120        assert!(list_features.contains(&"collection".to_string()));
1121        assert!(list_features.contains(&"list".to_string()));
1122    }
1123
1124    #[test]
1125    fn test_extract_features_functions() {
1126        let func_node = PythonNode::FuncDef {
1127            name: "foo".to_string(),
1128            args: vec![],
1129            body: vec![PythonNode::Pass],
1130        };
1131        let func_features = NautilusGenerator::extract_features(&func_node);
1132        assert!(func_features.contains(&"function_def".to_string()));
1133
1134        let call_node = PythonNode::Call {
1135            func: "print".to_string(),
1136            args: vec![],
1137        };
1138        let call_features = NautilusGenerator::extract_features(&call_node);
1139        assert!(call_features.contains(&"function_call".to_string()));
1140    }
1141
1142    #[test]
1143    fn test_extract_features_if_with_else() {
1144        let if_node = PythonNode::If {
1145            test: Box::new(PythonNode::BoolLit(true)),
1146            body: vec![PythonNode::Pass],
1147            orelse: vec![PythonNode::Pass],
1148        };
1149        let features = NautilusGenerator::extract_features(&if_node);
1150        assert!(features.contains(&"conditional".to_string()));
1151        assert!(features.contains(&"else_branch".to_string()));
1152    }
1153
1154    #[test]
1155    fn test_node_type_name() {
1156        assert_eq!(
1157            NautilusGenerator::node_type_name(&PythonNode::IntLit(1)),
1158            "IntLit"
1159        );
1160        assert_eq!(
1161            NautilusGenerator::node_type_name(&PythonNode::FloatLit(1.0)),
1162            "FloatLit"
1163        );
1164        assert_eq!(
1165            NautilusGenerator::node_type_name(&PythonNode::StrLit("x".to_string())),
1166            "StrLit"
1167        );
1168        assert_eq!(
1169            NautilusGenerator::node_type_name(&PythonNode::BoolLit(true)),
1170            "BoolLit"
1171        );
1172        assert_eq!(
1173            NautilusGenerator::node_type_name(&PythonNode::NoneLit),
1174            "NoneLit"
1175        );
1176        assert_eq!(
1177            NautilusGenerator::node_type_name(&PythonNode::Name("x".to_string())),
1178            "Name"
1179        );
1180        assert_eq!(NautilusGenerator::node_type_name(&PythonNode::Pass), "Pass");
1181        assert_eq!(
1182            NautilusGenerator::node_type_name(&PythonNode::Break),
1183            "Break"
1184        );
1185        assert_eq!(
1186            NautilusGenerator::node_type_name(&PythonNode::Continue),
1187            "Continue"
1188        );
1189    }
1190
1191    #[test]
1192    fn test_node_type_name_operators() {
1193        let binop = PythonNode::BinOp {
1194            left: Box::new(PythonNode::IntLit(1)),
1195            op: BinaryOp::Add,
1196            right: Box::new(PythonNode::IntLit(2)),
1197        };
1198        assert!(NautilusGenerator::node_type_name(&binop).starts_with("BinOp_"));
1199
1200        let unaryop = PythonNode::UnaryOp {
1201            op: UnaryOp::Neg,
1202            operand: Box::new(PythonNode::IntLit(1)),
1203        };
1204        assert!(NautilusGenerator::node_type_name(&unaryop).starts_with("UnaryOp_"));
1205
1206        let compare = PythonNode::Compare {
1207            left: Box::new(PythonNode::IntLit(1)),
1208            op: CompareOp::Lt,
1209            right: Box::new(PythonNode::IntLit(2)),
1210        };
1211        assert!(NautilusGenerator::node_type_name(&compare).starts_with("Compare_"));
1212    }
1213
1214    #[test]
1215    fn test_node_type_name_statements() {
1216        let assign = PythonNode::Assign {
1217            target: "x".to_string(),
1218            value: Box::new(PythonNode::IntLit(1)),
1219        };
1220        assert_eq!(NautilusGenerator::node_type_name(&assign), "Assign");
1221
1222        let ret = PythonNode::Return(Some(Box::new(PythonNode::IntLit(1))));
1223        assert_eq!(NautilusGenerator::node_type_name(&ret), "Return");
1224
1225        let if_node = PythonNode::If {
1226            test: Box::new(PythonNode::BoolLit(true)),
1227            body: vec![],
1228            orelse: vec![],
1229        };
1230        assert_eq!(NautilusGenerator::node_type_name(&if_node), "If");
1231
1232        let while_node = PythonNode::While {
1233            test: Box::new(PythonNode::BoolLit(true)),
1234            body: vec![],
1235        };
1236        assert_eq!(NautilusGenerator::node_type_name(&while_node), "While");
1237
1238        let for_node = PythonNode::For {
1239            target: "i".to_string(),
1240            iter: Box::new(PythonNode::List(vec![])),
1241            body: vec![],
1242        };
1243        assert_eq!(NautilusGenerator::node_type_name(&for_node), "For");
1244
1245        let func = PythonNode::FuncDef {
1246            name: "f".to_string(),
1247            args: vec![],
1248            body: vec![],
1249        };
1250        assert_eq!(NautilusGenerator::node_type_name(&func), "FuncDef");
1251
1252        let call = PythonNode::Call {
1253            func: "f".to_string(),
1254            args: vec![],
1255        };
1256        assert_eq!(NautilusGenerator::node_type_name(&call), "Call");
1257
1258        let list = PythonNode::List(vec![]);
1259        assert_eq!(NautilusGenerator::node_type_name(&list), "List");
1260
1261        let module = PythonNode::Module(vec![]);
1262        assert_eq!(NautilusGenerator::node_type_name(&module), "Module");
1263    }
1264
1265    #[test]
1266    fn test_coverage_stats_debug() {
1267        let stats = CoverageStats {
1268            total_coverage: 10,
1269            node_types_covered: 5,
1270            ast_paths_covered: 3,
1271            features_covered: 2,
1272            corpus_size: 100,
1273        };
1274        let debug = format!("{:?}", stats);
1275        assert!(debug.contains("CoverageStats"));
1276    }
1277
1278    #[test]
1279    fn test_coverage_stats_clone() {
1280        let stats = CoverageStats {
1281            total_coverage: 10,
1282            node_types_covered: 5,
1283            ast_paths_covered: 3,
1284            features_covered: 2,
1285            corpus_size: 100,
1286        };
1287        let cloned = stats.clone();
1288        assert_eq!(cloned.total_coverage, stats.total_coverage);
1289    }
1290
1291    #[test]
1292    fn test_nautilus_generator_debug() {
1293        let gen = NautilusGenerator::new(Language::Python, 2);
1294        let debug = format!("{:?}", gen);
1295        assert!(debug.contains("NautilusGenerator"));
1296    }
1297
1298    #[test]
1299    fn test_add_to_corpus_full() {
1300        let mut gen = NautilusGenerator::new(Language::Python, 2).with_max_corpus(2);
1301
1302        // Add entries up to max
1303        for i in 0..3 {
1304            let code = GeneratedCode {
1305                code: format!("x = {i}"),
1306                language: Language::Python,
1307                ast_depth: 1,
1308                features: vec![],
1309            };
1310            let mut coverage = CoverageMap::new();
1311            coverage.record_node(&format!("unique_node_{i}"));
1312            gen.add_to_corpus(code, coverage);
1313        }
1314
1315        // Corpus should stay at max size
1316        assert!(gen.corpus_size() <= 2);
1317    }
1318
1319    #[test]
1320    fn test_extract_features_binary_op() {
1321        let node = PythonNode::BinOp {
1322            left: Box::new(PythonNode::IntLit(1)),
1323            op: BinaryOp::Add,
1324            right: Box::new(PythonNode::IntLit(2)),
1325        };
1326        let features = NautilusGenerator::extract_features(&node);
1327        assert!(features.contains(&"binary_op".to_string()));
1328        assert!(features.iter().any(|f| f.starts_with("op_")));
1329    }
1330
1331    #[test]
1332    fn test_extract_features_unary_op() {
1333        let node = PythonNode::UnaryOp {
1334            op: UnaryOp::Neg,
1335            operand: Box::new(PythonNode::IntLit(1)),
1336        };
1337        let features = NautilusGenerator::extract_features(&node);
1338        assert!(features.contains(&"unary_op".to_string()));
1339        assert!(features.iter().any(|f| f.starts_with("op_")));
1340    }
1341
1342    #[test]
1343    fn test_extract_features_compare() {
1344        let node = PythonNode::Compare {
1345            left: Box::new(PythonNode::IntLit(1)),
1346            op: CompareOp::Lt,
1347            right: Box::new(PythonNode::IntLit(2)),
1348        };
1349        let features = NautilusGenerator::extract_features(&node);
1350        assert!(features.contains(&"comparison".to_string()));
1351        assert!(features.iter().any(|f| f.starts_with("cmp_")));
1352    }
1353
1354    #[test]
1355    fn test_extract_features_module() {
1356        let node = PythonNode::Module(vec![PythonNode::Pass]);
1357        let features = NautilusGenerator::extract_features(&node);
1358        assert!(features.contains(&"module".to_string()));
1359    }
1360
1361    #[test]
1362    fn test_compute_coverage() {
1363        let gen = NautilusGenerator::new(Language::Python, 2);
1364
1365        let coverage = gen.compute_coverage("def foo(): pass");
1366        assert!(coverage.node_types().contains("function_def"));
1367
1368        let coverage2 = gen.compute_coverage("if x: pass");
1369        assert!(coverage2.node_types().contains("if_stmt"));
1370
1371        let coverage3 = gen.compute_coverage("while True: pass");
1372        assert!(coverage3.node_types().contains("while_stmt"));
1373
1374        let coverage4 = gen.compute_coverage("for i in x: pass");
1375        assert!(coverage4.node_types().contains("for_stmt"));
1376
1377        let coverage5 = gen.compute_coverage("return 1");
1378        assert!(coverage5.node_types().contains("return_stmt"));
1379
1380        let coverage6 = gen.compute_coverage("[1, 2, 3]");
1381        assert!(coverage6.node_types().contains("list"));
1382    }
1383
1384    #[test]
1385    fn test_compute_coverage_from_ast() {
1386        let gen = NautilusGenerator::new(Language::Python, 2);
1387
1388        // Test with a complex AST that exercises all branches
1389        let ast = PythonNode::Module(vec![
1390            PythonNode::Assign {
1391                target: "x".to_string(),
1392                value: Box::new(PythonNode::BinOp {
1393                    left: Box::new(PythonNode::IntLit(1)),
1394                    op: BinaryOp::Add,
1395                    right: Box::new(PythonNode::IntLit(2)),
1396                }),
1397            },
1398            PythonNode::If {
1399                test: Box::new(PythonNode::Compare {
1400                    left: Box::new(PythonNode::Name("x".to_string())),
1401                    op: CompareOp::Lt,
1402                    right: Box::new(PythonNode::IntLit(5)),
1403                }),
1404                body: vec![PythonNode::Pass],
1405                orelse: vec![PythonNode::Pass],
1406            },
1407            PythonNode::While {
1408                test: Box::new(PythonNode::BoolLit(true)),
1409                body: vec![PythonNode::Break],
1410            },
1411            PythonNode::For {
1412                target: "i".to_string(),
1413                iter: Box::new(PythonNode::List(vec![PythonNode::IntLit(1)])),
1414                body: vec![PythonNode::Continue],
1415            },
1416            PythonNode::FuncDef {
1417                name: "foo".to_string(),
1418                args: vec!["a".to_string()],
1419                body: vec![PythonNode::Return(Some(Box::new(PythonNode::Name(
1420                    "a".to_string(),
1421                ))))],
1422            },
1423            PythonNode::UnaryOp {
1424                op: UnaryOp::Neg,
1425                operand: Box::new(PythonNode::IntLit(1)),
1426            },
1427            PythonNode::Call {
1428                func: "print".to_string(),
1429                args: vec![PythonNode::StrLit("hello".to_string())],
1430            },
1431        ]);
1432
1433        let coverage = gen.compute_coverage_from_ast(&ast);
1434
1435        assert!(coverage.node_types().contains("Module"));
1436        assert!(coverage.node_types().contains("Assign"));
1437        assert!(coverage.node_types().contains("If"));
1438        assert!(coverage.node_types().contains("While"));
1439        assert!(coverage.node_types().contains("For"));
1440        assert!(coverage.node_types().contains("FuncDef"));
1441        assert!(coverage.node_types().contains("Call"));
1442    }
1443
1444    #[test]
1445    fn test_select_entry_mut_empty() {
1446        let mut gen = NautilusGenerator::new(Language::Python, 2);
1447        let result = gen.select_entry_mut();
1448        assert!(result.is_none());
1449    }
1450
1451    #[test]
1452    fn test_select_entry_mut_zero_energy() {
1453        let mut gen = NautilusGenerator::new(Language::Python, 2).with_seed(42);
1454
1455        // Add an entry with zero energy manually
1456        let code = GeneratedCode {
1457            code: "x = 1".to_string(),
1458            language: Language::Python,
1459            ast_depth: 1,
1460            features: vec![],
1461        };
1462        let coverage = CoverageMap::new();
1463        let mut entry = CorpusEntry::new(code, coverage);
1464        entry.energy = 0.0;
1465        gen.corpus.push(entry);
1466
1467        // Should still select via random index
1468        let result = gen.select_entry_mut();
1469        assert!(result.is_some());
1470    }
1471
1472    #[test]
1473    fn test_generate_covers_mutations() {
1474        let mut gen = NautilusGenerator::new(Language::Python, 2).with_seed(42);
1475
1476        // Initialize with AST corpus
1477        gen.initialize_corpus_with_ast();
1478
1479        // Generate enough to trigger mutations
1480        let programs = gen.generate(20);
1481        assert!(!programs.is_empty());
1482    }
1483
1484    #[test]
1485    fn test_mutate_ast() {
1486        let mut gen = NautilusGenerator::new(Language::Python, 2).with_seed(42);
1487
1488        let ast = PythonNode::Assign {
1489            target: "x".to_string(),
1490            value: Box::new(PythonNode::IntLit(1)),
1491        };
1492
1493        // Mutate multiple times to exercise different mutation types
1494        for _ in 0..10 {
1495            let mutated = gen.mutate_ast(&ast);
1496            // Should produce some mutation
1497            assert!(mutated.is_some() || true); // May or may not mutate
1498        }
1499    }
1500
1501    #[test]
1502    fn test_mutate_operator() {
1503        let mut gen = NautilusGenerator::new(Language::Python, 2).with_seed(42);
1504
1505        // Test BinOp mutation
1506        let binop = PythonNode::BinOp {
1507            left: Box::new(PythonNode::IntLit(1)),
1508            op: BinaryOp::Add,
1509            right: Box::new(PythonNode::IntLit(2)),
1510        };
1511        let result = gen.mutate_operator(&binop);
1512        // May or may not mutate depending on random
1513        assert!(result.is_some() || result.is_none());
1514
1515        // Test UnaryOp mutation
1516        let unaryop = PythonNode::UnaryOp {
1517            op: UnaryOp::Neg,
1518            operand: Box::new(PythonNode::IntLit(1)),
1519        };
1520        let result2 = gen.mutate_operator(&unaryop);
1521        assert!(result2.is_some() || result2.is_none());
1522
1523        // Test Compare mutation
1524        let compare = PythonNode::Compare {
1525            left: Box::new(PythonNode::IntLit(1)),
1526            op: CompareOp::Lt,
1527            right: Box::new(PythonNode::IntLit(2)),
1528        };
1529        let result3 = gen.mutate_operator(&compare);
1530        assert!(result3.is_some() || result3.is_none());
1531    }
1532
1533    #[test]
1534    fn test_mutate_literal() {
1535        let mut gen = NautilusGenerator::new(Language::Python, 2).with_seed(42);
1536
1537        // Test IntLit mutation
1538        let int_node = PythonNode::IntLit(42);
1539        let result = gen.mutate_literal(&int_node);
1540        assert!(result.is_some() || result.is_none());
1541
1542        // Test FloatLit mutation
1543        let float_node = PythonNode::FloatLit(3.14);
1544        let result2 = gen.mutate_literal(&float_node);
1545        assert!(result2.is_some() || result2.is_none());
1546
1547        // Test StrLit mutation
1548        let str_node = PythonNode::StrLit("hello".to_string());
1549        let result3 = gen.mutate_literal(&str_node);
1550        assert!(result3.is_some() || result3.is_none());
1551
1552        // Test BoolLit mutation
1553        let bool_node = PythonNode::BoolLit(true);
1554        let result4 = gen.mutate_literal(&bool_node);
1555        assert!(result4.is_some() || result4.is_none());
1556    }
1557
1558    #[test]
1559    fn test_insert_wrapper() {
1560        let mut gen = NautilusGenerator::new(Language::Python, 2).with_seed(42);
1561
1562        // Test with IntLit
1563        let node = PythonNode::IntLit(1);
1564        let result = gen.insert_wrapper(&node);
1565        // Should wrap with unary op
1566        if let Some(wrapped) = result {
1567            assert!(matches!(wrapped, PythonNode::UnaryOp { .. }));
1568        }
1569
1570        // Test with Name
1571        let name_node = PythonNode::Name("x".to_string());
1572        let result2 = gen.insert_wrapper(&name_node);
1573        assert!(result2.is_some());
1574
1575        // Test with FloatLit
1576        let float_node = PythonNode::FloatLit(1.0);
1577        let result3 = gen.insert_wrapper(&float_node);
1578        assert!(result3.is_some());
1579
1580        // Test with BinOp
1581        let binop = PythonNode::BinOp {
1582            left: Box::new(PythonNode::IntLit(1)),
1583            op: BinaryOp::Add,
1584            right: Box::new(PythonNode::IntLit(2)),
1585        };
1586        let result4 = gen.insert_wrapper(&binop);
1587        assert!(result4.is_some());
1588    }
1589
1590    #[test]
1591    fn test_delete_subtree() {
1592        let mut gen = NautilusGenerator::new(Language::Python, 2).with_seed(42);
1593
1594        // Test BinOp deletion
1595        let binop = PythonNode::BinOp {
1596            left: Box::new(PythonNode::IntLit(1)),
1597            op: BinaryOp::Add,
1598            right: Box::new(PythonNode::IntLit(2)),
1599        };
1600        let result = gen.delete_subtree(&binop);
1601        assert!(result.is_some());
1602
1603        // Test UnaryOp deletion
1604        let unaryop = PythonNode::UnaryOp {
1605            op: UnaryOp::Neg,
1606            operand: Box::new(PythonNode::IntLit(1)),
1607        };
1608        let result2 = gen.delete_subtree(&unaryop);
1609        assert!(result2.is_some());
1610        assert!(matches!(result2.unwrap(), PythonNode::IntLit(1)));
1611
1612        // Test If deletion
1613        let if_node = PythonNode::If {
1614            test: Box::new(PythonNode::BoolLit(true)),
1615            body: vec![PythonNode::Pass],
1616            orelse: vec![],
1617        };
1618        let result3 = gen.delete_subtree(&if_node);
1619        assert!(result3.is_some());
1620
1621        // Test with node that can't be deleted
1622        let int_lit = PythonNode::IntLit(1);
1623        let result4 = gen.delete_subtree(&int_lit);
1624        assert!(result4.is_none());
1625    }
1626}