WIP

nac3artiq/src/codegen.rs

@@ -1,9 +1,12 @@
 use nac3core::{
     codegen::{
-        classes::{ListValue, NDArrayValue, RangeValue, UntypedArrayLikeAccessor},
+        classes::{
+            ArrayLikeIndexer, ArrayLikeValue, ArraySliceValue, ListValue, NDArrayType,
+            NDArrayValue, RangeValue, UntypedArrayLikeAccessor,
+        },
         expr::{destructure_range, gen_call},
         irrt::call_ndarray_calc_size,
-        llvm_intrinsics::{call_int_smax, call_stackrestore, call_stacksave},
+        llvm_intrinsics::{call_int_smax, call_memcpy_generic, call_stackrestore, call_stacksave},
         stmt::{gen_block, gen_for_callback_incrementing, gen_if_callback, gen_with},
         CodeGenContext, CodeGenerator,
     },
@@ -17,8 +20,8 @@ use nac3parser::ast::{Expr, ExprKind, Located, Stmt, StmtKind, StrRef};
 use inkwell::{
     context::Context,
     module::Linkage,
-    types::IntType,
-    values::{BasicValueEnum, StructValue},
+    types::{BasicType, IntType},
+    values::{BasicValueEnum, PointerValue, StructValue},
     AddressSpace, IntPredicate,
 };
 
@@ -422,7 +425,10 @@ fn gen_rpc_tag(
                 } else {
                     unreachable!()
                 };
-                assert!((0u64..=u64::from(u8::MAX)).contains(&ndarray_ndims));
+                assert!(
+                    (0u64..=u64::from(u8::MAX)).contains(&ndarray_ndims),
+                    "Only NDArrays of sizes between 0 and 255 can be RPCed"
+                );
 
                 buffer.push(b'a');
                 buffer.push((ndarray_ndims & 0xFF) as u8);
@@ -434,6 +440,95 @@ fn gen_rpc_tag(
     Ok(())
 }
 
+/// Formats an RPC argument to conform to the expected format required by `send_value`.
+///
+/// See `artiq/firmware/libproto_artiq/rpc_proto.rs` for the expected format.
+fn format_rpc_arg<'ctx>(
+    generator: &mut dyn CodeGenerator,
+    ctx: &mut CodeGenContext<'ctx, '_>,
+    (arg, arg_ty, arg_idx): (BasicValueEnum<'ctx>, Type, usize),
+) -> PointerValue<'ctx> {
+    let llvm_i8 = ctx.ctx.i8_type();
+    let llvm_pi8 = llvm_i8.ptr_type(AddressSpace::default());
+
+    let arg_slot = match &*ctx.unifier.get_ty_immutable(arg_ty) {
+        TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::NDArray.id() => {
+            // NAC3: NDArray = { usize, usize*, T* }
+            // libproto_artiq: NDArray = [data[..], dim_sz[..]]
+
+            let llvm_i1 = ctx.ctx.bool_type();
+            let llvm_usize = generator.get_size_type(ctx.ctx);
+
+            let (elem_ty, _) = unpack_ndarray_var_tys(&mut ctx.unifier, arg_ty);
+            let llvm_arg_ty =
+                NDArrayType::new(generator, ctx.ctx, ctx.get_llvm_type(generator, elem_ty));
+            let llvm_arg = NDArrayValue::from_ptr_val(arg.into_pointer_value(), llvm_usize, None);
+
+            let llvm_usize_sizeof = ctx
+                .builder
+                .build_int_truncate_or_bit_cast(llvm_arg_ty.size_type().size_of(), llvm_usize, "")
+                .unwrap();
+            let llvm_pdata_sizeof = ctx
+                .builder
+                .build_int_truncate_or_bit_cast(
+                    llvm_arg_ty.element_type().ptr_type(AddressSpace::default()).size_of(),
+                    llvm_usize,
+                    "",
+                )
+                .unwrap();
+
+            let dims_buf_sz =
+                ctx.builder.build_int_mul(llvm_arg.load_ndims(ctx), llvm_usize_sizeof, "").unwrap();
+
+            let buffer_size =
+                ctx.builder.build_int_add(dims_buf_sz, llvm_pdata_sizeof, "").unwrap();
+
+            let buffer = ctx.builder.build_array_alloca(llvm_i8, buffer_size, "rpc.arg").unwrap();
+            let buffer = ArraySliceValue::from_ptr_val(buffer, buffer_size, Some("rpc.arg"));
+
+            let ppdata =
+                generator.gen_var_alloc(ctx, llvm_arg_ty.element_type(), None).unwrap();
+            ctx.builder.build_store(ppdata, llvm_arg.data().base_ptr(ctx, generator)).unwrap();
+
+            call_memcpy_generic(
+                ctx,
+                buffer.base_ptr(ctx, generator),
+                ppdata,
+                llvm_pdata_sizeof,
+                llvm_i1.const_zero(),
+            );
+
+            let pbuffer_dims_begin =
+                unsafe { buffer.ptr_offset_unchecked(ctx, generator, &llvm_pdata_sizeof, None) };
+            call_memcpy_generic(
+                ctx,
+                pbuffer_dims_begin,
+                llvm_arg.dim_sizes().base_ptr(ctx, generator),
+                dims_buf_sz,
+                llvm_i1.const_zero(),
+            );
+
+            buffer.base_ptr(ctx, generator)
+        }
+
+        _ => {
+            let arg_slot = generator
+                .gen_var_alloc(ctx, arg.get_type(), Some(&format!("rpc.arg{arg_idx}")))
+                .unwrap();
+            ctx.builder.build_store(arg_slot, arg).unwrap();
+
+            ctx.builder
+                .build_bitcast(arg_slot, llvm_pi8, "rpc.arg")
+                .map(BasicValueEnum::into_pointer_value)
+                .unwrap()
+        }
+    };
+
+    debug_assert_eq!(arg_slot.get_type(), llvm_pi8);
+
+    arg_slot
+}
+
 fn rpc_codegen_callback_fn<'ctx>(
     ctx: &mut CodeGenContext<'ctx, '_>,
     obj: Option<(Type, ValueEnum<'ctx>)>,
@@ -441,10 +536,10 @@ fn rpc_codegen_callback_fn<'ctx>(
     args: Vec<(Option<StrRef>, ValueEnum<'ctx>)>,
     generator: &mut dyn CodeGenerator,
 ) -> Result<Option<BasicValueEnum<'ctx>>, String> {
-    let ptr_type = ctx.ctx.i8_type().ptr_type(AddressSpace::default());
-    let size_type = generator.get_size_type(ctx.ctx);
     let int8 = ctx.ctx.i8_type();
     let int32 = ctx.ctx.i32_type();
+    let size_type = generator.get_size_type(ctx.ctx);
+    let ptr_type = int8.ptr_type(AddressSpace::default());
     let tag_ptr_type = ctx.ctx.struct_type(&[ptr_type.into(), size_type.into()], false);
 
     let service_id = int32.const_int(fun.1 .0 as u64, false);
@@ -517,22 +612,25 @@ fn rpc_codegen_callback_fn<'ctx>(
         .0
         .args
         .iter()
-        .map(|arg| mapping.remove(&arg.name).unwrap().to_basic_value_enum(ctx, generator, arg.ty))
-        .collect::<Result<Vec<_>, _>>()?;
+        .map(|arg| {
+            mapping
+                .remove(&arg.name)
+                .unwrap()
+                .to_basic_value_enum(ctx, generator, arg.ty)
+                .map(|llvm_val| (llvm_val, arg.ty))
+        })
+        .collect::<Result<Vec<(_, _)>, _>>()?;
     if let Some(obj) = obj {
-        if let ValueEnum::Static(obj) = obj.1 {
-            real_params.insert(0, obj.get_const_obj(ctx, generator));
+        if let ValueEnum::Static(obj_val) = obj.1 {
+            real_params.insert(0, (obj_val.get_const_obj(ctx, generator), obj.0));
         } else {
             // should be an error here...
             panic!("only host object is allowed");
         }
     }
 
-    for (i, arg) in real_params.iter().enumerate() {
-        let arg_slot =
-            generator.gen_var_alloc(ctx, arg.get_type(), Some(&format!("rpc.arg{i}"))).unwrap();
-        ctx.builder.build_store(arg_slot, *arg).unwrap();
-        let arg_slot = ctx.builder.build_bitcast(arg_slot, ptr_type, "rpc.arg").unwrap();
+    for (i, (arg, arg_ty)) in real_params.iter().enumerate() {
+        let arg_slot = format_rpc_arg(generator, ctx, (*arg, *arg_ty, i));
         let arg_ptr = unsafe {
             ctx.builder.build_gep(
                 args_ptr,

nac3core/src/codegen/expr.rs

@@ -3025,14 +3025,19 @@ pub fn gen_expr<'ctx, G: CodeGenerator>(
                     let Some(val) = generator.gen_expr(ctx, value)? else { return Ok(None) };
 
                     // Handle Class Method calls
-                    let id = if let TypeEnum::TObj { obj_id, .. } =
-                        &*ctx.unifier.get_ty(value.custom.unwrap())
-                    {
+                    // The attribute will be `DefinitionId` of the method if the call is to one of the parent methods
+                    let func_id = attr.to_string().parse::<usize>(); 
+
+                    let id = if let TypeEnum::TObj { obj_id, .. } = &*ctx.unifier.get_ty(value.custom.unwrap()) {
                         *obj_id
                     } else {
                         unreachable!()
                     };
-                    let fun_id = {
+
+                    // Use the `DefinitionID` from attribute if it is available
+                    let fun_id = if func_id.is_ok() {
+                        DefinitionId(func_id.unwrap())
+                    } else {
                         let defs = ctx.top_level.definitions.read();
                         let obj_def = defs.get(id.0).unwrap().read();
                         let TopLevelDef::Class { methods, .. } = &*obj_def else { unreachable!() };

nac3core/src/toplevel/composer.rs

@@ -23,7 +23,7 @@ impl Default for ComposerConfig {
     }
 }
 
-type DefAst = (Arc<RwLock<TopLevelDef>>, Option<Stmt<()>>);
+pub type DefAst = (Arc<RwLock<TopLevelDef>>, Option<Stmt<()>>);
 pub struct TopLevelComposer {
     // list of top level definitions, same as top level context
     pub definition_ast_list: Vec<DefAst>,
@@ -1822,7 +1822,15 @@ impl TopLevelComposer {
                         if *name != init_str_id {
                             unreachable!("must be init function here")
                         }
-                        let all_inited = Self::get_all_assigned_field(body.as_slice())?;
+                        // Since AST stores class names without prepending `__module__.`, we split the name for search purposes
+                        let class_name_only = class_name.to_string();
+                        let (_, class_name_only) = class_name_only.split_once('.').unwrap();
+                        
+                        let all_inited = Self::get_all_assigned_field(
+                            class_name_only.into(),
+                            definition_ast_list,
+                            body.as_slice(),
+                        )?;
                         for (f, _, _) in fields {
                             if !all_inited.contains(f) {
                                 return Err(HashSet::from([

nac3core/src/toplevel/helper.rs

@@ -3,6 +3,7 @@ use std::convert::TryInto;
 use crate::symbol_resolver::SymbolValue;
 use crate::toplevel::numpy::unpack_ndarray_var_tys;
 use crate::typecheck::typedef::{into_var_map, iter_type_vars, Mapping, TypeVarId, VarMap};
+use ast::ExprKind;
 use nac3parser::ast::{Constant, Location};
 use strum::IntoEnumIterator;
 use strum_macros::EnumIter;
@@ -733,7 +734,11 @@ impl TopLevelComposer {
         )
     }
 
-    pub fn get_all_assigned_field(stmts: &[Stmt<()>]) -> Result<HashSet<StrRef>, HashSet<String>> {
+    pub fn get_all_assigned_field(
+        class_name: StrRef,
+        ast: &Vec<DefAst>,
+        stmts: &[Stmt<()>],
+    ) -> Result<HashSet<StrRef>, HashSet<String>> {
         let mut result = HashSet::new();
         for s in stmts {
             match &s.node {
@@ -769,30 +774,106 @@ impl TopLevelComposer {
                 // TODO: do not check for For and While?
                 ast::StmtKind::For { body, orelse, .. }
                 | ast::StmtKind::While { body, orelse, .. } => {
-                    result.extend(Self::get_all_assigned_field(body.as_slice())?);
-                    result.extend(Self::get_all_assigned_field(orelse.as_slice())?);
+                    result.extend(Self::get_all_assigned_field(class_name, ast, body.as_slice())?);
+                    result.extend(Self::get_all_assigned_field(
+                        class_name,
+                        ast,
+                        orelse.as_slice(),
+                    )?);
                 }
                 ast::StmtKind::If { body, orelse, .. } => {
-                    let inited_for_sure = Self::get_all_assigned_field(body.as_slice())?
-                        .intersection(&Self::get_all_assigned_field(orelse.as_slice())?)
-                        .copied()
-                        .collect::<HashSet<_>>();
+                    let inited_for_sure =
+                        Self::get_all_assigned_field(class_name, ast, body.as_slice())?
+                            .intersection(&Self::get_all_assigned_field(
+                                class_name,
+                                ast,
+                                orelse.as_slice(),
+                            )?)
+                            .copied()
+                            .collect::<HashSet<_>>();
                     result.extend(inited_for_sure);
                 }
                 ast::StmtKind::Try { body, orelse, finalbody, .. } => {
-                    let inited_for_sure = Self::get_all_assigned_field(body.as_slice())?
-                        .intersection(&Self::get_all_assigned_field(orelse.as_slice())?)
-                        .copied()
-                        .collect::<HashSet<_>>();
+                    let inited_for_sure =
+                        Self::get_all_assigned_field(class_name, ast, body.as_slice())?
+                            .intersection(&Self::get_all_assigned_field(
+                                class_name,
+                                ast,
+                                orelse.as_slice(),
+                            )?)
+                            .copied()
+                            .collect::<HashSet<_>>();
                     result.extend(inited_for_sure);
-                    result.extend(Self::get_all_assigned_field(finalbody.as_slice())?);
+                    result.extend(Self::get_all_assigned_field(
+                        class_name,
+                        ast,
+                        finalbody.as_slice(),
+                    )?);
                 }
                 ast::StmtKind::With { body, .. } => {
-                    result.extend(Self::get_all_assigned_field(body.as_slice())?);
+                    result.extend(Self::get_all_assigned_field(class_name, ast, body.as_slice())?);
+                }
+                // Variables Initiated in function calls
+                ast::StmtKind::Expr { value, .. } => {
+                    let ExprKind::Call { func, .. } = &value.node else {
+                        continue;
+                    };
+                    let ExprKind::Attribute { value, attr, .. } = &func.node else {
+                        continue;
+                    };
+                    let ExprKind::Name { id, .. } = &value.node else {
+                        continue;
+                    };
+                    // Need to conside the two cases:
+                    //  Case 1) Call to class function i.e. id = `self`
+                    //  Case 2) Call to class ancestor function i.e. id = ancestor_name
+                    // We leave checking whether function in case 2 belonged to class ancestor or not to type checker
+                    //
+                    // According to current handling of `self`, function definition are fixed and do not change regardless
+                    // of which object is passed as `self` i.e. virtual polymorphism is not supported
+                    // Therefore, we change class name for case 2 to reflect behavior of our compiler
+                    let new_class_name = if *id == "self".into() { class_name } else { *id };
+
+                    let method_body = ast.iter().find_map(|def| {
+                        let Some(ast::Located {
+                            node: ast::StmtKind::ClassDef { name, body, .. },
+                            ..
+                        }) = &def.1
+                        else {
+                            return None;
+                        };
+                        if *name == new_class_name {
+                            body.iter().find_map(|m| {
+                                let ast::StmtKind::FunctionDef { name, body, .. } = &m.node else {
+                                    return None;
+                                };
+                                if *name == *attr {
+                                    return Some(body.clone());
+                                }
+                                None
+                            })
+                        } else {
+                            None
+                        }
+                    });
+
+                    // If method body is none then method does not exist
+                    if let Some(method_body) = method_body {
+                        result.extend(Self::get_all_assigned_field(
+                            new_class_name,
+                            ast,
+                            method_body.as_slice(),
+                        )?);
+                    } else {
+                        return Err(HashSet::from([format!(
+                            "{}.{} not found in class {new_class_name} at {}",
+                            *id, *attr, value.location
+                        )]));
+                    }
                 }
                 ast::StmtKind::Pass { .. }
                 | ast::StmtKind::Assert { .. }
-                | ast::StmtKind::Expr { .. } => {}
+                | ast::StmtKind::AnnAssign { .. } => {}
 
                 _ => {
                     unimplemented!()

nac3core/src/typecheck/type_inferencer/mod.rs

@@ -12,6 +12,7 @@ use super::{
         RecordField, RecordKey, Type, TypeEnum, TypeVar, Unifier, VarMap,
     },
 };
+use crate::toplevel::type_annotation::TypeAnnotation;
 use crate::{
     symbol_resolver::{SymbolResolver, SymbolValue},
     toplevel::{
@@ -102,6 +103,7 @@ pub struct Inferencer<'a> {
 }
 
 type InferenceError = HashSet<String>;
+type OverrideResult = Result<Option<ast::Expr<Option<Type>>>, InferenceError>;
 
 struct NaiveFolder();
 impl Fold<()> for NaiveFolder {
@@ -1672,6 +1674,86 @@ impl<'a> Inferencer<'a> {
         Ok(None)
     }
 
+    /// Checks whether a class method is calling parent function
+    /// Returns [`None`] if its not a call to parent method, otherwise
+    /// returns a new `func` with class name replaced by `self` and method resolved to its `DefinitionID`
+    ///
+    /// e.g.  A.f1(self, ...) returns Some(self.DefintionID(f1))
+    fn check_overriding(&mut self, func: &ast::Expr<()>, args: &[ast::Expr<()>]) -> OverrideResult {
+        // `self` must be first argument for call to parent method
+        if let Some(Located { node: ExprKind::Name { id, .. }, .. }) = &args.first() {
+            if *id != "self".into() {
+                return Ok(None);
+            }
+        } else {
+            return Ok(None);
+        }
+
+        let Located {
+            node: ExprKind::Attribute { value, attr: method_name, ctx }, location, ..
+        } = func
+        else {
+            return Ok(None);
+        };
+        let ExprKind::Name { id: class_name, ctx: class_ctx } = &value.node else {
+            return Ok(None);
+        };
+        let zelf = &self.fold_expr(args[0].clone())?;
+
+        // Check whether the method belongs to class ancestors
+        let def_id = self.unifier.get_ty(zelf.custom.unwrap());
+        let TypeEnum::TObj { obj_id, .. } = def_id.as_ref() else { unreachable!() };
+        let defs = self.top_level.definitions.read();
+        let res = {
+            if let TopLevelDef::Class { ancestors, .. } = &*defs[obj_id.0].read() {
+                let res = ancestors.iter().find_map(|f| {
+                    let TypeAnnotation::CustomClass { id, .. } = f else { unreachable!() };
+                    let TopLevelDef::Class { name, methods, .. } = &*defs[id.0].read() else {
+                        unreachable!()
+                    };
+                    // Class names are stored as `__module__.class`
+                    let name = name.to_string();
+                    let (_, name) = name.split_once('.').unwrap();
+                    if name == class_name.to_string() {
+                        return methods.iter().find_map(|f| {
+                            if f.0 == *method_name {
+                                return Some(*f);
+                            }
+                            None
+                        });
+                    }
+                    None
+                });
+                res
+            } else {
+                None
+            }
+        };
+
+        match res {
+            Some(r) => {
+
+                let mut new_func = func.clone();
+                let mut new_value = value.clone();
+                new_value.node = ExprKind::Name { id: "self".into(), ctx: *class_ctx };
+                new_func.node =
+                    ExprKind::Attribute { value: new_value.clone(), attr: *method_name, ctx: *ctx };
+
+                let mut new_func = self.fold_expr(new_func)?;
+                
+                let ExprKind::Attribute { value, .. } = new_func.node else { unreachable!() };
+                new_func.node = ExprKind::Attribute { value, attr: r.2.0.to_string().into(), ctx: *ctx };
+                new_func.custom = Some(r.1);
+
+                Ok(Some(new_func))
+            }
+            None => report_error(
+                format!("Method {class_name}.{method_name} not found in ancestor list").as_str(),
+                *location,
+            ),
+        }
+    }
+
     fn fold_call(
         &mut self,
         location: Location,
@@ -1685,8 +1767,20 @@ impl<'a> Inferencer<'a> {
             return Ok(spec_call_func);
         }
 
-        let func = Box::new(self.fold_expr(func)?);
-        let args = args.into_iter().map(|v| self.fold_expr(v)).collect::<Result<Vec<_>, _>>()?;
+        // Check for call to parent method
+        let override_res = self.check_overriding(&func, &args)?;
+        let is_override = override_res.is_some();
+        let func = if is_override { override_res.unwrap() } else { self.fold_expr(func)? };
+        let func = Box::new(func);
+
+        let mut args =
+            args.into_iter().map(|v| self.fold_expr(v)).collect::<Result<Vec<_>, _>>()?;
+
+        // TODO: Handle passing of self to functions to allow runtime lookup of functions to be called
+        // Currently removing `self` and using compile time function definitions
+        if is_override {
+            args.remove(0);
+        }
         let keywords = keywords
             .into_iter()
             .map(|v| fold::fold_keyword(self, v))

nac3standalone/demo/src/inheritance.py

@@ -10,23 +10,58 @@ class A:
     def __init__(self, a: int32):
         self.a = a
     
-    def f1(self):
-        self.f2()
-    
-    def f2(self):
+    def output_all_fields(self):
         output_int32(self.a)
+    
+    def set_a(self, a: int32):
+        self.a = a
 
 class B(A):
     b: int32
 
     def __init__(self, b: int32):
-        self.a = b + 1
+        A.__init__(self, b + 1)
+        self.set_b(b)
+    
+    def output_parent_fields(self):
+        A.output_all_fields(self)
+
+    def output_all_fields(self):
+        A.output_all_fields(self)
+        output_int32(self.b)
+    
+    def set_b(self, b: int32):
         self.b = b
 
+class C(B):
+    c: int32
+
+    def __init__(self, c: int32):
+        B.__init__(self, c + 1)
+        self.c = c
+    
+    def output_parent_fields(self):
+        B.output_all_fields(self)
+
+    def output_all_fields(self):
+        B.output_all_fields(self)
+        output_int32(self.c)
+    
+    def set_c(self, c: int32):
+        self.c = c
 
 def run() -> int32:
-    aaa = A(5)
-    bbb = B(2)
-    aaa.f1()
-    bbb.f1()
+    ccc = C(10)
+    ccc.output_all_fields()
+    ccc.set_a(1)
+    ccc.set_b(2)
+    ccc.set_c(3)
+    ccc.output_all_fields()
+
+    bbb = B(10)
+    bbb.set_a(9)
+    bbb.set_b(8)
+    bbb.output_all_fields()
+    ccc.output_all_fields()
+
     return 0