Fix a bug with eliminated external variables

Fixes #339
JetBrains · May 7, 2024 · 8264da5 · 8264da5
1 parent 38fe801
commit 8264da5
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 11 deletions.
diff --git a/base/src/main/java/org/arend/typechecking/visitor/WhereVarsFixVisitor.java b/base/src/main/java/org/arend/typechecking/visitor/WhereVarsFixVisitor.java
@@ -155,17 +155,6 @@ record ParamData(TCDefReferable definition, Concrete.Parameter parameter, int in
           selfArgsMap.put(definition.getData(), selfArgs);
         }
 
-        if (!parametersOriginalDefinitions.isEmpty() && definition instanceof Concrete.BaseFunctionDefinition && !definition.getParameters().isEmpty()) {
-          Concrete.FunctionBody body = ((Concrete.BaseFunctionDefinition) definition).getBody();
-          if (body instanceof Concrete.ElimFunctionBody && body.getEliminatedReferences().isEmpty()) {
-            for (Concrete.Parameter parameter : definition.getParameters()) {
-              for (Referable referable : parameter.getReferableList()) {
-                ((Concrete.ElimFunctionBody) body).getEliminatedReferences().add(new Concrete.ReferenceExpression(definition.getData(), referable));
-              }
-            }
-          }
-        }
-
         Set<Referable> levelRefs = new HashSet<>();
         for (Concrete.Parameter param : newParams) {
           if (param.getType() != null) {
@@ -227,6 +216,22 @@ record ParamData(TCDefReferable definition, Concrete.Parameter parameter, int in
         }
         varsFixVisitor.visitParameters(pair.proj1, null);
         definition.addParameters(pair.proj1, pair.proj2);
+
+        if (definition instanceof Concrete.BaseFunctionDefinition) {
+          Concrete.FunctionBody body = ((Concrete.BaseFunctionDefinition) definition).getBody();
+          if (body instanceof Concrete.ElimFunctionBody && body.getEliminatedReferences().isEmpty()) {
+            for (Concrete.FunctionClause clause : body.getClauses()) {
+              addParametersToClause(pair.proj1, clause);
+            }
+          }
+        }
+
+        if (definition instanceof Concrete.DataDefinition dataDef && dataDef.getEliminatedReferences() != null && dataDef.getEliminatedReferences().isEmpty()) {
+          for (Concrete.ConstructorClause clause : dataDef.getConstructorClauses()) {
+            addParametersToClause(pair.proj1, clause);
+          }
+        }
+
         if (definition instanceof Concrete.CoClauseFunctionDefinition coClauseDef) {
           int n = coClauseDef.getNumberOfExternalParameters();
           for (Concrete.Parameter parameter : pair.proj1) {
@@ -239,6 +244,16 @@ record ParamData(TCDefReferable definition, Concrete.Parameter parameter, int in
     }
   }
 
+  private static void addParametersToClause(List<Concrete.Parameter> newParameters, Concrete.Clause clause) {
+    List<Concrete.Pattern> newPatterns = new ArrayList<>();
+    for (Concrete.Parameter parameter : newParameters) {
+      for (Referable referable : parameter.getReferableList()) {
+        newPatterns.add(new Concrete.NamePattern(clause.getData(), false, referable, null));
+      }
+    }
+    clause.getPatterns().addAll(0, newPatterns);
+  }
+
   private static void checkLevels(Set<TCDefReferable> defs, Concrete.LevelParameters parameters, ErrorReporter errorReporter, String kind, Concrete.SourceNode sourceNode) {
     if (defs.size() > 1 || !defs.isEmpty() && parameters != null) {
       errorReporter.report(new TypecheckingError("Definition refers to different " + kind + "-levels", parameters != null ? parameters : sourceNode));

diff --git a/src/test/java/org/arend/typechecking/VarsTest.java b/src/test/java/org/arend/typechecking/VarsTest.java
@@ -848,4 +848,36 @@ public void appTest() {
           \\func test => xxx {0}
       """);
   }
+
+  @Test
+  public void dataElimScopeTest() {
+    typeCheckModule("""
+      \\func f {X : \\Type} => X \\where {
+        \\data D (n : Nat) \\with
+          | 0 => cons X
+      }
+      """);
+  }
+
+  @Test
+  public void funcElimScopeTest() {
+    typeCheckModule("""
+      \\func f {m : Nat} => m \\where {
+        \\func g (n : Nat) : Nat \\with
+          | 0 => m
+          | suc n => n
+      }
+      """);
+  }
+
+  @Test
+  public void funcElimScopeTest2() {
+    typeCheckModule("""
+      \\func f {m : Nat} => m \\where {
+        \\func g {n : Nat} : Nat \\with
+          | {0} => m
+          | {suc n} => n
+      }
+      """);
+  }
 }