1 files changed, 44 insertions, 41 deletions

diff --git a/src/libexpr/pos-table.hh b/src/libexpr/pos-table.hh
index 1decf3c85..0b60c4f6d 100644
--- a/src/libexpr/pos-table.hh
+++ b/src/libexpr/pos-table.hh
@@ -7,6 +7,7 @@
 #include "chunked-vector.hh"
 #include "pos-idx.hh"
 #include "position.hh"
+#include "sync.hh"
 
 namespace nix {
 
@@ -17,66 +18,68 @@ public:
     {
         friend PosTable;
     private:
-        // must always be invalid by default, add() replaces this with the actual value.
-        // subsequent add() calls use this index as a token to quickly check whether the
-        // current origins.back() can be reused or not.
-        mutable uint32_t idx = std::numeric_limits<uint32_t>::max();
-
-        // Used for searching in PosTable::[].
-        explicit Origin(uint32_t idx)
-            : idx(idx)
-            , origin{std::monostate()}
-        {
-        }
+        uint32_t offset;
+
+        Origin(Pos::Origin origin, uint32_t offset, size_t size):
+            offset(offset), origin(origin), size(size)
+        {}
 
     public:
         const Pos::Origin origin;
+        const size_t size;
 
-        Origin(Pos::Origin origin)
-            : origin(origin)
+        uint32_t offsetOf(PosIdx p) const
         {
+            return p.id - 1 - offset;
         }
     };
 
-    struct Offset
+private:
+    using Lines = std::vector<uint32_t>;
+
+    std::map<uint32_t, Origin> origins;
+    mutable Sync<std::map<uint32_t, Lines>> lines;
+
+    const Origin * resolve(PosIdx p) const
     {
-        uint32_t line, column;
-    };
+        if (p.id == 0)
+            return nullptr;
 
-private:
-    std::vector<Origin> origins;
-    ChunkedVector<Offset, 8192> offsets;
+        const auto idx = p.id - 1;
+        /* we want the last key <= idx, so we'll take prev(first key > idx).
+            this is guaranteed to never rewind origin.begin because the first
+            key is always 0. */
+        const auto pastOrigin = origins.upper_bound(idx);
+        return &std::prev(pastOrigin)->second;
+    }
 
 public:
-    PosTable()
-        : offsets(1024)
+    Origin addOrigin(Pos::Origin origin, size_t size)
     {
-        origins.reserve(1024);
+        uint32_t offset = 0;
+        if (auto it = origins.rbegin(); it != origins.rend())
+            offset = it->first + it->second.size;
+        // +1 because all PosIdx are offset by 1 to begin with (because noPos == 0), and
+        // another +1 to ensure that all origins can point to EOF, eg on (invalid) empty inputs.
+        if (2 + offset + size < offset)
+            return Origin{origin, offset, 0};
+        return origins.emplace(offset, Origin{origin, offset, size}).first->second;
     }
 
-    PosIdx add(const Origin & origin, uint32_t line, uint32_t column)
+    PosIdx add(const Origin & origin, size_t offset)
     {
-        const auto idx = offsets.add({line, column}).second;
-        if (origins.empty() || origins.back().idx != origin.idx) {
-            origin.idx = idx;
-            origins.push_back(origin);
-        }
-        return PosIdx(idx + 1);
+        if (offset > origin.size)
+            return PosIdx();
+        return PosIdx(1 + origin.offset + offset);
     }
 
-    Pos operator[](PosIdx p) const
+    Pos operator[](PosIdx p) const;
+
+    Pos::Origin originOf(PosIdx p) const
     {
-        if (p.id == 0 || p.id > offsets.size())
-            return {};
-        const auto idx = p.id - 1;
-        /* we want the last key <= idx, so we'll take prev(first key > idx).
-           this is guaranteed to never rewind origin.begin because the first
-           key is always 0. */
-        const auto pastOrigin = std::upper_bound(
-            origins.begin(), origins.end(), Origin(idx), [](const auto & a, const auto & b) { return a.idx < b.idx; });
-        const auto origin = *std::prev(pastOrigin);
-        const auto offset = offsets[idx];
-        return {offset.line, offset.column, origin.origin};
+        if (auto o = resolve(p))
+            return o->origin;
+        return std::monostate{};
     }
 };