//===-- Immediate.cpp - the swift immediate mode --------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2015 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See http://swift.org/LICENSE.txt for license information
// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//
//
// This is the implementation of the swift interpreter, which takes a
// TranslationUnit and JITs it.
//
//===----------------------------------------------------------------------===//

#include "Completion.h"
#include "Immediate.h"
#include "Frontend.h"
#include "swift/Subsystems.h"
#include "swift/IRGen/Options.h"
#include "swift/Parse/Lexer.h"
#include "swift/AST/ASTContext.h"
#include "swift/AST/Component.h"
#include "swift/AST/Decl.h"
#include "swift/AST/Diagnostics.h"
#include "swift/AST/Module.h"
#include "swift/AST/NameLookup.h"
#include "swift/AST/Stmt.h"
#include "swift/AST/Types.h"
#include "swift/Basic/DiagnosticConsumer.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/StringSwitch.h"
#include "llvm/ExecutionEngine/JIT.h"
#include "llvm/Support/ConvertUTF.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/Host.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Support/Process.h"
#include "llvm/Support/SaveAndRestore.h"
#include "llvm/Support/Signals.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/system_error.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/Transforms/IPO/PassManagerBuilder.h"
#include "llvm/Transforms/IPO.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Module.h"
#include "llvm/Linker.h"
#include "llvm/PassManager.h"

#include <cmath>
#include <histedit.h>
#include <dlfcn.h>

using namespace swift;

namespace {
template<size_t N>
class ConvertForWcharSize;

template<>
class ConvertForWcharSize<2> {
public:
  static ConversionResult ConvertFromUTF8(const char** sourceStart,
                                          const char* sourceEnd,
                                          wchar_t** targetStart,
                                          wchar_t* targetEnd,
                                          ConversionFlags flags) {
    return ConvertUTF8toUTF16(reinterpret_cast<const UTF8**>(sourceStart),
                              reinterpret_cast<const UTF8*>(sourceEnd),
                              reinterpret_cast<UTF16**>(targetStart),
                              reinterpret_cast<UTF16*>(targetEnd),
                              flags);
  }
  
  static ConversionResult ConvertToUTF8(const wchar_t** sourceStart,
                                        const wchar_t* sourceEnd,
                                        char** targetStart,
                                        char* targetEnd,
                                        ConversionFlags flags) {
    return ConvertUTF16toUTF8(reinterpret_cast<const UTF16**>(sourceStart),
                              reinterpret_cast<const UTF16*>(sourceEnd),
                              reinterpret_cast<UTF8**>(targetStart),
                              reinterpret_cast<UTF8*>(targetEnd),
                              flags);
  }
};

template<>
class ConvertForWcharSize<4> {
public:
  static ConversionResult ConvertFromUTF8(const char** sourceStart,
                                          const char* sourceEnd,
                                          wchar_t** targetStart,
                                          wchar_t* targetEnd,
                                          ConversionFlags flags) {
    return ConvertUTF8toUTF32(reinterpret_cast<const UTF8**>(sourceStart),
                              reinterpret_cast<const UTF8*>(sourceEnd),
                              reinterpret_cast<UTF32**>(targetStart),
                              reinterpret_cast<UTF32*>(targetEnd),
                              flags);
  }
  
  static ConversionResult ConvertToUTF8(const wchar_t** sourceStart,
                                        const wchar_t* sourceEnd,
                                        char** targetStart,
                                        char* targetEnd,
                                        ConversionFlags flags) {
    return ConvertUTF32toUTF8(reinterpret_cast<const UTF32**>(sourceStart),
                              reinterpret_cast<const UTF32*>(sourceEnd),
                              reinterpret_cast<UTF8**>(targetStart),
                              reinterpret_cast<UTF8*>(targetEnd),
                              flags);
  }
};

using Convert = ConvertForWcharSize<sizeof(wchar_t)>;
  
static void convertFromUTF8(llvm::StringRef utf8,
                            llvm::SmallVectorImpl<wchar_t> &out) {
  size_t reserve = out.size() + utf8.size();
  out.reserve(reserve);
  const char *utf8_begin = utf8.begin();
  wchar_t *wide_begin = out.end();
  auto res = Convert::ConvertFromUTF8(&utf8_begin, utf8.end(),
                                      &wide_begin, out.data() + reserve,
                                      lenientConversion);
  assert(res == conversionOK && "utf8-to-wide conversion failed!");
  out.set_size(wide_begin - out.begin());
}
  
static void convertToUTF8(llvm::ArrayRef<wchar_t> wide,
                          llvm::SmallVectorImpl<char> &out) {
  size_t reserve = out.size() + wide.size()*4;
  out.reserve(reserve);
  const wchar_t *wide_begin = wide.begin();
  char *utf8_begin = out.end();
  auto res = Convert::ConvertToUTF8(&wide_begin, wide.end(),
                                    &utf8_begin, out.data() + reserve,
                                    lenientConversion);
  assert(res == conversionOK && "wide-to-utf8 conversion failed!");
  out.set_size(utf8_begin - out.begin());
}

} // end anonymous namespace

static void loadRuntimeLib(StringRef sharedLibName) {
  // FIXME: Need error-checking.
  llvm::sys::Path LibPath =
  llvm::sys::Path::GetMainExecutable(0, (void*)&swift::RunImmediately);
  LibPath.eraseComponent();
  LibPath.eraseComponent();
  LibPath.appendComponent("lib");
  LibPath.appendComponent(sharedLibName);
  dlopen(LibPath.c_str(), 0);
}

static void loadSwiftRuntime() {
  loadRuntimeLib("libswift_stdlib.dylib");
}

static bool IRGenImportedModules(TranslationUnit *TU,
                                 llvm::Module &Module,
                                 llvm::SmallPtrSet<TranslationUnit*, 8>
                                     &ImportedModules,
                                 SmallVectorImpl<llvm::Function*> &InitFns,
                                 irgen::Options &Options,
                                 bool IsREPL = true) {
  // IRGen the modules this module depends on.
  for (auto ModPair : TU->getImportedModules()) {
    if (isa<BuiltinModule>(ModPair.second) || isa<ClangModule>(ModPair.second))
      continue;

    TranslationUnit *SubTU = cast<TranslationUnit>(ModPair.second);
    if (!ImportedModules.insert(SubTU))
      continue;

    // For the moment, if we're in the REPL, don't bother to IRGen
    // swift.swift at all.
    // FIXME: Checking for "swift" explicitly is an ugly hack.
    if (SubTU->Name.str() == "swift")
      continue;

    // Recursively IRGen imported modules.
    IRGenImportedModules(SubTU, Module, ImportedModules, InitFns, Options);

    // FIXME: Need to check whether this is actually safe in general.
    llvm::Module SubModule(SubTU->Name.str(), Module.getContext());
    performCaptureAnalysis(SubTU);
    performIRGeneration(Options, &SubModule, SubTU);

    if (TU->Ctx.hadError())
      return true;

    std::string ErrorMessage;
    if (llvm::Linker::LinkModules(&Module, &SubModule,
                                  llvm::Linker::DestroySource,
                                  &ErrorMessage)) {
      llvm::errs() << "Error linking swift modules\n";
      llvm::errs() << ErrorMessage << "\n";
      return true;
    }

    // FIXME: This is an ugly hack; need to figure out how this should
    // actually work.
    SmallVector<char, 20> NameBuf;
    StringRef InitFnName = (SubTU->Name.str() + ".init").toStringRef(NameBuf);
    llvm::Function *InitFn = Module.getFunction(InitFnName);
    if (InitFn)
      InitFns.push_back(InitFn);

    // Load the shared library corresponding to this module.
    // FIXME: Swift and Clang modules alike need to record the dylibs against
    // which one needs to link when using the module. For now, just hardcode
    // the Swift libraries we care about.
    StringRef sharedLibName
      = llvm::StringSwitch<StringRef>(SubTU->Name.str())
          .Case("Foundation", "libswiftFoundation.dylib")
          .Case("ObjectiveC", "libswiftObjectiveC.dylib")
          .Case("AppKit",     "libswiftAppKit.dylib")
          .Default("");
    if (!sharedLibName.empty()) {
      loadRuntimeLib(sharedLibName);
    }
  }

  return false;
}

void swift::RunImmediately(TranslationUnit *TU, SILModule *SILMod) {
  ASTContext &Context = TU->Ctx;
  irgen::Options Options;
  Options.OutputFilename = "";
  Options.Triple = llvm::sys::getDefaultTargetTriple();
  Options.OptLevel = 2;
  Options.OutputKind = irgen::OutputKind::Module;
  Options.UseJIT = true;
  
  // IRGen the main module.
  llvm::LLVMContext LLVMContext;
  llvm::Module Module(TU->Name.str(), LLVMContext);
  performCaptureAnalysis(TU);
  performIRGeneration(Options, &Module, TU, SILMod);

  if (Context.hadError())
    return;

  SmallVector<llvm::Function*, 8> InitFns;
  llvm::SmallPtrSet<TranslationUnit*, 8> ImportedModules;
  if (IRGenImportedModules(TU, Module, ImportedModules, InitFns, Options,
                           /*IsREPL*/false))
    return;

  llvm::PassManagerBuilder PMBuilder;
  PMBuilder.OptLevel = 2;
  PMBuilder.Inliner = llvm::createFunctionInliningPass(200);
  llvm::PassManager ModulePasses;
  ModulePasses.add(new llvm::DataLayout(Module.getDataLayout()));
  PMBuilder.populateModulePassManager(ModulePasses);
  ModulePasses.run(Module);

  loadSwiftRuntime();

  // Build the ExecutionEngine.
  llvm::EngineBuilder builder(&Module);
  std::string ErrorMsg;
  llvm::TargetOptions TargetOpt;
  TargetOpt.NoFramePointerElimNonLeaf = true;
  builder.setTargetOptions(TargetOpt);
  builder.setErrorStr(&ErrorMsg);
  builder.setEngineKind(llvm::EngineKind::JIT);
  llvm::ExecutionEngine *EE = builder.create();
  if (!EE) {
    llvm::errs() << "Error loading JIT: " << ErrorMsg;
    return;
  }

  // Run the generated program.
  for (auto InitFn : InitFns)
    EE->runFunctionAsMain(InitFn, std::vector<std::string>(), 0);

  EE->runStaticConstructorsDestructors(false);
  llvm::Function *EntryFn = Module.getFunction("main");
  EE->runFunctionAsMain(EntryFn, std::vector<std::string>(), 0);
}

/// An arbitrary, otherwise-unused char value that editline interprets as
/// entering/leaving "literal mode", meaning it passes prompt characters through
/// to the terminal without affecting the line state. This prevents color
/// escape sequences from interfering with editline's internal state.
static constexpr wchar_t LITERAL_MODE_CHAR = L'\1';

/// Append a terminal escape sequence in "literal mode" so that editline
/// ignores it.
static void appendEscapeSequence(SmallVectorImpl<wchar_t> &dest,
                                 llvm::StringRef src)
{
  dest.push_back(LITERAL_MODE_CHAR);
  convertFromUTF8(src, dest);
  dest.push_back(LITERAL_MODE_CHAR);
}

/// RAII and Swift-specific setup wrapper for EditLine. All of its methods
/// must be usable from a separate thread and so shouldn't touch anything
/// outside of the EditLine, History, and member object state.
///
/// FIXME: Need the TU for completions--use a lock?
struct EditLineWrapper {
  TranslationUnit *TU;
  
  EditLine *e;
  HistoryW *h;
  size_t PromptContinuationLevel;
  bool NeedPromptContinuation;
  bool ShowColors;
  bool PromptedForLine;
  bool Outdented;
  Completions completions;
  
  llvm::SmallVector<wchar_t, 80> PromptString;

  EditLineWrapper(TranslationUnit *TU) : TU(TU) {
    // Only show colors if both stderr and stdout are displayed.
    ShowColors = llvm::errs().is_displayed() && llvm::outs().is_displayed();

    e = el_init("swift", stdin, stdout, stderr);
    h = history_winit();
    PromptContinuationLevel = 0;
    el_wset(e, EL_EDITOR, L"emacs");
    el_wset(e, EL_PROMPT_ESC, PromptFn, LITERAL_MODE_CHAR);
    el_wset(e, EL_CLIENTDATA, (void*)this);
    el_wset(e, EL_HIST, history, h);
    el_wset(e, EL_SIGNAL, 1);
    el_wset(e, EL_GETCFN, GetCharFn);
    
    // Provide special outdenting behavior for '}' and ':'.
    el_wset(e, EL_ADDFN, L"swift-close-brace", L"Reduce {} indentation level",
            BindingFn<&EditLineWrapper::onCloseBrace>);
    el_wset(e, EL_BIND, L"}", L"swift-close-brace", nullptr);
    
    el_wset(e, EL_ADDFN, L"swift-colon", L"Reduce label indentation level",
            BindingFn<&EditLineWrapper::onColon>);
    el_wset(e, EL_BIND, L":", L"swift-colon", nullptr);
    
    // Provide special indent/completion behavior for tab.
    el_wset(e, EL_ADDFN, L"swift-indent-or-complete",
           L"Indent line or trigger completion",
           BindingFn<&EditLineWrapper::onIndentOrComplete>);
    el_wset(e, EL_BIND, L"\t", L"swift-indent-or-complete", nullptr);

    el_wset(e, EL_ADDFN, L"swift-complete",
            L"Trigger completion",
            BindingFn<&EditLineWrapper::onComplete>);
    
    // Provide some common bindings to complement editline's defaults.
    // ^W should delete previous word, not the entire line.
    el_wset(e, EL_BIND, L"\x17", L"ed-delete-prev-word", nullptr);
    // ^_ should undo.
    el_wset(e, EL_BIND, L"\x1f", L"vi-undo", nullptr);
    
    HistEventW ev;
    history_w(h, &ev, H_SETSIZE, 800);
  }

  static wchar_t *PromptFn(EditLine *e) {
    void* clientdata;
    el_wget(e, EL_CLIENTDATA, &clientdata);
    return const_cast<wchar_t*>(((EditLineWrapper*)clientdata)->getPrompt());
  }
  
  const wchar_t *getPrompt() {
    PromptString.clear();

    if (ShowColors) {
      const char *colorCode =
        llvm::sys::Process::OutputColor(llvm::raw_ostream::YELLOW, false, false);
      if (colorCode)
        appendEscapeSequence(PromptString, colorCode);
    }
    
    static const wchar_t *PS1 = L"(swift) ";
    static const wchar_t *PS2 = L"        ";
    
    if (!NeedPromptContinuation)
      PromptString.insert(PromptString.end(), PS1, PS1 + wcslen(PS1));
    else {
      PromptString.insert(PromptString.end(), PS2, PS2 + wcslen(PS1));
      PromptString.append(2*PromptContinuationLevel, ' ');
    }
    
    if (ShowColors) {
      const char *colorCode = llvm::sys::Process::ResetColor();
      if (colorCode)
        appendEscapeSequence(PromptString, colorCode);
    }

    PromptedForLine = true;
    PromptString.push_back(L'\0');
    return PromptString.data();
  }

  /// Custom GETCFN to reset completion state after typing.
  /// NB: editline expects a char even in "wide" mode
  static int GetCharFn(EditLine *e, wchar_t *out) {
    void* clientdata;
    el_wget(e, EL_CLIENTDATA, &clientdata);
    EditLineWrapper *that = (EditLineWrapper*)clientdata;
    
    wint_t c;
    do {
      c = getwc(stdin);
      if (c == EOF) {
        if (feof(stdin)) {
          *out = '\0';
          return 0;
        }
        if (ferror(stdin)) {
          if (errno == EINTR)
            continue;
          *out = '\0';
          return -1;
        }
      }
    } while (false);
      
    // If the user typed anything other than tab, reset the completion state.
    if (c != L'\t')
      that->completions.reset();
    *out = wchar_t(c);
    return 1;
  }
  
  template<unsigned char (EditLineWrapper::*method)(int)>
  static unsigned char BindingFn(EditLine *e, int ch) {
    void *clientdata;
    el_wget(e, EL_CLIENTDATA, &clientdata);
    return (((EditLineWrapper*)clientdata)->*method)(ch);
  }
  
  bool isAtStartOfLine(LineInfoW const *line) {
    for (wchar_t c : llvm::makeArrayRef(line->buffer,
                                        line->cursor - line->buffer)) {
      if (!iswspace(c))
        return false;
    }
    return true;
  }

  // /^\s*\w+\s*:$/
  bool lineLooksLikeLabel(LineInfoW const *line) {
    wchar_t const *p = line->buffer;
    while (p != line->cursor && iswspace(*p))
      ++p;

    if (p == line->cursor)
      return false;
    
    do {
      ++p;
    } while (p != line->cursor && (iswalnum(*p) || *p == L'_'));
    
    while (p != line->cursor && iswspace(*p))
      ++p;

    return p+1 == line->cursor || *p == L':';
  }
  
  // /^\s*set\s*\(.*\)\s*:$/
  bool lineLooksLikeSetter(LineInfoW const *line) {
    wchar_t const *p = line->buffer;
    while (p != line->cursor && iswspace(*p))
      ++p;
    
    if (p == line->cursor || *p++ != L's')
      return false;
    if (p == line->cursor || *p++ != L'e')
      return false;
    if (p == line->cursor || *p++ != L't')
      return false;

    while (p != line->cursor && iswspace(*p))
      ++p;
    
    if (p == line->cursor || *p++ != L'(')
      return false;
    
    if (line->cursor - p < 2 || line->cursor[-1] != L':')
      return false;

    p = line->cursor - 1;
    while (iswspace(*--p));

    return *p == L')';
  }

  void outdent() {
    // If we didn't already outdent, do so.
    if (!Outdented) {
      if (PromptContinuationLevel > 0)
        --PromptContinuationLevel;
      Outdented = true;
    }
  }
  
  unsigned char onColon(int ch) {
    // Add the character to the string.
    wchar_t s[2] = {(wchar_t)ch, 0};
    el_winsertstr(e, s);
    
    LineInfoW const *line = el_wline(e);

    // Outdent if the line looks like a label.
    if (lineLooksLikeLabel(line))
      outdent();
    // Outdent if the line looks like a setter.
    if (lineLooksLikeSetter(line))
      outdent();
    
    return CC_REFRESH;
  }
  
  unsigned char onCloseBrace(int ch) {
    bool atStart = isAtStartOfLine(el_wline(e));
    
    // Add the character to the string.
    wchar_t s[2] = {(wchar_t)ch, 0};
    el_winsertstr(e, s);
    
    // Don't outdent if we weren't at the start of the line.
    if (!atStart) {
      return CC_REFRESH;
    }
    
    outdent();
    return CC_REFRESH;
  }
  
  unsigned char onIndentOrComplete(int ch) {
    LineInfoW const *line = el_wline(e);
    
    // FIXME: UTF-8? What's that?
    size_t cursorPos = line->cursor - line->buffer;
    
    // If there's nothing but whitespace before the cursor, indent to the next
    // 2-character tab stop.
    if (isAtStartOfLine(line)) {
      wchar_t const *indent = cursorPos & 1 ? L" " : L"  ";
      el_winsertstr(e, indent);
      return CC_REFRESH;
    }
    
    // Otherwise, look for completions.
    return onComplete(ch);
  }
  
  void insertStringRef(StringRef s) {
    if (s.empty())
      return;
    // Convert s to wchar_t* and null-terminate for el_winsertstr.
    SmallVector<wchar_t, 64> TmpStr;
    convertFromUTF8(s, TmpStr);
    TmpStr.push_back(L'\0');
    el_winsertstr(e, TmpStr.data());
  }
  
  void displayCompletions(llvm::ArrayRef<llvm::StringRef> list) {
    // FIXME: Do the print-completions-below-the-prompt thing bash does.
    llvm::outs() << '\n';
    // Trim the completion list to the terminal size.
    int lines_int = 0, columns_int = 0;
    // NB: EL_GETTC doesn't work with el_wget (?!)
    el_get(e, EL_GETTC, "li", &lines_int);
    el_get(e, EL_GETTC, "co", &columns_int);
    assert(lines_int > 0 && columns_int > 0 && "negative or zero screen size?!");
    
    auto lines = size_t(lines_int), columns = size_t(columns_int);
    size_t trimToColumns = columns > 2 ? columns - 2 : 0;
    
    size_t trimmed = 0;
    if (list.size() > lines - 1) {
      size_t trimToLines = lines > 2 ? lines - 2 : 0;
      trimmed = list.size() - trimToLines;
      list = list.slice(0, trimToLines);
    }
    
    for (StringRef completion : list) {
      if (completion.size() > trimToColumns)
        completion = completion.slice(0, trimToColumns);
      llvm::outs() << "  " << completion << '\n';
    }
    if (trimmed > 0)
      llvm::outs() << "  (and " << trimmed << " more)\n";
  }
  
  unsigned char onComplete(int ch) {
    LineInfoW const *line = el_wline(e);
    llvm::ArrayRef<wchar_t> wprefix(line->buffer, line->cursor - line->buffer);
    llvm::SmallString<16> prefix;
    convertToUTF8(wprefix, prefix);
    
    if (!completions) {
      // If we aren't currently working with a completion set, generate one.
      completions = Completions(TU, prefix);
      // Display the common root of the found completions and beep unless we
      // found a unique one.
      insertStringRef(completions.getRoot());
      return completions.isUnique()
        ? CC_REFRESH
        : CC_REFRESH_BEEP;
    }
    
    // Otherwise, advance through the completion state machine.
    switch (completions.getState()) {
    case CompletionState::CompletedRoot:
      // We completed the root. Next step is to display the completion list.
      displayCompletions(completions.getCompletionList());
      completions.setState(CompletionState::DisplayedCompletionList);
      return CC_REDISPLAY;

    case CompletionState::DisplayedCompletionList: {
      // Complete the next completion stem in the cycle.
      llvm::StringRef last = completions.getPreviousStem();
      el_wdeletestr(e, last.size());
      insertStringRef(completions.getNextStem());
      return CC_REFRESH;
    }
    
    case CompletionState::Empty:
    case CompletionState::Unique:
      // We already provided a definitive completion--nothing else to do.
      return CC_REFRESH_BEEP;

    case CompletionState::Invalid:
      llvm_unreachable("got an invalid completion set?!");
    }
  }

  ~EditLineWrapper() {
    el_end(e);
  }
  operator EditLine*() { return e; }
};

enum class PrintOrDump { Print, Dump };

static void printOrDumpDecl(Decl *d, PrintOrDump which) {
  if (which == PrintOrDump::Print) {
    d->print(llvm::outs());
    llvm::outs() << '\n';
  } else
    d->dump();
}

enum class REPLInputKind {
  /// The REPL got a "quit" signal.
  REPLQuit,
  /// Empty whitespace-only input.
  Empty,
  /// A REPL directive, such as ':help'.
  REPLDirective,
  /// Swift source code.
  SourceCode,
};

// NOTE: This code has to be able to run in its own thread. It should not touch
// anything other than the EditLineWrapper and Line objects.
static REPLInputKind getREPLInput(EditLineWrapper &e,
                                  llvm::SmallVectorImpl<char> &Line) {
  unsigned BraceCount = 0;
  bool HadLineContinuation = false;
  unsigned CurChunkLines = 0;

  Line.clear();
  
  do {
    // Read one line.
    e.PromptContinuationLevel = BraceCount;
    e.NeedPromptContinuation = BraceCount != 0 || HadLineContinuation;
    e.PromptedForLine = false;
    e.Outdented = false;
    int LineCount;
    size_t LineStart = Line.size();
    const wchar_t* WLine = el_wgets(e, &LineCount);
    if (!WLine) {
      // End-of-file.
      if (e.PromptedForLine)
        printf("\n");
      return REPLInputKind::REPLQuit;
    }
    
    size_t indent = e.PromptContinuationLevel*2;
    Line.append(indent, ' ');
    
    convertToUTF8(llvm::makeArrayRef(WLine, WLine + wcslen(WLine)), Line);

    // Special-case backslash for line continuations in the REPL.
    if (Line.size() > 2 && Line.end()[-1] == '\n' && Line.end()[-2] == '\\') {
      HadLineContinuation = true;
      Line.erase(Line.end() - 2);
    } else {
      HadLineContinuation = false;
    }

    // Enter the line into the line history.
    // FIXME: We should probably be a bit more clever here about which lines we
    // put into the history and when we put them in.
    HistEventW ev;
    history_w(e.h, &ev, H_ENTER, WLine);
    
    ++CurChunkLines;
    
    // If we detect a line starting with a colon, treat it as a special
    // REPL escape.
    char const *p = Line.data() + LineStart;
    while (p < Line.end() && isspace(*p)) {
      ++p;
    }
    if (p == Line.end())
      return REPLInputKind::Empty;
    
    if (CurChunkLines == 1 && BraceCount == 0 && *p == ':')
      return REPLInputKind::REPLDirective;
    
    // If we detect unbalanced braces, keep reading before
    // we start parsing.
    while (p < Line.end()) {
      if (*p == '{' || *p == '(' || *p == '[')
        ++BraceCount;
      else if (*p == '}' || *p == ')' || *p == ']')
        --BraceCount;
      ++p;
    }
  } while (BraceCount != 0 || HadLineContinuation);
  
  // The lexer likes null-terminated data.
  Line.push_back('\0');
  Line.pop_back();
  
  return REPLInputKind::SourceCode;
}

void swift::REPL(ASTContext &Context) {
  // FIXME: We should do something a bit more elaborate than
  // "allocate a 1MB buffer and hope it's enough".
  static const size_t BUFFER_SIZE = 1 << 20;
  llvm::MemoryBuffer *Buffer =
      llvm::MemoryBuffer::getNewMemBuffer(BUFFER_SIZE, "<REPL Buffer>");

  Component *Comp = new (Context.Allocate<Component>(1)) Component();
  unsigned BufferID =
    Context.SourceMgr.AddNewSourceBuffer(Buffer, llvm::SMLoc());
  Identifier ID = Context.getIdentifier("REPL");
  TranslationUnit *TU = new (Context) TranslationUnit(ID, Comp, Context,
                                                      /*IsMainModule=*/true,
                                                      /*IsReplModule=*/true);

  llvm::SmallPtrSet<TranslationUnit*, 8> ImportedModules;
  SmallVector<llvm::Function*, 8> InitFns;
  llvm::LLVMContext LLVMContext;
  llvm::Module Module("REPL", LLVMContext);
  llvm::Module DumpModule("REPL", LLVMContext);
  llvm::SmallString<128> DumpSource;

  loadSwiftRuntime();

  llvm::EngineBuilder builder(&Module);
  std::string ErrorMsg;
  llvm::TargetOptions TargetOpt;
  TargetOpt.NoFramePointerElimNonLeaf = true;
  builder.setTargetOptions(TargetOpt);
  builder.setErrorStr(&ErrorMsg);
  builder.setEngineKind(llvm::EngineKind::JIT);
  llvm::ExecutionEngine *EE = builder.create();

  irgen::Options Options;
  Options.OutputFilename = "";
  Options.Triple = llvm::sys::getDefaultTargetTriple();
  Options.OptLevel = 0;
  Options.OutputKind = irgen::OutputKind::Module;
  Options.UseJIT = true;
  
  EditLineWrapper e(TU);

  char* BufferStart = const_cast<char*>(Buffer->getBufferStart());
  
  unsigned CurTUElem = 0;
  unsigned CurIRGenElem = 0;

  // Force swift.swift to be parsed/type-checked immediately.  This forces
  // any errors to appear upfront, and helps eliminate some nasty lag after
  // the first statement is typed into the REPL.
  const char importstmt[] = "import swift\n";
  strcpy(BufferStart, importstmt);

  unsigned BufferOffset = 0;
  swift::appendToMainTranslationUnit(TU, BufferID, CurTUElem,
                                     /*startOffset*/ BufferOffset,
                                     /*endOffset*/ strlen(importstmt));
  if (Context.hadError())
    return;

  CurTUElem = CurIRGenElem = TU->Decls.size();

  if (llvm::sys::Process::StandardInIsUserInput())
    printf("%s", "Welcome to swift.  Type ':help' for assistance.\n");

  while (1) {
    // Get the next input from the REPL.
    llvm::SmallString<80> Line;
    REPLInputKind inputKind = getREPLInput(e, Line);

    Lexer L(Line, Context.SourceMgr, nullptr);
    switch (inputKind) {
    case REPLInputKind::REPLQuit:
      return;
        
    case REPLInputKind::Empty:
      break;
      
    case REPLInputKind::REPLDirective: {
      Token Tok;
      L.lex(Tok);
      assert(Tok.is(tok::colon));

      if (L.peekNextToken().getText() == "help") {
        printf("%s", "Available commands:\n"
               "  :quit - quit the interpreter (you can also use :exit"
               " or Control+D or exit(0))\n"
               "  :constraints (on|off) - turn on/off the constraint-"
               "based type checker\n"
               "  :constraints debug (on|off) - turn on/off the debug "
               "output for the constraint-based type checker\n"
               "  :dump_ir - dump the LLVM IR generated by the REPL\n"
               "  :dump_ast - dump the AST representation of"
               " the REPL input\n"
               "  :dump_decl <name> - dump the AST representation of the "
               "named declarations\n"
               "  :dump_source - dump the user input (ignoring"
               " lines with errors)\n"
               "  :print_decl <name> - print the AST representation of the "
               "named declarations\n"
               "API documentation etc. will be here eventually.\n");
      } else if (L.peekNextToken().getText() == "quit" ||
                 L.peekNextToken().getText() == "exit") {
        return;
      } else if (L.peekNextToken().getText() == "dump_ir") {
        DumpModule.dump();
      } else if (L.peekNextToken().getText() == "dump_ast") {
        TU->dump();
      } else if (L.peekNextToken().getText() == "dump_decl" ||
                 L.peekNextToken().getText() == "print_decl") {
        PrintOrDump doPrint = (L.peekNextToken().getText() == "print_decl")
          ? PrintOrDump::Print : PrintOrDump::Dump;
        L.lex(Tok);
        L.lex(Tok);
        UnqualifiedLookup lookup(Context.getIdentifier(Tok.getText()), TU);
        for (auto result : lookup.Results) {
          if (result.hasValueDecl()) {
            printOrDumpDecl(result.getValueDecl(), doPrint);
            
            if (auto typeDecl = dyn_cast<TypeDecl>(result.getValueDecl())) {
              if (auto typeAliasDecl = dyn_cast<TypeAliasDecl>(typeDecl)) {
                TypeDecl *origTypeDecl = typeAliasDecl->getUnderlyingType()
                  ->getNominalOrBoundGenericNominal();
                if (origTypeDecl) {
                  printOrDumpDecl(origTypeDecl, doPrint);
                  typeDecl = origTypeDecl;
                }
              }
              
              // FIXME: Hack!
              auto type = typeDecl->getDeclaredType();
              bool searchedClangModule = false;
              SmallVector<ExtensionDecl *, 4> extensions;
              for (auto ext : TU->lookupExtensions(type)) {
                extensions.push_back(ext);
              }
              
              llvm::SmallPtrSet<swift::Module *, 16> visited;
              for (auto &impEntry : TU->getImportedModules()) {
                if (!visited.insert(impEntry.second))
                  continue;
                
                // FIXME: Don't visit clang modules twice.
                if (isa<ClangModule>(impEntry.second)) {
                  if (searchedClangModule)
                    continue;
                  
                  searchedClangModule = true;
                }
                
                for (auto ext : impEntry.second->lookupExtensions(type)) {
                  extensions.push_back(ext);
                }
              }
              
              for (auto ext : extensions) {
                printOrDumpDecl(ext, doPrint);
              }
            }
          }
        }
      } else if (L.peekNextToken().getText() == "dump_source") {
        llvm::errs() << DumpSource;
      } else if (L.peekNextToken().getText() == "constraints") {
        L.lex(Tok);
        L.lex(Tok);
        if (Tok.getText() == "on") {
          TU->getASTContext().LangOpts.UseConstraintSolver = true;
        } else if (Tok.getText() == "off") {
          TU->getASTContext().LangOpts.UseConstraintSolver = false;
        } else if (Tok.getText() == "debug") {
          L.lex(Tok);
          if (Tok.getText() == "on") {
            TU->getASTContext().LangOpts.DebugConstraintSolver = true;
          } else if (Tok.getText() == "off") {
            TU->getASTContext().LangOpts.DebugConstraintSolver = false;
          } else {
            printf("%s", "Unknown :constraints debug command; try :help\n");
          }
        } else {
          printf("%s", "Unknown :constraints command; try :help\n");
        }
      } else {
        printf("%s", "Unknown interpreter escape; try :help\n");
      }
      break;
    }
      
    case REPLInputKind::SourceCode: {
      assert(Line.size() < BUFFER_SIZE &&
             "line too big for our stupid fixed-size repl buffer");
      strcpy(BufferStart, Line.c_str());
      BufferOffset = 0;
      // Parse the current line(s).
      bool ShouldRun =
        swift::appendToMainTranslationUnit(TU, BufferID, CurTUElem,
                                           BufferOffset, Line.size());

      if (Context.hadError()) {
        Context.Diags.resetHadAnyError();
        while (TU->Decls.size() > CurTUElem)
          TU->Decls.pop_back();
        TU->clearUnresolvedIdentifierTypes();

        // FIXME: Handling of "import" declarations?  Is there any other
        // state which needs to be reset?

        break;
      }

      CurTUElem = TU->Decls.size();
      
      DumpSource += Line;

      // If we didn't see an expression, statement, or decl which might have
      // side-effects, keep reading.
      if (!ShouldRun)
        continue;

      // IRGen the current line(s).
      llvm::Module LineModule("REPLLine", LLVMContext);
      performCaptureAnalysis(TU, CurIRGenElem);
      performIRGeneration(Options, &LineModule, TU, /*sil=*/nullptr,
                          CurIRGenElem);
      CurIRGenElem = CurTUElem;

      if (Context.hadError())
        return;

      std::string ErrorMessage;
      if (llvm::Linker::LinkModules(&Module, &LineModule,
                                    llvm::Linker::PreserveSource,
                                    &ErrorMessage)) {
        llvm::errs() << "Error linking swift modules\n";
        llvm::errs() << ErrorMessage << "\n";
        return;
      }
      if (llvm::Linker::LinkModules(&DumpModule, &LineModule,
                                    llvm::Linker::DestroySource,
                                    &ErrorMessage)) {
        llvm::errs() << "Error linking swift modules\n";
        llvm::errs() << ErrorMessage << "\n";
        return;
      }
      llvm::Function *DumpModuleMain = DumpModule.getFunction("main");
      DumpModuleMain->setName("repl.line");

      if (IRGenImportedModules(TU, Module, ImportedModules, InitFns, Options))
        return;

      for (auto InitFn : InitFns)
        EE->runFunctionAsMain(InitFn, std::vector<std::string>(), 0);
      InitFns.clear();

      // FIXME: The way we do this is really ugly... we should be able to
      // improve this.
      EE->runStaticConstructorsDestructors(&Module, false);
      llvm::Function *EntryFn = Module.getFunction("main");
      EE->runFunctionAsMain(EntryFn, std::vector<std::string>(), 0);
      EE->freeMachineCodeForFunction(EntryFn);
      EntryFn->eraseFromParent();

      break;
    }
    }
  }
}