diff --git a/nac3core/src/codegen/builtin_fns.rs b/nac3core/src/codegen/builtin_fns.rs index 311fd358..0aa89798 100644 --- a/nac3core/src/codegen/builtin_fns.rs +++ b/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, 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, diff --git a/nac3core/src/toplevel/builtins.rs b/nac3core/src/toplevel/builtins.rs index be8687ea..2fef5eb4 100644 --- a/nac3core/src/toplevel/builtins.rs +++ b/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, diff --git a/nac3core/src/typecheck/type_inferencer/mod.rs b/nac3core/src/typecheck/type_inferencer/mod.rs index 9ac503a1..a9c468f9 100644 --- a/nac3core/src/typecheck/type_inferencer/mod.rs +++ b/nac3core/src/typecheck/type_inferencer/mod.rs @@ -1069,6 +1069,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()) diff --git a/nac3standalone/demo/src/tuple.py b/nac3standalone/demo/src/tuple.py index f71d3f99..6947171d 100644 --- a/nac3standalone/demo/src/tuple.py +++ b/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 \ No newline at end of file