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::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
298 .session
299 .lock()
300 .unwrap_or_else(std::sync::PoisonError::into_inner);
301 let outputs = match _guard.run(ort::inputs![self.input_name.as_str() => tensor]) {
302 Ok(o) => o,
303 Err(e) => {
304 tracing::warn!("neural_score: ORT inference failed: {e}");
305 return 0.5;
306 }
307 };
308 let out = match outputs[self.output_name.as_str()].try_extract_tensor::<f32>() {
309 Ok((_, data)) => *data.first().unwrap_or(&0.5),
310 Err(e) => {
311 tracing::warn!("neural_score: output extraction failed: {e}");
312 0.5
313 }
314 };
315 out as f64
316 }
317
318 #[cfg(not(feature = "neural"))]
319 #[allow(clippy::unused_self)]
320 fn decision_tree_score(&self, features: &LineFeatures) -> f64 {
321 let f = features.to_array();
322
323 let mut score = 0.5;
324
325 if f[3] > 0.5 {
326 score += 0.3; }
328 if f[5] > 0.5 {
329 score -= 0.2; }
331 if f[6] > 0.5 {
332 score -= 0.3; }
334 if f[4] > 0.5 {
335 score -= 0.1; }
337 if f[9] > 0.5 {
338 score += 0.15; }
340
341 let pos = f[8];
342 let u_curve = if pos <= 0.5 {
343 1.0 - 0.6 * (2.0 * pos).powi(2)
344 } else {
345 1.0 - 0.6 * (2.0 * (1.0 - pos)).powi(2)
346 };
347 score *= u_curve;
348
349 score.clamp(0.0, 1.0)
350 }
351}
352
353pub fn score_all_lines(
354 lines: &[&str],
355 scorer: &NeuralLineScorer,
356 task_keywords: &[String],
357) -> Vec<f64> {
358 let n = lines.len();
359 let mut nesting_depth: usize = 0;
360
361 lines
362 .iter()
363 .enumerate()
364 .map(|(i, line)| {
365 let trimmed = line.trim();
366 nesting_depth = nesting_depth
367 .saturating_add(trimmed.matches('{').count())
368 .saturating_sub(trimmed.matches('}').count());
369
370 let prev_type = if i > 0 {
371 classify_type(lines[i - 1].trim())
372 } else {
373 0
374 };
375 let next_type = if i + 1 < n {
376 classify_type(lines[i + 1].trim())
377 } else {
378 0
379 };
380 let position = i as f64 / (n.max(1) - 1).max(1) as f64;
381
382 let context = LineContext {
383 nesting_depth,
384 prev_line_type: prev_type,
385 next_line_type: next_type,
386 };
387 let features = LineFeatures::from_line(line, position, &context);
388 scorer.score_from_features(&features, task_keywords)
389 })
390 .collect()
391}
392
393fn classify_type(line: &str) -> u8 {
394 if line.is_empty() {
395 return 0;
396 }
397 if LineFeatures::check_definition(line) {
398 return 1;
399 }
400 if LineFeatures::check_import(line) {
401 return 2;
402 }
403 if LineFeatures::check_comment(line) {
404 return 3;
405 }
406 if LineFeatures::check_closing(line) {
407 return 5;
408 }
409 4 }