examples/Kaleidoscope/Chapter3/toy.cpp

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