examples/Kaleidoscope/BuildingAJIT/Chapter3/toy.cpp

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