Skip to content

Commit ea2fc92

Browse files
support to print dataclasses and tuples in REPL
1 parent c58d31c commit ea2fc92

File tree

6 files changed

+234
-12
lines changed

6 files changed

+234
-12
lines changed

src/libasr/codegen/evaluator.cpp

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,11 @@ llvm::Function *LLVMModule::get_function(const std::string &fn_name) {
9595
return m->getFunction(fn_name);
9696
}
9797

98+
llvm::GlobalVariable *LLVMModule::get_global(const std::string &global_name) {
99+
llvm::Module *m = m_m.get();
100+
return m->getNamedGlobal(global_name);
101+
}
102+
98103
std::string LLVMModule::get_return_type(const std::string &fn_name)
99104
{
100105
llvm::Module *m = m_m.get();

src/libasr/codegen/evaluator.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ namespace llvm {
1818
class LLVMContext;
1919
class Module;
2020
class Function;
21+
class GlobalVariable;
2122
class TargetMachine;
2223
class DataLayout;
2324
namespace orc {
@@ -37,6 +38,7 @@ class LLVMModule
3738
// Return a function return type as a string (real / integer)
3839
std::string get_return_type(const std::string &fn_name);
3940
llvm::Function *get_function(const std::string &fn_name);
41+
llvm::GlobalVariable *get_global(const std::string &global_name);
4042
};
4143

4244
class LLVMEvaluator

src/libasr/pass/global_stmts.cpp

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,8 @@ void pass_wrap_global_stmts(Allocator &al,
5252
(ASRUtils::expr_type(value)->type == ASR::ttypeType::Complex) ||
5353
(ASRUtils::expr_type(value)->type == ASR::ttypeType::Character) ||
5454
(ASRUtils::expr_type(value)->type == ASR::ttypeType::List) ||
55-
(ASRUtils::expr_type(value)->type == ASR::ttypeType::Tuple)) {
55+
(ASRUtils::expr_type(value)->type == ASR::ttypeType::Tuple) ||
56+
(ASRUtils::expr_type(value)->type == ASR::ttypeType::StructType)) {
5657
s.from_str(al, fn_name_s + std::to_string(idx));
5758
var_name = s.c_str(al);
5859
type = ASRUtils::expr_type(value);
@@ -102,6 +103,13 @@ void pass_wrap_global_stmts(Allocator &al,
102103
unit.m_symtab->add_symbol(global_underscore_name, down_cast<ASR::symbol_t>(global_underscore));
103104
ASR::stmt_t* asr_stmt = ASRUtils::STMT(ASR::make_Assignment_t(al, loc, global_underscore_ref, return_var_ref, nullptr));
104105
body.push_back(al, asr_stmt);
106+
107+
if ((ASRUtils::expr_type(return_var_ref)->type == ASR::ttypeType::List) ||
108+
(ASRUtils::expr_type(return_var_ref)->type == ASR::ttypeType::Tuple) ||
109+
(ASRUtils::expr_type(return_var_ref)->type == ASR::ttypeType::StructType)) {
110+
return_var_ref = nullptr;
111+
return_var = nullptr;
112+
}
105113
}
106114

107115
ASR::asr_t *fn = ASRUtils::make_Function_t_util(

src/lpython/python_evaluator.cpp

Lines changed: 54 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
#ifdef HAVE_LFORTRAN_LLVM
1717
#include <libasr/codegen/evaluator.h>
1818
#include <libasr/codegen/asr_to_llvm.h>
19+
#include <llvm/IR/GlobalVariable.h>
1920
#include <llvm/IR/Function.h>
2021
#include <llvm/IR/Type.h>
2122
#include <llvm/IR/DataLayout.h>
@@ -123,7 +124,8 @@ Result<PythonCompiler::EvalResult> PythonCompiler::evaluate(
123124
result.llvm_ir = m->str();
124125
}
125126

126-
ASR::symbol_t *global_underscore_symbol = symbol_table->get_symbol("_" + run_fn);
127+
ASR::symbol_t *global_underscore_symbol = ASR::down_cast<ASR::Module_t>(symbol_table->resolve_symbol(module_name))
128+
->m_symtab->get_symbol("_" + run_fn);
127129
if (global_underscore_symbol) {
128130
global_underscore_name = "_" + run_fn;
129131
}
@@ -134,19 +136,24 @@ Result<PythonCompiler::EvalResult> PythonCompiler::evaluate(
134136
call_run_fn = true;
135137
}
136138

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

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

146-
llvm::Function *fn = m->get_function(run_fn);
147-
llvm::Type *llvm_type = fn->getReturnType();
148-
LCOMPILERS_ASSERT(llvm_type->isStructTy())
149-
compute_offsets(llvm_type, global_underscore_sym, result);
150+
llvm::GlobalVariable *g = m->get_global("_" + run_fn);
151+
LCOMPILERS_ASSERT(g)
152+
llvm::Type *llvm_type = g->getValueType();
153+
if (llvm_type->isStructTy()) {
154+
struct_to_print = true;
155+
compute_offsets(llvm_type, global_underscore_sym, result);
156+
}
150157
}
151158

152159
e->add_module(std::move(m));
@@ -246,7 +253,14 @@ Result<PythonCompiler::EvalResult> PythonCompiler::evaluate(
246253
}
247254
} else if (return_type == "void") {
248255
e->execfn<void>(run_fn);
249-
result.type = EvalResult::statement;
256+
if (global_underscore_sym && struct_to_print) {
257+
void *r = (void*)e->get_symbol_address("_" + run_fn);
258+
LCOMPILERS_ASSERT(r)
259+
result.structure.structure = r;
260+
result.type = EvalResult::struct_type;
261+
} else {
262+
result.type = EvalResult::statement;
263+
}
250264
} else if (return_type == "none") {
251265
result.type = EvalResult::none;
252266
} else {
@@ -259,16 +273,18 @@ Result<PythonCompiler::EvalResult> PythonCompiler::evaluate(
259273
->erase_symbol(run_fn);
260274
}
261275
if (global_underscore_symbol) {
262-
if (symbol_table->resolve_symbol("_")) {
263-
symbol_table->erase_symbol("_");
276+
if (ASR::down_cast<ASR::Module_t>(symbol_table->resolve_symbol(module_name))->m_symtab->resolve_symbol("_")) {
277+
ASR::down_cast<ASR::Module_t>(symbol_table->resolve_symbol(module_name))->m_symtab
278+
->erase_symbol("_");
264279
}
265280
ASR::Variable_t *a = ASR::down_cast<ASR::Variable_t>(global_underscore_symbol);
266281
ASR::Variable_t *b = al.make_new<ASR::Variable_t>();
267282
*b = *a;
268283
Str s;
269284
s.from_str(al, "_");
270285
b->m_name = s.c_str(al);
271-
symbol_table->add_symbol("_", ASR::down_cast<ASR::symbol_t>((ASR::asr_t*)b));
286+
ASR::down_cast<ASR::Module_t>(symbol_table->resolve_symbol(module_name))->m_symtab
287+
->add_symbol("_", ASR::down_cast<ASR::symbol_t>((ASR::asr_t*)b));
272288
}
273289

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

534+
} else if (asr_type->type == ASR::ttypeType::Tuple) {
535+
ASR::Tuple_t *tuple_type = ASR::down_cast<ASR::Tuple_t>(asr_type);
536+
result += "(";
537+
for (size_t i = 0; i < tuple_type->n_type - 1; i++) {
538+
print_type(tuple_type->m_type[i], ((char*)data)+offsets[i], result);
539+
result += ", ";
540+
}
541+
print_type(tuple_type->m_type[tuple_type->n_type - 1], ((char*)data)+offsets[tuple_type->n_type - 1], result);
542+
result += ")";
543+
544+
} else if (asr_type->type == ASR::ttypeType::StructType) {
545+
ASR::StructType_t *class_type = ASR::down_cast<ASR::StructType_t>(asr_type);
546+
ASR::Struct_t *struct_info = ASR::down_cast<ASR::Struct_t>(class_type->m_derived_type);
547+
LCOMPILERS_ASSERT(class_type->n_data_member_types == struct_info->n_members)
548+
result += struct_info->m_name;
549+
result += "(";
550+
for (size_t i = 0; i < struct_info->n_members - 1; i++) {
551+
result += struct_info->m_members[i];
552+
result += "=";
553+
print_type(class_type->m_data_member_types[i], ((char*)data)+offsets[i], result);
554+
result += ", ";
555+
}
556+
result += struct_info->m_members[struct_info->n_members - 1];
557+
result += "=";
558+
print_type(class_type->m_data_member_types[struct_info->n_members - 1], ((char*)data)+offsets[struct_info->n_members - 1], result);
559+
result += ")";
560+
518561
} else {
519562
throw LCompilersException("PythonCompiler::evaluate(): Return type not supported");
520563
}

src/lpython/semantics/python_ast_to_asr.cpp

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5097,6 +5097,16 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
50975097
// Erase the function in TranslationUnit
50985098
unit->m_symtab->erase_symbol(func_name);
50995099
}
5100+
ASR::symbol_t *g_sym = unit->m_symtab->get_symbol("_" + func_name);
5101+
if (g_sym) {
5102+
// Move the `global_underscore` variable into the
5103+
// Module from TranslationUnit
5104+
ASR::Variable_t *f = ASR::down_cast<ASR::Variable_t>(g_sym);
5105+
f->m_parent_symtab = mod->m_symtab;
5106+
mod->m_symtab->add_symbol("_" + func_name, (ASR::symbol_t *) f);
5107+
// Erase the function in TranslationUnit
5108+
unit->m_symtab->erase_symbol("_" + func_name);
5109+
}
51005110
items.p = nullptr;
51015111
items.n = 0;
51025112
}

src/lpython/tests/test_llvm.cpp

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1506,6 +1506,160 @@ TEST_CASE("PythonCompiler lists") {
15061506
CHECK(e.aggregate_type_to_string(r.result) == "[\"lfortran\", \"lpython\", \"lc\"]");
15071507
}
15081508

1509+
TEST_CASE("PythonCompiler tuples") {
1510+
CompilerOptions cu;
1511+
cu.po.disable_main = true;
1512+
cu.emit_debug_line_column = false;
1513+
cu.generate_object_code = false;
1514+
cu.interactive = true;
1515+
cu.po.runtime_library_dir = LCompilers::LPython::get_runtime_library_dir();
1516+
PythonCompiler e(cu);
1517+
LCompilers::Result<PythonCompiler::EvalResult>
1518+
1519+
r = e.evaluate2("(1, 2)");
1520+
CHECK(r.ok);
1521+
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
1522+
CHECK(LCompilers::ASRUtils::get_type_code(r.result.structure.ttype) == "tuple[i32, i32]");
1523+
CHECK(e.aggregate_type_to_string(r.result) == "(1, 2)");
1524+
1525+
r = e.evaluate2("(1, 2, 2.5)");
1526+
CHECK(r.ok);
1527+
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
1528+
CHECK(LCompilers::ASRUtils::get_type_code(r.result.structure.ttype) == "tuple[i32, i32, r64]");
1529+
CHECK(e.aggregate_type_to_string(r.result) == "(1, 2, 2.500000)");
1530+
1531+
r = e.evaluate2("(1, 2, 2.5, \"LPython\")");
1532+
CHECK(r.ok);
1533+
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
1534+
CHECK(LCompilers::ASRUtils::get_type_code(r.result.structure.ttype) == "tuple[i32, i32, r64, str]");
1535+
CHECK(e.aggregate_type_to_string(r.result) == "(1, 2, 2.500000, \"LPython\")");
1536+
1537+
r = e.evaluate2("(1, 2, 2.5, \"LPython\", True)");
1538+
CHECK(r.ok);
1539+
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
1540+
CHECK(LCompilers::ASRUtils::get_type_code(r.result.structure.ttype) == "tuple[i32, i32, r64, str, i1]");
1541+
CHECK(e.aggregate_type_to_string(r.result) == "(1, 2, 2.500000, \"LPython\", True)");
1542+
1543+
r = e.evaluate2("(i8(1), i16(1), i64(1))");
1544+
CHECK(r.ok);
1545+
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
1546+
CHECK(LCompilers::ASRUtils::get_type_code(r.result.structure.ttype) == "tuple[i8, i16, i64]");
1547+
CHECK(e.aggregate_type_to_string(r.result) == "(1, 1, 1)");
1548+
1549+
r = e.evaluate2("(f32(1.0),)");
1550+
CHECK(r.ok);
1551+
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
1552+
CHECK(LCompilers::ASRUtils::get_type_code(r.result.structure.ttype) == "tuple[r32]");
1553+
CHECK(e.aggregate_type_to_string(r.result) == "(1.000000)");
1554+
}
1555+
1556+
TEST_CASE("PythonCompiler classes") {
1557+
CompilerOptions cu;
1558+
cu.po.disable_main = true;
1559+
cu.emit_debug_line_column = false;
1560+
cu.generate_object_code = false;
1561+
cu.interactive = true;
1562+
cu.po.runtime_library_dir = LCompilers::LPython::get_runtime_library_dir();
1563+
PythonCompiler e(cu);
1564+
LCompilers::Result<PythonCompiler::EvalResult>
1565+
1566+
r = e.evaluate2(R"(
1567+
@dataclass
1568+
class MyClass1:
1569+
x: i32
1570+
)");
1571+
CHECK(r.ok);
1572+
CHECK(r.result.type == PythonCompiler::EvalResult::none);
1573+
1574+
r = e.evaluate2("c1: MyClass1 = MyClass1(12)");
1575+
CHECK(r.ok);
1576+
CHECK(r.result.type == PythonCompiler::EvalResult::statement);
1577+
1578+
r = e.evaluate2("c1");
1579+
CHECK(r.ok);
1580+
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
1581+
CHECK(e.aggregate_type_to_string(r.result) == "MyClass1(x=12)");
1582+
1583+
r = e.evaluate2(R"(
1584+
@dataclass
1585+
class MyClass2:
1586+
i: i32
1587+
f: f64
1588+
)");
1589+
CHECK(r.ok);
1590+
CHECK(r.result.type == PythonCompiler::EvalResult::none);
1591+
1592+
r = e.evaluate2("c2: MyClass2 = MyClass2(12, 2.5)");
1593+
CHECK(r.ok);
1594+
CHECK(r.result.type == PythonCompiler::EvalResult::statement);
1595+
1596+
r = e.evaluate2("c2");
1597+
CHECK(r.ok);
1598+
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
1599+
CHECK(e.aggregate_type_to_string(r.result) == "MyClass2(i=12, f=2.500000)");
1600+
1601+
r = e.evaluate2(R"(
1602+
@dataclass
1603+
class MyClass3:
1604+
i: i32
1605+
f: f64
1606+
s: str
1607+
)");
1608+
CHECK(r.ok);
1609+
CHECK(r.result.type == PythonCompiler::EvalResult::none);
1610+
1611+
r = e.evaluate2("c3: MyClass3 = MyClass3(12, 2.5, \"LPython\")");
1612+
CHECK(r.ok);
1613+
CHECK(r.result.type == PythonCompiler::EvalResult::statement);
1614+
1615+
r = e.evaluate2("c3");
1616+
CHECK(r.ok);
1617+
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
1618+
CHECK(e.aggregate_type_to_string(r.result) == "MyClass3(i=12, f=2.500000, s=\"LPython\")");
1619+
1620+
r = e.evaluate2(R"(
1621+
@dataclass
1622+
class MyClass4:
1623+
i_1: bool
1624+
i_8: i8
1625+
i_16: i16
1626+
i_32: i32
1627+
i_64: i64
1628+
)");
1629+
CHECK(r.ok);
1630+
CHECK(r.result.type == PythonCompiler::EvalResult::none);
1631+
1632+
r = e.evaluate2("c4: MyClass4 = MyClass4(True, i8(2), i16(3), i32(4), i64(5))");
1633+
CHECK(r.ok);
1634+
CHECK(r.result.type == PythonCompiler::EvalResult::statement);
1635+
1636+
r = e.evaluate2("c4");
1637+
CHECK(r.ok);
1638+
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
1639+
CHECK(e.aggregate_type_to_string(r.result) == "MyClass4(i_1=True, i_8=2, i_16=3, i_32=4, i_64=5)");
1640+
1641+
r = e.evaluate2(R"(
1642+
@dataclass
1643+
class MyClass5:
1644+
u_1: bool
1645+
u_8: u8
1646+
u_16: u16
1647+
u_32: u32
1648+
u_64: u64
1649+
)");
1650+
CHECK(r.ok);
1651+
CHECK(r.result.type == PythonCompiler::EvalResult::none);
1652+
1653+
r = e.evaluate2("c5: MyClass5 = MyClass5(False, u8(2), u16(3), u32(4), u64(5))");
1654+
CHECK(r.ok);
1655+
CHECK(r.result.type == PythonCompiler::EvalResult::statement);
1656+
1657+
r = e.evaluate2("c5");
1658+
CHECK(r.ok);
1659+
CHECK(r.result.type == PythonCompiler::EvalResult::struct_type);
1660+
CHECK(e.aggregate_type_to_string(r.result) == "MyClass5(u_1=False, u_8=2, u_16=3, u_32=4, u_64=5)");
1661+
}
1662+
15091663
TEST_CASE("PythonCompiler underscore 1") {
15101664
CompilerOptions cu;
15111665
cu.po.disable_main = true;

0 commit comments

Comments
 (0)