fix re-slice in chunk cause OOM. (zio#6807)

* fix re-slice in chunk cause OOM. * format with scalafmt
IMax153 · May 13, 2022 · 0c6c907 · 0c6c907
1 parent 2828be1
commit 0c6c907
Showing 1 changed file with 16 additions and 4 deletions.
diff --git a/core/shared/src/main/scala/zio/Chunk.scala b/core/shared/src/main/scala/zio/Chunk.scala
@@ -196,7 +196,10 @@ sealed abstract class Chunk[+A] extends ChunkLike[A] { self =>
     else
       self match {
         case Chunk.Slice(c, o, l) => Chunk.Slice(c, o + n, l - n)
-        case _                    => Chunk.Slice(self, n, len - n)
+        case Chunk.Concat(l, r) =>
+          if (n > l.length) r.drop(n - l.length)
+          else Chunk.Concat(l.drop(n), r)
+        case _ => Chunk.Slice(self, n, len - n)
       }
   }
 
@@ -211,7 +214,10 @@ sealed abstract class Chunk[+A] extends ChunkLike[A] { self =>
     else
       self match {
         case Chunk.Slice(c, o, l) => Chunk.Slice(c, o, l - n)
-        case _                    => Chunk.Slice(self, 0, len - n)
+        case Chunk.Concat(l, r) =>
+          if (n > r.length) l.dropRight(n - r.length)
+          else Chunk.Concat(l, r.dropRight(n))
+        case _ => Chunk.Slice(self, 0, len - n)
       }
   }
 
@@ -780,7 +786,10 @@ sealed abstract class Chunk[+A] extends ChunkLike[A] { self =>
     else
       self match {
         case Chunk.Slice(c, o, _) => Chunk.Slice(c, o, n)
-        case _                    => Chunk.Slice(self, 0, n)
+        case Chunk.Concat(l, r) =>
+          if (n > l.length) Chunk.Concat(l, r.take(n - l.length))
+          else l.take(n)
+        case _ => Chunk.Slice(self, 0, n)
       }
 
   /**
@@ -792,7 +801,10 @@ sealed abstract class Chunk[+A] extends ChunkLike[A] { self =>
     else
       self match {
         case Chunk.Slice(c, o, l) => Chunk.Slice(c, o + l - n, n)
-        case _                    => Chunk.Slice(self, length - n, n)
+        case Chunk.Concat(l, r) =>
+          if (n > r.length) Chunk.Concat(l.takeRight(n - r.length), r)
+          else r.takeRight(n)
+        case _ => Chunk.Slice(self, length - n, n)
       }
 
   /**