Improve module method checks

Author: FlandiaYingman

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

@@ -1092,13 +1092,14 @@ extends Importer:
       N
     case N => N
 
-  def fieldOrVarSym(k: TermDefKind, id: Ident)(using Ctx): LocalSymbol & NamedSymbol =
+  def fieldOrVarSym(k: TermDefKind, id: Ident)(using Ctx): TermSymbol | VarSymbol =
     if ctx.outer.isDefined then TermSymbol(k, ctx.outer, id)
     else VarSymbol(id)
 
-  def param(t: Tree, inUsing: Bool): Ctxl[Opt[Opt[Bool] -> Param]] = t match
+  def param(t: Tree, inUsing: Bool): Ctxl[Opt[Opt[Bool] -> Param]] =
+    def go(t: Tree, inUsing: Bool, flags: FldFlags): Ctxl[Opt[Opt[Bool] -> Param]] = t match
     case TypeDef(Mod, inner, N, N) =>
-      val ps = param(inner, inUsing).map(_.mapSecond(p => p.copy(flags = p.flags.copy(mod = true))))
+      val ps = go(inner, inUsing, flags.copy(mod = true))
       for p <- ps if p._2.flags.mod do p._2.sign match
         case N =>
           raise(ErrorReport(msg"Module parameters must have explicit types." -> t.toLoc :: Nil))
@@ -1107,10 +1108,21 @@ extends Importer:
         case _ => ()
       ps
     case TypeDef(Pat, inner, N, N) =>
-      param(inner, inUsing).map(_.mapSecond(p => p.copy(flags = p.flags.copy(pat = true))))
+      go(inner, inUsing, flags.copy(pat = true))
     case _ =>
       t.asParam(inUsing).map: (isSpd, p, t) =>
-        isSpd -> Param(FldFlags.empty, fieldOrVarSym(ParamBind, p), t.map(term(_)))
+        val sym = fieldOrVarSym(ParamBind, p)
+        val sign = t.map(term(_))
+        val param = Param(flags, sym, sign)
+        sym match
+          case sym: TermSymbol =>
+            // TODO: How can a TermSymbol accept a Declaration
+            // sym.defn = S(TermDefinition(ctx.outer, Fun, sym, Nil, Nil, sign, N, FlowSymbol(s"‹result of ${sym}›"), TermDefFlags.empty, Nil))
+          case sym: VarSymbol =>
+            sym.decl = S(param)
+        isSpd -> param
+    go(t, inUsing, FldFlags.empty)
+      
 
   def params(t: Tree): Ctxl[(ParamList, Ctx)] = t match
     case Tup(ps) =>

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

@@ -282,21 +282,50 @@ object ModuleChecker:
     .exists(_.isInstanceOf[TyParam])
 
   /** Checks if a term evaluates to a module value. */
-  def evalsToModule(t: Term): Bool = 
-    def isModule(t: Tree): Bool = t match
-      case Tree.TypeDef(Mod, _, _, _) => true
-      case _ => false
+  def evalsToModule(t: Term): Bool =
     def returnsModule(t: Tree.TermDef): Bool = t.annotatedResultType match
       case S(Tree.TypeDef(Mod, _, N, N)) => true
       case _ => false
+    def checkDecl(decl: Declaration): Bool = decl match
+      // All TypeLikeDef are not modules, except for modules themselves.
+      // Objects use ModuleDef but is not a module.
+      case ModuleDef(kind = Mod) =>
+        true
+      case _: TypeLikeDef =>
+        false
+      // Check Member/Local symbols
+      case defn: TermDefinition => 
+        defn.flags.isModTyped
+      case defn: Param =>
+        defn.flags.mod
+      case defn: TyParam =>
+        defn.flags.mod
+    def checkSym(sym: Symbol): Bool = sym match
+      case sym if sym.asMod.nonEmpty => true
+      case sym if sym.asBlkMember.flatMap(_.trmTree).exists(returnsModule) => true
+      case _: (BuiltinSymbol | TopLevelSymbol) => false
+      case sym: BlockLocalSymbol => sym.decl match
+        case S(decl) => checkDecl(decl)
+        case N =>
+          // Most local symbols are let-bindings 
+          // which do not have a definition at this point.
+          false
+      case sym: MemberSymbol[?] => sym.defn match
+        case S(defn) => checkDecl(defn)
+        case N =>
+          // At this point all member symbols should have definition,
+          // except for the class(-like) that are currently being elaborated.
+          // TODO: We will fix this by deferring the checks to the resolution stage.
+          false
+      case sym => 
+        lastWords(s"Unsupported symbol kind ${sym}")
     t match
       case Term.Blk(_, res) => evalsToModule(res)
       case Term.App(lhs, rhs) => lhs.symbol match
         case S(sym: BlockMemberSymbol) => sym.trmTree.exists(returnsModule)
         case _ => false
-      case t => t.symbol match
-        case S(sym: BlockMemberSymbol) => sym.modTree.exists(isModule)
-        case _ => false
+      case t: Term.Ref => checkSym(t.sym)
+      case t => t.symbol.exists(checkSym)
 
   /**
    * An extractor that extracts the (tree) definition of a module method.

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

@@ -263,14 +263,14 @@ final case class LetDecl(sym: LocalSymbol, annotations: Ls[Annot]) extends State
 
 final case class DefineVar(sym: LocalSymbol, rhs: Term) extends Statement
 
-final case class TermDefFlags(isModMember: Bool):
+final case class TermDefFlags(isModMember: Bool, isModTyped: Bool):
   def showDbg: Str = 
     val flags = Buffer.empty[String]
     if isModMember then flags += "module"
     flags.mkString(" ")
   override def toString: String = "‹" + showDbg + "›"
 
-object TermDefFlags { val empty: TermDefFlags = TermDefFlags(false) }
+object TermDefFlags { val empty: TermDefFlags = TermDefFlags(false, false) }
 
 final case class TermDefinition(
     owner: Opt[InnerSymbol],
@@ -494,7 +494,7 @@ final case class TyParam(flags: FldFlags, vce: Opt[Bool], sym: VarSymbol) extend
 
 
 final case class Param(flags: FldFlags, sym: LocalSymbol & NamedSymbol, sign: Opt[Term]) 
-extends AutoLocated:
+extends Declaration with AutoLocated:
   def subTerms: Ls[Term] = sign.toList
   override protected def children: List[Located] = subTerms
   // def children: Ls[Located] = self.value :: self.asc.toList ::: Nil

hkmc2/shared/src/main/scala/hkmc2/utils/utils.scala

@@ -22,6 +22,7 @@ extension (s: String)
 
 import hkmc2.semantics.TermDefFlags
 import hkmc2.semantics.FldFlags
+import hkmc2.semantics.ParamListFlags
 import scala.collection.mutable.Buffer
 import mlscript.utils.StringOps
 import hkmc2.semantics.CtxArg
@@ -42,9 +43,10 @@ extension (t: Product)
       case xs: List[_] => "Ls of \n" + xs.iterator.map(aux(_)).mkString("\n").indent("  ")
       case xs: Vector[_] => "Vector of \n" + xs.iterator.map(aux(_)).mkString("\n").indent("  ")
       case s: String => s.escaped
-      case TermDefFlags(mod) =>
+      case TermDefFlags(isModMember, isModTyped) =>
         val flags = Buffer.empty[String]
-        if mod then flags += "module"
+        if isModMember then flags += "modMember"
+        if isModMember then flags += "modTyped"
         flags.mkString("(", ", ", ")")
       case FldFlags(mut, spec, genGetter, mod, pat) =>
         val flags = Buffer.empty[String]
@@ -54,6 +56,10 @@ extension (t: Product)
         if mod then flags += "module"
         if pat then flags += "pat"
         flags.mkString("(", ", ", ")")
+      case ParamListFlags(ctx) =>
+        val flags = Buffer.empty[String]
+        if ctx then flags += "ctx"
+        flags.mkString("(", ", ", ")")
       case Loc(start, end, origin) =>
         val (sl, _, sc) = origin.fph.getLineColAt(start)
         val (el, _, ec) = origin.fph.getLineColAt(end)

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

@@ -68,14 +68,16 @@ fun f7(): module M
 //│ ╙──      	    ^^^^^^^^^^^^^^
 
 
-:todo // should be an error
 :e
 fun f8(module m: M) = m
+//│ ╔══[ERROR] Functions returning module values must have explicit return types.
+//│ ║  l.72: 	fun f8(module m: M) = m
+//│ ╙──      	    ^^^^^^^^^^^^^^^
 
 :e
 fun f9(module m: M): module M = m
 //│ ╔══[ERROR] Only module methods may return module values.
-//│ ║  l.76: 	fun f9(module m: M): module M = m
+//│ ║  l.78: 	fun f9(module m: M): module M = m
 //│ ╙──      	    ^^^^^^^^^^^^^^^^^^^^^^^^^
 
 module Test with

hkmc2/shared/src/test/mlscript/codegen/FieldSymbols.mls

@@ -60,7 +60,7 @@ case
 //│   stats = Nil
 //│   res = Lam:
 //│     params = ParamList:
-//│       flags = ParamListFlags of false
+//│       flags = ()
 //│       params = Ls of 
 //│         Param:
 //│           flags = ()
@@ -92,7 +92,7 @@ case
 //│     lhs = $block$res
 //│     rhs = Lam:
 //│       params = ParamList:
-//│         flags = ParamListFlags of false
+//│         flags = ()
 //│         params = Ls of 
 //│           Param:
 //│             flags = ()

hkmc2/shared/src/test/mlscript/codegen/Hygiene.mls

@@ -156,10 +156,10 @@ foo()
 
 :sjs
 module Whoops with
-  val v = this
+  val v: module Whoops = this
   fun f() = "Hello"
   module Whoops with
-    val w = this
+    val w: module Whoops = this
     fun g() = f()
 //│ JS (unsanitized):
 //│ let Whoops2;

hkmc2/shared/src/test/mlscript/codegen/SelfReferences.mls

@@ -20,10 +20,12 @@ Foo.self
 //│ = Foo
 
 
-:todo
 :e
 module Foo with
   val self = this
+//│ ╔══[ERROR] Functions returning module values must have explicit return types.
+//│ ║  l.25: 	  val self = this
+//│ ╙──      	      ^^^^
 
 Foo.self
 //│ = Foo

hkmc2/shared/src/test/mlscript/ucs/papers/OperatorSplit.mls

@@ -75,7 +75,7 @@ fun example(args) =
 //│       sym = member:example
 //│       params = Ls of 
 //│         ParamList:
-//│           flags = ParamListFlags of false
+//│           flags = ()
 //│           params = Ls of 
 //│             Param:
 //│               flags = ()

hkmc2/shared/src/test/mlscript/ucs/syntax/NestedOpSplits.mls

@@ -16,7 +16,7 @@ fun f(x) =
 //│       sym = member:f
 //│       params = Ls of 
 //│         ParamList:
-//│           flags = ParamListFlags of false
+//│           flags = ()
 //│           params = Ls of 
 //│             Param:
 //│               flags = ()