Stack safe recursion

core/shared/main/scala/utils/package.scala

@@ -226,7 +226,7 @@ package object utils {
   def TODO(msg: Any): Nothing = throw new NotImplementedError(
     msg.toString + s" (of class ${msg.getClass().getSimpleName()})")
   def TODO(msg: Any, cond: Bool): Unit = if (cond) TODO(msg)
-  def die: Nothing = lastWords("Program reached and unexpected state.")
+  def die: Nothing = lastWords("Program reached an unexpected state.")
   def lastWords(msg: String): Nothing = throw new Exception(s"Internal Error: $msg")
   def wat(msg: String, wat: Any): Nothing = lastWords(s"$msg ($wat)")
 

hkmc2/jvm/src/test/scala/hkmc2/JSBackendDiffMaker.scala

@@ -14,6 +14,7 @@ import utils.Scope
 import hkmc2.syntax.Tree.Ident
 import hkmc2.codegen.Path
 import hkmc2.Diagnostic.Source
+import hkmc2.syntax.Keyword.`then`
 
 abstract class JSBackendDiffMaker extends MLsDiffMaker:
 
@@ -27,6 +28,8 @@ abstract class JSBackendDiffMaker extends MLsDiffMaker:
   val handler = NullaryCommand("handler")
   val expect = Command("expect"): ln =>
     ln.trim
+  val stackSafe = Command("stackSafe"): ln =>
+    ln.trim
 
   private val baseScp: utils.Scope =
     utils.Scope.empty
@@ -49,21 +52,35 @@ abstract class JSBackendDiffMaker extends MLsDiffMaker:
   override def run(): Unit =
     try super.run() finally if hostCreated then host.terminate()
 
+  private val DEFAULT_STACK_LIMT = 500
+
   override def processTerm(blk: semantics.Term.Blk, inImport: Bool)(using Raise): Unit =
     super.processTerm(blk, inImport)
     val outerRaise: Raise = summon
     var showingJSYieldedCompileError = false
+    val stackLimit = stackSafe.get.map(_.toIntOption) match
+      case None => None
+      case Some(value) => value match
+        case None => Some(DEFAULT_STACK_LIMT)
+        case Some(value) =>
+          if value < 0 then
+            failures += 1
+            output("/!\\ Stack limit must be positive, but the stack limit here is set to " + value)
+            Some(DEFAULT_STACK_LIMT)
+          else
+            Some(value)
     if showJS.isSet then
       given Raise =
         case d @ ErrorReport(source = Source.Compilation) =>
           showingJSYieldedCompileError = true
           outerRaise(d)
         case d => outerRaise(d)
+      given Elaborator.Ctx = curCtx
       val low = ltl.givenIn:
         new codegen.Lowering
           with codegen.LoweringSelSanityChecks(instrument = false)
           with codegen.LoweringTraceLog(instrument = false)
-      given Elaborator.Ctx = curCtx
+          with codegen.LoweringHandler(handler.isSet, stackLimit)
       val jsb = new JSBuilder
         with JSBuilderArgNumSanityChecks(instrument = false)
       val le = low.program(blk)
@@ -74,12 +91,12 @@ abstract class JSBackendDiffMaker extends MLsDiffMaker:
       output(s"JS (unsanitized):")
       output(jsStr)
     if js.isSet && !showingJSYieldedCompileError then
+      given Elaborator.Ctx = curCtx
       val low = ltl.givenIn:
         new codegen.Lowering
           with codegen.LoweringSelSanityChecks(noSanityCheck.isUnset)
           with codegen.LoweringTraceLog(traceJS.isSet)
-          with codegen.LoweringHandler(handler.isSet)
-      given Elaborator.Ctx = curCtx
+          with codegen.LoweringHandler(handler.isSet, stackLimit)
       val jsb = new JSBuilder
         with JSBuilderArgNumSanityChecks(noSanityCheck.isUnset)
       val le = low.program(blk)

hkmc2/shared/src/main/scala/hkmc2/codegen/Block.scala

@@ -12,6 +12,8 @@ import hkmc2.semantics.{Term => st}
 import syntax.{Literal, Tree}
 import semantics.*
 import semantics.Term.*
+import scala.reflect.ClassTag
+import sem.Elaborator.State
 
 
 case class Program(
@@ -36,13 +38,13 @@ sealed abstract class Block extends Product with AutoLocated:
     case Break(_) => Set.empty
     case Continue(_) => Set.empty
     case Define(defn, rst) => rst.definedVars
-    case HandleBlock(lhs, res, cls, hdr, bod, rst) => bod.definedVars ++ rst.definedVars + lhs
+    case HandleBlock(lhs, res, par, cls, hdr, bod, rst) => bod.definedVars ++ rst.definedVars + lhs
     case HandleBlockReturn(_) => Set.empty
     case TryBlock(sub, fin, rst) => sub.definedVars ++ fin.definedVars ++ rst.definedVars
     case Label(lbl, bod, rst) => bod.definedVars ++ rst.definedVars
 
   lazy val size: Int = this match
-    case _: Return | _: Throw | _: End | _: Break | _: Continue => 1
+    case _: Return | _: Throw | _: End | _: Break | _: Continue | _: HandleBlockReturn => 1
     case Begin(sub, rst) => sub.size + rst.size
     case Assign(_, _, rst) => 1 + rst.size
     case AssignField(_, _, _, rst) => 1 + rst.size
@@ -51,15 +53,142 @@ sealed abstract class Block extends Product with AutoLocated:
     case Define(_, rst) => 1 + rst.size
     case TryBlock(sub, fin, rst) => 1 + sub.size + fin.size + rst.size
     case Label(_, bod, rst) => 1 + bod.size + rst.size
+    case HandleBlock(lhs, res, par, cls, handlers, bdy, rst) => 1 + handlers.map(_.body.size).sum + bdy.size + rst.size
+
+  // ignoring blocks inside functions and handle block
+  def map(f: Block => Block): Block = this match
+    case _: Return | _: Throw | _: End | _: Break | _: Continue | _: HandleBlockReturn => this
+    case Match(scrut, arms, dflt, rst) => Match(scrut, arms.map(_ -> f(_)), dflt.map(f), f(rst))
+    case Label(lbl, bod, rst) => Label(lbl, f(bod), f(rst))
+    case Begin(sub, rst) => Begin(f(sub), f(rst))
+    case TryBlock(sub, fin, rst) => TryBlock(f(sub), f(fin), f(rst))
+    case Assign(l, r, rst) => Assign(l, r, f(rst))
+    case b @ AssignField(l, n, r, rst) => AssignField(l, n, r, f(rst))(b.symbol)
+    case Define(defn, rst) => Define(defn, f(rst))
+    case HandleBlock(l, res, par, cls, hdr, bod, rst) => HandleBlock(l, res, par, cls, hdr, bod, f(rst))
+
+  def mapResult(f: Result => Opt[(Result => Block) => Block]): Block = this match
+    case Return(res, implct) => f(res).map(_(Return(_, implct))).getOrElse(this)
+    case Throw(exc) => f(exc).map(_(Throw(_))).getOrElse(this)
+    case Assign(l, r, rst) => f(r).map(_(Assign(l, _, rst))).getOrElse(this)
+    case b @ AssignField(l, n, r, rst) => f(r).map(_(AssignField(l, n, _, rst)(b.symbol))).getOrElse(this)
+    case HandleBlockReturn(res) => f(res).map(_(HandleBlockReturn(_))).getOrElse(this)
+    case _: End | _: Break | _: Continue | _: Match | _: Label | _: Begin | _: TryBlock | _: Define | _: HandleBlock | _: HandleBlockReturn => this
+
+  def mapPath(f: Path => Path): Block = this match
+    case Return(res: Path, implct) => Return(f(res), implct)
+    case Throw(exc: Path) => Throw(f(exc))
+    case Assign(l, r: Path, rst) => Assign(l, f(r), rst)
+    case b @ AssignField(l, n, r: Path, rst) => AssignField(l, n, f(r), rst)(b.symbol)
+    case Match(scrut: Path, arms, dflt, rst) => Match(f(scrut), arms, dflt, rst)
+    case Define(ValDefn(owner, k, sym, rhs: Path), rst) => Define(ValDefn(owner, k, sym, f(rhs)), rst)
+    case HandleBlockReturn(res: Path) => HandleBlockReturn(f(res))
+    case _: Return | _: Throw | _: Assign | _: AssignField | _: End | _: Break | _: Continue | _: Match | _: Label | _: Begin | _: TryBlock | _: Define | _: HandleBlock | _: HandleBlockReturn => this.mapResult {
+      case r @ Call(fun, args) => S(_(Call(f(fun), args.map(arg => Arg(arg.spread, f(arg.value))))(r.isMlsFun)))
+      case r @ Instantiate(cls, args) => S(_(Instantiate(f(cls), args.map(f))))
+      case r => N
+    }
+
+  def mapSyms(using State, SymbolSubst): Block =
+    def pMap(p: Param) = 
+      val newSym: LocalSymbol & NamedSymbol = p.sym.subst
+      Param(p.flags, p.sym.subst, p.sign)
+    def pListMap(p: ParamList) = ParamList(p.flags, p.params.map(pMap), p.restParam.map(pMap))
+    def pathMap(p: Path): Path = p match
+      case s @ Select(qual, name) => Select(pathMap(qual), name)(s.symbol.map(_.subst))
+      case Value.Ref(l) => Value.Ref(l.subst)
+      case Value.This(sym) => Value.This(sym.subst)
+      case Value.Lit(lit) => p
+      case Value.Lam(params, body) => Value.Lam(
+        pListMap(params),
+        body.mapSyms
+      )
+      case Value.Arr(elems) => Value.Arr(elems.map(a => Arg(a.spread, pathMap(a.value))))
+
+    def _resMap(r: Result) = r match
+      case p: Path => pathMap(p)
+      case c @ Call(fun, args) => Call(pathMap(fun), args.map(a => Arg(a.spread, pathMap(a.value))))(c.isMlsFun)
+      case Instantiate(cls, args) => Instantiate(pathMap(cls), args.map(pathMap))
+
+    def resMap[T <: Result : ClassTag](l: T): T = _resMap(l) match
+      case r : T => r
+      case _ => l
+
+    this match
+      case Return(res, implct) => Return(resMap(res), implct)
+      case Throw(exc: Path) => Throw(pathMap(exc))
+      case Assign(l, r, rst) => Assign(l.subst, resMap(r), rst.mapSyms)
+      case blk @ AssignField(l, n, r, rst) => AssignField(pathMap(l), n, _resMap(r), rst.mapSyms)(blk.symbol)
+      case Match(scrut, arms, dflt, rst) => 
+        val newArms = arms.map((cse, blk) =>
+          val newCse = cse match
+            case Case.Lit(lit) => cse
+            // HACK: cls.subst gives something other than ClassSymbol | ModuleSymbol and it doesn't type check...
+            case Case.Cls(cls, path) => cls match 
+              case s: ClassSymbol => Case.Cls(s.subst, pathMap(path))
+              case s: ModuleSymbol => Case.Cls(s.subst, pathMap(path))
+            case Case.Tup(len, inf) => cse
+          (newCse, blk.mapSyms)
+        )
+        Match(pathMap(scrut), newArms, dflt.map(_.mapSyms), rst.mapSyms)
+      case Define(d, rst) => 
+        val newDefn = d match
+          case FunDefn(sym, params, body) => 
+            FunDefn(sym.subst, params.map(pListMap), body.mapSyms)
+          case ValDefn(owner, k, sym, rhs) => 
+            ValDefn(owner.map(_.subst), k, sym.subst, pathMap(rhs))
+          case ClsLikeDefn(sym, k, parentPath, methods, privateFields, publicFields, preCtor, ctor) =>
+            ClsLikeDefn(
+              sym.subst, k, parentPath.map(pathMap),
+              methods.map(fDef => FunDefn(fDef.sym.subst, fDef.params.map(pListMap), fDef.body.mapSyms)),
+              privateFields.map(_.subst),
+              publicFields.map(t => TermDefinition(
+                t.owner.map(_.subst),
+                t.k,
+                t.sym.subst,
+                t.params.map(pListMap),
+                t.sign,
+                t.body,
+                t.resSym.subst,
+                t.flags,
+                t.annotations,
+              )),
+              preCtor.mapSyms,
+              ctor.mapSyms
+            )
+        Define(newDefn, rst.mapSyms)
+
+      case HandleBlockReturn(res: Path) => HandleBlockReturn(pathMap(res))
+      case Throw(res) => Throw(resMap(res))
+      case Label(lbl, body, rest) => Label(lbl.subst, body.mapSyms, rest.mapSyms)
+      case HandleBlockReturn(res) => HandleBlockReturn(resMap(res))
+      case Begin(sub, rest) => Begin(sub.mapSyms, rest.mapSyms)
+      case TryBlock(sub, fin, rest) => TryBlock(sub.mapSyms, fin.mapSyms, rest.mapSyms)
+      case HandleBlock(lhs, res, par, cls, handlers, body, rest) =>
+        HandleBlock(
+          lhs.subst, res.subst, pathMap(par), cls.subst, 
+          handlers.map(h => Handler(h.sym.subst, h.resumeSym.subst, h.params.map(pListMap), body.mapSyms)),
+          body.mapSyms,
+          rest.mapSyms
+        )
+      case Break(s) => Break(s.subst)
+      case Continue(s) => Continue(s.subst)
+      case End(_) => this
+
+  def mapValue(f: Value => Value): Block =
+    def go(p: Path): Path = p match
+      case sel: Select => Select(go(sel.qual), sel.name)(sel.symbol)
+      case v: Value => f(v)
+    this.mapPath(go)
 
   // TODO conserve if no changes
   def mapTail(f: BlockTail => BlockTail): Block = this match
     case b: BlockTail => f(b)
     case Begin(sub, rst) => Begin(sub, rst.mapTail(f))
     case Assign(lhs, rhs, rst) => Assign(lhs, rhs, rst.mapTail(f))
     case Define(defn, rst) => Define(defn, rst.mapTail(f))
-    case HandleBlock(lhs, res, cls, handlers, body, rest) =>
-      HandleBlock(lhs, res, cls, handlers.map(h => Handler(h.sym, h.resumeSym, h.params, h.body.mapTail(f))), body.mapTail(f), rest.mapTail(f))
+    case HandleBlock(lhs, res, par, cls, handlers, body, rest) =>
+      HandleBlock(lhs, res, par, cls, handlers.map(h => Handler(h.sym, h.resumeSym, h.params, h.body.mapTail(f))), body.mapTail(f), rest.mapTail(f))
     case Match(scrut, arms, dflt, rst: End) =>
       Match(scrut, arms.map(_ -> _.mapTail(f)), dflt.map(_.mapTail(f)), rst)
     case Match(scrut, arms, dflt, rst) =>
@@ -80,11 +209,81 @@ sealed abstract class Block extends Product with AutoLocated:
     case Assign(lhs, rhs, rest) => Set(lhs) ++ rhs.freeVars ++ rest.freeVars
     case AssignField(lhs, nme, rhs, rest) => lhs.freeVars ++ rhs.freeVars ++ rest.freeVars
     case Define(defn, rest) => defn.freeVars ++ rest.freeVars
-    case HandleBlock(lhs, res, cls, handlers, body, rest) =>
-      (body.freeVars - lhs) ++ rest.freeVars ++ handlers.flatMap(_.freeVars)
+    case HandleBlock(lhs, res, par, cls, hdr, bod, rst) =>
+      (bod.freeVars - lhs) ++ rst.freeVars ++ hdr.flatMap(_.freeVars)
     case HandleBlockReturn(res) => res.freeVars
     case End(msg) => Set.empty
 
+  def floatOutDefns(outerOnly: Bool) =
+    def rec(b: Block, acc: List[Defn]): (Block, List[Defn]) =
+      b match
+        case Match(scrut, arms, dflt, rest) =>
+          if outerOnly then
+            val (rstRes, rstDefns) = rec(rest, acc)
+            (Match(scrut, arms, dflt, rstRes), rstDefns)
+          else
+            val (armsRes, armsDefns) = arms.foldLeft[(List[(Case, Block)], List[Defn])](Nil, acc)(
+              (accc, d) =>
+                val (accCases, accDefns) = accc
+                val (cse, blk) = d
+                val (resBlk, resDefns) = rec(blk, accDefns)
+                ((cse, resBlk) :: accCases, resDefns)
+            )
+            dflt match
+              case None =>
+                val (rstRes, rstDefns) = rec(rest, armsDefns)
+                (Match(scrut, armsRes, None, rstRes), rstDefns)
+
+              case Some(dflt) =>
+                val (dfltRes, dfltDefns) = rec(dflt, armsDefns)
+                val (rstRes, rstDefns) = rec(rest, dfltDefns)
+                (Match(scrut, armsRes, S(dfltRes), rstRes), rstDefns)
+
+        case Return(res, implct) => (b, acc)
+        case Throw(exc) => (b, acc)
+        case Label(label, body, rest) =>
+          if outerOnly then
+            val (rstRes, rstDefns) = rec(rest, acc)
+            (Label(label, body, rstRes), rstDefns)
+          else
+            val (bodyRes, bodyDefns) = rec(body, acc)
+            val (rstRes, rstDefns) = rec(rest, bodyDefns)
+            (Label(label, bodyRes, rstRes), rstDefns)
+        case Break(label) => (b, acc)
+        case Continue(label) => (b, acc)
+        case Begin(sub, rest) => 
+          val (subRes, subDefns) = rec(sub, acc)
+          val (rstRes, rstDefns) = rec(rest, subDefns)
+          (Begin(subRes, rstRes), rstDefns)
+        case TryBlock(sub, finallyDo, rest) =>
+          if outerOnly then
+            val (rstRes, rstDefns) = rec(rest, acc)
+            (TryBlock(sub, finallyDo, rstRes), rstDefns)
+          else
+            val (subRes, subDefns) = rec(sub, acc)
+            val (finallyRes, finallyDefns) = rec(rest, subDefns)
+            val (rstRes, rstDefns) = rec(rest, finallyDefns)
+            (TryBlock(subRes, finallyRes, rstRes), rstDefns)
+        case Assign(lhs, rhs, rest) => 
+          val (rstRes, rstDefns) = rec(rest, acc)
+          (Assign(lhs, rhs, rstRes), rstDefns)
+        case a @ AssignField(path, nme, result, rest) =>
+          val (rstRes, rstDefns) = rec(rest, acc)
+          (AssignField(path, nme, result, rstRes)(a.symbol), rstDefns)
+        case Define(defn, rest) => defn match
+          case ValDefn(owner, k, sym, rhs) => 
+            val (rstRes, rstDefns) = rec(rest, acc)
+            (Define(defn, rstRes), rstDefns)
+          case _ =>
+            val (rstRes, rstDefns) = rec(rest, defn :: acc)
+            (rstRes, rstDefns)
+        case HandleBlock(lhs, res, par, cls, handlers, body, rest) =>
+          val (rstRes, rstDefns) = rec(rest, acc)
+          (HandleBlock(lhs, res, par, cls, handlers, body, rstRes), rstDefns)
+        case HandleBlockReturn(res) => (b, acc)
+        case End(msg) => (b, acc)
+    rec(this, Nil)
+
 end Block
 
 sealed abstract class BlockTail extends Block
@@ -119,7 +318,7 @@ case class AssignField(lhs: Path, nme: Tree.Ident, rhs: Result, rest: Block)(val
 
 case class Define(defn: Defn, rest: Block) extends Block with ProductWithTail
 
-case class HandleBlock(lhs: Local, res: Local, cls: Path, handlers: Ls[Handler], body: Block, rest: Block) extends Block with ProductWithTail
+case class HandleBlock(lhs: Local, res: Local, par: Path, cls: ClassSymbol, handlers: Ls[Handler], body: Block, rest: Block) extends Block with ProductWithTail
 case class HandleBlockReturn(res: Result) extends BlockTail
 
 sealed abstract class Defn:
@@ -149,7 +348,7 @@ final case class ValDefn(
 final case class ClsLikeDefn(
   sym: MemberSymbol[? <: ClassLikeDef],
   k: syntax.ClsLikeKind,
-  parentSym: Opt[Path],
+  parentPath: Opt[Path],
   methods: Ls[FunDefn],
   privateFields: Ls[TermSymbol],
   publicFields: Ls[TermDefinition],