diff --git a/CHANGELOG.md b/CHANGELOG.md index f087610f8..fdb67e4cc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ * Add a Kotlin PSI-based tree generator alongside the Kotlin tree-sitter generator * Fix plugin discovery and restore the JSON generator in the CLI distribution * Make DOT diff labels configurable and preserve their original content +* Tree-sitter generators are 2–10× faster, label multi-line nodes with their source text, and compute correct offsets after non-ASCII characters and byte order marks ## v4.1.0 * Registries discover annotated matchers, generators, and clients when first accessed diff --git a/gen.treesitter-ng/src/main/java/com/github/gumtreediff/gen/treesitterng/AbstractTreeSitterNgGenerator.java b/gen.treesitter-ng/src/main/java/com/github/gumtreediff/gen/treesitterng/AbstractTreeSitterNgGenerator.java index 5a9a1bcd0..b5b021634 100644 --- a/gen.treesitter-ng/src/main/java/com/github/gumtreediff/gen/treesitterng/AbstractTreeSitterNgGenerator.java +++ b/gen.treesitter-ng/src/main/java/com/github/gumtreediff/gen/treesitterng/AbstractTreeSitterNgGenerator.java @@ -22,18 +22,22 @@ import com.github.gumtreediff.tree.Tree; import com.github.gumtreediff.tree.TreeContext; import com.github.gumtreediff.tree.TypeSet; -import com.github.gumtreediff.utils.Pair; import org.treesitter.*; import org.yaml.snakeyaml.Yaml; -import java.io.BufferedReader; +import java.io.IOException; import java.io.Reader; import java.nio.charset.StandardCharsets; import java.util.*; +import java.util.concurrent.ConcurrentHashMap; +/** + * Converts a tree-sitter tree into a GumTree tree in a single tree cursor traversal. + * Positions and lengths are UTF-16 offsets (the offsets of Java strings), converted from the UTF-8 byte + * offsets of tree-sitter, and labels are the source text of the nodes. + */ public abstract class AbstractTreeSitterNgGenerator extends TreeGenerator { - private static final String LF = "\n"; private static final String RULES_FILE = "rules.yml"; private static final String YAML_IGNORED = "ignored"; @@ -43,17 +47,15 @@ public abstract class AbstractTreeSitterNgGenerator extends TreeGenerator { private static final Map> RULES; + private static final Map LANGUAGE_RULES = new ConcurrentHashMap<>(); + static { Yaml yaml = new Yaml(); RULES = yaml.load(Thread.currentThread().getContextClassLoader().getResourceAsStream(RULES_FILE)); } @Override - protected TreeContext generate(Reader r) throws java.io.IOException { - TSParser parser = new TSParser(); - TSLanguage language = getTreeSitterLanguage(); - parser.setLanguage(language); - + protected TreeContext generate(Reader r) throws IOException { StringBuilder sb = new StringBuilder(); char[] buf = new char[8192]; int n; @@ -64,71 +66,69 @@ protected TreeContext generate(Reader r) throws java.io.IOException { if (content.isEmpty()) return emptyContext(); - List contentLines = Arrays.asList(content.split(LF, -1)); - TSTree tree = parser.parseString(null, content); - Map currentRule = RULES.getOrDefault(getLanguageName(), new HashMap<>()); - return generateFromTreeSitterTree(contentLines, currentRule, tree); - } - - private static String getLabel(List contentLines, TSNode node) { - int startRow = node.getStartPoint().getRow(); - int startColumn = node.getStartPoint().getColumn(); - int endRow = node.getEndPoint().getRow(); - int endColumn = node.getEndPoint().getColumn(); - List substringLines; - // tree-sitter handles string by byte array, so we need this. - String startRowStr = contentLines.get(startRow); - byte[] startRowBytes = startRowStr.getBytes(StandardCharsets.UTF_8); - if (startRow == endRow) { - // endColumn == startRowBytes.length + 1 when the label in tree-sitter contains line separator - if (endColumn == startRowBytes.length + 1) { - substringLines = Collections.singletonList(startRowStr); - } - else { - substringLines = Collections.singletonList(new String( - startRowBytes, startColumn, endColumn - startColumn, StandardCharsets.UTF_8)); + byte[] bytes = content.getBytes(StandardCharsets.UTF_8); + // one char per byte means that the byte offsets are already char offsets + int[] charOffsets = bytes.length == content.length() ? null : charOffsets(content, bytes.length); + LanguageRules rules = LANGUAGE_RULES.computeIfAbsent(getLanguageName(), + name -> new LanguageRules(RULES.getOrDefault(name, Collections.emptyMap()))); + try (TSParser parser = new TSParser()) { + parser.setLanguage(getTreeSitterLanguage()); + try (TSTree tree = parse(parser, bytes)) { + return generateFromTreeSitterTree(content, charOffsets, rules, tree); } } - else { - substringLines = new ArrayList<>(); - String endRowStr = contentLines.get(endRow); - byte[] endRowBytes = endRowStr.getBytes(StandardCharsets.UTF_8); - String startLineSubstring; - if (startColumn > startRowBytes.length) { - // usually, line separator is not the start char of a tree node label - // if this situation happened, just put an empty string at start - startLineSubstring = ""; - } else { - startLineSubstring = new String(startRowBytes, 0, startColumn, StandardCharsets.UTF_8); - } - List middleLines = contentLines.subList(startRow + 1, endRow); - String endLineSubstring; - if (endColumn > endRowStr.length()) { - endLineSubstring = endRowStr; - } - else { - endLineSubstring = new String(endRowBytes, 0, endColumn, StandardCharsets.UTF_8); - } - substringLines.add(startLineSubstring); - substringLines.addAll(middleLines); - substringLines.add(endLineSubstring); - } - return String.join(LF, substringLines); } - private static int calculateOffset(List contentLines, TSPoint point) { - int startRow = point.getRow(); - int startColumn = point.getColumn(); - int offset = 0; - for (int i = 0; i < startRow; i++) { - // Each line in contentLines (except maybe the last) was terminated by LF (\n). - // If the original was CRLF, the CR (\r) is still at the end of the line string. - // offset must be a char (UTF-16 code unit) index, matching how consumers (e.g. - // VanillaDiffHtmlBuilder) walk the source, so use .length() and not byte length. - offset += contentLines.get(i).length() + 1; + /** + * Parses the UTF-8 bytes of the content. TSParser.parseString drops the U+FEFF characters of its input, + * which shifts the byte offsets of all the nodes after them. + */ + private static TSTree parse(TSParser parser, byte[] bytes) throws IOException { + byte[] readBuffer = new byte[bytes.length]; + TSTree tree = parser.parse(readBuffer, null, (buffer, offset, position) -> { + int length = Math.min(buffer.length, bytes.length - offset); + if (length <= 0) + return 0; + System.arraycopy(bytes, offset, buffer, 0, length); + return length; + }, TSInputEncoding.TSInputEncodingUTF8); + if (tree == null) + throw new IOException("tree-sitter could not parse the content"); + return tree; + } + + /** + * Computes the char (UTF-16) offset of each UTF-8 byte offset of the content, as encoded by + * String.getBytes, which replaces an unpaired surrogate with a single byte. + */ + private static int[] charOffsets(String content, int byteLength) { + int[] offsets = new int[byteLength + 1]; + int byteOffset = 0; + for (int i = 0; i < content.length(); i++) { + char c = content.charAt(i); + int length; + if (c < 0x80) + length = 1; + else if (c < 0x800) + length = 2; + else if (Character.isHighSurrogate(c) && i + 1 < content.length() + && Character.isLowSurrogate(content.charAt(i + 1))) + length = 4; + else if (Character.isSurrogate(c)) + length = 1; + else + length = 3; + Arrays.fill(offsets, byteOffset, byteOffset + length, i); + byteOffset += length; + if (length == 4) + i++; } - offset += startColumn; - return offset; + offsets[byteOffset] = content.length(); + return offsets; + } + + private static int charOffset(int[] charOffsets, int byteOffset) { + return charOffsets == null ? byteOffset : charOffsets[byteOffset]; } /** @@ -142,7 +142,6 @@ protected static String matchNodeOrAncestorTypes(Collection ruleSet, TSN int depth = 1; String nodeType = node.getType(); StringBuilder typesBuilder = new StringBuilder(nodeType); - TSNode ancestor = node.getParent(); int maxDepthInRuleSet = 1; for (String rule : ruleSet) { String[] ruleTypes = rule.split(" "); @@ -154,6 +153,8 @@ protected static String matchNodeOrAncestorTypes(Collection ruleSet, TSN return nodeType; } } + // the parent is computed from the root by tree-sitter, so it is only computed when needed + TSNode ancestor = node.getParent(); while (!ancestor.isNull() && depth <= maxDepthInRuleSet) { typesBuilder.insert(0, ' '); typesBuilder.insert(0, ancestor.getType()); @@ -167,84 +168,64 @@ protected static String matchNodeOrAncestorTypes(Collection ruleSet, TSN return null; } - @SuppressWarnings("unchecked") - protected static Pair tsNode2GumTree( - List contentLines, Map currentRule, TreeContext context, TSNode node) { - String type = node.getType(); - if (currentRule.containsKey(YAML_IGNORED)) { - List ignores = (List) currentRule.get(YAML_IGNORED); - if (matchNodeOrAncestorTypes(ignores, node) != null) { - return null; - } - } - boolean ignoreLabel = false; - if (currentRule.containsKey(YAML_LABEL_IGNORED)) { - List ignores = (List) currentRule.get(YAML_LABEL_IGNORED); - if (matchNodeOrAncestorTypes(ignores, node) != null) { - ignoreLabel = true; - } - } - if (currentRule.containsKey(YAML_ALIASED)) { - Map alias = (Map) currentRule.get(YAML_ALIASED); - String matchedTypes = matchNodeOrAncestorTypes(alias.keySet(), node); - if (matchedTypes != null) { - type = alias.get(matchedTypes); - } - } - boolean flatten = false; - if (currentRule.containsKey(YAML_FLATTENED)) { - List flattenList = (List) currentRule.get(YAML_FLATTENED); - if (matchNodeOrAncestorTypes(flattenList, node) != null) { - flatten = true; - } - } - Tree tree; - // attach label for non ignore-label leafs or flattened nodes - if ((node.getChildCount() == 0 && !ignoreLabel) || flatten) { - String label = getLabel(contentLines, node); - tree = context.createTree(TypeSet.type(type), label); - } - else { - tree = context.createTree(TypeSet.type(type)); - } - tree.setPos(calculateOffset(contentLines, node.getStartPoint())); - int endOffset = calculateOffset(contentLines, node.getEndPoint()); - tree.setLength(endOffset - tree.getPos()); - return new Pair<>(tree, flatten); - } - + /** + * Traverses the tree-sitter tree with a tree cursor, keeping the types of the ancestors of the current node + * for the rules, instead of asking tree-sitter for the parent of each node. + */ private static TreeContext generateFromTreeSitterTree( - List contentLines, Map currentRule, TSTree tree) { - TSNode rootNode = tree.getRootNode(); + String content, int[] charOffsets, LanguageRules rules, TSTree tree) { TreeContext context = new TreeContext(); - Pair rootPair = tsNode2GumTree(contentLines, currentRule, context, rootNode); - if (rootPair == null) { - return context; - } - context.setRoot(rootPair.first); - Stack tsNodeStack = new Stack<>(); - Stack> treeStack = new Stack<>(); - tsNodeStack.push(rootNode); - treeStack.push(rootPair); - while (!tsNodeStack.isEmpty()) { - TSNode tsNodeNow = tsNodeStack.pop(); - Pair treeNow = treeStack.pop(); - // if node was flattened, ignore children - if (treeNow.second) { - continue; - } - int childCount = tsNodeNow.getChildCount(); - for (int i = 0; i < childCount; i++) { - TSNode child = tsNodeNow.getChild(i); - Pair childTree = tsNode2GumTree(contentLines, currentRule, context, child); - if (childTree != null) { - treeNow.first.addChild(childTree.first); - tsNodeStack.push(child); - treeStack.push(childTree); + List ancestorTypes = new ArrayList<>(); + Deque ancestors = new ArrayDeque<>(); + try (TSTreeCursor cursor = new TSTreeCursor(tree.getRootNode())) { + while (true) { + TSNode node = cursor.currentNode(); + String nodeType = node.getType(); + boolean visitChildren = false; + Tree gumTree = null; + if (rules.ignored.match(nodeType, ancestorTypes) == null) { + boolean ignoreLabel = rules.labelIgnored.match(nodeType, ancestorTypes) != null; + String type = nodeType; + String aliasedTypes = rules.aliased.match(nodeType, ancestorTypes); + if (aliasedTypes != null) + type = rules.aliases.get(aliasedTypes); + boolean flatten = rules.flattened.match(nodeType, ancestorTypes) != null; + int start = charOffset(charOffsets, node.getStartByte()); + int end = charOffset(charOffsets, node.getEndByte()); + int childCount = node.getChildCount(); + // attach label for non ignore-label leafs or flattened nodes + if ((childCount == 0 && !ignoreLabel) || flatten) + gumTree = context.createTree(TypeSet.type(type), content.substring(start, end)); + else + gumTree = context.createTree(TypeSet.type(type)); + gumTree.setPos(start); + gumTree.setLength(end - start); + if (ancestors.isEmpty()) + context.setRoot(gumTree); + else + ancestors.peek().addChild(gumTree); + // if node was flattened, ignore children + visitChildren = !flatten && childCount > 0; + } + else if (ancestors.isEmpty()) { + // the root is ignored + return context; + } + if (visitChildren && cursor.gotoFirstChild()) { + ancestors.push(gumTree); + ancestorTypes.add(nodeType); + continue; + } + while (!cursor.gotoNextSibling()) { + if (ancestors.isEmpty() || !cursor.gotoParent()) + return context; + ancestors.pop(); + ancestorTypes.remove(ancestorTypes.size() - 1); + if (ancestors.isEmpty()) + return context; } } } - return context; } protected abstract TSLanguage getTreeSitterLanguage(); @@ -259,4 +240,62 @@ private TreeContext emptyContext() { ctx.setRoot(root); return ctx; } + + /** + * The rules of a language in rules.yml. + */ + private static class LanguageRules { + private final RuleSet ignored; + private final RuleSet labelIgnored; + private final RuleSet flattened; + private final RuleSet aliased; + private final Map aliases; + + @SuppressWarnings("unchecked") + private LanguageRules(Map rule) { + ignored = new RuleSet((Collection) rule.get(YAML_IGNORED)); + labelIgnored = new RuleSet((Collection) rule.get(YAML_LABEL_IGNORED)); + flattened = new RuleSet((Collection) rule.get(YAML_FLATTENED)); + aliases = (Map) rule.getOrDefault(YAML_ALIASED, Collections.emptyMap()); + aliased = new RuleSet(aliases.keySet()); + } + } + + /** + * A set of rules, each one a node type or the space-separated types of ancestors ending with the node type. + */ + private static class RuleSet { + private final Set rules; + private final Set nodeTypeRules = new HashSet<>(); + private int maxDepth = 1; + + private RuleSet(Collection rules) { + this.rules = rules == null ? Collections.emptySet() : new HashSet<>(rules); + for (String rule : this.rules) { + int depth = rule.split(" ").length; + maxDepth = Math.max(maxDepth, depth); + if (depth == 1) + nodeTypeRules.add(rule); + } + } + + /** + * Same as matchNodeOrAncestorTypes, where the last ancestor type is the type of the parent. + */ + private String match(String nodeType, List ancestorTypes) { + if (rules.isEmpty()) + return null; + if (nodeTypeRules.contains(nodeType)) + return nodeType; + StringBuilder typesBuilder = new StringBuilder(nodeType); + for (int depth = 1, i = ancestorTypes.size() - 1; i >= 0 && depth <= maxDepth; depth++, i--) { + typesBuilder.insert(0, ' '); + typesBuilder.insert(0, ancestorTypes.get(i)); + String typesString = typesBuilder.toString(); + if (rules.contains(typesString)) + return typesString; + } + return null; + } + } } diff --git a/gen.treesitter-ng/src/test/java/com/github/gumtreediff/gen/treesitterng/AbstractTreeSitterNgGeneratorTest.java b/gen.treesitter-ng/src/test/java/com/github/gumtreediff/gen/treesitterng/AbstractTreeSitterNgGeneratorTest.java index 65ea5fa63..b8b262ce0 100644 --- a/gen.treesitter-ng/src/test/java/com/github/gumtreediff/gen/treesitterng/AbstractTreeSitterNgGeneratorTest.java +++ b/gen.treesitter-ng/src/test/java/com/github/gumtreediff/gen/treesitterng/AbstractTreeSitterNgGeneratorTest.java @@ -74,6 +74,85 @@ public void testMultiByteOffsetConsistency() throws IOException { "Line 2 should start at char offset 5 after a surrogate-pair emoji and LF"); } + @Test + public void testMultiByteOffsetOnSameLine() throws IOException { + // 'é' is 2 UTF-8 bytes but 1 char, so y is at char offset 9 (byte offset 10) + String content = "s = \"é\"; y = 1"; + TreeContext ctx = generator.generateFrom().string(content); + assertSourceRanges(content, ctx.getRoot()); + Tree y = findFirst(ctx.getRoot(), "identifier", "y"); + assertNotNull(y); + assertEquals(9, y.getPos()); + assertEquals(content.length(), ctx.getRoot().getEndPos()); + } + + @Test + public void testSurrogatePairOffsetOnSameLine() throws IOException { + // '🐍' is 4 UTF-8 bytes but 2 chars + String content = "s = \"🐍\"; y = 1"; + TreeContext ctx = generator.generateFrom().string(content); + assertSourceRanges(content, ctx.getRoot()); + assertEquals(10, findFirst(ctx.getRoot(), "identifier", "y").getPos()); + } + + @Test + public void testMultiLineLabel() throws IOException { + String content = "def f():\n x = \"\"\"first\n second\"\"\"\n"; + TreeContext ctx = generator.generateFrom().string(content); + assertSourceRanges(content, ctx.getRoot()); + Tree string = findFirst(ctx.getRoot(), "string", null); + assertNotNull(string); + assertEquals("\"\"\"first\n second\"\"\"", string.getLabel()); + } + + @Test + public void testMultiLineLabelCrlf() throws IOException { + String content = "def f():\r\n x = \"\"\"é\r\n second\"\"\"\r\n"; + TreeContext ctx = generator.generateFrom().string(content); + assertSourceRanges(content, ctx.getRoot()); + assertEquals("\"\"\"é\r\n second\"\"\"", findFirst(ctx.getRoot(), "string", null).getLabel()); + } + + @Test + public void testByteOrderMark() throws IOException { + // tree-sitter-ng TSParser.parseString drops U+FEFF, which shifted all the following nodes + String content = "x = 1\ny = 2"; + TreeContext ctx = generator.generateFrom().string(content); + assertSourceRanges(content, ctx.getRoot()); + assertEquals(1, findFirst(ctx.getRoot(), "identifier", "x").getPos()); + assertEquals(7, findFirst(ctx.getRoot(), "identifier", "y").getPos()); + } + + @Test + public void testFlattenedAndAncestorRules() throws IOException { + // rules.yml flattens python strings, and ignores parentheses, commas and "default_parameter =" + String content = "def g(a=1):\n print(f\"a {b} c\", 'é')\n"; + TreeContext ctx = generator.generateFrom().string(content); + assertSourceRanges(content, ctx.getRoot()); + Tree string = findFirst(ctx.getRoot(), "string", null); + assertEquals("f\"a {b} c\"", string.getLabel()); + assertTrue(string.isLeaf()); + assertNull(findFirst(ctx.getRoot(), "(", null)); + assertNull(findFirst(ctx.getRoot(), ",", null)); + Tree defaultParameter = findFirst(ctx.getRoot(), "default_parameter", null); + assertEquals(2, defaultParameter.getChildren().size()); + } + + private static void assertSourceRanges(String content, Tree root) { + for (Tree t : root.preOrder()) { + assertTrue(t.getPos() >= 0 && t.getEndPos() <= content.length(), t.toString()); + if (t.isLeaf() && !t.getLabel().isEmpty()) + assertEquals(content.substring(t.getPos(), t.getEndPos()), t.getLabel(), t.toString()); + } + } + + private static Tree findFirst(Tree root, String type, String label) { + for (Tree t : root.preOrder()) + if (t.getType().name.equals(type) && (label == null || label.equals(t.getLabel()))) + return t; + return null; + } + @Test public void testMatchNodeOrAncestorTypes() { MockTypeOnlyTreeSitterNode root = new MockTypeOnlyTreeSitterNode();