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 IRBuilder<> Builder(getGlobalContext());
392 static std::map<std::string, Value *> NamedValues;
393 static std::unique_ptr<legacy::FunctionPassManager> TheFPM;
394 static std::unique_ptr<KaleidoscopeJIT> TheJIT;
395 static std::map<std::string, std::unique_ptr<PrototypeAST>> FunctionProtos;
396 
397 Value *LogErrorV(const char *Str) {
398   LogError(Str);
399   return nullptr;
400 }
401 
402 Function *getFunction(std::string Name) {
403   // First, see if the function has already been added to the current module.
404   if (auto *F = TheModule->getFunction(Name))
405     return F;
406 
407   // If not, check whether we can codegen the declaration from some existing
408   // prototype.
409   auto FI = FunctionProtos.find(Name);
410   if (FI != FunctionProtos.end())
411     return FI->second->codegen();
412 
413   // If no existing prototype exists, return null.
414   return nullptr;
415 }
416 
417 Value *NumberExprAST::codegen() {
418   return ConstantFP::get(getGlobalContext(), APFloat(Val));
419 }
420 
421 Value *VariableExprAST::codegen() {
422   // Look this variable up in the function.
423   Value *V = NamedValues[Name];
424   if (!V)
425     return LogErrorV("Unknown variable name");
426   return V;
427 }
428 
429 Value *BinaryExprAST::codegen() {
430   Value *L = LHS->codegen();
431   Value *R = RHS->codegen();
432   if (!L || !R)
433     return nullptr;
434 
435   switch (Op) {
436   case '+':
437     return Builder.CreateFAdd(L, R, "addtmp");
438   case '-':
439     return Builder.CreateFSub(L, R, "subtmp");
440   case '*':
441     return Builder.CreateFMul(L, R, "multmp");
442   case '<':
443     L = Builder.CreateFCmpULT(L, R, "cmptmp");
444     // Convert bool 0/1 to double 0.0 or 1.0
445     return Builder.CreateUIToFP(L, Type::getDoubleTy(getGlobalContext()),
446                                 "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(),
475                               Type::getDoubleTy(getGlobalContext()));
476   FunctionType *FT =
477       FunctionType::get(Type::getDoubleTy(getGlobalContext()), Doubles, false);
478 
479   Function *F =
480       Function::Create(FT, Function::ExternalLinkage, Name, TheModule.get());
481 
482   // Set names for all arguments.
483   unsigned Idx = 0;
484   for (auto &Arg : F->args())
485     Arg.setName(Args[Idx++]);
486 
487   return F;
488 }
489 
490 Function *FunctionAST::codegen() {
491   // Transfer ownership of the prototype to the FunctionProtos map, but keep a
492   // reference to it for use below.
493   auto &P = *Proto;
494   FunctionProtos[Proto->getName()] = std::move(Proto);
495   Function *TheFunction = getFunction(P.getName());
496   if (!TheFunction)
497     return nullptr;
498 
499   // Create a new basic block to start insertion into.
500   BasicBlock *BB = BasicBlock::Create(getGlobalContext(), "entry", TheFunction);
501   Builder.SetInsertPoint(BB);
502 
503   // Record the function arguments in the NamedValues map.
504   NamedValues.clear();
505   for (auto &Arg : TheFunction->args())
506     NamedValues[Arg.getName()] = &Arg;
507 
508   if (Value *RetVal = Body->codegen()) {
509     // Finish off the function.
510     Builder.CreateRet(RetVal);
511 
512     // Validate the generated code, checking for consistency.
513     verifyFunction(*TheFunction);
514 
515     // Run the optimizer on the function.
516     TheFPM->run(*TheFunction);
517 
518     return TheFunction;
519   }
520 
521   // Error reading body, remove function.
522   TheFunction->eraseFromParent();
523   return nullptr;
524 }
525 
526 //===----------------------------------------------------------------------===//
527 // Top-Level parsing and JIT Driver
528 //===----------------------------------------------------------------------===//
529 
530 static void InitializeModuleAndPassManager() {
531   // Open a new module.
532   TheModule = llvm::make_unique<Module>("my cool jit", getGlobalContext());
533   TheModule->setDataLayout(TheJIT->getTargetMachine().createDataLayout());
534 
535   // Create a new pass manager attached to it.
536   TheFPM = llvm::make_unique<legacy::FunctionPassManager>(TheModule.get());
537 
538   // Do simple "peephole" optimizations and bit-twiddling optzns.
539   TheFPM->add(createInstructionCombiningPass());
540   // Reassociate expressions.
541   TheFPM->add(createReassociatePass());
542   // Eliminate Common SubExpressions.
543   TheFPM->add(createGVNPass());
544   // Simplify the control flow graph (deleting unreachable blocks, etc).
545   TheFPM->add(createCFGSimplificationPass());
546 
547   TheFPM->doInitialization();
548 }
549 
550 static void HandleDefinition() {
551   if (auto FnAST = ParseDefinition()) {
552     if (auto *FnIR = FnAST->codegen()) {
553       fprintf(stderr, "Read function definition:");
554       FnIR->dump();
555       TheJIT->addModule(std::move(TheModule));
556       InitializeModuleAndPassManager();
557     }
558   } else {
559     // Skip token for error recovery.
560     getNextToken();
561   }
562 }
563 
564 static void HandleExtern() {
565   if (auto ProtoAST = ParseExtern()) {
566     if (auto *FnIR = ProtoAST->codegen()) {
567       fprintf(stderr, "Read extern: ");
568       FnIR->dump();
569       FunctionProtos[ProtoAST->getName()] = std::move(ProtoAST);
570     }
571   } else {
572     // Skip token for error recovery.
573     getNextToken();
574   }
575 }
576 
577 static void HandleTopLevelExpression() {
578   // Evaluate a top-level expression into an anonymous function.
579   if (auto FnAST = ParseTopLevelExpr()) {
580     if (FnAST->codegen()) {
581 
582       // JIT the module containing the anonymous expression, keeping a handle so
583       // we can free it later.
584       auto H = TheJIT->addModule(std::move(TheModule));
585       InitializeModuleAndPassManager();
586 
587       // Search the JIT for the __anon_expr symbol.
588       auto ExprSymbol = TheJIT->findSymbol("__anon_expr");
589       assert(ExprSymbol && "Function not found");
590 
591       // Get the symbol's address and cast it to the right type (takes no
592       // arguments, returns a double) so we can call it as a native function.
593       double (*FP)() = (double (*)())(intptr_t)ExprSymbol.getAddress();
594       fprintf(stderr, "Evaluated to %f\n", FP());
595 
596       // Delete the anonymous expression module from the JIT.
597       TheJIT->removeModule(H);
598     }
599   } else {
600     // Skip token for error recovery.
601     getNextToken();
602   }
603 }
604 
605 /// top ::= definition | external | expression | ';'
606 static void MainLoop() {
607   while (1) {
608     fprintf(stderr, "ready> ");
609     switch (CurTok) {
610     case tok_eof:
611       return;
612     case ';': // ignore top-level semicolons.
613       getNextToken();
614       break;
615     case tok_def:
616       HandleDefinition();
617       break;
618     case tok_extern:
619       HandleExtern();
620       break;
621     default:
622       HandleTopLevelExpression();
623       break;
624     }
625   }
626 }
627 
628 //===----------------------------------------------------------------------===//
629 // "Library" functions that can be "extern'd" from user code.
630 //===----------------------------------------------------------------------===//
631 
632 /// putchard - putchar that takes a double and returns 0.
633 extern "C" double putchard(double X) {
634   fputc((char)X, stderr);
635   return 0;
636 }
637 
638 /// printd - printf that takes a double prints it as "%f\n", returning 0.
639 extern "C" double printd(double X) {
640   fprintf(stderr, "%f\n", X);
641   return 0;
642 }
643 
644 //===----------------------------------------------------------------------===//
645 // Main driver code.
646 //===----------------------------------------------------------------------===//
647 
648 int main() {
649   InitializeNativeTarget();
650   InitializeNativeTargetAsmPrinter();
651   InitializeNativeTargetAsmParser();
652 
653   // Install standard binary operators.
654   // 1 is lowest precedence.
655   BinopPrecedence['<'] = 10;
656   BinopPrecedence['+'] = 20;
657   BinopPrecedence['-'] = 20;
658   BinopPrecedence['*'] = 40; // highest.
659 
660   // Prime the first token.
661   fprintf(stderr, "ready> ");
662   getNextToken();
663 
664   TheJIT = llvm::make_unique<KaleidoscopeJIT>();
665 
666   InitializeModuleAndPassManager();
667 
668   // Run the main "interpreter loop" now.
669   MainLoop();
670 
671   return 0;
672 }
673