refactor speedy.Compiler.freeVars to be functional

I was trying to track down what looks like a new source of flakiness in the ANF work (#6440), which manifested in a stack overflow on the recursion of the `go` function in speedy.Compiler. I spent a bit of time refactoring this to be functional in order to then trampoline it, only to realize afterwards that I had worked on `freeVars/go` whereas the one blowing up the stack was `validate/go`. Obviously, it did nothing to reduce the rate of flakiness on the stack overflow issue I was trying to address, as it's a separate function. So there isn't really a good reason for this change, except I had done it and I do believe the code looks marginally nicer, so might as well submit it. CHANGELOG_BEGIN CHANGELOG_END
digital-asset · Jun 30, 2020 · 6502a5c · 6502a5c
1 parent d3a69b3
commit 6502a5c
Showing 1 changed file with 21 additions and 38 deletions.
diff --git a/daml-lf/interpreter/src/main/scala/com/digitalasset/daml/lf/speedy/Compiler.scala b/daml-lf/interpreter/src/main/scala/com/digitalasset/daml/lf/speedy/Compiler.scala
@@ -1125,67 +1125,50 @@ private[lf] final case class Compiler(
     * The returned free variables are de bruijn indices
     * adjusted to the stack of the caller. */
   def freeVars(expr: SExpr, initiallyBound: Int): Set[Int] = {
-    var bound = initiallyBound
-    var free = Set.empty[Int]
-
-    def go(expr: SExpr): Unit =
+    def go(expr: SExpr, bound: Int): Set[Int] =
       expr match {
         case SEVar(i) =>
-          if (i > bound)
-            free += i - bound /* adjust to caller's environment */
-        case _: SEVal => ()
-        case _: SEBuiltin => ()
-        case _: SEValue => ()
-        case _: SEBuiltinRecursiveDefinition => ()
+          if (i > bound) Set(i - bound) else Set() /* adjust to caller's environment */
+        case _: SEVal => Set()
+        case _: SEBuiltin => Set()
+        case _: SEValue => Set()
+        case _: SEBuiltinRecursiveDefinition => Set()
         case SELocation(_, body) =>
-          go(body)
+          go(body, bound)
         case SEAppGeneral(fun, args) =>
-          go(fun)
-          args.foreach(go)
+          go(fun, bound) ++ args.map(go(_, bound)).reduce(_ ++ _)
         case SEAppAtomicFun(fun, args) =>
-          go(fun)
-          args.foreach(go)
+          go(fun, bound) ++ args.map(go(_, bound)).reduce(_ ++ _)
         case SEAppSaturatedBuiltinFun(_, args) =>
-          args.foreach(go)
+          args.map(go(_, bound)).reduce(_ ++ _)
         case SEAbs(n, body) =>
-          bound += n
-          go(body)
-          bound -= n
+          go(body, bound + n)
         case x: SELoc =>
           throw CompilationError(s"freeVars: unexpected SELoc: $x")
         case x: SEMakeClo =>
           throw CompilationError(s"freeVars: unexpected SEMakeClo: $x")
         case SECase(scrut, alts) =>
-          go(scrut)
-          alts.foreach {
-            case SCaseAlt(pat, body) =>
-              val n = patternNArgs(pat)
-              bound += n; go(body); bound -= n
-          }
+          go(scrut, bound) ++ alts
+            .map { case SCaseAlt(pat, body) => go(body, bound + patternNArgs(pat)) }
+            .reduce(_ ++ _)
         case SELet(bounds, body) =>
-          bounds.foreach { e =>
-            go(e)
-            bound += 1
-          }
-          go(body)
-          bound -= bounds.size
+          (bounds ++ Seq(body)).zipWithIndex
+            .map { case (expr, idx) => go(expr, bound + idx) }
+            .reduce(_ ++ _)
         case SECatch(body, handler, fin) =>
-          go(body)
-          go(handler)
-          go(fin)
+          go(body, bound) ++ go(handler, bound) ++ go(fin, bound)
         case SELabelClosure(_, expr) =>
-          go(expr)
+          go(expr, bound)
         case x: SEWronglyTypeContractId =>
           throw CompilationError(s"unexpected SEWronglyTypeContractId: $x")
         case x: SEImportValue =>
           throw CompilationError(s"unexpected SEImportValue: $x")
       }
-    go(expr)
-    free
+    go(expr, initiallyBound)
   }
 
   /** Validate variable references in a speedy expression */
-  // valiate that we correctly captured all free-variables, and so reference to them is
+  // validate that we correctly captured all free-variables, and so reference to them is
   // via the surrounding closure, instead of just finding them higher up on the stack
   def validate(expr0: SExpr): SExpr = {