Skip to content

Commit

Permalink
Cache Treesitter to map to reduce memory use (openpilot-hub#67)
Browse files Browse the repository at this point in the history
* tree sitter cache

* cache parser to map

* check style
  • Loading branch information
xiangtianyu authored Aug 19, 2024
1 parent 8ca64ff commit fe0f159
Showing 1 changed file with 21 additions and 14 deletions.
35 changes: 21 additions & 14 deletions src/main/java/com/zhongan/devpilot/treesitter/TreeSitterParser.java
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
import com.zhongan.devpilot.util.LanguageUtil;

import java.util.Locale;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.treesitter.TSLanguage;
import org.treesitter.TSParser;
Expand All @@ -14,6 +16,15 @@
public class TreeSitterParser {
private final TSLanguage language;

private final static Map<String, TreeSitterParser> parserMap = new ConcurrentHashMap<>();

static {
parserMap.put("default", new TreeSitterParser(null));
parserMap.put("java", new TreeSitterParser(new TreeSitterJava()));
parserMap.put("go", new TreeSitterParser(new TreeSitterGo()));
parserMap.put("python", new TreeSitterParser(new TreeSitterPython()));
}

public TreeSitterParser(TSLanguage language) {
this.language = language;
}
Expand Down Expand Up @@ -89,23 +100,19 @@ public static TreeSitterParser getInstance(String extension) {
var language = LanguageUtil.getLanguageByExtension(extension);

if (language == null) {
return new TreeSitterParser(null);
return getDefaultParser();
}

TSLanguage tsLanguage = null;

switch (language.getLanguageName().toLowerCase(Locale.ROOT)) {
case "java":
tsLanguage = new TreeSitterJava();
break;
case "go":
tsLanguage = new TreeSitterGo();
break;
case "python":
tsLanguage = new TreeSitterPython();
break;
var parser = parserMap.get(language.getLanguageName().toLowerCase(Locale.ROOT));

if (parser == null) {
return getDefaultParser();
}

return new TreeSitterParser(tsLanguage);
return parser;
}

private static TreeSitterParser getDefaultParser() {
return parserMap.get("default");
}
}

0 comments on commit fe0f159

Please sign in to comment.