Experiment: use JSDoc types in JS runtime

effekt/shared/src/main/scala/effekt/generator/js/JavaScript.scala

@@ -78,7 +78,7 @@ class JavaScript(additionalFeatureFlags: List[String] = Nil) extends Compiler[St
    */
   lazy val CompileLSP = CPSTransformed map {
     case (mainSymbol, mainFile, core, cps) =>
-      TransformerCps.compileLSP(cps, core)
+      TransformerCps.compileLSP(cps, core, mainSymbol)
   }
 
   private def pretty(stmts: List[js.Stmt]): Document =

effekt/shared/src/main/scala/effekt/generator/js/PrettyPrinter.scala

@@ -62,7 +62,10 @@ object PrettyPrinter extends ParenPrettyPrinter {
     case Let(id, expr)                 => "let" <+> toDoc(id) <+> "=" <+> toDoc(expr) <> ";"
     case Destruct(ids, expr)           => "const" <+> braces(hsep(ids.map(toDoc), comma)) <+> "=" <+> toDoc(expr) <> ";"
     case Assign(target, expr)          => toDoc(target) <+> "=" <+> toDoc(expr) <> ";"
-    case Function(name, params, stmts) => "function" <+> toDoc(name) <> parens(params map toDoc) <+> jsBlock(stmts map toDoc)
+    case Function(name, params, stmts, None) =>
+      "function" <+> toDoc(name) <> parens(params map toDoc) <+> jsBlock(stmts map toDoc)
+    case Function(name, params, stmts, Some(docComment)) => toDoc(docComment) <> line <>
+      "function" <+> toDoc(name) <> parens(params map toDoc) <+> jsBlock(stmts map toDoc)
     case Class(name, methods)          => "class" <+> toDoc(name) <+> jsBlock(methods.map(jsMethod))
     case If(cond, thn, Block(Nil))     => "if" <+> parens(toDoc(cond)) <+> toDocBlock(thn)
     case If(cond, thn, els)            => "if" <+> parens(toDoc(cond)) <+> toDocBlock(thn) <+> "else" <+> toDocBlock(els)
@@ -78,6 +81,12 @@ object PrettyPrinter extends ParenPrettyPrinter {
     case Switch(sc, branches, default) => "switch" <+> parens(toDoc(sc)) <+> jsBlock(branches.map {
       case (tag, stmts) => "case" <+> toDoc(tag) <> ":" <+> nested(stmts map toDoc)
     } ++ default.toList.map { stmts => "default:" <+> nested(stmts map toDoc) })
+
+    case LineComment(contents) => "//" <+> contents
+    case DocComment(lines) =>
+      "/**" <> line <>
+      vcat(lines.map { l => " *" <+> l }) <> line <>
+      " */"
   }
 
   def toDocBlock(stmt: Stmt): Doc = stmt match {
@@ -87,7 +96,7 @@ object PrettyPrinter extends ParenPrettyPrinter {
   }
 
   def jsMethod(c: js.Function): Doc = c match {
-    case js.Function(name, params, stmts) =>
+    case js.Function(name, params, stmts, _docComment) =>
       toDoc(name) <> parens(params map toDoc) <+> jsBlock(stmts.map(toDoc))
   }
 

effekt/shared/src/main/scala/effekt/generator/js/TransformerCps.scala

@@ -36,6 +36,8 @@ object TransformerCps extends Transformer {
     directStyle: Option[ContinuationInfo],
     // the current direct-style metacontinuation
     metacont: Option[Id],
+    // the main symbol (entrypoint)
+    mainSymbol: symbols.TermSymbol,
     // the original declaration context (used to compile pattern matching)
     declarations: DeclarationContext,
     // the usual compiler context
@@ -52,9 +54,9 @@ object TransformerCps extends Transformer {
       js.Return(Call(RUN_TOPLEVEL, nameRef(mainSymbol))))))
 
     given DeclarationContext = new DeclarationContext(coreModule.declarations, coreModule.externs)
-    toJS(input, exports)
+    toJS(input, exports, mainSymbol)
 
-  def toJS(module: cps.ModuleDecl, exports: List[js.Export])(using D: DeclarationContext, C: Context): js.Module =
+  def toJS(module: cps.ModuleDecl, exports: List[js.Export], mainSymbol: symbols.TermSymbol)(using D: DeclarationContext, C: Context): js.Module =
     module match {
       case cps.ModuleDecl(path, includes, declarations, externs, definitions, _) =>
         given TransformerContext(
@@ -63,6 +65,7 @@ object TransformerCps extends Transformer {
           None,
           None,
           None,
+          mainSymbol,
           D, C)
 
         val name      = JSName(jsModuleName(module.path))
@@ -73,22 +76,23 @@ object TransformerCps extends Transformer {
         js.Module(name, Nil, exports, jsDecls ++ jsExterns ++ stmts)
     }
 
-  def compileLSP(input: cps.ModuleDecl, coreModule: core.ModuleDecl)(using C: Context): List[js.Stmt] =
+  def compileLSP(input: cps.ModuleDecl, coreModule: core.ModuleDecl, mainSymbol: symbols.TermSymbol)(using C: Context): List[js.Stmt] =
     val D = new DeclarationContext(coreModule.declarations, coreModule.externs)
     given TransformerContext(
           false,
           input.externs.collect { case d: Extern.Def => (d.id, d) }.toMap,
           None,
           None,
           None,
+          mainSymbol,
           D, C)
 
     input.definitions.map(toJS)
 
 
-  def toJS(d: cps.ToplevelDefinition)(using TransformerContext): js.Stmt = d match {
+  def toJS(d: cps.ToplevelDefinition)(using C: TransformerContext): js.Stmt = d match {
     case cps.ToplevelDefinition.Def(id, block) =>
-      js.Const(nameDef(id), requiringThunk { toJS(id, block) })
+      js.Const(nameDef(id), requiringThunk { toJS(id, block) }, isMainSymbol = C.mainSymbol == id)
     case cps.ToplevelDefinition.Val(id, ks, k, binding) =>
       js.Const(nameDef(id), Call(RUN_TOPLEVEL, js.Lambda(List(nameDef(ks), nameDef(k)), toJS(binding).stmts)))
     case cps.ToplevelDefinition.Let(id, binding) =>

effekt/shared/src/main/scala/effekt/generator/js/Tree.scala

@@ -31,6 +31,7 @@ case class Module(name: JSName, imports: List[Import], exports: List[Export], st
    * Generates the Javascript module skeleton for whole program compilation
    */
   def commonjs: List[Stmt] = {
+    val typecheckAnnotation = js.LineComment("@ts-check")
     val effekt = js.Const(JSName("$effekt"), js.Object())
 
     val importStmts = imports.map {
@@ -43,11 +44,12 @@ case class Module(name: JSName, imports: List[Import], exports: List[Export], st
         js.Destruct(names, js.Call(Variable(JSName("require")), List(JsString(s"./${ file }"))))
     }
 
+    val ignoreTypesOfExport = js.LineComment("@ts-ignore - Universal module pattern for Node/Browser compatibility")
     val exportStatement = js.Assign(RawExpr(s"(typeof module != \"undefined\" && module !== null ? module : {}).exports = ${name.name}"),
       js.Object(exports.map { e => e.name -> e.expr })
     )
 
-    List(effekt) ++ importStmts ++ stmts ++ List(exportStatement)
+    List(typecheckAnnotation, effekt) ++ importStmts ++ stmts ++ List(ignoreTypesOfExport, exportStatement)
   }
 
   /**
@@ -62,6 +64,7 @@ case class Module(name: JSName, imports: List[Import], exports: List[Export], st
    * }}}
    */
   def virtual : List[Stmt] = {
+    val typecheckAnnotation = js.LineComment("@ts-check")
     val effekt = js.Const(JSName("$effekt"), js.Object())
 
     val importStmts = imports.map {
@@ -78,7 +81,7 @@ case class Module(name: JSName, imports: List[Import], exports: List[Export], st
 
     // module.exports = { EXPORTS }
     val exportStatement = js.Assign(RawExpr("module.exports"), js.Object(exports.map { e => e.name -> e.expr }))
-    List(effekt) ++ importStmts ++ List(declaration) ++ stmts ++ List(exportStatement)
+    List(typecheckAnnotation, effekt) ++ importStmts ++ List(declaration) ++ stmts ++ List(exportStatement)
   }
 }
 
@@ -160,8 +163,14 @@ enum Stmt {
   // e.g. switch (sc) { case <EXPR>: <STMT>; ...; default: <STMT> }
   case Switch(scrutinee: Expr, branches: List[(Expr, List[Stmt])], default: Option[List[Stmt]]) // TODO maybe flatten?
 
-  // e.g. function <NAME>(x, y) { <STMT>* }
-  case Function(name: JSName, params: List[JSName], stmts: List[Stmt])
+  // e.g.
+  // ```js
+  // /**
+  //   * My doc comment
+  //   */
+  // function <NAME>(x, y) { <STMT>* }
+  // ```
+  case Function(name: JSName, params: List[JSName], stmts: List[Stmt], docComment: Option[Stmt.DocComment] = None)
 
   // e.g. class <NAME> {
   //        <NAME>(x, y) { <STMT>* }...
@@ -188,17 +197,37 @@ enum Stmt {
 
   // e.g. <EXPR>;
   case ExprStmt(expr: Expr)
+
+  // e.g. `// This is my comment`
+  case LineComment(contents: String)
+
+  // e.g.
+  //
+  // /**
+  //  * This is my
+  //  * comment
+  //  */
+  case DocComment(lines: List[String])
 }
 export Stmt.*
 
 
 // Smart constructors
 // ------------------
 
-def Const(name: JSName, binding: Expr): Stmt = binding match {
-  case Expr.Lambda(params, Block(stmts)) => js.Function(name, params, stmts)
-  case Expr.Lambda(params, stmt) => js.Function(name, params, List(stmt))
-  case _ => js.Const(Pattern.Variable(name), binding)
+def Const(name: JSName, binding: Expr, isMainSymbol: Boolean = false): Stmt = {
+  def docCommentFor(params: List[JSName]): Option[DocComment] = Option.when(isMainSymbol) {
+    params match {
+      case ks :: k :: Nil => DocComment(List(s"@param {MetaContinuation} ${ks.name}", s"@param {Continuation} ${k.name}"))
+      case _ => sys error s"Assumed that the JS entrypoint has exactly two params, but found ${params.length} instead"
+    }
+  }
+
+  binding match {
+    case Expr.Lambda(params, Block(stmts)) => js.Function(name, params, stmts, docCommentFor(params))
+    case Expr.Lambda(params, stmt) => js.Function(name, params, List(stmt), docCommentFor(params))
+    case _ => js.Const(Pattern.Variable(name), binding)
+  }
 }
 
 def Let(name: JSName, binding: Expr): Stmt = js.Let(Pattern.Variable(name), binding)

libraries/js/effekt_builtins.js

@@ -1,3 +1,7 @@
+/**
+ * @param {*} obj
+ * @returns {string}
+ */
 $effekt.show = function(obj) {
   if (!!obj && !!obj.__reflect) {
     const meta = obj.__reflect()
@@ -10,6 +14,11 @@ $effekt.show = function(obj) {
   }
 }
 
+/**
+ * @param {*} obj1
+ * @param {*} obj2
+ * @returns {boolean}
+ */
 $effekt.equals = function(obj1, obj2) {
   if (!!obj1.__equals) {
     return obj1.__equals(obj2)
@@ -18,12 +27,21 @@ $effekt.equals = function(obj1, obj2) {
   }
 }
 
+/**
+ * @param {*} n1
+ * @param {*} n2
+ */
 function compare$prim(n1, n2) {
   if (n1 == n2) { return 0; }
   else if (n1 > n2) { return 1; }
   else { return -1; }
 }
 
+/**
+ * @param {*} obj1
+ * @param {*} obj2
+ * @returns {-1 | 0 | 1} - -1 if obj1 < obj2, 0 if equal, 1 if obj1 > obj2
+ */
 $effekt.compare = function(obj1, obj2) {
   if ($effekt.equals(obj1, obj2)) { return 0; }
 
@@ -50,12 +68,34 @@ $effekt.compare = function(obj1, obj2) {
   return compare$prim(obj1, obj2);
 }
 
+/**
+ * @typedef {Object} Unit
+ * @property {true} __unit
+ */
+
+/**
+ * Unit singleton value (Effekt's `()`)
+ * @type {Unit}
+ */
+$effekt.unit = { __unit: true };
+
+/**
+ * @param {string} str
+ * @returns {Unit}
+ */
 $effekt.println = function println$impl(str) {
   console.log(str); return $effekt.unit;
 }
 
-$effekt.unit = { __unit: true }
-
+/**
+ * Throws an error for incomplete pattern matches
+ * @throws {Error}
+ */
 $effekt.emptyMatch = function() { throw "empty match" }
 
-$effekt.hole = function(pos) { throw pos + " not implemented yet" }
+/**
+ * Placeholder for unimplemented code
+ * @param {string} pos - Source position (already formatted)
+ * @throws {Error}
+ */
+$effekt.hole = function(pos) { throw pos + " not implemented yet" }