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