Enforce rules on functions returning modules

Author: FlandiaYingman

hkmc2/shared/src/main/scala/hkmc2/semantics/Elaborator.scala

@@ -489,12 +489,47 @@ extends Importer:
                   case ((pss, ctx), ps) => 
                     val (qs, newCtx) = params(ps)(using ctx)
                     (pss :+ ParamList(ParamListFlags.empty, qs), newCtx)
+              // * Elaborate signature
+              val st = td.signature.orElse(newSignatureTrees.get(id.name))
+              val s = st.map(term(_)(using newCtx))
               val b = rhs.map(term(_)(using newCtx))
               val r = FlowSymbol(s"‹result of ${sym}›", nextUid)
-              val tdf = TermDefinition(owner, k, sym, pss,
-                td.signature.orElse(newSignatureTrees.get(id.name)).map(term), b, r,
-                TermDefFlags(isModMember))
+              val tdf = TermDefinition(owner, k, sym, pss, s, b, r, 
+                TermDefFlags.empty.copy(mod = isModMember))
               sym.defn = S(tdf)
+
+              // the return type of the function
+              val result = td.head match
+                case InfixApp(_, Keyword.`:`, rhs) => S(term(rhs)(using newCtx))
+                case _ => N
+              
+              // indicates if the function really returns a module
+              val em = b.fold(false)(ModuleChecker.evalsToModule)
+              // indicates if the function marks its result as "module"
+              val mm = st match
+                case Some(TypeDef(Mod, _, N, N)) => true
+                case _ => false
+
+              // checks rules regarding module methods
+              s match
+                case N if em => raise:
+                  ErrorReport:
+                    msg"Function returning module values must have explicit return types." ->
+                    td.head.toLoc :: Nil
+                case S(t) if em && ModuleChecker.isTypeParam(t) => raise:
+                    ErrorReport:
+                      msg"Function returning module values must have concrete return types." ->
+                      td.head.toLoc :: Nil
+                case S(_) if em && !mm => raise:
+                  ErrorReport:
+                    msg"The return type of functions returning module values must be prefixed with module keyword." ->
+                    td.head.toLoc :: Nil
+                case S(_) if mm && !isModMember => raise:
+                  ErrorReport:
+                    msg"Only module methods may return module values." ->
+                    td.head.toLoc :: Nil
+                case _ => ()
+              
               tdf
             go(sts, tdf :: acc)
           case L(d) =>
@@ -697,6 +732,21 @@ extends Importer:
       .filter(_.isInstanceOf[VarSymbol])
       .flatMap(_.asInstanceOf[VarSymbol].decl)
       .fold(false)(_.isInstanceOf[TyParam])
+    
+    /** Checks if a term evaluates to a module value. */
+    def evalsToModule(t: Term): Bool = 
+      def returnsModule(t: TermDef): Bool = t.signature match
+        case S(TypeDef(Mod, _, N, N)) => true
+        case _ => false
+      t match
+        case Term.Blk(_, res) => evalsToModule(res)
+        case Term.App(lhs, rhs) => evalsToModule(lhs)
+        case t => t.symbol match
+          case S(sym: BlockMemberSymbol) => sym match
+            case _ if sym.modTree.isDefined => true
+            case _ if sym.trmTree.isDefined => returnsModule(sym.trmTree.get)
+            case _ => false
+          case _ => false
 
   class VarianceTraverser(var changed: Bool = true) extends Traverser:
     override def traverseType(pol: Pol)(trm: Term): Unit = trm match

hkmc2/shared/src/test/mlscript/basics/ModuleMethods.mls

@@ -26,19 +26,25 @@ fun f2[T](module m: T)
 module N with {
   fun f3(): M = M
 }
-//│ FAILURE: Unexpected lack of type error
+//│ ╔══[ERROR] The return type of functions returning module values must be prefixed with module keyword.
+//│ ║  l.27: 	  fun f3(): M = M
+//│ ╙──      	      ^^^^^^^
 
 :e
 module N with {
   fun f4[T](): module T = M
 }
-//│ FAILURE: Unexpected lack of type error
+//│ ╔══[ERROR] Function returning module values must have explicit return types.
+//│ ║  l.35: 	  fun f4[T](): module T = M
+//│ ╙──      	      ^^^^^^^^^^^^^^^^^
 
 :e
 module N with {
   fun f5(): M = M
 }
-//│ FAILURE: Unexpected lack of type error
+//│ ╔══[ERROR] The return type of functions returning module values must be prefixed with module keyword.
+//│ ║  l.43: 	  fun f5(): M = M
+//│ ╙──      	      ^^^^^^^
 
 
 fun f6(m: M)
@@ -53,7 +59,9 @@ f6(M.self())
 
 :e
 fun f7(): module M
-//│ FAILURE: Unexpected lack of type error
+//│ ╔══[ERROR] Only module methods may return module values.
+//│ ║  l.61: 	fun f7(): module M
+//│ ╙──      	    ^^^^^^^^^^^^^^
 
 
 fun ok1(module m: M)