feat: avante repo map rust crate (yetone#628)

ogfunkycold · Sep 25, 2024 · 0d90c04 · 0d90c04
1 parent 5461342
commit 0d90c04
Show file tree

Hide file tree

Showing 23 changed files with 1,399 additions and 737 deletions.
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -29,6 +29,18 @@ jobs:
         uses: lunarmodules/luacheck@v1
         with:
           args: ./lua/
+  rust-tests:
+    name: Run Rust tests
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: Swatinem/rust-cache@v2
+      - uses: dtolnay/rust-toolchain@master
+        with:
+          toolchain: stable
+          components: clippy, rustfmt
+      - name: Run rust tests
+        run: cargo test --features luajit
   rust:
     name: Check Rust style
     runs-on: ubuntu-latest

diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml
@@ -74,6 +74,7 @@ jobs:
           fi
           cp target/release/libavante_templates.$EXT results/avante_templates.$EXT
           cp target/release/libavante_tokenizers.$EXT results/avante_tokenizers.$EXT
+          cp target/release/libavante_repo_map.$EXT results/avante_repo_map.$EXT
 
           cd results
           tar zcvf avante_lib-${{ matrix.os }}-${{ matrix.feature }}.tar.gz *.${EXT}
@@ -85,6 +86,7 @@ jobs:
 
           Copy-Item -Path "target\release\avante_templates.dll" -Destination "results\avante_templates.dll"
           Copy-Item -Path "target\release\avante_tokenizers.dll" -Destination "results\avante_tokenizers.dll"
+          Copy-Item -Path "target\release\avante_repo_map.dll" -Destination "results\avante_repo_map.dll"
 
           Set-Location -Path results
 

diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -11,6 +11,7 @@ version = "0.1.0"
 [workspace.dependencies]
 avante-tokenizers = { path = "crates/avante-tokenizers" }
 avante-templates = { path = "crates/avante-templates" }
+avante-repo-map = { path = "crates/avante-repo-map" }
 minijinja = { version = "2.2.0", features = [
   "loader",
   "json",

diff --git a/Makefile b/Makefile
@@ -22,13 +22,15 @@ all: luajit
 define make_definitions
 ifeq ($(BUILD_FROM_SOURCE),true)
 ifeq ($(TARGET_LIBRARY), all)
-$1: $(BUILD_DIR)/libAvanteTokenizers-$1.$(EXT) $(BUILD_DIR)/libAvanteTemplates-$1.$(EXT)
+$1: $(BUILD_DIR)/libAvanteTokenizers-$1.$(EXT) $(BUILD_DIR)/libAvanteTemplates-$1.$(EXT) $(BUILD_DIR)/libAvanteRepoMap-$1.$(EXT)
 else ifeq ($(TARGET_LIBRARY), tokenizers)
 $1: $(BUILD_DIR)/libAvanteTokenizers-$1.$(EXT)
 else ifeq ($(TARGET_LIBRARY), templates)
 $1: $(BUILD_DIR)/libAvanteTemplates-$1.$(EXT)
+else ifeq ($(TARGET_LIBRARY), repo-map)
+$1: $(BUILD_DIR)/libAvanteRepoMap-$1.$(EXT)
 else
-	$$(error TARGET_LIBRARY must be one of all, tokenizers, templates)
+	$$(error TARGET_LIBRARY must be one of all, tokenizers, templates, repo-map)
 endif
 else
 $1:
@@ -41,16 +43,18 @@ $(foreach lua_version,$(LUA_VERSIONS),$(eval $(call make_definitions,$(lua_versi
 define build_package
 $1-$2:
 	cargo build --release --features=$1 -p avante-$2
-	cp target/release/libavante_$2.$(EXT) $(BUILD_DIR)/avante_$2.$(EXT)
+	cp target/release/libavante_$(shell echo $2 | tr - _).$(EXT) $(BUILD_DIR)/avante_$(shell echo $2 | tr - _).$(EXT)
 endef
 
 define build_targets
 $(BUILD_DIR)/libAvanteTokenizers-$1.$(EXT): $(BUILD_DIR) $1-tokenizers
 $(BUILD_DIR)/libAvanteTemplates-$1.$(EXT): $(BUILD_DIR) $1-templates
+$(BUILD_DIR)/libAvanteRepoMap-$1.$(EXT): $(BUILD_DIR) $1-repo-map
 endef
 
 $(foreach lua_version,$(LUA_VERSIONS),$(eval $(call build_package,$(lua_version),tokenizers)))
 $(foreach lua_version,$(LUA_VERSIONS),$(eval $(call build_package,$(lua_version),templates)))
+$(foreach lua_version,$(LUA_VERSIONS),$(eval $(call build_package,$(lua_version),repo-map)))
 $(foreach lua_version,$(LUA_VERSIONS),$(eval $(call build_targets,$(lua_version))))
 
 $(BUILD_DIR):

diff --git a/crates/avante-repo-map/Cargo.toml b/crates/avante-repo-map/Cargo.toml
@@ -0,0 +1,38 @@
+[lib]
+crate-type = ["cdylib"]
+
+[package]
+name = "avante-repo-map"
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+version.workspace = true
+
+[build-dependencies]
+cc="*"
+
+[dependencies]
+mlua = { workspace = true }
+minijinja = { workspace = true }
+serde = { workspace = true, features = ["derive"] }
+tree-sitter = "0.23"
+tree-sitter-language = "0.1"
+tree-sitter-rust = "0.23"
+tree-sitter-python = "0.23"
+tree-sitter-javascript = "0.23"
+tree-sitter-typescript = "0.23"
+tree-sitter-go = "0.23"
+tree-sitter-c = "0.23"
+tree-sitter-cpp = "0.23"
+tree-sitter-lua = "0.2"
+tree-sitter-ruby = "0.23"
+
+[lints]
+workspace = true
+
+[features]
+lua51 = ["mlua/lua51"]
+lua52 = ["mlua/lua52"]
+lua53 = ["mlua/lua53"]
+lua54 = ["mlua/lua54"]
+luajit = ["mlua/luajit"]
diff --git a/crates/avante-repo-map/queries/tree-sitter-c-defs.scm b/crates/avante-repo-map/queries/tree-sitter-c-defs.scm
@@ -0,0 +1,11 @@
+;; Capture extern functions, variables, public classes, and methods
+(function_definition
+  (storage_class_specifier) @extern
+) @function
+(class_specifier
+  (public) @class
+  (function_definition) @method
+) @class
+(declaration
+  (storage_class_specifier) @extern
+) @variable
diff --git a/crates/avante-repo-map/queries/tree-sitter-cpp-defs.scm b/crates/avante-repo-map/queries/tree-sitter-cpp-defs.scm
@@ -0,0 +1,11 @@
+;; Capture extern functions, variables, public classes, and methods
+(function_definition
+  (storage_class_specifier) @extern
+) @function
+(class_specifier
+  (public) @class
+  (function_definition) @method
+) @class
+(declaration
+  (storage_class_specifier) @extern
+) @variable
diff --git a/crates/avante-repo-map/queries/tree-sitter-go-defs.scm b/crates/avante-repo-map/queries/tree-sitter-go-defs.scm
@@ -0,0 +1,18 @@
+;; Capture top-level functions and struct definitions
+(var_declaration
+  (var_spec) @variable
+)
+(const_declaration
+  (const_spec) @variable
+)
+(function_declaration) @function
+(type_declaration
+  (type_spec (struct_type)) @class
+)
+(type_declaration
+  (type_spec
+    (struct_type
+      (field_declaration_list
+        (field_declaration) @class_variable)))
+)
+(method_declaration) @method
diff --git a/crates/avante-repo-map/queries/tree-sitter-javascript-defs.scm b/crates/avante-repo-map/queries/tree-sitter-javascript-defs.scm
@@ -0,0 +1,23 @@
+;; Capture exported functions, arrow functions, variables, classes, and method definitions
+(export_statement
+  declaration: (lexical_declaration
+    (variable_declarator) @variable
+  )
+)
+(export_statement
+  declaration: (function_declaration) @function
+)
+(export_statement
+  declaration: (class_declaration
+    body: (class_body
+      (field_definition) @class_variable
+    )
+  )
+)
+(export_statement
+  declaration: (class_declaration
+    body: (class_body
+      (method_definition) @method
+    )
+  )
+)
diff --git a/crates/avante-repo-map/queries/tree-sitter-lua-defs.scm b/crates/avante-repo-map/queries/tree-sitter-lua-defs.scm
@@ -0,0 +1,3 @@
+;; Capture function and method definitions
+(variable_list) @variable
+(function_declaration) @function
diff --git a/crates/avante-repo-map/queries/tree-sitter-python-defs.scm b/crates/avante-repo-map/queries/tree-sitter-python-defs.scm
@@ -0,0 +1,25 @@
+;; Capture top-level functions, class, and method definitions
+(module
+  (expression_statement
+    (assignment) @assignment
+  )
+)
+(module
+  (function_definition) @function
+)
+(module
+  (class_definition
+    body: (block
+      (expression_statement
+        (assignment) @class_assignment
+      )
+    )
+  )
+)
+(module
+  (class_definition
+    body: (block
+      (function_definition) @method
+    )
+  )
+)
diff --git a/crates/avante-repo-map/queries/tree-sitter-ruby-defs.scm b/crates/avante-repo-map/queries/tree-sitter-ruby-defs.scm
@@ -0,0 +1,16 @@
+;; Capture top-level methods, class definitions, and methods within classes
+(program
+  (class
+    (body_statement
+      (call) @class_call
+      (assignment) @class_assignment
+      (method) @method
+    )
+  ) @class
+)
+(program
+  (method) @function
+)
+(program
+  (assignment) @assignment
+)