Enforce rules on functions returning modules

Author: FlandiaYingman

hkmc2/shared/src/main/scala/hkmc2/codegen/js/JSBuilder.scala

@@ -187,7 +187,7 @@ class JSBuilder extends CodeBuilder:
                     } + ")""""
                 }; }"""
               } #}  # }"
-            if clsDefn.kind is syntax.Mod then
+            if (clsDefn.kind is syntax.Mod) || (clsDefn.kind is syntax.Obj) then
               val clsTmp = summon[Scope].allocateName(new semantics.TempSymbol(0/*TODO rm this useless param*/, N, sym.nme+"$"+"class"))
               clsDefn.owner match
               case S(owner) =>

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

@@ -519,19 +519,54 @@ 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(isModMember = 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) =>
             raise(d)
             go(sts, acc)
       case (td @ TypeDef(k, head, extension, body)) :: sts =>
-        assert((k is Als) || (k is Cls) || (k is Mod), k)
+        assert((k is Als) || (k is Cls) || (k is Mod) || (k is Obj), k)
         val nme = td.name match
           case R(id) => id
           case L(d) =>
@@ -583,7 +618,7 @@ extends Importer:
               semantics.TypeDef(alsSym, tps, extension.map(term), N)
             alsSym.defn = S(d)
             d
-        case Mod =>
+        case k: (Mod.type | Obj.type) =>
           val clsSym = td.symbol.asInstanceOf[ModuleSymbol] // TODO: improve `asInstanceOf`
           val owner = ctx.outer
           newCtx.nest(S(clsSym)).givenIn:
@@ -594,7 +629,7 @@ extends Importer:
                 // case S(t) => block(t :: Nil)
                 case S(t) => ???
                 case N => (new Term.Blk(Nil, Term.Lit(UnitLit(true))), ctx)
-              ModuleDef(owner, clsSym, tps, ps, ObjBody(bod))
+              ModuleDef(owner, clsSym, tps, ps, k, ObjBody(bod))
             clsSym.defn = S(cd)
             cd
         case Cls =>
@@ -618,7 +653,6 @@ extends Importer:
         // TODO: pass abstract to `go`
         go(body :: sts, acc)
       case Modified(Keyword.`declare`, absLoc, body) :: sts =>
-        ???
         // TODO: pass declare to `go`
         go(body :: sts, acc)
       case (result: Tree) :: Nil =>
@@ -723,6 +757,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 isModule(t: Tree): Bool = t match
+        case TypeDef(Mod, _, _, _) => true
+        case _ => false
+      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) => lhs.symbol match
+          case S(sym: BlockMemberSymbol) => sym.trmTree.fold(false)(returnsModule)
+          case _ => false
+        case t => t.symbol match
+          case S(sym: BlockMemberSymbol) => sym.modTree.fold(false)(isModule)
+          case _ => false
 
   class VarianceTraverser(var changed: Bool = true) extends Traverser:
     override def traverseType(pol: Pol)(trm: Term): Unit = trm match

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

@@ -96,7 +96,7 @@ class BlockMemberSymbol(val nme: Str, val trees: Ls[Tree]) extends MemberSymbol[
   def clsTree: Opt[Tree.TypeDef] = trees.collectFirst:
     case t: Tree.TypeDef if t.k is Cls => t
   def modTree: Opt[Tree.TypeDef] = trees.collectFirst:
-    case t: Tree.TypeDef if t.k is Mod => t
+    case t: Tree.TypeDef if (t.k is Mod) || (t.k is Obj) => t
   def alsTree: Opt[Tree.TypeDef] = trees.collectFirst:
     case t: Tree.TypeDef if t.k is Als => t
   def trmTree: Opt[Tree.TermDef] = trees.collectFirst:

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

@@ -236,9 +236,14 @@ sealed abstract class ClassLikeDef extends TypeLikeDef:
   val body: ObjBody
 
 
-case class ModuleDef(owner: Opt[InnerSymbol], sym: ModuleSymbol, tparams: Ls[TyParam], paramsOpt: Opt[Ls[Param]], body: ObjBody) extends ClassLikeDef with Companion:
-  self =>
-  val kind: ClsLikeKind = Mod
+case class ModuleDef(
+  owner: Opt[InnerSymbol], 
+  sym: ModuleSymbol, 
+  tparams: Ls[TyParam], 
+  paramsOpt: Opt[Ls[Param]], 
+  kind: ClsLikeKind,
+  body: ObjBody,
+) extends ClassLikeDef with Companion
 
 
 sealed abstract class ClassDef extends ClassLikeDef:

hkmc2/shared/src/main/scala/hkmc2/syntax/Keyword.scala

@@ -94,6 +94,7 @@ object Keyword:
   val `new` = Keyword("new", N, curPrec) // TODO: check the prec
   // val `namespace` = Keyword("namespace", N, N)
   val `module` = Keyword("module", N, curPrec)
+  val `object` = Keyword("object", N, curPrec)
   val `open` = Keyword("open", N, curPrec)
   val `type` = Keyword("type", N, N)
   val `where` = Keyword("where", N, N)

hkmc2/shared/src/main/scala/hkmc2/syntax/ParseRule.scala

@@ -271,6 +271,7 @@ object ParseRule:
     Kw(`class`)(typeDeclBody(Cls)),
     Kw(`trait`)(typeDeclBody(Trt)),
     Kw(`module`)(typeDeclBody(Mod)),
+    Kw(`object`)(typeDeclBody(Obj)),
     Kw(`open`):
       ParseRule("'open' keyword")(
         exprOrBlk(ParseRule("'open' declaration")(End(()))){

hkmc2/shared/src/main/scala/hkmc2/syntax/Tree.scala

@@ -203,6 +203,7 @@ case object Trt extends TypeDefKind("trait") with ObjDefKind
 case object Mxn extends TypeDefKind("mixin")
 case object Als extends TypeDefKind("type alias")
 case object Mod extends TypeDefKind("module") with ClsLikeKind
+case object Obj extends TypeDefKind("object") with ClsLikeKind
 
 
 
@@ -275,7 +276,7 @@ trait TypeDefImpl extends TypeOrTermDef:
 
   lazy val symbol = k match
     case Cls => semantics.ClassSymbol(this, name.getOrElse(Ident("<error>")))
-    case Mod => semantics.ModuleSymbol(this, name.getOrElse(Ident("<error>")))
+    case Mod | Obj => semantics.ModuleSymbol(this, name.getOrElse(Ident("<error>")))
     case Als => semantics.TypeAliasSymbol(name.getOrElse(Ident("<error>")))
     case Trt | Mxn => ???
 

hkmc2/shared/src/test/mlscript-compile/Option.mls

@@ -9,7 +9,7 @@ open Predef
 module Option with ...
 
 class Some(value)
-module None
+object None
 
 fun isDefined(x) = if x is
   Some then true

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

@@ -20,5 +20,39 @@ fun f2[T](module m: T)
 //│ ║  l.18: 	fun f2[T](module m: T)
 //│ ╙──      	                 ^^^^
 
+:e
+module N with {
+  fun f3() = M
+}
+//│ ╔══[ERROR] Function returning module values must have explicit return types.
+//│ ║  l.25: 	  fun f3() = M
+//│ ╙──      	      ^^^^
+
+:e
+module N with {
+  fun f4[T](): module T = M
+}
+//│ ╔══[ERROR] Function returning module values must have explicit return types.
+//│ ║  l.33: 	  fun f4[T](): module T = M
+//│ ╙──      	      ^^^^^^^^^^^^^^^^^
+
+:e
+module N with {
+  fun f5(): M = M
+}
+//│ ╔══[ERROR] The return type of functions returning module values must be prefixed with module keyword.
+//│ ║  l.41: 	  fun f5(): M = M
+//│ ╙──      	      ^^^^^^^
+
+:e
+fun f7(): module M
+//│ ╔══[ERROR] Only module methods may return module values.
+//│ ║  l.48: 	fun f7(): module M
+//│ ╙──      	    ^^^^^^^^^^^^^^
+
 
 fun ok1(module m: M)
+
+module N with {
+  fun ok2(): module M = M
+}

hkmc2/shared/src/test/mlscript/bbml/bbBorrowing.mls

@@ -61,7 +61,7 @@ letreg of r =>
     123
     if next(it) > 0 then () => 0 else () => clear(b)
   k()
-//│ /!!!\ Uncaught error: scala.MatchError: Cons(Branch(Ref($scrut@102),Lit(BoolLit(true)),Else(Lam(List(),Lit(IntLit(0))))),Else(Lam(List(),App(Sel(Ref(globalThis:block#5),Ident(clear)),Tup(List(Fld(‹›,Ref(b@91),None))))))) (of class hkmc2.semantics.Split$Cons)
+//│ /!!!\ Uncaught error: scala.MatchError: Cons(Branch(Ref($scrut@123),Lit(BoolLit(true)),Else(Lam(List(),Lit(IntLit(0))))),Else(Lam(List(),App(Sel(Ref(globalThis:block#5),Ident(clear)),Tup(List(Fld(‹›,Ref(b@112),None))))))) (of class hkmc2.semantics.Split$Cons)
 
 :e
 letreg of r =>