From 13501295deca09f0b1d6bf3fbbbe27514559063d Mon Sep 17 00:00:00 2001
From: valis <valery.isaev@gmail.com>
Date: Tue, 7 May 2024 16:58:22 +0300
Subject: [PATCH] Fix a bug with external variables and coclauses

---
 .../org/arend/term/concrete/Concrete.java     |  3 ++
 .../visitor/DefinitionTypechecker.java        | 28 ++++++++++++-------
 .../java/org/arend/typechecking/VarsTest.java | 28 +++++++++++++++++++
 3 files changed, 49 insertions(+), 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
index d04ff3f73..e17b711a1 100644
--- 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
index 3150b14d9..db0ebc5ed 100644
--- a/base/src/main/java/org/arend/typechecking/visitor/DefinitionTypechecker.java
+++ b/base/src/main/java/org/arend/typechecking/visitor/DefinitionTypechecker.java
@@ -1212,13 +1212,6 @@ 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) {
@@ -1226,9 +1219,24 @@ private Pair<Expression,ClassCallExpression> typecheckCoClauses(FunctionDefiniti
       } 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
index 7a61128ba..a33eeca3c 100644
--- 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
+          }
+      }
+      """);
+  }
 }