Fix contraction tree

include/contraction_tree.hpp

@@ -8,97 +8,145 @@
 #include <vector>
 #include <map>
 #include <string>
+#include <memory>
+#include <iostream>  // Add for debug output
 
 #ifdef BACKEND_TORCH
 #else
 namespace cytnx {
   /// @cond
-  class Node {
+  class Node : public std::enable_shared_from_this<Node> {
    public:
-    UniTensor utensor;  // don't worry about copy, because everything are references in cytnx!
+    UniTensor utensor;
     bool is_assigned;
-    Node *left;
-    Node *right;
+    std::shared_ptr<Node> left;
+    std::shared_ptr<Node> right;
+    std::weak_ptr<Node> root;
     std::string name;
-    Node *root;
-
-    Node() : is_assigned(false), left(nullptr), right(nullptr), root(nullptr){};
-    Node(const Node &rhs) {
-      this->left = rhs.left;
-      this->right = rhs.right;
-      this->root = rhs.root;
-      this->utensor = rhs.utensor;
-      this->is_assigned = rhs.is_assigned;
+
+    Node() : is_assigned(false) {}
+
+    Node(const Node& rhs)
+        : utensor(rhs.utensor),
+          is_assigned(rhs.is_assigned),
+          left(rhs.left),
+          right(rhs.right),
+          name(rhs.name) {
+      // Only copy root if it exists
+      if (auto r = rhs.root.lock()) {
+        root = r;
+      }
     }
-    Node &operator==(const Node &rhs) {
-      this->left = rhs.left;
-      this->right = rhs.right;
-      this->root = rhs.root;
-      this->utensor = rhs.utensor;
-      this->is_assigned = rhs.is_assigned;
+
+    Node& operator=(const Node& rhs) {
+      if (this != &rhs) {
+        utensor = rhs.utensor;
+        is_assigned = rhs.is_assigned;
+        left = rhs.left;
+        right = rhs.right;
+        name = rhs.name;
+        if (auto r = rhs.root.lock()) {
+          root = r;
+        }
+      }
       return *this;
     }
-    Node(Node *in_left, Node *in_right, const UniTensor &in_uten = UniTensor())
-        : is_assigned(false), left(nullptr), right(nullptr), root(nullptr) {
-      this->left = in_left;
-      this->right = in_right;
-      in_left->root = this;
-      in_right->root = this;
-      if (in_uten.uten_type() != UTenType.Void) this->utensor = in_uten;
+
+    Node(std::shared_ptr<Node> in_left, std::shared_ptr<Node> in_right,
+         const UniTensor& in_uten = UniTensor())
+        : is_assigned(false), left(in_left), right(in_right) {
+      // Set name based on children
+      if (left && right) {
+        name = "(" + left->name + "," + right->name + ")";
+      }
+
+      if (in_uten.uten_type() != UTenType.Void) {
+        utensor = in_uten;
+      }
     }
-    void assign_utensor(const UniTensor &in_uten) {
-      this->utensor = in_uten;
-      this->is_assigned = true;
+
+    void set_root_ptrs() {
+      try {
+        auto self = shared_from_this();
+
+        if (left) {
+          std::cout << "Setting root for left child of " << name << std::endl;
+          left->root = self;
+          left->set_root_ptrs();
+        }
+
+        if (right) {
+          std::cout << "Setting root for right child of " << name << std::endl;
+          right->root = self;
+          right->set_root_ptrs();
+        }
+      } catch (const std::bad_weak_ptr& e) {
+        std::cerr << "Failed to set root ptrs for node " << name << ": " << e.what() << std::endl;
+        throw;
+      }
     }
+
     void clear_utensor() {
-      this->is_assigned = false;
-      this->utensor = UniTensor();
+      if (left) {
+        left->clear_utensor();
+        left->root.reset();
+      }
+      if (right) {
+        right->clear_utensor();
+        right->root.reset();
+      }
+      is_assigned = false;
+      utensor = UniTensor();
+    }
+
+    void assign_utensor(const UniTensor& in_uten) {
+      utensor = in_uten;
+      is_assigned = true;
     }
   };
 
   class ContractionTree {
    public:
-    std::vector<Node> nodes_container;  // this contains intermediate layer.
-    std::vector<Node> base_nodes;  // this is the button layer.
+    std::vector<std::shared_ptr<Node>> nodes_container;  // intermediate layer
+    std::vector<std::shared_ptr<Node>> base_nodes;  // bottom layer
 
-    ContractionTree(){};
-    ContractionTree(const ContractionTree &rhs) {
-      this->nodes_container = rhs.nodes_container;
-      this->base_nodes = rhs.base_nodes;
-    }
-    ContractionTree &operator==(const ContractionTree &rhs) {
-      this->nodes_container = rhs.nodes_container;
-      this->base_nodes = rhs.base_nodes;
-      return *this;
-    }
+    ContractionTree() = default;
+    ContractionTree(const ContractionTree&) = default;
+    ContractionTree& operator=(const ContractionTree&) = default;
 
-    // clear all the elements in the whole tree.
     void clear() {
       nodes_container.clear();
       base_nodes.clear();
-      // nodes_container.reserve(1024);
     }
-    // clear all the intermediate layer, leave all the base_nodes intact.
-    // and reset the root pointer on the base ondes
+
     void reset_contraction_order() {
-      nodes_container.clear();
-      for (cytnx_uint64 i = 0; i < base_nodes.size(); i++) {
-        base_nodes[i].root = nullptr;
+      // First clear all root pointers
+      for (auto& node : base_nodes) {
+        if (node) node->root.reset();
       }
-      // nodes_container.reserve(1024);
+      // Then clear the container
+      nodes_container.clear();
     }
+
     void reset_nodes() {
-      // reset all nodes but keep the skeleton
-      for (cytnx_uint64 i = 0; i < this->nodes_container.size(); i++) {
-        this->nodes_container[i].clear_utensor();
+      // Clear from root down if we have nodes
+      if (!nodes_container.empty() && nodes_container.back()) {
+        nodes_container.back()->clear_utensor();
       }
-      for (cytnx_uint64 i = 0; i < this->base_nodes.size(); i++) {
-        this->base_nodes[i].clear_utensor();
+      nodes_container.clear();
+
+      // Reset base nodes
+      for (auto& node : base_nodes) {
+        if (node) {
+          node->is_assigned = false;
+          node->utensor = UniTensor();
+        }
       }
     }
+
     void build_default_contraction_tree();
-    void build_contraction_tree_by_tokens(const std::map<std::string, cytnx_uint64> &name2pos,
-                                          const std::vector<std::string> &tokens);
+    void build_contraction_tree_by_tokens(const std::map<std::string, cytnx_uint64>& name2pos,
+                                          const std::vector<std::string>& tokens);
   };
   /// @endcond
 }  // namespace cytnx

include/search_tree.hpp

@@ -7,98 +7,83 @@
 #include <vector>
 #include <map>
 #include <string>
+#include <bitset>
+#include <memory>
+#include <unordered_map>
 #include "UniTensor.hpp"
 
-#ifdef BACKEND_TORCH
-#else
-
 namespace cytnx {
-  /// @cond
-  class PsudoUniTensor {
+
+  using IndexSet = std::bitset<128>;
+
+  class PseudoUniTensor {
    public:
-    // UniTensor utensor; //don't worry about copy, because everything are references in cytnx!
+    bool isLeaf;
+
+    // Leaf node data
     std::vector<std::string> labels;
     std::vector<cytnx_uint64> shape;
     bool is_assigned;
-    PsudoUniTensor *left;
-    PsudoUniTensor *right;
-    PsudoUniTensor *root;
+    cytnx_uint64 tensorIndex;
+
+    // Internal node data
+    std::unique_ptr<PseudoUniTensor> left;
+    std::unique_ptr<PseudoUniTensor> right;
+
     cytnx_float cost;
     cytnx_uint64 ID;
-
     std::string accu_str;
 
-    PsudoUniTensor()
-        : is_assigned(false), left(nullptr), right(nullptr), root(nullptr), cost(0), ID(0){};
-    PsudoUniTensor(const PsudoUniTensor &rhs) {
-      this->left = rhs.left;
-      this->right = rhs.right;
-      this->root = rhs.root;
-      this->labels = rhs.labels;
-      this->shape = rhs.shape;
-      this->is_assigned = rhs.is_assigned;
-      this->cost = rhs.cost;
-      this->accu_str = rhs.accu_str;
-      this->ID = rhs.ID;
-    }
-    PsudoUniTensor &operator==(const PsudoUniTensor &rhs) {
-      this->left = rhs.left;
-      this->right = rhs.right;
-      this->root = rhs.root;
-      this->labels = rhs.labels;
-      this->shape = rhs.shape;
-      this->is_assigned = rhs.is_assigned;
-      this->cost = rhs.cost;
-      this->accu_str = rhs.accu_str;
-      this->ID = rhs.ID;
-      return *this;
-    }
-    void from_utensor(const UniTensor &in_uten) {
-      this->labels = in_uten.labels();
-      this->shape = in_uten.shape();
-      this->is_assigned = true;
-    }
-    void clear_utensor() {
-      this->is_assigned = false;
-      this->labels.clear();
-      this->shape.clear();
-      this->ID = 0;
-      this->cost = 0;
-      this->accu_str = "";
-    }
-    void set_ID(const cytnx_int64 &ID) { this->ID = ID; }
+    // Constructors
+    explicit PseudoUniTensor(cytnx_uint64 index = 0)
+        : isLeaf(true),
+          tensorIndex(index),
+          is_assigned(false),
+          cost(0),
+          ID(1ULL << index),
+          accu_str(std::to_string(index)) {}
+
+    PseudoUniTensor(std::unique_ptr<PseudoUniTensor> l, std::unique_ptr<PseudoUniTensor> r)
+        : isLeaf(false), left(std::move(l)), right(std::move(r)), cost(0), ID(0) {}
+
+    // Copy and move constructors and assignment operators
+    PseudoUniTensor(const PseudoUniTensor& rhs);
+    PseudoUniTensor(PseudoUniTensor&& rhs) noexcept;
+    PseudoUniTensor& operator=(const PseudoUniTensor& rhs);
+    PseudoUniTensor& operator=(PseudoUniTensor&& rhs) noexcept;
+    ~PseudoUniTensor() = default;
+
+    void from_utensor(const UniTensor& in_uten);
+    void clear_utensor();
   };
 
+  namespace OptimalTreeSolver {
+    std::unique_ptr<PseudoUniTensor> solve(const std::vector<PseudoUniTensor>& tensors,
+                                           bool verbose = false);
+  }
+
   class SearchTree {
    public:
-    std::vector<std::vector<PsudoUniTensor>> nodes_container;
-    // std::vector<PsudoUniTensor> nodes_container; // this contains intermediate layer.
-    std::vector<PsudoUniTensor> base_nodes;  // this is the button layer.
-
-    SearchTree(){};
-    SearchTree(const SearchTree &rhs) {
-      this->nodes_container = rhs.nodes_container;
-      this->base_nodes = rhs.base_nodes;
-    }
-    SearchTree &operator==(const SearchTree &rhs) {
-      this->nodes_container = rhs.nodes_container;
-      this->base_nodes = rhs.base_nodes;
-      return *this;
-    }
+    std::vector<PseudoUniTensor> base_nodes;
 
-    // clear all the elements in the whole tree.
+    SearchTree() = default;
     void clear() {
-      nodes_container.clear();
+      root_ptr.reset();
       base_nodes.clear();
-      // nodes_container.reserve(1024);
     }
-    // clear all the intermediate layer, leave all the base_nodes intact.
-    // and reset the root pointer on the base ondes
-    void reset_search_order() { nodes_container.clear(); }
+    void reset_search_order() { root_ptr.reset(); }
     void search_order();
+
+    std::vector<std::vector<PseudoUniTensor*>> get_root() const { return {{root_ptr.get()}}; }
+
+   private:
+    std::unique_ptr<PseudoUniTensor> root_ptr;
   };
-  /// @endcond
+
+  // Helper functions declarations
+  cytnx_float get_cost(const PseudoUniTensor& t1, const PseudoUniTensor& t2);
+  PseudoUniTensor pContract(PseudoUniTensor& t1, PseudoUniTensor& t2);
+
 }  // namespace cytnx
-#endif
 
 #endif  // CYTNX_SEARCH_TREE_H_