PseudoCodeStreams.scala

// staged stream emitter (w/o unimportant details), to show how we ought to
// do region management.

// "NEW" comments indicate where the implementation differs from the current
// region-unaware version.


def emit(dstR: Code[Region], ir: IR, env: Env): EmTr =
  ...

def deepCopy(dstR: Code[Region], t: EmTr): EmTr =
  ...

def emitS(ir: IR, env: Env): Stream[EmTr] = ir match {

  case StreamRange(lenIR) => new Stream[EmTr] {
    type S = (Code[Int], Code[Int])

    // invariant: strR will never be cleared for the duration of the stream
    def init(strR: Code[Region])(k: InitCont) = {
      val len = emit(strR, lenIR, env)
      len.m.mux(
        k(Missing),
        k(Start((len.v, 0))))
    }

    // dstR is only valid until this function returns; i.e. consumers may clear it
    // whenever they like
    def step(dstR: Code[Region], s: S)(k: StepCont) = {
      val (nLeft, i) = s
        (nLeft > 0).mux(
          k(Yield(present(i), (nLeft - 1, i + 1))),
          k(EOS))
    }
  }


  case ArrayMap(stream, name, bodyIR) => new Stream[EmTr] {
    val xv = newField
    val child = emitS(stream, env)
    type S = child.S

    def init(strR: Code[Region])(k: InitCont) =
      child.init(strR)(k)

    def step(dstR: Code[Region], s: S)(k: StepCont) =
      child.step(dstR, s) {
        case EOS => k(EOS)
        case Yield(x, s1) =>
          val body = emit(dstR, bodyIR, env.bind(name -> xv))
          Code(
            xv := x.v,
            k(Yield(present(body.v), s1)))
      }
  }


  case ArrayFilter(stream, name, condIR) => new Stream[EmTr] {
    val xv = newField
    val child = emitS(stream, env)
    type S = child.S

    def init(strR: Code[Region])(k: InitCont) =
      child.init(strR)(k)

    def step(dstR: Code[Region], s: S)(k: StepCont) = {
      val loop = joinPoint[S]()
      // invariant: dstR is empty at each entry into this joinpoint
      loop.define(child.step(dstR, _) {
        case EOS => k(EOS)
        case Yield(x, s1) =>
          val cond = emit(dstR, condIR, env.bind(name -> xv))
          Code(
            xv := x.v,
            cond.v.mux(
              k(Yield(present(xv), s1)),
              Code(
                dstR.clear(), // NEW
                loop(s1))))
      })
      loop(s)
    }
  }


  case ArrayFlatMap(outerStream, name, innerStream) => new Stream[EmTr] {
    val xv = newField
    val innR = newField[Region] // NEW
    val outer = emitS(outerStream, env)
    val inner = emitS(innerStream, env.bind(name -> xv))
    type S = (outer.S, inner.S)

    def init(strR: Code[Region])(k: InitCont) =
      outer.init(strR) {
        case Missing => k(Missing)
        case Start(outS) =>
          Code(
            innR := new Region(), // NEW
            k(Start((outS, inner.emptyS))))
      }

    def step(dstR: Code[Region], s: S)(k: StepCont) = {
      val innLoop = joinPoint[(outer.S, inner.S)]()
      val outLoop = joinPoint[outer.S]()

      innLoop.define { case (outS, innS) =>
        inner.step(dstR, innS) {
          case EOS => outLoop(outS)
          case Yield(y, innS1) => k(Yield(y, (outS, innS1)))
        }
      }

      outLoop.define { outS => Code(
        innR.clear(), // NEW
        outer.step(innR, outS) {
          case EOS => k(EOS)
          case Yield(x, outS1) =>
            Code(
              xv := x.v,
              inner.init(innR) {
                case Missing => outLoop(outS1)
                case Start(innS1) => innLoop((outS1, innS1))
              })
        })
      }

      innLoop(s)
    }
  }


  case ArrayScan(stream, zeroIR, eltName, accName, bodyIR) => new Stream[EmTr] {
    val xv = newField
    val av = newField
    val curR = newField[Region]
    val tmpR = newField[Region]
    val child = emitS(stream, env)
    type S = (child.S, Code[Accum], Code[Boolean])

    def init(strR: Code[Region])(k: InitCont) =
      child.init(strR) {
        case Missing => k(Missing)
        case Start(chS) =>
          val zero = emit(strR, zeroIR, env)
          Code(
            // NEW
            curR := new Region,
            tmpR := new Region,
            k(Start((chS, zero.v, true))))
      }

    def step(dstR: Code[Region], s: S)(k: StepCont) = {
      val (chS, acc, firstRun) = s
      firstRun.mux(
        k(Yield(deepCopy(dstR, acc), (chS, acc, false))),
        child.step(tmpR, chS) {
          case EOS => k(EOS)
          case Yield(x, chS1) =>
            val body = emit(tmpR, bodyIR, env.bind(eltName -> xv, accName -> av))
            val bv = newLocal
            Code(
              av := deepCopy(tmpR, acc), // NEW
              xv := x.v,
              bv := body.v,
              // NEW
              curR.clear(),
              swap(curR, tmpR),
              k(Yield(deepCopy(dstR, bv), (chS1, bv, false))))
        })
    }
  }

    //
}