feat: devShell: add bacon to flake

chore(deps): bump tracing-subscriber from 0.3.19 to 0.3.20

.github/dependabot.yml

@@ -0,0 +1,8 @@
+version: 2
+updates:
+  - package-ecosystem: "cargo" # See documentation for possible values
+    directory: "/" # Location of package manifests
+    schedule:
+      interval: "weekly"
+    open-pull-requests-limit: 4
+

.github/workflows/release.yml

@@ -0,0 +1,26 @@
+on:
+  release:
+    types: [created]
+
+jobs:
+  release:
+    name: release ${{ matrix.target }}
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - target: x86_64-pc-windows-gnu
+            archive: zip
+          - target: x86_64-unknown-linux-musl
+            archive: tar.zst
+    steps:
+      - uses: actions/checkout@master
+      - name: Compile and release
+        uses: rust-build/rust-build.action@v1.4.5
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        with:
+          RUSTTARGET: ${{ matrix.target }}
+          ARCHIVE_TYPES: ${{ matrix.archive }}
+          TOOLCHAIN_VERSION: stable

.github/workflows/rust.yml

@@ -0,0 +1,22 @@
+name: Build and test
+
+on:
+  push:
+    branches: [ "master" ]
+  pull_request:
+    branches: [ "master" ]
+
+env:
+  CARGO_TERM_COLOR: always
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Build
+      run: cargo build --verbose
+    - name: Run tests
+      run: cargo test --verbose

.helix/languages.toml

@@ -1,6 +1,11 @@
 [language-server.silos]
-command = "silos"
+command = "./target/debug/silos"
+args = ["lsp"]
 
 [[language]]
 name = "go"
-language-servers = [ { name = "silos" } ]
+language-servers = [ { name = "silos" }, "gopls" ]
+
+[[language]]
+name = "rust"
+language-servers = [ ]

.vscode/settings.json

@@ -9,6 +9,7 @@
             ],
             "command": [
                 "silos"
+                "lsp"
             ]
         }
     ]

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "silos"
-version = "4.0.0"
+version = "6.0.0"
 edition = "2024"
 
 [dependencies]
@@ -8,20 +8,20 @@ anyhow = "1.0.98"
 candle-core = "0.9.1"
 candle-nn = "0.9.1"
 candle-transformers = "0.9.1"
-clap = { version = "4.5.39", features = ["derive"] }
+clap = { version = "4.5.45", features = ["derive"] }
 derive_more = { version = "2.0.1", features = ["display", "error"] }
-glob = "0.3.2"
 hf-hub = "0.4.2"
 hora = "0.1.1"
 kdl = "6.3.4"
-serde = "1.0.219"
 serde_json = "1.0.140"
-tokenizers = "0.21.1"
+tokenizers = "0.21.4"
 tracing = "0.1.41"
-tracing-subscriber = "0.3.19"
-tree-sitter = "0.25.6"
+tracing-subscriber = "0.3.20"
+tree-sitter = "0.25.8"
 tree-sitter-c = "0.24.1"
 tree-sitter-go = "0.23.4"
 tree-sitter-rust = "0.24.0"
 tokio = { version = "1.45.1", features = ["io-std", "macros", "rt", "rt-multi-thread"] }
 tower-lsp = "0.20.0"
+tree-sitter-javascript = "0.25.0"
+tree-sitter-cpp = "0.23.4"

README.md

@@ -2,11 +2,11 @@
 
 Dumb, proomptable modular snippet search.
 
+![preview](./assets/preview.gif)
+
 ## Installation
 
-### Binary releases
-
-There are no binary releases yet.
+You can download a binary from releases tab or build the project from source.
 
 ### From source
 
@@ -76,13 +76,15 @@ This API parses code into an AST (Abstract Syntax Tree) via tree-sitter and can
 - C
 - Rust
 - Go
+- Javascript
+- C++
 
 ### Defining mutation collections
 
 ``` kdl
 description "describes the mutation collection"
 mutation {
-  expression "some ((beautiful) @adjective) AST expression"
+  expression "(some ((beautiful) @adjective) AST expression) @root"
   substitute {
     literal "hello"
     capture "adjective"
@@ -91,7 +93,7 @@ mutation {
 }
 
 mutation {
-  expression "another"
+  expression "(another) @root"
   substitute {
     literal "multiple mutations work"
     literal "as long as their expression"
@@ -102,14 +104,30 @@ mutation {
 
 - `description`: A textual description of the mutation collection.
 - `mutation`:  Defines individual code changes.
-  - `expression`: Uses tree-sitter to match and capture AST nodes with `@` prefixes, The special `@root` node is reserved for the entire expression.
+  - `expression`: Uses tree-sitter to match and capture AST nodes with `@` prefixes,
+  - The special `@root` node must be specify the expression to be replaced.
   - `substitute`:  Constructs the modified code using literals and captured arguments.
 
-See the example mutation collection in `./snippets/v2/go/mutations.kdl`.
+See the example mutation collection in `./snippets/v2/go/filepath-parent.kdl`.
 
 - The API performs a single-pass substitution based on the closest matching mutation.
 - Captured groups are used within the `substitute` block and the mutated code is returned.
 
+> Every capture group must contain the largest atom to be operated on.
+For example: if you wish to operate on elements of an array, capture each identifier inside the array
+
+Correct way: Here the `array` and `identifier` only hints about where the expression `root` lies.
+
+```
+(array (identifier @root))
+```
+
+Incorrect way: Here the root expression matches the block all the array elements inside the braces, not each element.
+
+```
+(array ((identifier)*) @entire-block-capture) @root
+```
+
 **Further reading**
 
 - [tree-sitter query snytax](https://tree-sitter.github.io/tree-sitter/using-parsers/queries/1-syntax.html) to create mutation expressions.

flake.lock

@@ -2,11 +2,11 @@
   "nodes": {
     "nixpkgs": {
       "locked": {
-        "lastModified": 1749776303,
-        "narHash": "sha256-OHibOvVwKqO1qvRg0r3agtd1EagW4THBcoWT7QGgcNo=",
+        "lastModified": 1755020227,
+        "narHash": "sha256-gGmm+h0t6rY88RPTaIm3su95QvQIVjAJx558YUG4Id8=",
         "owner": "NixOS",
         "repo": "nixpkgs",
-        "rev": "6e7721e37bf00fa7ea44ac3cfc9d2411284ec3ef",
+        "rev": "695d5db1b8b20b73292501683a524e0bd79074fb",
         "type": "github"
       },
       "original": {

flake.nix

@@ -13,6 +13,7 @@
           packages = with pkgs; [
             stdenv.cc.cc
             pkg-config
+            bacon
           ];
 
           libraries = with pkgs; [

snippets/refactor/go/base64.kdl

@@ -0,0 +1,13 @@
+description "base64 import"
+mutation {
+	expression "(import_spec_list ((import_spec)* @spec)) @root"
+	substitute {
+		literal "("
+		literal "\n"
+		capture "spec"
+		literal "\n"
+		literal #""base64""#
+		literal "\n"
+		literal ")"
+	}
+}

snippets/refactor/go/filepath-parent.kdl

@@ -4,7 +4,7 @@ mutation {
 	(call_expression
 		function: (_) @func (#eq? @func "filepath.Base")
 		arguments: (_) @args
-	)
+	) @root
 	"""
 	substitute {
 		literal "filepath.Base(filepath.Dir(filepath.Clean"

src/args.rs

@@ -1,8 +1,15 @@
-use clap::Parser;
+use clap::{Args, Parser, Subcommand};
+use std::path::PathBuf;
 
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
-pub(crate) struct Args {
+pub(crate) struct Cli {
+    #[command(subcommand)]
+    pub command: Command,
+}
+
+#[derive(Args, Debug)]
+pub(crate) struct Lsp {
     /// Run on the Nth GPU device.
     #[arg(long)]
     pub(crate) gpu: Option<usize>,
@@ -14,9 +21,50 @@ pub(crate) struct Args {
     /// Revision or branch.
     #[arg(long)]
     pub(crate) revision: Option<String>,
+
+    /// Path to the directory containing `generate` and `refactor` snippets.
+    #[arg(long, default_value = "./snippets")]
+    pub(crate) snippets: std::path::PathBuf,
 }
 
-impl Args {
+#[derive(Args, Debug)]
+pub struct DumpExpression {
+    pub path: PathBuf,
+}
+
+#[derive(Args, Debug)]
+pub struct ShowCaptures {
+    pub path: PathBuf,
+    pub expression: String,
+}
+
+#[derive(Args, Debug)]
+pub struct DryRun {
+    pub path: PathBuf,
+    pub edit_file: PathBuf,
+}
+
+#[derive(Subcommand, Debug)]
+pub enum Ast {
+    /// Dump the S expression for a given source file
+    DumpExpression(DumpExpression),
+    /// Show what parts of a source file gets captured by an S expression
+    ShowCaptures(ShowCaptures),
+
+    /// Test your edit snippets on a sample file
+    DryRun(DryRun),
+}
+
+#[derive(Subcommand, Debug)]
+pub enum Command {
+    /// quick actions to dump, modify and verify abstract syntax trees
+    #[command(subcommand)]
+    Ast(Ast),
+    /// spawn a language server for use with a text editor
+    Lsp(Lsp),
+}
+
+impl Lsp {
     pub(crate) fn resolve_model_and_revision(&self) -> (String, String) {
         let default_model = "sentence-transformers/all-MiniLM-L6-v2".to_string();
         let default_revision = "refs/pr/21".to_string();

src/embed.rs

@@ -7,17 +7,24 @@ use hf_hub::Repo;
 use hf_hub::RepoType;
 use hf_hub::api::sync::Api;
 use std::path::PathBuf;
-use tokenizers::Tokenizer;
-use tokenizers::TokenizerImpl;
+use tokenizers::DecoderWrapper;
 use tokenizers::ModelWrapper;
 use tokenizers::NormalizerWrapper;
-use tokenizers::PreTokenizerWrapper;
 use tokenizers::PostProcessorWrapper;
-use tokenizers::DecoderWrapper;
+use tokenizers::PreTokenizerWrapper;
+use tokenizers::Tokenizer;
+use tokenizers::TokenizerImpl;
 
 pub struct Embed {
     model: BertModel,
-    tokenizer: TokenizerImpl<ModelWrapper, NormalizerWrapper, PreTokenizerWrapper, PostProcessorWrapper, DecoderWrapper>,
+    pub hidden_size: usize,
+    tokenizer: TokenizerImpl<
+        ModelWrapper,
+        NormalizerWrapper,
+        PreTokenizerWrapper,
+        PostProcessorWrapper,
+        DecoderWrapper,
+    >,
 }
 
 impl Embed {
@@ -41,9 +48,14 @@ impl Embed {
         let tokenizer = tokenizer
             .with_padding(None)
             .with_truncation(None)
-            .map_err(E::msg)?.clone();
+            .map_err(E::msg)?
+            .clone();
 
-        Ok(Embed { model, tokenizer })
+        Ok(Embed {
+            model,
+            tokenizer,
+            hidden_size: config.hidden_size,
+        })
     }
 
     fn download_model_files(model_id: &str, revision: &str) -> Result<(PathBuf, PathBuf, PathBuf)> {
@@ -58,7 +70,8 @@ impl Embed {
     }
 
     pub(crate) fn embed(&self, prompt: &str) -> Result<Vec<f32>> {
-        let tokens = self.tokenizer
+        let tokens = self
+            .tokenizer
             .encode(prompt, true)
             .map_err(E::msg)?
             .get_ids()
@@ -68,9 +81,9 @@ impl Embed {
         let token_type_ids = token_ids.zeros_like()?;
 
         let embeddings = self.model.forward(&token_ids, &token_type_ids, None)?;
-        let (_n_sentence, n_tokens, _hidden_size) = embeddings.dims3()?;
-        let embeddings = (embeddings.sum(1)? / (n_tokens as f64))?;
-        let embeddings = normalize_l2(&embeddings)?.reshape(384)?.to_vec1::<f32>()?;
+        let embeddings = normalize_l2(&embeddings.sum(1)?)?
+            .reshape(self.hidden_size)?
+            .to_vec1::<f32>()?;
         Ok(embeddings)
     }
 }

src/lsp.rs

@@ -6,7 +6,7 @@ use tower_lsp::{Client, LanguageServer};
 
 pub struct Backend {
     pub client: Client,
-    pub body: Arc<Mutex<String>>,
+    pub body: Arc<Mutex<HashMap<Url, String>>>,
     pub appstate: crate::State,
 }
 
@@ -60,13 +60,18 @@ impl LanguageServer for Backend {
     }
 
     async fn did_open(&self, params: DidOpenTextDocumentParams) {
-        // TODO: build an index for multiple documents in workdir
-        *self.body.lock().await = params.text_document.text;
+        self.body
+            .lock()
+            .await
+            .insert(params.text_document.uri, params.text_document.text);
     }
 
     async fn did_change(&self, params: DidChangeTextDocumentParams) {
         if let Some(body) = params.content_changes.into_iter().next() {
-            *self.body.lock().await = body.text;
+            self.body
+                .lock()
+                .await
+                .insert(params.text_document.uri, body.text);
         }
     }
 
@@ -77,14 +82,20 @@ impl LanguageServer for Backend {
         let uri = params.text_document.uri;
         let Some(lang) = url_extension(&uri) else {
             self.client
-                .log_message(MessageType::ERROR, "unable to determine filetype, file has no extension")
+                .log_message(
+                    MessageType::ERROR,
+                    "unable to determine filetype, file has no extension",
+                )
                 .await;
             return Ok(None);
         };
 
-        let body = self.body.lock().await.to_string();
+        let body_locked = self.body.lock().await;
+        let Some(body) = body_locked.get(&uri) else {
+            return Ok(None);
+        };
         let mut range = params.range;
-        let selected_text = string_range_index(&body, range);
+        let selected_text = string_range_index(body, range);
 
         let Some(comment) = ParsedAction::new(selected_text) else {
             return Ok(None);
@@ -93,14 +104,15 @@ impl LanguageServer for Backend {
         let action_response = match comment.action {
             Action::Generate => {
                 range.start = range.end;
-                self.appstate.generate(&lang, comment.description, 1)
+                self.appstate
+                    .generate(&lang, comment.description, 1)
                     .map(|v| v.into_iter().map(|s| format!("{s}\n")).collect())
                     .map_err(|e| e.to_string())
             }
-            Action::Refactor => {
-                self.appstate.refactor(&lang, comment.description, selected_text, 1)
-                    .map_err(|e| e.to_string())
-            }
+            Action::Refactor => self
+                .appstate
+                .refactor(&lang, comment.description, selected_text, 1)
+                .map_err(|e| e.to_string()),
         };
 
         let closest_matches = match action_response {

src/main.rs

@@ -1,9 +1,9 @@
-use anyhow::{Context, Error as E, Result, bail};
+use anyhow::{Context, Error as E, Result};
 use clap::Parser;
 use hora::core::{ann_index::ANNIndex, metrics::Metric::Euclidean};
 use hora::index::hnsw_idx::HNSWIndex;
 use kdl::KdlDocument;
-use state::State;
+use state::{State, dump_expression};
 use std::collections::HashMap;
 use std::sync::Arc;
 use tokio::sync::Mutex;
@@ -12,53 +12,75 @@ use tower_lsp::{LspService, Server};
 mod args;
 mod embed;
 mod lsp;
-mod state;
 mod mutation;
-
-fn path_to_parent_base(p: &std::path::Path) -> Result<String> {
-    let Some(parent) = p
-        .parent()
-        .and_then(|v| v.file_name())
-        .and_then(|v| v.to_str())
-        .map(|v| v.to_string())
-    else {
-        bail!("failed to parse snippets path, make sure the directory structure is valid");
-    };
-    Ok(parent)
-}
+mod sources;
+mod state;
 
 #[tokio::main]
 async fn main() -> Result<()> {
     tracing_subscriber::fmt::init();
-    let args = args::Args::parse();
+    let args = match args::Cli::parse().command {
+        args::Command::Ast(ast) => {
+            match ast {
+                args::Ast::DumpExpression(source_file) => {
+                    println!("{}", dump_expression(&source_file.path)?);
+                }
+                args::Ast::ShowCaptures(show_captures) => {
+                    let source_bytes = std::fs::read(&show_captures.path)?;
+                    let langfn = state::lang_from_file_extension(&show_captures.path)?;
+                    let tree = state::parse_into_tree(&source_bytes, &langfn)?;
+                    let root_node = tree.root_node();
+                    let cooked = mutation::query(
+                        root_node,
+                        &show_captures.expression,
+                        &langfn,
+                        &source_bytes,
+                    );
+                    println!("{:#?}", cooked);
+                }
+                args::Ast::DryRun(dry_run) => {
+                    let mutation_collection = mutation::from_path(dry_run.edit_file)?;
+                    let source_bytes = std::fs::read(&dry_run.path)?;
+                    let langfn = state::lang_from_file_extension(&dry_run.path)?;
+                    let tree = state::parse_into_tree(&source_bytes, &langfn)?;
+                    let root_node = tree.root_node();
+                    let cooked =
+                        mutation::apply(langfn, &source_bytes, root_node, &mutation_collection)?;
+                    println!("{cooked}");
+                }
+            }
+            return Ok(());
+        }
+        args::Command::Lsp(lsp) => lsp,
+    };
+
     let (model_id, revision) = args.resolve_model_and_revision();
+
     let embed = embed::Embed::new(args.gpu, &model_id, &revision)?;
     let mut dict = HashMap::default();
-    let dimensions = 384;
+    let dimensions = embed.hidden_size;
 
-    let paths = glob::glob("./snippets/v1/*/*.kdl")?;
-    for path in paths {
-        let path = path?;
-        let parent = path_to_parent_base(&path)?;
+    for (language, paths) in sources::rule_files(args.snippets.join("generate"))? {
+        for path in paths {
+            let current_lang_index = dict
+                .entry(language.clone())
+                .or_insert_with(|| HNSWIndex::new(dimensions, &Default::default()));
 
-        let current_lang_index = dict
-            .entry(parent)
-            .or_insert_with(|| HNSWIndex::new(dimensions, &Default::default()));
+            let doc_str = std::fs::read_to_string(&path)?;
+            let doc: KdlDocument = doc_str
+                .parse()
+                .context(format!("failed to parse KDL: {}", path.display()))?;
 
-        let doc_str = std::fs::read_to_string(&path)?;
-        let doc: KdlDocument = doc_str
-            .parse()
-            .context(format!("failed to parse KDL: {}", path.display()))?;
-
-        let Some(desc) = doc.get_arg("desc").and_then(|v| v.as_string()) else {
-            continue;
-        };
-        let Some(body) = doc.get_arg("body").and_then(|v| v.as_string()) else {
-            continue;
-        };
-        current_lang_index
-            .add(&embed.embed(desc)?, body.to_string())
-            .map_err(E::msg)?;
+            let Some(desc) = doc.get_arg("desc").and_then(|v| v.as_string()) else {
+                continue;
+            };
+            let Some(body) = doc.get_arg("body").and_then(|v| v.as_string()) else {
+                continue;
+            };
+            current_lang_index
+                .add(&embed.embed(desc)?, body.to_string())
+                .map_err(E::msg)?;
+        }
     }
 
     for index in dict.values_mut() {
@@ -67,45 +89,45 @@ async fn main() -> Result<()> {
             .map_err(E::msg)?;
     }
 
-    // v2
-    let paths = glob::glob("./snippets/v2/*/*.kdl")?;
-    let mut v2_dict = HashMap::new();
-    let mut v2_mutations_collection = vec![];
-    for (i, path) in paths.enumerate() {
-        let path = path?;
-        let parent = path_to_parent_base(&path)?;
+    let mut refactor_dict = HashMap::new();
+    let mut mutations_collection = vec![];
+    for (language, paths) in sources::rule_files(args.snippets.join("refactor"))? {
+        for path in paths {
+            let mutations = mutation::from_path(path)?;
+            let current_lang_index = refactor_dict
+                .entry(language.clone())
+                .or_insert_with(|| HNSWIndex::new(dimensions, &Default::default()));
 
-        let mutations = mutation::from_path(path)?;
-        let current_lang_index = v2_dict
-            .entry(parent)
-            .or_insert_with(|| HNSWIndex::new(dimensions, &Default::default()));
-
-        current_lang_index
-            .add(&embed.embed(&mutations.description)?, i)
-            .map_err(E::msg)?;
-        v2_mutations_collection.push(mutations);
+            current_lang_index
+                .add(
+                    &embed.embed(&mutations.description)?,
+                    mutations_collection.len(),
+                )
+                .map_err(E::msg)?;
+            mutations_collection.push(mutations);
+        }
     }
 
-    for index in v2_dict.values_mut() {
+    for index in refactor_dict.values_mut() {
         index.build(Euclidean).map_err(E::msg)?;
     }
 
-    let appstate = State {
+    let appstate = State::new(
         embed,
-        generate: state::Generate { dict },
-        refactor: state::Refactor {
-            dict: v2_dict,
-            mutations_collection: v2_mutations_collection,
+        state::Generate { dict },
+        state::Refactor {
+            dict: refactor_dict,
+            mutations_collection,
         },
-    };
+    );
 
     let stdin = tokio::io::stdin();
     let stdout = tokio::io::stdout();
 
     let (service, socket) = LspService::new(|client| lsp::Backend {
         client,
-        body: Arc::new(Mutex::new(String::default())),
-        appstate
+        body: Arc::new(Mutex::new(HashMap::default())),
+        appstate,
     });
     Server::new(stdin, stdout, socket).serve(service).await;
     Ok(())

src/mutation.rs

@@ -72,10 +72,8 @@ pub fn from_path<P: AsRef<Path>>(path: P) -> Result<MutationCollection> {
             substitute.push(substitutor);
         }
 
-        let expression = format!("({expression}) @root");
-
         mutations.push(Mutation {
-            expression,
+            expression: expression.to_string(),
             substitute,
         })
     }
@@ -127,7 +125,7 @@ pub fn apply(
 }
 
 #[derive(Debug)]
-struct QueryCooked {
+pub struct QueryCooked {
     captures: HashMap<String, String>,
     end: usize,
     start: usize,
@@ -152,7 +150,7 @@ fn split_at_indices<'a>(c: &'a [u8], idx: &[usize]) -> SplitMap<'a> {
     SplitMap { values, indices }
 }
 
-fn query<'a>(
+pub fn query<'a>(
     node: Node<'a>,
     expr: &'a str,
     lang: &Language,
@@ -164,6 +162,7 @@ fn query<'a>(
     let mut query_matches = qc.matches(&query, node, source_bytes);
 
     let capture_names = query.capture_names();
+    // println!("names: {capture_names:#?}");
 
     let mut cooked = vec![];
 
@@ -171,19 +170,35 @@ fn query<'a>(
         let mut capture_cooked = HashMap::new();
         let mut start = 0;
         let mut end = 0;
-        for cap in matcha.captures {
-            let Some(name) = capture_names.get(cap.index as usize) else {
+        if matcha.captures.is_empty() {
+            continue;
+        }
+        //     println!("match {:#?}", matcha.id());
+
+        for (ix, name) in capture_names.iter().enumerate() {
+            let nodes = matcha.nodes_for_capture_index(ix.try_into().unwrap());
+            let mut start_pos = None;
+            let mut end_pos = None;
+            debug!("matches for {name}");
+            for node in nodes {
+                start_pos.get_or_insert(node.start_byte());
+                end_pos.replace(node.end_byte());
+                debug!("hit {node:#?}");
+            }
+
+            let (Some(start_pos), Some(end_pos)) = (start_pos, end_pos) else {
                 continue;
             };
+
             if *name == "root" {
-                start = cap.node.start_byte();
-                end = cap.node.end_byte();
-                continue;
+                start = start_pos;
+                end = end_pos;
             }
-            capture_cooked.insert(
-                name.to_string(),
-                cap.node.utf8_text(source_bytes).unwrap().to_string(),
-            );
+
+            let text_bytes = &source_bytes[start_pos..end_pos];
+            let text = std::str::from_utf8(text_bytes).unwrap();
+            //         println!("text: {text}");
+            capture_cooked.insert(name.to_string(), text.to_string());
         }
         cooked.push(QueryCooked {
             start,

src/sources.rs

@@ -0,0 +1,34 @@
+use std::{
+    collections::HashMap,
+    fs, io,
+    path::{Path, PathBuf},
+};
+
+pub fn rule_files<P: AsRef<Path>>(path: P) -> io::Result<HashMap<String, Vec<PathBuf>>> {
+    let per_language_dirs: Vec<_> = fs::read_dir(path)?
+        .filter_map(|res| res.ok())
+        .map(|direntry| direntry.path())
+        .filter(|dir| dir.is_dir())
+        .collect();
+
+    let mut basename_to_paths = HashMap::new();
+
+    for language_dir in per_language_dirs {
+        let Some(dirname) = language_dir
+            .file_stem()
+            .and_then(|v| v.to_str())
+            .map(|v| v.to_string())
+        else {
+            continue;
+        };
+        let rule_file_paths: Vec<_> = fs::read_dir(&language_dir)?
+            .filter_map(|res| res.ok())
+            .map(|entry| entry.path())
+            .filter(|file| file.is_file() && file.extension().is_some_and(|ext| ext == "kdl"))
+            .map(|path| path.to_path_buf())
+            .collect();
+        basename_to_paths.insert(dirname, rule_file_paths);
+    }
+    Ok(basename_to_paths)
+}
+// fn prebuilt_index();

src/state.rs

@@ -1,10 +1,11 @@
+use crate::mutation;
 use derive_more::Display;
 use derive_more::Error;
-use tree_sitter::Parser;
-use std::collections::HashMap;
-use hora::index::hnsw_idx::HNSWIndex;
 use hora::core::ann_index::ANNIndex;
-use crate::mutation;
+use hora::index::hnsw_idx::HNSWIndex;
+use std::collections::HashMap;
+use std::path::Path;
+use tree_sitter::Parser;
 
 #[derive(Debug, Display, Error)]
 pub enum Error {
@@ -22,16 +23,6 @@ pub struct Refactor {
 }
 
 impl Refactor {
-    fn get_lang(s: &str) -> Result<tree_sitter::Language, Error> {
-        Ok(match s {
-            "go" => tree_sitter_go::LANGUAGE,
-            "c" => tree_sitter_c::LANGUAGE,
-            "rs" => tree_sitter_rust::LANGUAGE,
-            _ => return Err(Error::UnknownLang),
-        }
-        .into())
-    }
-
     pub fn search(
         &self,
         lang: &str,
@@ -39,17 +30,9 @@ impl Refactor {
         body: &str,
         top_k: usize,
     ) -> Result<Vec<String>, Error> {
-        let langfn = Self::get_lang(lang)?;
-        let mut parser = Parser::new();
-        parser
-            .set_language(&langfn)
-            .map_err(|_| Error::UnknownLang)?;
-
-        let source_code = body;
-        let source_bytes = source_code.as_bytes();
-        let tree = parser
-            .parse(source_code, None)
-            .ok_or(Error::SnippetParsing)?;
+        let langfn = lang_from_name(lang)?;
+        let source_bytes = body.as_bytes();
+        let tree = parse_into_tree(source_bytes, &langfn)?;
         let root_node = tree.root_node();
 
         // search for k nearest neighbors
@@ -68,7 +51,8 @@ impl Refactor {
                     Err(e) => {
                         tracing::error!(
                             collection_index = index,
-                            "failed to apply mutations from collection {}", e
+                            "failed to apply mutations from collection {}",
+                            e
                         );
                         None
                     }
@@ -78,8 +62,50 @@ impl Refactor {
         Ok(collected)
     }
 }
+
+pub fn lang_from_name(s: &str) -> Result<tree_sitter::Language, Error> {
+    Ok(match s {
+        "go" => tree_sitter_go::LANGUAGE,
+        "c" | "h" => tree_sitter_c::LANGUAGE,
+        "cpp" | "hpp" => tree_sitter_cpp::LANGUAGE,
+        "js" | "ts" => tree_sitter_javascript::LANGUAGE,
+        "rs" => tree_sitter_rust::LANGUAGE,
+        _ => return Err(Error::UnknownLang),
+    }
+    .into())
+}
+
+pub fn lang_from_file_extension(path: &Path) -> Result<tree_sitter::Language, Error> {
+    let Some(lang) = path.extension() else {
+        return Err(Error::UnknownLang);
+    };
+    let lang = lang.to_str().ok_or(Error::UnknownLang)?;
+    lang_from_name(lang)
+}
+
+// parses `body` written in the language `langfn` into tree sitter AST
+pub fn parse_into_tree(
+    body: &[u8],
+    langfn: &tree_sitter::Language,
+) -> Result<tree_sitter::Tree, Error> {
+    let mut parser = Parser::new();
+    parser
+        .set_language(langfn)
+        .map_err(|_| Error::UnknownLang)?;
+    let tree = parser.parse(body, None).ok_or(Error::SnippetParsing)?;
+    Ok(tree)
+}
+
+pub fn dump_expression(path: &Path) -> Result<String, Error> {
+    let source_bytes = std::fs::read(path).map_err(|_| Error::SnippetParsing)?;
+
+    let tree = parse_into_tree(&source_bytes, &lang_from_file_extension(path)?)?;
+
+    Ok(tree.root_node().to_sexp().to_string())
+}
+
 pub struct Generate {
-    pub dict: HashMap<String, HNSWIndex<f32, String>>
+    pub dict: HashMap<String, HNSWIndex<f32, String>>,
 }
 
 impl Generate {
@@ -92,14 +118,19 @@ impl Generate {
 }
 
 pub struct State {
-    // TODO: create new constructor and private these fields
-    pub embed: crate::embed::Embed,
-    pub generate: Generate,
-    pub refactor: Refactor,
+    embed: crate::embed::Embed,
+    generate: Generate,
+    refactor: Refactor,
 }
 
 impl State {
-    
+    pub fn new(embed: crate::embed::Embed, generate: Generate, refactor: Refactor) -> Self {
+        Self {
+            embed,
+            generate,
+            refactor,
+        }
+    }
     pub fn generate(&self, lang: &str, prompt: &str, top_k: usize) -> Result<Vec<String>, Error> {
         let Ok(target) = self.embed.embed(prompt) else {
             return Err(Error::EmbedFailed);
@@ -108,7 +139,13 @@ impl State {
         self.generate.search(lang, &target, top_k)
     }
 
-    pub fn refactor(&self, lang: &str, prompt: &str, body: &str, top_k: usize) -> Result<Vec<String>, Error> {
+    pub fn refactor(
+        &self,
+        lang: &str,
+        prompt: &str,
+        body: &str,
+        top_k: usize,
+    ) -> Result<Vec<String>, Error> {
         let Ok(target) = self.embed.embed(prompt) else {
             return Err(Error::EmbedFailed);
         };