core/typecheck: Support tuple arg type in len()

nac3core/src/codegen/builtin_fns.rs

@@ -1,17 +1,20 @@
 use inkwell::types::BasicTypeEnum;
-use inkwell::values::{BasicValue, BasicValueEnum, PointerValue};
+use inkwell::values::{BasicValue, BasicValueEnum, IntValue, PointerValue};
 use inkwell::{FloatPredicate, IntPredicate, OptimizationLevel};
 use itertools::Itertools;
 
 use crate::codegen::classes::{
-    NDArrayValue, ProxyValue, UntypedArrayLikeAccessor, UntypedArrayLikeMutator,
+    ArrayLikeValue, NDArrayValue, ProxyValue, RangeValue, TypedArrayLikeAccessor,
+    UntypedArrayLikeAccessor, UntypedArrayLikeMutator,
 };
+use crate::codegen::expr::destructure_range;
+use crate::codegen::irrt::calculate_len_for_slice_range;
 use crate::codegen::numpy::ndarray_elementwise_unaryop_impl;
 use crate::codegen::stmt::gen_for_callback_incrementing;
 use crate::codegen::{extern_fns, irrt, llvm_intrinsics, numpy, CodeGenContext, CodeGenerator};
 use crate::toplevel::helper::PrimDef;
 use crate::toplevel::numpy::unpack_ndarray_var_tys;
-use crate::typecheck::typedef::Type;
+use crate::typecheck::typedef::{Type, TypeEnum};
 
 /// Shorthand for [`unreachable!()`] when a type of argument is not supported.
 ///
@@ -23,6 +26,67 @@ fn unsupported_type(ctx: &CodeGenContext<'_, '_>, fn_name: &str, tys: &[Type]) -
     )
 }
 
+/// Invokes the `len` builtin function.
+pub fn call_len<'ctx, G: CodeGenerator + ?Sized>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, '_>,
+    n: (Type, BasicValueEnum<'ctx>),
+) -> Result<IntValue<'ctx>, String> {
+    let llvm_i32 = ctx.ctx.i32_type();
+    let range_ty = ctx.primitives.range;
+    let (arg_ty, arg) = n;
+
+    Ok(if ctx.unifier.unioned(arg_ty, range_ty) {
+        let arg = RangeValue::from_ptr_val(arg.into_pointer_value(), Some("range"));
+        let (start, end, step) = destructure_range(ctx, arg);
+        calculate_len_for_slice_range(generator, ctx, start, end, step)
+    } else {
+        match &*ctx.unifier.get_ty_immutable(arg_ty) {
+            TypeEnum::TTuple { ty, .. } => llvm_i32.const_int(ty.len() as u64, false),
+            TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::List.id() => {
+                let zero = llvm_i32.const_zero();
+                let len = ctx
+                    .build_gep_and_load(
+                        arg.into_pointer_value(),
+                        &[zero, llvm_i32.const_int(1, false)],
+                        None,
+                    )
+                    .into_int_value();
+                ctx.builder.build_int_truncate_or_bit_cast(len, llvm_i32, "len").unwrap()
+            }
+            TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::NDArray.id() => {
+                let llvm_usize = generator.get_size_type(ctx.ctx);
+
+                let arg = NDArrayValue::from_ptr_val(arg.into_pointer_value(), llvm_usize, None);
+
+                let ndims = arg.dim_sizes().size(ctx, generator);
+                ctx.make_assert(
+                    generator,
+                    ctx.builder
+                        .build_int_compare(IntPredicate::NE, ndims, llvm_usize.const_zero(), "")
+                        .unwrap(),
+                    "0:TypeError",
+                    "len() of unsized object",
+                    [None, None, None],
+                    ctx.current_loc,
+                );
+
+                let len = unsafe {
+                    arg.dim_sizes().get_typed_unchecked(
+                        ctx,
+                        generator,
+                        &llvm_usize.const_zero(),
+                        None,
+                    )
+                };
+
+                ctx.builder.build_int_truncate_or_bit_cast(len, llvm_i32, "len").unwrap()
+            }
+            _ => unreachable!(),
+        }
+    })
+}
+
 /// Invokes the `int32` builtin function.
 pub fn call_int32<'ctx, G: CodeGenerator + ?Sized>(
     generator: &mut G,

nac3core/src/toplevel/builtins.rs

@@ -14,9 +14,7 @@ use strum::IntoEnumIterator;
 use crate::{
     codegen::{
         builtin_fns,
-        classes::{ArrayLikeValue, NDArrayValue, ProxyValue, RangeValue, TypedArrayLikeAccessor},
-        expr::destructure_range,
-        irrt::*,
+        classes::{ProxyValue, RangeValue},
         numpy::*,
         stmt::exn_constructor,
     },
@@ -1466,36 +1464,21 @@ impl<'a> BuiltinBuilder<'a> {
     fn build_len_function(&mut self) -> TopLevelDef {
         let prim = PrimDef::FunLen;
 
-        let PrimitiveStore { uint64, int32, .. } = *self.primitives;
+        // Type handled in [`Inferencer::try_fold_special_call`]
+        let arg_tvar = self.unifier.get_dummy_var();
 
-        let tvar = self.unifier.get_fresh_var(Some("L".into()), None);
-        let list = self
-            .unifier
-            .subst(
-                self.primitives.list,
-                &into_var_map([TypeVar { id: self.list_tvar.id, ty: tvar.ty }]),
-            )
-            .unwrap();
-        let ndims = self.unifier.get_fresh_const_generic_var(uint64, Some("N".into()), None);
-        let ndarray = make_ndarray_ty(self.unifier, self.primitives, Some(tvar.ty), Some(ndims.ty));
-
-        let arg_ty = self.unifier.get_fresh_var_with_range(
-            &[list, ndarray, self.primitives.range],
-            Some("I".into()),
-            None,
-        );
         TopLevelDef::Function {
             name: prim.name().into(),
             simple_name: prim.simple_name().into(),
             signature: self.unifier.add_ty(TypeEnum::TFunc(FunSignature {
                 args: vec![FuncArg {
-                    name: "ls".into(),
-                    ty: arg_ty.ty,
+                    name: "obj".into(),
+                    ty: arg_tvar.ty,
                     default_value: None,
                     is_vararg: false,
                 }],
-                ret: int32,
-                vars: into_var_map([tvar, arg_ty]),
+                ret: self.primitives.int32,
+                vars: into_var_map([arg_tvar]),
             })),
             var_id: Vec::default(),
             instance_to_symbol: HashMap::default(),
@@ -1503,86 +1486,10 @@ impl<'a> BuiltinBuilder<'a> {
             resolver: None,
             codegen_callback: Some(Arc::new(GenCall::new(Box::new(
                 move |ctx, _, fun, args, generator| {
-                    let range_ty = ctx.primitives.range;
                     let arg_ty = fun.0.args[0].ty;
                     let arg = args[0].1.clone().to_basic_value_enum(ctx, generator, arg_ty)?;
-                    Ok(if ctx.unifier.unioned(arg_ty, range_ty) {
-                        let arg = RangeValue::from_ptr_val(arg.into_pointer_value(), Some("range"));
-                        let (start, end, step) = destructure_range(ctx, arg);
-                        Some(calculate_len_for_slice_range(generator, ctx, start, end, step).into())
-                    } else {
-                        match &*ctx.unifier.get_ty_immutable(arg_ty) {
-                            TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::List.id() => {
-                                let int32 = ctx.ctx.i32_type();
-                                let zero = int32.const_zero();
-                                let len = ctx
-                                    .build_gep_and_load(
-                                        arg.into_pointer_value(),
-                                        &[zero, int32.const_int(1, false)],
-                                        None,
-                                    )
-                                    .into_int_value();
-                                if len.get_type().get_bit_width() == 32 {
-                                    Some(len.into())
-                                } else {
-                                    Some(
-                                        ctx.builder
-                                            .build_int_truncate(len, int32, "len2i32")
-                                            .map(Into::into)
-                                            .unwrap(),
-                                    )
-                                }
-                            }
-                            TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::NDArray.id() => {
-                                let llvm_i32 = ctx.ctx.i32_type();
-                                let llvm_usize = generator.get_size_type(ctx.ctx);
 
-                                let arg = NDArrayValue::from_ptr_val(
-                                    arg.into_pointer_value(),
-                                    llvm_usize,
-                                    None,
-                                );
-
-                                let ndims = arg.dim_sizes().size(ctx, generator);
-                                ctx.make_assert(
-                                    generator,
-                                    ctx.builder
-                                        .build_int_compare(
-                                            IntPredicate::NE,
-                                            ndims,
-                                            llvm_usize.const_zero(),
-                                            "",
-                                        )
-                                        .unwrap(),
-                                    "0:TypeError",
-                                    &format!("{name}() of unsized object", name = prim.name()),
-                                    [None, None, None],
-                                    ctx.current_loc,
-                                );
-
-                                let len = unsafe {
-                                    arg.dim_sizes().get_typed_unchecked(
-                                        ctx,
-                                        generator,
-                                        &llvm_usize.const_zero(),
-                                        None,
-                                    )
-                                };
-
-                                if len.get_type().get_bit_width() == 32 {
-                                    Some(len.into())
-                                } else {
-                                    Some(
-                                        ctx.builder
-                                            .build_int_truncate(len, llvm_i32, "len")
-                                            .map(Into::into)
-                                            .unwrap(),
-                                    )
-                                }
-                            }
-                            _ => unreachable!(),
-                        }
-                    })
+                    builtin_fns::call_len(generator, ctx, (arg_ty, arg)).map(|ret| Some(ret.into()))
                 },
             )))),
             loc: None,

nac3core/src/typecheck/type_inferencer/mod.rs

@@ -1070,6 +1070,58 @@ impl<'a> Inferencer<'a> {
             }));
         }
 
+        if id == &"len".into() && args.len() == 1 {
+            let obj = self.fold_expr(args.remove(0))?;
+            let obj_ty = obj.custom.unwrap();
+
+            match &*self.unifier.get_ty(obj_ty) {
+                TypeEnum::TObj { obj_id, .. }
+                    if *obj_id == self.primitives.range.obj_id(self.unifier).unwrap() => {}
+                TypeEnum::TObj { obj_id, .. }
+                    if *obj_id == self.primitives.list.obj_id(self.unifier).unwrap() => {}
+                TypeEnum::TObj { obj_id, .. }
+                    if *obj_id == self.primitives.ndarray.obj_id(self.unifier).unwrap() => {}
+                TypeEnum::TTuple { .. } => {}
+                _ => {
+                    return report_error(
+                        format!(
+                            "len() only accepts range, list, ndarray, or tuple. Got {}",
+                            self.unifier.stringify(obj_ty)
+                        )
+                        .as_str(),
+                        obj.location,
+                    )
+                }
+            }
+
+            let ret_ty = self.primitives.int32;
+
+            let func_ty = self.unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                args: vec![FuncArg {
+                    name: "obj".into(),
+                    ty: obj_ty,
+                    default_value: None,
+                    is_vararg: false,
+                }],
+                ret: ret_ty,
+                vars: VarMap::new(),
+            }));
+
+            return Ok(Some(Located {
+                location,
+                custom: Some(ret_ty),
+                node: ExprKind::Call {
+                    func: Box::new(Located {
+                        custom: Some(func_ty),
+                        location: func.location,
+                        node: ExprKind::Name { id: *id, ctx: *ctx },
+                    }),
+                    args: vec![obj],
+                    keywords: vec![],
+                },
+            }));
+        }
+
         if ["int32", "float", "bool", "round", "round64", "np_isnan", "np_isinf"]
             .iter()
             .any(|fun_id| id == &(*fun_id).into())

nac3standalone/demo/src/tuple.py

@@ -25,5 +25,12 @@ def run() -> int32:
         output_int32(tl[0][0])
         output_int32(tl[0][1])
         output_int32(tl[1])
-    
+
+    output_int32(len(()))
+    output_int32(len((1,)))
+    output_int32(len((1, 2)))
+    output_int32(len((1, 2, 3)))
+    output_int32(len((1, 2, 3, 4)))
+    output_int32(len((1, 2, 3, 4, 5)))
+
     return 0