use std::path::Path;
#[cfg(feature = "neural")]
use std::sync::Mutex;
pub(crate) struct NeuralLineScorer {
#[cfg(feature = "neural")]
session: Mutex<ort::session::Session>,
#[cfg(feature = "neural")]
input_name: String,
#[cfg(feature = "neural")]
output_name: String,
#[cfg(not(feature = "neural"))]
_phantom: (),
}
#[derive(Debug, Clone)]
pub(crate) struct LineFeatures {
pub line_length: f64,
pub indentation_level: f64,
pub token_diversity: f64,
pub is_definition: f64,
pub is_import: f64,
pub is_comment: f64,
pub is_closing: f64,
pub keyword_density: f64,
pub position_normalized: f64,
pub has_type_annotation: f64,
pub nesting_depth: f64,
pub prev_line_type: f64,
pub next_line_type: f64,
}
impl LineFeatures {
pub(crate) fn from_line(line: &str, position: f64, context: &LineContext) -> Self {
let trimmed = line.trim();
let leading = (line.len() - line.trim_start().len()) as f64;
Self {
line_length: trimmed.len() as f64,
indentation_level: leading / 4.0,
token_diversity: Self::compute_token_diversity(trimmed),
is_definition: if Self::check_definition(trimmed) {
1.0
} else {
0.0
},
is_import: if Self::check_import(trimmed) {
1.0
} else {
0.0
},
is_comment: if Self::check_comment(trimmed) {
1.0
} else {
0.0
},
is_closing: if Self::check_closing(trimmed) {
1.0
} else {
0.0
},
keyword_density: Self::compute_keyword_density(trimmed),
position_normalized: position,
has_type_annotation: if Self::check_type_annotation(trimmed) {
1.0
} else {
0.0
},
nesting_depth: context.nesting_depth as f64,
prev_line_type: context.prev_line_type as f64,
next_line_type: context.next_line_type as f64,
}
}
pub(crate) fn to_array(&self) -> [f64; 13] {
[
self.line_length,
self.indentation_level,
self.token_diversity,
self.is_definition,
self.is_import,
self.is_comment,
self.is_closing,
self.keyword_density,
self.position_normalized,
self.has_type_annotation,
self.nesting_depth,
self.prev_line_type,
self.next_line_type,
]
}
fn compute_token_diversity(line: &str) -> f64 {
let tokens: Vec<&str> = line.split_whitespace().collect();
if tokens.is_empty() {
return 0.0;
}
let unique: std::collections::HashSet<&str> = tokens.iter().copied().collect();
unique.len() as f64 / tokens.len() as f64
}
fn check_definition(line: &str) -> bool {
const STARTERS: &[&str] = &[
"fn ",
"pub fn ",
"async fn ",
"pub async fn ",
"def ",
"async def ",
"function ",
"export function ",
"async function ",
"class ",
"export class ",
"struct ",
"pub struct ",
"enum ",
"pub enum ",
"trait ",
"pub trait ",
"impl ",
"type ",
"pub type ",
"interface ",
"export interface ",
];
STARTERS.iter().any(|s| line.starts_with(s))
}
fn check_import(line: &str) -> bool {
line.starts_with("import ")
|| line.starts_with("use ")
|| line.starts_with("from ")
|| line.starts_with("#include")
|| line.starts_with("require(")
}
fn check_comment(line: &str) -> bool {
line.starts_with("//")
|| line.starts_with('#')
|| line.starts_with("/*")
|| line.starts_with('*')
|| line.starts_with("///")
}
fn check_closing(line: &str) -> bool {
matches!(line, "}" | "};" | "})" | "]" | ");" | "end")
}
fn check_type_annotation(line: &str) -> bool {
line.contains("->")
|| line.contains("=>")
|| line.contains(": ")
|| line.contains("Result<")
|| line.contains("Option<")
}
fn compute_keyword_density(line: &str) -> f64 {
const KEYWORDS: &[&str] = &[
"fn",
"let",
"mut",
"pub",
"use",
"impl",
"struct",
"enum",
"match",
"if",
"else",
"for",
"while",
"return",
"async",
"await",
"trait",
"where",
"def",
"class",
"import",
"from",
"function",
"export",
"const",
"var",
"type",
"interface",
"try",
"catch",
"throw",
"yield",
"raise",
];
let tokens: Vec<&str> = line.split_whitespace().collect();
if tokens.is_empty() {
return 0.0;
}
let hits = tokens
.iter()
.filter(|t| {
let clean = t.trim_end_matches(|c: char| !c.is_alphanumeric());
KEYWORDS.contains(&clean)
})
.count();
hits as f64 / tokens.len() as f64
}
}
#[derive(Debug, Clone, Default)]
pub(crate) struct LineContext {
pub nesting_depth: usize,
pub prev_line_type: u8,
pub next_line_type: u8,
}
impl NeuralLineScorer {
#[cfg(feature = "neural")]
pub(crate) fn load(model_path: &Path) -> anyhow::Result<Self> {
let eps = crate::core::ort_execution_providers::execution_providers();
let num_cpus = std::thread::available_parallelism().map_or(4, |n| n.get().max(1));
crate::core::ort_environment::ensure_ort_env(&eps)?;
let session = ort::session::Session::builder()
.map_err(|e| anyhow::anyhow!("ORT builder: {e}"))?
.with_intra_threads(num_cpus)
.map_err(|e| anyhow::anyhow!("ORT intra threads: {e}"))?
.with_optimization_level(ort::session::builder::GraphOptimizationLevel::All)
.map_err(|e| anyhow::anyhow!("ORT optimization: {e}"))?
.commit_from_file(model_path)
.map_err(|e| anyhow::anyhow!("ORT load model: {e}"))?;
let input_name = session
.inputs()
.first()
.map(|i| i.name().to_string())
.ok_or_else(|| anyhow::anyhow!("Neural model has no named inputs"))?;
let output_name = session
.outputs()
.first()
.map(|o| o.name().to_string())
.ok_or_else(|| anyhow::anyhow!("Neural model has no named outputs"))?;
Ok(Self {
session: Mutex::new(session),
input_name,
output_name,
})
}
#[cfg(not(feature = "neural"))]
pub(crate) fn load(_model_path: &Path) -> anyhow::Result<Self> {
anyhow::bail!("Neural feature not enabled. Compile with --features neural")
}
pub(crate) fn score_line(&self, line: &str, position: f64, task_keywords: &[String]) -> f64 {
let context = LineContext::default();
let features = LineFeatures::from_line(line, position, &context);
self.score_from_features(&features, task_keywords)
}
pub(crate) fn score_from_features(
&self,
features: &LineFeatures,
_task_keywords: &[String],
) -> f64 {
#[cfg(feature = "neural")]
{
self.neural_score(features)
}
#[cfg(not(feature = "neural"))]
{
self.decision_tree_score(features)
}
}
#[cfg(feature = "neural")]
fn neural_score(&self, features: &LineFeatures) -> f64 {
let input_data = features.to_array();
let float_data: Vec<f32> = input_data.iter().map(|&x| x as f32).collect();
let array = match ndarray::Array2::from_shape_vec((1, 13), float_data) {
Ok(a) => a,
Err(e) => {
tracing::warn!("neural_score: array creation failed: {e}");
return 0.5;
}
};
let tensor = match ort::value::Tensor::from_array(array) {
Ok(t) => t,
Err(e) => {
tracing::warn!("neural_score: tensor creation failed: {e}");
return 0.5;
}
};
let mut _guard = self
.session
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let outputs = match _guard.run(ort::inputs![self.input_name.as_str() => tensor]) {
Ok(o) => o,
Err(e) => {
tracing::warn!("neural_score: ORT inference failed: {e}");
return 0.5;
}
};
let out = match outputs[self.output_name.as_str()].try_extract_tensor::<f32>() {
Ok((_, data)) => *data.first().unwrap_or(&0.5),
Err(e) => {
tracing::warn!("neural_score: output extraction failed: {e}");
0.5
}
};
out as f64
}
#[cfg(not(feature = "neural"))]
#[allow(clippy::unused_self)]
fn decision_tree_score(&self, features: &LineFeatures) -> f64 {
let f = features.to_array();
let mut score = 0.5;
if f[3] > 0.5 {
score += 0.3; }
if f[5] > 0.5 {
score -= 0.2; }
if f[6] > 0.5 {
score -= 0.3; }
if f[4] > 0.5 {
score -= 0.1; }
if f[9] > 0.5 {
score += 0.15; }
let pos = f[8];
let u_curve = if pos <= 0.5 {
1.0 - 0.6 * (2.0 * pos).powi(2)
} else {
1.0 - 0.6 * (2.0 * (1.0 - pos)).powi(2)
};
score *= u_curve;
score.clamp(0.0, 1.0)
}
}
pub(crate) fn score_all_lines(
lines: &[&str],
scorer: &NeuralLineScorer,
task_keywords: &[String],
) -> Vec<f64> {
let n = lines.len();
let mut nesting_depth: usize = 0;
lines
.iter()
.enumerate()
.map(|(i, line)| {
let trimmed = line.trim();
nesting_depth = nesting_depth
.saturating_add(trimmed.matches('{').count())
.saturating_sub(trimmed.matches('}').count());
let prev_type = if i > 0 {
classify_type(lines[i - 1].trim())
} else {
0
};
let next_type = if i + 1 < n {
classify_type(lines[i + 1].trim())
} else {
0
};
let position = i as f64 / (n.max(1) - 1).max(1) as f64;
let context = LineContext {
nesting_depth,
prev_line_type: prev_type,
next_line_type: next_type,
};
let features = LineFeatures::from_line(line, position, &context);
scorer.score_from_features(&features, task_keywords)
})
.collect()
}
fn classify_type(line: &str) -> u8 {
if line.is_empty() {
return 0;
}
if LineFeatures::check_definition(line) {
return 1;
}
if LineFeatures::check_import(line) {
return 2;
}
if LineFeatures::check_comment(line) {
return 3;
}
if LineFeatures::check_closing(line) {
return 5;
}
4 }