Fix a bug with external variables and coclauses

JetBrains · May 7, 2024 · 1350129 · 1350129
1 parent 8264da5
commit 1350129
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 10 deletions.
diff --git a/base/src/main/java/org/arend/term/concrete/Concrete.java b/base/src/main/java/org/arend/term/concrete/Concrete.java
@@ -826,8 +826,11 @@ public interface CoClauseElement extends ClassElement {
   }
 
   public static class CoClauseFunctionReference extends ClassFieldImpl {
+    public final TCDefReferable functionReference;
+
     public CoClauseFunctionReference(Object data, Referable implementedField, TCDefReferable functionReference, boolean isDefault) {
       super(data, implementedField, new ReferenceExpression(data, functionReference), null, isDefault);
+      this.functionReference = functionReference;
     }
 
     public CoClauseFunctionReference(Referable implementedField, TCDefReferable functionReference, boolean isDefault) {

diff --git a/base/src/main/java/org/arend/typechecking/visitor/DefinitionTypechecker.java b/base/src/main/java/org/arend/typechecking/visitor/DefinitionTypechecker.java
@@ -1212,23 +1212,31 @@ private boolean typecheckFunctionHeader(FunctionDefinition typedDef, Concrete.Ba
   }
 
   private Pair<Expression,ClassCallExpression> typecheckCoClauses(FunctionDefinition typedDef, Concrete.BaseFunctionDefinition def, FunctionKind kind, List<Concrete.CoClauseElement> elements) {
-    List<Concrete.Argument> arguments = new ArrayList<>();
-    for (Concrete.Parameter parameter : def.getParameters()) {
-      for (Referable referable : parameter.getReferableList()) {
-        arguments.add(new Concrete.Argument(new Concrete.ReferenceExpression(def.getData(), referable), false));
-      }
-    }
-
     List<Concrete.ClassFieldImpl> classFieldImpls = new ArrayList<>(elements.size());
     for (Concrete.CoClauseElement element : elements) {
       if (element instanceof Concrete.ClassFieldImpl) {
         classFieldImpls.add((Concrete.ClassFieldImpl) element);
       } else {
         throw new IllegalStateException();
       }
-      if (element instanceof Concrete.CoClauseFunctionReference && !def.getParameters().isEmpty()) {
-        Concrete.Expression oldImpl = ((Concrete.CoClauseFunctionReference) element).implementation;
-        ((Concrete.CoClauseFunctionReference) element).implementation = Concrete.AppExpression.make(oldImpl.getData(), oldImpl, arguments);
+      if (element instanceof Concrete.CoClauseFunctionReference coClauseRef && !def.getParameters().isEmpty()) {
+        Set<Pair<TCDefReferable,Integer>> allowedExternalParameters = new HashSet<>();
+        if (coClauseRef.functionReference.getTypechecked() instanceof TopLevelDefinition topLevel) {
+          allowedExternalParameters.addAll(topLevel.getParametersOriginalDefinitions());
+        }
+
+        List<Concrete.Argument> arguments = new ArrayList<>();
+        int index = 0;
+        for (Concrete.Parameter parameter : def.getParameters()) {
+          for (Referable referable : parameter.getReferableList()) {
+            var externalVars = typedDef.getParametersOriginalDefinitions();
+            if (index >= externalVars.size() || allowedExternalParameters.contains(externalVars.get(index))) {
+              arguments.add(new Concrete.Argument(new Concrete.ReferenceExpression(def.getData(), referable), false));
+            }
+            index++;
+          }
+        }
+        coClauseRef.implementation = Concrete.AppExpression.make(coClauseRef.implementation.getData(), coClauseRef.implementation, arguments);
       }
     }
 

diff --git a/src/test/java/org/arend/typechecking/VarsTest.java b/src/test/java/org/arend/typechecking/VarsTest.java
@@ -880,4 +880,32 @@ public void funcElimScopeTest2() {
       }
       """);
   }
+
+  @Test
+  public void coclauseElimScopeTest() {
+    typeCheckModule("""
+      \\record R (x : Nat) (field : Nat -> Nat)
+      \\func f (m : Nat) => m \\where {
+        \\func g : R m \\cowith
+          | field (n : Nat) : Nat \\with {
+            | 0 => 0
+            | suc n => n
+          }
+      }
+      """);
+  }
+
+  @Test
+  public void coclauseElimScopeTest2() {
+    typeCheckModule("""
+      \\record R (x : Nat) (field : Nat -> Nat)
+      \\func f (m : Nat) => m \\where {
+        \\func g : R m \\cowith
+          | field (n : Nat) : Nat \\with {
+            | 0 => m
+            | suc n => n
+          }
+      }
+      """);
+  }
 }