examples/Kaleidoscope/Chapter4/toy.cpp

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