Perf: lazily compute all/specified units in LKQLContext

Author: raph-amiard

lkql_jit/language/src/main/java/com/adacore/lkql_jit/LKQLContext.java

@@ -26,6 +26,7 @@
 import java.io.File;
 import java.util.*;
 import java.util.function.BiConsumer;
+import java.util.stream.Stream;
 import org.json.JSONObject;
 import org.json.JSONTokener;
 
@@ -94,21 +95,6 @@ public final class LKQLContext {
      */
     private List<String> allSourceFiles;
 
-    /** Whether the source files were parsed. */
-    private boolean parsed;
-
-    /**
-     * The user-specified units to analyze. If not explicitly specified, those will be the units of
-     * the root project.
-     */
-    private LangkitSupport.AnalysisUnit[] specifiedUnits;
-
-    /** All the units of the project, including those of its non-externally-built dependencies. */
-    private LangkitSupport.AnalysisUnit[] allUnits;
-
-    /** The root nodes of all the analysis units of the project. */
-    private LangkitSupport.NodeInterface[] allUnitsRoots;
-
     // ----- Checker attributes -----
 
     /** A cache for all rule arguments to avoid evaluating twice the same argument source. */
@@ -160,7 +146,6 @@ public LKQLContext(TruffleLanguage.Env env, GlobalScope global, LKQLLanguage lan
         this.global = global;
         this.specifiedSourceFiles = new ArrayList<>();
         this.allSourceFiles = new ArrayList<>();
-        this.parsed = false;
         this.language = language;
     }
 
@@ -186,36 +171,34 @@ public GlobalScope getGlobal() {
         return this.global;
     }
 
-    public LangkitSupport.AnalysisUnit[] getSpecifiedUnits() {
-        if (!this.parsed) {
-            this.parseSources();
-        }
-        return this.specifiedUnits;
+    public Stream<LangkitSupport.AnalysisUnit> getSpecifiedUnits() {
+        String[] ignores = this.getIgnores();
+        return this.specifiedSourceFiles.stream()
+            .filter(source -> Arrays.stream(ignores).noneMatch(source::contains))
+            .map(f -> analysisContext.getUnitFromFile(f));
     }
 
-    public LangkitSupport.AnalysisUnit[] getAllUnits() {
-        if (!this.parsed) {
-            this.parseSources();
-        }
-        return this.allUnits;
+    public Stream<LangkitSupport.AnalysisUnit> getAllUnits() {
+        String[] ignores = this.getIgnores();
+        return this.allSourceFiles.stream()
+            .map(f -> {
+                return analysisContext.getUnitFromFile(f);
+            });
     }
 
-    public LangkitSupport.NodeInterface[] getAllUnitsRoots() {
-        if (!this.parsed) {
-            this.parseSources();
-        }
-        return this.allUnitsRoots;
+    @CompilerDirectives.TruffleBoundary
+    public LangkitSupport.NodeInterface[] allUnitsRoots() {
+        return getAllUnits().map(u -> u.getRoot()).toArray(LangkitSupport.NodeInterface[]::new);
     }
 
     public boolean hasRewritingContext() {
         return this.rewritingContext != null;
     }
 
+    @CompilerDirectives.TruffleBoundary
     public LangkitSupport.RewritingContextInterface getRewritingContext() {
         if (this.rewritingContext == null) {
-            if (!this.parsed) {
-                this.parseSources();
-            }
+            var ignore = getAllUnits().toArray(LangkitSupport.AnalysisUnit[]::new);
             this.rewritingContext = this.analysisContext.startRewriting();
         }
         return this.rewritingContext;
@@ -420,52 +403,11 @@ public CheckerUtils.DiagnosticEmitter getDiagnosticEmitter() {
         return this.emitter;
     }
 
-    // ----- Project analysis methods -----
-
-    /** Parse the ada source files and store analysis units and root nodes. */
-    @CompilerDirectives.TruffleBoundary
-    public void parseSources() {
-        // Filter the Ada source file list
-        String[] ignores = this.getIgnores();
-        String[] usedSources =
-            this.specifiedSourceFiles.stream()
-                .filter(source -> {
-                    for (String ignore : ignores) {
-                        if (source.contains(ignore)) return false;
-                    }
-                    return true;
-                })
-                .toArray(String[]::new);
-
-        // For each specified source file, store its corresponding analysis unit in the list of
-        // specified units
-        this.specifiedUnits = new LangkitSupport.AnalysisUnit[usedSources.length];
-        for (int i = 0; i < usedSources.length; i++) {
-            this.specifiedUnits[i] = this.analysisContext.getUnitFromFile(usedSources[i]);
-        }
-
-        // For each source file of the project, store its corresponding analysis unit in the list of
-        // all
-        // the units
-        // of the project, as well as their root nodes.
-        this.allUnits = new LangkitSupport.AnalysisUnit[this.allSourceFiles.size()];
-        this.allUnitsRoots = new LangkitSupport.NodeInterface[this.allSourceFiles.size()];
-
-        for (int i = 0; i < this.allUnits.length; i++) {
-            this.allUnits[i] = this.analysisContext.getUnitFromFile(this.allSourceFiles.get(i));
-            this.allUnitsRoots[i] = this.allUnits[i].getRoot();
-        }
-
-        // All source files are now parsed
-        this.parsed = true;
-    }
-
     /** Initialize the ada sources. */
     public void initSources() {
         // Reset the context fields
         this.specifiedSourceFiles.clear();
         this.allSourceFiles.clear();
-        this.parsed = false;
 
         // Add all the user-specified files to process after verifying they exist
         for (String file : this.getFiles()) {

lkql_jit/language/src/main/java/com/adacore/lkql_jit/built_ins/BuiltInFunctions.java

@@ -5,6 +5,7 @@
 
 package com.adacore.lkql_jit.built_ins;
 
+import com.adacore.langkit_support.LangkitSupport;
 import com.adacore.lkql_jit.LKQLLanguage;
 import com.adacore.lkql_jit.LKQLTypeSystemGen;
 import com.adacore.lkql_jit.annotations.*;
@@ -537,17 +538,26 @@ protected Object onLKQLValue(LKQLValue value) {
     abstract static class UnitsExpr extends BuiltInBody {
 
         @Specialization
+        @CompilerDirectives.TruffleBoundary
         protected LKQLList alwaysTrue() {
-            return new LKQLList(LKQLLanguage.getContext(this).getAllUnits());
+            return new LKQLList(
+                LKQLLanguage.getContext(this)
+                    .getAllUnits()
+                    .toArray(LangkitSupport.AnalysisUnit[]::new)
+            );
         }
     }
 
     @BuiltInFunction(name = "specified_units", doc = "Return a list of units specified by the user")
     abstract static class SpecifiedUnitsExpr extends BuiltInBody {
 
         @Specialization
+        @CompilerDirectives.TruffleBoundary
         protected LKQLList alwaysTrue() {
-            return new LKQLList(LKQLLanguage.getContext(this).getSpecifiedUnits());
+            var units = LKQLLanguage.getContext(this)
+                .getSpecifiedUnits()
+                .toArray(LangkitSupport.AnalysisUnit[]::new);
+            return new LKQLList(units);
         }
     }
 }

lkql_jit/language/src/main/java/com/adacore/lkql_jit/nodes/expressions/Query.java

@@ -138,7 +138,7 @@ else if (LKQLTypeSystemGen.isLKQLList(fromObject)) {
         }
         // Else, there is no "from" expression, we get the default roots
         else {
-            fromNodes = LKQLLanguage.getContext(this.pattern).getAllUnitsRoots();
+            fromNodes = LKQLLanguage.getContext(this.pattern).allUnitsRoots();
         }
 
         // If the query mode is all