4 files changed, 9 insertions, 5 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];
+    }
 };
 }