// Deterministic Java AST facts for todo2code. Run in JDK source-file mode: // java --add-modules jdk.compiler java/JavaAstExtract.java --max-file-bytes 524288 import com.sun.source.tree.*; import com.sun.source.util.JavacTask; import com.sun.source.util.SourcePositions; import com.sun.source.util.TreePathScanner; import com.sun.source.util.Trees; import java.io.IOException; import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.Paths; import java.util.*; import javax.tools.*; public final class JavaAstExtract { private static final Set IGNORED = Set.of( ".git", ".intent", "node_modules", "dist", "build", "target", "coverage", "vendor"); public static void main(String[] args) throws Exception { Path root = Paths.get(".").toAbsolutePath().normalize(); long maxBytes = 524_288; for (int index = 0; index < args.length; index++) { if ("--max-file-bytes".equals(args[index]) && index + 1 < args.length) { maxBytes = Long.parseLong(args[++index]); } else if (!args[index].startsWith("-")) { root = Paths.get(args[index]).toAbsolutePath().normalize(); } } List> facts = new ArrayList<>(); List warnings = new ArrayList<>(); JavaCompiler compiler = ToolProvider.getSystemJavaCompiler(); if (compiler == null) { warnings.add("JDK compiler is unavailable; a full JDK is required for Java AST extraction"); emit(facts, warnings); return; } List files = collect(root); for (Path file : files) { try { long size = Files.size(file); String relative = slash(root.relativize(file).toString()); if (size > maxBytes) { warnings.add(relative + ": skipped, " + size + " bytes exceeds limit " + maxBytes); continue; } parseFile(compiler, root, file, facts, warnings); } catch (Exception error) { warnings.add(slash(root.relativize(file).toString()) + ": " + error.getMessage()); } } emit(facts, warnings); } private static List collect(Path root) throws IOException { List output = new ArrayList<>(); try (var paths = Files.walk(root)) { paths.filter(Files::isRegularFile) .filter(path -> path.getFileName().toString().endsWith(".java")) .filter(path -> !containsIgnored(root.relativize(path))) .forEach(output::add); } output.sort(Comparator.comparing(Path::toString)); return output; } private static boolean containsIgnored(Path relative) { for (Path part : relative) if (IGNORED.contains(part.toString())) return true; return false; } private static void parseFile( JavaCompiler compiler, Path root, Path file, List> facts, List warnings) throws Exception { String source = Files.readString(file, StandardCharsets.UTF_8); DiagnosticCollector diagnostics = new DiagnosticCollector<>(); try (StandardJavaFileManager manager = compiler.getStandardFileManager(diagnostics, Locale.ROOT, StandardCharsets.UTF_8)) { Iterable extends JavaFileObject> units = manager.getJavaFileObjects(file.toFile()); JavacTask task = (JavacTask) compiler.getTask(null, manager, diagnostics, List.of("-proc:none"), null, units); Iterable extends CompilationUnitTree> parsed = task.parse(); Trees trees = Trees.instance(task); SourcePositions positions = trees.getSourcePositions(); String relative = slash(root.relativize(file).toString()); for (CompilationUnitTree unit : parsed) { new Collector(relative, source, unit, positions, facts).scan(unit, null); } } for (Diagnostic extends JavaFileObject> diagnostic : diagnostics.getDiagnostics()) { if (diagnostic.getKind() == Diagnostic.Kind.ERROR) { warnings.add(slash(root.relativize(file).toString()) + ":" + diagnostic.getLineNumber() + ": parse error: " + diagnostic.getMessage(Locale.ROOT)); } } } private static final class Collector extends TreePathScanner { private final String relative; private final String source; private final CompilationUnitTree unit; private final SourcePositions positions; private final List> facts; private final Deque types = new ArrayDeque<>(); private final Deque methods = new ArrayDeque<>(); Collector(String relative, String source, CompilationUnitTree unit, SourcePositions positions, List> facts) { this.relative = relative; this.source = source; this.unit = unit; this.positions = positions; this.facts = facts; } @Override public Void visitCompilationUnit(CompilationUnitTree node, Void unused) { if (node.getPackageName() != null) { String name = node.getPackageName().toString(); add(node.getPackageName(), "java_package_fact", "declare", name, name, null, map("kind", "package")); } return super.visitCompilationUnit(node, unused); } @Override public Void visitImport(ImportTree node, Void unused) { add(node, "java_import_fact", "depend_on", node.getQualifiedIdentifier().toString(), null, null, map("static", node.isStatic())); return null; } @Override public Void visitClass(ClassTree node, Void unused) { String name = node.getSimpleName().toString(); if (name.isEmpty()) return super.visitClass(node, unused); String qualified = types.isEmpty() ? name : types.peekLast() + "." + name; add(node, "java_type_fact", "declare", qualified, qualified, types.peekLast(), map( "kind", node.getKind().name().toLowerCase(Locale.ROOT), "modifiers", node.getModifiers().getFlags().stream().map(Object::toString).sorted().toList())); types.addLast(qualified); super.visitClass(node, unused); types.removeLast(); return null; } @Override public Void visitMethod(MethodTree node, Void unused) { String raw = node.getName().toString(); String name = "".equals(raw) && !types.isEmpty() ? types.peekLast().substring(types.peekLast().lastIndexOf('.') + 1) : raw; String symbol = types.isEmpty() ? name : types.peekLast() + "." + name; add(node, "java_method_fact", "declare", symbol, symbol, types.peekLast(), map( "kind", "method", "parameterCount", node.getParameters().size(), "modifiers", node.getModifiers().getFlags().stream().map(Object::toString).sorted().toList())); methods.addLast(symbol); super.visitMethod(node, unused); methods.removeLast(); return null; } @Override public Void visitVariable(VariableTree node, Void unused) { if (methods.isEmpty() && !types.isEmpty()) { String symbol = types.peekLast() + "." + node.getName(); add(node, "java_field_fact", "declare", symbol, symbol, types.peekLast(), map("kind", "field")); } return super.visitVariable(node, unused); } @Override public Void visitMethodInvocation(MethodInvocationTree node, Void unused) { String callee = node.getMethodSelect().toString(); add(node, "java_call_fact", "call", callee, methods.peekLast(), methods.peekLast(), map("callee", callee, "argumentCount", node.getArguments().size())); return super.visitMethodInvocation(node, unused); } @Override public Void visitNewClass(NewClassTree node, Void unused) { String callee = "new " + node.getIdentifier(); add(node, "java_call_fact", "call", callee, methods.peekLast(), methods.peekLast(), map("callee", callee, "argumentCount", node.getArguments().size(), "constructor", true)); return super.visitNewClass(node, unused); } private void add(Tree node, String kind, String action, String object, String symbol, String subject, Map metadata) { long rawStart = positions.getStartPosition(unit, node); long rawEnd = positions.getEndPosition(unit, node); int start = rawStart < 0 ? 0 : (int) Math.min(rawStart, source.length()); int end = rawEnd < start ? start : (int) Math.min(rawEnd, source.length()); int lineStart = rawStart < 0 ? 1 : (int) unit.getLineMap().getLineNumber(rawStart); int lineEnd = rawEnd < 0 ? lineStart : (int) unit.getLineMap().getLineNumber(Math.max(rawStart, rawEnd - 1)); String excerpt = source.substring(start, Math.min(end, start + 2000)); Map fact = new LinkedHashMap<>(); fact.put("path", relative); fact.put("lineStart", Math.max(1, lineStart)); fact.put("lineEnd", Math.max(lineStart, lineEnd)); fact.put("kind", kind); fact.put("action", action); fact.put("object", object); fact.put("symbol", symbol); fact.put("subject", subject); fact.put("excerpt", excerpt); fact.put("contentHash", ""); fact.put("metadata", metadata); facts.add(fact); } } private static Map map(Object... values) { Map output = new LinkedHashMap<>(); for (int index = 0; index + 1 < values.length; index += 2) { output.put(String.valueOf(values[index]), values[index + 1]); } return output; } private static void emit(List> facts, List warnings) { System.out.println(json(map("facts", facts, "warnings", warnings))); } private static String json(Object value) { if (value == null) return "null"; if (value instanceof Boolean || value instanceof Number) return value.toString(); if (value instanceof String text) return "\"" + escape(text) + "\""; if (value instanceof Map, ?> map) { List fields = new ArrayList<>(); for (var entry : map.entrySet()) fields.add(json(String.valueOf(entry.getKey())) + ":" + json(entry.getValue())); return "{" + String.join(",", fields) + "}"; } if (value instanceof Iterable> items) { List encoded = new ArrayList<>(); for (Object item : items) encoded.add(json(item)); return "[" + String.join(",", encoded) + "]"; } return json(value.toString()); } private static String escape(String value) { StringBuilder output = new StringBuilder(); for (int index = 0; index < value.length(); index++) { char character = value.charAt(index); switch (character) { case '\\' -> output.append("\\\\"); case '"' -> output.append("\\\""); case '\n' -> output.append("\\n"); case '\r' -> output.append("\\r"); case '\t' -> output.append("\\t"); default -> { if (character < 0x20) output.append(String.format("\\u%04x", (int) character)); else output.append(character); } } } return output.toString(); } private static String slash(String value) { return value.replace('\\', '/'); } }