test: make rbtree test use the optimizer

aya-prover · Jan 4, 2025 · 3999ede · 3999ede
1 parent 1a5b222
commit 3999ede
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/jit-compiler/src/main/java/org/aya/compiler/free/morphism/free/FreeOptimizer.java b/jit-compiler/src/main/java/org/aya/compiler/free/morphism/free/FreeOptimizer.java
@@ -29,13 +29,18 @@ case Clazz(var metadata, var owner, var nested, var superclass, var members) ->
       case FreeStmt.Switch(var elim, var cases, var branch, var defaultCase) -> {
         branch = branch.map(it -> it.flatMap(FreeOptimizer::optimize));
         defaultCase = defaultCase.flatMap(FreeOptimizer::optimize);
-        var defaultPanic = defaultCase.sizeEquals(1) && defaultCase.getFirst() == FreeStmt.Unreachable.INSTANCE;
-        if (defaultPanic && branch.sizeEquals(1)) {
-          yield branch.getFirst();
+        if (branch.isEmpty()) yield defaultCase;
+        if (defaultCase.sizeEquals(1) && defaultCase.getFirst() == FreeStmt.Unreachable.INSTANCE) {
+          if (branch.sizeEquals(1)) {
+            yield branch.getFirst();
+          } else if (branch.sizeGreaterThan(1)) {
+            yield ImmutableSeq.of(new FreeStmt.Switch(elim, cases.dropLast(1), branch.dropLast(1), branch.getLast()));
+          }
         }
         yield ImmutableSeq.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) ->
+        ImmutableSeq.of(new FreeStmt.Breakable(stmts.flatMap(FreeOptimizer::optimize)));
       case FreeStmt.IfThenElse(var cond, var thenBlock, var elseBlock) -> {
         var newThenBlock = thenBlock.flatMap(FreeOptimizer::optimize);
         var newElseBlock = elseBlock == null ? null : elseBlock.flatMap(FreeOptimizer::optimize);

diff --git a/jit-compiler/src/test/java/CompileTest.java b/jit-compiler/src/test/java/CompileTest.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.
 
 import kala.collection.immutable.ImmutableSeq;
@@ -104,7 +104,7 @@ public record TyckResult(@NotNull ImmutableSeq<TyckDef> defs, @NotNull ResolveIn
 
   public static @NotNull String serializeFrom(@NotNull TyckResult result) {
     return new ModuleSerializer(result.info.shapeFactory())
-      .serialize(SourceFreeJavaBuilder.create(), new ModuleSerializer.ModuleResult(
+      .serializeWithBestBuilder(new ModuleSerializer.ModuleResult(
         DumbModuleLoader.DUMB_MODULE_NAME, result.defs.filterIsInstance(TopLevelDef.class)));
   }