1
0
mirror of https://github.com/exaloop/codon.git synced 2025-06-03 15:03:52 +08:00

Automatic @export and wrapping for toplevel functions

This commit is contained in:
Ibrahim Numanagić 2023-02-05 16:51:42 -08:00
parent 33933f44e1
commit fea42475b9
5 changed files with 43 additions and 15 deletions

View File

@ -673,6 +673,8 @@ void LLVMVisitor::writeToPythonExtension(
llvmFunc = createPyTryCatchWrapper(llvmFunc);
auto name = original->getUnmangledName();
if (ast::startswith(name, "._py_"))
name = name.substr(5);
auto *nameVar = new llvm::GlobalVariable(
*M, llvm::ArrayType::get(B->getInt8Ty(), name.length() + 1),
/*isConstant=*/true, llvm::GlobalValue::PrivateLinkage,

View File

@ -92,8 +92,10 @@ void PythonExtensionLowering::run(Module *module) {
if (!util::hasAttribute(f, EXPORT_ATTR))
continue;
if (auto *g = generateExtensionFunc(f))
if (auto *g = generateExtensionFunc(f)) {
LOG("[pyext] exporting {}", f->getName());
extFuncs.emplace_back(f, g);
}
}
}
}

View File

@ -166,6 +166,8 @@ struct SimplifyContext : public Context<SimplifyItem> {
bool allowTypeOf;
/// Set if all assignments should not be dominated later on.
bool avoidDomination = false;
/// Canonical names of functions that should be exported (e.g., for Python use)
std::unordered_set<std::string> makeExport;
public:
SimplifyContext(std::string filename, Cache *cache);

View File

@ -314,6 +314,10 @@ void SimplifyVisitor::visit(FunctionStmt *stmt) {
} else {
resultStmt = f;
}
if (ctx->isGlobal() && ctx->getModule().empty()) {
ctx->makeExport.insert(f->name);
}
}
/// Make a capturing anonymous function with the provided suite and argument names.

View File

@ -32,6 +32,7 @@ SimplifyVisitor::apply(Cache *cache, const StmtPtr &node, const std::string &fil
auto preamble = std::make_shared<std::vector<StmtPtr>>();
seqassertn(cache->module, "cache's module is not set");
#define N std::make_shared
// Load standard library if it has not been loaded
if (!in(cache->imports, STDLIB_IMPORT)) {
// Load the internal.__init__
@ -63,11 +64,9 @@ SimplifyVisitor::apply(Cache *cache, const StmtPtr &node, const std::string &fil
// Load early compile-time defines (for standard library)
preamble->push_back(
SimplifyVisitor(stdlib, preamble)
.transform(std::make_shared<AssignStmt>(
std::make_shared<IdExpr>(d.first),
std::make_shared<IntExpr>(d.second),
std::make_shared<IndexExpr>(std::make_shared<IdExpr>("Static"),
std::make_shared<IdExpr>("int")))));
.transform(
N<AssignStmt>(N<IdExpr>(d.first), N<IntExpr>(d.second),
N<IndexExpr>(N<IdExpr>("Static"), N<IdExpr>("int")))));
}
preamble->push_back(SimplifyVisitor(stdlib, preamble)
.transform(parseFile(stdlib->cache, stdlibPath->path)));
@ -86,22 +85,21 @@ SimplifyVisitor::apply(Cache *cache, const StmtPtr &node, const std::string &fil
ctx->moduleName = {ImportFile::PACKAGE, file, MODULE_MAIN};
// Prepare the code
auto suite = std::make_shared<SuiteStmt>();
auto suite = N<SuiteStmt>();
for (auto &d : defines) {
// Load compile-time defines (e.g., codon run -DFOO=1 ...)
suite->stmts.push_back(std::make_shared<AssignStmt>(
std::make_shared<IdExpr>(d.first), std::make_shared<IntExpr>(d.second),
std::make_shared<IndexExpr>(std::make_shared<IdExpr>("Static"),
std::make_shared<IdExpr>("int"))));
suite->stmts.push_back(
N<AssignStmt>(N<IdExpr>(d.first), N<IntExpr>(d.second),
N<IndexExpr>(N<IdExpr>("Static"), N<IdExpr>("int"))));
}
// Set up __name__
suite->stmts.push_back(std::make_shared<AssignStmt>(
std::make_shared<IdExpr>("__name__"), std::make_shared<StringExpr>(MODULE_MAIN)));
suite->stmts.push_back(
N<AssignStmt>(N<IdExpr>("__name__"), N<StringExpr>(MODULE_MAIN)));
suite->stmts.push_back(node);
auto n = SimplifyVisitor(ctx, preamble).transform(suite);
suite = std::make_shared<SuiteStmt>();
suite->stmts.push_back(std::make_shared<SuiteStmt>(*preamble));
suite = N<SuiteStmt>();
suite->stmts.push_back(N<SuiteStmt>(*preamble));
// Add dominated assignment declarations
if (in(ctx->scope.stmts, ctx->scope.blocks.back()))
suite->stmts.insert(suite->stmts.end(),
@ -109,6 +107,26 @@ SimplifyVisitor::apply(Cache *cache, const StmtPtr &node, const std::string &fil
ctx->scope.stmts[ctx->scope.blocks.back()].end());
suite->stmts.push_back(n);
auto exports = ctx->makeExport;
for (auto &fn : exports) {
std::vector<Param> args;
std::vector<ExprPtr> callArgs;
auto t = N<IdExpr>("pyobj");
t->markType();
for (auto &a : ctx->cache->functions[fn].ast->args) {
args.push_back(
Param{ctx->cache->rev(a.name), a.type ? a.type->clone() : t->clone()});
callArgs.push_back(N<IdExpr>(ctx->cache->rev(a.name)));
}
// TODO: what to do in case of overrides?
auto ast = N<FunctionStmt>("._py_" + ctx->cache->rev(fn), nullptr, args,
N<SuiteStmt>(N<ReturnStmt>(N<CallExpr>(
N<IdExpr>(ctx->cache->rev(fn)), callArgs))),
Attr({Attr::Export}));
suite->stmts.push_back(SimplifyVisitor(ctx, preamble).transform(ast));
}
#undef N
if (!ctx->cache->errors.empty())
throw exc::ParserException();