XX Wrap symbols from global scope into a function

Author: Thirumalai-Shaktivel

examples/expr2.py

@@ -1,9 +1,15 @@
-def main0():
-    x: i32
-    x = (2+3)*5
-    print(x)
+from ltypes import i32
 
-main0()
+x: list[i32]
+x = [1,2,3]
+print(x, x[0])
+
+# def main0():
+#     x: list[i32] = [1]
+#     print(x, x[0])
+#     print(x[0])
+
+# main0()
 
 # Not implemented yet in LPython:
 #if __name__ == "__main__":

src/libasr/ASR.asdl

@@ -28,8 +28,10 @@
 
 module ASR {
 
+-- `main` in `TranslationUnit` designates a function name similar to the main
+-- function in C (temporary solution)
 unit
-    = TranslationUnit(symbol_table global_scope, node* items)
+    = TranslationUnit(symbol_table global_scope, identifier? main, node* items)
 
 -- # Documentation for the symbol type
 

src/libasr/asr_scopes.cpp

@@ -134,4 +134,27 @@ std::string SymbolTable::get_unique_name(const std::string &name) {
     return unique_name;
 }
 
+void SymbolTable::move_symbols_from_global_scope(SymbolTable *fn_scope,
+        std::vector<std::string> &syms) {
+    for (auto &a : scope) {
+        switch (a.second->type) {
+            case (ASR::symbolType::Variable) : {
+                ASR::Variable_t *v = ASR::down_cast<ASR::Variable_t>(a.second);
+                if (ASR::is_a<ASR::List_t>(*v->m_type)
+                 || ASR::is_a<ASR::Tuple_t>(*v->m_type)
+                 || ASR::is_a<ASR::Dict_t>(*v->m_type)) {
+                    v->m_parent_symtab = fn_scope;
+                    fn_scope->add_symbol(a.first, (ASR::symbol_t *)v);
+                    syms.push_back(a.first);
+                    break;
+                }
+            }
+            default : {
+                LCompilersException("Moving the symbol:`"
+                    "` from global scope is not implemented yet");
+            };
+        }
+    }
+}
+
 } // namespace LCompilers

src/libasr/asr_scopes.h

@@ -80,6 +80,9 @@ struct SymbolTable {
         size_t n_scope_names, char **m_scope_names);
 
     std::string get_unique_name(const std::string &name);
+
+    void move_symbols_from_global_scope(SymbolTable *fn_scope,
+        std::vector<std::string> &syms);
 };
 
 } // namespace LCompilers

src/libasr/asr_utils.cpp

@@ -180,7 +180,7 @@ ASR::Module_t* load_module(Allocator &al, SymbolTable *symtab,
     ASR::asr_t *orig_asr_owner = symtab->asr_owner;
     ASR::TranslationUnit_t *tu
         = ASR::down_cast2<ASR::TranslationUnit_t>(ASR::make_TranslationUnit_t(al, loc,
-            symtab, nullptr, 0));
+            symtab, "_global_statements", nullptr, 0));
 
     // Load any dependent modules recursively
     bool rerun = true;

src/libasr/asr_utils.h

@@ -2090,6 +2090,46 @@ class ReplaceReturnWithGotoVisitor: public ASR::BaseStmtReplacer<ReplaceReturnWi
 
 };
 
+class VarExprReplacer: public ASR::BaseExprReplacer<VarExprReplacer>
+{
+private:
+    Allocator &al;
+    SymbolTable *current_scope;
+
+public:
+    VarExprReplacer(Allocator& al_, SymbolTable *current_scope_) :
+        al(al_), current_scope(current_scope_)
+    {}
+
+    void replace_Var(ASR::Var_t *x) {
+        std::string sym_name = ASRUtils::symbol_name(x->m_v);
+        ASR::symbol_t *s = current_scope->resolve_symbol(sym_name);
+        LCOMPILERS_ASSERT(s != nullptr)
+        *current_expr = ASRUtils::EXPR(ASR::make_Var_t(al,
+            x->base.base.loc, s));
+    }
+};
+
+class VarExprVisitor : public ASR::CallReplacerOnExpressionsVisitor
+     <VarExprVisitor>
+{
+private:
+
+    VarExprReplacer replacer;
+
+public:
+
+    VarExprVisitor(Allocator& al_, SymbolTable *current_scope) :
+        replacer(al_, current_scope)
+        { }
+
+    void call_replacer() {
+        replacer.current_expr = current_expr;
+        replacer.replace_expr(*current_expr);
+    }
+
+};
+
 // Singleton LabelGenerator so that it generates
 // unique labels for different statements, from
 // whereever it is called (be it ASR passes, be it

src/libasr/asr_verify.cpp

@@ -533,14 +533,20 @@ class VerifyVisitor : public BaseWalkVisitor<VerifyVisitor>
                 "ExternalSymbol::m_original_name must match external->m_name");
             ASR::Module_t *m = ASRUtils::get_sym_module(x.m_external);
             ASR::StructType_t* sm = nullptr;
+            ASR::Function_t* fn = nullptr;
             bool is_valid_owner = false;
             is_valid_owner = m != nullptr;
             std::string asr_owner_name = "";
             if( !is_valid_owner ) {
+                is_valid_owner = true;
                 ASR::symbol_t* asr_owner_sym = ASRUtils::get_asr_owner(x.m_external);
-                is_valid_owner = ASR::is_a<ASR::StructType_t>(*asr_owner_sym);
-                sm = ASR::down_cast<ASR::StructType_t>(asr_owner_sym);
-                asr_owner_name = sm->m_name;
+                if ( ASR::is_a<ASR::StructType_t>(*asr_owner_sym) ) {
+                    sm = ASR::down_cast<ASR::StructType_t>(asr_owner_sym);
+                    asr_owner_name = sm->m_name;
+                } else if (ASR::is_a<ASR::Function_t>(*asr_owner_sym)) {
+                    fn = ASR::down_cast<ASR::Function_t>(asr_owner_sym);
+                    asr_owner_name = fn->m_name;
+                }
             } else {
                 asr_owner_name = m->m_name;
             }
@@ -554,6 +560,8 @@ class VerifyVisitor : public BaseWalkVisitor<VerifyVisitor>
                 s = m->m_symtab->find_scoped_symbol(x.m_original_name, x.n_scope_names, x.m_scope_names);
             } else if( sm ) {
                 s = sm->m_symtab->resolve_symbol(std::string(x.m_original_name));
+            } else if( fn ) {
+                s = fn->m_symtab->resolve_symbol(std::string(x.m_original_name));
             }
             require(s != nullptr,
                 "ExternalSymbol::m_original_name ('"

src/libasr/pass/global_stmts.cpp

@@ -162,4 +162,38 @@ void pass_wrap_global_stmts_into_function(Allocator &al,
     }
 }
 
+void pass_wrap_global_vars_into_function(Allocator &al, ASR::TranslationUnit_t &unit) {
+    std::vector<std::string> symbols;
+    ASR::symbol_t *fn_s = unit.m_global_scope->get_symbol(unit.m_main);
+
+    // Move symbols that fails in the global scope
+    unit.m_global_scope->move_symbols_from_global_scope(
+        ASR::down_cast<ASR::Function_t>(fn_s)->m_symtab, symbols);
+
+    // Create ExternalSymbol and erase the actual symbols
+    // that are moved into the function
+    for (auto &sym: symbols) {
+        ASR::symbol_t *var_s = unit.m_global_scope->get_symbol(sym);
+        ASR::Variable_t *v = ASR::down_cast<ASR::Variable_t>(var_s);
+        var_s = (ASR::symbol_t *) ASR::make_ExternalSymbol_t(al,
+            v->base.base.loc, unit.m_global_scope, v->m_name, var_s,
+            ASRUtils::symbol_name(ASRUtils::get_asr_owner(var_s)),
+            nullptr, 0, v->m_name, ASR::accessType::Public);
+        unit.m_global_scope->erase_symbol(sym);
+        unit.m_global_scope->add_symbol(v->m_name, var_s);
+    }
+
+    // Replace the Var_t with correct symbol_table
+    ASRUtils::VarExprVisitor replace_sym(al, unit.m_global_scope);
+    replace_sym.visit_TranslationUnit(unit);
+
+#if defined(WITH_LFORTRAN_ASSERT)
+        diag::Diagnostics diagnostics;
+        if (!asr_verify(unit, true, diagnostics)) {
+            std::cerr << diagnostics.render2();
+            throw LCompilersException("Verify failed");
+        };
+#endif
+}
+
 } // namespace LCompilers

src/libasr/pass/global_stmts.h

@@ -8,6 +8,8 @@ namespace LCompilers {
 
     void pass_wrap_global_stmts_into_function(Allocator &al, ASR::TranslationUnit_t &unit,
                                               const LCompilers::PassOptions& pass_options);
+    void pass_wrap_global_vars_into_function(Allocator &al,
+        ASR::TranslationUnit_t &unit);
 
 } // namespace LCompilers
 

src/libasr/pass/global_stmts_program.cpp

@@ -26,6 +26,7 @@ void pass_wrap_global_stmts_into_program(Allocator &al,
     prog_body.reserve(al, 1);
     if (unit.n_items > 0) {
         pass_wrap_global_stmts_into_function(al, unit, pass_options);
+        pass_wrap_global_vars_into_function(al, unit);
         ASR::symbol_t *fn = unit.m_global_scope->get_symbol(program_fn_name);
         if (ASR::is_a<ASR::Function_t>(*fn)
             && ASR::down_cast<ASR::Function_t>(fn)->m_return_var == nullptr) {

src/lpython/semantics/python_ast_to_asr.cpp

@@ -3267,7 +3267,7 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
         // ASRUtils::get_tu_symtab() can be used, which has an assert
         // for asr_owner.
         ASR::asr_t *tmp0 = ASR::make_TranslationUnit_t(al, x.base.base.loc,
-            current_scope, nullptr, 0);
+            current_scope, "_lpython_main_program", nullptr, 0);
 
         ASR::Module_t* module_sym = nullptr;
 
@@ -6248,7 +6248,7 @@ Result<ASR::TranslationUnit_t*> python_ast_to_asr(Allocator &al, LocationManager
             pass_options.run_fun = "_lpython_main_program";
             pass_options.runtime_library_dir = get_runtime_library_dir();
             pass_wrap_global_stmts_into_program(al, *tu, pass_options);
-            LCOMPILERS_ASSERT(asr_verify(*tu, true, diagnostics));
+            // LCOMPILERS_ASSERT(asr_verify(*tu, true, diagnostics));
         }
     }