From aebfd0a37d9b8640a6bdc4e9c67bd388a704d833 Mon Sep 17 00:00:00 2001
From: ychenfo <yc@m-labs.hk>
Date: Sat, 26 Mar 2022 01:44:50 +0800
Subject: [PATCH] nac3core/artiq: raise ValueError for unwrap none

---
 nac3artiq/demo/min_artiq.py       |  2 ++
 nac3core/src/codegen/expr.rs      | 20 ++++++++++++++++++++
 nac3core/src/toplevel/builtins.rs |  7 +------
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/nac3artiq/demo/min_artiq.py b/nac3artiq/demo/min_artiq.py
index bea9cc816..c361d55b9 100644
--- a/nac3artiq/demo/min_artiq.py
+++ b/nac3artiq/demo/min_artiq.py
@@ -46,6 +46,8 @@ class Option(Generic[T]):
         return not self.is_none()
     
     def unwrap(self):
+        if self.is_none():
+            raise ValueError("unwrap on none")
         return self._nac3_option
 
     def __repr__(self) -> str:
diff --git a/nac3core/src/codegen/expr.rs b/nac3core/src/codegen/expr.rs
index cf6cb33c9..72bc2cd65 100644
--- a/nac3core/src/codegen/expr.rs
+++ b/nac3core/src/codegen/expr.rs
@@ -1309,6 +1309,26 @@ pub fn gen_expr<'ctx, 'a, G: CodeGenerator>(
                             unreachable!()
                         }
                     };
+                    // directly generate code for option.unwrap
+                    // since it needs location information from ast
+                    if attr == &"unwrap".into()
+                        && id == ctx.primitives.option.get_obj_id(&ctx.unifier)
+                    {
+                        if let BasicValueEnum::PointerValue(ptr) = val.to_basic_value_enum(ctx, generator)? {
+                            let not_null = ctx.builder.build_is_not_null(ptr, "unwrap_not_null");
+                            ctx.make_assert(
+                                generator,
+                                not_null,
+                                "0:ValueError",
+                                "unwrap on none",
+                                [None, None, None],
+                                expr.location,
+                            );
+                            return Ok(Some(ctx.builder.build_load(ptr, "unwrap_some").into()))
+                        } else {
+                            unreachable!("option must be ptr")
+                        }
+                    }
                     return Ok(generator
                         .gen_call(
                             ctx,
diff --git a/nac3core/src/toplevel/builtins.rs b/nac3core/src/toplevel/builtins.rs
index 0367b5745..2a00797e2 100644
--- a/nac3core/src/toplevel/builtins.rs
+++ b/nac3core/src/toplevel/builtins.rs
@@ -263,12 +263,7 @@ pub fn get_builtins(primitives: &mut (PrimitiveStore, Unifier)) -> BuiltinInfo {
             resolver: None,
             codegen_callback: Some(Arc::new(GenCall::new(Box::new(
                 |ctx, obj, _, _, generator| {
-                    let obj_val = obj.unwrap().1.clone().to_basic_value_enum(ctx, generator)?;
-                    if let BasicValueEnum::PointerValue(ptr) = obj_val {
-                        Ok(Some(ctx.builder.build_load(ptr, "unwrap_some")))
-                    } else {
-                        unreachable!("option must be ptr")
-                    }
+                    unreachable!("handled in gen_expr")
                 },
             )))),
             loc: None,