Do not enlarge types of functions defined by terms when unnecessary

JetBrains · Jun 9, 2024 · 762b4c8 · 762b4c8
1 parent 53aa4e6
commit 762b4c8
Showing 1 changed file with 20 additions and 1 deletion.
diff --git a/base/src/main/java/org/arend/typechecking/visitor/DefinitionTypechecker.java b/base/src/main/java/org/arend/typechecking/visitor/DefinitionTypechecker.java
@@ -1709,7 +1709,26 @@ private List<ExtElimClause> typecheckFunctionBody(FunctionDefinition typedDef, C
           }
 
           if (!def.isRecursive()) {
-            typedDef.setResultType(termResult.type);
+            Expression newType = termResult.type;
+            if ((typedDef.isSFunc() || kind == FunctionKind.CONS) && typedDef.getResultType() != null) {
+              Expression normNewType = newType.normalize(NormalizationMode.WHNF);
+              Expression oldType = typedDef.getResultType().normalize(NormalizationMode.WHNF);
+              if (oldType instanceof ClassCallExpression oldClassCall && normNewType instanceof ClassCallExpression newClassCall) {
+                Map<ClassField, Expression> impls = new LinkedHashMap<>();
+                for (Map.Entry<ClassField, Expression> entry : newClassCall.getImplementedHere().entrySet()) {
+                  if (oldClassCall.isImplemented(entry.getKey())) {
+                    impls.put(entry.getKey(), entry.getValue());
+                  }
+                }
+                if (impls.size() != newClassCall.getImplementedHere().size()) {
+                  newClassCall = new ClassCallExpression(newClassCall.getDefinition(), newClassCall.getLevels(), impls, newClassCall.getDefinition().getSort(), newClassCall.getDefinition().getUniverseKind());
+                  newClassCall.updateHasUniverses();
+                  typechecker.fixClassExtSort(newClassCall, def.getResultType());
+                  newType = newClassCall;
+                }
+              }
+            }
+            typedDef.setResultType(newType);
           }
           if (termResult.expression != null) {
             typedDef.setBody(termResult.expression);