Fix sub-imports

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};

codon/parser/visitors/typecheck/typecheck.cpp

@@ -23,8 +23,10 @@ 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) {
+    // LOG("{}", so->toString(2));
     v.error("cannot typecheck the program");
   }
   if (s->getSuite())
@@ -283,13 +285,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 +324,8 @@ int TypecheckVisitor::canCall(const types::FuncTypePtr &fn,
       score = -1;
     }
   }
+  if (score >= 0)
+    score += (real_gi == fn->funcGenerics.size());
   return score;
 }
 
@@ -336,6 +342,8 @@ TypecheckVisitor::findMatchingMethods(const types::ClassTypePtr &typ,
       continue; // avoid overloads that have not been seen yet
     auto method = ctx->instantiate(mi, typ)->getFunc();
     int score = canCall(method, args);
+    // LOG("{}: {} {} :: {} :: {}", getSrcInfo(), method->debugString(2), args, score,
+    //     method->ast->getSrcInfo());
     if (score != -1) {
       results.push_back(mi);
     }

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