use anyhow::Result;
use async_trait::async_trait;
use std::collections::HashMap;
use tracing::{info, debug};
use crate::core::CrateContext;
use crate::providers::{OpenAIProvider, GenerationRequest, LegacyLLMProvider};
use super::Stage;
pub struct OptimizationStage {
openai_provider: OpenAIProvider,
}
impl OptimizationStage {
pub fn new(openai_provider: OpenAIProvider) -> Self {
Self { openai_provider }
}
async fn optimize_code(&self, file_path: &str, content: &str) -> Result<String> {
let prompt = format!(
"Optimize this Rust code for performance, memory usage, and idiomatic patterns:\n\n{}",
content
);
let request = GenerationRequest {
prompt,
context: format!("File: {}", file_path),
max_tokens: Some(2048),
temperature: Some(0.3),
model: None,
};
let response = self.openai_provider.generate(request).await?;
Ok(response.content)
}
async fn optimize_dependencies(&self, dependencies: &[String]) -> Result<Vec<String>> {
let prompt = format!(
"Optimize these Rust dependencies for minimal compilation time and binary size: {:?}.
Suggest alternatives or feature flags to reduce bloat.",
dependencies
);
let request = GenerationRequest {
prompt,
context: "Dependency optimization".to_string(),
max_tokens: Some(1024),
temperature: Some(0.4),
model: None,
};
let response = self.openai_provider.generate(request).await?;
Ok(dependencies.to_vec())
}
async fn add_performance_annotations(&self, context: &CrateContext) -> Result<HashMap<String, String>> {
let mut optimized_files = HashMap::new();
for (path, content) in &context.file_structure {
if path.ends_with(".rs") {
let optimized_content = self.optimize_code(path, content).await?;
optimized_files.insert(path.clone(), optimized_content);
} else {
optimized_files.insert(path.clone(), content.clone());
}
}
Ok(optimized_files)
}
}
#[async_trait]
impl Stage for OptimizationStage {
type Input = CrateContext;
type Output = CrateContext;
async fn execute(&self, input: &Self::Input) -> Result<Self::Output> {
info!("Starting optimization for crate: {}", input.spec.name);
let optimized_files = self.add_performance_annotations(input).await?;
let optimized_dependencies = self.optimize_dependencies(&input.dependencies).await?;
let mut optimized_context = input.clone();
optimized_context.file_structure = optimized_files;
optimized_context.dependencies = optimized_dependencies;
optimized_context.metadata.insert(
"optimized_at".to_string(),
chrono::Utc::now().to_rfc3339()
);
debug!("Optimization complete for crate: {}", input.spec.name);
Ok(optimized_context)
}
fn name(&self) -> &str {
"optimization"
}
}