Fix instance axiom bug (#2439)

* Closes #2438

src/Juvix/Compiler/Core/Translation/FromInternal.hs

@@ -101,12 +101,8 @@ goModuleNoVisit ::
   Sem r ()
 goModuleNoVisit (Internal.ModuleIndex m) = do
   mapM_ goImport (m ^. Internal.moduleBody . Internal.moduleImports)
-  mapM_ go (m ^. Internal.moduleBody . Internal.moduleStatements)
+  mapM_ goMutualBlock (m ^. Internal.moduleBody . Internal.moduleStatements)
   where
-    go :: Internal.Statement -> Sem r ()
-    go = \case
-      Internal.StatementAxiom a -> goAxiomInductive a >> goAxiomDef a
-      Internal.StatementMutual f -> goMutualBlock f
     goImport :: Internal.Import -> Sem r ()
     goImport (Internal.Import i) = visit i
 
@@ -238,14 +234,17 @@ goMutualBlock (Internal.MutualBlock m) = preMutual m >>= goMutual
     preMutual :: NonEmpty Internal.MutualStatement -> Sem r PreMutual
     preMutual stmts = do
       let (inds, funs) = partition isInd (toList stmts)
-      -- inductives must be pre-registered first to avoid crashing on unknown
-      -- inductive types when pre-registering functions
+      -- types must be pre-registered first to avoid crashing on unknown types
+      -- when pre-registering functions/axioms
       execState (PreMutual [] []) $ mapM_ step (inds ++ funs)
       where
         isInd :: Internal.MutualStatement -> Bool
         isInd = \case
           Internal.StatementInductive {} -> True
           Internal.StatementFunction {} -> False
+          Internal.StatementAxiom Internal.AxiomDef {..}
+            | Internal.ExpressionUniverse {} <- _axiomType -> True
+            | otherwise -> False
 
         step :: Internal.MutualStatement -> Sem (State PreMutual ': r) ()
         step = \case
@@ -255,6 +254,9 @@ goMutualBlock (Internal.MutualBlock m) = preMutual m >>= goMutual
           Internal.StatementInductive i -> do
             p <- preInductiveDef i
             modify' (over preInductives (p :))
+          Internal.StatementAxiom a -> do
+            goAxiomInductive a
+            goAxiomDef a
 
     goMutual :: PreMutual -> Sem r ()
     goMutual PreMutual {..} = do

src/Juvix/Compiler/Internal/Data/InfoTable.hs

@@ -79,7 +79,7 @@ computeTable recurIntoImports (ModuleIndex m) = compute
     mutuals :: [MutualStatement]
     mutuals =
       [ d
-        | StatementMutual (MutualBlock b) <- ss,
+        | MutualBlock b <- ss,
           d <- toList b
       ]
 
@@ -119,7 +119,7 @@ computeTable recurIntoImports (ModuleIndex m) = compute
     _infoAxioms =
       HashMap.fromList
         [ (d ^. axiomName, AxiomInfo d)
-          | StatementAxiom d <- ss
+          | StatementAxiom d <- mutuals
         ]
 
     _infoInstances :: InstanceTable
@@ -137,7 +137,7 @@ computeTable recurIntoImports (ModuleIndex m) = compute
           | otherwise =
               Nothing
 
-    ss :: [Statement]
+    ss :: [MutualBlock]
     ss = m ^. moduleBody . moduleStatements
 
 lookupConstructor :: forall r. (Member (Reader InfoTable) r) => Name -> Sem r ConstructorInfo

src/Juvix/Compiler/Internal/Extra/Base.hs

@@ -183,6 +183,18 @@ instance HasExpressions MutualStatement where
   leafExpressions f = \case
     StatementFunction d -> StatementFunction <$> leafExpressions f d
     StatementInductive d -> StatementInductive <$> leafExpressions f d
+    StatementAxiom d -> StatementAxiom <$> leafExpressions f d
+
+instance HasExpressions AxiomDef where
+  leafExpressions f AxiomDef {..} = do
+    ty' <- leafExpressions f _axiomType
+    pure
+      AxiomDef
+        { _axiomType = ty',
+          _axiomName,
+          _axiomBuiltin,
+          _axiomPragmas
+        }
 
 instance HasExpressions InductiveParameter where
   leafExpressions _ param@InductiveParameter {} = do

src/Juvix/Compiler/Internal/Extra/DependencyBuilder.hs

@@ -89,7 +89,7 @@ goModuleNoVisited :: forall r. (Members '[Reader ExportsTable, State DependencyG
 goModuleNoVisited (ModuleIndex m) = do
   checkStartNode (m ^. moduleName)
   let b = m ^. moduleBody
-  mapM_ (goStatement (m ^. moduleName)) (b ^. moduleStatements)
+  mapM_ (goMutual (m ^. moduleName)) (b ^. moduleStatements)
   mapM_ goImport (b ^. moduleImports)
 
 goImport :: (Members '[Reader ExportsTable, State DependencyGraph, State StartNodes, Visit ModuleIndex] r) => Import -> Sem r ()
@@ -105,11 +105,6 @@ goPreModule m = do
   -- added from definitions in M to definitions in N)
   mapM_ goImport (b ^. moduleImports)
 
-goStatement :: forall r. (Members '[Reader ExportsTable, State DependencyGraph, State StartNodes] r) => Name -> Statement -> Sem r ()
-goStatement parentModule = \case
-  StatementAxiom ax -> goAxiom parentModule ax
-  StatementMutual f -> goMutual parentModule f
-
 goMutual :: forall r. (Members '[Reader ExportsTable, State DependencyGraph, State StartNodes] r) => Name -> MutualBlock -> Sem r ()
 goMutual parentModule (MutualBlock s) = mapM_ go s
   where
@@ -117,6 +112,7 @@ goMutual parentModule (MutualBlock s) = mapM_ go s
     go = \case
       StatementInductive i -> goInductive parentModule i
       StatementFunction i -> goTopFunctionDef parentModule i
+      StatementAxiom ax -> goAxiom parentModule ax
 
 goPreLetStatement ::
   forall r.

src/Juvix/Compiler/Internal/Language.hs

@@ -17,11 +17,11 @@ import Juvix.Data.Universe hiding (smallUniverse)
 import Juvix.Data.WithLoc
 import Juvix.Prelude
 
-type Module = Module' Statement
+type Module = Module' MutualBlock
 
 type PreModule = Module' PreStatement
 
-type ModuleBody = ModuleBody' Statement
+type ModuleBody = ModuleBody' MutualBlock
 
 type PreModuleBody = ModuleBody' PreStatement
 
@@ -52,14 +52,10 @@ data ModuleBody' stmt = ModuleBody
   }
   deriving stock (Data)
 
-data Statement
-  = StatementMutual MutualBlock
-  | StatementAxiom AxiomDef
-  deriving stock (Data)
-
 data MutualStatement
   = StatementInductive InductiveDef
   | StatementFunction FunctionDef
+  | StatementAxiom AxiomDef
   deriving stock (Generic, Data)
 
 newtype MutualBlock = MutualBlock

src/Juvix/Compiler/Internal/Pretty/Base.hs

@@ -275,6 +275,7 @@ instance PrettyCode MutualStatement where
   ppCode = \case
     StatementInductive d -> ppCode d
     StatementFunction d -> ppCode d
+    StatementAxiom d -> ppCode d
 
 instance PrettyCode MutualBlock where
   ppCode (MutualBlock funs) = ppMutual funs
@@ -283,11 +284,6 @@ instance PrettyCode MutualBlockLet where
   ppCode (MutualBlockLet funs) =
     vsep2 <$> mapM ppCode funs
 
-instance PrettyCode Statement where
-  ppCode = \case
-    StatementMutual f -> ppCode f
-    StatementAxiom f -> ppCode f
-
 instance PrettyCode ModuleBody where
   ppCode m = do
     includes <- mapM ppCode (m ^. moduleImports)

src/Juvix/Compiler/Internal/Translation/FromConcrete.hs

@@ -260,28 +260,28 @@ fromPreModuleBody b = do
   let moduleStatements' = map goSCC sccs
   return (set Internal.moduleStatements moduleStatements' b)
   where
-    goSCC :: SCC Internal.PreStatement -> Internal.Statement
+    goSCC :: SCC Internal.PreStatement -> Internal.MutualBlock
     goSCC = \case
       AcyclicSCC s -> goAcyclic s
       CyclicSCC c -> goCyclic (nonEmpty' c)
       where
-        goCyclic :: NonEmpty Internal.PreStatement -> Internal.Statement
-        goCyclic c = Internal.StatementMutual (Internal.MutualBlock (goMutual <$> c))
+        goCyclic :: NonEmpty Internal.PreStatement -> Internal.MutualBlock
+        goCyclic c = Internal.MutualBlock (goMutual <$> c)
           where
             goMutual :: Internal.PreStatement -> Internal.MutualStatement
             goMutual = \case
               Internal.PreInductiveDef i -> Internal.StatementInductive i
               Internal.PreFunctionDef i -> Internal.StatementFunction i
-              _ -> impossible
+              Internal.PreAxiomDef i -> Internal.StatementAxiom i
 
-        goAcyclic :: Internal.PreStatement -> Internal.Statement
+        goAcyclic :: Internal.PreStatement -> Internal.MutualBlock
         goAcyclic = \case
           Internal.PreInductiveDef i -> one (Internal.StatementInductive i)
           Internal.PreFunctionDef i -> one (Internal.StatementFunction i)
-          Internal.PreAxiomDef i -> Internal.StatementAxiom i
+          Internal.PreAxiomDef i -> one (Internal.StatementAxiom i)
           where
-            one :: Internal.MutualStatement -> Internal.Statement
-            one = Internal.StatementMutual . Internal.MutualBlock . pure
+            one :: Internal.MutualStatement -> Internal.MutualBlock
+            one = Internal.MutualBlock . pure
 
 goModuleBody ::
   forall r.

src/Juvix/Compiler/Internal/Translation/FromInternal/Analysis/ArityChecking/Checker.hs

@@ -37,7 +37,7 @@ checkModuleBody ::
   ModuleBody ->
   Sem r ModuleBody
 checkModuleBody ModuleBody {..} = do
-  _moduleStatements' <- mapM checkStatement _moduleStatements
+  _moduleStatements' <- mapM checkMutualBlock _moduleStatements
   _moduleImports' <- mapM checkImport _moduleImports
   return
     ModuleBody
@@ -57,14 +57,6 @@ checkImport ::
   Sem r Import
 checkImport = traverseOf importModule checkModuleIndex
 
-checkStatement ::
-  (Members '[Reader InfoTable, NameIdGen, Error ArityCheckerError] r) =>
-  Statement ->
-  Sem r Statement
-checkStatement s = case s of
-  StatementMutual b -> StatementMutual <$> checkMutualBlock b
-  StatementAxiom a -> StatementAxiom <$> checkAxiom a
-
 checkInductive :: forall r. (Members '[Reader InfoTable, NameIdGen, Error ArityCheckerError] r) => InductiveDef -> Sem r InductiveDef
 checkInductive d = do
   let _inductiveName = d ^. inductiveName
@@ -111,6 +103,7 @@ checkMutualStatement ::
 checkMutualStatement = \case
   StatementFunction f -> StatementFunction <$> checkFunctionDef f
   StatementInductive f -> StatementInductive <$> checkInductive f
+  StatementAxiom a -> StatementAxiom <$> checkAxiom a
 
 checkMutualBlockLet ::
   (Members '[Reader InfoTable, NameIdGen, Error ArityCheckerError] r) =>

src/Juvix/Compiler/Internal/Translation/FromInternal/Analysis/Reachability.hs

@@ -41,7 +41,7 @@ goModuleNoCache (ModuleIndex m) = do
   where
     goBody :: ModuleBody -> Sem r ModuleBody
     goBody body = do
-      _moduleStatements <- mapMaybeM goStatement (body ^. moduleStatements)
+      _moduleStatements <- mapMaybeM goMutual (body ^. moduleStatements)
       _moduleImports <- mapM goImport (body ^. moduleImports)
       return ModuleBody {..}
 
@@ -51,16 +51,12 @@ goModule = cacheGet . ModuleIndex
 goModuleIndex :: (Members '[Reader NameDependencyInfo, MCache] r) => ModuleIndex -> Sem r ModuleIndex
 goModuleIndex = fmap ModuleIndex . cacheGet
 
-goStatement :: forall r. (Member (Reader NameDependencyInfo) r) => Statement -> Sem r (Maybe Statement)
-goStatement s = case s of
-  StatementMutual m -> fmap StatementMutual <$> goMutual m
-  StatementAxiom ax -> returnIfReachable (ax ^. axiomName) s
-  where
-    -- note that the first mutual statement is reachable iff all are reachable
-    goMutual :: MutualBlock -> Sem r (Maybe MutualBlock)
-    goMutual b@(MutualBlock (m :| _)) = case m of
-      StatementFunction f -> returnIfReachable (f ^. funDefName) b
-      StatementInductive f -> returnIfReachable (f ^. inductiveName) b
+-- note that the first mutual statement is reachable iff all are reachable
+goMutual :: forall r. (Member (Reader NameDependencyInfo) r) => MutualBlock -> Sem r (Maybe MutualBlock)
+goMutual b@(MutualBlock (m :| _)) = case m of
+  StatementFunction f -> returnIfReachable (f ^. funDefName) b
+  StatementInductive f -> returnIfReachable (f ^. inductiveName) b
+  StatementAxiom ax -> returnIfReachable (ax ^. axiomName) b
 
 goImport :: forall r. (Members '[Reader NameDependencyInfo, MCache] r) => Import -> Sem r Import
 goImport i = do

src/Juvix/Compiler/Internal/Translation/FromInternal/Analysis/Termination/Checker.hs

@@ -113,12 +113,7 @@ scanModule ::
 scanModule m = scanModuleBody (m ^. moduleBody)
 
 scanModuleBody :: (Members '[State CallMap] r) => ModuleBody -> Sem r ()
-scanModuleBody body = mapM_ scanStatement (body ^. moduleStatements)
-
-scanStatement :: (Members '[State CallMap] r) => Statement -> Sem r ()
-scanStatement = \case
-  StatementAxiom a -> scanAxiom a
-  StatementMutual m -> scanMutual m
+scanModuleBody body = mapM_ scanMutual (body ^. moduleStatements)
 
 scanMutual :: (Members '[State CallMap] r) => MutualBlock -> Sem r ()
 scanMutual (MutualBlock ss) = mapM_ scanMutualStatement ss
@@ -135,6 +130,7 @@ scanMutualStatement :: (Members '[State CallMap] r) => MutualStatement -> Sem r
 scanMutualStatement = \case
   StatementInductive i -> scanInductive i
   StatementFunction i -> scanFunctionDef i
+  StatementAxiom a -> scanAxiom a
 
 scanAxiom :: (Members '[State CallMap] r) => AxiomDef -> Sem r ()
 scanAxiom = scanTopExpression . (^. axiomType)

src/Juvix/Compiler/Internal/Translation/FromInternal/Analysis/TypeChecking/Checker.hs

@@ -72,7 +72,7 @@ checkModuleBody ::
   Sem r ModuleBody
 checkModuleBody ModuleBody {..} = do
   _moduleImports' <- mapM checkImport _moduleImports
-  _moduleStatements' <- mapM checkStatement _moduleStatements
+  _moduleStatements' <- mapM checkMutualBlock _moduleStatements
   return
     ModuleBody
       { _moduleStatements = _moduleStatements',
@@ -85,15 +85,11 @@ checkImport ::
   Sem r Import
 checkImport = traverseOf importModule checkModuleIndex
 
-checkStatement ::
+checkMutualBlock ::
   (Members '[HighlightBuilder, Reader EntryPoint, State NegativeTypeParameters, Reader InfoTable, Error TypeCheckerError, NameIdGen, State TypesTable, State FunctionsTable, Output Example, Builtins, Termination] r) =>
-  Statement ->
-  Sem r Statement
-checkStatement s = case s of
-  StatementMutual mut -> StatementMutual <$> runReader emptyLocalVars (checkTopMutualBlock mut)
-  StatementAxiom ax -> do
-    registerNameIdType (ax ^. axiomName . nameId) (ax ^. axiomType)
-    return s
+  MutualBlock ->
+  Sem r MutualBlock
+checkMutualBlock s = runReader emptyLocalVars (checkTopMutualBlock s)
 
 checkInductiveDef ::
   forall r.
@@ -175,6 +171,9 @@ checkMutualStatement ::
 checkMutualStatement = \case
   StatementFunction f -> StatementFunction <$> resolveInstanceHoles (checkFunctionDef f)
   StatementInductive f -> StatementInductive <$> resolveInstanceHoles (checkInductiveDef f)
+  StatementAxiom ax -> do
+    registerNameIdType (ax ^. axiomName . nameId) (ax ^. axiomType)
+    return $ StatementAxiom ax
 
 checkFunctionDef ::
   (Members '[HighlightBuilder, Reader LocalVars, Reader InfoTable, Error TypeCheckerError, NameIdGen, State TypesTable, State FunctionsTable, Output Example, Builtins, Inference, Termination, Output TypedHole] r) =>

test/Reachability/Positive.hs

@@ -46,14 +46,14 @@ getNames m =
   concatMap getDeclName (m ^. Internal.moduleBody . Internal.moduleStatements)
     <> concatMap (getNames . (^. Internal.importModule . Internal.moduleIxModule)) (m ^. Internal.moduleBody . Internal.moduleImports)
   where
-    getDeclName :: Internal.Statement -> [Text]
+    getDeclName :: Internal.MutualBlock -> [Text]
     getDeclName = \case
-      Internal.StatementMutual (Internal.MutualBlock f) -> map getMutualName (toList f)
-      Internal.StatementAxiom ax -> [ax ^. (Internal.axiomName . Internal.nameText)]
+      (Internal.MutualBlock f) -> map getMutualName (toList f)
     getMutualName :: Internal.MutualStatement -> Text
     getMutualName = \case
       Internal.StatementFunction f -> f ^. Internal.funDefName . Internal.nameText
       Internal.StatementInductive f -> f ^. Internal.inductiveName . Internal.nameText
+      Internal.StatementAxiom ax -> ax ^. (Internal.axiomName . Internal.nameText)
 
 allTests :: TestTree
 allTests =

test/Typecheck/Positive.hs

@@ -296,7 +296,11 @@ tests =
     posTest
       "Hole in type parameter"
       $(mkRelDir ".")
-      $(mkRelFile "HoleTypeParameter.juvix")
+      $(mkRelFile "HoleTypeParameter.juvix"),
+    posTest
+      "Instance axiom"
+      $(mkRelDir ".")
+      $(mkRelFile "InstanceAxiom.juvix")
   ]
     <> [ compilationTest t | t <- Compilation.tests
        ]

tests/positive/InstanceAxiom.juvix

@@ -0,0 +1,9 @@
+module InstanceAxiom;
+
+trait
+type T := t;
+
+axiom <body> : T;
+
+instance
+inst : T := <body>;