Make 'units()' primitive lazy

Author: Killian Perlin

lkql_checker/share/lkql/kp/kp_20229.lkql

@@ -50,7 +50,7 @@ fun kp_20229(node) =
         if output then
             stdlib.any([
                 from r.ref through parent select first ReturnStmt
-                for r in output.p_find_all_references(units())
+                for r in output.p_find_all_references(units().to_list)
             ])
         else false
     }

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

@@ -17,6 +17,8 @@
 import com.adacore.lkql_jit.runtime.values.interfaces.Indexable;
 import com.adacore.lkql_jit.runtime.values.interfaces.Iterable;
 import com.adacore.lkql_jit.runtime.values.interfaces.Iterator;
+import com.adacore.lkql_jit.runtime.values.lists.BaseLKQLList;
+import com.adacore.lkql_jit.runtime.values.lists.LKQLLazyListStreamWrapper;
 import com.adacore.lkql_jit.runtime.values.lists.LKQLList;
 import com.adacore.lkql_jit.utils.Constants;
 import com.adacore.lkql_jit.utils.LKQLTypesHelper;
@@ -556,12 +558,8 @@ abstract static class UnitsExpr extends BuiltInBody {
 
         @Specialization
         @CompilerDirectives.TruffleBoundary
-        protected LKQLList alwaysTrue() {
-            return new LKQLList(
-                LKQLLanguage.getContext(this)
-                    .getAllUnits()
-                    .toArray(LangkitSupport.AnalysisUnit[]::new)
-            );
+        protected BaseLKQLList alwaysTrue() {
+            return new LKQLLazyListStreamWrapper(LKQLLanguage.getContext(this).getAllUnits());
         }
     }
 

lkql_jit/language/src/main/java/com/adacore/lkql_jit/runtime/ListStorage.java

@@ -24,6 +24,7 @@ public ListStorage(int startSize) {
 
     public void append(T item) {
         if (current_size >= capacity) {
+            capacity = Math.max(capacity, 1);
             capacity *= 2;
             var new_storage = new Object[capacity];
             for (int i = 0; i < current_size; i++) {

lkql_jit/language/src/main/java/com/adacore/lkql_jit/runtime/values/lists/LKQLLazyListStreamWrapper.java

@@ -0,0 +1,38 @@
+//
+//  Copyright (C) 2005-2025, AdaCore
+//  SPDX-License-Identifier: GPL-3.0-or-later
+//
+
+package com.adacore.lkql_jit.runtime.values.lists;
+
+import com.adacore.lkql_jit.runtime.ListStorage;
+import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
+import java.util.Iterator;
+import java.util.stream.Stream;
+
+/**
+ * This class is a wrapper from a Java Stream to an LKQLLazyList.
+ * It is currently only used to wrap the units() builtin.
+ */
+public class LKQLLazyListStreamWrapper extends BaseLKQLLazyList {
+
+    // ----- Attributes -----
+
+    private Iterator<? extends Object> iterator;
+
+    // ----- Constructors -----
+
+    public LKQLLazyListStreamWrapper(Stream<? extends Object> source) {
+        super(new ListStorage<>(0));
+        iterator = source.iterator();
+    }
+
+    // ----- Instance methods -----
+
+    @TruffleBoundary
+    protected void initCacheTo(long n) {
+        while ((n < 0 || cache.size() <= n) && iterator.hasNext()) {
+            cache.append(iterator.next());
+        }
+    }
+}

testsuite/tests/gnatcheck/scenario_variables/rules/display_units.lkql

@@ -1,3 +1,3 @@
 @unit_check
-fun display_units(unit) = [{message: "UNITS: " & units().img,
+fun display_units(unit) = [{message: "UNITS: " & units().to_list.img,
                             loc: unit.root}]

testsuite/tests/miscellaneous/units/script.lkql

@@ -1 +1 @@
-print(units())
+print(units().to_list)