Enforce functions with module parameters to have an explicit return type

Author: FlandiaYingman

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

@@ -516,12 +516,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

@@ -344,6 +344,9 @@ idMod(IntMT)
 // Error: functions taking module parameters must have an explicit result type.
 :e
 fun f3(x: M) = x
+//│ ╔══[ERROR] Functions with module parameters must have a explicit return type.
+//│ ║  l.346: 	fun f3(x: M) = x
+//│ ╙──       	    ^^^^^^^^
 //│ Elaborated tree:
 //│ Blk:
 //│   stats = Ls of 
@@ -366,7 +369,6 @@ fun f3(x: M) = x
 //│       resSym = ‹result of member:f3›@66
 //│       flags = ()
 //│   res = Lit of UnitLit of true
-//│ FAILURE: Unexpected lack of type error
 
 // OK: t is not a module parameter and doesn't need to have a explicit result type.
 fun f4(x: Int) = x