Skip to main content

cha_core/
wasm.rs

1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3
4use streaming_iterator::StreamingIterator;
5use wasmtime::component::{Component, HasSelf, Linker};
6use wasmtime::{Engine, Store};
7use wasmtime_wasi::{ResourceTable, WasiCtx, WasiCtxBuilder, WasiCtxView, WasiView};
8
9use crate::model::ArmValue;
10use crate::plugin::{Finding, Location, Severity, SmellCategory};
11use crate::{AnalysisContext, Plugin};
12
13mod bindings {
14    wasmtime::component::bindgen!({
15        path: "wit/plugin.wit",
16        world: "analyzer",
17    });
18}
19
20use bindings::Analyzer;
21use bindings::cha::plugin::tree_query;
22pub use bindings::cha::plugin::types as wit;
23
24struct HostState {
25    wasi: WasiCtx,
26    table: ResourceTable,
27    tree: Option<tree_sitter::Tree>,
28    source: Vec<u8>,
29    ts_language: Option<tree_sitter::Language>,
30    query_cache: HashMap<String, tree_sitter::Query>,
31}
32
33impl WasiView for HostState {
34    fn ctx(&mut self) -> WasiCtxView<'_> {
35        WasiCtxView {
36            ctx: &mut self.wasi,
37            table: &mut self.table,
38        }
39    }
40}
41
42fn new_host_state(
43    tree: Option<tree_sitter::Tree>,
44    source: Vec<u8>,
45    ts_language: Option<tree_sitter::Language>,
46) -> HostState {
47    let wasi = WasiCtxBuilder::new().build();
48    HostState {
49        wasi,
50        table: ResourceTable::new(),
51        tree,
52        source,
53        ts_language,
54        query_cache: HashMap::new(),
55    }
56}
57
58impl bindings::cha::plugin::types::Host for HostState {}
59
60impl tree_query::Host for HostState {
61    fn run_query(&mut self, pattern: String) -> Vec<Vec<tree_query::QueryMatch>> {
62        self.execute_query(&pattern)
63    }
64
65    fn run_queries(&mut self, patterns: Vec<String>) -> Vec<Vec<Vec<tree_query::QueryMatch>>> {
66        patterns.iter().map(|p| self.execute_query(p)).collect()
67    }
68
69    fn node_at(&mut self, line: u32, col: u32) -> Option<tree_query::QueryMatch> {
70        let tree = self.tree.as_ref()?;
71        let point = tree_sitter::Point::new(line as usize, col as usize);
72        let node = tree.root_node().descendant_for_point_range(point, point)?;
73        Some(node_to_query_match(&node, &self.source, ""))
74    }
75
76    fn nodes_in_range(&mut self, start_line: u32, end_line: u32) -> Vec<tree_query::QueryMatch> {
77        let tree = match &self.tree {
78            Some(t) => t,
79            None => return vec![],
80        };
81        let mut results = vec![];
82        let mut cursor = tree.root_node().walk();
83        for child in tree.root_node().children(&mut cursor) {
84            let node_start = child.start_position().row as u32;
85            let node_end = child.end_position().row as u32;
86            if node_end < start_line {
87                continue;
88            }
89            if node_start > end_line {
90                break;
91            }
92            if child.is_named() {
93                results.push(node_to_query_match(&child, &self.source, ""));
94            }
95        }
96        results
97    }
98}
99
100impl HostState {
101    fn execute_query(&mut self, pattern: &str) -> Vec<Vec<tree_query::QueryMatch>> {
102        let (tree, ts_lang) = match (&self.tree, &self.ts_language) {
103            (Some(t), Some(l)) => (t, l),
104            _ => return vec![],
105        };
106
107        if !self.query_cache.contains_key(pattern) {
108            let q = match tree_sitter::Query::new(ts_lang, pattern) {
109                Ok(q) => q,
110                Err(_) => return vec![],
111            };
112            self.query_cache.insert(pattern.to_string(), q);
113        }
114        let query = self.query_cache.get(pattern).unwrap();
115        let capture_names: Vec<&str> = query.capture_names().to_vec();
116
117        let mut cursor = tree_sitter::QueryCursor::new();
118        let mut results = vec![];
119        let mut matches = cursor.matches(query, tree.root_node(), self.source.as_slice());
120        while let Some(m) = StreamingIterator::next(&mut matches) {
121            let captures: Vec<_> = m
122                .captures
123                .iter()
124                .map(|c| {
125                    let name: &str = capture_names.get(c.index as usize).copied().unwrap_or("");
126                    node_to_query_match(&c.node, &self.source, name)
127                })
128                .collect();
129            results.push(captures);
130        }
131        results
132    }
133}
134
135fn node_to_query_match(
136    node: &tree_sitter::Node,
137    source: &[u8],
138    capture_name: &str,
139) -> tree_query::QueryMatch {
140    let text = node.utf8_text(source).unwrap_or("").to_string();
141    tree_query::QueryMatch {
142        capture_name: capture_name.to_string(),
143        node_kind: node.kind().to_string(),
144        text,
145        start_line: node.start_position().row as u32,
146        start_col: node.start_position().column as u32,
147        end_line: node.end_position().row as u32,
148        end_col: node.end_position().column as u32,
149    }
150}
151
152/// Adapter that loads a WASM component and wraps it as a Plugin.
153pub struct WasmPlugin {
154    engine: Engine,
155    component: Component,
156    plugin_name: String,
157    plugin_version: String,
158    plugin_description: String,
159    plugin_authors: Vec<String>,
160    plugin_smells: Vec<String>,
161    options: Vec<(String, wit::OptionValue)>,
162}
163
164impl WasmPlugin {
165    pub fn load(path: &Path) -> wasmtime::Result<Self> {
166        let engine = Engine::default();
167        let bytes = std::fs::read(path)?;
168        let component = Component::from_binary(&engine, &bytes)?;
169
170        let mut linker = Linker::<HostState>::new(&engine);
171        wasmtime_wasi::p2::add_to_linker_sync(&mut linker)?;
172        Analyzer::add_to_linker::<HostState, HasSelf<HostState>>(&mut linker, |s| s)?;
173
174        let mut store = Store::new(&engine, new_host_state(None, vec![], None));
175        let instance = Analyzer::instantiate(&mut store, &component, &linker)?;
176        let name = instance.call_name(&mut store)?;
177        let version = instance.call_version(&mut store)?;
178        let description = instance.call_description(&mut store)?;
179        let authors = instance.call_authors(&mut store)?;
180        let smells = instance.call_smells(&mut store)?;
181
182        Ok(Self {
183            engine,
184            component,
185            plugin_name: name,
186            plugin_version: version,
187            plugin_description: description,
188            plugin_authors: authors,
189            plugin_smells: smells,
190            options: vec![],
191        })
192    }
193
194    /// Set plugin options from config.
195    pub fn set_options(&mut self, options: Vec<(String, wit::OptionValue)>) {
196        self.options = options;
197    }
198}
199
200impl Plugin for WasmPlugin {
201    fn name(&self) -> &str {
202        &self.plugin_name
203    }
204
205    fn version(&self) -> &str {
206        &self.plugin_version
207    }
208
209    fn description(&self) -> &str {
210        &self.plugin_description
211    }
212
213    fn authors(&self) -> Vec<String> {
214        self.plugin_authors.clone()
215    }
216
217    fn smells(&self) -> Vec<String> {
218        self.plugin_smells.clone()
219    }
220
221    fn analyze(&self, ctx: &AnalysisContext) -> Vec<Finding> {
222        let result = (|| -> wasmtime::Result<Vec<Finding>> {
223            let mut linker = Linker::<HostState>::new(&self.engine);
224            wasmtime_wasi::p2::add_to_linker_sync(&mut linker)?;
225            Analyzer::add_to_linker::<HostState, HasSelf<HostState>>(&mut linker, |s| s)?;
226
227            let (tree, ts_lang) = match (ctx.tree, ctx.ts_language) {
228                (Some(t), Some(l)) => (Some(t.clone()), Some(l.clone())),
229                _ => (None, None),
230            };
231            let source = ctx.file.content.as_bytes().to_vec();
232            let mut store = Store::new(&self.engine, new_host_state(tree, source, ts_lang));
233            let instance = Analyzer::instantiate(&mut store, &self.component, &linker)?;
234            let input = to_wit_input(ctx, &self.options);
235            let results = instance.call_analyze(&mut store, &input)?;
236            Ok(results.into_iter().map(from_wit_finding).collect())
237        })();
238
239        result.unwrap_or_else(|e| {
240            eprintln!("wasm plugin error: {}", e);
241            vec![]
242        })
243    }
244}
245
246fn to_wit_input(
247    ctx: &AnalysisContext,
248    options: &[(String, wit::OptionValue)],
249) -> wit::AnalysisInput {
250    wit::AnalysisInput {
251        path: ctx.file.path.to_string_lossy().into(),
252        content: ctx.file.content.clone(),
253        language: ctx.model.language.clone(),
254        total_lines: ctx.model.total_lines as u32,
255        role: infer_file_role(&ctx.file.path),
256        functions: convert_functions(&ctx.model.functions),
257        classes: convert_classes(&ctx.model.classes),
258        imports: convert_imports(&ctx.model.imports),
259        comments: convert_comments(&ctx.model.comments),
260        type_aliases: ctx
261            .model
262            .type_aliases
263            .iter()
264            .map(|(k, v)| (k.clone(), v.clone()))
265            .collect(),
266        options: options.to_vec(),
267    }
268}
269
270fn infer_file_role(path: &Path) -> wit::FileRole {
271    let s = path.to_string_lossy();
272    if s.contains("/test") || s.contains("_test.") || s.contains("/tests/") || s.contains("/spec/")
273    {
274        return wit::FileRole::Test;
275    }
276    if s.contains("/generated/") || s.contains(".generated.") || s.contains(".gen.") {
277        return wit::FileRole::Generated;
278    }
279    match path.extension().and_then(|e| e.to_str()) {
280        Some("md" | "txt" | "rst" | "adoc") => return wit::FileRole::Doc,
281        Some("toml" | "json" | "yaml" | "yml" | "ini" | "cfg") => return wit::FileRole::Config,
282        _ => {}
283    }
284    wit::FileRole::Source
285}
286
287/// Generic slice converter to avoid duplicate map-collect patterns.
288fn convert_slice<T, U>(items: &[T], f: impl Fn(&T) -> U) -> Vec<U> {
289    items.iter().map(f).collect()
290}
291
292fn convert_functions(funcs: &[crate::model::FunctionInfo]) -> Vec<wit::FunctionInfo> {
293    convert_slice(funcs, |f| wit::FunctionInfo {
294        name: f.name.clone(),
295        start_line: f.start_line as u32,
296        end_line: f.end_line as u32,
297        name_col: f.name_col as u32,
298        name_end_col: f.name_end_col as u32,
299        line_count: f.line_count as u32,
300        complexity: f.complexity as u32,
301        parameter_count: f.parameter_count as u32,
302        parameter_types: f.parameter_types.iter().map(to_wit_type_ref).collect(),
303        parameter_names: f.parameter_names.clone(),
304        chain_depth: f.chain_depth as u32,
305        switch_arms: f.switch_arms as u32,
306        switch_arm_values: f.switch_arm_values.iter().map(to_wit_arm_value).collect(),
307        external_refs: f.external_refs.clone(),
308        is_delegating: f.is_delegating,
309        is_exported: f.is_exported,
310        comment_lines: f.comment_lines as u32,
311        referenced_fields: f.referenced_fields.clone(),
312        null_check_fields: f.null_check_fields.clone(),
313        switch_dispatch_target: f.switch_dispatch_target.clone(),
314        optional_param_count: f.optional_param_count as u32,
315        called_functions: f.called_functions.clone(),
316        cognitive_complexity: f.cognitive_complexity as u32,
317        body_hash: f.body_hash.map(|h| format!("{h:016x}")),
318        return_type: f.return_type.as_ref().map(to_wit_type_ref),
319    })
320}
321
322fn to_wit_arm_value(v: &ArmValue) -> wit::ArmValue {
323    match v {
324        ArmValue::Str(s) => wit::ArmValue::StrLit(s.clone()),
325        ArmValue::Int(i) => wit::ArmValue::IntLit(*i),
326        ArmValue::Char(c) => wit::ArmValue::CharLit(*c as u32),
327        ArmValue::Other => wit::ArmValue::Other,
328    }
329}
330
331fn to_wit_type_ref(t: &crate::model::TypeRef) -> wit::TypeRef {
332    wit::TypeRef {
333        name: t.name.clone(),
334        raw: t.raw.clone(),
335        origin: match &t.origin {
336            crate::model::TypeOrigin::Local => wit::TypeOrigin::ProjectLocal,
337            crate::model::TypeOrigin::External(m) => wit::TypeOrigin::External(m.clone()),
338            crate::model::TypeOrigin::Primitive => wit::TypeOrigin::Primitive,
339            crate::model::TypeOrigin::Unknown => wit::TypeOrigin::Unknown,
340        },
341    }
342}
343
344fn convert_classes(classes: &[crate::model::ClassInfo]) -> Vec<wit::ClassInfo> {
345    convert_slice(classes, |c| wit::ClassInfo {
346        name: c.name.clone(),
347        start_line: c.start_line as u32,
348        end_line: c.end_line as u32,
349        name_col: c.name_col as u32,
350        name_end_col: c.name_end_col as u32,
351        method_count: c.method_count as u32,
352        line_count: c.line_count as u32,
353        delegating_method_count: c.delegating_method_count as u32,
354        field_count: c.field_count as u32,
355        field_names: c.field_names.clone(),
356        field_types: c.field_types.clone(),
357        is_exported: c.is_exported,
358        has_behavior: c.has_behavior,
359        is_interface: c.is_interface,
360        parent_name: c.parent_name.clone(),
361        override_count: c.override_count as u32,
362        self_call_count: c.self_call_count as u32,
363        has_listener_field: c.has_listener_field,
364        has_notify_method: c.has_notify_method,
365    })
366}
367
368fn convert_imports(imports: &[crate::model::ImportInfo]) -> Vec<wit::ImportInfo> {
369    convert_slice(imports, |i| wit::ImportInfo {
370        source: i.source.clone(),
371        line: i.line as u32,
372        col: i.col as u32,
373        is_module_decl: i.is_module_decl,
374    })
375}
376
377fn convert_comments(comments: &[crate::model::CommentInfo]) -> Vec<wit::CommentInfo> {
378    convert_slice(comments, |c| wit::CommentInfo {
379        text: c.text.clone(),
380        line: c.line as u32,
381    })
382}
383
384fn from_wit_finding(f: wit::Finding) -> Finding {
385    Finding {
386        smell_name: f.smell_name,
387        category: convert_category(f.category),
388        severity: convert_severity(f.severity),
389        location: Location {
390            path: PathBuf::from(&f.location.path),
391            start_line: f.location.start_line as usize,
392            start_col: f.location.start_col as usize,
393            end_line: f.location.end_line as usize,
394            end_col: f.location.end_col as usize,
395            name: f.location.name,
396        },
397        message: f.message,
398        suggested_refactorings: f.suggested_refactorings,
399        actual_value: f.actual_value,
400        threshold: f.threshold,
401        risk_score: None,
402    }
403}
404
405fn convert_severity(s: wit::Severity) -> Severity {
406    match s {
407        wit::Severity::Hint => Severity::Hint,
408        wit::Severity::Warning => Severity::Warning,
409        wit::Severity::Error => Severity::Error,
410    }
411}
412
413fn convert_category(c: wit::SmellCategory) -> SmellCategory {
414    match c {
415        wit::SmellCategory::Bloaters => SmellCategory::Bloaters,
416        wit::SmellCategory::OoAbusers => SmellCategory::OoAbusers,
417        wit::SmellCategory::ChangePreventers => SmellCategory::ChangePreventers,
418        wit::SmellCategory::Dispensables => SmellCategory::Dispensables,
419        wit::SmellCategory::Couplers => SmellCategory::Couplers,
420        wit::SmellCategory::Security => SmellCategory::Security,
421    }
422}
423
424/// Scan plugin directories and load all .wasm components.
425pub fn load_wasm_plugins(project_dir: &Path) -> Vec<WasmPlugin> {
426    let mut plugins: Vec<WasmPlugin> = Vec::new();
427    let mut seen = HashMap::new();
428
429    let project_plugins = project_dir.join(".cha").join("plugins");
430    let global_plugins = home_dir().join(".cha").join("plugins");
431
432    for dir in [&project_plugins, &global_plugins] {
433        load_plugins_from_dir(dir, &mut seen, &mut plugins);
434    }
435
436    plugins
437}
438
439/// Load .wasm plugins from a single directory, skipping duplicates by filename.
440fn load_plugins_from_dir(
441    dir: &Path,
442    seen: &mut HashMap<String, bool>,
443    plugins: &mut Vec<WasmPlugin>,
444) {
445    let entries = match std::fs::read_dir(dir) {
446        Ok(e) => e,
447        Err(_) => return,
448    };
449    for entry in entries.flatten() {
450        let path = entry.path();
451        if path.extension().is_none_or(|e| e != "wasm") {
452            continue;
453        }
454        let filename = path.file_name().unwrap().to_string_lossy().to_string();
455        if seen.contains_key(&filename) {
456            continue;
457        }
458        match WasmPlugin::load(&path) {
459            Ok(p) => {
460                seen.insert(filename, true);
461                plugins.push(p);
462            }
463            Err(e) => {
464                eprintln!("failed to load wasm plugin {}: {}", path.display(), e);
465            }
466        }
467    }
468}
469
470fn home_dir() -> PathBuf {
471    std::env::var("HOME")
472        .map(PathBuf::from)
473        .unwrap_or_else(|_| PathBuf::from("."))
474}
475
476/// Convert a TOML value to a WIT option-value.
477pub fn toml_to_option_value(v: &toml::Value) -> Option<wit::OptionValue> {
478    match v {
479        toml::Value::String(s) => Some(wit::OptionValue::Str(s.clone())),
480        toml::Value::Integer(i) => Some(wit::OptionValue::Int(*i)),
481        toml::Value::Float(f) => Some(wit::OptionValue::Float(*f)),
482        toml::Value::Boolean(b) => Some(wit::OptionValue::Boolean(*b)),
483        toml::Value::Array(arr) => {
484            let strs: Vec<String> = arr
485                .iter()
486                .filter_map(|v| v.as_str().map(String::from))
487                .collect();
488            Some(wit::OptionValue::ListStr(strs))
489        }
490        _ => None,
491    }
492}