From 513d30152b19fb210cabcdde66df3d9bd0083aa6 Mon Sep 17 00:00:00 2001
From: abdul124 <ar@m-labs.hk>
Date: Wed, 10 Jul 2024 17:34:40 +0800
Subject: [PATCH] core: support raise exception short form

---
 nac3core/src/codegen/expr.rs                  | 24 ++++++++++++++++++-
 nac3core/src/typecheck/type_inferencer/mod.rs |  5 +++-
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/nac3core/src/codegen/expr.rs b/nac3core/src/codegen/expr.rs
index 8c5429a49..c42c8444c 100644
--- a/nac3core/src/codegen/expr.rs
+++ b/nac3core/src/codegen/expr.rs
@@ -2457,7 +2457,29 @@ pub fn gen_expr<'ctx, G: CodeGenerator>(
             Some((_, Some(static_value), _)) => ValueEnum::Static(static_value.clone()),
             None => {
                 let resolver = ctx.resolver.clone();
-                resolver.get_symbol_value(*id, ctx).unwrap()
+                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));
+                }
             }
         },
         ExprKind::List { elts, .. } => {
diff --git a/nac3core/src/typecheck/type_inferencer/mod.rs b/nac3core/src/typecheck/type_inferencer/mod.rs
index 4bff17c4c..d9380ab18 100644
--- a/nac3core/src/typecheck/type_inferencer/mod.rs
+++ b/nac3core/src/typecheck/type_inferencer/mod.rs
@@ -398,7 +398,10 @@ impl<'a> Fold<()> for Inferencer<'a> {
                 }
                 if let Some(exc) = exc {
                     self.virtual_checks.push((
-                        exc.custom.unwrap(),
+                        match &*self.unifier.get_ty(exc.custom.unwrap()) {
+                            TypeEnum::TFunc(sign) => sign.ret,
+                            _ => exc.custom.unwrap(),
+                        },
                         self.primitives.exception,
                         exc.location,
                     ));