Adding type parameters to functions without return values

src/libasr/asr_utils.h

@@ -1387,6 +1387,19 @@ static inline ASR::ttype_t* get_contained_type(ASR::ttype_t* asr_type) {
     }
 }
 
+static inline ASR::ttype_t* get_type_parameter(ASR::ttype_t* t) {
+    switch (t->type) {
+        case ASR::ttypeType::TypeParameter: {
+            return t;
+        }
+        case ASR::ttypeType::List: {
+            ASR::List_t *tl = ASR::down_cast<ASR::List_t>(t);
+            return get_type_parameter(tl->m_type);
+        }
+        default: throw LCompilersException("Cannot get type parameter from this type.");
+    }
+}
+
 class ReplaceArgVisitor: public ASR::BaseExprReplacer<ReplaceArgVisitor> {
 
     private:

src/lpython/semantics/python_ast_to_asr.cpp

@@ -2271,6 +2271,8 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
         bool current_procedure_interface = false;
         bool overload = false;
         std::set<std::string> ps;
+        Vec<ASR::ttype_t*> tps;
+        tps.reserve(al, x.m_args.n_args);
         bool vectorize = false;
         if (x.n_decorator_list > 0) {
             for(size_t i=0; i<x.n_decorator_list; i++) {
@@ -2309,6 +2311,8 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
             if (ASRUtils::is_generic(*arg_type)) {
                 std::string param_name = ASRUtils::get_parameter_name(arg_type);
                 ps.insert(param_name);
+                ASR::ttype_t *tp = ASRUtils::duplicate_type(al, ASRUtils::get_type_parameter(arg_type));
+                tps.push_back(al, tp);
             }
 
             std::string arg_s = arg;
@@ -2374,43 +2378,19 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
                         ASR::down_cast<ASR::symbol_t>(return_var));
                 ASR::asr_t *return_var_ref = ASR::make_Var_t(al, x.base.base.loc,
                     current_scope->get_symbol(return_var_name));
-                if (ps.size() > 0) {
-                    Vec<ASR::ttype_t*> type_params;
-                    type_params.reserve(al, ps.size());
-                    for (auto &p: ps) {
-                        std::string param = p;
-                        ASR::ttype_t *type_p = ASRUtils::TYPE(ASR::make_TypeParameter_t(al,
-                                x.base.base.loc, s2c(al, p), nullptr, 0));
-                        type_params.push_back(al, type_p);
-                    }
-                    tmp = ASR::make_Function_t(
-                        al, x.base.base.loc,
-                        /* a_symtab */ current_scope,
-                        /* a_name */ s2c(al, sym_name),
-                        /* a_args */ args.p,
-                        /* n_args */ args.size(),
-                        /* a_type_params */ type_params.p,
-                        /* n_type_params */ type_params.size(),
-                        /* a_body */ nullptr,
-                        /* n_body */ 0,
-                        /* a_return_var */ ASRUtils::EXPR(return_var_ref),
-                        current_procedure_abi_type,
-                        s_access, deftype, bindc_name, vectorize, false, false);
-                } else {
-                    tmp = ASR::make_Function_t(
-                        al, x.base.base.loc,
-                        /* a_symtab */ current_scope,
-                        /* a_name */ s2c(al, sym_name),
-                        /* a_args */ args.p,
-                        /* n_args */ args.size(),
-                        /* a_type_params */ nullptr,
-                        /* n_type_params */ 0,
-                        /* a_body */ nullptr,
-                        /* n_body */ 0,
-                        /* a_return_var */ ASRUtils::EXPR(return_var_ref),
-                        current_procedure_abi_type,
-                        s_access, deftype, bindc_name, vectorize, false, false);
-                }
+                tmp = ASR::make_Function_t(
+                    al, x.base.base.loc,
+                    /* a_symtab */ current_scope,
+                    /* a_name */ s2c(al, sym_name),
+                    /* a_args */ args.p,
+                    /* n_args */ args.size(),
+                    /* a_type_params */ tps.p,
+                    /* n_type_params */ tps.size(),
+                    /* a_body */ nullptr,
+                    /* n_body */ 0,
+                    /* a_return_var */ ASRUtils::EXPR(return_var_ref),
+                    current_procedure_abi_type,
+                    s_access, deftype, bindc_name, vectorize, false, false);
             } else {
                 throw SemanticError("Return variable must be an identifier (Name AST node) or an array (Subscript AST node)",
                     x.m_returns->base.loc);
@@ -2423,7 +2403,8 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
                 /* a_name */ s2c(al, sym_name),
                 /* a_args */ args.p,
                 /* n_args */ args.size(),
-                nullptr, 0,
+                /* a_type_params */ tps.p, 
+                /* n_type_params */ tps.size(),
                 /* a_body */ nullptr,
                 /* n_body */ 0,
                 nullptr,

tests/reference/asr-generics_01-d616074.json

@@ -6,7 +6,7 @@
     "outfile": null,
     "outfile_hash": null,
     "stdout": "asr-generics_01-d616074.stdout",
-    "stdout_hash": "9f748f982d651e6706e00bbf46909c4661cf2be40e7ebd25b30ca7fe",
+    "stdout_hash": "996d4e1de29da6f9e90fadfb1687c6d81f5e9eaee27420db1a2692aa",
     "stderr": null,
     "stderr_hash": null,
     "returncode": 0

tests/reference/asr-generics_01-d616074.stdout

@@ -1 +1 @@
-(TranslationUnit (SymbolTable 1 {T: (Variable 1 T Local () () Default (TypeParameter T []) Source Public Required .false.), __lpython_generic_f_0: (Function (SymbolTable 3 {_lpython_return_variable: (Variable 3 _lpython_return_variable ReturnVar () () Default (Integer 4 []) Source Public Required .false.), x: (Variable 3 x In () () Default (Integer 4 []) Source Public Required .false.), y: (Variable 3 y In () () Default (Integer 4 []) Source Public Required .false.)}) __lpython_generic_f_0 [(Var 3 x) (Var 3 y)] [] [(= (Var 3 _lpython_return_variable) (IntegerBinOp (Var 3 x) Add (Var 3 y) (Integer 4 []) ()) ()) (Return)] (Var 3 _lpython_return_variable) Source Public Implementation () .false. .false. .false.), __lpython_generic_f_1: (Function (SymbolTable 4 {_lpython_return_variable: (Variable 4 _lpython_return_variable ReturnVar () () Default (Character 1 1 () []) Source Public Required .false.), x: (Variable 4 x In () () Default (Character 1 1 () []) Source Public Required .false.), y: (Variable 4 y In () () Default (Character 1 1 () []) Source Public Required .false.)}) __lpython_generic_f_1 [(Var 4 x) (Var 4 y)] [] [(= (Var 4 _lpython_return_variable) (StringConcat (Var 4 x) (Var 4 y) (Character 1 2 () []) ()) ()) (Return)] (Var 4 _lpython_return_variable) Source Public Implementation () .false. .false. .false.), _lpython_main_program: (Function (SymbolTable 6 {}) _lpython_main_program [] [] [(Print () [(FunctionCall 1 __lpython_generic_f_0 () [((IntegerConstant 1 (Integer 4 []))) ((IntegerConstant 2 (Integer 4 [])))] (Integer 4 []) () ())] () ()) (Print () [(FunctionCall 1 __lpython_generic_f_1 () [((StringConstant "a" (Character 1 1 () []))) ((StringConstant "b" (Character 1 1 () [])))] (Character 1 1 () []) () ())] () ()) (Print () [(FunctionCall 1 __lpython_generic_f_1 () [((StringConstant "c" (Character 1 1 () []))) ((StringConstant "d" (Character 1 1 () [])))] (Character 1 1 () []) () ())] () ())] () Source Public Implementation () .false. .false. .false.), f: (Function (SymbolTable 2 {_lpython_return_variable: (Variable 2 _lpython_return_variable ReturnVar () () Default (TypeParameter T []) Source Public Required .false.), x: (Variable 2 x In () () Default (TypeParameter T []) Source Public Required .false.), y: (Variable 2 y In () () Default (TypeParameter T []) Source Public Required .false.)}) f [(Var 2 x) (Var 2 y)] [(TypeParameter T [])] [(= (Var 2 _lpython_return_variable) (TemplateBinOp (Var 2 x) Add (Var 2 y) (TypeParameter T []) ()) ()) (Return)] (Var 2 _lpython_return_variable) Source Public Implementation () .false. .false. .false.), main_program: (Program (SymbolTable 5 {}) main_program [] [(SubroutineCall 1 _lpython_main_program () [] ())])}) [])
+(TranslationUnit (SymbolTable 1 {T: (Variable 1 T Local () () Default (TypeParameter T []) Source Public Required .false.), __lpython_generic_f_0: (Function (SymbolTable 3 {_lpython_return_variable: (Variable 3 _lpython_return_variable ReturnVar () () Default (Integer 4 []) Source Public Required .false.), x: (Variable 3 x In () () Default (Integer 4 []) Source Public Required .false.), y: (Variable 3 y In () () Default (Integer 4 []) Source Public Required .false.)}) __lpython_generic_f_0 [(Var 3 x) (Var 3 y)] [] [(= (Var 3 _lpython_return_variable) (IntegerBinOp (Var 3 x) Add (Var 3 y) (Integer 4 []) ()) ()) (Return)] (Var 3 _lpython_return_variable) Source Public Implementation () .false. .false. .false.), __lpython_generic_f_1: (Function (SymbolTable 4 {_lpython_return_variable: (Variable 4 _lpython_return_variable ReturnVar () () Default (Character 1 1 () []) Source Public Required .false.), x: (Variable 4 x In () () Default (Character 1 1 () []) Source Public Required .false.), y: (Variable 4 y In () () Default (Character 1 1 () []) Source Public Required .false.)}) __lpython_generic_f_1 [(Var 4 x) (Var 4 y)] [] [(= (Var 4 _lpython_return_variable) (StringConcat (Var 4 x) (Var 4 y) (Character 1 2 () []) ()) ()) (Return)] (Var 4 _lpython_return_variable) Source Public Implementation () .false. .false. .false.), _lpython_main_program: (Function (SymbolTable 6 {}) _lpython_main_program [] [] [(Print () [(FunctionCall 1 __lpython_generic_f_0 () [((IntegerConstant 1 (Integer 4 []))) ((IntegerConstant 2 (Integer 4 [])))] (Integer 4 []) () ())] () ()) (Print () [(FunctionCall 1 __lpython_generic_f_1 () [((StringConstant "a" (Character 1 1 () []))) ((StringConstant "b" (Character 1 1 () [])))] (Character 1 1 () []) () ())] () ()) (Print () [(FunctionCall 1 __lpython_generic_f_1 () [((StringConstant "c" (Character 1 1 () []))) ((StringConstant "d" (Character 1 1 () [])))] (Character 1 1 () []) () ())] () ())] () Source Public Implementation () .false. .false. .false.), f: (Function (SymbolTable 2 {_lpython_return_variable: (Variable 2 _lpython_return_variable ReturnVar () () Default (TypeParameter T []) Source Public Required .false.), x: (Variable 2 x In () () Default (TypeParameter T []) Source Public Required .false.), y: (Variable 2 y In () () Default (TypeParameter T []) Source Public Required .false.)}) f [(Var 2 x) (Var 2 y)] [(TypeParameter T []) (TypeParameter T [])] [(= (Var 2 _lpython_return_variable) (TemplateBinOp (Var 2 x) Add (Var 2 y) (TypeParameter T []) ()) ()) (Return)] (Var 2 _lpython_return_variable) Source Public Implementation () .false. .false. .false.), main_program: (Program (SymbolTable 5 {}) main_program [] [(SubroutineCall 1 _lpython_main_program () [] ())])}) [])

tests/reference/asr-generics_array_01-682b1b2.json

@@ -6,7 +6,7 @@
     "outfile": null,
     "outfile_hash": null,
     "stdout": "asr-generics_array_01-682b1b2.stdout",
-    "stdout_hash": "e953b223a6c2898132158d1f3bb381f02d5f1c03bc50d444d266cda2",
+    "stdout_hash": "417bc995c9624154c582f31d2ac95e2add9584b1f38373c4c90c9093",
     "stderr": null,
     "stderr_hash": null,
     "returncode": 0

tests/reference/asr-generics_array_01-682b1b2.stdout

@@ -1 +1 @@
-(TranslationUnit (SymbolTable 1 {T: (Variable 1 T Local () () Default (TypeParameter T []) Source Public Required .false.), _lpython_main_program: (Function (SymbolTable 6 {}) _lpython_main_program [] [] [(SubroutineCall 1 use_array () [] ())] () Source Public Implementation () .false. .false. .false.), f: (Function (SymbolTable 2 {_lpython_return_variable: (Variable 2 _lpython_return_variable ReturnVar () () Default (TypeParameter T []) Source Public Required .false.), i: (Variable 2 i In () () Default (TypeParameter T []) Source Public Required .false.), lst: (Variable 2 lst InOut () () Default (TypeParameter T [(() ())]) Source Public Required .false.)}) f [(Var 2 lst) (Var 2 i)] [(TypeParameter T [])] [(= (ArrayItem (Var 2 lst) [(() (IntegerConstant 0 (Integer 4 [])) ())] (TypeParameter T []) ()) (Var 2 i) ()) (= (Var 2 _lpython_return_variable) (ArrayItem (Var 2 lst) [(() (IntegerConstant 0 (Integer 4 [])) ())] (TypeParameter T []) ()) ()) (Return)] (Var 2 _lpython_return_variable) Source Public Implementation () .false. .false. .false.), main_program: (Program (SymbolTable 5 {}) main_program [] [(SubroutineCall 1 _lpython_main_program () [] ())]), use_array: (Function (SymbolTable 3 {__lpython_generic_f_0: (Function (SymbolTable 4 {_lpython_return_variable: (Variable 4 _lpython_return_variable ReturnVar () () Default (Integer 4 []) Source Public Required .false.), i: (Variable 4 i In () () Default (Integer 4 []) Source Public Required .false.), lst: (Variable 4 lst InOut () () Default (Integer 4 [(() ())]) Source Public Required .false.)}) __lpython_generic_f_0 [(Var 4 lst) (Var 4 i)] [] [(= (ArrayItem (Var 4 lst) [(() (IntegerConstant 0 (Integer 4 [])) ())] (Integer 4 []) ()) (Var 4 i) ()) (= (Var 4 _lpython_return_variable) (ArrayItem (Var 4 lst) [(() (IntegerConstant 0 (Integer 4 [])) ())] (Integer 4 []) ()) ()) (Return)] (Var 4 _lpython_return_variable) Source Public Implementation () .false. .false. .false.), array: (Variable 3 array Local () () Default (Integer 4 [((IntegerConstant 0 (Integer 4 [])) (IntegerConstant 1 (Integer 4 [])))]) Source Public Required .false.), x: (Variable 3 x Local () () Default (Integer 4 []) Source Public Required .false.)}) use_array [] [] [(= (Var 3 x) (IntegerConstant 69 (Integer 4 [])) ()) (Print () [(FunctionCall 3 __lpython_generic_f_0 () [((Var 3 array)) ((Var 3 x))] (Integer 4 []) () ())] () ())] () Source Public Implementation () .false. .false. .false.)}) [])
+(TranslationUnit (SymbolTable 1 {T: (Variable 1 T Local () () Default (TypeParameter T []) Source Public Required .false.), _lpython_main_program: (Function (SymbolTable 6 {}) _lpython_main_program [] [] [(SubroutineCall 1 use_array () [] ())] () Source Public Implementation () .false. .false. .false.), f: (Function (SymbolTable 2 {_lpython_return_variable: (Variable 2 _lpython_return_variable ReturnVar () () Default (TypeParameter T []) Source Public Required .false.), i: (Variable 2 i In () () Default (TypeParameter T []) Source Public Required .false.), lst: (Variable 2 lst InOut () () Default (TypeParameter T [(() ())]) Source Public Required .false.)}) f [(Var 2 lst) (Var 2 i)] [(TypeParameter T [(() ())]) (TypeParameter T [])] [(= (ArrayItem (Var 2 lst) [(() (IntegerConstant 0 (Integer 4 [])) ())] (TypeParameter T []) ()) (Var 2 i) ()) (= (Var 2 _lpython_return_variable) (ArrayItem (Var 2 lst) [(() (IntegerConstant 0 (Integer 4 [])) ())] (TypeParameter T []) ()) ()) (Return)] (Var 2 _lpython_return_variable) Source Public Implementation () .false. .false. .false.), main_program: (Program (SymbolTable 5 {}) main_program [] [(SubroutineCall 1 _lpython_main_program () [] ())]), use_array: (Function (SymbolTable 3 {__lpython_generic_f_0: (Function (SymbolTable 4 {_lpython_return_variable: (Variable 4 _lpython_return_variable ReturnVar () () Default (Integer 4 []) Source Public Required .false.), i: (Variable 4 i In () () Default (Integer 4 []) Source Public Required .false.), lst: (Variable 4 lst InOut () () Default (Integer 4 [(() ())]) Source Public Required .false.)}) __lpython_generic_f_0 [(Var 4 lst) (Var 4 i)] [] [(= (ArrayItem (Var 4 lst) [(() (IntegerConstant 0 (Integer 4 [])) ())] (Integer 4 []) ()) (Var 4 i) ()) (= (Var 4 _lpython_return_variable) (ArrayItem (Var 4 lst) [(() (IntegerConstant 0 (Integer 4 [])) ())] (Integer 4 []) ()) ()) (Return)] (Var 4 _lpython_return_variable) Source Public Implementation () .false. .false. .false.), array: (Variable 3 array Local () () Default (Integer 4 [((IntegerConstant 0 (Integer 4 [])) (IntegerConstant 1 (Integer 4 [])))]) Source Public Required .false.), x: (Variable 3 x Local () () Default (Integer 4 []) Source Public Required .false.)}) use_array [] [] [(= (Var 3 x) (IntegerConstant 69 (Integer 4 [])) ()) (Print () [(FunctionCall 3 __lpython_generic_f_0 () [((Var 3 array)) ((Var 3 x))] (Integer 4 []) () ())] () ())] () Source Public Implementation () .false. .false. .false.)}) [])