infer type for deopt + bugfixes

Author: Jakobeha

fir-intellij-plugin/src/main/java/org/prlprg/fir/intellij/FirNavBarModelExtension.java

@@ -0,0 +1,35 @@
+package org.prlprg.fir.intellij;
+
+import com.intellij.ide.navigationToolbar.StructureAwareNavBarModelExtension;
+import com.intellij.lang.Language;
+import javax.swing.Icon;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+/** Adds FIR function declarations and versions to the navigation bar. */
+public final class FirNavBarModelExtension extends StructureAwareNavBarModelExtension {
+  @Override
+  protected @NotNull Language getLanguage() {
+    return FirLanguage.INSTANCE;
+  }
+
+  @Override
+  public @Nullable String getPresentableText(Object object) {
+    if (object instanceof FirFunDeclaration fun) return fun.getHeaderText();
+    if (object instanceof FirVersion ver) return ver.getSignatureText();
+    return null;
+  }
+
+  @Override
+  public @Nullable Icon getIcon(Object object) {
+    if (object instanceof FirFunDeclaration fun) {
+      var pres = fun.getPresentation();
+      if (pres != null) return pres.getIcon(false);
+    }
+    if (object instanceof FirVersion ver) {
+      var pres = ver.getPresentation();
+      if (pres != null) return pres.getIcon(false);
+    }
+    return null;
+  }
+}

fir-intellij-plugin/src/main/resources/META-INF/plugin.xml

@@ -39,5 +39,7 @@
 
     <breadcrumbsInfoProvider
         implementation="org.prlprg.fir.intellij.FirBreadcrumbsProvider"/>
+
+    <navbar implementation="org.prlprg.fir.intellij.FirNavBarModelExtension"/>
   </extensions>
 </idea-plugin>

server/src/main/java/org/prlprg/bc2fir/BC2FirCFGCompiler.java

@@ -533,12 +533,16 @@ case StartFor(var _, var elemName, var step) -> {
                     seq));
         var init = new Constant(SEXPs.integer(0));
         var index = insertAndReturn("_idx", new Aea(init));
+        // Must be pushed for deopt in body
+        push(length);
         push(index);
         setJump(goto_(stepBb));
 
         // For loop step
         moveTo(stepBb);
         // Increment the index
+        index = pop();
+        length = top();
         var index1 =
             insertAndReturn(
                 "_idx",
@@ -548,7 +552,7 @@ case StartFor(var _, var elemName, var step) -> {
                         ImmutableList.of(Type.BOXED_INTEGER, Type.BOXED_INTEGER),
                         Type.BOXED_INTEGER,
                         Effects.NONE),
-                    pop(),
+                    index,
                     new Constant(SEXPs.integer(1))));
         push(index1);
         // Compare the index to the length
@@ -651,7 +655,8 @@ case StepFor(var body) -> {
       }
       case EndFor() -> {
         compileEndLoop(LoopType.FOR);
-        // Pop index and replace with `NULL` (a SEXP, in case we use it)
+        // Pop index and length and push `NULL` (loop's "output")
+        pop();
         pop();
         push(new Constant(SEXPs.NULL));
       }
@@ -661,31 +666,34 @@ case EndFor() -> {
       case LdTrue() -> push(new Constant(SEXPs.TRUE));
       case LdFalse() -> push(new Constant(SEXPs.FALSE));
       case GetVar(var name) -> {
+        var preStack = ImmutableList.copyOf(stack);
         pushInsert(getStr(name), new Load(LoadType.LOCAL_VAR, getVar(name)));
-        tryAddCheckpoint(false);
+        tryAddCheckpoint(false, preStack);
+
         pushInsert(getStr(name), new Force(true, pop()));
         insert(intrinsic("checkMissing", top()));
-        tryAddCheckpoint(true);
+        tryAddCheckpoint(true, stack);
       }
       case DdVal(var name) -> {
+        var preStack = ImmutableList.copyOf(stack);
         var ddIndex = NamedVariable.ddNum(get(name).ddNum());
         pushInsert(getStr(name), new Load(LoadType.LOCAL_VAR, ddIndex));
-        tryAddCheckpoint(false);
+        tryAddCheckpoint(false, preStack);
+
         pushInsert(getStr(name), new Force(true, pop()));
         insert(intrinsic("checkMissing", top()));
-        tryAddCheckpoint(true);
+        tryAddCheckpoint(true, stack);
       }
       case SetVar(var name) -> insert(new Store(StoreType.LOCAL_VAR, getVar(name), top()));
       case GetFun(var name) -> {
-        tryAddCheckpoint(false);
         var fun = insertAndReturn(getStr(name), new Load(LoadType.LOCAL_FUN, getVar(name)));
         pushCall(fun);
-        tryAddCheckpoint(true);
+        tryAddCheckpoint(true, stack);
       }
       case GetGlobFun(var name) -> {
         var fun = insertAndReturn(getStr(name), new Load(LoadType.GLOBAL_FUN, getVar(name)));
         pushCall(fun);
-        tryAddCheckpoint(true);
+        tryAddCheckpoint(true, stack);
       }
       case GetBuiltin(var name) -> pushCall(new Builtin(get(name).name()));
       case GetIntlBuiltin(var name) -> pushCall(new Builtin(get(name).name()));
@@ -890,14 +898,18 @@ case Or2nd(var _) -> {
         insert(this::goto_);
       }
       case GetVarMissOk(var name) -> {
+        var preStack = ImmutableList.copyOf(stack);
         pushInsert(getStr(name), new Load(LoadType.LOCAL_VAR, getVar(name)));
-        tryAddCheckpoint(false);
+        tryAddCheckpoint(false, preStack);
+
         pushInsertThenCp(getStr(name), new Force(true, pop()));
       }
       case DdValMissOk(var name) -> {
+        var preStack = ImmutableList.copyOf(stack);
         var ddIndex = get(name).ddNum();
         pushInsert(getStr(name), new Load(LoadType.LOCAL_VAR, NamedVariable.ddNum(ddIndex)));
-        tryAddCheckpoint(false);
+        tryAddCheckpoint(false, preStack);
+
         pushInsertThenCp(getStr(name), new Force(true, pop()));
       }
       case Visible() -> insert(intrinsic("setVisible"));
@@ -1196,9 +1208,9 @@ private Argument compilePromise(SEXP codeSexp) {
     return insertAndReturn("_p", new Promise(Type.ANY_VALUE_SEXP, Effects.REFLECT, cfg));
   }
 
-  /// End the previously-compiled for loop instruction (the latest [#pushWhileOrRepeatLoop(BB,
-  /// BB)] or [#pushForLoop(BB,BB,BB)]): pop its data, assert that the loop is of the correct
-  /// type, and that the cursor is right after its end.
+  /// End the previously-compiled for loop instruction (the latest
+  /// [#pushWhileOrRepeatLoop(BB, BB)] or [#pushForLoop(BB,BB,BB)]): pop its data, assert that
+  /// the loop is of the correct type, and that the cursor is right after its end.
   private void compileEndLoop(LoopType type) {
     var loop = popLoop();
     if (loop.type != type) {
@@ -1288,27 +1300,33 @@ private void compileCall(Call call) {
   // endregion compile instructions
 
   // region checkpoints
-  void pushInsertThenCp(String name, Expression expression) {
+  private void pushInsertThenCp(String name, Expression expression) {
     pushInsert(name, expression);
-    tryAddCheckpoint(true);
+    tryAddCheckpoint(true, stack);
   }
 
-  void pushInsertThenCp(Expression expression) {
+  private void pushInsertThenCp(Expression expression) {
     pushInsert(expression);
-    tryAddCheckpoint(true);
+    tryAddCheckpoint(true, stack);
   }
 
-  void tryAddCheckpoint(boolean afterInstruction) {
-    var deoptBcPos = afterInstruction ? bcPos + 1 : bcPos;
+  /// Add a checkpoint at the current FIŘ position, before or after the current bytecode
+  /// position
+  ///
+  /// - If `afterBcInstr = false`, `stack` must be the stack before the bytecode was compiled,
+  ///   and every FIŘ instruction emitted for the bytecode so far must be pure.
+  /// - If `afterBcInstr = true`, `stack` must be the current stack, and no more FIŘ
+  ///   instructions can be emitted for this bytecode instruction.
+  private void tryAddCheckpoint(boolean afterBcInstr, List<Argument> stack) {
+    var deoptBcPos = afterBcInstr ? bcPos + 1 : bcPos;
     var deopt = cfg.addBB("D" + numDeopts++);
-    deopt.setJump(new Deopt(deoptBcPos, ImmutableList.copyOf(stack)));
+    deopt.setJump(new Deopt(deoptBcPos, stack));
 
     // Don't add phis because they're never necessary.
     // Don't use `insert` because it adds phis.
     var success = cfg.addBB();
     cursor.bb().setJump(new Checkpoint(new Target(success), new Target(deopt)));
     cursor.moveToStart(success);
-    // TODO: Is the below correct?
     bbsWithPhis.add(success);
   }
 
@@ -1537,8 +1555,8 @@ private void pushInsert(String name, Expression expression) {
     push(insertAndReturn(name, expression));
   }
 
-  /// Push a value onto the "virtual stack" so that the next call to [#pop(Class)] or
-  /// [#top(Class)] will return it.
+  /// Push a value onto the "virtual stack" so that the next call to [#pop()] or
+  /// [#top()] will return it.
   ///
   /// The bytecode is stack-based and IR is SSA-form. To convert properly, when the compiler
   /// encounters `BcInstr.Ld...` and other instructions that push real values onto the
@@ -1767,8 +1785,8 @@ private Expression warning(String message) {
         new Constant(SEXPs.NULL));
   }
 
-  /// Stub for function guard (TODO what does [#BcInstr.CheckFun] do again?), which can fallback to
-  /// an intrinsic but usually gets caught and turned into [DynamicCallee].
+  /// Stub for function guard (TODO what does [BcInstr.CheckFun] do again?), which can fallback
+  /// to an intrinsic but usually gets caught and turned into [DynamicCallee].
   private Expression checkFun(Argument fun) {
     return intrinsic("checkFun", fun);
   }

server/src/main/java/org/prlprg/fir/analyze/type/InferType.java

@@ -31,6 +31,7 @@
 import org.prlprg.fir.ir.expression.Store;
 import org.prlprg.fir.ir.expression.SubscriptRead;
 import org.prlprg.fir.ir.expression.SubscriptWrite;
+import org.prlprg.fir.ir.instruction.Deopt;
 import org.prlprg.fir.ir.instruction.Return;
 import org.prlprg.fir.ir.type.Concreteness;
 import org.prlprg.fir.ir.type.Kind.PrimitiveVector;
@@ -56,8 +57,12 @@ public InferType(Abstraction scope) {
   public @Nullable Type of(CFG cfg) {
     Type result = null;
     for (var bb : cfg.bbs()) {
-      if (bb.jump() instanceof Return(_, var value)) {
-        result = Type.union(result, of(value));
+      switch (bb.jump()) {
+        case Return(_, var value) -> result = Type.union(result, of(value));
+        case Deopt _ -> {
+          return Type.ANY_VALUE_SEXP;
+        }
+        default -> {}
       }
     }
     return result;

server/src/main/java/org/prlprg/fir/interpret/internal/Builtins.java

@@ -486,7 +486,7 @@ private static SEXP applyBinaryPreservingInt(
       }
       default -> {}
     }
-    throw interpreter.fail("`" + ctx + "` generic requires logical or numeric args");
+    throw interpreter.failUnsupported("`" + ctx + "` generic requires logical or numeric args");
   }
 
   /// Apply unary math op on SEXP, preserving int type (for +, -)
@@ -520,7 +520,9 @@ private static SEXP applyUnaryPreservingInt(
                   .map(i -> lv.get(i) == Logical.NA ? Constants.NA_INT : lv.get(i).toInt())
                   .map(j -> (int) fn.applyAsDouble(j))
                   .toArray());
-      default -> throw interpreter.fail("`" + ctx + "` unary requires a logical or numeric arg");
+      default ->
+          throw interpreter.failUnsupported(
+              "`" + ctx + "` unary requires a logical or numeric arg");
     };
   }
 
@@ -2586,7 +2588,7 @@ private static double sexpToDouble(SEXP sexp, InternalInterpreter interpreter, S
           : sexp.asScalarInteger().get();
     if (sexp.asScalarReal().isPresent()) return sexp.asScalarReal().get();
     if (sexp.asScalarLogical().isPresent()) return sexp.asScalarLogical().get().toInt();
-    throw interpreter.fail(ctx + " requires a numeric scalar");
+    throw interpreter.failUnsupported("Mock " + ctx + " only implemented for a numeric scalar");
   }
 
   private static @Nullable Double sexpToDoubleOpt(SEXP sexp) {
@@ -2606,7 +2608,7 @@ private static int sexpToInt(SEXP sexp, InternalInterpreter interpreter, String
           ? Constants.NA_INT
           : (int) sexp.asScalarReal().get().doubleValue();
     if (sexp.asScalarLogical().isPresent()) return sexp.asScalarLogical().get().toInt();
-    throw interpreter.fail(ctx + " requires a numeric scalar");
+    throw interpreter.failUnsupported("Mock " + ctx + " only implemented for numeric scalar");
   }
 
   private static Logical sexpToLogical(SEXP sexp, InternalInterpreter interpreter, String ctx) {
@@ -2621,7 +2623,8 @@ private static Logical sexpToLogical(SEXP sexp, InternalInterpreter interpreter,
       if (Double.isNaN(v)) return Logical.NA;
       return v == 0.0 ? Logical.FALSE : Logical.TRUE;
     }
-    throw interpreter.failUnsupported(ctx + " requires a scalar convertible to logical");
+    throw interpreter.failUnsupported(
+        "Mock " + ctx + " only implemented for scalar convertible to logical");
   }
 
   private static String sexpToString(SEXP sexp) {

server/src/main/java/org/prlprg/fir/interpret/internal/InternalInterpreter.java

@@ -2,6 +2,7 @@
 
 import static org.prlprg.fir.GlobalModules.BOX_FUN;
 import static org.prlprg.fir.GlobalModules.UNBOX_FUN;
+import static org.prlprg.fir.ir.cfg.iterator.BbDfs.bbDfs;
 import static org.prlprg.sexp.ArgumentMatcher.matchArguments;
 
 import com.google.common.collect.ImmutableList;
@@ -1168,11 +1169,13 @@ private CFGCursor restoreDeopt(int pc, List<SEXP> deoptStack, @Nullable CFG deop
     }
 
     // Get corresponding deopt BC and run sanity checks
+    // Corresponding = same bytecode, and if there are multiple, the one with the dominating
+    // predecessor (the other(s) do speculative pure instructions that are redundant in
+    // bytecode, but may not be here).
     var deoptBc =
-        deoptRestoreCfg.bbs().stream()
+        com.google.common.collect.Streams.stream(bbDfs(deoptRestoreCfg))
             .filter(bb -> bb.jump() instanceof Deopt(_, var pc1, _) && pc == pc1)
-            .collect(
-                Streams.zeroOneOrThrow(() -> fail("multiple deopt branches with the same pc?")))
+            .findFirst()
             .orElseThrow(
                 () ->
                     fail(

server/src/main/java/org/prlprg/fir/ir/abstraction/Abstraction.java

@@ -324,7 +324,7 @@ case Consume(var register) -> {
   }
 
   /// Converts the string into a register name that is valid and doesn't already exist.
-  public Register resemblance(String prefix) {
+  private Register resemblance(String prefix) {
     return Variable.register(
         nextLocalDisambiguator.disambiguate(Register.resemblance(prefix).name()));
   }

server/src/main/java/org/prlprg/fir/ir/abstraction/substitute/InlineSubstituter.java

@@ -58,8 +58,8 @@ protected void commitAffectLocals() {
       var type = scope.removeLocal(oldRegister).type();
       newLocals.add(new Local(newRegister, type));
     }
-    // Must add new locals after,
-    // otherwise we may add a local that already exists, but will be removed later.
+    // Must add new locals after
+    // Otherwise we may add a local that initially existed, but was removed
     for (var local : newLocals) {
       scope.addLocal(local);
     }

server/src/main/java/org/prlprg/fir/opt/Inline.java

@@ -287,8 +287,13 @@ private void tryInlineCall(
       // substitute parameters with arguments, and substitute locals with new locals.
       var body = new Abstraction(scope.module(), List.of());
       assert body.cfg() != null;
+      var tempLocalsForDisambiguation = new ArrayList<Register>();
       for (var parameter : callee.parameters()) {
         body.addLocal(new Local(parameter.variable(), parameter.type()));
+        if (!scope.contains(parameter.variable())) {
+          scope.addLocal(new Local(parameter.variable(), Type.ANY_SEXP));
+          tempLocalsForDisambiguation.add(parameter.variable());
+        }
       }
       for (var local : callee.locals()) {
         if (!(local.variable() instanceof Register _)) {
@@ -307,6 +312,9 @@ private void tryInlineCall(
         localSubstituter.stage(oldVariable, disambiguatedVariable);
       }
       localSubstituter.commit();
+      for (var tempLocalForDisambiguation : tempLocalsForDisambiguation) {
+        scope.removeLocal(tempLocalForDisambiguation);
+      }
       var parameterSubstituter = new Substituter(body);
       for (int i = 0; i < callee.parameters().size(); i++) {
         var parameter = callee.parameters().get(i).variable();