1use std::path::Path;
10#[cfg(feature = "neural")]
11use std::sync::Mutex;
12
13pub struct NeuralLineScorer {
14 #[cfg(feature = "neural")]
15 session: Mutex<ort::session::Session>,
16 #[cfg(feature = "neural")]
17 input_name: String,
18 #[cfg(feature = "neural")]
19 output_name: String,
20 #[cfg(not(feature = "neural"))]
21 _phantom: (),
22}
23
24#[derive(Debug, Clone)]
25pub struct LineFeatures {
26 pub line_length: f64,
27 pub indentation_level: f64,
28 pub token_diversity: f64,
29 pub is_definition: f64,
30 pub is_import: f64,
31 pub is_comment: f64,
32 pub is_closing: f64,
33 pub keyword_density: f64,
34 pub position_normalized: f64,
35 pub has_type_annotation: f64,
36 pub nesting_depth: f64,
37 pub prev_line_type: f64,
38 pub next_line_type: f64,
39}
40
41impl LineFeatures {
42 pub fn from_line(line: &str, position: f64, context: &LineContext) -> Self {
43 let trimmed = line.trim();
44 let leading = (line.len() - line.trim_start().len()) as f64;
45
46 Self {
47 line_length: trimmed.len() as f64,
48 indentation_level: leading / 4.0,
49 token_diversity: Self::compute_token_diversity(trimmed),
50 is_definition: if Self::check_definition(trimmed) {
51 1.0
52 } else {
53 0.0
54 },
55 is_import: if Self::check_import(trimmed) {
56 1.0
57 } else {
58 0.0
59 },
60 is_comment: if Self::check_comment(trimmed) {
61 1.0
62 } else {
63 0.0
64 },
65 is_closing: if Self::check_closing(trimmed) {
66 1.0
67 } else {
68 0.0
69 },
70 keyword_density: Self::compute_keyword_density(trimmed),
71 position_normalized: position,
72 has_type_annotation: if Self::check_type_annotation(trimmed) {
73 1.0
74 } else {
75 0.0
76 },
77 nesting_depth: context.nesting_depth as f64,
78 prev_line_type: context.prev_line_type as f64,
79 next_line_type: context.next_line_type as f64,
80 }
81 }
82
83 pub fn to_array(&self) -> [f64; 13] {
84 [
85 self.line_length,
86 self.indentation_level,
87 self.token_diversity,
88 self.is_definition,
89 self.is_import,
90 self.is_comment,
91 self.is_closing,
92 self.keyword_density,
93 self.position_normalized,
94 self.has_type_annotation,
95 self.nesting_depth,
96 self.prev_line_type,
97 self.next_line_type,
98 ]
99 }
100
101 fn compute_token_diversity(line: &str) -> f64 {
102 let tokens: Vec<&str> = line.split_whitespace().collect();
103 if tokens.is_empty() {
104 return 0.0;
105 }
106 let unique: std::collections::HashSet<&str> = tokens.iter().copied().collect();
107 unique.len() as f64 / tokens.len() as f64
108 }
109
110 fn check_definition(line: &str) -> bool {
111 const STARTERS: &[&str] = &[
112 "fn ",
113 "pub fn ",
114 "async fn ",
115 "pub async fn ",
116 "def ",
117 "async def ",
118 "function ",
119 "export function ",
120 "async function ",
121 "class ",
122 "export class ",
123 "struct ",
124 "pub struct ",
125 "enum ",
126 "pub enum ",
127 "trait ",
128 "pub trait ",
129 "impl ",
130 "type ",
131 "pub type ",
132 "interface ",
133 "export interface ",
134 ];
135 STARTERS.iter().any(|s| line.starts_with(s))
136 }
137
138 fn check_import(line: &str) -> bool {
139 line.starts_with("import ")
140 || line.starts_with("use ")
141 || line.starts_with("from ")
142 || line.starts_with("#include")
143 || line.starts_with("require(")
144 }
145
146 fn check_comment(line: &str) -> bool {
147 line.starts_with("//")
148 || line.starts_with('#')
149 || line.starts_with("/*")
150 || line.starts_with('*')
151 || line.starts_with("///")
152 }
153
154 fn check_closing(line: &str) -> bool {
155 matches!(line, "}" | "};" | "})" | "]" | ");" | "end")
156 }
157
158 fn check_type_annotation(line: &str) -> bool {
159 line.contains("->")
160 || line.contains("=>")
161 || line.contains(": ")
162 || line.contains("Result<")
163 || line.contains("Option<")
164 }
165
166 fn compute_keyword_density(line: &str) -> f64 {
167 const KEYWORDS: &[&str] = &[
168 "fn",
169 "let",
170 "mut",
171 "pub",
172 "use",
173 "impl",
174 "struct",
175 "enum",
176 "match",
177 "if",
178 "else",
179 "for",
180 "while",
181 "return",
182 "async",
183 "await",
184 "trait",
185 "where",
186 "def",
187 "class",
188 "import",
189 "from",
190 "function",
191 "export",
192 "const",
193 "var",
194 "type",
195 "interface",
196 "try",
197 "catch",
198 "throw",
199 "yield",
200 "raise",
201 ];
202 let tokens: Vec<&str> = line.split_whitespace().collect();
203 if tokens.is_empty() {
204 return 0.0;
205 }
206 let hits = tokens
207 .iter()
208 .filter(|t| {
209 let clean = t.trim_end_matches(|c: char| !c.is_alphanumeric());
210 KEYWORDS.contains(&clean)
211 })
212 .count();
213 hits as f64 / tokens.len() as f64
214 }
215}
216
217#[derive(Debug, Clone, Default)]
218pub struct LineContext {
219 pub nesting_depth: usize,
220 pub prev_line_type: u8,
221 pub next_line_type: u8,
222}
223
224impl NeuralLineScorer {
225 #[cfg(feature = "neural")]
226 pub fn load(model_path: &Path) -> anyhow::Result<Self> {
227 let eps = crate::core::ort_execution_providers::gpu_execution_providers();
228 let num_cpus = std::thread::available_parallelism().map_or(4, |n| n.get().max(1));
229 crate::core::ort_environment::ensure_ort_env(&eps)?;
230 let session = ort::session::Session::builder()
231 .map_err(|e| anyhow::anyhow!("ORT builder: {e}"))?
232 .with_intra_threads(num_cpus)
233 .map_err(|e| anyhow::anyhow!("ORT intra threads: {e}"))?
234 .with_optimization_level(ort::session::builder::GraphOptimizationLevel::All)
235 .map_err(|e| anyhow::anyhow!("ORT optimization: {e}"))?
236 .commit_from_file(model_path)
237 .map_err(|e| anyhow::anyhow!("ORT load model: {e}"))?;
238
239 let input_name = session
240 .inputs()
241 .first()
242 .map(|i| i.name().to_string())
243 .ok_or_else(|| anyhow::anyhow!("Neural model has no named inputs"))?;
244 let output_name = session
245 .outputs()
246 .first()
247 .map(|o| o.name().to_string())
248 .ok_or_else(|| anyhow::anyhow!("Neural model has no named outputs"))?;
249
250 Ok(Self {
251 session: Mutex::new(session),
252 input_name,
253 output_name,
254 })
255 }
256
257 #[cfg(not(feature = "neural"))]
258 pub fn load(_model_path: &Path) -> anyhow::Result<Self> {
259 anyhow::bail!("Neural feature not enabled. Compile with --features neural")
260 }
261
262 pub fn score_line(&self, line: &str, position: f64, task_keywords: &[String]) -> f64 {
263 let context = LineContext::default();
264 let features = LineFeatures::from_line(line, position, &context);
265 self.score_from_features(&features, task_keywords)
266 }
267
268 pub fn score_from_features(&self, features: &LineFeatures, _task_keywords: &[String]) -> f64 {
269 #[cfg(feature = "neural")]
270 {
271 self.neural_score(features)
272 }
273 #[cfg(not(feature = "neural"))]
274 {
275 self.decision_tree_score(features)
276 }
277 }
278
279 #[cfg(feature = "neural")]
280 fn neural_score(&self, features: &LineFeatures) -> f64 {
281 let input_data = features.to_array();
282 let float_data: Vec<f32> = input_data.iter().map(|&x| x as f32).collect();
283 let array = match ndarray::Array2::from_shape_vec((1, 13), float_data) {
284 Ok(a) => a,
285 Err(e) => {
286 tracing::warn!("neural_score: array creation failed: {e}");
287 return 0.5;
288 }
289 };
290 let tensor = match ort::value::Tensor::from_array(array) {
291 Ok(t) => t,
292 Err(e) => {
293 tracing::warn!("neural_score: tensor creation failed: {e}");
294 return 0.5;
295 }
296 };
297 let mut _guard = self.session.lock().unwrap();
298 let outputs = match _guard.run(ort::inputs![self.input_name.as_str() => tensor]) {
299 Ok(o) => o,
300 Err(e) => {
301 tracing::warn!("neural_score: ORT inference failed: {e}");
302 return 0.5;
303 }
304 };
305 let out = match outputs[self.output_name.as_str()].try_extract_tensor::<f32>() {
306 Ok((_, data)) => *data.first().unwrap_or(&0.5),
307 Err(e) => {
308 tracing::warn!("neural_score: output extraction failed: {e}");
309 0.5
310 }
311 };
312 out as f64
313 }
314
315 #[cfg(not(feature = "neural"))]
316 #[allow(clippy::unused_self)]
317 fn decision_tree_score(&self, features: &LineFeatures) -> f64 {
318 let f = features.to_array();
319
320 let mut score = 0.5;
321
322 if f[3] > 0.5 {
323 score += 0.3; }
325 if f[5] > 0.5 {
326 score -= 0.2; }
328 if f[6] > 0.5 {
329 score -= 0.3; }
331 if f[4] > 0.5 {
332 score -= 0.1; }
334 if f[9] > 0.5 {
335 score += 0.15; }
337
338 let pos = f[8];
339 let u_curve = if pos <= 0.5 {
340 1.0 - 0.6 * (2.0 * pos).powi(2)
341 } else {
342 1.0 - 0.6 * (2.0 * (1.0 - pos)).powi(2)
343 };
344 score *= u_curve;
345
346 score.clamp(0.0, 1.0)
347 }
348}
349
350pub fn score_all_lines(
351 lines: &[&str],
352 scorer: &NeuralLineScorer,
353 task_keywords: &[String],
354) -> Vec<f64> {
355 let n = lines.len();
356 let mut nesting_depth: usize = 0;
357
358 lines
359 .iter()
360 .enumerate()
361 .map(|(i, line)| {
362 let trimmed = line.trim();
363 nesting_depth = nesting_depth
364 .saturating_add(trimmed.matches('{').count())
365 .saturating_sub(trimmed.matches('}').count());
366
367 let prev_type = if i > 0 {
368 classify_type(lines[i - 1].trim())
369 } else {
370 0
371 };
372 let next_type = if i + 1 < n {
373 classify_type(lines[i + 1].trim())
374 } else {
375 0
376 };
377 let position = i as f64 / (n.max(1) - 1).max(1) as f64;
378
379 let context = LineContext {
380 nesting_depth,
381 prev_line_type: prev_type,
382 next_line_type: next_type,
383 };
384 let features = LineFeatures::from_line(line, position, &context);
385 scorer.score_from_features(&features, task_keywords)
386 })
387 .collect()
388}
389
390fn classify_type(line: &str) -> u8 {
391 if line.is_empty() {
392 return 0;
393 }
394 if LineFeatures::check_definition(line) {
395 return 1;
396 }
397 if LineFeatures::check_import(line) {
398 return 2;
399 }
400 if LineFeatures::check_comment(line) {
401 return 3;
402 }
403 if LineFeatures::check_closing(line) {
404 return 5;
405 }
406 4 }