examples/Kaleidoscope/Chapter5/toy.cpp

   1 #include "llvm/ADT/STLExtras.h"
   2 #include "llvm/Analysis/BasicAliasAnalysis.h"
   3 #include "llvm/Analysis/Passes.h"
   4 #include "llvm/IR/IRBuilder.h"
   5 #include "llvm/IR/LLVMContext.h"
   6 #include "llvm/IR/LegacyPassManager.h"
   7 #include "llvm/IR/Module.h"
   8 #include "llvm/IR/Verifier.h"
   9 #include "llvm/Support/TargetSelect.h"
  10 #include "llvm/Transforms/Scalar.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   // control
  39   tok_if = -6,
  40   tok_then = -7,
  41   tok_else = -8,
  42   tok_for = -9,
  43   tok_in = -10
  44 };
  45
  46 static std::string IdentifierStr; // Filled in if tok_identifier
  47 static double NumVal;             // Filled in if tok_number
  48
  49 /// gettok - Return the next token from standard input.
  50 static int gettok() {
  51   static int LastChar = ' ';
  52
  53   // Skip any whitespace.
  54   while (isspace(LastChar))
  55     LastChar = getchar();
  56
  57   if (isalpha(LastChar)) { // identifier: [a-zA-Z][a-zA-Z0-9]*
  58     IdentifierStr = LastChar;
  59     while (isalnum((LastChar = getchar())))
  60       IdentifierStr += LastChar;
  61
  62     if (IdentifierStr == "def")
  63       return tok_def;
  64     if (IdentifierStr == "extern")
  65       return tok_extern;
  66     if (IdentifierStr == "if")
  67       return tok_if;
  68     if (IdentifierStr == "then")
  69       return tok_then;
  70     if (IdentifierStr == "else")
  71       return tok_else;
  72     if (IdentifierStr == "for")
  73       return tok_for;
  74     if (IdentifierStr == "in")
  75       return tok_in;
  76     return tok_identifier;
  77   }
  78
  79   if (isdigit(LastChar) || LastChar == '.') { // Number: [0-9.]+
  80     std::string NumStr;
  81     do {
  82       NumStr += LastChar;
  83       LastChar = getchar();
  84     } while (isdigit(LastChar) || LastChar == '.');
  85
  86     NumVal = strtod(NumStr.c_str(), 0);
  87     return tok_number;
  88   }
  89
  90   if (LastChar == '#') {
  91     // Comment until end of line.
  92     do
  93       LastChar = getchar();
  94     while (LastChar != EOF && LastChar != '\n' && LastChar != '\r');
  95
  96     if (LastChar != EOF)
  97       return gettok();
  98   }
  99
 100   // Check for end of file.  Don't eat the EOF.
 101   if (LastChar == EOF)
 102     return tok_eof;
 103
 104   // Otherwise, just return the character as its ascii value.
 105   int ThisChar = LastChar;
 106   LastChar = getchar();
 107   return ThisChar;
 108 }
 109
 110 //===----------------------------------------------------------------------===//
 111 // Abstract Syntax Tree (aka Parse Tree)
 112 //===----------------------------------------------------------------------===//
 113 namespace {
 114 /// ExprAST - Base class for all expression nodes.
 115 class ExprAST {
 116 public:
 117   virtual ~ExprAST() {}
 118   virtual Value *codegen() = 0;
 119 };
 120
 121 /// NumberExprAST - Expression class for numeric literals like "1.0".
 122 class NumberExprAST : public ExprAST {
 123   double Val;
 124
 125 public:
 126   NumberExprAST(double Val) : Val(Val) {}
 127   Value *codegen() override;
 128 };
 129
 130 /// VariableExprAST - Expression class for referencing a variable, like "a".
 131 class VariableExprAST : public ExprAST {
 132   std::string Name;
 133
 134 public:
 135   VariableExprAST(const std::string &Name) : Name(Name) {}
 136   Value *codegen() override;
 137 };
 138
 139 /// BinaryExprAST - Expression class for a binary operator.
 140 class BinaryExprAST : public ExprAST {
 141   char Op;
 142   std::unique_ptr<ExprAST> LHS, RHS;
 143
 144 public:
 145   BinaryExprAST(char Op, std::unique_ptr<ExprAST> LHS,
 146                 std::unique_ptr<ExprAST> RHS)
 147       : Op(Op), LHS(std::move(LHS)), RHS(std::move(RHS)) {}
 148   Value *codegen() override;
 149 };
 150
 151 /// CallExprAST - Expression class for function calls.
 152 class CallExprAST : public ExprAST {
 153   std::string Callee;
 154   std::vector<std::unique_ptr<ExprAST>> Args;
 155
 156 public:
 157   CallExprAST(const std::string &Callee,
 158               std::vector<std::unique_ptr<ExprAST>> Args)
 159       : Callee(Callee), Args(std::move(Args)) {}
 160   Value *codegen() override;
 161 };
 162
 163 /// IfExprAST - Expression class for if/then/else.
 164 class IfExprAST : public ExprAST {
 165   std::unique_ptr<ExprAST> Cond, Then, Else;
 166
 167 public:
 168   IfExprAST(std::unique_ptr<ExprAST> Cond, std::unique_ptr<ExprAST> Then,
 169             std::unique_ptr<ExprAST> Else)
 170       : Cond(std::move(Cond)), Then(std::move(Then)), Else(std::move(Else)) {}
 171   Value *codegen() override;
 172 };
 173
 174 /// ForExprAST - Expression class for for/in.
 175 class ForExprAST : public ExprAST {
 176   std::string VarName;
 177   std::unique_ptr<ExprAST> Start, End, Step, Body;
 178
 179 public:
 180   ForExprAST(const std::string &VarName, std::unique_ptr<ExprAST> Start,
 181              std::unique_ptr<ExprAST> End, std::unique_ptr<ExprAST> Step,
 182              std::unique_ptr<ExprAST> Body)
 183       : VarName(VarName), Start(std::move(Start)), End(std::move(End)),
 184         Step(std::move(Step)), Body(std::move(Body)) {}
 185   Value *codegen() override;
 186 };
 187
 188 /// PrototypeAST - This class represents the "prototype" for a function,
 189 /// which captures its name, and its argument names (thus implicitly the number
 190 /// of arguments the function takes).
 191 class PrototypeAST {
 192   std::string Name;
 193   std::vector<std::string> Args;
 194
 195 public:
 196   PrototypeAST(const std::string &Name, std::vector<std::string> Args)
 197       : Name(Name), Args(std::move(Args)) {}
 198   Function *codegen();
 199   const std::string &getName() const { return Name; }
 200 };
 201
 202 /// FunctionAST - This class represents a function definition itself.
 203 class FunctionAST {
 204   std::unique_ptr<PrototypeAST> Proto;
 205   std::unique_ptr<ExprAST> Body;
 206
 207 public:
 208   FunctionAST(std::unique_ptr<PrototypeAST> Proto,
 209               std::unique_ptr<ExprAST> Body)
 210       : Proto(std::move(Proto)), Body(std::move(Body)) {}
 211   Function *codegen();
 212 };
 213 } // end anonymous namespace
 214
 215 //===----------------------------------------------------------------------===//
 216 // Parser
 217 //===----------------------------------------------------------------------===//
 218
 219 /// CurTok/getNextToken - Provide a simple token buffer.  CurTok is the current
 220 /// token the parser is looking at.  getNextToken reads another token from the
 221 /// lexer and updates CurTok with its results.
 222 static int CurTok;
 223 static int getNextToken() { return CurTok = gettok(); }
 224
 225 /// BinopPrecedence - This holds the precedence for each binary operator that is
 226 /// defined.
 227 static std::map<char, int> BinopPrecedence;
 228
 229 /// GetTokPrecedence - Get the precedence of the pending binary operator token.
 230 static int GetTokPrecedence() {
 231   if (!isascii(CurTok))
 232     return -1;
 233
 234   // Make sure it's a declared binop.
 235   int TokPrec = BinopPrecedence[CurTok];
 236   if (TokPrec <= 0)
 237     return -1;
 238   return TokPrec;
 239 }
 240
 241 /// Error* - These are little helper functions for error handling.
 242 std::unique_ptr<ExprAST> Error(const char *Str) {
 243   fprintf(stderr, "Error: %s\n", Str);
 244   return nullptr;
 245 }
 246 std::unique_ptr<PrototypeAST> ErrorP(const char *Str) {
 247   Error(Str);
 248   return nullptr;
 249 }
 250
 251 static std::unique_ptr<ExprAST> ParseExpression();
 252
 253 /// numberexpr ::= number
 254 static std::unique_ptr<ExprAST> ParseNumberExpr() {
 255   auto Result = llvm::make_unique<NumberExprAST>(NumVal);
 256   getNextToken(); // consume the number
 257   return std::move(Result);
 258 }
 259
 260 /// parenexpr ::= '(' expression ')'
 261 static std::unique_ptr<ExprAST> ParseParenExpr() {
 262   getNextToken(); // eat (.
 263   auto V = ParseExpression();
 264   if (!V)
 265     return nullptr;
 266
 267   if (CurTok != ')')
 268     return Error("expected ')'");
 269   getNextToken(); // eat ).
 270   return V;
 271 }
 272
 273 /// identifierexpr
 274 ///   ::= identifier
 275 ///   ::= identifier '(' expression* ')'
 276 static std::unique_ptr<ExprAST> ParseIdentifierExpr() {
 277   std::string IdName = IdentifierStr;
 278
 279   getNextToken(); // eat identifier.
 280
 281   if (CurTok != '(') // Simple variable ref.
 282     return llvm::make_unique<VariableExprAST>(IdName);
 283
 284   // Call.
 285   getNextToken(); // eat (
 286   std::vector<std::unique_ptr<ExprAST>> Args;
 287   if (CurTok != ')') {
 288     while (1) {
 289       if (auto Arg = ParseExpression())
 290         Args.push_back(std::move(Arg));
 291       else
 292         return nullptr;
 293
 294       if (CurTok == ')')
 295         break;
 296
 297       if (CurTok != ',')
 298         return Error("Expected ')' or ',' in argument list");
 299       getNextToken();
 300     }
 301   }
 302
 303   // Eat the ')'.
 304   getNextToken();
 305
 306   return llvm::make_unique<CallExprAST>(IdName, std::move(Args));
 307 }
 308
 309 /// ifexpr ::= 'if' expression 'then' expression 'else' expression
 310 static std::unique_ptr<ExprAST> ParseIfExpr() {
 311   getNextToken(); // eat the if.
 312
 313   // condition.
 314   auto Cond = ParseExpression();
 315   if (!Cond)
 316     return nullptr;
 317
 318   if (CurTok != tok_then)
 319     return Error("expected then");
 320   getNextToken(); // eat the then
 321
 322   auto Then = ParseExpression();
 323   if (!Then)
 324     return nullptr;
 325
 326   if (CurTok != tok_else)
 327     return Error("expected else");
 328
 329   getNextToken();
 330
 331   auto Else = ParseExpression();
 332   if (!Else)
 333     return nullptr;
 334
 335   return llvm::make_unique<IfExprAST>(std::move(Cond), std::move(Then),
 336                                       std::move(Else));
 337 }
 338
 339 /// forexpr ::= 'for' identifier '=' expr ',' expr (',' expr)? 'in' expression
 340 static std::unique_ptr<ExprAST> ParseForExpr() {
 341   getNextToken(); // eat the for.
 342
 343   if (CurTok != tok_identifier)
 344     return Error("expected identifier after for");
 345
 346   std::string IdName = IdentifierStr;
 347   getNextToken(); // eat identifier.
 348
 349   if (CurTok != '=')
 350     return Error("expected '=' after for");
 351   getNextToken(); // eat '='.
 352
 353   auto Start = ParseExpression();
 354   if (!Start)
 355     return nullptr;
 356   if (CurTok != ',')
 357     return Error("expected ',' after for start value");
 358   getNextToken();
 359
 360   auto End = ParseExpression();
 361   if (!End)
 362     return nullptr;
 363
 364   // The step value is optional.
 365   std::unique_ptr<ExprAST> Step;
 366   if (CurTok == ',') {
 367     getNextToken();
 368     Step = ParseExpression();
 369     if (!Step)
 370       return nullptr;
 371   }
 372
 373   if (CurTok != tok_in)
 374     return Error("expected 'in' after for");
 375   getNextToken(); // eat 'in'.
 376
 377   auto Body = ParseExpression();
 378   if (!Body)
 379     return nullptr;
 380
 381   return llvm::make_unique<ForExprAST>(IdName, std::move(Start), std::move(End),
 382                                        std::move(Step), std::move(Body));
 383 }
 384
 385 /// primary
 386 ///   ::= identifierexpr
 387 ///   ::= numberexpr
 388 ///   ::= parenexpr
 389 ///   ::= ifexpr
 390 ///   ::= forexpr
 391 static std::unique_ptr<ExprAST> ParsePrimary() {
 392   switch (CurTok) {
 393   default:
 394     return Error("unknown token when expecting an expression");
 395   case tok_identifier:
 396     return ParseIdentifierExpr();
 397   case tok_number:
 398     return ParseNumberExpr();
 399   case '(':
 400     return ParseParenExpr();
 401   case tok_if:
 402     return ParseIfExpr();
 403   case tok_for:
 404     return ParseForExpr();
 405   }
 406 }
 407
 408 /// binoprhs
 409 ///   ::= ('+' primary)*
 410 static std::unique_ptr<ExprAST> ParseBinOpRHS(int ExprPrec,
 411                                               std::unique_ptr<ExprAST> LHS) {
 412   // If this is a binop, find its precedence.
 413   while (1) {
 414     int TokPrec = GetTokPrecedence();
 415
 416     // If this is a binop that binds at least as tightly as the current binop,
 417     // consume it, otherwise we are done.
 418     if (TokPrec < ExprPrec)
 419       return LHS;
 420
 421     // Okay, we know this is a binop.
 422     int BinOp = CurTok;
 423     getNextToken(); // eat binop
 424
 425     // Parse the primary expression after the binary operator.
 426     auto RHS = ParsePrimary();
 427     if (!RHS)
 428       return nullptr;
 429
 430     // If BinOp binds less tightly with RHS than the operator after RHS, let
 431     // the pending operator take RHS as its LHS.
 432     int NextPrec = GetTokPrecedence();
 433     if (TokPrec < NextPrec) {
 434       RHS = ParseBinOpRHS(TokPrec + 1, std::move(RHS));
 435       if (!RHS)
 436         return nullptr;
 437     }
 438
 439     // Merge LHS/RHS.
 440     LHS =
 441         llvm::make_unique<BinaryExprAST>(BinOp, std::move(LHS), std::move(RHS));
 442   }
 443 }
 444
 445 /// expression
 446 ///   ::= primary binoprhs
 447 ///
 448 static std::unique_ptr<ExprAST> ParseExpression() {
 449   auto LHS = ParsePrimary();
 450   if (!LHS)
 451     return nullptr;
 452
 453   return ParseBinOpRHS(0, std::move(LHS));
 454 }
 455
 456 /// prototype
 457 ///   ::= id '(' id* ')'
 458 static std::unique_ptr<PrototypeAST> ParsePrototype() {
 459   if (CurTok != tok_identifier)
 460     return ErrorP("Expected function name in prototype");
 461
 462   std::string FnName = IdentifierStr;
 463   getNextToken();
 464
 465   if (CurTok != '(')
 466     return ErrorP("Expected '(' in prototype");
 467
 468   std::vector<std::string> ArgNames;
 469   while (getNextToken() == tok_identifier)
 470     ArgNames.push_back(IdentifierStr);
 471   if (CurTok != ')')
 472     return ErrorP("Expected ')' in prototype");
 473
 474   // success.
 475   getNextToken(); // eat ')'.
 476
 477   return llvm::make_unique<PrototypeAST>(FnName, std::move(ArgNames));
 478 }
 479
 480 /// definition ::= 'def' prototype expression
 481 static std::unique_ptr<FunctionAST> ParseDefinition() {
 482   getNextToken(); // eat def.
 483   auto Proto = ParsePrototype();
 484   if (!Proto)
 485     return nullptr;
 486
 487   if (auto E = ParseExpression())
 488     return llvm::make_unique<FunctionAST>(std::move(Proto), std::move(E));
 489   return nullptr;
 490 }
 491
 492 /// toplevelexpr ::= expression
 493 static std::unique_ptr<FunctionAST> ParseTopLevelExpr() {
 494   if (auto E = ParseExpression()) {
 495     // Make an anonymous proto.
 496     auto Proto = llvm::make_unique<PrototypeAST>("__anon_expr",
 497                                                  std::vector<std::string>());
 498     return llvm::make_unique<FunctionAST>(std::move(Proto), std::move(E));
 499   }
 500   return nullptr;
 501 }
 502
 503 /// external ::= 'extern' prototype
 504 static std::unique_ptr<PrototypeAST> ParseExtern() {
 505   getNextToken(); // eat extern.
 506   return ParsePrototype();
 507 }
 508
 509 //===----------------------------------------------------------------------===//
 510 // Code Generation
 511 //===----------------------------------------------------------------------===//
 512
 513 static std::unique_ptr<Module> TheModule;
 514 static IRBuilder<> Builder(getGlobalContext());
 515 static std::map<std::string, Value *> NamedValues;
 516 static std::unique_ptr<legacy::FunctionPassManager> TheFPM;
 517 static std::unique_ptr<KaleidoscopeJIT> TheJIT;
 518 static std::map<std::string, std::unique_ptr<PrototypeAST>> FunctionProtos;
 519
 520 Value *ErrorV(const char *Str) {
 521   Error(Str);
 522   return nullptr;
 523 }
 524
 525 Function *getFunction(std::string Name) {
 526   // First, see if the function has already been added to the current module.
 527   if (auto *F = TheModule->getFunction(Name))
 528     return F;
 529
 530   // If not, check whether we can codegen the declaration from some existing
 531   // prototype.
 532   auto FI = FunctionProtos.find(Name);
 533   if (FI != FunctionProtos.end())
 534     return FI->second->codegen();
 535
 536   // If no existing prototype exists, return null.
 537   return nullptr;
 538 }
 539
 540 Value *NumberExprAST::codegen() {
 541   return ConstantFP::get(getGlobalContext(), APFloat(Val));
 542 }
 543
 544 Value *VariableExprAST::codegen() {
 545   // Look this variable up in the function.
 546   Value *V = NamedValues[Name];
 547   if (!V)
 548     return ErrorV("Unknown variable name");
 549   return V;
 550 }
 551
 552 Value *BinaryExprAST::codegen() {
 553   Value *L = LHS->codegen();
 554   Value *R = RHS->codegen();
 555   if (!L || !R)
 556     return nullptr;
 557
 558   switch (Op) {
 559   case '+':
 560     return Builder.CreateFAdd(L, R, "addtmp");
 561   case '-':
 562     return Builder.CreateFSub(L, R, "subtmp");
 563   case '*':
 564     return Builder.CreateFMul(L, R, "multmp");
 565   case '<':
 566     L = Builder.CreateFCmpULT(L, R, "cmptmp");
 567     // Convert bool 0/1 to double 0.0 or 1.0
 568     return Builder.CreateUIToFP(L, Type::getDoubleTy(getGlobalContext()),
 569                                 "booltmp");
 570   default:
 571     return ErrorV("invalid binary operator");
 572   }
 573 }
 574
 575 Value *CallExprAST::codegen() {
 576   // Look up the name in the global module table.
 577   Function *CalleeF = getFunction(Callee);
 578   if (!CalleeF)
 579     return ErrorV("Unknown function referenced");
 580
 581   // If argument mismatch error.
 582   if (CalleeF->arg_size() != Args.size())
 583     return ErrorV("Incorrect # arguments passed");
 584
 585   std::vector<Value *> ArgsV;
 586   for (unsigned i = 0, e = Args.size(); i != e; ++i) {
 587     ArgsV.push_back(Args[i]->codegen());
 588     if (!ArgsV.back())
 589       return nullptr;
 590   }
 591
 592   return Builder.CreateCall(CalleeF, ArgsV, "calltmp");
 593 }
 594
 595 Value *IfExprAST::codegen() {
 596   Value *CondV = Cond->codegen();
 597   if (!CondV)
 598     return nullptr;
 599
 600   // Convert condition to a bool by comparing equal to 0.0.
 601   CondV = Builder.CreateFCmpONE(
 602       CondV, ConstantFP::get(getGlobalContext(), APFloat(0.0)), "ifcond");
 603
 604   Function *TheFunction = Builder.GetInsertBlock()->getParent();
 605
 606   // Create blocks for the then and else cases.  Insert the 'then' block at the
 607   // end of the function.
 608   BasicBlock *ThenBB =
 609       BasicBlock::Create(getGlobalContext(), "then", TheFunction);
 610   BasicBlock *ElseBB = BasicBlock::Create(getGlobalContext(), "else");
 611   BasicBlock *MergeBB = BasicBlock::Create(getGlobalContext(), "ifcont");
 612
 613   Builder.CreateCondBr(CondV, ThenBB, ElseBB);
 614
 615   // Emit then value.
 616   Builder.SetInsertPoint(ThenBB);
 617
 618   Value *ThenV = Then->codegen();
 619   if (!ThenV)
 620     return nullptr;
 621
 622   Builder.CreateBr(MergeBB);
 623   // Codegen of 'Then' can change the current block, update ThenBB for the PHI.
 624   ThenBB = Builder.GetInsertBlock();
 625
 626   // Emit else block.
 627   TheFunction->getBasicBlockList().push_back(ElseBB);
 628   Builder.SetInsertPoint(ElseBB);
 629
 630   Value *ElseV = Else->codegen();
 631   if (!ElseV)
 632     return nullptr;
 633
 634   Builder.CreateBr(MergeBB);
 635   // Codegen of 'Else' can change the current block, update ElseBB for the PHI.
 636   ElseBB = Builder.GetInsertBlock();
 637
 638   // Emit merge block.
 639   TheFunction->getBasicBlockList().push_back(MergeBB);
 640   Builder.SetInsertPoint(MergeBB);
 641   PHINode *PN =
 642       Builder.CreatePHI(Type::getDoubleTy(getGlobalContext()), 2, "iftmp");
 643
 644   PN->addIncoming(ThenV, ThenBB);
 645   PN->addIncoming(ElseV, ElseBB);
 646   return PN;
 647 }
 648
 649 // Output for-loop as:
 650 //   ...
 651 //   start = startexpr
 652 //   goto loop
 653 // loop:
 654 //   variable = phi [start, loopheader], [nextvariable, loopend]
 655 //   ...
 656 //   bodyexpr
 657 //   ...
 658 // loopend:
 659 //   step = stepexpr
 660 //   nextvariable = variable + step
 661 //   endcond = endexpr
 662 //   br endcond, loop, endloop
 663 // outloop:
 664 Value *ForExprAST::codegen() {
 665   // Emit the start code first, without 'variable' in scope.
 666   Value *StartVal = Start->codegen();
 667   if (!StartVal)
 668     return nullptr;
 669
 670   // Make the new basic block for the loop header, inserting after current
 671   // block.
 672   Function *TheFunction = Builder.GetInsertBlock()->getParent();
 673   BasicBlock *PreheaderBB = Builder.GetInsertBlock();
 674   BasicBlock *LoopBB =
 675       BasicBlock::Create(getGlobalContext(), "loop", TheFunction);
 676
 677   // Insert an explicit fall through from the current block to the LoopBB.
 678   Builder.CreateBr(LoopBB);
 679
 680   // Start insertion in LoopBB.
 681   Builder.SetInsertPoint(LoopBB);
 682
 683   // Start the PHI node with an entry for Start.
 684   PHINode *Variable = Builder.CreatePHI(Type::getDoubleTy(getGlobalContext()),
 685                                         2, VarName.c_str());
 686   Variable->addIncoming(StartVal, PreheaderBB);
 687
 688   // Within the loop, the variable is defined equal to the PHI node.  If it
 689   // shadows an existing variable, we have to restore it, so save it now.
 690   Value *OldVal = NamedValues[VarName];
 691   NamedValues[VarName] = Variable;
 692
 693   // Emit the body of the loop.  This, like any other expr, can change the
 694   // current BB.  Note that we ignore the value computed by the body, but don't
 695   // allow an error.
 696   if (!Body->codegen())
 697     return nullptr;
 698
 699   // Emit the step value.
 700   Value *StepVal = nullptr;
 701   if (Step) {
 702     StepVal = Step->codegen();
 703     if (!StepVal)
 704       return nullptr;
 705   } else {
 706     // If not specified, use 1.0.
 707     StepVal = ConstantFP::get(getGlobalContext(), APFloat(1.0));
 708   }
 709
 710   Value *NextVar = Builder.CreateFAdd(Variable, StepVal, "nextvar");
 711
 712   // Compute the end condition.
 713   Value *EndCond = End->codegen();
 714   if (!EndCond)
 715     return nullptr;
 716
 717   // Convert condition to a bool by comparing equal to 0.0.
 718   EndCond = Builder.CreateFCmpONE(
 719       EndCond, ConstantFP::get(getGlobalContext(), APFloat(0.0)), "loopcond");
 720
 721   // Create the "after loop" block and insert it.
 722   BasicBlock *LoopEndBB = Builder.GetInsertBlock();
 723   BasicBlock *AfterBB =
 724       BasicBlock::Create(getGlobalContext(), "afterloop", TheFunction);
 725
 726   // Insert the conditional branch into the end of LoopEndBB.
 727   Builder.CreateCondBr(EndCond, LoopBB, AfterBB);
 728
 729   // Any new code will be inserted in AfterBB.
 730   Builder.SetInsertPoint(AfterBB);
 731
 732   // Add a new entry to the PHI node for the backedge.
 733   Variable->addIncoming(NextVar, LoopEndBB);
 734
 735   // Restore the unshadowed variable.
 736   if (OldVal)
 737     NamedValues[VarName] = OldVal;
 738   else
 739     NamedValues.erase(VarName);
 740
 741   // for expr always returns 0.0.
 742   return Constant::getNullValue(Type::getDoubleTy(getGlobalContext()));
 743 }
 744
 745 Function *PrototypeAST::codegen() {
 746   // Make the function type:  double(double,double) etc.
 747   std::vector<Type *> Doubles(Args.size(),
 748                               Type::getDoubleTy(getGlobalContext()));
 749   FunctionType *FT =
 750       FunctionType::get(Type::getDoubleTy(getGlobalContext()), Doubles, false);
 751
 752   Function *F =
 753       Function::Create(FT, Function::ExternalLinkage, Name, TheModule.get());
 754
 755   // Set names for all arguments.
 756   unsigned Idx = 0;
 757   for (auto &Arg : F->args())
 758     Arg.setName(Args[Idx++]);
 759
 760   return F;
 761 }
 762
 763 Function *FunctionAST::codegen() {
 764   // Transfer ownership of the prototype to the FunctionProtos map, but keep a
 765   // reference to it for use below.
 766   auto &P = *Proto;
 767   FunctionProtos[Proto->getName()] = std::move(Proto);
 768   Function *TheFunction = getFunction(P.getName());
 769   if (!TheFunction)
 770     return nullptr;
 771
 772   // Create a new basic block to start insertion into.
 773   BasicBlock *BB = BasicBlock::Create(getGlobalContext(), "entry", TheFunction);
 774   Builder.SetInsertPoint(BB);
 775
 776   // Record the function arguments in the NamedValues map.
 777   NamedValues.clear();
 778   for (auto &Arg : TheFunction->args())
 779     NamedValues[Arg.getName()] = &Arg;
 780
 781   if (Value *RetVal = Body->codegen()) {
 782     // Finish off the function.
 783     Builder.CreateRet(RetVal);
 784
 785     // Validate the generated code, checking for consistency.
 786     verifyFunction(*TheFunction);
 787
 788     // Run the optimizer on the function.
 789     TheFPM->run(*TheFunction);
 790
 791     return TheFunction;
 792   }
 793
 794   // Error reading body, remove function.
 795   TheFunction->eraseFromParent();
 796   return nullptr;
 797 }
 798
 799 //===----------------------------------------------------------------------===//
 800 // Top-Level parsing and JIT Driver
 801 //===----------------------------------------------------------------------===//
 802
 803 static void InitializeModuleAndPassManager() {
 804   // Open a new module.
 805   TheModule = llvm::make_unique<Module>("my cool jit", getGlobalContext());
 806   TheModule->setDataLayout(TheJIT->getTargetMachine().createDataLayout());
 807
 808   // Create a new pass manager attached to it.
 809   TheFPM = llvm::make_unique<legacy::FunctionPassManager>(TheModule.get());
 810
 811   // Provide basic AliasAnalysis support for GVN.
 812   TheFPM->add(createBasicAliasAnalysisPass());
 813   // Do simple "peephole" optimizations and bit-twiddling optzns.
 814   TheFPM->add(createInstructionCombiningPass());
 815   // Reassociate expressions.
 816   TheFPM->add(createReassociatePass());
 817   // Eliminate Common SubExpressions.
 818   TheFPM->add(createGVNPass());
 819   // Simplify the control flow graph (deleting unreachable blocks, etc).
 820   TheFPM->add(createCFGSimplificationPass());
 821
 822   TheFPM->doInitialization();
 823 }
 824
 825 static void HandleDefinition() {
 826   if (auto FnAST = ParseDefinition()) {
 827     if (auto *FnIR = FnAST->codegen()) {
 828       fprintf(stderr, "Read function definition:");
 829       FnIR->dump();
 830       TheJIT->addModule(std::move(TheModule));
 831       InitializeModuleAndPassManager();
 832     }
 833   } else {
 834     // Skip token for error recovery.
 835     getNextToken();
 836   }
 837 }
 838
 839 static void HandleExtern() {
 840   if (auto ProtoAST = ParseExtern()) {
 841     if (auto *FnIR = ProtoAST->codegen()) {
 842       fprintf(stderr, "Read extern: ");
 843       FnIR->dump();
 844       FunctionProtos[ProtoAST->getName()] = std::move(ProtoAST);
 845     }
 846   } else {
 847     // Skip token for error recovery.
 848     getNextToken();
 849   }
 850 }
 851
 852 static void HandleTopLevelExpression() {
 853   // Evaluate a top-level expression into an anonymous function.
 854   if (auto FnAST = ParseTopLevelExpr()) {
 855     if (FnAST->codegen()) {
 856
 857       // JIT the module containing the anonymous expression, keeping a handle so
 858       // we can free it later.
 859       auto H = TheJIT->addModule(std::move(TheModule));
 860       InitializeModuleAndPassManager();
 861
 862       // Search the JIT for the __anon_expr symbol.
 863       auto ExprSymbol = TheJIT->findSymbol("__anon_expr");
 864       assert(ExprSymbol && "Function not found");
 865
 866       // Get the symbol's address and cast it to the right type (takes no
 867       // arguments, returns a double) so we can call it as a native function.
 868       double (*FP)() = (double (*)())(intptr_t)ExprSymbol.getAddress();
 869       fprintf(stderr, "Evaluated to %f\n", FP());
 870
 871       // Delete the anonymous expression module from the JIT.
 872       TheJIT->removeModule(H);
 873     }
 874   } else {
 875     // Skip token for error recovery.
 876     getNextToken();
 877   }
 878 }
 879
 880 /// top ::= definition | external | expression | ';'
 881 static void MainLoop() {
 882   while (1) {
 883     fprintf(stderr, "ready> ");
 884     switch (CurTok) {
 885     case tok_eof:
 886       return;
 887     case ';': // ignore top-level semicolons.
 888       getNextToken();
 889       break;
 890     case tok_def:
 891       HandleDefinition();
 892       break;
 893     case tok_extern:
 894       HandleExtern();
 895       break;
 896     default:
 897       HandleTopLevelExpression();
 898       break;
 899     }
 900   }
 901 }
 902
 903 //===----------------------------------------------------------------------===//
 904 // "Library" functions that can be "extern'd" from user code.
 905 //===----------------------------------------------------------------------===//
 906
 907 /// putchard - putchar that takes a double and returns 0.
 908 extern "C" double putchard(double X) {
 909   fputc((char)X, stderr);
 910   return 0;
 911 }
 912
 913 /// printd - printf that takes a double prints it as "%f\n", returning 0.
 914 extern "C" double printd(double X) {
 915   fprintf(stderr, "%f\n", X);
 916   return 0;
 917 }
 918
 919 //===----------------------------------------------------------------------===//
 920 // Main driver code.
 921 //===----------------------------------------------------------------------===//
 922
 923 int main() {
 924   InitializeNativeTarget();
 925   InitializeNativeTargetAsmPrinter();
 926   InitializeNativeTargetAsmParser();
 927
 928   // Install standard binary operators.
 929   // 1 is lowest precedence.
 930   BinopPrecedence['<'] = 10;
 931   BinopPrecedence['+'] = 20;
 932   BinopPrecedence['-'] = 20;
 933   BinopPrecedence['*'] = 40; // highest.
 934
 935   // Prime the first token.
 936   fprintf(stderr, "ready> ");
 937   getNextToken();
 938
 939   TheJIT = llvm::make_unique<KaleidoscopeJIT>();
 940
 941   InitializeModuleAndPassManager();
 942
 943   // Run the main "interpreter loop" now.
 944   MainLoop();
 945
 946   return 0;
 947 }