LazyChunk.scala

//> using scala "3.7.1-RC2"

//> using dep io.getkyo::kyo-core:0.19.0+40-6e6cea21+20250603-1734-SNAPSHOT
//> using dep io.getkyo::kyo-prelude:0.19.0+40-6e6cea21+20250603-1734-SNAPSHOT
//> using dep io.getkyo::kyo-direct:0.19.0+40-6e6cea21+20250603-1734-SNAPSHOT

//> using dep org.scalatest::scalatest:3.2.19


//> using option -Wvalue-discard
//> using option -Wnonunit-statement
//> using option -Wconf:msg=(unused.*value|discarded.*value|pure.*statement):error
//> using option -language:strictEquality

import kyo.*

sealed trait LazyChunk[+V, -S] derives CanEqual:
  def map[V2](f: V => V2)(using Frame): LazyChunk[V2, S] =
    this match
      case LazyChunk.Empty => LazyChunk.Empty
      case LazyChunk.Cons(values, next) =>
        LazyChunk.Cons(values.map(f), defer(next.now.map(f)))
  end map

  def foldLeft[B, S2](z: B)(f: (B, V) => B < S2)(using Frame): B < (S2 & S) =
    Loop(this, z) {
      case (LazyChunk.Empty, z) => Loop.done(z)
      case (LazyChunk.Cons(values, next), z) =>
        val nextZ: B < S2 = Kyo.foldLeft(values)(z)(f)
        nextZ.map(z => next.map(lazyChunk => Loop.continue(lazyChunk,z)))
    }
  end foldLeft

  def foreach[S2](f: V => Unit < S2): Unit < (S2 & S) =
    foreachChunk(v => Kyo.foreachDiscard(v)(f))

  def foreachChunk[S2](f: Chunk[V] => Unit < S2): Unit < (S2 & S) =
    Loop(this):
      case LazyChunk.Empty => Loop.done
      case LazyChunk.Cons(values, next) =>
        f(values).andThen(next.map(Loop.continue))



  def existsChunk(f: Chunk[V] => Boolean): Boolean < S =
    Loop(this):
      case LazyChunk.Empty => Loop.done(false)
      case LazyChunk.Cons(values, _) if f(values) => Loop.done(true)
      case LazyChunk.Cons(_, next) => next.map(Loop.continue)


  def isEmpty: Boolean < S = existsChunk(_.nonEmpty).map(!_)


  def toStream[B >: V : Tag]: Stream[B, S] =
    Stream:
      foreachChunk(chunk => Emit.value(chunk))

object LazyChunk:
  case object Empty extends LazyChunk[Nothing, Any]

  case class Cons[V, S](values: Chunk[V], next: LazyChunk[V, S] < S) extends LazyChunk[V, S]

  def fromIterator[V](it: => Iterator[V], bufferSize: Int = 32)
                     (using Frame): LazyChunk[V, IO & Memo] =
    val builder: Iterator[V] => LazyChunk[V, IO & Memo] < (IO & Memo) =
      Memo: it =>
        IO:
          val builder = ChunkBuilder.init[V]
          val next: Int => Chunk[V] < (IO & Memo) =
            Memo: i =>
              IO:
                var count = 0
                while count < bufferSize && it.hasNext do
                  builder.addOne(it.next())
                  count += 1

                val res = builder.result()
                builder.clear()
                res



          def lazyChunk(idx: Int): LazyChunk[V, IO & Memo] < (IO & Memo) =
            next(idx).map: chunk =>
              if chunk.isEmpty then Empty
              else Cons(chunk, lazyChunk(idx + 1))

          lazyChunk(0)

    LazyChunk.Cons(Chunk.empty, builder(it))
  end fromIterator

end LazyChunk


// Tests
class LazyChunkTest extends Test:

  def runMemo(v: Assertion < (IO & Memo)) = run(Memo.run(v))

  "LazyChunk" - {
    "empty iterator" in runMemo {
      LazyChunk.fromIterator(Iterator.empty).toStream.run.map(x => assert(x == Chunk.empty))
    }

    "choice" in runMemo {

      val it = Iterator(1, 2, 3, 4)

      def f(i: Int): Int < Choice = Choice.eval(Seq(true, false)).andThen(i)

      val e = defer:
        val lazyChunk = LazyChunk.fromIterator(it)

        val stream = lazyChunk.toStream.map(i => f(i).later)

        stream.run.now == stream.run.now

      defer:
        assert(Choice.run(e).now.forall(identity))

    }

    "other rerun" in runMemo {
      val it = Iterator(1, 2, 3, 4)
      val lazyChunk: LazyChunk[Int, IO & Memo] = LazyChunk.fromIterator(it)

      defer:
        assert(lazyChunk.toStream.run.now == lazyChunk.toStream.run.now)

    }

    "continually" in runMemo {
      var count = 0
      val it = Iterator.continually {
        count += 1;
        count
      }

      val lazyChunk: LazyChunk[Int, IO & Memo] = LazyChunk.fromIterator(it, bufferSize = 2)

      defer:
        assert(lazyChunk.toStream.take(4).run.now == lazyChunk.toStream.take(4).run.now)

    }
  }


package kyo:

  abstract class Test extends org.scalatest.freespec.AsyncFreeSpec
    with org.scalatest.NonImplicitAssertions
    with kyo.internal.BaseKyoCoreTest:

    type Assertion = org.scalatest.Assertion

    def assertionSuccess = succeed

    def assertionFailure(msg: String) = fail(msg)

    override given executionContext: scala.concurrent.ExecutionContext = kyo.kernel.Platform.executionContext
  end Test
end kyo