1#![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#[derive(Debug, Clone, Default)]
23pub struct CoverageMap {
24 node_types: HashSet<String>,
26 ast_paths: HashSet<(String, String)>,
28 feature_combos: HashSet<String>,
30}
31
32impl CoverageMap {
33 #[must_use]
35 pub fn new() -> Self {
36 Self::default()
37 }
38
39 pub fn record_node(&mut self, node_type: &str) {
41 self.node_types.insert(node_type.to_string());
42 }
43
44 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 pub fn record_feature(&mut self, feature: &str) {
52 self.feature_combos.insert(feature.to_string());
53 }
54
55 #[must_use]
57 pub fn has_new_coverage(&self, existing: &Self) -> bool {
58 for node in &self.node_types {
60 if !existing.node_types.contains(node) {
61 return true;
62 }
63 }
64
65 for path in &self.ast_paths {
67 if !existing.ast_paths.contains(path) {
68 return true;
69 }
70 }
71
72 for feature in &self.feature_combos {
74 if !existing.feature_combos.contains(feature) {
75 return true;
76 }
77 }
78
79 false
80 }
81
82 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 #[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 #[must_use]
98 pub fn node_types(&self) -> &HashSet<String> {
99 &self.node_types
100 }
101
102 #[must_use]
104 pub fn ast_paths(&self) -> &HashSet<(String, String)> {
105 &self.ast_paths
106 }
107}
108
109#[derive(Debug, Clone)]
111pub struct CorpusEntry {
112 pub code: GeneratedCode,
114 pub coverage: CoverageMap,
116 pub energy: f64,
118 pub selection_count: usize,
120 pub ast: Option<PythonNode>,
122}
123
124impl CorpusEntry {
125 #[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 #[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 pub fn update_energy(&mut self, global_coverage: &CoverageMap) {
151 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 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#[derive(Debug)]
173pub struct NautilusGenerator {
174 corpus: Vec<CorpusEntry>,
176 global_coverage: CoverageMap,
178 max_corpus_size: usize,
180 max_depth: usize,
182 language: Language,
184 seed: u64,
186 rng_state: u64,
188}
189
190impl NautilusGenerator {
191 #[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 #[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 #[must_use]
215 pub fn with_max_corpus(mut self, size: usize) -> Self {
216 self.max_corpus_size = size;
217 self
218 }
219
220 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 fn random_float(&mut self) -> f64 {
232 (self.next_random() as f64) / (u64::MAX as f64)
233 }
234
235 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 pub fn initialize_corpus_with_ast(&mut self) {
251 let enumerator = PythonEnumerator::new(self.max_depth.min(2));
252
253 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 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 fn compute_coverage(&self, code: &str) -> CoverageMap {
353 let mut coverage = CoverageMap::new();
354
355 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 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 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 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 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 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 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 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 pub fn generate(&mut self, count: usize) -> Vec<GeneratedCode> {
577 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; while results.len() < count && iterations < max_iterations {
587 iterations += 1;
588
589 if let Some(idx) = self.select_entry_mut() {
591 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 let should_mutate = self.random_float() < 0.7 && has_ast;
598
599 if should_mutate {
600 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 self.add_to_corpus(program.clone(), coverage);
615 results.push(program);
616 }
617 }
618 } else {
619 results.push(code_clone);
621 }
622 } else {
623 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 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 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 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 fn insert_wrapper(&mut self, node: &PythonNode) -> Option<PythonNode> {
700 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 fn delete_subtree(&mut self, node: &PythonNode) -> Option<PythonNode> {
719 match node {
720 PythonNode::BinOp { left, right, .. } => {
721 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 body.first().cloned()
732 }
733 _ => None,
734 }
735 }
736
737 #[must_use]
739 pub fn corpus_size(&self) -> usize {
740 self.corpus.len()
741 }
742
743 #[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#[derive(Debug, Clone)]
758pub struct CoverageStats {
759 pub total_coverage: usize,
761 pub node_types_covered: usize,
763 pub ast_paths_covered: usize,
765 pub features_covered: usize,
767 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 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 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 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 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 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 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 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 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 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 gen.initialize_corpus_with_ast();
1478
1479 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 for _ in 0..10 {
1495 let mutated = gen.mutate_ast(&ast);
1496 assert!(mutated.is_some() || true); }
1499 }
1500
1501 #[test]
1502 fn test_mutate_operator() {
1503 let mut gen = NautilusGenerator::new(Language::Python, 2).with_seed(42);
1504
1505 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 assert!(result.is_some() || result.is_none());
1514
1515 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 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 let int_node = PythonNode::IntLit(42);
1539 let result = gen.mutate_literal(&int_node);
1540 assert!(result.is_some() || result.is_none());
1541
1542 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 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 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 let node = PythonNode::IntLit(1);
1564 let result = gen.insert_wrapper(&node);
1565 if let Some(wrapped) = result {
1567 assert!(matches!(wrapped, PythonNode::UnaryOp { .. }));
1568 }
1569
1570 let name_node = PythonNode::Name("x".to_string());
1572 let result2 = gen.insert_wrapper(&name_node);
1573 assert!(result2.is_some());
1574
1575 let float_node = PythonNode::FloatLit(1.0);
1577 let result3 = gen.insert_wrapper(&float_node);
1578 assert!(result3.is_some());
1579
1580 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 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 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 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 let int_lit = PythonNode::IntLit(1);
1623 let result4 = gen.delete_subtree(&int_lit);
1624 assert!(result4.is_none());
1625 }
1626}