nac3artiq/core: host option object support

the None token is parsed as a special ast::constant instead of an ast::name, and we use None now as a special value of type Option[Any]

nac3artiq/demo/min_artiq.py

@@ -11,6 +11,7 @@ from embedding_map import EmbeddingMap
 
 __all__ = [
     "Kernel", "KernelInvariant", "virtual",
+    "Option", "Some",
     "round64", "floor64", "ceil64",
     "extern", "kernel", "portable", "nac3",
     "rpc", "ms", "us", "ns",
@@ -32,6 +33,36 @@ class KernelInvariant(Generic[T]):
 class virtual(Generic[T]):
     pass
 
+class Option(Generic[T]):
+    _nac3_option: T
+
+    def __init__(self, v: T):
+        self._nac3_option = v
+
+    def is_none(self):
+        return self._nac3_option is None
+    
+    def is_some(self):
+        return not self.is_none()
+    
+    def unwrap(self):
+        return self._nac3_option
+
+    def __repr__(self) -> str:
+        if self.is_none():
+            return "Option(None)"
+        else:
+            return "Some({})".format(repr(self._nac3_option))
+    
+    def __str__(self) -> str:
+        if self.is_none():
+            return "None"
+        else:
+            return "Some({})".format(str(self._nac3_option))
+
+def Some(v: T) -> Option[T]:
+    return Option(v)
+
 
 def round64(x):
     return round(x)

nac3artiq/src/lib.rs

@@ -71,6 +71,7 @@ pub struct PrimitivePythonId {
     exception: u64,
     generic_alias: (u64, u64),
     virtual_id: u64,
+    option: u64,
 }
 
 type TopLevelComponent = (Stmt, String, PyObject);
@@ -351,6 +352,7 @@ impl Nac3 {
 
         let builtins_mod = PyModule::import(py, "builtins").unwrap();
         let id_fn = builtins_mod.getattr("id").unwrap();
+        let type_fn = builtins_mod.getattr("type").unwrap();
         let numpy_mod = PyModule::import(py, "numpy").unwrap();
         let typing_mod = PyModule::import(py, "typing").unwrap();
         let types_mod = PyModule::import(py, "types").unwrap();
@@ -372,7 +374,11 @@ impl Nac3 {
                 get_attr_id(typing_mod, "_GenericAlias"),
                 get_attr_id(types_mod, "GenericAlias"),
             ),
-            none: get_attr_id(builtins_mod, "None"),
+            none: id_fn
+                .call1((type_fn.call1((builtins_mod.getattr("None").unwrap(),)).unwrap(),))
+                .unwrap()
+                .extract()
+                .unwrap(),
             typevar: get_attr_id(typing_mod, "TypeVar"),
             int: get_attr_id(builtins_mod, "int"),
             int32: get_attr_id(numpy_mod, "int32"),
@@ -384,6 +390,17 @@ impl Nac3 {
             list: get_attr_id(builtins_mod, "list"),
             tuple: get_attr_id(builtins_mod, "tuple"),
             exception: get_attr_id(builtins_mod, "Exception"),
+            option: id_fn
+                .call1((builtins_mod
+                    .getattr("globals")
+                    .unwrap()
+                    .call0()
+                    .unwrap()
+                    .get_item("Option")
+                    .unwrap(),))
+                .unwrap()
+                .extract()
+                .unwrap(),
         };
 
         let working_directory = tempfile::Builder::new().prefix("nac3-").tempdir().unwrap();

nac3artiq/src/symbol_resolver.rs

@@ -260,6 +260,27 @@ impl InnerResolver {
         } else if ty_id == self.primitive_ids.tuple {
             // do not handle type var param and concrete check here
             Ok(Ok((unifier.add_ty(TypeEnum::TTuple { ty: vec![] }), false)))
+        } else if ty_id == self.primitive_ids.option {
+            Ok(Ok((primitives.option, false)))
+        } else if ty_id == self.primitive_ids.none {
+            if let TypeEnum::TObj { params, .. } =
+                unifier.get_ty_immutable(primitives.option).as_ref()
+            {
+                let var_map = params
+                    .iter()
+                    .map(|(id_var, ty)| {
+                        if let TypeEnum::TVar { id, range, name, loc, .. } = &*unifier.get_ty(*ty) {
+                            assert_eq!(*id, *id_var);
+                            (*id, unifier.get_fresh_var_with_range(range, *name, *loc).0)
+                        } else {
+                            unreachable!()
+                        }
+                    })
+                    .collect::<HashMap<_, _>>();
+                Ok(Ok((unifier.subst(primitives.option, &var_map).unwrap(), true)))
+            } else {
+                unreachable!("must be tobj")
+            }
         } else if let Some(def_id) = self.pyid_to_def.read().get(&ty_id).cloned() {
             let def = defs[def_id.0].read();
             if let TopLevelDef::Class { object_id, type_vars, fields, methods, .. } = &*def {
@@ -538,6 +559,34 @@ impl InnerResolver {
                 let types = types?;
                 Ok(types.map(|types| unifier.add_ty(TypeEnum::TTuple { ty: types })))
             }
+            // special handling for option type since its class member layout in python side
+            // is special and cannot be mapped directly to a nac3 type as below
+            (TypeEnum::TObj { obj_id, params, .. }, false)
+                if *obj_id == primitives.option.get_obj_id(unifier) =>
+            {
+                let field_data = match obj.getattr("_nac3_option") {
+                    Ok(d) => d,
+                    // None should be already handled above
+                    Err(_) => unreachable!("cannot be None")
+                };
+                let field_obj_id: u64 = self.helper.id_fn.call1(py, (field_data,))?.extract(py)?;
+                let zelf_obj_id: u64 = self.helper.id_fn.call1(py, (obj,))?.extract(py)?;
+                if field_obj_id == zelf_obj_id {
+                    return Ok(Err("self recursive option type is not allowed".into()))
+                }
+                let ty = match self.get_obj_type(py, field_data, unifier, defs, primitives)? {
+                    Ok(t) => t,
+                    Err(e) => {
+                        return Ok(Err(format!(
+                            "error when getting type of the option object ({})",
+                            e
+                        )))
+                    }
+                };
+                let new_var_map: HashMap<_, _> = params.iter().map(|(id, _)| (*id, ty)).collect();
+                let res = unifier.subst(extracted_ty, &new_var_map).unwrap_or(extracted_ty);
+                Ok(Ok(res))
+            }
             (TypeEnum::TObj { params: var_map, fields, .. }, false) => {
                 self.pyid_to_type.write().insert(ty_id, extracted_ty);
                 let mut instantiate_obj = || {
@@ -756,6 +805,37 @@ impl InnerResolver {
             let global = ctx.module.add_global(ty, Some(AddressSpace::Generic), &id_str);
             global.set_initializer(&val);
             Ok(Some(global.as_pointer_value().into()))
+        } else if ty_id == self.primitive_ids.option {
+            match self
+                .get_obj_value(py, obj.getattr("_nac3_option").unwrap(), ctx, generator)
+                .map_err(|e| {
+                    super::CompileError::new_err(format!(
+                        "Error getting value of Option object: {}",
+                        e
+                    ))
+                })? {
+                Some(v) => {
+                    let global_str = format!("{}_option", id);
+                    {
+                        if self.global_value_ids.read().contains(&id) {
+                            let global = ctx.module.get_global(&global_str).unwrap_or_else(|| {
+                                ctx.module.add_global(v.get_type(), Some(AddressSpace::Generic), &global_str)
+                            });
+                            return Ok(Some(global.as_pointer_value().into()));
+                        } else {
+                            self.global_value_ids.write().insert(id);
+                        }
+                    }
+                    let global = ctx.module.add_global(v.get_type(), Some(AddressSpace::Generic), &global_str);
+                    global.set_initializer(&v);
+                    Ok(Some(global.as_pointer_value().into()))
+                },
+                None => Ok(None),
+            }
+        } else if ty_id == self.primitive_ids.none {
+            // for option type, just a null ptr, whose type needs to be casted in codegen
+            // according to the type info attached in the ast
+            Ok(Some(ctx.ctx.i8_type().ptr_type(AddressSpace::Generic).const_null().into()))
         } else {
             let id_str = id.to_string();
 
@@ -900,7 +980,15 @@ impl SymbolResolver for Resolver {
                             }
                         }
                         if let Ok(t) = sym_ty {
-                            self.0.pyid_to_type.write().insert(*id, t);
+                            // do not cache for option type since None can have same pyid but different type
+                            match unifier.get_ty_immutable(t).as_ref() {
+                                TypeEnum::TObj { obj_id, .. }
+                                    if *obj_id == primitives.option.get_obj_id(unifier) =>
+                                {
+                                    None
+                                }
+                                _ => self.0.pyid_to_type.write().insert(*id, t),
+                            };
                         }
                         Ok(sym_ty)
                     })

nac3core/src/codegen/expr.rs

@@ -138,6 +138,7 @@ impl<'ctx, 'a> CodeGenContext<'ctx, 'a> {
             &mut self.unifier,
             self.top_level,
             &mut self.type_cache,
+            &self.primitives,
             ty,
         )
     }
@@ -199,6 +200,22 @@ impl<'ctx, 'a> CodeGenContext<'ctx, 'a> {
                     val
                 }
             }
+            Constant::None => {
+                match (
+                    self.unifier.get_ty(ty).as_ref(),
+                    self.unifier.get_ty(self.primitives.option).as_ref(),
+                ) {
+                    (
+                        TypeEnum::TObj { obj_id, params, .. },
+                        TypeEnum::TObj { obj_id: opt_id, .. },
+                    ) if *obj_id == *opt_id => self
+                        .get_llvm_type(generator, *params.iter().next().unwrap().1)
+                        .ptr_type(AddressSpace::Generic)
+                        .const_null()
+                        .into(),
+                    _ => unreachable!("must be option type"),
+                }
+            }
             _ => unreachable!(),
         }
     }
@@ -935,16 +952,27 @@ pub fn gen_expr<'ctx, 'a, G: CodeGenerator>(
                 let resolver = ctx.resolver.clone();
                 let val = resolver.get_symbol_value(*id, ctx).unwrap();
                 // if is tuple, need to deref it to handle tuple as value
-                if let (TypeEnum::TTuple { .. }, BasicValueEnum::PointerValue(ptr)) = (
+                // if is option, need to cast pointer to handle None
+                match (
                     &*ctx.unifier.get_ty(expr.custom.unwrap()),
                     resolver
                         .get_symbol_value(*id, ctx)
                         .unwrap()
                         .to_basic_value_enum(ctx, generator)?,
                 ) {
-                    ctx.builder.build_load(ptr, "tup_val").into()
-                } else {
-                    val
+                    (TypeEnum::TTuple { .. }, BasicValueEnum::PointerValue(ptr)) => {
+                        ctx.builder.build_load(ptr, "tup_val").into()
+                    }
+                    (TypeEnum::TObj { obj_id, params, .. }, BasicValueEnum::PointerValue(ptr))
+                        if *obj_id == ctx.primitives.option.get_obj_id(&ctx.unifier) => {
+                            let actual_ptr_ty = ctx.get_llvm_type(
+                                generator,
+                                *params.iter().next().unwrap().1,
+                            )
+                            .ptr_type(AddressSpace::Generic);
+                            ctx.builder.build_bitcast(ptr, actual_ptr_ty, "option_ptr_cast").into()
+                        }
+                    _ => val,
                 }
             }
         },

nac3core/src/codegen/mod.rs

@@ -259,6 +259,7 @@ fn get_llvm_type<'ctx>(
     unifier: &mut Unifier,
     top_level: &TopLevelContext,
     type_cache: &mut HashMap<Type, BasicTypeEnum<'ctx>>,
+    primitives: &PrimitiveStore,
     ty: Type,
 ) -> BasicTypeEnum<'ctx> {
     use TypeEnum::*;
@@ -268,9 +269,28 @@ fn get_llvm_type<'ctx>(
         let ty_enum = unifier.get_ty(ty);
         let result = match &*ty_enum {
             TObj { obj_id, fields, .. } => {
-                // check to avoid treating primitives as classes
-                if obj_id.0 <= 7 {
-                    unreachable!();
+                // check to avoid treating primitives other than Option as classes
+                if obj_id.0 <= 10 {
+                    match (unifier.get_ty(ty).as_ref(), unifier.get_ty(primitives.option).as_ref())
+                    {
+                        (
+                            TypeEnum::TObj { obj_id, params, .. },
+                            TypeEnum::TObj { obj_id: opt_id, .. },
+                        ) if *obj_id == *opt_id => {
+                            return get_llvm_type(
+                                ctx,
+                                generator,
+                                unifier,
+                                top_level,
+                                type_cache,
+                                primitives,
+                                *params.iter().next().unwrap().1,
+                            )
+                            .ptr_type(AddressSpace::Generic)
+                            .into();
+                        }
+                        _ => unreachable!("must be option type"),
+                    }
                 }
                 // a struct with fields in the order of declaration
                 let top_level_defs = top_level.definitions.read();
@@ -289,6 +309,7 @@ fn get_llvm_type<'ctx>(
                                 unifier,
                                 top_level,
                                 type_cache,
+                                primitives,
                                 fields[&f.0].0,
                             )
                         })
@@ -304,14 +325,14 @@ fn get_llvm_type<'ctx>(
                 // a struct with fields in the order present in the tuple
                 let fields = ty
                     .iter()
-                    .map(|ty| get_llvm_type(ctx, generator, unifier, top_level, type_cache, *ty))
+                    .map(|ty| get_llvm_type(ctx, generator, unifier, top_level, type_cache, primitives, *ty))
                     .collect_vec();
                 ctx.struct_type(&fields, false).into()
             }
             TList { ty } => {
                 // a struct with an integer and a pointer to an array
                 let element_type =
-                    get_llvm_type(ctx, generator, unifier, top_level, type_cache, *ty);
+                    get_llvm_type(ctx, generator, unifier, top_level, type_cache, primitives, *ty);
                 let fields = [
                     element_type.ptr_type(AddressSpace::Generic).into(),
                     generator.get_size_type(ctx).into(),
@@ -384,6 +405,7 @@ pub fn gen_func<'ctx, G: CodeGenerator>(
         range: unifier.get_representative(primitives.range),
         str: unifier.get_representative(primitives.str),
         exception: unifier.get_representative(primitives.exception),
+        option: unifier.get_representative(primitives.option),
     };
 
     let mut type_cache: HashMap<_, _> = [
@@ -416,6 +438,8 @@ pub fn gen_func<'ctx, G: CodeGenerator>(
         exception.set_body(&fields, false);
         exception.ptr_type(AddressSpace::Generic).into()
     });
+    // NOTE: special handling of option cannot use this type cache since it contains type var,
+    // handled inside get_llvm_type instead
 
     let (args, ret) = if let ConcreteTypeEnum::TFunc { args, ret, .. } =
         task.store.get(task.signature)
@@ -436,7 +460,7 @@ pub fn gen_func<'ctx, G: CodeGenerator>(
     let ret_type = if unifier.unioned(ret, primitives.none) {
         None
     } else {
-        Some(get_llvm_type(context, generator, &mut unifier, top_level_ctx.as_ref(), &mut type_cache, ret))
+        Some(get_llvm_type(context, generator, &mut unifier, top_level_ctx.as_ref(), &mut type_cache, &primitives, ret))
     };
 
     let has_sret = ret_type.map_or(false, |ty| need_sret(context, ty));
@@ -449,6 +473,7 @@ pub fn gen_func<'ctx, G: CodeGenerator>(
                 &mut unifier,
                 top_level_ctx.as_ref(),
                 &mut type_cache,
+                &primitives,
                 arg.ty,
             )
             .into()
@@ -496,6 +521,7 @@ pub fn gen_func<'ctx, G: CodeGenerator>(
                 &mut unifier,
                 top_level_ctx.as_ref(),
                 &mut type_cache,
+                &primitives,
                 arg.ty,
             ),
             &arg.name.to_string(),

nac3core/src/symbol_resolver.rs

@@ -145,12 +145,11 @@ pub trait SymbolResolver {
 }
 
 thread_local! {
-    static IDENTIFIER_ID: [StrRef; 12] = [
+    static IDENTIFIER_ID: [StrRef; 11] = [
         "int32".into(),
         "int64".into(),
         "float".into(),
         "bool".into(),
-        "None".into(),
         "virtual".into(),
         "list".into(),
         "tuple".into(),
@@ -175,14 +174,13 @@ pub fn parse_type_annotation<T>(
     let int64_id = ids[1];
     let float_id = ids[2];
     let bool_id = ids[3];
-    let none_id = ids[4];
-    let virtual_id = ids[5];
-    let list_id = ids[6];
-    let tuple_id = ids[7];
-    let str_id = ids[8];
-    let exn_id = ids[9];
-    let uint32_id = ids[10];
-    let uint64_id = ids[11];
+    let virtual_id = ids[4];
+    let list_id = ids[5];
+    let tuple_id = ids[6];
+    let str_id = ids[7];
+    let exn_id = ids[8];
+    let uint32_id = ids[9];
+    let uint64_id = ids[10];
 
     let name_handling = |id: &StrRef, loc: Location, unifier: &mut Unifier| {
         if *id == int32_id {
@@ -197,8 +195,6 @@ pub fn parse_type_annotation<T>(
             Ok(primitives.float)
         } else if *id == bool_id {
             Ok(primitives.bool)
-        } else if *id == none_id {
-            Ok(primitives.none)
         } else if *id == str_id {
             Ok(primitives.str)
         } else if *id == exn_id {

nac3core/src/toplevel/builtins.rs

@@ -105,6 +105,19 @@ pub fn get_builtins(primitives: &mut (PrimitiveStore, Unifier)) -> BuiltinInfo {
         ("__param2__".into(), int64, true),
     ];
 
+    // for Option, is_some and is_none share the same type
+    let (is_some_ty, unwrap_ty, (option_ty_var, option_ty_var_id)) =
+        if let TypeEnum::TObj { fields, params, .. } =
+            primitives.1.get_ty(primitives.0.option).as_ref()
+        {
+            (
+                *fields.get(&"is_some".into()).unwrap(),
+                *fields.get(&"unwrap".into()).unwrap(),
+                (*params.iter().next().unwrap().1, *params.iter().next().unwrap().0),
+            )
+        } else {
+            unreachable!()
+        };
     let top_level_def_list = vec![
         Arc::new(RwLock::new(TopLevelComposer::make_top_level_class_def(
             0,
@@ -180,6 +193,86 @@ pub fn get_builtins(primitives: &mut (PrimitiveStore, Unifier)) -> BuiltinInfo {
             None,
             None,
         ))),
+        Arc::new(RwLock::new({
+            TopLevelDef::Class {
+                name: "Option".into(),
+                object_id: DefinitionId(10),
+                type_vars: vec![option_ty_var],
+                fields: vec![],
+                methods: vec![
+                    ("is_some".into(), is_some_ty.0, DefinitionId(11)),
+                    ("is_none".into(), is_some_ty.0, DefinitionId(12)),
+                    ("unwrap".into(), unwrap_ty.0, DefinitionId(13)),
+                ],
+                ancestors: vec![TypeAnnotation::CustomClass {
+                    id: DefinitionId(10),
+                    params: Default::default(),
+                }],
+                constructor: None,
+                resolver: None,
+                loc: None,
+            }
+        })),
+        Arc::new(RwLock::new(TopLevelDef::Function {
+            name: "Option.is_some".into(),
+            simple_name: "is_some".into(),
+            signature: is_some_ty.0,
+            var_id: vec![option_ty_var_id],
+            instance_to_symbol: Default::default(),
+            instance_to_stmt: Default::default(),
+            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_is_not_null(ptr, "is_some").into()))
+                    } else {
+                        unreachable!("option must be ptr")
+                    }
+                },
+            )))),
+            loc: None,
+        })),
+        Arc::new(RwLock::new(TopLevelDef::Function {
+            name: "Option.is_none".into(),
+            simple_name: "is_none".into(),
+            signature: is_some_ty.0,
+            var_id: vec![option_ty_var_id],
+            instance_to_symbol: Default::default(),
+            instance_to_stmt: Default::default(),
+            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_is_null(ptr, "is_none").into()))
+                    } else {
+                        unreachable!("option must be ptr")
+                    }
+                },
+            )))),
+            loc: None,
+        })),
+        Arc::new(RwLock::new(TopLevelDef::Function {
+            name: "Option.unwrap".into(),
+            simple_name: "unwrap".into(),
+            signature: unwrap_ty.0,
+            var_id: vec![option_ty_var_id],
+            instance_to_symbol: Default::default(),
+            instance_to_stmt: Default::default(),
+            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")
+                    }
+                },
+            )))),
+            loc: None,
+        })),
         Arc::new(RwLock::new(TopLevelDef::Function {
             name: "int32".into(),
             simple_name: "int32".into(),
@@ -1098,6 +1191,28 @@ pub fn get_builtins(primitives: &mut (PrimitiveStore, Unifier)) -> BuiltinInfo {
             )))),
             loc: None,
         })),
+        Arc::new(RwLock::new(TopLevelDef::Function {
+            name: "Some".into(),
+            simple_name: "Some".into(),
+            signature: primitives.1.add_ty(TypeEnum::TFunc(FunSignature {
+                args: vec![FuncArg { name: "n".into(), ty: option_ty_var, default_value: None }],
+                ret: primitives.0.option,
+                vars: HashMap::from([(option_ty_var_id, option_ty_var)]),
+            })),
+            var_id: vec![option_ty_var_id],
+            instance_to_symbol: Default::default(),
+            instance_to_stmt: Default::default(),
+            resolver: None,
+            codegen_callback: Some(Arc::new(GenCall::new(Box::new(
+                |ctx, _, _fun, args, generator| {
+                    let arg_val = args[0].1.clone().to_basic_value_enum(ctx, generator)?;
+                    let alloca = ctx.builder.build_alloca(arg_val.get_type(), "alloca_some");
+                    ctx.builder.build_store(alloca, arg_val);
+                    Ok(Some(alloca.into()))
+                },
+            )))),
+            loc: None,
+        })),
     ];
 
     let ast_list: Vec<Option<ast::Stmt<()>>> =
@@ -1123,6 +1238,7 @@ pub fn get_builtins(primitives: &mut (PrimitiveStore, Unifier)) -> BuiltinInfo {
             "min",
             "max",
             "abs",
+            "Some",
         ],
     )
 }

nac3core/src/toplevel/composer.rs

@@ -74,6 +74,8 @@ impl TopLevelComposer {
             "self".into(),
             "Kernel".into(),
             "KernelInvariant".into(),
+            "Some".into(),
+            "Option".into(),
         ]);
         let defined_names: HashSet<String> = Default::default();
         let method_class: HashMap<DefinitionId, DefinitionId> = Default::default();
@@ -92,7 +94,6 @@ impl TopLevelComposer {
             } else if let TopLevelDef::Class { name, constructor, object_id, type_vars, .. } = &*def
             {
                 assert!(id == object_id.0);
-                assert!(type_vars.is_empty());
                 if let Some(constructor) = constructor {
                     builtin_ty.insert(*name, *constructor);
                 }
@@ -1747,9 +1748,7 @@ impl TopLevelComposer {
                             })
                         };
                         let mut identifiers = {
-                            // NOTE: none and function args?
                             let mut result: HashSet<_> = HashSet::new();
-                            result.insert("None".into());
                             if self_type.is_some() {
                                 result.insert("self".into());
                             }
@@ -1772,9 +1771,7 @@ impl TopLevelComposer {
                             },
                             unifier,
                             variable_mapping: {
-                                // NOTE: none and function args?
                                 let mut result: HashMap<StrRef, Type> = HashMap::new();
-                                result.insert("None".into(), primitives_ty.none);
                                 if let Some(self_ty) = self_type {
                                     result.insert("self".into(), self_ty);
                                 }

nac3core/src/toplevel/helper.rs

@@ -107,7 +107,43 @@ impl TopLevelComposer {
             fields: HashMap::new(),
             params: HashMap::new(),
         });
-        let primitives = PrimitiveStore { int32, int64, float, bool, none, range, str, exception, uint32, uint64 };
+
+        let option_type_var = unifier.get_fresh_var(Some("option_type_var".into()), None);
+        let is_some_type_fun_ty = unifier.add_ty(TypeEnum::TFunc(FunSignature {
+            args: vec![],
+            ret: bool,
+            vars: HashMap::from([(option_type_var.1, option_type_var.0)]),
+        }));
+        let unwrap_fun_ty = unifier.add_ty(TypeEnum::TFunc(FunSignature {
+            args: vec![],
+            ret: option_type_var.0,
+            vars: HashMap::from([(option_type_var.1, option_type_var.0)]),
+        }));
+        let option = unifier.add_ty(TypeEnum::TObj {
+            obj_id: DefinitionId(10),
+            fields: vec![
+                ("is_some".into(), (is_some_type_fun_ty, true)),
+                ("is_none".into(), (is_some_type_fun_ty, true)),
+                ("unwrap".into(), (unwrap_fun_ty, true)),
+            ]
+            .into_iter()
+            .collect::<HashMap<_, _>>(),
+            params: HashMap::from([(option_type_var.1, option_type_var.0)]),
+        });
+
+        let primitives = PrimitiveStore {
+            int32,
+            int64,
+            float,
+            bool,
+            none,
+            range,
+            str,
+            exception,
+            uint32,
+            uint64,
+            option,
+        };
         crate::typecheck::magic_methods::set_primitives_magic_methods(&primitives, &mut unifier);
         (primitives, unifier)
     }

nac3core/src/toplevel/snapshots/nac3core__toplevel__test__test_analyze__generic_class.snap

@@ -1,13 +1,13 @@
 ---
 source: nac3core/src/toplevel/test.rs
-assertion_line: 540
+assertion_line: 549
 expression: res_vec
 
 ---
 [
     "Class {\nname: \"Generic_A\",\nancestors: [\"{class: Generic_A, params: [\\\"V\\\"]}\", \"{class: B, params: []}\"],\nfields: [\"aa\", \"a\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"foo\", \"fn[[b:T], none]\"), (\"fun\", \"fn[[a:int32], V]\")],\ntype_vars: [\"V\"]\n}\n",
-    "Function {\nname: \"Generic_A.__init__\",\nsig: \"fn[[], none]\",\nvar_id: [6]\n}\n",
-    "Function {\nname: \"Generic_A.fun\",\nsig: \"fn[[a:int32], V]\",\nvar_id: [6, 17]\n}\n",
+    "Function {\nname: \"Generic_A.__init__\",\nsig: \"fn[[], none]\",\nvar_id: [7]\n}\n",
+    "Function {\nname: \"Generic_A.fun\",\nsig: \"fn[[a:int32], V]\",\nvar_id: [7, 18]\n}\n",
     "Class {\nname: \"B\",\nancestors: [\"{class: B, params: []}\"],\nfields: [\"aa\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"foo\", \"fn[[b:T], none]\")],\ntype_vars: []\n}\n",
     "Function {\nname: \"B.__init__\",\nsig: \"fn[[], none]\",\nvar_id: []\n}\n",
     "Function {\nname: \"B.foo\",\nsig: \"fn[[b:T], none]\",\nvar_id: []\n}\n",

nac3core/src/toplevel/snapshots/nac3core__toplevel__test__test_analyze__inheritance_override.snap

@@ -1,6 +1,6 @@
 ---
 source: nac3core/src/toplevel/test.rs
-assertion_line: 540
+assertion_line: 549
 expression: res_vec
 
 ---
@@ -9,9 +9,9 @@ expression: res_vec
     "Function {\nname: \"A.__init__\",\nsig: \"fn[[t:T], none]\",\nvar_id: []\n}\n",
     "Function {\nname: \"A.fun\",\nsig: \"fn[[a:int32, b:T], list[virtual[B[bool]]]]\",\nvar_id: []\n}\n",
     "Function {\nname: \"A.foo\",\nsig: \"fn[[c:C], none]\",\nvar_id: []\n}\n",
-    "Class {\nname: \"B\",\nancestors: [\"{class: B, params: [\\\"var6\\\"]}\", \"{class: A, params: [\\\"float\\\"]}\"],\nfields: [\"a\", \"b\", \"c\", \"d\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"fun\", \"fn[[a:int32, b:T], list[virtual[B[bool]]]]\"), (\"foo\", \"fn[[c:C], none]\")],\ntype_vars: [\"var6\"]\n}\n",
-    "Function {\nname: \"B.__init__\",\nsig: \"fn[[], none]\",\nvar_id: [6]\n}\n",
-    "Function {\nname: \"B.fun\",\nsig: \"fn[[a:int32, b:T], list[virtual[B[bool]]]]\",\nvar_id: [6]\n}\n",
+    "Class {\nname: \"B\",\nancestors: [\"{class: B, params: [\\\"var7\\\"]}\", \"{class: A, params: [\\\"float\\\"]}\"],\nfields: [\"a\", \"b\", \"c\", \"d\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"fun\", \"fn[[a:int32, b:T], list[virtual[B[bool]]]]\"), (\"foo\", \"fn[[c:C], none]\")],\ntype_vars: [\"var7\"]\n}\n",
+    "Function {\nname: \"B.__init__\",\nsig: \"fn[[], none]\",\nvar_id: [7]\n}\n",
+    "Function {\nname: \"B.fun\",\nsig: \"fn[[a:int32, b:T], list[virtual[B[bool]]]]\",\nvar_id: [7]\n}\n",
     "Class {\nname: \"C\",\nancestors: [\"{class: C, params: []}\", \"{class: B, params: [\\\"bool\\\"]}\", \"{class: A, params: [\\\"float\\\"]}\"],\nfields: [\"a\", \"b\", \"c\", \"d\", \"e\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"fun\", \"fn[[a:int32, b:T], list[virtual[B[bool]]]]\"), (\"foo\", \"fn[[c:C], none]\")],\ntype_vars: []\n}\n",
     "Function {\nname: \"C.__init__\",\nsig: \"fn[[], none]\",\nvar_id: []\n}\n",
 ]

nac3core/src/toplevel/snapshots/nac3core__toplevel__test__test_analyze__list_tuple_generic.snap

@@ -1,14 +1,14 @@
 ---
 source: nac3core/src/toplevel/test.rs
-assertion_line: 540
+assertion_line: 549
 expression: res_vec
 
 ---
 [
     "Function {\nname: \"foo\",\nsig: \"fn[[a:list[int32], b:tuple[T, float]], A[B, bool]]\",\nvar_id: []\n}\n",
     "Class {\nname: \"A\",\nancestors: [\"{class: A, params: [\\\"T\\\", \\\"V\\\"]}\"],\nfields: [\"a\", \"b\"],\nmethods: [(\"__init__\", \"fn[[v:V], none]\"), (\"fun\", \"fn[[a:T], V]\")],\ntype_vars: [\"T\", \"V\"]\n}\n",
-    "Function {\nname: \"A.__init__\",\nsig: \"fn[[v:V], none]\",\nvar_id: [18, 19]\n}\n",
-    "Function {\nname: \"A.fun\",\nsig: \"fn[[a:T], V]\",\nvar_id: [19, 24]\n}\n",
+    "Function {\nname: \"A.__init__\",\nsig: \"fn[[v:V], none]\",\nvar_id: [19, 20]\n}\n",
+    "Function {\nname: \"A.fun\",\nsig: \"fn[[a:T], V]\",\nvar_id: [20, 25]\n}\n",
     "Function {\nname: \"gfun\",\nsig: \"fn[[a:A[int32, list[float]]], none]\",\nvar_id: []\n}\n",
     "Class {\nname: \"B\",\nancestors: [\"{class: B, params: []}\"],\nfields: [],\nmethods: [(\"__init__\", \"fn[[], none]\")],\ntype_vars: []\n}\n",
     "Function {\nname: \"B.__init__\",\nsig: \"fn[[], none]\",\nvar_id: []\n}\n",

nac3core/src/toplevel/snapshots/nac3core__toplevel__test__test_analyze__self1.snap

@@ -1,13 +1,13 @@
 ---
 source: nac3core/src/toplevel/test.rs
-assertion_line: 540
+assertion_line: 549
 expression: res_vec
 
 ---
 [
-    "Class {\nname: \"A\",\nancestors: [\"{class: A, params: [\\\"var5\\\", \\\"var6\\\"]}\"],\nfields: [\"a\", \"b\"],\nmethods: [(\"__init__\", \"fn[[a:A[bool, float], b:B], none]\"), (\"fun\", \"fn[[a:A[bool, float]], A[bool, int32]]\")],\ntype_vars: [\"var5\", \"var6\"]\n}\n",
-    "Function {\nname: \"A.__init__\",\nsig: \"fn[[a:A[bool, float], b:B], none]\",\nvar_id: [6]\n}\n",
-    "Function {\nname: \"A.fun\",\nsig: \"fn[[a:A[bool, float]], A[bool, int32]]\",\nvar_id: [6]\n}\n",
+    "Class {\nname: \"A\",\nancestors: [\"{class: A, params: [\\\"var6\\\", \\\"var7\\\"]}\"],\nfields: [\"a\", \"b\"],\nmethods: [(\"__init__\", \"fn[[a:A[bool, float], b:B], none]\"), (\"fun\", \"fn[[a:A[bool, float]], A[bool, int32]]\")],\ntype_vars: [\"var6\", \"var7\"]\n}\n",
+    "Function {\nname: \"A.__init__\",\nsig: \"fn[[a:A[bool, float], b:B], none]\",\nvar_id: [7]\n}\n",
+    "Function {\nname: \"A.fun\",\nsig: \"fn[[a:A[bool, float]], A[bool, int32]]\",\nvar_id: [7]\n}\n",
     "Class {\nname: \"B\",\nancestors: [\"{class: B, params: []}\", \"{class: A, params: [\\\"int64\\\", \\\"bool\\\"]}\"],\nfields: [\"a\", \"b\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"fun\", \"fn[[a:A[bool, float]], A[bool, int32]]\"), (\"foo\", \"fn[[b:B], B]\"), (\"bar\", \"fn[[a:A[int32, list[B]]], tuple[A[bool, virtual[A[B, int32]]], B]]\")],\ntype_vars: []\n}\n",
     "Function {\nname: \"B.__init__\",\nsig: \"fn[[], none]\",\nvar_id: []\n}\n",
     "Function {\nname: \"B.foo\",\nsig: \"fn[[b:B], B]\",\nvar_id: []\n}\n",

nac3core/src/toplevel/snapshots/nac3core__toplevel__test__test_analyze__simple_class_compose.snap

@@ -1,6 +1,6 @@
 ---
 source: nac3core/src/toplevel/test.rs
-assertion_line: 540
+assertion_line: 549
 expression: res_vec
 
 ---
@@ -8,12 +8,12 @@ expression: res_vec
     "Class {\nname: \"A\",\nancestors: [\"{class: A, params: []}\"],\nfields: [\"a\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"fun\", \"fn[[b:B], none]\"), (\"foo\", \"fn[[a:T, b:V], none]\")],\ntype_vars: []\n}\n",
     "Function {\nname: \"A.__init__\",\nsig: \"fn[[], none]\",\nvar_id: []\n}\n",
     "Function {\nname: \"A.fun\",\nsig: \"fn[[b:B], none]\",\nvar_id: []\n}\n",
-    "Function {\nname: \"A.foo\",\nsig: \"fn[[a:T, b:V], none]\",\nvar_id: [25]\n}\n",
+    "Function {\nname: \"A.foo\",\nsig: \"fn[[a:T, b:V], none]\",\nvar_id: [26]\n}\n",
     "Class {\nname: \"B\",\nancestors: [\"{class: B, params: []}\", \"{class: C, params: []}\", \"{class: A, params: []}\"],\nfields: [\"a\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"fun\", \"fn[[b:B], none]\"), (\"foo\", \"fn[[a:T, b:V], none]\")],\ntype_vars: []\n}\n",
     "Function {\nname: \"B.__init__\",\nsig: \"fn[[], none]\",\nvar_id: []\n}\n",
     "Class {\nname: \"C\",\nancestors: [\"{class: C, params: []}\", \"{class: A, params: []}\"],\nfields: [\"a\"],\nmethods: [(\"__init__\", \"fn[[], none]\"), (\"fun\", \"fn[[b:B], none]\"), (\"foo\", \"fn[[a:T, b:V], none]\")],\ntype_vars: []\n}\n",
     "Function {\nname: \"C.__init__\",\nsig: \"fn[[], none]\",\nvar_id: []\n}\n",
     "Function {\nname: \"C.fun\",\nsig: \"fn[[b:B], none]\",\nvar_id: []\n}\n",
     "Function {\nname: \"foo\",\nsig: \"fn[[a:A], none]\",\nvar_id: []\n}\n",
-    "Function {\nname: \"ff\",\nsig: \"fn[[a:T], V]\",\nvar_id: [33]\n}\n",
+    "Function {\nname: \"ff\",\nsig: \"fn[[a:T], V]\",\nvar_id: [34]\n}\n",
 ]

nac3core/src/toplevel/type_annotation.rs

@@ -72,8 +72,6 @@ pub fn parse_ast_to_type_annotation_kinds<T>(
             Ok(TypeAnnotation::Primitive(primitives.float))
         } else if id == &"bool".into() {
             Ok(TypeAnnotation::Primitive(primitives.bool))
-        } else if id == &"None".into() {
-            Ok(TypeAnnotation::Primitive(primitives.none))
         } else if id == &"str".into() {
             Ok(TypeAnnotation::Primitive(primitives.str))
         } else if id == &"Exception".into() {
@@ -223,6 +221,29 @@ pub fn parse_ast_to_type_annotation_kinds<T>(
             Ok(TypeAnnotation::List(def_ann.into()))
         }
 
+        // option
+        ast::ExprKind::Subscript { value, slice, .. }
+            if {
+                matches!(&value.node, ast::ExprKind::Name { id, .. } if id == &"Option".into())
+            } =>
+        {
+            let def_ann = parse_ast_to_type_annotation_kinds(
+                resolver,
+                top_level_defs,
+                unifier,
+                primitives,
+                slice.as_ref(),
+                locked,
+            )?;
+            let id =
+                if let TypeEnum::TObj { obj_id, .. } = unifier.get_ty(primitives.option).as_ref() {
+                    *obj_id
+                } else {
+                    unreachable!()
+                };
+            Ok(TypeAnnotation::CustomClass { id, params: vec![def_ann] })
+        }
+
         // tuple
         ast::ExprKind::Subscript { value, slice, .. }
             if {

nac3core/src/typecheck/type_inferencer/mod.rs

@@ -40,6 +40,7 @@ pub struct PrimitiveStore {
     pub range: Type,
     pub str: Type,
     pub exception: Type,
+    pub option: Type,
 }
 
 pub struct FunctionData {
@@ -932,6 +933,17 @@ impl<'a> Inferencer<'a> {
                 Ok(self.unifier.add_ty(TypeEnum::TTuple { ty: ty? }))
             }
             ast::Constant::Str(_) => Ok(self.primitives.str),
+            ast::Constant::None => {
+                let option_ty = self.primitives.option;
+                let new_mapping = if let TypeEnum::TObj { params, .. } = &*self.unifier.get_ty_immutable(option_ty) {
+                    let (id, _) = params.iter().next().unwrap();
+                    // None can be Option[Any]
+                    vec![(*id, self.unifier.get_fresh_var(None, None).0)].into_iter().collect()
+                } else {
+                    unreachable!("option must be tobj")
+                };
+                Ok(self.unifier.subst(option_ty, &new_mapping).unwrap())
+            }
             _ => report_error("not supported", *loc),
         }
     }

nac3core/src/typecheck/type_inferencer/test.rs

@@ -129,7 +129,24 @@ impl TestEnvironment {
             fields: HashMap::new(),
             params: HashMap::new(),
         });
-        let primitives = PrimitiveStore { int32, int64, float, bool, none, range, str, exception, uint32, uint64 };
+        let option = unifier.add_ty(TypeEnum::TObj {
+            obj_id: DefinitionId(10),
+            fields: HashMap::new(),
+            params: HashMap::new(),
+        });
+        let primitives = PrimitiveStore {
+            int32,
+            int64,
+            float,
+            bool,
+            none,
+            range,
+            str,
+            exception,
+            uint32,
+            uint64,
+            option,
+        };
         set_primitives_magic_methods(&primitives, &mut unifier);
 
         let id_to_name = [
@@ -237,6 +254,11 @@ impl TestEnvironment {
             fields: HashMap::new(),
             params: HashMap::new(),
         });
+        let option = unifier.add_ty(TypeEnum::TObj {
+            obj_id: DefinitionId(10),
+            fields: HashMap::new(),
+            params: HashMap::new(),
+        });
         identifier_mapping.insert("None".into(), none);
         for (i, name) in ["int32", "int64", "float", "bool", "none", "range", "str", "Exception"]
             .iter()
@@ -259,7 +281,19 @@ impl TestEnvironment {
         }
         let defs = 7;
 
-        let primitives = PrimitiveStore { int32, int64, float, bool, none, range, str, exception, uint32, uint64 };
+        let primitives = PrimitiveStore {
+            int32,
+            int64,
+            float,
+            bool,
+            none,
+            range,
+            str,
+            exception,
+            uint32,
+            uint64,
+            option,
+        };
 
         let (v0, id) = unifier.get_dummy_var();
 

nac3core/src/typecheck/typedef/mod.rs

@@ -54,6 +54,18 @@ pub enum RecordKey {
     Int(i32),
 }
 
+impl Type {
+    // a wrapper function for cleaner code so that we don't need to
+    // write this long pattern matching just to get the field `obj_id`
+    pub fn get_obj_id(self, unifier: &Unifier) -> DefinitionId {
+        if let TypeEnum::TObj { obj_id, .. } = unifier.get_ty_immutable(self).as_ref() {
+            *obj_id
+        } else {
+            unreachable!("expect a object type")
+        }
+    }
+}
+
 impl From<&RecordKey> for StrRef {
     fn from(r: &RecordKey) -> Self {
         match r {