Skip to content

Commit

Permalink
support to print dataclasses and tuples in REPL
Browse files Browse the repository at this point in the history
  • Loading branch information
Vipul-Cariappa committed Jul 22, 2024
1 parent c58d31c commit ea2fc92
Show file tree
Hide file tree
Showing 6 changed files with 234 additions and 12 deletions.
5 changes: 5 additions & 0 deletions src/libasr/codegen/evaluator.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,11 @@ llvm::Function *LLVMModule::get_function(const std::string &fn_name) {
return m->getFunction(fn_name);
}

llvm::GlobalVariable *LLVMModule::get_global(const std::string &global_name) {
llvm::Module *m = m_m.get();
return m->getNamedGlobal(global_name);
}

std::string LLVMModule::get_return_type(const std::string &fn_name)
{
llvm::Module *m = m_m.get();
Expand Down
2 changes: 2 additions & 0 deletions src/libasr/codegen/evaluator.h
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ namespace llvm {
class LLVMContext;
class Module;
class Function;
class GlobalVariable;
class TargetMachine;
class DataLayout;
namespace orc {
Expand All @@ -37,6 +38,7 @@ class LLVMModule
// Return a function return type as a string (real / integer)
std::string get_return_type(const std::string &fn_name);
llvm::Function *get_function(const std::string &fn_name);
llvm::GlobalVariable *get_global(const std::string &global_name);
};

class LLVMEvaluator
Expand Down
10 changes: 9 additions & 1 deletion src/libasr/pass/global_stmts.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,8 @@ void pass_wrap_global_stmts(Allocator &al,
(ASRUtils::expr_type(value)->type == ASR::ttypeType::Complex) ||
(ASRUtils::expr_type(value)->type == ASR::ttypeType::Character) ||
(ASRUtils::expr_type(value)->type == ASR::ttypeType::List) ||
(ASRUtils::expr_type(value)->type == ASR::ttypeType::Tuple)) {
(ASRUtils::expr_type(value)->type == ASR::ttypeType::Tuple) ||
(ASRUtils::expr_type(value)->type == ASR::ttypeType::StructType)) {
s.from_str(al, fn_name_s + std::to_string(idx));
var_name = s.c_str(al);
type = ASRUtils::expr_type(value);
Expand Down Expand Up @@ -102,6 +103,13 @@ void pass_wrap_global_stmts(Allocator &al,
unit.m_symtab->add_symbol(global_underscore_name, down_cast<ASR::symbol_t>(global_underscore));
ASR::stmt_t* asr_stmt = ASRUtils::STMT(ASR::make_Assignment_t(al, loc, global_underscore_ref, return_var_ref, nullptr));
body.push_back(al, asr_stmt);

if ((ASRUtils::expr_type(return_var_ref)->type == ASR::ttypeType::List) ||
(ASRUtils::expr_type(return_var_ref)->type == ASR::ttypeType::Tuple) ||
(ASRUtils::expr_type(return_var_ref)->type == ASR::ttypeType::StructType)) {
return_var_ref = nullptr;
return_var = nullptr;
}
}

ASR::asr_t *fn = ASRUtils::make_Function_t_util(
Expand Down
65 changes: 54 additions & 11 deletions src/lpython/python_evaluator.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
#ifdef HAVE_LFORTRAN_LLVM
#include <libasr/codegen/evaluator.h>
#include <libasr/codegen/asr_to_llvm.h>
#include <llvm/IR/GlobalVariable.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/Type.h>
#include <llvm/IR/DataLayout.h>
Expand Down Expand Up @@ -123,7 +124,8 @@ Result<PythonCompiler::EvalResult> PythonCompiler::evaluate(
result.llvm_ir = m->str();
}

ASR::symbol_t *global_underscore_symbol = symbol_table->get_symbol("_" + run_fn);
ASR::symbol_t *global_underscore_symbol = ASR::down_cast<ASR::Module_t>(symbol_table->resolve_symbol(module_name))
->m_symtab->get_symbol("_" + run_fn);
if (global_underscore_symbol) {
global_underscore_name = "_" + run_fn;
}
Expand All @@ -134,19 +136,24 @@ Result<PythonCompiler::EvalResult> PythonCompiler::evaluate(
call_run_fn = true;
}

ASR::symbol_t *global_underscore_sym = symbol_table->get_symbol(global_underscore_name);
if ((return_type == "struct") && (global_underscore_sym)) {
bool struct_to_print = false;
ASR::symbol_t *global_underscore_sym = ASR::down_cast<ASR::Module_t>(symbol_table->resolve_symbol(module_name))
->m_symtab->get_symbol("_" + run_fn);
if ((return_type == "void") && (global_underscore_sym)) {
// we compute the offsets of the struct's attribute here
// we will be using it later in aggregate_type_to_string to print the struct

// we compute the offsets here instead of computing it in aggregate_type_to_string
// because once we call `e->add_module`, internally LLVM may deallocate all the
// type info after compiling the IR into machine code

llvm::Function *fn = m->get_function(run_fn);
llvm::Type *llvm_type = fn->getReturnType();
LCOMPILERS_ASSERT(llvm_type->isStructTy())
compute_offsets(llvm_type, global_underscore_sym, result);
llvm::GlobalVariable *g = m->get_global("_" + run_fn);
LCOMPILERS_ASSERT(g)
llvm::Type *llvm_type = g->getValueType();
if (llvm_type->isStructTy()) {
struct_to_print = true;
compute_offsets(llvm_type, global_underscore_sym, result);
}
}

e->add_module(std::move(m));
Expand Down Expand Up @@ -246,7 +253,14 @@ Result<PythonCompiler::EvalResult> PythonCompiler::evaluate(
}
} else if (return_type == "void") {
e->execfn<void>(run_fn);
result.type = EvalResult::statement;
if (global_underscore_sym && struct_to_print) {
void *r = (void*)e->get_symbol_address("_" + run_fn);
LCOMPILERS_ASSERT(r)
result.structure.structure = r;
result.type = EvalResult::struct_type;
} else {
result.type = EvalResult::statement;
}
} else if (return_type == "none") {
result.type = EvalResult::none;
} else {
Expand All @@ -259,16 +273,18 @@ Result<PythonCompiler::EvalResult> PythonCompiler::evaluate(
->erase_symbol(run_fn);
}
if (global_underscore_symbol) {
if (symbol_table->resolve_symbol("_")) {
symbol_table->erase_symbol("_");
if (ASR::down_cast<ASR::Module_t>(symbol_table->resolve_symbol(module_name))->m_symtab->resolve_symbol("_")) {
ASR::down_cast<ASR::Module_t>(symbol_table->resolve_symbol(module_name))->m_symtab
->erase_symbol("_");
}
ASR::Variable_t *a = ASR::down_cast<ASR::Variable_t>(global_underscore_symbol);
ASR::Variable_t *b = al.make_new<ASR::Variable_t>();
*b = *a;
Str s;
s.from_str(al, "_");
b->m_name = s.c_str(al);
symbol_table->add_symbol("_", ASR::down_cast<ASR::symbol_t>((ASR::asr_t*)b));
ASR::down_cast<ASR::Module_t>(symbol_table->resolve_symbol(module_name))->m_symtab
->add_symbol("_", ASR::down_cast<ASR::symbol_t>((ASR::asr_t*)b));
}

eval_count++;
Expand Down Expand Up @@ -515,6 +531,33 @@ std::string PythonCompiler::aggregate_type_to_string(const struct EvalResult &r)
print_type(element_ttype, ((char*)array)+((size - 1)*element_size), result);
result += "]";

} else if (asr_type->type == ASR::ttypeType::Tuple) {
ASR::Tuple_t *tuple_type = ASR::down_cast<ASR::Tuple_t>(asr_type);
result += "(";
for (size_t i = 0; i < tuple_type->n_type - 1; i++) {
print_type(tuple_type->m_type[i], ((char*)data)+offsets[i], result);
result += ", ";
}
print_type(tuple_type->m_type[tuple_type->n_type - 1], ((char*)data)+offsets[tuple_type->n_type - 1], result);
result += ")";

} else if (asr_type->type == ASR::ttypeType::StructType) {
ASR::StructType_t *class_type = ASR::down_cast<ASR::StructType_t>(asr_type);
ASR::Struct_t *struct_info = ASR::down_cast<ASR::Struct_t>(class_type->m_derived_type);
LCOMPILERS_ASSERT(class_type->n_data_member_types == struct_info->n_members)
result += struct_info->m_name;
result += "(";
for (size_t i = 0; i < struct_info->n_members - 1; i++) {
result += struct_info->m_members[i];
result += "=";
print_type(class_type->m_data_member_types[i], ((char*)data)+offsets[i], result);
result += ", ";
}
result += struct_info->m_members[struct_info->n_members - 1];
result += "=";
print_type(class_type->m_data_member_types[struct_info->n_members - 1], ((char*)data)+offsets[struct_info->n_members - 1], result);
result += ")";

} else {
throw LCompilersException("PythonCompiler::evaluate(): Return type not supported");
}
Expand Down
10 changes: 10 additions & 0 deletions src/lpython/semantics/python_ast_to_asr.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -5097,6 +5097,16 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
// Erase the function in TranslationUnit
unit->m_symtab->erase_symbol(func_name);
}
ASR::symbol_t *g_sym = unit->m_symtab->get_symbol("_" + func_name);
if (g_sym) {
// Move the `global_underscore` variable into the
// Module from TranslationUnit
ASR::Variable_t *f = ASR::down_cast<ASR::Variable_t>(g_sym);
f->m_parent_symtab = mod->m_symtab;
mod->m_symtab->add_symbol("_" + func_name, (ASR::symbol_t *) f);
// Erase the function in TranslationUnit
unit->m_symtab->erase_symbol("_" + func_name);
}
items.p = nullptr;
items.n = 0;
}
Expand Down
154 changes: 154 additions & 0 deletions src/lpython/tests/test_llvm.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1506,6 +1506,160 @@ TEST_CASE("PythonCompiler lists") {
CHECK(e.aggregate_type_to_string(r.result) == "[\"lfortran\", \"lpython\", \"lc\"]");
}

TEST_CASE("PythonCompiler tuples") {
CompilerOptions cu;
cu.po.disable_main = true;
cu.emit_debug_line_column = false;
cu.generate_object_code = false;
cu.interactive = true;
cu.po.runtime_library_dir = LCompilers::LPython::get_runtime_library_dir();
PythonCompiler e(cu);
LCompilers::Result<PythonCompiler::EvalResult>

r = e.evaluate2("(1, 2)");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
CHECK(LCompilers::ASRUtils::get_type_code(r.result.structure.ttype) == "tuple[i32, i32]");
CHECK(e.aggregate_type_to_string(r.result) == "(1, 2)");

r = e.evaluate2("(1, 2, 2.5)");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
CHECK(LCompilers::ASRUtils::get_type_code(r.result.structure.ttype) == "tuple[i32, i32, r64]");
CHECK(e.aggregate_type_to_string(r.result) == "(1, 2, 2.500000)");

r = e.evaluate2("(1, 2, 2.5, \"LPython\")");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
CHECK(LCompilers::ASRUtils::get_type_code(r.result.structure.ttype) == "tuple[i32, i32, r64, str]");
CHECK(e.aggregate_type_to_string(r.result) == "(1, 2, 2.500000, \"LPython\")");

r = e.evaluate2("(1, 2, 2.5, \"LPython\", True)");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
CHECK(LCompilers::ASRUtils::get_type_code(r.result.structure.ttype) == "tuple[i32, i32, r64, str, i1]");
CHECK(e.aggregate_type_to_string(r.result) == "(1, 2, 2.500000, \"LPython\", True)");

r = e.evaluate2("(i8(1), i16(1), i64(1))");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
CHECK(LCompilers::ASRUtils::get_type_code(r.result.structure.ttype) == "tuple[i8, i16, i64]");
CHECK(e.aggregate_type_to_string(r.result) == "(1, 1, 1)");

r = e.evaluate2("(f32(1.0),)");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
CHECK(LCompilers::ASRUtils::get_type_code(r.result.structure.ttype) == "tuple[r32]");
CHECK(e.aggregate_type_to_string(r.result) == "(1.000000)");
}

TEST_CASE("PythonCompiler classes") {
CompilerOptions cu;
cu.po.disable_main = true;
cu.emit_debug_line_column = false;
cu.generate_object_code = false;
cu.interactive = true;
cu.po.runtime_library_dir = LCompilers::LPython::get_runtime_library_dir();
PythonCompiler e(cu);
LCompilers::Result<PythonCompiler::EvalResult>

r = e.evaluate2(R"(
@dataclass
class MyClass1:
x: i32
)");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::none);

r = e.evaluate2("c1: MyClass1 = MyClass1(12)");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::statement);

r = e.evaluate2("c1");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
CHECK(e.aggregate_type_to_string(r.result) == "MyClass1(x=12)");

r = e.evaluate2(R"(
@dataclass
class MyClass2:
i: i32
f: f64
)");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::none);

r = e.evaluate2("c2: MyClass2 = MyClass2(12, 2.5)");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::statement);

r = e.evaluate2("c2");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
CHECK(e.aggregate_type_to_string(r.result) == "MyClass2(i=12, f=2.500000)");

r = e.evaluate2(R"(
@dataclass
class MyClass3:
i: i32
f: f64
s: str
)");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::none);

r = e.evaluate2("c3: MyClass3 = MyClass3(12, 2.5, \"LPython\")");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::statement);

r = e.evaluate2("c3");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
CHECK(e.aggregate_type_to_string(r.result) == "MyClass3(i=12, f=2.500000, s=\"LPython\")");

r = e.evaluate2(R"(
@dataclass
class MyClass4:
i_1: bool
i_8: i8
i_16: i16
i_32: i32
i_64: i64
)");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::none);

r = e.evaluate2("c4: MyClass4 = MyClass4(True, i8(2), i16(3), i32(4), i64(5))");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::statement);

r = e.evaluate2("c4");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
CHECK(e.aggregate_type_to_string(r.result) == "MyClass4(i_1=True, i_8=2, i_16=3, i_32=4, i_64=5)");

r = e.evaluate2(R"(
@dataclass
class MyClass5:
u_1: bool
u_8: u8
u_16: u16
u_32: u32
u_64: u64
)");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::none);

r = e.evaluate2("c5: MyClass5 = MyClass5(False, u8(2), u16(3), u32(4), u64(5))");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::statement);

r = e.evaluate2("c5");
CHECK(r.ok);
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
CHECK(e.aggregate_type_to_string(r.result) == "MyClass5(u_1=False, u_8=2, u_16=3, u_32=4, u_64=5)");
}

TEST_CASE("PythonCompiler underscore 1") {
CompilerOptions cu;
cu.po.disable_main = true;
Expand Down

0 comments on commit ea2fc92

Please sign in to comment.