Merge branch 'develop' into alloc-hoist-opt

exaloop · Feb 7, 2024 · 4143887 · 4143887
2 parents c2f7b88 + c16dfb5
commit 4143887
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 7 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -53,7 +53,7 @@ jobs:
 
       - name: Cache Dependencies
         id: cache-deps
-        uses: actions/cache@v3
+        uses: actions/cache@v4
         with:
           path: llvm
           key: manylinux-llvm
@@ -133,7 +133,7 @@ jobs:
 
       - name: Cache Dependencies
         id: cache-deps
-        uses: actions/cache@v3
+        uses: actions/cache@v4
         with:
           path: llvm
           key: ${{ runner.os }}-llvm

diff --git a/codon/parser/visitors/simplify/access.cpp b/codon/parser/visitors/simplify/access.cpp
@@ -230,10 +230,11 @@ SimplifyVisitor::getImport(const std::vector<std::string> &chain) {
 
   // Find the longest prefix that corresponds to the existing import
   // (e.g., `a.b.c.d` -> `a.b.c` if there is `import a.b.c`)
-  SimplifyContext::Item val = nullptr;
+  SimplifyContext::Item val = nullptr, importVal = nullptr;
   for (auto i = chain.size(); i-- > 0;) {
     val = ctx->find(join(chain, "/", 0, i + 1));
     if (val && val->isImport()) {
+      importVal = val;
       importName = val->importPath, importEnd = i + 1;
       break;
     }
@@ -254,6 +255,14 @@ SimplifyVisitor::getImport(const std::vector<std::string> &chain) {
         return {importEnd, val};
       } else {
         val = fctx->find(join(chain, ".", importEnd, i + 1));
+        if (val && i + 1 != chain.size() && val->isImport()) {
+          importVal = val;
+          importName = val->importPath;
+          importEnd = i + 1;
+          fctx = ctx->cache->imports[importName].ctx;
+          i = chain.size();
+          continue;
+        }
         if (val && (importName.empty() || val->isType() || !val->isConditional())) {
           itemName = val->canonicalName, itemEnd = i + 1;
           break;
@@ -264,10 +273,23 @@ SimplifyVisitor::getImport(const std::vector<std::string> &chain) {
       if (ctx->getBase()->pyCaptures)
         return {1, nullptr};
       E(Error::IMPORT_NO_MODULE, getSrcInfo(), chain[importEnd]);
-    }
-    if (itemName.empty())
+    } else if (itemName.empty()) {
+      if (!ctx->isStdlibLoading && endswith(importName, "__init__.codon")) {
+        auto import = ctx->cache->imports[importName];
+        auto file =
+            getImportFile(ctx->cache->argv0, chain[importEnd], importName, false,
+                          ctx->cache->module0, ctx->cache->pluginImportPaths);
+        if (file) {
+          auto s = SimplifyVisitor(import.ctx, preamble)
+                       .transform(N<ImportStmt>(N<IdExpr>(chain[importEnd]), nullptr));
+          prependStmts->push_back(s);
+          return getImport(chain);
+        }
+      }
+
       E(Error::IMPORT_NO_NAME, getSrcInfo(), chain[importEnd],
         ctx->cache->imports[importName].moduleName);
+    }
     importEnd = itemEnd;
   }
   return {importEnd, val};

diff --git a/codon/parser/visitors/typecheck/typecheck.cpp b/codon/parser/visitors/typecheck/typecheck.cpp
@@ -23,7 +23,8 @@ StmtPtr TypecheckVisitor::apply(Cache *cache, const StmtPtr &stmts) {
   if (!cache->typeCtx)
     cache->typeCtx = std::make_shared<TypeContext>(cache);
   TypecheckVisitor v(cache->typeCtx);
-  auto s = v.inferTypes(clone(stmts), true);
+  auto so = clone(stmts);
+  auto s = v.inferTypes(so, true);
   if (!s) {
     v.error("cannot typecheck the program");
   }
@@ -283,13 +284,15 @@ int TypecheckVisitor::canCall(const types::FuncTypePtr &fn,
         return 0;
       },
       [](error::Error, const SrcInfo &, const std::string &) { return -1; });
-  for (int ai = 0, mai = 0, gi = 0; score != -1 && ai < reordered.size(); ai++) {
+  int ai = 0, mai = 0, gi = 0, real_gi = 0;
+  for (; score != -1 && ai < reordered.size(); ai++) {
     auto expectTyp = fn->ast->args[ai].status == Param::Normal
                          ? fn->getArgTypes()[mai++]
                          : fn->funcGenerics[gi++].type;
     auto [argType, argTypeIdx] = reordered[ai];
     if (!argType)
       continue;
+    real_gi += fn->ast->args[ai].status != Param::Normal;
     if (fn->ast->args[ai].status != Param::Normal) {
       // Check if this is a good generic!
       if (expectTyp && expectTyp->isStaticType()) {
@@ -320,6 +323,8 @@ int TypecheckVisitor::canCall(const types::FuncTypePtr &fn,
       score = -1;
     }
   }
+  if (score >= 0)
+    score += (real_gi == fn->funcGenerics.size());
   return score;
 }
 

diff --git a/test/parser/a/sub/__init__.codon b/test/parser/a/sub/__init__.codon
@@ -0,0 +1,4 @@
+print('a.sub')
+
+def foo():
+    print('a.sub.foo')
diff --git a/test/parser/simplify_stmt.codon b/test/parser/simplify_stmt.codon
@@ -474,6 +474,13 @@ a.ha()  #: B
 
 print par  #: x
 
+#%% import_subimport,barebones
+import a as xa  #: a
+
+xa.foo()  #: a.foo
+#: a.sub
+xa.sub.foo()  #: a.sub.foo
+
 #%% import_order,barebones
 def foo():
     import a