From f50ad0a180ba6906c856fc82e3e51596ec086807 Mon Sep 17 00:00:00 2001
From: Daniel Urban <urban.dani@gmail.com>
Date: Sat, 13 Jan 2024 02:25:06 +0100
Subject: [PATCH] ByteStack: detect overflow immediately

---
 .../main/scala/cats/effect/ByteStack.scala    |  6 ++--
 .../main/scala/cats/effect/ByteStack.scala    |  6 ++--
 .../scala/cats/effect/ByteStackSpec.scala     | 36 +++++++++++++++++++
 3 files changed, 44 insertions(+), 4 deletions(-)
 create mode 100644 tests/shared/src/test/scala/cats/effect/ByteStackSpec.scala

diff --git a/core/js/src/main/scala/cats/effect/ByteStack.scala b/core/js/src/main/scala/cats/effect/ByteStack.scala
index c4e387f749..7a341a7f30 100644
--- a/core/js/src/main/scala/cats/effect/ByteStack.scala
+++ b/core/js/src/main/scala/cats/effect/ByteStack.scala
@@ -18,6 +18,8 @@ package cats.effect
 
 import scala.scalajs.js
 
+import java.lang.Math.addExact
+
 private object ByteStack {
 
   type T = js.Array[Int]
@@ -28,7 +30,7 @@ private object ByteStack {
   }
 
   @inline final def growIfNeeded(stack: js.Array[Int], count: Int): js.Array[Int] = {
-    if ((1 + ((count + 1) >> 3)) < stack.length) {
+    if ((1 + (addExact(count, 1) >> 3)) < stack.length) {
       stack
     } else {
       stack.push(0)
@@ -42,7 +44,7 @@ private object ByteStack {
     val s = (c >> 3) + 1 // current slot in `use`
     val shift = (c & 7) << 2 // BEGIN MAGIC
     use(s) = (use(s) & ~(0xffffffff << shift)) | (op << shift) // END MAGIC
-    use(0) += 1 // write the new count
+    use(0) = addExact(c, 1) // write the new count
     use
   }
 
diff --git a/core/jvm-native/src/main/scala/cats/effect/ByteStack.scala b/core/jvm-native/src/main/scala/cats/effect/ByteStack.scala
index 4d2bc16ef4..d2921fe0fe 100644
--- a/core/jvm-native/src/main/scala/cats/effect/ByteStack.scala
+++ b/core/jvm-native/src/main/scala/cats/effect/ByteStack.scala
@@ -16,6 +16,8 @@
 
 package cats.effect
 
+import java.lang.Math.addExact
+
 private object ByteStack {
 
   type T = Array[Int]
@@ -42,7 +44,7 @@ private object ByteStack {
     new Array[Int](1 + 1 + ((initialMaxOps - 1) >> 3)) // count-slot + 1 for each set of 8 ops
 
   final def growIfNeeded(stack: Array[Int], count: Int): Array[Int] = {
-    if ((1 + ((count + 1) >> 3)) < stack.length) {
+    if ((1 + (addExact(count, 1) >> 3)) < stack.length) {
       stack
     } else {
       val bigger = new Array[Int](stack.length << 1)
@@ -57,7 +59,7 @@ private object ByteStack {
     val s = (c >> 3) + 1 // current slot in `use`
     val shift = (c & 7) << 2 // BEGIN MAGIC
     use(s) = (use(s) & ~(0xffffffff << shift)) | (op << shift) // END MAGIC
-    use(0) += 1 // write the new count
+    use(0) = addExact(c, 1) // write the new count
     use
   }
 
diff --git a/tests/shared/src/test/scala/cats/effect/ByteStackSpec.scala b/tests/shared/src/test/scala/cats/effect/ByteStackSpec.scala
new file mode 100644
index 0000000000..05aedef810
--- /dev/null
+++ b/tests/shared/src/test/scala/cats/effect/ByteStackSpec.scala
@@ -0,0 +1,36 @@
+/*
+ * Copyright 2020-2023 Typelevel
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package cats.effect
+
+import scala.util.Try
+
+class ByteStackSpec extends BaseSpec {
+
+  "ByteStack" should {
+
+    "immediately detect when count overflows" in {
+      var bs = ByteStack.create(8)
+      for (_ <- 1 to Int.MaxValue) {
+        bs = ByteStack.push(bs, 13.toByte)
+      }
+      ByteStack.size(bs) must beEqualTo(Int.MaxValue)
+      Try {
+        ByteStack.push(bs, 13.toByte)
+      } must beAFailedTry.withThrowable[ArithmeticException]
+    }
+  }
+}