From b0f775ca9ffb21a67462d9f8a000d35a8a4ba86e Mon Sep 17 00:00:00 2001
From: Fabian Ruffy <5960321+fruffy@users.noreply.github.com>
Date: Thu, 12 Dec 2024 19:25:18 +0000
Subject: [PATCH] Add checks to annotation getters. (#5052)

* Add annotation error messages.

Signed-off-by: fruffy <fruffy@nyu.edu>

* Check for KVAnnotation instead of deferring to annotationKind().

Signed-off-by: fruffy <fruffy@nyu.edu>

* Use try-catch instead of BUG_CHECK.

Signed-off-by: fruffy <fruffy@nyu.edu>

---------

Signed-off-by: fruffy <fruffy@nyu.edu>
---
 ir/base.def | 59 +++++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 50 insertions(+), 9 deletions(-)

diff --git a/ir/base.def b/ir/base.def
index aaaa1f4b4c9..101515ecf08 100644
--- a/ir/base.def
+++ b/ir/base.def
@@ -332,17 +332,58 @@ class Annotation {
             Vector<Expression>,
             IndexedVector<NamedExpression>> body;
 
-    inline auto &getUnparsed() { return std::get<UnparsedAnnotation>(body); }
-    inline const auto &getUnparsed() const { return std::get<UnparsedAnnotation>(body); }
-    inline auto &getExpr() { return std::get<ExpressionAnnotation>(body); }
-    inline const auto &getExpr() const { return std::get<ExpressionAnnotation>(body); }
+    inline auto &getUnparsed() {
+        try {
+            return std::get<UnparsedAnnotation>(body);
+        } catch (const std::bad_variant_access &) {
+            BUG("Annotation has been parsed already.");
+        }
+    }
+    inline const auto &getUnparsed() const {
+        try {
+            return std::get<UnparsedAnnotation>(body);
+        } catch (const std::bad_variant_access &) {
+            BUG("Annotation has been parsed already.");
+        }
+    }
+    inline auto &getExpr() {
+        try {
+            return std::get<ExpressionAnnotation>(body);
+        } catch (const std::bad_variant_access &) {
+            BUG("Annotation does not contain an expression list.");
+        }
+     }
+    inline const auto &getExpr() const {
+        try {
+            return std::get<ExpressionAnnotation>(body);
+        } catch (const std::bad_variant_access &) {
+            BUG("Annotation does not contain an expression list.");
+        }
+    }
     inline Expression getExpr(size_t idx) const {
-        const auto &expr = getExpr();
-        BUG_CHECK(idx < expr.size(), "invalid annotation expression index");
-        return expr[idx];
+        try {
+            const auto &expr = getExpr();
+            return expr[idx];
+        } catch (const std::out_of_range &) {
+            BUG("invalid annotation expression index");
+        } catch (const std::bad_variant_access &) {
+            BUG("Annotation does not contain an expression list.");
+        }
+    }
+    inline auto &getKV() {
+        try {
+            return std::get<KVAnnotation>(body);
+        } catch (const std::bad_variant_access &) {
+            BUG("Annotation does not contain a key-value list.");
+        }
+    }
+    inline const auto &getKV() const {
+        try {
+            return std::get<KVAnnotation>(body);
+        } catch (const std::bad_variant_access &) {
+            BUG("Annotation does not contain a key-value list.");
+        }
     }
-    inline auto &getKV() { return std::get<KVAnnotation>(body); }
-    inline const auto &getKV() const { return std::get<KVAnnotation>(body); }
 
     /// If this is true this is a structured annotation, and there are some
     /// constraints on its contents.