From 26cb4721d91d388056b55de13adccc7432bc0679 Mon Sep 17 00:00:00 2001
From: lyken <lyken@m-labs.hk>
Date: Tue, 13 Aug 2024 10:27:27 +0800
Subject: [PATCH] core/typecheck: Support tuple arg type in len()

---
 nac3core/src/toplevel/builtins.rs             | 25 ++--------
 nac3core/src/typecheck/type_inferencer/mod.rs | 47 +++++++++++++++++++
 2 files changed, 52 insertions(+), 20 deletions(-)

diff --git a/nac3core/src/toplevel/builtins.rs b/nac3core/src/toplevel/builtins.rs
index fc25f20b..1c30769e 100644
--- a/nac3core/src/toplevel/builtins.rs
+++ b/nac3core/src/toplevel/builtins.rs
@@ -1399,31 +1399,16 @@ 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, default_value: None }],
-                ret: int32,
-                vars: into_var_map([tvar, arg_ty]),
+                args: vec![FuncArg { name: "obj".into(), ty: arg_tvar.ty, default_value: None }],
+                ret: self.primitives.int32,
+                vars: into_var_map([arg_tvar]),
             })),
             var_id: Vec::default(),
             instance_to_symbol: HashMap::default(),
diff --git a/nac3core/src/typecheck/type_inferencer/mod.rs b/nac3core/src/typecheck/type_inferencer/mod.rs
index d9380ab1..79b1ae5e 100644
--- a/nac3core/src/typecheck/type_inferencer/mod.rs
+++ b/nac3core/src/typecheck/type_inferencer/mod.rs
@@ -1072,6 +1072,53 @@ 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 }],
+                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())