Fix #1151

base/src/main/java/org/aya/primitive/PrimFactory.java

@@ -8,9 +8,11 @@
 import kala.control.Option;
 import kala.tuple.Tuple;
 import org.aya.normalize.Normalizer;
+import org.aya.syntax.compile.JitPrim;
 import org.aya.syntax.concrete.stmt.decl.PrimDecl;
 import org.aya.syntax.core.Closure;
 import org.aya.syntax.core.def.PrimDef;
+import org.aya.syntax.core.def.PrimDefLike;
 import org.aya.syntax.core.term.*;
 import org.aya.syntax.core.term.call.PrimCall;
 import org.aya.syntax.core.term.repr.StringTerm;
@@ -33,7 +35,7 @@
 
 public class PrimFactory {
   private final @NotNull Map<@NotNull ID, @NotNull PrimSeed> seeds;
-  private final @NotNull EnumMap<@NotNull ID, @NotNull PrimDef> defs = new EnumMap<>(ID.class);
+  private final @NotNull EnumMap<@NotNull ID, @NotNull PrimDefLike> defs = new EnumMap<>(ID.class);
 
   public PrimFactory() {
     seeds = ImmutableMap.from(ImmutableSeq.of(
@@ -45,6 +47,11 @@
     ).map(seed -> Tuple.of(seed.name, seed)));
   }
 
+  public void definePrim(PrimDefLike prim) {
+    assert !isForbiddenRedefinition(prim.id(), prim instanceof JitPrim);
+    defs.put(prim.id(), prim);
+  }
+
   @FunctionalInterface
   public interface Unfolder extends BiFunction<@NotNull PrimCall, @NotNull TyckState, @NotNull Term> { }
 
@@ -96,7 +103,7 @@
   }
 
   final @NotNull PrimSeed stringType =
-    new PrimSeed(ID.STRING, this::primCall,
+    new PrimSeed(ID.STRING, (prim, _) -> prim,
       ref -> new PrimDef(ref, Type0, ID.STRING), ImmutableSeq.empty());
 
   final @NotNull PrimSeed stringConcat =
@@ -153,41 +160,48 @@
   }
   */
 
-  private @NotNull PrimCall primCall(@NotNull PrimCall prim, @NotNull TyckState tyckState) { return prim; }
-
   public final @NotNull PrimSeed intervalType = new PrimSeed(ID.I,
     ((_, _) -> DimTyTerm.INSTANCE),
     ref -> new PrimDef(ref, SortTerm.ISet, ID.I),
     ImmutableSeq.empty());
 
-  public @NotNull PrimDef factory(@NotNull ID name, @NotNull DefVar<PrimDef, PrimDecl> ref) {
-    assert suppressRedefinition() || !have(name);
-    var rst = seeds.get(name).supply(ref);
-    defs.put(name, rst);
+  public @NotNull PrimDefLike factory(@NotNull ID name, @NotNull DefVar<PrimDef, PrimDecl> ref) {
+    var rst = new PrimDef.Delegate(seeds.get(name).supply(ref).ref);
+    definePrim(rst);
     return rst;
   }
 
   public @NotNull PrimCall getCall(@NotNull ID id, @NotNull ImmutableSeq<Term> args) {
-    return new PrimCall(getOption(id).get().ref(), 0, args);
+    return new PrimCall(getOption(id).get(), 0, args);
   }
 
   public @NotNull PrimCall getCall(@NotNull ID id) {
-    return getCall(id, ImmutableSeq.empty());
+    return new PrimCall(getOption(id).get());
   }
 
-  public @NotNull Option<PrimDef> getOption(@NotNull ID name) {
+  public @NotNull Option<PrimDefLike> getOption(@NotNull ID name) {
     return Option.ofNullable(defs.get(name));
   }
 
   public boolean have(@NotNull ID name) {
     return defs.containsKey(name);
   }
 
-  /** whether redefinition should be treated as error */
-  @ForLSP public boolean suppressRedefinition() { return false; }
-
-  public @NotNull PrimDef getOrCreate(@NotNull ID name, @NotNull DefVar<PrimDef, PrimDecl> ref) {
-    return getOption(name).getOrElse(() -> factory(name, ref));
+  /**
+   * Whether this definition is a redefinition that should be treated as error.
+   * There are two cases where a redefinition is allowed:
+   * <ul>
+   *   <li>When we are working in an LSP, and users can reload a file to redefine things.</li>
+   *   <li>When we are serializing a file, which we will deserialize immediately, and this will
+   *     replace the existing PrimDefs with their JIT-compiled version.</li>
+   * </ul>
+   *
+   * @return true if redefinition is forbidden.
+   */
+  @ForLSP public boolean isForbiddenRedefinition(@NotNull PrimDef.ID id, boolean isJit) {
+    if (isJit)
+      return have(id) && defs.get(id) instanceof JitPrim;
+    else return have(id);
   }
 
   public @NotNull Option<ImmutableSeq<@NotNull ID>> checkDependency(@NotNull ID name) {

base/src/main/java/org/aya/resolve/ResolveInfo.java

@@ -116,6 +116,10 @@ public void open(@NotNull ResolveInfo other, @NotNull SourcePos sourcePos, @NotN
     });
   }
 
+  public @NotNull Reporter reporter() {
+    return this.opSet.reporter;
+  }
+
   @Debug.Renderer(text = "opInfo.name()")
   public record RenamedOpDecl(@NotNull OpInfo opInfo) implements OpDecl { }
 }

base/src/main/java/org/aya/resolve/module/ModuleCallback.java

@@ -9,6 +9,6 @@
 
 @FunctionalInterface
 public interface ModuleCallback<E extends Exception> {
-  void onModuleTycked(@NotNull ResolveInfo moduleResolve, @NotNull ImmutableSeq<TyckDef> defs)
+  void onModuleTycked(@NotNull ResolveInfo resolveInfo, @NotNull ImmutableSeq<TyckDef> defs)
     throws E;
 }

base/src/main/java/org/aya/resolve/visitor/StmtBinder.java

@@ -13,6 +13,7 @@
 import org.aya.syntax.concrete.stmt.decl.DataDecl;
 import org.aya.syntax.concrete.stmt.decl.FnDecl;
 import org.aya.syntax.concrete.stmt.decl.PrimDecl;
+import org.aya.syntax.ref.AnyDefVar;
 import org.aya.syntax.ref.DefVar;
 import org.aya.util.binop.OpDecl;
 import org.aya.util.error.Panic;
@@ -37,15 +38,15 @@ public void bind(@NotNull Context ctx, @NotNull BindBlock bindBlock, OpDecl self
       bind(self, ctx, OpDecl.BindPred.Tighter, tighter)));
   }
 
-  private @Nullable DefVar<?, ?> bind(
+  private @Nullable AnyDefVar bind(
     @NotNull OpDecl self, @NotNull Context ctx,
     @NotNull OpDecl.BindPred pred, @NotNull QualifiedID id
   ) throws Context.ResolvingInterruptedException {
     var var = ctx.get(id);
     var opDecl = info.resolveOpDecl(var);
     if (opDecl != null) {
       info.opSet().bind(self, pred, opDecl, id.sourcePos());
-      return var instanceof DefVar<?, ?> defVar ? defVar : null;
+      return var instanceof AnyDefVar defVar ? defVar : null;
     }
 
     // make compiler happy 😥

base/src/main/java/org/aya/resolve/visitor/StmtPreResolver.java

@@ -132,7 +132,7 @@ public ImmutableSeq<ResolvingStmt> resolveStmt(@NotNull ImmutableSeq<Stmt> stmts
         var lack = factory.checkDependency(primID);
         if (lack.isNotEmpty() && lack.get().isNotEmpty())
           context.reportAndThrow(new PrimResolveError.Dependency(name, lack.get(), sourcePos));
-        else if (factory.have(primID) && !factory.suppressRedefinition())
+        else if (factory.isForbiddenRedefinition(primID, false))
           context.reportAndThrow(new PrimResolveError.Redefinition(name, sourcePos));
         factory.factory(primID, decl.ref);
         var resolvedCtx = resolveTopLevelDecl(decl, context);

base/src/main/java/org/aya/unify/TermComparator.java

@@ -4,10 +4,10 @@
 
 import kala.collection.immutable.ImmutableSeq;
 import org.aya.generic.Renamer;
+import org.aya.generic.term.DTKind;
 import org.aya.generic.term.SortKind;
 import org.aya.prettier.AyaPrettierOptions;
 import org.aya.syntax.core.term.*;
-import org.aya.generic.term.DTKind;
 import org.aya.syntax.core.term.call.*;
 import org.aya.syntax.core.term.marker.Formation;
 import org.aya.syntax.core.term.repr.IntegerTerm;
@@ -339,7 +339,7 @@ case ProjTerm(var lof, var ldx) -> {
       // We already compare arguments in compareApprox, if we arrive here,
       // it means their arguments don't match (even the ref don't match),
       // so we are unable to do more if we can't normalize them.
-      case FnCall _ -> null;
+      case FnCall _, PrimCall _ -> null;
 
       default -> throw noRules(lhs);
     };

base/src/test/java/org/aya/tyck/TyckTest.java

@@ -222,7 +222,7 @@ public static TyckResult tyck(@Language("Aya") @NotNull String code) {
     var moduleLoader = SyntaxTestUtil.moduleLoader();
     var callback = new ModuleCallback<RuntimeException>() {
       ImmutableSeq<TyckDef> ok;
-      @Override public void onModuleTycked(@NotNull ResolveInfo x, @NotNull ImmutableSeq<TyckDef> defs) { ok = defs; }
+      @Override public void onModuleTycked(@NotNull ResolveInfo resolveInfo, @NotNull ImmutableSeq<TyckDef> defs) { ok = defs; }
     };
     var info = moduleLoader.tyckModule(moduleLoader.resolve(SyntaxTestUtil.parse(code)), callback);
     return new TyckResult(callback.ok, info);

cli-impl/src/main/java/org/aya/cli/interactive/ReplCompiler.java

@@ -31,6 +31,7 @@
 import org.aya.syntax.GenericAyaFile;
 import org.aya.syntax.concrete.Expr;
 import org.aya.syntax.concrete.stmt.Stmt;
+import org.aya.syntax.core.def.PrimDef;
 import org.aya.syntax.core.def.TyckDef;
 import org.aya.syntax.core.term.Term;
 import org.aya.syntax.literate.CodeOptions.NormalizeMode;
@@ -75,7 +76,9 @@ public ReplCompiler(
     reporter = CountingReporter.delegate(delegateReporter);
     this.locator = locator != null ? locator : new SourceFileLocator.Module(this.modulePaths);
     this.primFactory = new PrimFactory() {
-      @Override public boolean suppressRedefinition() { return true; }
+      @Override public boolean isForbiddenRedefinition(PrimDef.@NotNull ID id, boolean isJit) {
+        return false;
+      }
     };
     this.shapeFactory = new ReplShapeFactory();
     this.opSet = new AyaBinOpSet(reporter);

cli-impl/src/main/java/org/aya/cli/library/LibraryModuleLoader.java

@@ -74,15 +74,22 @@
     assert program != null;
     var context = new EmptyContext(reporter, sourcePath).derive(mod);
     var resolveInfo = resolveModule(states.primFactory, context, program, recurseLoader);
-    source.resolveInfo().set(resolveInfo);
-    return tyckModule(resolveInfo, (moduleResolve, defs) -> {
+    tyckModule(resolveInfo, (moduleResolve, defs) -> {
       source.notifyTycked(moduleResolve, defs);
-      if (reporter.noError()) saveCompiledCore(source, moduleResolve, defs);
+      if (reporter.noError()) saveCompiledCore(source, moduleResolve, defs, recurseLoader);
     });
+
+    @Nullable var tyckedInfo = source.resolveInfo().get();
+    // I know we can set resolveInfo before the tyckModule, but that doesn't tell the goal of the code.
+    if (tyckedInfo == null) {
+      // this happens if [saveCompiledCore] meets an exception, use fallback ResolveInfo
+      tyckedInfo = resolveInfo;
+    }
+
+    return tyckedInfo;
   }
 
-  @Override
-  public boolean existsFileLevelModule(@NotNull ModulePath path) {
+  @Override public boolean existsFileLevelModule(@NotNull ModulePath path) {
     return owner.findModule(path) != null;
   }
 
@@ -94,11 +101,11 @@
   }
 
   private void saveCompiledCore(
-    @NotNull LibrarySource file,
-    @NotNull ResolveInfo resolveInfo,
-    @NotNull ImmutableSeq<TyckDef> defs
+    @NotNull LibrarySource file, @NotNull ResolveInfo resolveInfo,
+    @NotNull ImmutableSeq<TyckDef> defs, @NotNull ModuleLoader recurseLoader
   ) {
-    advisor.saveCompiledCore(file, resolveInfo, defs);
+    var info = advisor.saveCompiledCore(file, resolveInfo, defs, recurseLoader);
+    file.resolveInfo().set(info);
   }
 
   record United(@NotNull PrimFactory primFactory) {

cli-impl/src/main/java/org/aya/cli/library/incremental/CompilerAdvisor.java

@@ -69,11 +69,19 @@
     @NotNull ModuleLoader recurseLoader
   ) throws IOException, ClassNotFoundException;
 
-  void doSaveCompiledCore(
+  /**
+   * Save the well-typed defs as compiled defs
+   *
+   * @param file        the source file
+   * @param resolveInfo the resolve info of the module
+   * @param defs        the well-typed definitions
+   */
+  @NotNull ResolveInfo doSaveCompiledCore(
     @NotNull LibrarySource file,
     @NotNull ResolveInfo resolveInfo,
-    @NotNull ImmutableSeq<TyckDef> defs
-  ) throws IOException;
+    @NotNull ImmutableSeq<TyckDef> defs,
+    @NotNull ModuleLoader recurseLoader
+  ) throws IOException, ClassNotFoundException;
 
   @ApiStatus.NonExtendable
   default @Nullable ResolveInfo loadCompiledCore(
@@ -91,17 +99,25 @@
     }
   }
 
+  /**
+   * @return possibly the compiled version of the {@link ResolveInfo}
+   * @see #doSaveCompiledCore
+   */
   @ApiStatus.NonExtendable
-  default void saveCompiledCore(
+  default @NotNull ResolveInfo saveCompiledCore(
     @NotNull LibrarySource file,
     @NotNull ResolveInfo resolveInfo,
-    @NotNull ImmutableSeq<TyckDef> defs
+    @NotNull ImmutableSeq<TyckDef> defs,
+    @NotNull ModuleLoader recurseLoader
   ) {
+    assert recurseLoader instanceof CachedModuleLoader<?>;
     try {
-      doSaveCompiledCore(file, resolveInfo, defs);
+      var info = doSaveCompiledCore(file, resolveInfo, defs, recurseLoader);
       updateLastModified(file);
-    } catch (IOException e) {
+      return info;
+    } catch (IOException | ClassNotFoundException e) {
       e.printStackTrace();
+      return resolveInfo;
     }
   }
 

cli-impl/src/main/java/org/aya/cli/library/incremental/DelegateCompilerAdvisor.java

@@ -56,9 +56,9 @@ public void notifyIncrementalJob(@NotNull ImmutableSeq<LibrarySource> modified,
     return delegate.doLoadCompiledCore(reporter, owner, mod, sourcePath, corePath, recurseLoader);
   }
 
-  @Override public void
-  doSaveCompiledCore(@NotNull LibrarySource file, @NotNull ResolveInfo resolveInfo, @NotNull ImmutableSeq<TyckDef> defs) throws IOException {
-    delegate.doSaveCompiledCore(file, resolveInfo, defs);
+  @Override public @NotNull ResolveInfo
+  doSaveCompiledCore(@NotNull LibrarySource file, @NotNull ResolveInfo resolveInfo, @NotNull ImmutableSeq<TyckDef> defs, @NotNull ModuleLoader recurseLoader) throws IOException, ClassNotFoundException {
+    return delegate.doSaveCompiledCore(file, resolveInfo, defs, recurseLoader);
   }
   @Override public void close() throws Exception {
     delegate.close();