How to aggregate elements of one Akka stream based on elements of another?

Example scenario: group bytes of a stream into chunks of sizes determined by another stream (of integers).

def partition[A, B, C](
  first:Source[A, NotUsed],
  second:Source[B, NotUsed],
  aggregate:(Int => Seq[A], B) => C
):Source[C, NotUsed] = ???

val bytes:Source[Byte, NotUsed] = ???
val sizes:Source[Int, NotUsed] = ???

val chunks:Source[ByteString, NotUsed] =
  partition(bytes, sizes, (grab, count) => ByteString(grab(count)))

My initial attempt includes a combination of Flow#scan and Flow#prefixAndTail, but it doesn't feel quite right (see below). I also took a look at Framing, but it doesn't seem to be applicable to the example scenario above (nor is it general enough to accommodate non-bytestring streams). I'm guessing my only option is to use Graphs (or the more general FlowOps#transform), but I'm not nearly proficient enough (yet) with Akka streams to attempt that.

Here's what I was able to come up with so far (specific to the example scenario):

val chunks:Source[ByteString, NotUsed] = sizes
  .scan(bytes prefixAndTail 0) {
    (grouped, count) => grouped flatMapConcat {
      case (chunk, remainder) => remainder prefixAndTail count
    }
  }
  .flatMapConcat(identity)
  .collect { case (chunk, _) if chunk.nonEmpty => ByteString(chunk:_*) }