use ast_grep_config::RuleCore;
use ast_grep_core::pinned::{NodeData, PinnedNodeData};
use ast_grep_core::{AstGrep, NodeMatch};
use ast_grep_language::SupportLang;
use ignore::{WalkBuilder, WalkParallel, WalkState};
use napi::anyhow::{anyhow, Context, Result as Ret};
use napi::bindgen_prelude::*;
use napi::threadsafe_function::{ErrorStrategy, ThreadsafeFunction, ThreadsafeFunctionCallMode};
use napi::{JsNumber, Task};
use napi_derive::napi;
use std::collections::HashMap;
use std::sync::atomic::{AtomicU32, Ordering};
use crate::doc::{JsDoc, NapiConfig};
use crate::napi_lang::{build_files, Lang, LangOption};
use crate::sg_node::{SgNode, SgRoot};
pub struct ParseAsync {
pub src: String,
pub lang: Lang,
}
impl Task for ParseAsync {
type Output = SgRoot;
type JsValue = SgRoot;
fn compute(&mut self) -> Result<Self::Output> {
let src = std::mem::take(&mut self.src);
let doc = JsDoc::new(src, self.lang.into());
Ok(SgRoot(AstGrep::doc(doc), "anonymous".into()))
}
fn resolve(&mut self, _env: Env, output: Self::Output) -> Result<Self::JsValue> {
Ok(output)
}
}
type Entry = std::result::Result<ignore::DirEntry, ignore::Error>;
pub struct IterateFiles<D> {
walk: WalkParallel,
lang_option: LangOption,
tsfn: D,
producer: fn(&D, Entry, &LangOption) -> Ret<bool>,
}
impl<T: 'static + Send + Sync> Task for IterateFiles<T> {
type Output = u32;
type JsValue = JsNumber;
fn compute(&mut self) -> Result<Self::Output> {
let tsfn = &self.tsfn;
let file_count = AtomicU32::new(0);
let producer = self.producer;
let walker = std::mem::replace(&mut self.walk, WalkBuilder::new(".").build_parallel());
walker.run(|| {
let file_count = &file_count;
let lang_option = &self.lang_option;
Box::new(move |entry| match producer(tsfn, entry, lang_option) {
Ok(succeed) => {
if succeed {
file_count.fetch_add(1, Ordering::AcqRel);
}
WalkState::Continue
}
Err(_) => WalkState::Skip,
})
});
Ok(file_count.load(Ordering::Acquire))
}
fn resolve(&mut self, env: Env, output: Self::Output) -> Result<Self::JsValue> {
env.create_uint32(output)
}
}
const THREAD_FUNC_QUEUE_SIZE: usize = 1000;
type ParseFiles = IterateFiles<ThreadsafeFunction<SgRoot, ErrorStrategy::CalleeHandled>>;
#[napi(object)]
pub struct FileOption {
pub paths: Vec<String>,
pub language_globs: HashMap<String, Vec<String>>,
}
#[napi(ts_return_type = "Promise<number>")]
pub fn parse_files(
paths: Either<Vec<String>, FileOption>,
#[napi(ts_arg_type = "(err: null | Error, result: SgRoot) => void")] callback: JsFunction,
) -> Result<AsyncTask<ParseFiles>> {
let tsfn: ThreadsafeFunction<SgRoot, ErrorStrategy::CalleeHandled> =
callback.create_threadsafe_function(THREAD_FUNC_QUEUE_SIZE, |ctx| Ok(vec![ctx.value]))?;
let (paths, globs) = match paths {
Either::A(v) => (v, HashMap::new()),
Either::B(FileOption {
paths,
language_globs,
}) => (paths, Lang::lang_globs(language_globs)),
};
let walk = build_files(paths, &globs)?;
Ok(AsyncTask::new(ParseFiles {
walk,
tsfn,
lang_option: LangOption::infer(&globs),
producer: call_sg_root,
}))
}
fn call_sg_root(
tsfn: &ThreadsafeFunction<SgRoot, ErrorStrategy::CalleeHandled>,
entry: std::result::Result<ignore::DirEntry, ignore::Error>,
lang_option: &LangOption,
) -> Ret<bool> {
let entry = entry?;
if !entry
.file_type()
.context("could not use stdin as file")?
.is_file()
{
return Ok(false);
}
let (root, path) = get_root(entry, lang_option)?;
let sg = SgRoot(root, path);
tsfn.call(Ok(sg), ThreadsafeFunctionCallMode::Blocking);
Ok(true)
}
fn get_root(entry: ignore::DirEntry, lang_option: &LangOption) -> Ret<(AstGrep<JsDoc>, String)> {
let path = entry.into_path();
let file_content = std::fs::read_to_string(&path)?;
let lang = lang_option
.get_lang(&path)
.context(anyhow!("file not recognized"))?;
let doc = JsDoc::new(file_content, lang);
Ok((AstGrep::doc(doc), path.to_string_lossy().into()))
}
pub type FindInFiles = IterateFiles<(
ThreadsafeFunction<PinnedNodes, ErrorStrategy::CalleeHandled>,
RuleCore<SupportLang>,
)>;
pub struct PinnedNodes(
PinnedNodeData<JsDoc, Vec<NodeMatch<'static, JsDoc>>>,
String,
);
unsafe impl Send for PinnedNodes {}
unsafe impl Sync for PinnedNodes {}
#[napi(object)]
pub struct FindConfig {
pub paths: Vec<String>,
pub matcher: NapiConfig,
pub language_globs: Option<Vec<String>>,
}
pub fn find_in_files_impl(
lang: Lang,
config: FindConfig,
callback: JsFunction,
) -> Result<AsyncTask<FindInFiles>> {
let tsfn = callback.create_threadsafe_function(THREAD_FUNC_QUEUE_SIZE, |ctx| {
from_pinned_data(ctx.value, ctx.env)
})?;
let FindConfig {
paths,
matcher,
language_globs,
} = config;
let rule = matcher.parse_with(lang)?;
let walk = lang.find_files(paths, language_globs)?;
Ok(AsyncTask::new(FindInFiles {
walk,
tsfn: (tsfn, rule),
lang_option: LangOption::Specified(lang),
producer: call_sg_node,
}))
}
fn from_pinned_data(pinned: PinnedNodes, env: napi::Env) -> Result<Vec<Vec<SgNode>>> {
let (root, nodes) = pinned.0.into_raw();
let sg_root = SgRoot(AstGrep { inner: root }, pinned.1);
let reference = SgRoot::into_reference(sg_root, env)?;
let mut v = vec![];
for mut node in nodes {
let root_ref = reference.clone(env)?;
let sg_node = SgNode {
inner: root_ref.share_with(env, |root| {
let r = &root.0.inner;
node.visit_nodes(|n| unsafe { r.readopt(n) });
Ok(node)
})?,
};
v.push(sg_node);
}
Ok(vec![v])
}
fn call_sg_node(
(tsfn, rule): &(
ThreadsafeFunction<PinnedNodes, ErrorStrategy::CalleeHandled>,
RuleCore<SupportLang>,
),
entry: std::result::Result<ignore::DirEntry, ignore::Error>,
lang_option: &LangOption,
) -> Ret<bool> {
let entry = entry?;
if !entry
.file_type()
.context("could not use stdin as file")?
.is_file()
{
return Ok(false);
}
let (root, path) = get_root(entry, lang_option)?;
let mut pinned = PinnedNodeData::new(root.inner, |r| r.root().find_all(rule).collect());
let hits: &Vec<_> = pinned.get_data();
if hits.is_empty() {
return Ok(false);
}
let pinned = PinnedNodes(pinned, path);
tsfn.call(Ok(pinned), ThreadsafeFunctionCallMode::Blocking);
Ok(true)
}