frontend: Further improve frontend performance (focused on allocations)

java-annotations/main/vadl/javaannotations/ast/ChildAnnotationProcessor.java

@@ -176,18 +176,16 @@ private void generateNodeChildrenRegistry() throws IOException {
       out.println("// Generated code from %s".formatted(this.getClass().getName()));
       out.println("package " + packageName + ";");
       out.println();
-      out.println("import java.util.List;");
-      out.println("import java.util.ArrayList;");
-      out.println("import java.util.Collections;");
       out.println("import java.util.Map;");
       out.println("import java.util.HashMap;");
-      out.println("import java.util.function.Function;");
+      out.println("import java.util.function.BiConsumer;");
+      out.println("import java.util.function.Consumer;");
       out.println();
 
       // Generate registry class
       out.println("public final class NodeChildrenRegistry {");
       out.println(
-          "    private static final Map<Class<?>, Function<Node, List<Node>>> COLLECTORS = "
+          "    private static final Map<Class<?>, BiConsumer<Node, Consumer<Node>>> COLLECTORS = "
               + "new HashMap<>();");
       out.println();
 
@@ -202,91 +200,63 @@ private void generateNodeChildrenRegistry() throws IOException {
         String className = classElement.getQualifiedName().toString();
         List<VariableElement> childFields = entry.getValue();
 
-        // Optimization for when all children are already in a single list
-        if (childFields.size() == 1 && isFieldList(childFields.getFirst())) {
-          var field = childFields.getFirst();
-          out.println(
-              "        COLLECTORS.put(" + className + ".class, (node) -> (List<Node>) (List<?>) (("
-                  + className + ") node)." + fieldAccessor(field) + ");");
-          continue;
-        }
-
-        out.println("        COLLECTORS.put(" + className + ".class, (node) -> {");
+        out.println("        COLLECTORS.put(" + className + ".class, (node, action) -> {");
         out.println("            " + className + " n = (" + className + ") node;");
 
-        // Calculate the capacity to optimize allocations
-        out.println("            int capacity = 0");
-        for (VariableElement field : childFields) {
-          String fieldName = field.getSimpleName().toString();
-          if (isFieldList(field)) {
-            out.println(
-                "                  + n.%s.size() // %s".formatted(fieldAccessor(field), fieldName));
-          } else {
-            out.println("                  + 1 // %s".formatted(fieldName));
-          }
-        }
-        out.println("                 ;");
-
-
-        out.println("            List<Node> children = new ArrayList<>(capacity);");
-
-        // Add code to collect children for this node type
         for (VariableElement field : childFields) {
-          // Handle different field types
           if (isFieldList(field)) {
-            out.println("            if ( n." + fieldAccessor(field) + " != null) {");
-            out.println("               for (var child: n." + fieldAccessor(field) + ") {");
-            out.println("                   children.add((Node) child);");
-            out.println("               }");
+            out.println("            if (n." + fieldAccessor(field) + " != null) {");
+            out.println("                for (var child : n." + fieldAccessor(field) + ") {");
+            out.println("                    action.accept((Node) child);");
+            out.println("                }");
             out.println("            }");
           } else {
-            out.println("            if ( n." + fieldAccessor(field) + " != null) {");
-            out.println("                children.add((Node) n." + fieldAccessor(field) + ");");
+            out.println("            if (n." + fieldAccessor(field) + " != null) {");
+            out.println("                action.accept((Node) n." + fieldAccessor(field) + ");");
             out.println("            }");
           }
         }
 
-        out.println("            return children;");
         out.println("        });");
       }
 
       out.println("    }");
       out.println();
 
-      // Method to get children for any node
-      out.println("    public static List<Node> getChildren(Node node) {");
+      // Method to iterate children for any node
+      out.println(
+          "    public static void forEachChild(Node node, Consumer<Node> action) {");
       out.println(
-          "        Function<Node, List<Node>> collector = COLLECTORS.get(node.getClass());");
+          "        BiConsumer<Node, Consumer<Node>> collector = COLLECTORS.get(node.getClass());");
       out.println("        if (collector != null) {");
-      out.println("            return collector.apply(node);");
+      out.println("            collector.accept(node, action);");
       out.println("        }");
-      out.println("        return Collections.emptyList();");
       out.println("    }");
       out.println("");
 
-
-      // Method to get all children but specify the exact class, only used for edgecases
+      // Method to iterate children but specify the exact class, only used for edge cases
       out.println("    /**");
       out.println("     * Specify the class directly as which it should be loaded.");
       out.println(
           "     * This should only be used when you know what you do, like if you want "
               + "to get the children");
       out.println("     * from your superclass.");
       out.println("     *");
-      out.println("     * @param node from which the children are loaded.");
+      out.println("     * @param node from which the children are iterated.");
       out.println("     * @param nodeType as which the node should be interpreted.");
-      out.println("     * @return the children.");
+      out.println("     * @param action called for each child.");
       out.println("     */");
       out.println(
-          "    public static List<Node> unsafeGetChildrenDirect(Node node, "
-              + "Class<? extends Node> nodeType) {");
-      out.println("        Function<Node, List<Node>> collector = COLLECTORS.get(nodeType);");
-      out.println("            if (collector == null) {");
+          "    public static void unsafeForEachChildDirect(Node node, "
+              + "Class<? extends Node> nodeType, Consumer<Node> action) {");
+      out.println(
+          "        BiConsumer<Node, Consumer<Node>> collector = COLLECTORS.get(nodeType);");
+      out.println("        if (collector == null) {");
       out.println(
-          "                throw new IllegalArgumentException(\"Node type \" + nodeType + \" "
+          "            throw new IllegalArgumentException(\"Node type \" + nodeType + \" "
               + "not supported\");");
-      out.println("            }");
-      out.println("        return collector.apply(node);");
+      out.println("        }");
+      out.println("        collector.accept(node, action);");
       out.println("    }");
       out.println("}");
     }

vadl-cli/main/vadl/cli/BaseCommand.java

@@ -266,20 +266,20 @@ private Specification parseToVIAM() {
     var ast = parseToAst();
     printSpecStats(ast, new DiskVirtualFileSystem());
     printSpecStatsCsv(ast, new DiskVirtualFileSystem());
-    ast.timingRecorder.passTimings.forEach(
+    ast.timingRecorder.passTimings.sequencedValues().forEach(
         t -> timings.add(new Timing(t.description(), t.durationNS())));
     ast.timingRecorder.passTimings.clear();
     dumpExpaned(ast);
     dumpUntyped(ast);
 
     TypeChecker.verify(ast);
-    ast.timingRecorder.passTimings.forEach(
+    ast.timingRecorder.passTimings.sequencedValues().forEach(
         t -> timings.add(new Timing(t.description(), t.durationNS())));
     ast.timingRecorder.passTimings.clear();
     dumpTyped(ast);
 
     var spec = ViamLowering.generate(ast);
-    ast.timingRecorder.passTimings.forEach(
+    ast.timingRecorder.passTimings.sequencedValues().forEach(
         t -> timings.add(new Timing(t.description(), t.durationNS())));
 
     return spec;

vadl/main/vadl/ast/AnnotationTable.java

@@ -35,6 +35,7 @@
 import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import javax.annotation.Nullable;
 import vadl.error.Diagnostic;
 import vadl.error.DiagnosticBuilder;
@@ -418,12 +419,12 @@ static List<String> availableAnnotationNames(Class<? extends Definition> klass)
    */
   static Map<AnnotationGroupProvider, List<Annotation>> groupings(
       Definition definition) {
-    return groupings(definition.annotations.stream().map(d -> d.annotation).toList());
+    return groupings(definition.annotations.stream().map(d -> d.annotation));
   }
 
   private static Map<AnnotationGroupProvider, List<Annotation>> groupings(
-      List<Annotation> annotations) {
-    return annotations.stream().collect(Collectors.groupingBy(a -> a.groupProvider));
+      Stream<Annotation> annotations) {
+    return annotations.collect(Collectors.groupingBy(a -> a.groupProvider));
   }
 
   /**

vadl/main/vadl/ast/AsmGrammarDefaultRules.java

@@ -169,12 +169,12 @@ public static FunctionDefinition asmNegFunctionDefinition(WithLocation locatable
             new Parameter(
                 new Identifier("x", loc),
                 new TypeLiteral(new Identifier("SInt", loc),
-                    List.of(new IntegerLiteral("64", loc)),
+                    List.of(new IntegerLiteral(64, loc)),
                     loc)
             )
         )),
         new TypeLiteral(new Identifier("SInt", loc),
-            List.of(new IntegerLiteral("64", loc)),
+            List.of(new IntegerLiteral(64, loc)),
             loc),
         new UnaryExpr(new UnOp(UnaryOperator.NEGATIVE, loc),
             new Identifier("x", loc)),

vadl/main/vadl/ast/Ast.java

@@ -22,6 +22,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -210,6 +211,16 @@ static boolean isBlockLayout(Node n) {
 
   abstract void prettyPrint(int indent, StringBuilder builder);
 
+  /**
+   * Invokes {@code action} for each child owned by this node.
+   *
+   * <p>Either automatically generated by annotating children with the
+   * {@link vadl.javaannotations.ast.Child} annotation or by overriding this method.
+   */
+  void forEachChild(Consumer<Node> action) {
+    NodeChildrenRegistry.forEachChild(this, action);
+  }
+
   /**
    * Returns all children "owned" by this node.
    *
@@ -218,8 +229,10 @@ static boolean isBlockLayout(Node n) {
    *
    * @return a list of children.
    */
-  List<Node> children() {
-    return NodeChildrenRegistry.getChildren(this);
+  final List<Node> children() {
+    var result = new ArrayList<Node>();
+    forEachChild(result::add);
+    return result;
   }
 
   @Override

vadl/main/vadl/ast/AstDumper.java

@@ -42,6 +42,24 @@ public String dump(Ast ast) {
     return builder.toString();
   }
 
+  static <T extends Definition> String debugNode(T node) {
+    var dumper = new AstDumper();
+    node.accept(dumper);
+    return dumper.builder.toString();
+  }
+
+  static <T extends Expr> String debugNode(T node) {
+    var dumper = new AstDumper();
+    node.accept(dumper);
+    return dumper.builder.toString();
+  }
+
+  static <T extends Statement> String debuNode(T node) {
+    var dumper = new AstDumper();
+    node.accept(dumper);
+    return dumper.builder.toString();
+  }
+
   private String indentString() {
     var indentBy = 2;
     var indentCharacters = ". : ' | ";

vadl/main/vadl/ast/AstUtils.java

@@ -18,6 +18,8 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import vadl.types.BuiltInTable;
@@ -77,9 +79,8 @@ static BuiltInTable.BuiltIn getOperatorBuiltIn(Operator operator, List<Type> arg
     }
 
     String finalOperatorSymbol = symbol;
-    var builtIns = operatorLookupTable.getOrDefault(finalOperatorSymbol, List.of()).stream()
-        .filter(b -> b.signature().argTypeClasses().size() == argTypes.size())
-        .toList();
+    var builtIns = operatorLookupTable.getOrDefault(finalOperatorSymbol, List.of());
+    builtIns.removeIf(b -> b.signature().argTypeClasses().size() != argTypes.size());
 
     // Sometimes there are a signed and unsigned version of builtin operation
     return switch (builtIns.size()) {
@@ -124,8 +125,40 @@ static BuiltInTable.BuiltIn getOperatorBuiltIn(Operator operator, List<Type> arg
     };
   }
 
-  static List<Expr> flatArguments(List<CallIndexExpr.Arguments> args) {
-    return args.stream().flatMap(a -> a.values.stream()).collect(Collectors.toList());
+  static List<Expr> flatArguments(List<CallIndexExpr.Arguments> argGroups) {
+    return argGroups.stream().flatMap(a -> a.values.stream()).collect(Collectors.toList());
   }
 
+  static void forEachArgument(List<CallIndexExpr.Arguments> argGroups, Consumer<Expr> consumer) {
+    argGroups.forEach(a -> a.values.forEach(consumer));
+  }
+
+  static int argumentCount(List<CallIndexExpr.Arguments> argGroups) {
+    int cnt = 0;
+    for (var args : argGroups) {
+      cnt += args.values.size();
+    }
+    return cnt;
+  }
+
+  static boolean isFullyExpanded(Node node) {
+    if (node instanceof ModelDefinition
+        || node instanceof PlaceholderNode
+        || node instanceof PlaceholderDefinition
+        || node instanceof PlaceholderStatement
+        || node instanceof PlaceholderExpr
+        || node instanceof AsIdExpr
+        || node instanceof AsStrExpr) {
+      return false;
+    }
+
+    AtomicBoolean areChildrenExpanded = new AtomicBoolean(true);
+    node.forEachChild(child -> {
+      if (!isFullyExpanded(child)) {
+        areChildrenExpanded.set(false);
+      }
+    });
+    return areChildrenExpanded.get();
+
+  }
 }