Enforce functions with module parameters to have an explicit return type

Author: FlandiaYingman

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

@@ -513,12 +513,29 @@ 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)
+
+              // Restrictions regarding functions receiving module parameters
+              
+              // not working because sometimes the signature is N 
+              // even when the function has a explicit return type
+              // val explicitReturnType = s.isDefined
+              
+              // FIXME: this is magic; we shall determine the return type based on the signature
+              val explicitReturnType = 
+                td.head.isInstanceOf[InfixApp] && td.head.asInstanceOf[InfixApp]._2 == Keyword.`:`
+              
+              if pss.flatMap(_.params).exists(_.flags.mod) && !explicitReturnType then
+                raise:
+                  ErrorReport:
+                    msg"Functions with module parameters must have a explicit return type." ->
+                    td.head.toLoc :: Nil
+
               tdf
             go(sts, tdf :: acc)
           case L(d) =>

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

@@ -387,6 +387,10 @@ idMod(IntM2).f(0) // OK
 
 // Error: functions taking module parameters must have an explicit result type.
 fun idMod[T](t: M2[T]) = t
+//│ FAILURE: Unexpected type error
+//│ ╔══[ERROR] Functions with module parameters must have a explicit return type.
+//│ ║  l.389: 	fun idMod[T](t: M2[T]) = t
+//│ ╙──       	    ^^^^^^^^^^^^^^^^^^
 //│ Elaborated tree:
 //│ Blk:
 //│   stats = Ls of