Allow module self-reference (#3240)

* Closes #3031 --------- Co-authored-by: Jan Mas Rovira <[email protected]>
anoma · Dec 18, 2024 · cb10fe3 · cb10fe3
1 parent af19142
commit cb10fe3
Show file tree

Hide file tree

Showing 5 changed files with 103 additions and 17 deletions.
diff --git a/src/Juvix/Compiler/Concrete/Data/NameSpace.hs b/src/Juvix/Compiler/Concrete/Data/NameSpace.hs
@@ -36,6 +36,14 @@ type family NameSpaceEntryType s = res | res -> s where
   NameSpaceEntryType 'NameSpaceModules = ModuleSymbolEntry
   NameSpaceEntryType 'NameSpaceFixities = FixitySymbolEntry
 
+entryName :: forall ns. (SingI ns) => Lens' (NameSpaceEntryType ns) S.Name
+entryName = case sing :: SNameSpace ns of
+  SNameSpaceSymbols -> \f -> \case
+    PreSymbolAlias (Alias n) -> PreSymbolAlias . Alias <$> f n
+    PreSymbolFinal (SymbolEntry n) -> PreSymbolFinal . SymbolEntry <$> f n
+  SNameSpaceModules -> moduleEntry
+  SNameSpaceFixities -> fixityEntry
+
 exportNameSpace :: forall ns. (SingI ns) => Lens' ExportInfo (HashMap C.Symbol (NameSpaceEntryType ns))
 exportNameSpace = case sing :: SNameSpace ns of
   SNameSpaceSymbols -> exportSymbols

diff --git a/src/Juvix/Compiler/Concrete/Translation/FromParsed/Analysis/Scoping.hs b/src/Juvix/Compiler/Concrete/Translation/FromParsed/Analysis/Scoping.hs
@@ -696,6 +696,37 @@ getModuleExportInfo m = fromMaybeM err (gets (^? scoperModules . at (m ^. module
             <> ppTrace ms
         )
 
+lookupLocalSymbolAux ::
+  forall r.
+  (Members '[State ScoperState, State Scope, Output ModuleSymbolEntry, Output PreSymbolEntry, Output FixitySymbolEntry] r) =>
+  (S.WhyInScope -> Bool) ->
+  [Symbol] ->
+  Symbol ->
+  Sem r ()
+lookupLocalSymbolAux whyInScope modules final =
+  case modules of
+    [] ->
+      lookHere
+    p : ps -> do
+      entries <- gets (^.. scopeModuleSymbols . at p . _Just . symbolInfo . each)
+      let entries' = filter (whyInScope . (^. moduleEntry . S.nameWhyInScope)) entries
+      mapM_ (getModuleExportInfo >=> lookInExport final ps) entries'
+  where
+    lookHere :: Sem r ()
+    lookHere = do
+      let helper ::
+            forall ns r'.
+            (SingI ns, Members '[Output (NameSpaceEntryType ns), State Scope] r') =>
+            Proxy ns ->
+            Sem r' ()
+          helper Proxy = do
+            entries <- gets (^.. scopeNameSpace @ns . at final . _Just . symbolInfo . each)
+            let entries' = filter (whyInScope . (^. entryName . S.nameWhyInScope)) entries
+            mapM_ output entries'
+      helper (Proxy @'NameSpaceSymbols)
+      helper (Proxy @'NameSpaceModules)
+      helper (Proxy @'NameSpaceFixities)
+
 -- | Do not call directly. Looks for a symbol in (possibly) nested local modules
 lookupSymbolAux ::
   forall r.
@@ -707,22 +738,31 @@ lookupSymbolAux modules final = do
   hereOrInLocalModule
   importedTopModule
   where
-    hereOrInLocalModule :: Sem r () =
-      case modules of
-        [] -> do
-          let helper ::
-                forall ns r'.
-                (SingI ns, Members '[Output (NameSpaceEntryType ns), State Scope] r') =>
-                Proxy ns ->
-                Sem r' ()
-              helper Proxy =
-                gets (^.. scopeNameSpace @ns . at final . _Just . symbolInfo . each) >>= mapM_ output
-          helper (Proxy @'NameSpaceSymbols)
-          helper (Proxy @'NameSpaceModules)
-          helper (Proxy @'NameSpaceFixities)
-        p : ps ->
-          gets (^.. scopeModuleSymbols . at p . _Just . symbolInfo . each)
-            >>= mapM_ (getModuleExportInfo >=> lookInExport final ps)
+    hereOrInLocalModule :: Sem r ()
+    hereOrInLocalModule = do
+      path0 <- gets (^. scopePath)
+      let topPath = path0 ^. S.absTopModulePath
+          path1 = topPath ^. modulePathDir ++ [topPath ^. modulePathName]
+          path2 = path0 ^. S.absLocalPath
+          pref = commonPrefix path2 modules
+      when (isPrefixOf path1 modules) $ do
+        let modules' = drop (length path1) modules
+            pref' = commonPrefix path2 modules'
+        lookPrefix pref' path2 modules'
+      when (notNull pref) $
+        lookPrefix pref path2 modules
+      lookupLocalSymbolAux (const True) modules final
+
+    lookPrefix :: [Symbol] -> [Symbol] -> [Symbol] -> Sem r ()
+    lookPrefix pref path modules' = do
+      let prefLen = length pref
+          inheritDepth = length path - prefLen
+          modules'' = drop prefLen modules'
+      lookupLocalSymbolAux
+        (== iterate S.BecauseInherited S.BecauseDefined !! inheritDepth)
+        modules''
+        final
+
     importedTopModule :: Sem r ()
     importedTopModule = do
       tbl <- gets (^. scopeTopModules)

diff --git a/test/Compilation/Positive.hs b/test/Compilation/Positive.hs
@@ -520,5 +520,10 @@ tests =
         "Test089: Loop invariant code motion"
         $(mkRelDir ".")
         $(mkRelFile "test089.juvix")
-        $(mkRelFile "out/test089.out")
+        $(mkRelFile "out/test089.out"),
+      posTest
+        "Test090: Module self-reference"
+        $(mkRelDir ".")
+        $(mkRelFile "test090.juvix")
+        $(mkRelFile "out/test090.out")
     ]
diff --git a/tests/Compilation/positive/out/test090.out b/tests/Compilation/positive/out/test090.out
@@ -0,0 +1 @@
+31
diff --git a/tests/Compilation/positive/test090.juvix b/tests/Compilation/positive/test090.juvix
@@ -0,0 +1,32 @@
+-- Module self-reference
+module test090;
+
+import Stdlib.Data.Nat as Delta;
+
+x : Delta.Nat := 9;
+
+module Resource0;
+  x : Delta.Nat := 10;
+end;
+
+module Resource;
+  x : Delta.Nat := 1;
+
+  module Gamma;
+    x : Delta.Nat := 2;
+  end;
+
+  module Delta;
+    open Delta using {+; *};
+
+    x : Delta.Nat := Resource.x + Resource0.x;
+    a : Delta.Nat := x * Resource.Gamma.x + test090.x;
+  end;
+
+  open Resource.Delta;
+
+  a : Delta.Nat := a;
+end;
+
+-- result: 31
+main : Delta.Nat := Resource.a;