From 552dbc02f1bbf8dbc43966f11ed069d5abc1c851 Mon Sep 17 00:00:00 2001
From: David Mak <chmakac@connect.ust.hk>
Date: Mon, 26 Aug 2024 14:29:14 +0800
Subject: [PATCH] [core] Fix resolution of exception names in raise short form

Previous implementation fails as `resolver.get_identifier_def` in ARTIQ
would return the exception __init__ function rather than the class.

We fix this by limiting the exception class resolution to only include
raise statements, and to force the exception name to always be treated
as a class.

Fixes #501.
---
 nac3core/src/codegen/expr.rs | 24 +-----------------------
 nac3core/src/codegen/stmt.rs | 25 ++++++++++++++++++++++++-
 2 files changed, 25 insertions(+), 24 deletions(-)

diff --git a/nac3core/src/codegen/expr.rs b/nac3core/src/codegen/expr.rs
index 50d70eb7..730caa5d 100644
--- a/nac3core/src/codegen/expr.rs
+++ b/nac3core/src/codegen/expr.rs
@@ -2873,29 +2873,7 @@ pub fn gen_expr<'ctx, G: CodeGenerator>(
             Some((_, Some(static_value), _)) => ValueEnum::Static(static_value.clone()),
             None => {
                 let resolver = ctx.resolver.clone();
-                if let Some(res) = resolver.get_symbol_value(*id, ctx) {
-                    res
-                } else {
-                    // Allow "raise Exception" short form
-                    let def_id = resolver.get_identifier_def(*id).map_err(|e| {
-                        format!("{} (at {})", e.iter().next().unwrap(), expr.location)
-                    })?;
-                    let def = ctx.top_level.definitions.read();
-                    if let TopLevelDef::Class { constructor, .. } = *def[def_id.0].read() {
-                        let TypeEnum::TFunc(signature) =
-                            ctx.unifier.get_ty(constructor.unwrap()).as_ref().clone()
-                        else {
-                            return Err(format!(
-                                "Failed to resolve symbol {} (at {})",
-                                id, expr.location
-                            ));
-                        };
-                        return Ok(generator
-                            .gen_call(ctx, None, (&signature, def_id), Vec::default())?
-                            .map(Into::into));
-                    }
-                    return Err(format!("Failed to resolve symbol {} (at {})", id, expr.location));
-                }
+                resolver.get_symbol_value(*id, ctx).unwrap()
             }
         },
         ExprKind::List { elts, .. } => {
diff --git a/nac3core/src/codegen/stmt.rs b/nac3core/src/codegen/stmt.rs
index 081a5cef..ab53a8f4 100644
--- a/nac3core/src/codegen/stmt.rs
+++ b/nac3core/src/codegen/stmt.rs
@@ -1760,7 +1760,30 @@ pub fn gen_stmt<G: CodeGenerator>(
         StmtKind::Try { .. } => gen_try(generator, ctx, stmt)?,
         StmtKind::Raise { exc, .. } => {
             if let Some(exc) = exc {
-                let exc = if let Some(v) = generator.gen_expr(ctx, exc)? {
+                let exn = if let ExprKind::Name { id, .. } = &exc.node {
+                    // Handle "raise Exception" short form
+                    let def_id = ctx.resolver.get_identifier_def(*id).map_err(|e| {
+                        format!("{} (at {})", e.iter().next().unwrap(), exc.location)
+                    })?;
+                    let def = ctx.top_level.definitions.read();
+                    let TopLevelDef::Class { constructor, .. } = *def[def_id.0].read() else {
+                        return Err(format!("Failed to resolve symbol {id} (at {})", exc.location));
+                    };
+
+                    let TypeEnum::TFunc(signature) =
+                        ctx.unifier.get_ty(constructor.unwrap()).as_ref().clone()
+                    else {
+                        return Err(format!("Failed to resolve symbol {id} (at {})", exc.location));
+                    };
+
+                    generator
+                        .gen_call(ctx, None, (&signature, def_id), Vec::default())?
+                        .map(Into::into)
+                } else {
+                    generator.gen_expr(ctx, exc)?
+                };
+
+                let exc = if let Some(v) = exn {
                     v.to_basic_value_enum(ctx, generator, exc.custom.unwrap())?
                 } else {
                     return Ok(());