Fix: Clean reset and correct alias patch

Author: amritamishra01

src/lpython/semantics/python_ast_to_asr.cpp

@@ -4872,49 +4872,50 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
         }
     }
 
-    // --- 1. RESTORED & FIXED: visit_ImportFrom ---
     void visit_ImportFrom(const AST::ImportFrom_t &x) {
         if (!x.m_module) {
             throw SemanticError("Not implemented: The import statement must currently specify the module name", x.base.base.loc);
         }
         std::string msym = x.m_module; // Module name
 
         // Get the module, for now assuming it is not loaded, so we load it:
-        ASR::symbol_t *t = nullptr; 
-        std::string rl_path = get_runtime_library_dir();
-        std::vector<std::string> paths;
-        for (auto &path:import_paths) {
-            paths.push_back(path);
-        }
-        paths.push_back(rl_path);
-        paths.push_back(parent_dir);
-
-        // FIX: Initialize bools to false to avoid undefined behavior
-        bool lpython = false, enum_py = false, copy = false, sympy = false;
-        
-        set_module_symbol(msym, paths);
-        t = (ASR::symbol_t*)(load_module(al, global_scope,
-            msym, x.base.base.loc, diag, lm, false, paths, lpython, enum_py, copy, sympy,
-            [&](const std::string &msg, const Location &loc) { throw SemanticError(msg, loc); },
-            allow_implicit_casting));
-
-        if (lpython || enum_py || copy || sympy) {
-            tmp = nullptr;
-            return;
-        }
+        ASR::symbol_t *t = nullptr; // current_scope->parent->resolve_symbol(msym);
         if (!t) {
-            throw SemanticError("The module '" + msym + "' cannot be loaded",
-                    x.base.base.loc);
-        }
-        if( msym == "__init__" ) {
-            for( auto item: ASRUtils::symbol_symtab(t)->get_scope() ) {
-                if( ASR::is_a<ASR::ExternalSymbol_t>(*item.second) ) {
-                    current_module_dependencies.push_back(al,
-                        ASR::down_cast<ASR::ExternalSymbol_t>(item.second)->m_module_name);
+            std::string rl_path = get_runtime_library_dir();
+            std::vector<std::string> paths;
+            for (auto &path:import_paths) {
+                paths.push_back(path);
+            }
+            paths.push_back(rl_path);
+            paths.push_back(parent_dir);
+
+            bool lpython, enum_py, copy, sympy;
+            set_module_symbol(msym, paths);
+            t = (ASR::symbol_t*)(load_module(al, global_scope,
+                msym, x.base.base.loc, diag, lm, false, paths, lpython, enum_py, copy, sympy,
+                [&](const std::string &msg, const Location &loc) { throw SemanticError(msg, loc); },
+                allow_implicit_casting));
+            if (lpython || enum_py || copy || sympy) {
+                // TODO: For now we skip lpython import completely. Later on we should note what symbols
+                // got imported from it, and give an error message if an annotation is used without
+                // importing it.
+                tmp = nullptr;
+                return;
+            }
+            if (!t) {
+                throw SemanticError("The module '" + msym + "' cannot be loaded",
+                        x.base.base.loc);
+            }
+            if( msym == "__init__" ) {
+                for( auto item: ASRUtils::symbol_symtab(t)->get_scope() ) {
+                    if( ASR::is_a<ASR::ExternalSymbol_t>(*item.second) ) {
+                        current_module_dependencies.push_back(al,
+                            ASR::down_cast<ASR::ExternalSymbol_t>(item.second)->m_module_name);
+                    }
                 }
+            } else {
+                current_module_dependencies.push_back(al, s2c(al, msym));
             }
-        } else {
-            current_module_dependencies.push_back(al, s2c(al, msym));
         }
 
         ASR::Module_t *m = ASR::down_cast<ASR::Module_t>(t);
@@ -4929,29 +4930,26 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
             if (x.m_names[j].m_asname) {
                 new_sym_name = ASRUtils::get_mangled_name(m, x.m_names[j].m_asname);
             }
-            
-            // FIX: Renamed 't' to 'imported_sym' to avoid shadowing the outer module pointer 't'
-            ASR::symbol_t *imported_sym = import_from_module(al, m, current_scope, msym,
-                                        remote_sym, new_sym_name, x.m_names[i].loc, true);
-                                        
+            ASR::symbol_t *t = import_from_module(al, m, current_scope, msym,
+                                remote_sym, new_sym_name, x.m_names[i].loc, true);
             if (current_scope->get_scope().find(new_sym_name) != current_scope->get_scope().end()) {
                 ASR::symbol_t *old_sym = current_scope->get_scope().find(new_sym_name)->second;
                 diag.add(diag::Diagnostic(
                     "The symbol '" + new_sym_name + "' imported from " + std::string(m->m_name) +" will shadow the existing symbol '" + new_sym_name + "'",
                     diag::Level::Warning, diag::Stage::Semantic, {
                         diag::Label("old symbol", {old_sym->base.loc}),
-                        diag::Label("new symbol", {imported_sym->base.loc}),
+                        diag::Label("new symbol", {t->base.loc}),
                     })
                 );
-                current_scope->overwrite_symbol(new_sym_name, imported_sym);
+                current_scope->overwrite_symbol(new_sym_name, t);
             } else {
-                current_scope->add_symbol(new_sym_name, imported_sym);
+                current_scope->add_symbol(new_sym_name, t);
             }
         }
+
         tmp = nullptr;
     }
 
-    // --- 2. FIXED: visit_Import (Safe Bools + No Else) ---
     void visit_Import(const AST::Import_t &x) {
         ASR::symbol_t *t = nullptr;
         std::string rl_path = get_runtime_library_dir();
@@ -4966,7 +4964,7 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
             std::string mod_sym = x.m_names[i].m_name;
             char* alias = x.m_names[i].m_asname; 
 
-            // FIX: Initialize bools
+            // Initialize bools
             bool lpython = false, enum_py = false, copy = false, sympy = false;
 
             set_module_symbol(mod_sym, paths);
@@ -4981,7 +4979,7 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
                 continue;
             }
             if (!t) {
-                throw SemanticError("The module '" + mod_sym + "' cannot be loaded",
+                throw SemanticError("The module " + mod_sym + " cannot be loaded",
                         x.base.base.loc);
             }
 
@@ -5000,23 +4998,98 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
             if (alias) {
                 std::string alias_str = std::string(alias);
                 if (current_scope->get_symbol(alias_str) == nullptr) {
-                    // Note: Using the signature that matches your local libasr version
                     ASR::asr_t *ext_sym = ASR::make_ExternalSymbol_t(
-                        al, 
-                        x.base.base.loc, 
-                        current_scope,
-                        s2c(al, alias_str), 
-                        t, 
-                        s2c(al, mod_sym),
-                        nullptr, 
-                        0, 
-                        s2c(al, mod_sym), 
-                        ASR::accessType::Public
+                        al, x.base.base.loc, current_scope,
+                        s2c(al, alias_str), t, s2c(al, mod_sym),
+                        nullptr, 0, s2c(al, mod_sym), ASR::accessType::Public
                     );
                     current_scope->add_symbol(alias_str, ASR::down_cast<ASR::symbol_t>(ext_sym));
                 }
             }
-            // NO ELSE BLOCK: Standard imports are handled by load_module logic implicitly.
+        }
+    }
+
+    void visit_AugAssign(const AST::AugAssign_t &/*x*/) {
+        // We skip this in the SymbolTable visitor, but visit it in the BodyVisitor
+    }
+
+    void visit_AnnAssign(const AST::AnnAssign_t &/*x*/) {
+        // We skip this in the SymbolTable visitor, but visit it in the BodyVisitor
+    }
+
+    void visit_Assign(const AST::Assign_t &x) {
+        /**
+         *  Type variables have to be put in the symbol table before checking function definitions.
+         *  This describes a different treatment for Assign in the form of `T = TypeVar('T', ...)`
+         */
+        if (x.n_targets == 1 && AST::is_a<AST::Call_t>(*x.m_value)) {
+            AST::Call_t *rh = AST::down_cast<AST::Call_t>(x.m_value);
+            if (AST::is_a<AST::Name_t>(*rh->m_func)) {
+                AST::Name_t *tv = AST::down_cast<AST::Name_t>(rh->m_func);
+                std::string f_name = tv->m_id;
+                if (strcmp(s2c(al, f_name), "TypeVar") == 0 &&
+                        rh->n_args > 0 && AST::is_a<AST::ConstantStr_t>(*rh->m_args[0])) {
+                    if (AST::is_a<AST::Name_t>(*x.m_targets[0])) {
+                        std::string tvar_name = AST::down_cast<AST::Name_t>(x.m_targets[0])->m_id;
+                        // Check if the type variable name is a reserved type keyword
+                        const char* type_list[15]
+                            = { "list", "set", "dict", "tuple",
+                                "i8", "i16", "i32", "i64", "f32",
+                                "f64", "c32", "c64", "str", "bool", "i1"};
+                        for (int i = 0; i < 15; i++) {
+                            if (strcmp(s2c(al, tvar_name), type_list[i]) == 0) {
+                                throw SemanticError(tvar_name + " is a reserved type, consider a different type variable name",
+                                    x.base.base.loc);
+                            }
+                        }
+                        // Check if the type variable is already defined
+                        if (current_scope->get_scope().find(tvar_name) !=
+                                current_scope->get_scope().end()) {
+                            ASR::symbol_t *orig_decl = current_scope->get_symbol(tvar_name);
+                            throw SemanticError(diag::Diagnostic(
+                                "Variable " + tvar_name + " is already declared in the same scope",
+                                diag::Level::Error, diag::Stage::Semantic, {
+                                    diag::Label("original declaration", {orig_decl->base.loc}, false),
+                                    diag::Label("redeclaration", {x.base.base.loc}),
+                            }));
+                        }
+
+                        // Build ttype
+                        Vec<ASR::dimension_t> dims;
+                        dims.reserve(al, 4);
+
+                        ASR::ttype_t *type = ASRUtils::TYPE(ASR::make_TypeParameter_t(al, x.base.base.loc,
+                            s2c(al, tvar_name)));
+                        type = ASRUtils::make_Array_t_util(al, x.base.base.loc, type, dims.p, dims.size());
+
+                        ASR::expr_t *value = nullptr;
+                        ASR::expr_t *init_expr = nullptr;
+                        ASR::intentType s_intent = ASRUtils::intent_local;
+                        ASR::storage_typeType storage_type = ASR::storage_typeType::Default;
+                        ASR::abiType current_procedure_abi_type = ASR::abiType::Source;
+                        ASR::accessType s_access = ASR::accessType::Public;
+                        ASR::presenceType s_presence = ASR::presenceType::Required;
+                        bool value_attr = false;
+
+                        SetChar variable_dependencies_vec;
+                        variable_dependencies_vec.reserve(al, 1);
+                        ASRUtils::collect_variable_dependencies(al, variable_dependencies_vec, type, init_expr, value);
+                        // Build the variable and add it to the scope
+                        ASR::asr_t *v = ASR::make_Variable_t(al, x.base.base.loc, current_scope,
+                            s2c(al, tvar_name), variable_dependencies_vec.p, variable_dependencies_vec.size(),
+                            s_intent, init_expr, value, storage_type, type, nullptr, current_procedure_abi_type,
+                            s_access, s_presence, value_attr, false, false, nullptr, false, false);
+                        current_scope->add_symbol(tvar_name, ASR::down_cast<ASR::symbol_t>(v));
+
+                        tmp = nullptr;
+
+                        return;
+                    } else {
+                        // This error might need to be further elaborated
+                        throw SemanticError("Type variable must be a variable", x.base.base.loc);
+                    }
+                }
+            }
         }
     }