Fix binding of labels to return early if Assembler is in error state
PiperOrigin-RevId: 430962785
diff --git a/src/jit/aarch32-assembler.cc b/src/jit/aarch32-assembler.cc
index d503262..01f3bb1 100644
--- a/src/jit/aarch32-assembler.cc
+++ b/src/jit/aarch32-assembler.cc
@@ -122,6 +122,10 @@
}
void Assembler::bind(Label& l) {
+ if (error_ != Error::kNoError) {
+ return;
+ }
+
if (l.bound) {
error_ = Error::kLabelAlreadyBound;
return;
diff --git a/src/jit/aarch64-assembler.cc b/src/jit/aarch64-assembler.cc
index 85e9447..af7dd7d 100644
--- a/src/jit/aarch64-assembler.cc
+++ b/src/jit/aarch64-assembler.cc
@@ -600,6 +600,10 @@
}
void Assembler::bind(Label& l) {
+ if (error_ != Error::kNoError) {
+ return;
+ }
+
if (l.bound) {
error_ = Error::kLabelAlreadyBound;
return;
diff --git a/test/aarch32-assembler.cc b/test/aarch32-assembler.cc
index 3c1cbb4..2e8ae37 100644
--- a/test/aarch32-assembler.cc
+++ b/test/aarch32-assembler.cc
@@ -423,6 +423,24 @@
ASSERT_EQ(xnn_status_success, xnn_release_code_memory(&b));
}
+TEST(AArch32Assembler, BoundOverflow) {
+ xnn_code_buffer b;
+ xnn_allocate_code_memory(&b, 4);
+ Assembler a(&b);
+ Label l1;
+ a.add(r0, r0, 2);
+ EXPECT_EQ(Error::kNoError, a.error());
+
+ // This is out of bounds, not written.
+ a.bhi(l1);
+ EXPECT_EQ(Error::kOutOfMemory, a.error());
+
+ a.bind(l1);
+ EXPECT_EQ(false, l1.bound);
+
+ ASSERT_EQ(xnn_status_success, xnn_release_code_memory(&b));
+}
+
#if XNN_ARCH_ARM && XNN_PLATFORM_JIT
TEST(AArch32Assembler, JitAllocCodeBuffer) {
typedef uint32_t (*Func)(uint32_t);
diff --git a/test/aarch64-assembler.cc b/test/aarch64-assembler.cc
index 872f672..061611c 100644
--- a/test/aarch64-assembler.cc
+++ b/test/aarch64-assembler.cc
@@ -458,6 +458,23 @@
ASSERT_EQ(xnn_status_success, xnn_release_code_memory(&b));
}
+TEST(AArch64Assembler, BindOverflow) {
+ xnn_code_buffer b;
+ xnn_allocate_code_memory(&b, 4);
+ Assembler a(&b);
+ Label l1;
+ a.add(x0, x0, 2);
+ EXPECT_EQ(Error::kNoError, a.error());
+
+ // This is out of bounds, not written.
+ a.tbz(x1, 1, l1);
+ EXPECT_EQ(Error::kOutOfMemory, a.error());
+
+ a.bind(l1);
+ ASSERT_EQ(false, l1.bound);
+
+ ASSERT_EQ(xnn_status_success, xnn_release_code_memory(&b));
+}
} // namespace aarch64
} // namespace xnnpack
