From 2b2b2dbf8f37ca3d9a0181d7297eb150baaa21c2 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 18b490d7..882af5a9 100644
--- a/nac3core/src/codegen/expr.rs
+++ b/nac3core/src/codegen/expr.rs
@@ -2876,29 +2876,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 da4b6ee0..61a3d50e 100644
--- a/nac3core/src/codegen/stmt.rs
+++ b/nac3core/src/codegen/stmt.rs
@@ -1762,7 +1762,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(());