Fix joins, rename, union

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

3 files changed, 105 insertions, 13 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);
     };