diff --git a/chapter3.py b/chapter3.py index 307ecca..e46c17e 100644 --- a/chapter3.py +++ b/chapter3.py @@ -4,6 +4,7 @@ from collections import namedtuple from enum import Enum import llvmlite.ir as ir +import llvmlite.binding as llvm # Each token is a tuple of kind and value. kind is one of the enumeration values # in TokenKind. value is the textual value of the token in the input. @@ -340,6 +341,10 @@ class LLVMCodeGenerator(object): # names to ir.Value. self.func_symtab = {} + # Counter to disambiguate anonymous functions created from toplevel + # expressions. + self.anon_function_counter = 0 + def generate_code(self, node): assert isinstance(node, (PrototypeAST, FunctionAST)) self._codegen(node) @@ -351,55 +356,100 @@ class LLVMCodeGenerator(object): For AST node of class Foo, calls self._codegen_Foo. Each visitor is expected to return a llvmlite.ir.Value. """ - method = 'visit_' + node.__class__.__name__ + method = '_codegen_' + node.__class__.__name__ return getattr(self, method)(node) def _codegen_NumberExprAST(self, node): return self.builder.constant(ir.DoubleType(), float(node.val)) def _codegen_VariableExprAST(self, node): - return self.symtab[node.name] + return self.func_symtab[node.name] + def _codegen_BinaryExprAST(self, node): + lhs = self._codegen(node.lhs) + rhs = self._codegen(node.rhs) + + if node.op == '+': + return self.builder.fadd(lhs, rhs, 'addtmp') + elif node.op == '-': + return self.builder.fsub(lhs, rhs, 'subtmp') + elif node.op == '*': + return self.builder.fmul(lhs, rhs, 'multmp') + elif node.op == '<': + cmp = self.builder.fcmp_unordered('<', lhs, rhs, 'cmptmp') + return self.builder.uitofp(cmp, ir.DoubleType(), 'booltmp') + else: + raise CodegenError('Unknown binary operator', node.op) + + def _codegen_CallExprAST(self, node): + callee_func = self.module.globals.get(node.callee, None) + if callee_func is None or not isinstance(callee_func, ir.Function): + raise CodegenError('Call to unknown function', node.callee) + if len(callee_func.args) != len(node.args): + raise CodegenError('Call argument length mismatch', node.callee) + call_args = [self._codegen(arg) for arg in node.args] + return self.builder.call(callee_func, call_args, 'calltmp') + def _codegen_PrototypeAST(self, node): + funcname = node.name + if not funcname: + funcname = 'anon{0}'.format(self.anon_function_counter) + self.anon_function_counter += 1 + # Create a function type func_ty = ir.FunctionType(ir.DoubleType(), - [ir.DoubleType() * len(node.argnames)]) + [ir.DoubleType()] * len(node.argnames)) # If a function with this name already exists in the module... - if node.name in self.module.globals: + if funcname in self.module.globals: # We only allow the case in which a declaration exists and now the # function is defined (or redeclared) with the same number of args. - existing_func = self.module[node.name] + existing_func = self.module[funcname] if not isinstance(existing_func, ir.Function): - raise CodegenError('Function/Global name collision', node.name) + raise CodegenError('Function/Global name collision', funcname) if not existing_func.is_declaration(): - raise CodegenError('Redifinition of {0}'.format(node.name)) + raise CodegenError('Redifinition of {0}'.format(funcname)) if len(existing_func.function_type.args) != len(func_ty.args): raise CodegenError( 'Redifinition with different number of arguments') - func = self.module.globals[node.name] + func = self.module.globals[funcname] else: # Otherwise create a new function - func = ir.Function(self.module, func_ty, node.name) + func = ir.Function(self.module, func_ty, funcname) # Set function argument names from AST - for i in range(len(func.args)): - func.args[i].name = node.argnames[i] + for i, arg in enumerate(func.args): + arg.name = node.argnames[i] + self.func_symtab[arg.name] = arg return func def _codegen_FunctionAST(self, node): + # Reset the symbol table. Prototype generation will pre-populate it with + # function arguments. + self.func_symtab = {} # Create the function skeleton from the prototype. func = self._codegen(node.proto) # Create the entry BB in the function and set the builder to it. bb_entry = func.append_basic_block('entry') self.builder = ir.IRBuilder(bb_entry) - - # We're going to generate the function body now. Reset the symbol table. - self.func_symtab = {} retval = self._codegen(node.body) - irbuilder.ret(retval) + self.builder.ret(retval) return func if __name__ == '__main__': - pass + def parse(s): + return Parser(s).parse_toplevel() + dfoo = parse('def foo(a b) a + 4.1414 * (b < a)') + dcos = parse('extern cos(a)') + dcl = parse('2 + foo(3.14, cos(9.91))') + + cg = LLVMCodeGenerator() + mod = cg.generate_code(dfoo) + mod = cg.generate_code(dcos) + mod = cg.generate_code(dcl) + + print(mod) + + llvmmod = llvm.parse_assembly(str(mod)) + llvmmod.verify()