llvm/examples/Kaleidoscope/Chapter6/toy.cpp

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