merge: JIT compiler should generate fewer things (#1278)

Fix 3 issues in a row.

jit-compiler/src/main/java/org/aya/compiler/free/morphism/free/FreeOptimizer.java

@@ -2,6 +2,7 @@
 // Use of this source code is governed by the MIT license that can be found in the LICENSE.md file.
 package org.aya.compiler.free.morphism.free;
 
+import kala.collection.SeqView;
 import kala.collection.immutable.ImmutableSeq;
 import org.aya.compiler.free.morphism.free.FreeDecl.Clazz;
 import org.aya.compiler.free.morphism.free.FreeDecl.ConstantField;
@@ -20,33 +21,50 @@ case Clazz(var metadata, var owner, var nested, var superclass, var members) ->
         yield new Clazz(metadata, owner, nested, superclass, newMembers);
       }
       case ConstantField field -> field;
-      case Method(var signature, var body) -> new Method(signature, body.flatMap(FreeOptimizer::optimize));
+      case Method(var signature, var body) -> new Method(signature, optimizeBlock(body, false));
     };
   }
 
-  static @NotNull ImmutableSeq<FreeStmt> optimize(FreeStmt stmt) {
+  static @NotNull ImmutableSeq<FreeStmt> optimizeBlock(ImmutableSeq<FreeStmt> block, boolean endOfBreakable) {
+    if (!endOfBreakable) return block.flatMap(it -> optimize(it, false));
+    if (block.isEmpty()) return block;
+    var exceptLast = block.view().dropLast(1).flatMap(it -> optimize(it, false));
+    var last = optimize(block.getLast(), true);
+    return exceptLast.concat(last).toImmutableSeq();
+  }
+
+  static @NotNull SeqView<FreeStmt> optimize(FreeStmt stmt, boolean endOfBreakable) {
     return switch (stmt) {
       case FreeStmt.Switch(var elim, var cases, var branch, var defaultCase) -> {
-        branch = branch.map(it -> it.flatMap(FreeOptimizer::optimize));
-        defaultCase = defaultCase.flatMap(FreeOptimizer::optimize);
-        if (branch.isEmpty()) yield defaultCase;
+        branch = branch.map(it -> optimizeBlock(it, endOfBreakable));
+        defaultCase = optimizeBlock(defaultCase, endOfBreakable);
+        if (branch.isEmpty()) yield defaultCase.view();
         if (defaultCase.sizeEquals(1) && defaultCase.getFirst() == FreeStmt.Unreachable.INSTANCE) {
           if (branch.sizeEquals(1)) {
-            yield branch.getFirst();
+            yield branch.getFirst().view();
+          } else if (branch.sizeEquals(2)) {
+            yield SeqView.of(new FreeStmt.IfThenElse(
+              new FreeStmt.Condition.IsIntEqual(new FreeExpr.RefVariable(elim), cases.getFirst()),
+              branch.getFirst(),
+              branch.getLast()
+            ));
           } else if (branch.sizeGreaterThan(1)) {
-            yield ImmutableSeq.of(new FreeStmt.Switch(elim, cases.dropLast(1), branch.dropLast(1), branch.getLast()));
+            yield SeqView.of(new FreeStmt.Switch(elim, cases.dropLast(1), branch.dropLast(1), branch.getLast()));
           }
         }
-        yield ImmutableSeq.of(new FreeStmt.Switch(elim, cases, branch, defaultCase));
+        yield SeqView.of(new FreeStmt.Switch(elim, cases, branch, defaultCase));
       }
-      case FreeStmt.Breakable(var stmts) ->
-        ImmutableSeq.of(new FreeStmt.Breakable(stmts.flatMap(FreeOptimizer::optimize)));
+      case FreeStmt.Breakable(var stmts) -> SeqView.of(new FreeStmt.Breakable(optimizeBlock(stmts, true)));
       case FreeStmt.IfThenElse(var cond, var thenBlock, var elseBlock) -> {
-        var newThenBlock = thenBlock.flatMap(FreeOptimizer::optimize);
-        var newElseBlock = elseBlock == null ? null : elseBlock.flatMap(FreeOptimizer::optimize);
-        yield ImmutableSeq.of(new FreeStmt.IfThenElse(cond, newThenBlock, newElseBlock));
+        thenBlock = optimizeBlock(thenBlock, endOfBreakable);
+        if (elseBlock != null) {
+          elseBlock = optimizeBlock(elseBlock, endOfBreakable);
+          if (elseBlock.isEmpty()) elseBlock = null;
+        }
+        yield SeqView.of(new FreeStmt.IfThenElse(cond, thenBlock, elseBlock));
       }
-      default -> ImmutableSeq.of(stmt);
+      case FreeStmt.Break _ when endOfBreakable -> SeqView.empty();
+      default -> SeqView.of(stmt);
     };
   }
 }

jit-compiler/src/main/java/org/aya/compiler/free/morphism/free/FreeVariable.java

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2024 Tesla (Yinsen) Zhang.
+// Copyright (c) 2020-2025 Tesla (Yinsen) Zhang.
 // Use of this source code is governed by the MIT license that can be found in the LICENSE.md file.
 package org.aya.compiler.free.morphism.free;
 
@@ -8,16 +8,9 @@
 
 public sealed interface FreeVariable extends LocalVariable {
   record Local(int index) implements FreeVariable {
-    @Override
-    public @NotNull FreeJavaExpr ref() {
-      return new FreeExpr.RefVariable(this);
-    }
+    @Override public @NotNull FreeJavaExpr ref() { return new FreeExpr.RefVariable(this); }
   }
-
   record Arg(int nth) implements FreeVariable {
-    @Override
-    public @NotNull FreeJavaExpr ref() {
-      return new FreeExpr.RefVariable(this);
-    }
+    @Override public @NotNull FreeJavaExpr ref() { return new FreeExpr.RefVariable(this); }
   }
 }

jit-compiler/src/main/java/org/aya/compiler/serializers/ConSerializer.java

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2024 Tesla (Yinsen) Zhang.
+// Copyright (c) 2020-2025 Tesla (Yinsen) Zhang.
 // Use of this source code is governed by the MIT license that can be found in the LICENSE.md file.
 package org.aya.compiler.serializers;
 
@@ -42,27 +42,21 @@ public ConSerializer(ModuleSerializer.@NotNull MatchyRecorder recorder) {
     ));
   }
 
-  /**
-   * @see JitCon#isAvailable(Seq)
-   */
+  /// @param unit must be indexed, otherwise it should use the default impl.
+  /// @see JitCon#isAvailable
   private void buildIsAvailable(@NotNull FreeCodeBuilder builder, ConDef unit, @NotNull LocalVariable argsTerm) {
     FreeJavaExpr matchResult;
     var termSeq = builder.invoke(Constants.SEQ_TOIMMSEQ, argsTerm.ref(), ImmutableSeq.empty());
-    if (unit.pats.isEmpty()) {
-      // not indexed data type, this constructor is always available
-      matchResult = builder.invoke(Constants.RESULT_OK, ImmutableSeq.of(termSeq));
-    } else {
-      // It is too stupid to serialize pat meta solving, so we just call PatMatcher
-      var patsTerm = unit.pats.map(x -> new PatternExprializer(builder, true, recorder).serialize(x));
-      var patsSeq = AbstractExprializer.makeImmutableSeq(builder, Pat.class, patsTerm);
-      var id = builder.invoke(Constants.CLOSURE_ID, ImmutableSeq.empty());
-      var matcherTerm = builder.mkNew(PatMatcher.class, ImmutableSeq.of(
-        builder.iconst(true), id
-      ));
+    // It is too stupid to serialize pat meta solving, so we just call PatMatcher
+    var patsTerm = unit.pats.map(x -> new PatternExprializer(builder, true, recorder).serialize(x));
+    var patsSeq = AbstractExprializer.makeImmutableSeq(builder, Pat.class, patsTerm);
+    var id = builder.invoke(Constants.CLOSURE_ID, ImmutableSeq.empty());
+    var matcherTerm = builder.mkNew(PatMatcher.class, ImmutableSeq.of(
+      builder.iconst(true), id
+    ));
 
-      matchResult = builder.invoke(Constants.PATMATCHER_APPLY, matcherTerm,
-        ImmutableSeq.of(patsSeq, termSeq));
-    }
+    matchResult = builder.invoke(Constants.PATMATCHER_APPLY, matcherTerm,
+      ImmutableSeq.of(patsSeq, termSeq));
 
     builder.returnWith(matchResult);
   }
@@ -90,10 +84,10 @@ private void buildEquality(
 
   @Override public @NotNull ConSerializer serialize(@NotNull FreeClassBuilder builder0, ConDef unit) {
     buildFramework(builder0, unit, builder -> {
-      builder.buildMethod(
+      if (unit.pats.isNotEmpty()) builder.buildMethod(
         FreeUtil.fromClass(Result.class),
         "isAvailable",
-        ImmutableSeq.of(Constants.CD_Seq),
+        ImmutableSeq.of(Constants.CD_ImmutableSeq),
         (ap, builder1) ->
           buildIsAvailable(builder1, unit, ap.arg(0)));
 

jit-compiler/src/main/java/org/aya/compiler/serializers/JitTeleSerializer.java

@@ -59,7 +59,7 @@ protected JitTeleSerializer(
     @NotNull Consumer<FreeClassBuilder> continuation
   ) {
     super.buildFramework(builder, unit, nestBuilder -> {
-      nestBuilder.buildMethod(
+      if (unit.telescope().isNotEmpty()) nestBuilder.buildMethod(
         Constants.CD_Term,
         "telescope",
         ImmutableSeq.of(ConstantDescs.CD_int, Constants.CD_Seq),
@@ -105,7 +105,7 @@ protected void buildTelescope(@NotNull FreeCodeBuilder builder, @NotNull T unit,
 
         cb.returnWith(result);
       },
-      builder1 -> builder1.unreachable());
+      FreeCodeBuilder::unreachable);
   }
 
   /**

syntax/src/main/java/org/aya/syntax/compile/JitCon.java

@@ -1,23 +1,20 @@
-// Copyright (c) 2020-2024 Tesla (Yinsen) Zhang.
+// Copyright (c) 2020-2025 Tesla (Yinsen) Zhang.
 // Use of this source code is governed by the MIT license that can be found in the LICENSE.md file.
 package org.aya.syntax.compile;
 
 import kala.collection.Seq;
 import kala.collection.immutable.ImmutableSeq;
 import kala.collection.mutable.MutableArrayList;
-import kala.collection.mutable.MutableList;
 import kala.control.Result;
 import org.aya.generic.State;
 import org.aya.syntax.core.def.ConDefLike;
 import org.aya.syntax.core.def.DataDefLike;
 import org.aya.syntax.core.term.FreeTerm;
 import org.aya.syntax.core.term.Param;
 import org.aya.syntax.core.term.Term;
-import org.aya.syntax.ref.GenerateKind;
 import org.aya.syntax.ref.LocalVar;
 import org.aya.syntax.telescope.JitTele;
 import org.aya.util.error.Panic;
-import org.aya.util.error.SourcePos;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.UnknownNullability;
 
@@ -37,13 +34,14 @@ protected JitCon(
     this.selfTeleSize = selfTeleSize;
   }
 
-  /**
-   * Whether this constructor is available of data type
-   *
-   * @param args the argument to the data type
-   * @return a match result, a sequence of substitution if success
-   */
-  public abstract @NotNull Result<ImmutableSeq<Term>, State> isAvailable(@NotNull Seq<Term> args);
+  /// Whether this constructor is available of data type.
+  /// The default impl is for non-indexed constructors, where it's always available.
+  ///
+  /// @param args the argument to the data type
+  /// @return a match result, a sequence of substitution if success
+  public @NotNull Result<ImmutableSeq<Term>, State> isAvailable(@NotNull ImmutableSeq<Term> args) {
+    return Result.ok(args);
+  }
 
   @Override public boolean hasEq() { return hasEq; }
   @Override public @NotNull Term equality(Seq<Term> args, boolean is0) { throw new Panic("Not an HIT"); }

syntax/src/main/java/org/aya/syntax/telescope/AbstractTele.java

@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2024 Tesla (Yinsen) Zhang.
+// Copyright (c) 2020-2025 Tesla (Yinsen) Zhang.
 // Use of this source code is governed by the MIT license that can be found in the LICENSE.md file.
 package org.aya.syntax.telescope;
 
@@ -15,26 +15,24 @@
 import org.aya.syntax.core.term.Param;
 import org.aya.syntax.core.term.Term;
 import org.aya.syntax.ref.LocalVar;
+import org.aya.util.error.Panic;
 import org.jetbrains.annotations.NotNull;
 
-/**
- * Index-safe telescope
- */
+/// Index-safe telescope
 public interface AbstractTele {
-  /**
-   * @param teleArgs the arguments before {@param i}, for constructor, it also contains the arguments to the data
-   */
+  /// @param teleArgs the arguments before {@param i}, for constructor, it also contains the arguments to the data
   default @NotNull Term telescope(int i, Term[] teleArgs) {
     return telescope(i, ArraySeq.wrap(teleArgs));
   }
 
-  /**
-   * Get the type of {@param i}-th (count from {@code 0}) parameter.
-   *
-   * @param teleArgs the arguments to the former parameters
-   * @return the type of {@param i}-th parameter.
-   */
-  @NotNull Term telescope(int i, Seq<Term> teleArgs);
+  /// Get the type of {@param i}-th (count from `0`) parameter.
+  /// The default implementation is for empty telescope, because there are many such cases.
+  ///
+  /// @param teleArgs the arguments to the former parameters
+  /// @return the type of {@param i}-th parameter.
+  default @NotNull Term telescope(int i, Seq<Term> teleArgs) {
+    return Panic.unreachable();
+  }
 
   /**
    * Get the result of this signature