21 files changed, 115 insertions, 25 deletions
diff --git a/include/cps/CodeGen.h b/include/cps/CodeGen.h
index 0962753..c3a7b19 100644
--- a/include/cps/CodeGen.h
+++ b/include/cps/CodeGen.h
@@ -106,6 +106,7 @@ class CodeGen {
         llvm::AllocaInst *Alive;
         llvm::AllocaInst *Count; // Non-null for an array; Addr then holds its buffer pointer.
         int Line;
+        size_t OwnerUnit;
         bool Temporary = false;
     };
     std::vector<StackObject> StackObjects;
diff --git a/include/cps/Lexer.h b/include/cps/Lexer.h
index 05fb857..3d057a3 100644
--- a/include/cps/Lexer.h
+++ b/include/cps/Lexer.h
@@ -128,7 +128,6 @@ class Lexer {
     int CurrentLine = 1;
     int TokenLine = 1;
     bool HadError = false;
-    bool Library = false;
     bool LineStart = true;
     struct TokenData {
         int Kind, Line;
@@ -158,7 +157,6 @@ public:
     int peekToken(unsigned Offset = 0);
     int peekLine(unsigned Offset = 0);
     int getLine() const { return TokenLine; }
-    bool isLibrary() const { return Library; }
     bool hadError() const { return HadError; }
     // Human-readable spelling of Tok for diagnostics; Lex supplies identifier text.
     static std::string tokenSpelling(int Tok, const Lexer &Lex);
diff --git a/include/cps/RuntimeCheck.h b/include/cps/RuntimeCheck.h
index d2ee1ea..71943d4 100644
--- a/include/cps/RuntimeCheck.h
+++ b/include/cps/RuntimeCheck.h
@@ -3,6 +3,8 @@
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include <cstdint>
+#include <functional>
+#include <utility>
 #include <string>
 
 namespace cps {
@@ -18,7 +20,7 @@ class RuntimeCheck {
     llvm::FunctionCallee StderrFunc;   // Windows: __acrt_iob_func(2)
     llvm::GlobalVariable *StderrVar;   // elsewhere: the C library's stderr
     ExceptionHandler &Exceptions;
-    bool LibraryMode = false;
+    std::function<bool(int)> LibraryLookup;
     llvm::GlobalVariable *UserLine;
     llvm::Value *DivZeroMsg;
     llvm::Value *OutOfBoundsMsg;
@@ -33,7 +35,7 @@ public:
     RuntimeCheck(llvm::Module &M, llvm::LLVMContext &C, llvm::IRBuilder<> &B,
                  ExceptionHandler &EH);
 
-    void setLibraryMode(bool Enabled) { LibraryMode = Enabled; }
+    void setLibraryLookup(std::function<bool(int)> Lookup) { LibraryLookup = std::move(Lookup); }
     llvm::Value *lineValue(int Line);
     void setUserLine(int Line);
 
diff --git a/include/cps/SourceUnit.h b/include/cps/SourceUnit.h
index 34266ac..153d020 100644
--- a/include/cps/SourceUnit.h
+++ b/include/cps/SourceUnit.h
@@ -9,6 +9,9 @@ struct SourceUnit {
     std::string Path;
     const std::vector<std::unique_ptr<StmtAST>> &Statements;
     std::vector<bool> Visible;
-    bool IsLibrary = false;
+    std::vector<bool> LibraryLines;
+    bool isLibrary(int Line) const {
+        return Line > 0 && static_cast<size_t>(Line) <= LibraryLines.size() && LibraryLines[Line - 1];
+    }
 };
 }
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;
diff --git a/pseudocode-lib/memory.cps b/pseudocode-lib/memory.cps
index 26287c3..b02ca2a 100644
--- a/pseudocode-lib/memory.cps
+++ b/pseudocode-lib/memory.cps
@@ -1,4 +1,4 @@
-#LIB
+PP_DEFINE __LIB
 INCLUDE_OBJ "memory.o"
 TYPE _MEMORY_ADDRESS = C_TYPE("void*")
 TYPE _MEMORY_SIZE = C_TYPE("unsigned long long")
diff --git a/pseudocode-lib/std.cps b/pseudocode-lib/std.cps
index 6068bd0..6aa411b 100644
--- a/pseudocode-lib/std.cps
+++ b/pseudocode-lib/std.cps
@@ -1,3 +1,3 @@
-#LIB
+PP_DEFINE __LIB
 INCLUDE "string.cps"
 INCLUDE "memory.cps"
diff --git a/pseudocode-lib/string.cps b/pseudocode-lib/string.cps
index cb06758..135f5d7 100644
--- a/pseudocode-lib/string.cps
+++ b/pseudocode-lib/string.cps
@@ -1,4 +1,4 @@
-#LIB
+PP_DEFINE __LIB
 INCLUDE_OBJ "string.o"
 TYPE _STRING_HANDLE = C_TYPE("void*")
 TYPE _STRING_C_INTEGER = C_TYPE("long long")
diff --git a/tests/i/library-lines/library.cps b/tests/i/library-lines/library.cps
index a35a2e9..d54f948 100644
--- a/tests/i/library-lines/library.cps
+++ b/tests/i/library-lines/library.cps
@@ -1,4 +1,4 @@
-#LIB
+PP_DEFINE __LIB
 PROCEDURE Check()
     CALL Inner()
 ENDPROCEDURE
diff --git a/tests/i/library-regions/case.json b/tests/i/library-regions/case.json
new file mode 100644
index 0000000..9f95525
--- /dev/null
+++ b/tests/i/library-regions/case.json
@@ -0,0 +1,6 @@
+{
+  "sources": [
+    "program.cps",
+    "regions.cps"
+  ]
+}
diff --git a/tests/i/library-regions/expected-stderr.txt b/tests/i/library-regions/expected-stderr.txt
new file mode 100644
index 0000000..89aee0f
--- /dev/null
+++ b/tests/i/library-regions/expected-stderr.txt
@@ -0,0 +1,6 @@
+[Fatal] line 2: Division by zero
+[Fatal] line 10: Division by zero
+[Fatal] line 15: Division by zero
+[Fatal] line 23: Division by zero
+[Fatal] line 4: Division by zero
+[Fatal] line 9: Division by zero
diff --git a/tests/i/library-regions/expected.txt b/tests/i/library-regions/expected.txt
new file mode 100644
index 0000000..b414108
--- /dev/null
+++ b/tests/i/library-regions/expected.txt
@@ -0,0 +1,6 @@
+1
+2
+3
+4
+5
+6
diff --git a/tests/i/library-regions/notes.md b/tests/i/library-regions/notes.md
new file mode 100644
index 0000000..dbe4faf
--- /dev/null
+++ b/tests/i/library-regions/notes.md
@@ -0,0 +1,5 @@
+# Trace notes
+
+Review status: pending human review.
+
+Mixed has an ordinary header but enables __LIB inside its body. Its first division uses the caller's line 2. After UNDEF, its own division uses regions.cps line 10 and a nested library call uses line 15. A disabled DEFINE does not enable __LIB, so the next division uses line 23. The library template retains its definition's mode during deferred emission and reports its caller's line 4. DeleteOnExit destroys its local object on return; the automatic library destructor retains the ordinary caller's line 9. Every error is caught and prints the next integer from 1 through 6. These are compiler extensions.
diff --git a/tests/i/library-regions/program.cps b/tests/i/library-regions/program.cps
new file mode 100644
index 0000000..315dbe6
--- /dev/null
+++ b/tests/i/library-regions/program.cps
@@ -0,0 +1,12 @@
+INCLUDE "regions.cps"
+CALL Mixed()
+TRY
+    CALL TemplateError<INTEGER>()
+CATCH
+    OUTPUT 5
+ENDTRY
+TRY
+    CALL DeleteOnExit()
+CATCH
+    OUTPUT 6
+ENDTRY
diff --git a/tests/i/library-regions/regions.cps b/tests/i/library-regions/regions.cps
new file mode 100644
index 0000000..310014a
--- /dev/null
+++ b/tests/i/library-regions/regions.cps
@@ -0,0 +1,47 @@
+PROCEDURE Mixed()
+PP_DEFINE __LIB
+TRY
+    OUTPUT 1 DIV 0
+CATCH
+    OUTPUT 1
+ENDTRY
+PP_UNDEF __LIB
+TRY
+    OUTPUT 1 DIV 0
+CATCH
+    OUTPUT 2
+ENDTRY
+TRY
+    CALL LibraryError()
+CATCH
+    OUTPUT 3
+ENDTRY
+PP_IFDEF Never
+PP_DEFINE __LIB
+PP_ENDIF
+TRY
+    OUTPUT 1 DIV 0
+CATCH
+    OUTPUT 4
+ENDTRY
+ENDPROCEDURE
+PP_DEFINE __LIB
+PROCEDURE LibraryError()
+    CALL NestedError()
+ENDPROCEDURE
+PROCEDURE NestedError()
+    OUTPUT 1 DIV 0
+ENDPROCEDURE
+TEMPLATE <T>
+PROCEDURE TemplateError()
+    OUTPUT 1 DIV 0
+ENDPROCEDURE
+CLASS DeleteError
+    PROCEDURE DELETE()
+        OUTPUT 1 DIV 0
+    ENDPROCEDURE
+ENDCLASS
+PROCEDURE DeleteOnExit()
+    DECLARE X : DeleteError
+ENDPROCEDURE
+PP_UNDEF __LIB
diff --git a/tests/iii/legacy-lib/case.json b/tests/iii/legacy-lib/case.json
new file mode 100644
index 0000000..d593f34
--- /dev/null
+++ b/tests/iii/legacy-lib/case.json
@@ -0,0 +1 @@
+{"errors":[{"reason":"Unknown compiler directive","line":1}]}
diff --git a/tests/iii/legacy-lib/program.cps b/tests/iii/legacy-lib/program.cps
new file mode 100644
index 0000000..8abe9d9
--- /dev/null
+++ b/tests/iii/legacy-lib/program.cps
@@ -0,0 +1,2 @@
+#LIB
+OUTPUT 1
diff --git a/tools/driver/main.cc b/tools/driver/main.cc
index e3d4c28..d0e7c07 100644
--- a/tools/driver/main.cc
+++ b/tools/driver/main.cc
@@ -200,7 +200,7 @@ bool parseUnits(const std::vector<std::string> &InputPaths, size_t UserInputCoun
         cps::Parser Parser(Lex);
         Programs[I] = Parser.Parse();
         if (Lex.hadError() || Parser.hadError()) return false;
-        Units.push_back({InputPath, Programs[I], std::vector<bool>(InputPaths.size(), false), Lex.isLibrary()});
+        Units.push_back({InputPath, Programs[I], std::vector<bool>(InputPaths.size(), false), std::move(Source.LibraryLines)});
         Units.back().Visible[I] = true;
         if (!NoStd && I < UserInputCount)
             for (size_t J = UserInputCount; J < InputPaths.size(); ++J)