4 files changed, 16 insertions, 15 deletions
diff --git a/lib/CodeGen/CodeGen.cc b/lib/CodeGen/CodeGen.cc
index 49b470c..ce5fa00 100644
--- a/lib/CodeGen/CodeGen.cc
+++ b/lib/CodeGen/CodeGen.cc
@@ -149,9 +149,12 @@ void CodeGen::selectUnit(size_t Unit) {
     CurrentUnit = Unit;
     DiagnosticSource = (*Units)[Unit].Path;
     FuncGen->setCurrentUnit(Unit, (*Units)[Unit].Path);
-    RuntimeChecker->setLibraryMode((*Units)[Unit].IsLibrary);
-    FuncGen->setCallHook([this](const FuncSig &Sig) {
-        if (!(*Units)[CurrentUnit].IsLibrary && (*Units)[Sig.OwnerUnit].IsLibrary)
+    RuntimeChecker->setLibraryLookup([this](int Line) {
+        return (*Units)[CurrentUnit].isLibrary(Line);
+    });
+    FuncGen->setCallHook([this](const FuncSig &) {
+        // The callee may enable __LIB inside its body even if its header is ordinary code.
+        if (!(*Units)[CurrentUnit].isLibrary(DiagnosticLine))
             RuntimeChecker->setUserLine(DiagnosticLine);
     });
 }
@@ -195,7 +198,8 @@ void CodeGen::compile(const std::vector<SourceUnit> &Sources) {
         if (!Added && !Compatible && (It->second.Unit != CurrentUnit ||
             ((IsTemplate || It->second.IsTemplate) && !(IsFunction && It->second.IsFunction)))) {
             // Default library units follow user inputs; point at the user's conflicting definition.
-            bool UserConflict = Sources[CurrentUnit].IsLibrary && !Sources[It->second.Unit].IsLibrary;
+            bool UserConflict = Sources[CurrentUnit].isLibrary(DiagnosticLine) &&
+                                !Sources[It->second.Unit].isLibrary(It->second.Line);
             SaveAndRestore<int> Location(DiagnosticLine, UserConflict ? It->second.Line : DiagnosticLine);
             SaveAndRestore<std::string> Source(DiagnosticSource,
                 UserConflict ? Sources[It->second.Unit].Path : DiagnosticSource);
diff --git a/lib/CodeGen/ConstructorGen.cc b/lib/CodeGen/ConstructorGen.cc
index b029385..da89f59 100644
--- a/lib/CodeGen/ConstructorGen.cc
+++ b/lib/CodeGen/ConstructorGen.cc
@@ -284,7 +284,7 @@ size_t CodeGen::registerStackObject(Value *Addr, const std::string &TypeName, in
         Count = Entry.CreateAlloca(Entry.getInt64Ty(), nullptr, "object_count");
         Entry.CreateStore(Entry.getInt64(0), Count);
     }
-    StackObjects.push_back({Addr, TypeName, Alive, Count, Line});
+    StackObjects.push_back({Addr, TypeName, Alive, Count, Line, CurrentUnit});
     // This DECLARE can execute repeatedly, including after a backward loop edge.
     emitStackObjectCleanup(StackObjects.back());
     return StackObjects.size() - 1;
@@ -298,6 +298,8 @@ void CodeGen::activateStackObject(size_t Index, Value *Count) {
 }
 
 void CodeGen::emitStackObjectCleanup(const StackObject &Object) {
+    size_t SavedUnit = CurrentUnit;
+    selectUnit(Object.OwnerUnit);
     auto *F = Builder->GetInsertBlock()->getParent();
     auto *Destroy = BasicBlock::Create(*TheContext, "destroy_object", F);
     auto *Done = BasicBlock::Create(*TheContext, "destroyed_object", F);
@@ -339,6 +341,7 @@ void CodeGen::emitStackObjectCleanup(const StackObject &Object) {
     }
     Builder->CreateBr(Done);
     Builder->SetInsertPoint(Done);
+    selectUnit(SavedUnit);
 }
 
 void CodeGen::finishStackObjects(Function *F) {
diff --git a/lib/CodeGen/RuntimeCheck.cc b/lib/CodeGen/RuntimeCheck.cc
index e15c10b..3991bf2 100644
--- a/lib/CodeGen/RuntimeCheck.cc
+++ b/lib/CodeGen/RuntimeCheck.cc
@@ -17,8 +17,9 @@ RuntimeCheck::RuntimeCheck(Module &M, LLVMContext &C, IRBuilder<> &B, ExceptionH
 }
 
 Value *RuntimeCheck::lineValue(int Line) {
-    return LibraryMode ? Builder.CreateLoad(Builder.getInt32Ty(), UserLine, "user_line")
-                       : static_cast<Value *>(Builder.getInt32(Line));
+    return LibraryLookup && LibraryLookup(Line)
+        ? Builder.CreateLoad(Builder.getInt32Ty(), UserLine, "user_line")
+        : static_cast<Value *>(Builder.getInt32(Line));
 }
 
 void RuntimeCheck::setUserLine(int Line) {
diff --git a/lib/Lexer/Lexer.cc b/lib/Lexer/Lexer.cc
index 5bf2230..40bfce3 100644
--- a/lib/Lexer/Lexer.cc
+++ b/lib/Lexer/Lexer.cc
@@ -220,17 +220,10 @@ int Lexer::readToken() {
 
     TokenLine = CurrentLine;
     if (LastChar == '#' && LineStart) {
-        std::string Directive;
         do {
-            Directive += static_cast<char>(LastChar);
             LastChar = Input.get();
         } while (LastChar != EOF && LastChar != '\n' && LastChar != '\r');
-        while (!Directive.empty() && isspace(static_cast<unsigned char>(Directive.back())))
-            Directive.pop_back();
-        if (Directive == "#LIB") Library = true;
-        else {
-            reportError(TokenLine, "Unknown compiler directive");
-        }
+        reportError(TokenLine, "Unknown compiler directive");
         return readToken();
     }
     LineStart = false;