Enforce rules on functions returning modules

Author: FlandiaYingman

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

@@ -489,12 +489,40 @@ extends Importer:
                   case ((pss, ctx), ps) => 
                     val (qs, newCtx) = params(ps)(using ctx)
                     (pss :+ ParamList(ParamListFlags.empty, qs), newCtx)
+              // * Elaborate signature
+              val s = td.signature.orElse(newSignatureTrees.get(id.name)).map(term)
               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(mod))
+              val tdf = TermDefinition(owner, k, sym, pss, s, b, r, TermDefFlags(mod))
               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)
+
+              // checks rules regarding module methods
+              result match
+                case N if em => raise:
+                  ErrorReport:
+                    msg"Function returning module values must have explicit return types." ->
+                    td.head.toLoc :: Nil
+                case S(ret) if em && ModuleChecker.isTypeParam(ret) => raise:
+                    ErrorReport:
+                      msg"Function returning module values must have concrete return types." ->
+                      td.head.toLoc :: Nil
+                case S(ret) if em && !ret.isInstanceOf[Term.Mod] => raise:
+                  ErrorReport:
+                    msg"The return type of functions returning module values must be prefixed with module keyword." ->
+                    td.head.toLoc :: Nil
+                case S(Term.Mod(_)) if !mod => 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 +725,23 @@ 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: Tree): Bool = t match
+        case InfixApp(_, Keyword.`:`, rhs) => rhs match
+          case TypeDef(Mod, _, _, _) => true
+          case _ => false
+        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.head)
+            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

@@ -24,19 +24,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.25: 	  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 concrete return types.
+//│ ║  l.33: 	  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.41: 	  fun f5(): M = M
+//│ ╙──      	      ^^^^^^^
 
 :e
 fun f6(module m: M)
@@ -45,7 +51,9 @@ f6(new C)
 
 :e
 fun f7(): module M
-//│ FAILURE: Unexpected lack of type error
+//│ ╔══[ERROR] Only module methods may return module values.
+//│ ║  l.53: 	fun f7(): module M
+//│ ╙──      	    ^^^^^^^^^^^^^^
 
 
 fun ok1(module m: M)