Fix the JIT code for the Kaleidoscope tutorial.
[oota-llvm.git] / examples / Kaleidoscope / Chapter4 / toy.cpp
1 #include "llvm/Analysis/Passes.h"
2 #include "llvm/ExecutionEngine/ExecutionEngine.h"
3 #include "llvm/ExecutionEngine/MCJIT.h"
4 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
5 #include "llvm/IR/DataLayout.h"
6 #include "llvm/IR/DerivedTypes.h"
7 #include "llvm/IR/IRBuilder.h"
8 #include "llvm/IR/LLVMContext.h"
9 #include "llvm/IR/Module.h"
10 #include "llvm/IR/Verifier.h"
11 #include "llvm/PassManager.h"
12 #include "llvm/Support/TargetSelect.h"
13 #include "llvm/Transforms/Scalar.h"
14 #include <cctype>
15 #include <cstdio>
16 #include <map>
17 #include <string>
18 #include <vector>
19 using namespace llvm;
20
21 //===----------------------------------------------------------------------===//
22 // Lexer
23 //===----------------------------------------------------------------------===//
24
25 // The lexer returns tokens [0-255] if it is an unknown character, otherwise one
26 // of these for known things.
27 enum Token {
28   tok_eof = -1,
29
30   // commands
31   tok_def = -2, tok_extern = -3,
32
33   // primary
34   tok_identifier = -4, tok_number = -5
35 };
36
37 static std::string IdentifierStr;  // Filled in if tok_identifier
38 static double NumVal;              // Filled in if tok_number
39
40 /// gettok - Return the next token from standard input.
41 static int gettok() {
42   static int LastChar = ' ';
43
44   // Skip any whitespace.
45   while (isspace(LastChar))
46     LastChar = getchar();
47
48   if (isalpha(LastChar)) { // identifier: [a-zA-Z][a-zA-Z0-9]*
49     IdentifierStr = LastChar;
50     while (isalnum((LastChar = getchar())))
51       IdentifierStr += LastChar;
52
53     if (IdentifierStr == "def") return tok_def;
54     if (IdentifierStr == "extern") return tok_extern;
55     return tok_identifier;
56   }
57
58   if (isdigit(LastChar) || LastChar == '.') {   // Number: [0-9.]+
59     std::string NumStr;
60     do {
61       NumStr += LastChar;
62       LastChar = getchar();
63     } while (isdigit(LastChar) || LastChar == '.');
64
65     NumVal = strtod(NumStr.c_str(), 0);
66     return tok_number;
67   }
68
69   if (LastChar == '#') {
70     // Comment until end of line.
71     do LastChar = getchar();
72     while (LastChar != EOF && LastChar != '\n' && LastChar != '\r');
73     
74     if (LastChar != EOF)
75       return gettok();
76   }
77   
78   // Check for end of file.  Don't eat the EOF.
79   if (LastChar == EOF)
80     return tok_eof;
81
82   // Otherwise, just return the character as its ascii value.
83   int ThisChar = LastChar;
84   LastChar = getchar();
85   return ThisChar;
86 }
87
88 //===----------------------------------------------------------------------===//
89 // Abstract Syntax Tree (aka Parse Tree)
90 //===----------------------------------------------------------------------===//
91 namespace {
92 /// ExprAST - Base class for all expression nodes.
93 class ExprAST {
94 public:
95   virtual ~ExprAST() {}
96   virtual Value *Codegen() = 0;
97 };
98
99 /// NumberExprAST - Expression class for numeric literals like "1.0".
100 class NumberExprAST : public ExprAST {
101   double Val;
102 public:
103   NumberExprAST(double val) : Val(val) {}
104   virtual Value *Codegen();
105 };
106
107 /// VariableExprAST - Expression class for referencing a variable, like "a".
108 class VariableExprAST : public ExprAST {
109   std::string Name;
110 public:
111   VariableExprAST(const std::string &name) : Name(name) {}
112   virtual Value *Codegen();
113 };
114
115 /// BinaryExprAST - Expression class for a binary operator.
116 class BinaryExprAST : public ExprAST {
117   char Op;
118   ExprAST *LHS, *RHS;
119 public:
120   BinaryExprAST(char op, ExprAST *lhs, ExprAST *rhs) 
121     : Op(op), LHS(lhs), RHS(rhs) {}
122   virtual Value *Codegen();
123 };
124
125 /// CallExprAST - Expression class for function calls.
126 class CallExprAST : public ExprAST {
127   std::string Callee;
128   std::vector<ExprAST*> Args;
129 public:
130   CallExprAST(const std::string &callee, std::vector<ExprAST*> &args)
131     : Callee(callee), Args(args) {}
132   virtual Value *Codegen();
133 };
134
135 /// PrototypeAST - This class represents the "prototype" for a function,
136 /// which captures its name, and its argument names (thus implicitly the number
137 /// of arguments the function takes).
138 class PrototypeAST {
139   std::string Name;
140   std::vector<std::string> Args;
141 public:
142   PrototypeAST(const std::string &name, const std::vector<std::string> &args)
143     : Name(name), Args(args) {}
144   
145   Function *Codegen();
146 };
147
148 /// FunctionAST - This class represents a function definition itself.
149 class FunctionAST {
150   PrototypeAST *Proto;
151   ExprAST *Body;
152 public:
153   FunctionAST(PrototypeAST *proto, ExprAST *body)
154     : Proto(proto), Body(body) {}
155   
156   Function *Codegen();
157 };
158 } // end anonymous namespace
159
160 //===----------------------------------------------------------------------===//
161 // Parser
162 //===----------------------------------------------------------------------===//
163
164 /// CurTok/getNextToken - Provide a simple token buffer.  CurTok is the current
165 /// token the parser is looking at.  getNextToken reads another token from the
166 /// lexer and updates CurTok with its results.
167 static int CurTok;
168 static int getNextToken() {
169   return CurTok = gettok();
170 }
171
172 /// BinopPrecedence - This holds the precedence for each binary operator that is
173 /// defined.
174 static std::map<char, int> BinopPrecedence;
175
176 /// GetTokPrecedence - Get the precedence of the pending binary operator token.
177 static int GetTokPrecedence() {
178   if (!isascii(CurTok))
179     return -1;
180   
181   // Make sure it's a declared binop.
182   int TokPrec = BinopPrecedence[CurTok];
183   if (TokPrec <= 0) return -1;
184   return TokPrec;
185 }
186
187 /// Error* - These are little helper functions for error handling.
188 ExprAST *Error(const char *Str) { fprintf(stderr, "Error: %s\n", Str);return 0;}
189 PrototypeAST *ErrorP(const char *Str) { Error(Str); return 0; }
190 FunctionAST *ErrorF(const char *Str) { Error(Str); return 0; }
191
192 static ExprAST *ParseExpression();
193
194 /// identifierexpr
195 ///   ::= identifier
196 ///   ::= identifier '(' expression* ')'
197 static ExprAST *ParseIdentifierExpr() {
198   std::string IdName = IdentifierStr;
199   
200   getNextToken();  // eat identifier.
201   
202   if (CurTok != '(') // Simple variable ref.
203     return new VariableExprAST(IdName);
204   
205   // Call.
206   getNextToken();  // eat (
207   std::vector<ExprAST*> Args;
208   if (CurTok != ')') {
209     while (1) {
210       ExprAST *Arg = ParseExpression();
211       if (!Arg) return 0;
212       Args.push_back(Arg);
213
214       if (CurTok == ')') break;
215
216       if (CurTok != ',')
217         return Error("Expected ')' or ',' in argument list");
218       getNextToken();
219     }
220   }
221
222   // Eat the ')'.
223   getNextToken();
224   
225   return new CallExprAST(IdName, Args);
226 }
227
228 /// numberexpr ::= number
229 static ExprAST *ParseNumberExpr() {
230   ExprAST *Result = new NumberExprAST(NumVal);
231   getNextToken(); // consume the number
232   return Result;
233 }
234
235 /// parenexpr ::= '(' expression ')'
236 static ExprAST *ParseParenExpr() {
237   getNextToken();  // eat (.
238   ExprAST *V = ParseExpression();
239   if (!V) return 0;
240   
241   if (CurTok != ')')
242     return Error("expected ')'");
243   getNextToken();  // eat ).
244   return V;
245 }
246
247 /// primary
248 ///   ::= identifierexpr
249 ///   ::= numberexpr
250 ///   ::= parenexpr
251 static ExprAST *ParsePrimary() {
252   switch (CurTok) {
253   default: return Error("unknown token when expecting an expression");
254   case tok_identifier: return ParseIdentifierExpr();
255   case tok_number:     return ParseNumberExpr();
256   case '(':            return ParseParenExpr();
257   }
258 }
259
260 /// binoprhs
261 ///   ::= ('+' primary)*
262 static ExprAST *ParseBinOpRHS(int ExprPrec, ExprAST *LHS) {
263   // If this is a binop, find its precedence.
264   while (1) {
265     int TokPrec = GetTokPrecedence();
266     
267     // If this is a binop that binds at least as tightly as the current binop,
268     // consume it, otherwise we are done.
269     if (TokPrec < ExprPrec)
270       return LHS;
271     
272     // Okay, we know this is a binop.
273     int BinOp = CurTok;
274     getNextToken();  // eat binop
275     
276     // Parse the primary expression after the binary operator.
277     ExprAST *RHS = ParsePrimary();
278     if (!RHS) return 0;
279     
280     // If BinOp binds less tightly with RHS than the operator after RHS, let
281     // the pending operator take RHS as its LHS.
282     int NextPrec = GetTokPrecedence();
283     if (TokPrec < NextPrec) {
284       RHS = ParseBinOpRHS(TokPrec+1, RHS);
285       if (RHS == 0) return 0;
286     }
287     
288     // Merge LHS/RHS.
289     LHS = new BinaryExprAST(BinOp, LHS, RHS);
290   }
291 }
292
293 /// expression
294 ///   ::= primary binoprhs
295 ///
296 static ExprAST *ParseExpression() {
297   ExprAST *LHS = ParsePrimary();
298   if (!LHS) return 0;
299   
300   return ParseBinOpRHS(0, LHS);
301 }
302
303 /// prototype
304 ///   ::= id '(' id* ')'
305 static PrototypeAST *ParsePrototype() {
306   if (CurTok != tok_identifier)
307     return ErrorP("Expected function name in prototype");
308
309   std::string FnName = IdentifierStr;
310   getNextToken();
311   
312   if (CurTok != '(')
313     return ErrorP("Expected '(' in prototype");
314   
315   std::vector<std::string> ArgNames;
316   while (getNextToken() == tok_identifier)
317     ArgNames.push_back(IdentifierStr);
318   if (CurTok != ')')
319     return ErrorP("Expected ')' in prototype");
320   
321   // success.
322   getNextToken();  // eat ')'.
323   
324   return new PrototypeAST(FnName, ArgNames);
325 }
326
327 /// definition ::= 'def' prototype expression
328 static FunctionAST *ParseDefinition() {
329   getNextToken();  // eat def.
330   PrototypeAST *Proto = ParsePrototype();
331   if (Proto == 0) return 0;
332
333   if (ExprAST *E = ParseExpression())
334     return new FunctionAST(Proto, E);
335   return 0;
336 }
337
338 /// toplevelexpr ::= expression
339 static FunctionAST *ParseTopLevelExpr() {
340   if (ExprAST *E = ParseExpression()) {
341     // Make an anonymous proto.
342     PrototypeAST *Proto = new PrototypeAST("", std::vector<std::string>());
343     return new FunctionAST(Proto, E);
344   }
345   return 0;
346 }
347
348 /// external ::= 'extern' prototype
349 static PrototypeAST *ParseExtern() {
350   getNextToken();  // eat extern.
351   return ParsePrototype();
352 }
353
354 //===----------------------------------------------------------------------===//
355 // Code Generation
356 //===----------------------------------------------------------------------===//
357
358 static Module *TheModule;
359 static IRBuilder<> Builder(getGlobalContext());
360 static std::map<std::string, Value*> NamedValues;
361 static FunctionPassManager *TheFPM;
362
363 Value *ErrorV(const char *Str) { Error(Str); return 0; }
364
365 Value *NumberExprAST::Codegen() {
366   return ConstantFP::get(getGlobalContext(), APFloat(Val));
367 }
368
369 Value *VariableExprAST::Codegen() {
370   // Look this variable up in the function.
371   Value *V = NamedValues[Name];
372   return V ? V : ErrorV("Unknown variable name");
373 }
374
375 Value *BinaryExprAST::Codegen() {
376   Value *L = LHS->Codegen();
377   Value *R = RHS->Codegen();
378   if (L == 0 || R == 0) return 0;
379   
380   switch (Op) {
381   case '+': return Builder.CreateFAdd(L, R, "addtmp");
382   case '-': return Builder.CreateFSub(L, R, "subtmp");
383   case '*': return Builder.CreateFMul(L, R, "multmp");
384   case '<':
385     L = Builder.CreateFCmpULT(L, R, "cmptmp");
386     // Convert bool 0/1 to double 0.0 or 1.0
387     return Builder.CreateUIToFP(L, Type::getDoubleTy(getGlobalContext()),
388                                 "booltmp");
389   default: return ErrorV("invalid binary operator");
390   }
391 }
392
393 Value *CallExprAST::Codegen() {
394   // Look up the name in the global module table.
395   Function *CalleeF = TheModule->getFunction(Callee);
396   if (CalleeF == 0)
397     return ErrorV("Unknown function referenced");
398   
399   // If argument mismatch error.
400   if (CalleeF->arg_size() != Args.size())
401     return ErrorV("Incorrect # arguments passed");
402
403   std::vector<Value*> ArgsV;
404   for (unsigned i = 0, e = Args.size(); i != e; ++i) {
405     ArgsV.push_back(Args[i]->Codegen());
406     if (ArgsV.back() == 0) return 0;
407   }
408   
409   return Builder.CreateCall(CalleeF, ArgsV, "calltmp");
410 }
411
412 Function *PrototypeAST::Codegen() {
413   // Make the function type:  double(double,double) etc.
414   std::vector<Type*> Doubles(Args.size(),
415                              Type::getDoubleTy(getGlobalContext()));
416   FunctionType *FT = FunctionType::get(Type::getDoubleTy(getGlobalContext()),
417                                        Doubles, false);
418   
419   Function *F = Function::Create(FT, Function::ExternalLinkage, Name, TheModule);
420   
421   // If F conflicted, there was already something named 'Name'.  If it has a
422   // body, don't allow redefinition or reextern.
423   if (F->getName() != Name) {
424     // Delete the one we just made and get the existing one.
425     F->eraseFromParent();
426     F = TheModule->getFunction(Name);
427     
428     // If F already has a body, reject this.
429     if (!F->empty()) {
430       ErrorF("redefinition of function");
431       return 0;
432     }
433     
434     // If F took a different number of args, reject.
435     if (F->arg_size() != Args.size()) {
436       ErrorF("redefinition of function with different # args");
437       return 0;
438     }
439   }
440   
441   // Set names for all arguments.
442   unsigned Idx = 0;
443   for (Function::arg_iterator AI = F->arg_begin(); Idx != Args.size();
444        ++AI, ++Idx) {
445     AI->setName(Args[Idx]);
446     
447     // Add arguments to variable symbol table.
448     NamedValues[Args[Idx]] = AI;
449   }
450   
451   return F;
452 }
453
454 Function *FunctionAST::Codegen() {
455   NamedValues.clear();
456   
457   Function *TheFunction = Proto->Codegen();
458   if (TheFunction == 0)
459     return 0;
460   
461   // Create a new basic block to start insertion into.
462   BasicBlock *BB = BasicBlock::Create(getGlobalContext(), "entry", TheFunction);
463   Builder.SetInsertPoint(BB);
464   
465   if (Value *RetVal = Body->Codegen()) {
466     // Finish off the function.
467     Builder.CreateRet(RetVal);
468
469     // Validate the generated code, checking for consistency.
470     verifyFunction(*TheFunction);
471
472     // Optimize the function.
473     TheFPM->run(*TheFunction);
474     
475     return TheFunction;
476   }
477   
478   // Error reading body, remove function.
479   TheFunction->eraseFromParent();
480   return 0;
481 }
482
483 //===----------------------------------------------------------------------===//
484 // Top-Level parsing and JIT Driver
485 //===----------------------------------------------------------------------===//
486
487 static ExecutionEngine *TheExecutionEngine;
488
489 static void HandleDefinition() {
490   if (FunctionAST *F = ParseDefinition()) {
491     if (Function *LF = F->Codegen()) {
492       fprintf(stderr, "Read function definition:");
493       LF->dump();
494     }
495   } else {
496     // Skip token for error recovery.
497     getNextToken();
498   }
499 }
500
501 static void HandleExtern() {
502   if (PrototypeAST *P = ParseExtern()) {
503     if (Function *F = P->Codegen()) {
504       fprintf(stderr, "Read extern: ");
505       F->dump();
506     }
507   } else {
508     // Skip token for error recovery.
509     getNextToken();
510   }
511 }
512
513 static void HandleTopLevelExpression() {
514   // Evaluate a top-level expression into an anonymous function.
515   if (FunctionAST *F = ParseTopLevelExpr()) {
516     if (Function *LF = F->Codegen()) {
517       TheExecutionEngine->finalizeObject();
518       // JIT the function, returning a function pointer.
519       void *FPtr = TheExecutionEngine->getPointerToFunction(LF);
520       
521       // Cast it to the right type (takes no arguments, returns a double) so we
522       // can call it as a native function.
523       double (*FP)() = (double (*)())(intptr_t)FPtr;
524       fprintf(stderr, "Evaluated to %f\n", FP());
525     }
526   } else {
527     // Skip token for error recovery.
528     getNextToken();
529   }
530 }
531
532 /// top ::= definition | external | expression | ';'
533 static void MainLoop() {
534   while (1) {
535     fprintf(stderr, "ready> ");
536     switch (CurTok) {
537     case tok_eof:    return;
538     case ';':        getNextToken(); break;  // ignore top-level semicolons.
539     case tok_def:    HandleDefinition(); break;
540     case tok_extern: HandleExtern(); break;
541     default:         HandleTopLevelExpression(); break;
542     }
543   }
544 }
545
546 //===----------------------------------------------------------------------===//
547 // "Library" functions that can be "extern'd" from user code.
548 //===----------------------------------------------------------------------===//
549
550 /// putchard - putchar that takes a double and returns 0.
551 extern "C" 
552 double putchard(double X) {
553   putchar((char)X);
554   return 0;
555 }
556
557 //===----------------------------------------------------------------------===//
558 // Main driver code.
559 //===----------------------------------------------------------------------===//
560
561 int main() {
562   InitializeNativeTarget();
563   InitializeNativeTargetAsmPrinter();
564   InitializeNativeTargetAsmParser();
565   LLVMContext &Context = getGlobalContext();
566
567   // Install standard binary operators.
568   // 1 is lowest precedence.
569   BinopPrecedence['<'] = 10;
570   BinopPrecedence['+'] = 20;
571   BinopPrecedence['-'] = 20;
572   BinopPrecedence['*'] = 40;  // highest.
573
574   // Prime the first token.
575   fprintf(stderr, "ready> ");
576   getNextToken();
577
578   // Make the module, which holds all the code.
579   std::unique_ptr<Module> Owner = make_unique<Module>("my cool jit", Context);
580   TheModule = Owner.get();
581
582   // Create the JIT.  This takes ownership of the module.
583   std::string ErrStr;
584   TheExecutionEngine = EngineBuilder(std::move(Owner))
585                            .setErrorStr(&ErrStr)
586                            .setMCJITMemoryManager(new SectionMemoryManager())
587                            .create();
588   if (!TheExecutionEngine) {
589     fprintf(stderr, "Could not create ExecutionEngine: %s\n", ErrStr.c_str());
590     exit(1);
591   }
592
593   FunctionPassManager OurFPM(TheModule);
594
595   // Set up the optimizer pipeline.  Start with registering info about how the
596   // target lays out data structures.
597   TheModule->setDataLayout(TheExecutionEngine->getDataLayout());
598   OurFPM.add(new DataLayoutPass());
599   // Provide basic AliasAnalysis support for GVN.
600   OurFPM.add(createBasicAliasAnalysisPass());
601   // Do simple "peephole" optimizations and bit-twiddling optzns.
602   OurFPM.add(createInstructionCombiningPass());
603   // Reassociate expressions.
604   OurFPM.add(createReassociatePass());
605   // Eliminate Common SubExpressions.
606   OurFPM.add(createGVNPass());
607   // Simplify the control flow graph (deleting unreachable blocks, etc).
608   OurFPM.add(createCFGSimplificationPass());
609
610   OurFPM.doInitialization();
611
612   // Set the global so the code gen can use this.
613   TheFPM = &OurFPM;
614
615   // Run the main "interpreter loop" now.
616   MainLoop();
617
618   TheFPM = 0;
619
620   // Print out all of the generated code.
621   TheModule->dump();
622
623   return 0;
624 }