Added go parsing for semantic search, and added preceeding comments on go and rust.

Co-authored-by: Alex <alexviscreanu@gmail.com>
Co-authored-by: maxbrunsfeld <max@zed.dev>
This commit is contained in:
KCaverly 2023-07-13 14:33:31 -04:00
parent a56d454a07
commit 5eab628580
4 changed files with 55 additions and 54 deletions

View file

@ -525,7 +525,6 @@ pub struct EmbeddingConfig {
pub item_capture_ix: u32, pub item_capture_ix: u32,
pub name_capture_ix: u32, pub name_capture_ix: u32,
pub context_capture_ix: Option<u32>, pub context_capture_ix: Option<u32>,
pub extra_context_capture_ix: Option<u32>,
} }
struct InjectionConfig { struct InjectionConfig {
@ -1246,14 +1245,12 @@ impl Language {
let mut item_capture_ix = None; let mut item_capture_ix = None;
let mut name_capture_ix = None; let mut name_capture_ix = None;
let mut context_capture_ix = None; let mut context_capture_ix = None;
let mut extra_context_capture_ix = None;
get_capture_indices( get_capture_indices(
&query, &query,
&mut [ &mut [
("item", &mut item_capture_ix), ("item", &mut item_capture_ix),
("name", &mut name_capture_ix), ("name", &mut name_capture_ix),
("context", &mut context_capture_ix), ("context", &mut context_capture_ix),
("context.extra", &mut extra_context_capture_ix),
], ],
); );
if let Some((item_capture_ix, name_capture_ix)) = item_capture_ix.zip(name_capture_ix) { if let Some((item_capture_ix, name_capture_ix)) = item_capture_ix.zip(name_capture_ix) {
@ -1262,7 +1259,6 @@ impl Language {
item_capture_ix, item_capture_ix,
name_capture_ix, name_capture_ix,
context_capture_ix, context_capture_ix,
extra_context_capture_ix,
}); });
} }
Ok(self) Ok(self)

View file

@ -53,7 +53,7 @@ impl CodeContextRetriever {
.ok_or_else(|| anyhow!("parsing failed"))?; .ok_or_else(|| anyhow!("parsing failed"))?;
let mut documents = Vec::new(); let mut documents = Vec::new();
let mut context_spans = Vec::new(); let mut document_texts = Vec::new();
// Iterate through query matches // Iterate through query matches
for mat in self.cursor.matches( for mat in self.cursor.matches(
@ -61,11 +61,10 @@ impl CodeContextRetriever {
tree.root_node(), tree.root_node(),
content.as_bytes(), content.as_bytes(),
) { ) {
// log::info!("-----MATCH-----");
let mut name: Vec<&str> = vec![]; let mut name: Vec<&str> = vec![];
let mut item: Option<&str> = None; let mut item: Option<&str> = None;
let mut offset: Option<usize> = None; let mut offset: Option<usize> = None;
let mut context_spans: Vec<&str> = vec![];
for capture in mat.captures { for capture in mat.captures {
if capture.index == embedding_config.item_capture_ix { if capture.index == embedding_config.item_capture_ix {
offset = Some(capture.node.byte_range().start); offset = Some(capture.node.byte_range().start);
@ -79,25 +78,21 @@ impl CodeContextRetriever {
if let Some(context_capture_ix) = embedding_config.context_capture_ix { if let Some(context_capture_ix) = embedding_config.context_capture_ix {
if capture.index == context_capture_ix { if capture.index == context_capture_ix {
if let Some(context) = content.get(capture.node.byte_range()) { if let Some(context) = content.get(capture.node.byte_range()) {
name.push(context); context_spans.push(context);
} }
} }
} }
} }
if item.is_some() && offset.is_some() && name.len() > 0 { if item.is_some() && offset.is_some() && name.len() > 0 {
let context_span = CODE_CONTEXT_TEMPLATE let item = format!("{}\n{}", context_spans.join("\n"), item.unwrap());
let document_text = CODE_CONTEXT_TEMPLATE
.replace("<path>", pending_file.relative_path.to_str().unwrap()) .replace("<path>", pending_file.relative_path.to_str().unwrap())
.replace("<language>", &pending_file.language.name().to_lowercase()) .replace("<language>", &pending_file.language.name().to_lowercase())
.replace("<item>", item.unwrap()); .replace("<item>", item.as_str());
let mut truncated_span = context_span.clone(); document_texts.push(document_text);
truncated_span.truncate(100);
// log::info!("Name: {:?}", name);
// log::info!("Span: {:?}", truncated_span);
context_spans.push(context_span);
documents.push(Document { documents.push(Document {
name: name.join(" "), name: name.join(" "),
offset: offset.unwrap(), offset: offset.unwrap(),
@ -112,7 +107,7 @@ impl CodeContextRetriever {
mtime: pending_file.modified_time, mtime: pending_file.modified_time,
documents, documents,
}, },
context_spans, document_texts,
)); ));
} }
} }

View file

@ -0,0 +1,24 @@
(
(comment)* @context
.
(type_declaration
(type_spec
name: (_) @name)
) @item
)
(
(comment)* @context
.
(function_declaration
name: (_) @name
) @item
)
(
(comment)* @context
.
(method_declaration
name: (_) @name
) @item
)

View file

@ -1,36 +1,22 @@
(struct_item (
(visibility_modifier)? @context (line_comment)* @context
"struct" @context .
name: (_) @name) @item [
(enum_item (enum_item
(visibility_modifier)? @context
"enum" @context
name: (_) @name) @item name: (_) @name) @item
(struct_item
name: (_) @name) @item
(impl_item (impl_item
"impl" @context
trait: (_)? @name trait: (_)? @name
"for"? @context "for"? @name
type: (_) @name) @item type: (_) @name) @item
(trait_item (trait_item
(visibility_modifier)? @context
"trait" @context
name: (_) @name) @item name: (_) @name) @item
(function_item (function_item
(visibility_modifier)? @context
(function_modifiers)? @context
"fn" @context
name: (_) @name) @item name: (_) @name) @item
(function_signature_item
(visibility_modifier)? @context
(function_modifiers)? @context
"fn" @context
name: (_) @name) @item
(macro_definition (macro_definition
. "macro_rules!" @context
name: (_) @name) @item name: (_) @name) @item
(function_signature_item
name: (_) @name) @item
]
)