Fix joins, rename, union

Co-authored-by: Marcos Magueta <maguetamarcos@gmail.com>

5 files changed, 300 insertions, 18 deletions

diff --git a/include/RNT_C_API.h b/include/RNT_C_API.h
index 1a59785..19f6165 100644
--- a/include/RNT_C_API.h
+++ b/include/RNT_C_API.h
@@ -397,6 +397,7 @@ typedef struct {
 typedef struct {
     rnt_plan_t left;
     rnt_plan_t right;
+    const char** attrs;
 } PlanArgsJoin;
 
 /**
@@ -422,6 +423,38 @@ typedef struct {
 } PlanArgsProject;
 
 /**
+ * MATERIALIZE context: caches @p source on its first scan, then replays the
+ * cache on later scans without touching @p source again.
+ *
+ * Put this on the inner side of a JOIN so the inner relation is read once
+ * instead of once per outer tuple. Only safe when @p source does not depend on
+ * the outer tuple (no Var SCAN leaves bound by the JOIN), since the cache is
+ * reused as-is for every outer tuple. Takes ownership of @p source; on failure
+ * @p source is freed.
+ */
+typedef struct {
+    rnt_plan_t source;
+} PlanArgsMaterialize;
+
+/**
+ * RENAME context: returns @p source, with the keys of @p pairs the values of it.
+ *
+ * @p pairs is a NULL-terminated array of pairs containing the original attribute
+ * name together with it's new name.
+ */
+
+typedef struct {
+    rnt_plan_t source;
+    const char** pairs;
+} PlanArgsRename;
+
+/** UNION context. */
+
+typedef struct {
+    rnt_plan_t *sources;
+} PlanArgsUnion;
+
+/**
  * A single plan-construction request. @p operation selects which context member
  * is read; the others are ignored. The members are laid out side by side rather
  * than overlapped in a union so the struct maps cleanly onto OCaml ctypes, which
@@ -433,6 +466,9 @@ typedef struct {
     PlanArgsJoin join;
     PlanArgsTake take;
     PlanArgsProject project;
+    PlanArgsMaterialize materialize;
+    PlanArgsRename rename;
+    PlanArgsUnion union_;
 } PlanAction;
 
 /**
diff --git a/include/Types.h b/include/Types.h
index 6956e09..08eebcf 100644
--- a/include/Types.h
+++ b/include/Types.h
@@ -2,6 +2,8 @@
 
 #include <string>
 #include <vector>
+#include <unordered_map>
+#include <ranges>
 
 /**
  * @file Types.h
@@ -132,28 +134,41 @@ namespace nt {
      * one attribute at a time. Returns nullptr when all attributes are exhausted.
      */
     class Tuple {
-      public:
-        explicit Tuple(std::vector<Attribute> attributes) : attributes_(std::move(attributes)) {
+    public:
+        explicit Tuple(std::vector<Attribute> attributes) {
+            for (auto &attr : attributes) attributes_[attr.name] = attr.value;
+            this->Reset();
         }
 
         /** @brief Returns the next attribute, or nullptr when exhausted. */
         const Attribute* Next() {
-            if (position_ < attributes_.size())
-                return &attributes_[position_++];
-            return nullptr;
+            if (iter_ == attributes_.end()) return nullptr;
+            current_attr_ = {iter_->first, iter_->second};
+            iter_++;
+            return &current_attr_;
         }
 
         void Reset() {
-            position_ = 0;
+            iter_ = attributes_.begin();
         }
 
-        /** @brief Direct read-only access to all attributes without affecting position. */
-        const std::vector<Attribute>& attrs() const {
-            return attributes_;
+        const std::string operator[](std::string key) {
+            return attributes_[key];
         }
 
-      private:
-        std::vector<Attribute> attributes_;
-        size_t position_ = 0;
+        const auto attrs() const {
+          auto view =
+              attributes_ |
+              std::views::transform([](std::pair<std::string, std::string> p) {
+                  Attribute attr { p.first, p.second };
+                  return attr;
+              });
+          return view;
+        }
+
+    private:
+        std::unordered_map<std::string, std::string> attributes_{};
+        std::unordered_map<std::string, std::string>::iterator iter_;
+        Attribute current_attr_;
     };
 } // namespace nt
diff --git a/include/VM.h b/include/VM.h
index 3d97779..3f2f665 100644
--- a/include/VM.h
+++ b/include/VM.h
@@ -81,7 +81,10 @@ namespace nt {
         FOL_OPERATION_SCAN = 1,
         FOL_OPERATION_JOIN,
         FOL_OPERATION_TAKE,
-        FOL_OPERATION_PROJECT
+        FOL_OPERATION_PROJECT,
+        FOL_OPERATION_MATERIALIZE,
+        FOL_OPERATION_RENAME,
+        FOL_OPERATION_UNION
     };
 
     /**
@@ -99,6 +102,16 @@ namespace nt {
      * - TAKE  : passes at most take_limit tuples through, then returns nullptr. Use
      *           this to bound scans over AlephZero ephemeral relations.
      * - PROJECT: filters each input tuple to the named attributes in project_attrs.
+     * - MATERIALIZE: on its first scan it pulls the whole child and keeps a
+     *           private copy of every tuple. On later scans it replays that copy
+     *           and never reads the child again. Sits on the inner side of a
+     *           nested-loop JOIN so the inner relation is read once instead of
+     *           once per outer tuple. This only works when the child does not
+     *           depend on the outer tuple, meaning it has no Var SCAN leaves that
+     *           the JOIN binds, because the same cached tuples are handed back for
+     *           every outer tuple. Rewind does not pass the outer tuple through
+     *           this node, so it is up to the planner to place it only where the
+     *           child is independent of the outer side.
      */
     struct PlanNode {
         Operation op;
@@ -142,6 +155,26 @@ namespace nt {
          */
         Tuple* join_left = nullptr;
         std::optional<Tuple> join_buffer;
+        std::unordered_set<std::string> join_attrs;
+
+        /**
+         * MATERIALIZE runtime state.
+         *
+         * mat_buffer holds owned copies of the child's tuples, filled on the
+         * first scan. mat_done marks the cache complete (child exhausted).
+         * mat_pos is the replay cursor, reset to 0 by Rewind on each later scan.
+         */
+        std::vector<Tuple> mat_buffer;
+        std::size_t mat_pos = 0;
+        bool mat_done = false;
+
+        /** RENAME runtime state. */
+        PlanNode* upstream = nullptr;
+        std::unordered_map<std::string, std::string> attrs;
+
+        /** UNION runtime state. */
+        std::vector<PlanNode*> nodes;
+        std::vector<PlanNode*>::iterator current_node;
     };
 
     class VM {
@@ -184,6 +217,14 @@ namespace nt {
         /** Resets the inner cursor state and writes resolved args into it. */
         static void ResetInner(CursorManager::cursor* c, std::vector<std::string> args);
 
+        /**
+         * Restarts the inner (right) subtree of a JOIN for a new outer tuple.
+         * Walks down to every SCAN leaf, rebinds its scan_args to the outer
+         * tuple, and resets its cursor. On the way it also clears the operator
+         * state it passes through, such as TAKE counters and nested JOIN buffers.
+         */
+        void Rewind(PlanNode* node, Tuple* outer);
+
         /** Merges left and right tuple attributes into node->join_buffer. */
         static Tuple* MergeInto(PlanNode* node, Tuple* left, Tuple* right);
     };
diff --git a/src/RNT_C_API.cpp b/src/RNT_C_API.cpp
index abd683c..a632ab9 100644
--- a/src/RNT_C_API.cpp
+++ b/src/RNT_C_API.cpp
@@ -503,6 +503,10 @@ namespace {
         node->left = l->root;
         node->right = r->root;
 
+        if (a.attrs)
+            for (const char** p = a.attrs; *p; ++p)
+                node->join_attrs.emplace(*p);
+
         auto* pw = new PlanWrapper();
         pw->root = node.get();
         pw->nodes.push_back(std::move(node));
@@ -577,6 +581,85 @@ namespace {
         delete s;
         return pw;
     }
+
+    PlanWrapper* build_materialize(const PlanArgsMaterialize& a) {
+        auto* s = static_cast<PlanWrapper*>(a.source);
+        if (!s)
+            return nullptr;
+
+        auto node = std::make_unique<nt::PlanNode>();
+        node->op = nt::FOL_OPERATION_MATERIALIZE;
+        node->left = s->root;
+
+        auto* pw = new PlanWrapper();
+        pw->root = node.get();
+        pw->nodes.push_back(std::move(node));
+        for (auto& n : s->nodes)
+            pw->nodes.push_back(std::move(n));
+        for (auto* c : s->cursors)
+            pw->cursors.push_back(c);
+        for (auto* h : s->handles)
+            pw->handles.push_back(h);
+
+        delete s;
+        return pw;
+    }
+
+    PlanWrapper* build_rename(const PlanArgsRename& a) {
+        auto* s = static_cast<PlanWrapper*>(a.source);
+        if (!s)
+            return nullptr;
+
+        auto node = std::make_unique<nt::PlanNode>();
+        node->op = nt::FOL_OPERATION_RENAME;
+        node->upstream = s->root;
+
+        for (const char** k = a.pairs; *k; k++)
+            node->attrs[*k] = *++k;
+
+        auto* pw = new PlanWrapper();
+        pw->root = node.get();
+        pw->nodes.push_back(std::move(node));
+
+        for (auto& n : s->nodes)
+            pw->nodes.push_back(std::move(n));
+        for (auto* c : s->cursors)
+            pw->cursors.push_back(c);
+        for (auto* h : s->handles)
+            pw->handles.push_back(h);
+
+        delete s;
+        return pw;
+    }
+
+    PlanWrapper* build_union(const PlanArgsUnion& a) {
+        if (!*a.sources)
+            return nullptr;
+
+        auto node = std::make_unique<nt::PlanNode>();
+        node->op = nt::FOL_OPERATION_UNION;
+
+        auto* pw = new PlanWrapper();
+        for (auto* s = a.sources; *s; s++) {
+            auto* source = static_cast<PlanWrapper*>(*s);
+
+            node->nodes.push_back(source->root);
+
+            for (auto& n : source->nodes)
+                pw->nodes.push_back(std::move(n));
+            for (auto* c : source->cursors)
+                pw->cursors.push_back(std::move(c));
+            for (auto* h : source->handles)
+                pw->handles.push_back(std::move(h));
+
+            delete source;
+        }
+
+        pw->root = node.get();
+        pw->nodes.push_back(std::move(node));
+
+        return pw;
+    }
 } // namespace
 
 // ---------------------------------------------------------------------------
@@ -1019,6 +1102,12 @@ rnt_plan_t rnt_plan_assemble(PlanAction action) {
         return build_take(action.take);
     case nt::FOL_OPERATION_PROJECT:
         return build_project(action.project);
+    case nt::FOL_OPERATION_MATERIALIZE:
+        return build_materialize(action.materialize);
+    case nt::FOL_OPERATION_RENAME:
+        return build_rename(action.rename);
+    case nt::FOL_OPERATION_UNION:
+        return build_union(action.union_);
     }
 
     return nullptr;
diff --git a/src/VM.cpp b/src/VM.cpp
index 5f66f83..1f509ff 100644
--- a/src/VM.cpp
+++ b/src/VM.cpp
@@ -32,6 +32,49 @@ namespace nt {
         c->exhausted = false;
     }
 
+    void VM::Rewind(PlanNode* node, Tuple* outer) {
+        if (node == nullptr)
+            return;
+
+        switch (node->op) {
+        case Operation::FOL_OPERATION_SCAN:
+            ResetInner(node->scan_cursor, ResolveArgs(node->scan_args, outer));
+            break;
+
+        case Operation::FOL_OPERATION_PROJECT:
+            Rewind(node->left, outer);
+            break;
+
+        case Operation::FOL_OPERATION_TAKE:
+            node->take_count = 0;
+            Rewind(node->left, outer);
+            break;
+
+        case Operation::FOL_OPERATION_JOIN:
+            node->join_left = nullptr;
+            Rewind(node->left, outer);
+            Rewind(node->right, outer);
+            break;
+
+        case Operation::FOL_OPERATION_MATERIALIZE:
+            // Rewind the replay cursor only, and don't pass `outer` to the child.
+            // The point of this node is to reuse the cached tuples for every
+            // outer tuple, so it must not re-run or rebind what it cached.
+            node->mat_pos = 0;
+            break;
+
+        case Operation::FOL_OPERATION_RENAME:
+            Rewind(node->upstream, outer);
+            break;
+
+        case Operation::FOL_OPERATION_UNION:
+            node->current_node = node->nodes.begin();
+            for (auto& node : node->nodes)
+                Rewind(node, outer);
+            break;
+        }
+    }
+
     Tuple* VM::MergeInto(PlanNode* node, Tuple* left, Tuple* right) {
         std::vector<Attribute> merged;
         for (const auto& a : left->attrs())
@@ -67,18 +110,22 @@ namespace nt {
 
         case Operation::FOL_OPERATION_JOIN: {
             while (true) {
+            again:
                 if (node->join_left == nullptr) {
                     node->join_left = Next(node->left);
-                    if (node->join_left == nullptr)
-                        return nullptr;
+                    if (node->join_left == nullptr) return nullptr;
 
-                    auto args = ResolveArgs(node->right->scan_args, node->join_left);
-                    ResetInner(node->right->scan_cursor, std::move(args));
+                    Rewind(node->right, node->join_left);
                 }
 
                 Tuple* right = Next(node->right);
-                if (right != nullptr)
+                if (right != nullptr) {
+                    for (auto &attr : node->join_attrs)
+                        if ((*node->join_left)[attr] != (*right)[attr])
+                            goto again;
+
                     return MergeInto(node, node->join_left, right);
+                }
 
                 node->join_left = nullptr;
             }
@@ -92,6 +139,60 @@ namespace nt {
                 ++node->take_count;
             return t;
         }
+
+        case Operation::FOL_OPERATION_MATERIALIZE: {
+            if (!node->mat_done) {
+                // First scan: pull a tuple and keep our own copy. We rebuild it
+                // from attrs() rather than copying the Tuple, because a Tuple
+                // holds an iterator into its own map and a plain copy would leave
+                // that iterator dangling. The child also overwrites its output
+                // buffer on the next pull, so we cannot just keep its pointer.
+                Tuple* t = Next(node->left);
+                if (t != nullptr) {
+                    std::vector<Attribute> attrs;
+                    for (const auto& a : t->attrs())
+                        attrs.push_back(a);
+                    node->mat_buffer.emplace_back(std::move(attrs));
+                    return &node->mat_buffer.back();
+                }
+                node->mat_done = true;
+                return nullptr; // child exhausted; do not fall into replay
+            }
+
+            // Later scans: hand back cached tuples until the cache runs out.
+            if (node->mat_pos < node->mat_buffer.size())
+                return &node->mat_buffer[node->mat_pos++];
+            return nullptr;
+        }
+
+        case Operation::FOL_OPERATION_RENAME: {
+            Tuple* t = Next(node->upstream);
+            if (!t)
+                return nullptr;
+
+            std::vector<Attribute> attrs;
+            for (const auto& a : t->attrs()) {
+                std::string name = a.name;
+                if (node->attrs.contains(name))
+                    name = node->attrs[name];
+                Attribute new_attr { name, a.value };
+                attrs.push_back(new_attr);
+            }
+            delete t;
+            return new Tuple(attrs);
+        }
+
+        case Operation::FOL_OPERATION_UNION: {
+            if (node->current_node == node->nodes.end())
+                return nullptr;
+
+            Tuple* t = Next(*node->current_node);
+            if (!t) {
+                node->current_node++;
+                return Next(node);
+            }
+            return t;
+        }
         }
 
         return nullptr;