feat: args expansion

crates/erg_common/opcode309.rs

@@ -113,8 +113,10 @@ impl_u8_enum! {Opcode309;
     BUILD_TUPLE_UNPACK_WITH_CALL = 158,
     LOAD_METHOD = 160,
     CALL_METHOD = 161,
-    CALL_FINALLY = 162,
+    LIST_EXTEND = 162,
-    POP_FINALLY = 163,
+    SET_UPDATE = 163,
+    DICT_MERGE = 164,
+    DICT_UPDATE = 165,
     // Erg-specific opcodes (must have a unary `ERG_`)
     // Define in descending order from 219, 255
     ERG_POP_NTH = 196,

crates/erg_common/opcode310.rs

@@ -111,6 +111,9 @@ impl_u8_enum! {Opcode310;
     LOAD_METHOD = 160,
     CALL_METHOD = 161,
     LIST_EXTEND = 162,
+    SET_UPDATE = 163,
+    DICT_MERGE = 164,
+    DICT_UPDATE = 165,
     // Erg-specific opcodes (must have a unary `ERG_`)
     // Define in descending order from 219, 255
     ERG_POP_NTH = 196,

crates/erg_common/opcode311.rs

@@ -101,6 +101,9 @@ impl_u8_enum! {Opcode311;
     BUILD_STRING = 157,
     LOAD_METHOD = 160,
     LIST_EXTEND = 162,
+    SET_UPDATE = 163,
+    DICT_MERGE = 164,
+    DICT_UPDATE = 165,
     PRECALL = 166,
     CALL = 171,
     KW_NAMES = 172,

crates/erg_compiler/codegen.rs

@@ -609,7 +609,7 @@ impl PyCodeGenerator {
     }
 
     fn stack_dec_n(&mut self, n: usize) {
-        if n > 0 && self.stack_len() == 0 {
+        if n as u32 > self.stack_len() {
             let lasti = self.lasti();
             let last = self.cur_block_codeobj().code.last().unwrap();
             self.crash(&format!(
@@ -2825,6 +2825,19 @@ impl PyCodeGenerator {
             self.write_instr(Opcode310::LIST_TO_TUPLE);
             self.write_arg(0);
         }
+        self.stack_dec();
+    }
+
+    fn emit_kw_var_args_311(&mut self, pos_len: usize, kw_var: &PosArg) {
+        self.write_instr(BUILD_TUPLE);
+        self.write_arg(pos_len);
+        self.stack_dec_n(pos_len.saturating_sub(1));
+        self.write_instr(BUILD_MAP);
+        self.write_arg(0);
+        self.emit_expr(kw_var.expr.clone());
+        self.write_instr(Opcode311::DICT_MERGE);
+        self.write_arg(1);
+        self.stack_dec();
     }
 
     fn emit_var_args_308(&mut self, pos_len: usize, var_args: &PosArg) {
@@ -2839,6 +2852,14 @@ impl PyCodeGenerator {
         }
     }
 
+    fn emit_kw_var_args_308(&mut self, pos_len: usize, kw_var: &PosArg) {
+        self.write_instr(BUILD_TUPLE);
+        self.write_arg(pos_len);
+        self.emit_expr(kw_var.expr.clone());
+        self.stack_dec_n(pos_len.saturating_sub(1));
+        self.stack_dec();
+    }
+
     fn emit_args_311(&mut self, mut args: Args, kind: AccessKind) {
         let argc = args.len();
         let pos_len = args.pos_args.len();
@@ -2857,6 +2878,13 @@ impl PyCodeGenerator {
             kws.push(ValueObj::Str(arg.keyword.content));
             self.emit_expr(arg.expr);
         }
+        if let Some(kw_var) = &args.kw_var {
+            if self.py_version.minor >= Some(10) {
+                self.emit_kw_var_args_311(pos_len, kw_var);
+            } else {
+                self.emit_kw_var_args_308(pos_len, kw_var);
+            }
+        }
         let kwsc = if !kws.is_empty() {
             self.emit_call_kw_instr(argc, kws);
             #[allow(clippy::bool_to_int_with_if)]
@@ -2866,9 +2894,9 @@ impl PyCodeGenerator {
                 1
             }
         } else {
-            if args.var_args.is_some() {
+            if args.var_args.is_some() || args.kw_var.is_some() {
                 self.write_instr(CALL_FUNCTION_EX);
-                if kws.is_empty() {
+                if kws.is_empty() && args.kw_var.is_none() {
                     self.write_arg(0);
                 } else {
                     self.write_arg(1);

crates/erg_compiler/context/inquire.rs

@@ -1081,12 +1081,14 @@ impl Context {
     }
 
     // returns callee's type, not the return type
+    #[allow(clippy::too_many_arguments)]
     fn search_callee_info(
         &self,
         obj: &hir::Expr,
         attr_name: &Option<Identifier>,
         pos_args: &[hir::PosArg],
         kw_args: &[hir::KwArg],
+        (var_args, kw_var_args): (Option<&hir::PosArg>, Option<&hir::PosArg>),
         input: &Input,
         namespace: &Context,
     ) -> SingleTyCheckResult<VarInfo> {
@@ -1107,10 +1109,24 @@ impl Context {
         if let Some(attr_name) = attr_name.as_ref() {
             let mut vi =
                 self.search_method_info(obj, attr_name, pos_args, kw_args, input, namespace)?;
-            vi.t = self.resolve_overload(obj, vi.t, pos_args, kw_args, attr_name)?;
+            vi.t = self.resolve_overload(
+                obj,
+                vi.t,
+                pos_args,
+                kw_args,
+                (var_args, kw_var_args),
+                attr_name,
+            )?;
             Ok(vi)
         } else {
-            let t = self.resolve_overload(obj, obj.t(), pos_args, kw_args, obj)?;
+            let t = self.resolve_overload(
+                obj,
+                obj.t(),
+                pos_args,
+                kw_args,
+                (var_args, kw_var_args),
+                obj,
+            )?;
             Ok(VarInfo {
                 t,
                 ..VarInfo::default()
@@ -1155,6 +1171,7 @@ impl Context {
         instance: Type,
         pos_args: &[hir::PosArg],
         kw_args: &[hir::KwArg],
+        (var_args, kw_var_args): (Option<&hir::PosArg>, Option<&hir::PosArg>),
         loc: &impl Locational,
     ) -> SingleTyCheckResult<Type> {
         let intersecs = instance.intersection_types();
@@ -1195,7 +1212,15 @@ impl Context {
                 if self.subtype_of(ty, &input_t) {
                     if let Ok(instance) = self.instantiate(ty.clone(), obj) {
                         let subst = self
-                            .substitute_call(obj, &None, &instance, pos_args, kw_args, self)
+                            .substitute_call(
+                                obj,
+                                &None,
+                                &instance,
+                                pos_args,
+                                kw_args,
+                                (var_args, kw_var_args),
+                                self,
+                            )
                             .is_ok();
                         let eval = self.eval_t_params(instance, self.level, obj).is_ok();
                         if subst && eval {
@@ -1641,7 +1666,7 @@ impl Context {
                 )
             })?;
         let op = hir::Expr::Accessor(hir::Accessor::private(symbol, t));
-        self.get_call_t(&op, &None, args, &[], input, namespace)
+        self.get_call_t(&op, &None, args, &[], (None, None), input, namespace)
             .map_err(|(_, errs)| {
                 let hir::Expr::Accessor(hir::Accessor::Ident(op_ident)) = op else {
                     return errs;
@@ -1695,7 +1720,7 @@ impl Context {
                 )
             })?;
         let op = hir::Expr::Accessor(hir::Accessor::private(symbol, vi));
-        self.get_call_t(&op, &None, args, &[], input, namespace)
+        self.get_call_t(&op, &None, args, &[], (None, None), input, namespace)
             .map_err(|(_, errs)| {
                 let hir::Expr::Accessor(hir::Accessor::Ident(op_ident)) = op else {
                     return errs;
@@ -1787,6 +1812,7 @@ impl Context {
     /// ↓ don't substitute `Int` to `self`
     /// substitute_call(obj: Int, instance: ((self: Int, other: Int) -> Int), [1, 2]) => instance: (Int, Int) -> Int
     /// ```
+    #[allow(clippy::too_many_arguments)]
     fn substitute_call(
         &self,
         obj: &hir::Expr,
@@ -1794,6 +1820,7 @@ impl Context {
         instance: &Type,
         pos_args: &[hir::PosArg],
         kw_args: &[hir::KwArg],
+        (var_args, kw_var_args): (Option<&hir::PosArg>, Option<&hir::PosArg>),
         namespace: &Context,
     ) -> TyCheckResult<SubstituteResult> {
         match instance {
@@ -1803,6 +1830,7 @@ impl Context {
                 fv.unsafe_crack(),
                 pos_args,
                 kw_args,
+                (var_args, kw_var_args),
                 namespace,
             ),
             Type::FreeVar(fv) => {
@@ -1816,12 +1844,24 @@ impl Context {
                         if instance.is_quantified_subr() {
                             let instance = self.instantiate(instance.clone(), obj)?;
                             self.substitute_call(
-                                obj, attr_name, &instance, pos_args, kw_args, namespace,
+                                obj,
+                                attr_name,
+                                &instance,
+                                pos_args,
+                                kw_args,
+                                (var_args, kw_var_args),
+                                namespace,
                             )?;
                             return Ok(SubstituteResult::Coerced(instance));
                         } else if get_hash(instance) != hash {
                             return self.substitute_call(
-                                obj, attr_name, instance, pos_args, kw_args, namespace,
+                                obj,
+                                attr_name,
+                                instance,
+                                pos_args,
+                                kw_args,
+                                (var_args, kw_var_args),
+                                namespace,
                             );
                         }
                     }
@@ -1852,17 +1892,36 @@ impl Context {
                     Ok(SubstituteResult::Ok)
                 }
             }
-            Type::Refinement(refine) => {
+            Type::Refinement(refine) => self.substitute_call(
-                self.substitute_call(obj, attr_name, &refine.t, pos_args, kw_args, namespace)
+                obj,
-            }
+                attr_name,
+                &refine.t,
+                pos_args,
+                kw_args,
+                (var_args, kw_var_args),
+                namespace,
+            ),
             // instance must be instantiated
             Type::Quantified(_) => unreachable_error!(TyCheckErrors, TyCheckError, self),
             Type::Subr(subr) => {
-                let res = self.substitute_subr_call(obj, attr_name, subr, pos_args, kw_args);
+                let res = self.substitute_subr_call(
+                    obj,
+                    attr_name,
+                    subr,
+                    pos_args,
+                    kw_args,
+                    (var_args, kw_var_args),
+                );
                 // TODO: change polymorphic type syntax
                 if res.is_err() && subr.return_t.is_class_type() {
                     self.substitute_dunder_call(
-                        obj, attr_name, instance, pos_args, kw_args, namespace,
+                        obj,
+                        attr_name,
+                        instance,
+                        pos_args,
+                        kw_args,
+                        (var_args, kw_var_args),
+                        namespace,
                     )
                     .or(res)
                 } else {
@@ -1870,14 +1929,34 @@ impl Context {
                 }
             }
             Type::And(_, _) => {
-                let instance =
+                let instance = self.resolve_overload(
-                    self.resolve_overload(obj, instance.clone(), pos_args, kw_args, &obj)?;
+                    obj,
-                self.substitute_call(obj, attr_name, &instance, pos_args, kw_args, namespace)
+                    instance.clone(),
+                    pos_args,
+                    kw_args,
+                    (var_args, kw_var_args),
+                    &obj,
+                )?;
+                self.substitute_call(
+                    obj,
+                    attr_name,
+                    &instance,
+                    pos_args,
+                    kw_args,
+                    (var_args, kw_var_args),
+                    namespace,
+                )
             }
             Type::Failure => Ok(SubstituteResult::Ok),
-            _ => {
+            _ => self.substitute_dunder_call(
-                self.substitute_dunder_call(obj, attr_name, instance, pos_args, kw_args, namespace)
+                obj,
-            }
+                attr_name,
+                instance,
+                pos_args,
+                kw_args,
+                (var_args, kw_var_args),
+                namespace,
+            ),
         }
     }
 
@@ -1888,6 +1967,7 @@ impl Context {
         subr: &SubrType,
         pos_args: &[hir::PosArg],
         kw_args: &[hir::KwArg],
+        (var_args, kw_var_args): (Option<&hir::PosArg>, Option<&hir::PosArg>),
     ) -> TyCheckResult<SubstituteResult> {
         let mut errs = TyCheckErrors::empty();
         // method: obj: 1, subr: (self: Int, other: Int) -> Int
@@ -1911,7 +1991,10 @@ impl Context {
         } else {
             subr.non_default_params.len() + subr.default_params.len()
         };
-        if (params_len < pos_args.len() || params_len < pos_args.len() + kw_args.len())
+        let there_var = var_args.is_some() || kw_var_args.is_some();
+        if (params_len < pos_args.len()
+            || params_len < pos_args.len() + kw_args.len()
+            || (params_len == pos_args.len() + kw_args.len() && there_var))
             && subr.is_no_var()
         {
             return Err(self.gen_too_many_args_error(&callee, subr, is_method, pos_args, kw_args));
@@ -2047,7 +2130,59 @@ impl Context {
                         .into()
                 })
                 .collect::<Vec<_>>();
-            if !missing_params.is_empty() {
+            if let Some(var_args) = var_args {
+                if !self.subtype_of(
+                    var_args.expr.ref_t(),
+                    &poly("Iterable", vec![TyParam::t(Obj)]),
+                ) {
+                    let err = TyCheckError::type_mismatch_error(
+                        self.cfg.input.clone(),
+                        line!() as usize,
+                        var_args.loc(),
+                        self.caused_by(),
+                        "_",
+                        None,
+                        &poly("Iterable", vec![TyParam::t(Obj)]),
+                        var_args.expr.ref_t(),
+                        None,
+                        None,
+                    );
+                    errs.push(err);
+                }
+            }
+            if let Some(kw_var_args) = kw_var_args {
+                if !self.subtype_of(
+                    kw_var_args.expr.ref_t(),
+                    &poly("Mapping", vec![TyParam::t(Obj), TyParam::t(Obj)]),
+                ) {
+                    let err = TyCheckError::type_mismatch_error(
+                        self.cfg.input.clone(),
+                        line!() as usize,
+                        kw_var_args.loc(),
+                        self.caused_by(),
+                        "_",
+                        None,
+                        &poly("Mapping", vec![TyParam::t(Obj), TyParam::t(Obj)]),
+                        kw_var_args.expr.ref_t(),
+                        None,
+                        None,
+                    );
+                    errs.push(err);
+                }
+            }
+            if missing_params.is_empty() && (var_args.is_some() || kw_var_args.is_some()) {
+                return Err(TyCheckErrors::from(TyCheckError::too_many_args_error(
+                    self.cfg.input.clone(),
+                    line!() as usize,
+                    callee.loc(),
+                    &callee.to_string(),
+                    self.caused_by(),
+                    params_len,
+                    pos_args.len(),
+                    kw_args.len(),
+                )));
+            }
+            if !missing_params.is_empty() && var_args.is_none() && kw_var_args.is_none() {
                 return Err(TyCheckErrors::from(TyCheckError::args_missing_error(
                     self.cfg.input.clone(),
                     line!() as usize,
@@ -2068,6 +2203,7 @@ impl Context {
         }
     }
 
+    #[allow(clippy::too_many_arguments)]
     fn substitute_dunder_call(
         &self,
         obj: &hir::Expr,
@@ -2075,6 +2211,7 @@ impl Context {
         instance: &Type,
         pos_args: &[hir::PosArg],
         kw_args: &[hir::KwArg],
+        (var_args, kw_var_args): (Option<&hir::PosArg>, Option<&hir::PosArg>),
         namespace: &Context,
     ) -> TyCheckResult<SubstituteResult> {
         let ctxs = self
@@ -2121,9 +2258,15 @@ impl Context {
                         )?;
                     }
                     let instance = self.instantiate_def_type(&call_vi.t)?;
-                    let instance = match self
+                    let instance = match self.substitute_call(
-                        .substitute_call(obj, attr_name, &instance, pos_args, kw_args, namespace)?
+                        obj,
-                    {
+                        attr_name,
+                        &instance,
+                        pos_args,
+                        kw_args,
+                        (var_args, kw_var_args),
+                        namespace,
+                    )? {
                         SubstituteResult::__Call__(instance)
                         | SubstituteResult::Coerced(instance) => instance,
                         SubstituteResult::Ok => instance,
@@ -2138,9 +2281,15 @@ impl Context {
                 })
             {
                 let instance = self.instantiate_def_type(&call_vi.t)?;
-                let instance = match self
+                let instance = match self.substitute_call(
-                    .substitute_call(obj, attr_name, &instance, pos_args, kw_args, namespace)?
+                    obj,
-                {
+                    attr_name,
+                    &instance,
+                    pos_args,
+                    kw_args,
+                    (var_args, kw_var_args),
+                    namespace,
+                )? {
                     SubstituteResult::__Call__(instance) | SubstituteResult::Coerced(instance) => {
                         instance
                     }
@@ -2469,12 +2618,14 @@ impl Context {
         Ok(res)
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub(crate) fn get_call_t(
         &self,
         obj: &hir::Expr,
         attr_name: &Option<Identifier>,
         pos_args: &[hir::PosArg],
         kw_args: &[hir::KwArg],
+        (var_args, kw_var_args): (Option<&hir::PosArg>, Option<&hir::PosArg>),
         input: &Input,
         namespace: &Context,
     ) -> FailableOption<VarInfo> {
@@ -2492,7 +2643,15 @@ impl Context {
             }
         }
         let found = self
-            .search_callee_info(obj, attr_name, pos_args, kw_args, input, namespace)
+            .search_callee_info(
+                obj,
+                attr_name,
+                pos_args,
+                kw_args,
+                (var_args, kw_var_args),
+                input,
+                namespace,
+            )
             .map_err(|err| (None, TyCheckErrors::from(err)))?;
         log!(
             "Found:\ncallee: {obj}{}\nfound: {found}",
@@ -2507,7 +2666,15 @@ impl Context {
             fmt_slice(kw_args)
         );
         let instance = match self
-            .substitute_call(obj, attr_name, &instance, pos_args, kw_args, namespace)
+            .substitute_call(
+                obj,
+                attr_name,
+                &instance,
+                pos_args,
+                kw_args,
+                (var_args, kw_var_args),
+                namespace,
+            )
             .map_err(|errs| {
                 (
                     Some(VarInfo {

crates/erg_compiler/hir.rs

@@ -290,7 +290,8 @@ impl Args {
     pub fn len(&self) -> usize {
         #[allow(clippy::bool_to_int_with_if)]
         let var_argc = if self.var_args.is_none() { 0 } else { 1 };
-        self.pos_args.len() + var_argc + self.kw_args.len()
+        let kw_var_argc = if self.kw_var.is_none() { 0 } else { 1 };
+        self.pos_args.len() + var_argc + self.kw_args.len() + kw_var_argc
     }
 
     #[inline]

crates/erg_compiler/lower.rs

@@ -1409,6 +1409,7 @@ impl<A: ASTBuildable> GenericASTLowerer<A> {
             &call.attr_name,
             &hir_args.pos_args,
             &hir_args.kw_args,
+            (hir_args.var_args.as_deref(), hir_args.kw_var.as_deref()),
             &self.cfg.input,
             &self.module.context,
         ) {
@@ -1566,6 +1567,7 @@ impl<A: ASTBuildable> GenericASTLowerer<A> {
             &Some(attr_name.clone()),
             &args,
             &[],
+            (None, None),
             &self.cfg.input,
             &self.module.context,
         ) {

tests/should_err/args_expansion.er

@@ -0,0 +1,8 @@
+foo first: Int, second: Int = log first, second
+
+foo(1, 2, *[1, 2]) # ERR
+data = {"first": 1, "second": 2}
+foo(1, 2, **data) # ERR
+
+foo(1, *2) # ERR
+foo(1, **[2]) # ERR

tests/should_ok/args_expansion.er

@@ -0,0 +1,12 @@
+foo first: Int, second: Int, third: Int = log first, second, third
+
+foo(*[1, 2, 3])
+foo(1, *[2, 3])
+foo(1, 2, *[3])
+data = {"first": 1, "second": 2, "third": 3}
+foo(**data)
+foo(1, **{"second": 2, "third": 3})
+# FIXME:
+# foo(1, 2, **{"third": 3})
+
+print! "OK"

tests/test.rs

@@ -16,6 +16,11 @@ fn exec_advanced_type_spec() -> Result<(), ()> {
     expect_success("tests/should_ok/advanced_type_spec.er", 5)
 }
 
+#[test]
+fn exec_args_expansion() -> Result<(), ()> {
+    expect_success("tests/should_ok/args_expansion.er", 0)
+}
+
 #[test]
 fn exec_list_test() -> Result<(), ()> {
     expect_success("tests/should_ok/list.er", 0)
@@ -502,6 +507,11 @@ fn exec_args() -> Result<(), ()> {
     expect_failure("tests/should_err/args.er", 0, 19)
 }
 
+#[test]
+fn exec_args_expansion_err() -> Result<(), ()> {
+    expect_failure("tests/should_err/args_expansion.er", 0, 4)
+}
+
 #[test]
 fn exec_list_err() -> Result<(), ()> {
     expect_failure("examples/list.er", 0, 1)