Merge pull request #311 from czgdp1807/inline_pass_fix

Fixed inline function calls pass to inline simple functions

Author: certik

src/bin/lpython.cpp

@@ -24,6 +24,7 @@
 #include <libasr/pass/arr_slice.h>
 #include <libasr/pass/print_arr.h>
 #include <libasr/pass/unused_functions.h>
+#include <libasr/pass/inline_function_calls.h>
 #include <libasr/asr_utils.h>
 #include <libasr/asr_verify.h>
 #include <libasr/modfile.h>
@@ -49,6 +50,7 @@ enum Backend {
 enum ASRPass {
     do_loops, global_stmts, implied_do_loops, array_op,
     arr_slice, print_arr, class_constructor, unused_functions,
+    inline_function_calls
 };
 
 std::string remove_extension(const std::string& filename) {
@@ -722,6 +724,8 @@ int main(int argc, char *argv[])
                 passes.push_back(ASRPass::implied_do_loops);
             } else if (arg_pass == "array_op") {
                 passes.push_back(ASRPass::array_op);
+            } else if (arg_pass == "inline_function_calls") {
+                passes.push_back(ASRPass::inline_function_calls);
             } else if (arg_pass == "class_constructor") {
                 passes.push_back(ASRPass::class_constructor);
             } else if (arg_pass == "print_arr") {

src/libasr/asr_utils.h

@@ -385,7 +385,10 @@ static inline bool is_intrinsic_function2(const ASR::Function_t *fn) {
     ASR::symbol_t *sym = (ASR::symbol_t*)fn;
     ASR::Module_t *m = get_sym_module0(sym);
     if (m != nullptr) {
-        if (m->m_intrinsic) return true;
+        if (m->m_intrinsic ||
+            fn->m_abi == ASR::abiType::Intrinsic) {
+                return true;
+        }
     }
     return false;
 }

src/libasr/pass/inline_function_calls.cpp

@@ -44,6 +44,7 @@ class InlineFunctionCallVisitor : public PassUtils::PassVisitor<InlineFunctionCa
     ASR::expr_t* function_result_var;
 
     bool from_inline_function_call, inlining_function;
+    bool fixed_duplicated_expr_stmt;
 
     // Stores the local variables corresponding to the ones
     // present in function symbol table.
@@ -56,16 +57,18 @@ class InlineFunctionCallVisitor : public PassUtils::PassVisitor<InlineFunctionCa
 
     ASR::ExprStmtDuplicator node_duplicator;
 
+    SymbolTable* current_routine_scope;
+
 public:
 
     bool function_inlined;
 
     InlineFunctionCallVisitor(Allocator &al_, const std::string& rl_path_, bool inline_external_symbol_calls_)
     : PassVisitor(al_, nullptr),
     rl_path(rl_path_), function_result_var(nullptr),
-    from_inline_function_call(false), inlining_function(false),
+    from_inline_function_call(false), inlining_function(false), fixed_duplicated_expr_stmt(false),
     current_routine(""), inline_external_symbol_calls(inline_external_symbol_calls_),
-    node_duplicator(al_), function_inlined(false)
+    node_duplicator(al_), current_routine_scope(nullptr), function_inlined(false)
     {
         pass_result.reserve(al, 1);
     }
@@ -85,14 +88,28 @@ class InlineFunctionCallVisitor : public PassUtils::PassVisitor<InlineFunctionCa
 
     void visit_Var(const ASR::Var_t& x) {
         ASR::Var_t& xx = const_cast<ASR::Var_t&>(x);
-        ASR::Variable_t* x_var = ASR::down_cast<ASR::Variable_t>(x.m_v);
-        std::string x_var_name = std::string(x_var->m_name);
-        if( arg2value.find(x_var_name) != arg2value.end() ) {
-            x_var = ASR::down_cast<ASR::Variable_t>(arg2value[x_var_name]);
-            if( current_scope->scope.find(std::string(x_var->m_name)) != current_scope->scope.end() ) {
-                xx.m_v = arg2value[x_var_name];
+        std::string x_var_name = std::string(ASRUtils::symbol_name(x.m_v));
+
+        // If anything is not local to a function being inlined
+        // then do not inline the function by setting
+        // fixed_duplicated_expr_stmt to false.
+        // To be supported later.
+        if( current_routine_scope &&
+            current_routine_scope->scope.find(x_var_name) == current_routine_scope->scope.end() ) {
+            fixed_duplicated_expr_stmt = false;
+            return ;
+        }
+        if( x.m_v->type == ASR::symbolType::Variable ) {
+            ASR::Variable_t* x_var = ASR::down_cast<ASR::Variable_t>(x.m_v);
+            if( arg2value.find(x_var_name) != arg2value.end() ) {
+                x_var = ASR::down_cast<ASR::Variable_t>(arg2value[x_var_name]);
+                if( current_scope->scope.find(std::string(x_var->m_name)) != current_scope->scope.end() ) {
+                    xx.m_v = arg2value[x_var_name];
+                }
+                x_var = ASR::down_cast<ASR::Variable_t>(x.m_v);
             }
-            x_var = ASR::down_cast<ASR::Variable_t>(x.m_v);
+        } else {
+            fixed_duplicated_expr_stmt = false;
         }
     }
 
@@ -166,10 +183,13 @@ class InlineFunctionCallVisitor : public PassUtils::PassVisitor<InlineFunctionCa
 
         // Avoid inlining current function call if its a recursion.
         ASR::Function_t* func = ASR::down_cast<ASR::Function_t>(routine);
-        if( std::string(func->m_name) == current_routine ) {
+        if( ASRUtils::is_intrinsic_function2(func) ||
+            std::string(func->m_name) == current_routine ) {
             return ;
         }
 
+        current_routine_scope = func->m_symtab;
+
         ASR::expr_t* return_var = nullptr;
         // The following prepares arg2value map for inlining the
         // current function call. Variables are created in the current
@@ -259,7 +279,12 @@ class InlineFunctionCallVisitor : public PassUtils::PassVisitor<InlineFunctionCa
         // in the current scope. See, `visit_Var` to know how replacement occurs.
         for( size_t i = 0; i < exprs_to_be_visited.size() && success; i++ ) {
             ASR::expr_t* value = exprs_to_be_visited[i].first;
+            fixed_duplicated_expr_stmt = true;
             visit_expr(*value);
+            if( !fixed_duplicated_expr_stmt ) {
+                success = false;
+                break;
+            }
             ASR::symbol_t* variable = exprs_to_be_visited[i].second;
             ASR::expr_t* var = LFortran::ASRUtils::EXPR(ASR::make_Var_t(al, variable->base.loc, variable));
             ASR::stmt_t* assign_stmt = ASRUtils::STMT(ASR::make_Assignment_t(al, var->base.loc, var, value, nullptr));
@@ -280,32 +305,45 @@ class InlineFunctionCallVisitor : public PassUtils::PassVisitor<InlineFunctionCa
         }
 
         if( success ) {
-            // If duplication is successfull then fill the
-            // pass result with assignment statements
-            // (for local variables in the loop just below)
-            // and the function body (the next loop).
-            for( size_t i = 0; i < pass_result_local.size(); i++ ) {
-                pass_result.push_back(al, pass_result_local[i]);
-            }
             // Set inlining_function to true so that we inline
             // only one function at a time.
             inlining_function = true;
-            for( size_t i = 0; i < func->n_body; i++ ) {
+            for( size_t i = 0; i < func->n_body && success; i++ ) {
+                fixed_duplicated_expr_stmt = true;
                 visit_stmt(*func_copy[i]);
-                pass_result.push_back(al, func_copy[i]);
+                success = success && fixed_duplicated_expr_stmt;
+            }
+
+            if( success ) {
+                // If duplication is successfull then fill the
+                // pass result with assignment statements
+                // (for local variables in the loop just below)
+                // and the function body (the next loop).
+                for( size_t i = 0; i < pass_result_local.size(); i++ ) {
+                    pass_result.push_back(al, pass_result_local[i]);
+                }
+
+                for( size_t i = 0; i < func->n_body; i++ ) {
+                    pass_result.push_back(al, func_copy[i]);
+                }
             }
             inlining_function = false;
+            current_routine_scope = nullptr;
             function_result_var = return_var;
-        } else {
+        }
+
+        if (!success) {
             // If not successfull then delete all the local variables
             // created for the purpose of inlining the current function call.
             for( auto& itr : arg2value ) {
                 ASR::Variable_t* auxiliary_var = ASR::down_cast<ASR::Variable_t>(itr.second);
                 current_scope->scope.erase(std::string(auxiliary_var->m_name));
             }
+            function_result_var = nullptr;
         }
         // At least one function is inlined
         function_inlined = success;
+        success = false;
         // Clear up the arg2value to avoid corruption
         // of any kind.
         arg2value.clear();