Transform external interface to a modular and extensible design

CMakeLists.txt

@@ -367,6 +367,8 @@ set(SeQuant_symb_src
         SeQuant/core/tensor_network/v2.hpp
         SeQuant/core/tensor_network/v3.cpp
         SeQuant/core/tensor_network/v3.hpp
+        SeQuant/core/tree_index.cpp
+        SeQuant/core/tree_index.hpp
         SeQuant/core/utility/aggregate.hpp
         SeQuant/core/utility/atomic.hpp
         SeQuant/core/utility/context.hpp

SeQuant/core/attr.hpp

@@ -11,6 +11,7 @@
 #include <cstdlib>
 #include <ostream>
 #include <string>
+#include <type_traits>
 
 namespace sequant {
 
@@ -130,10 +131,10 @@ enum class BraKetPos {
 ///
 /// @note This does not include slot bundles, like braket, etc.
 enum class SlotType {
-  Bra,
-  Ket,
-  Aux,
-  Proto,
+  Bra = 0b1,
+  Ket = 0b10,
+  Aux = 0b100,
+  Proto = 0b1000,
 };
 
 template <typename CharT, typename Traits>
@@ -156,6 +157,33 @@ std::basic_ostream<CharT, Traits>& operator<<(
   return stream;
 }
 
+struct SlotTypes {
+  std::underlying_type_t<SlotType> active = 0;
+
+  constexpr SlotTypes(SlotType type)
+      : active(static_cast<decltype(active)>(type)) {}
+  constexpr SlotTypes(decltype(active) val) : active(val) {}
+
+  constexpr bool operator==(const SlotTypes&) const = default;
+  constexpr auto operator<=>(const SlotTypes&) const = default;
+
+  constexpr bool operator&(SlotType type) const {
+    return active & static_cast<decltype(active)>(type);
+  }
+
+  constexpr SlotTypes operator|(SlotType type) const {
+    return {active | static_cast<decltype(active)>(type)};
+  }
+};
+
+constexpr SlotTypes operator|(SlotType lhs, SlotType rhs) {
+  using IntType = std::underlying_type_t<SlotType>;
+  return {static_cast<IntType>(lhs) | static_cast<IntType>(rhs)};
+}
+
+static constexpr const SlotTypes AnySlotType =
+    SlotType::Bra | SlotType::Ket | SlotType::Aux | SlotType::Proto;
+
 enum class Statistics {
   FermiDirac,
   BoseEinstein,

SeQuant/core/export/export.hpp

@@ -10,6 +10,7 @@
 #include <SeQuant/core/export/expression_group.hpp>
 #include <SeQuant/core/export/generator.hpp>
 #include <SeQuant/core/expr.hpp>
+#include <SeQuant/core/io/latex/latex.hpp>
 #include <SeQuant/core/logger.hpp>
 #include <SeQuant/core/utility/macros.hpp>
 #include <SeQuant/core/utility/tensor.hpp>
@@ -798,7 +799,8 @@ void preprocess_and_maybe_log(ExportNode<T> &tree, PreprocessResult &result,
     std::cout << "Tree before preprocessing:\n"
               << tree.tikz(
                      [](const ExportNode<T> &node) {
-                       return "$" + toUtf8(to_latex(node->expr())) + "$";
+                       return "$" + toUtf8(io::latex::to_string(node->expr())) +
+                              "$";
                      },
                      [](const ExportNode<T>) -> std::string { return ""; })
               << "\n";
@@ -811,7 +813,8 @@ void preprocess_and_maybe_log(ExportNode<T> &tree, PreprocessResult &result,
     std::cout << "Tree after pre-processing:\n"
               << tree.tikz(
                      [](const ExportNode<T> &node) {
-                       return "$" + toUtf8(to_latex(node->expr())) + "$";
+                       return "$" + toUtf8(io::latex::to_string(node->expr())) +
+                              "$";
                      },
                      [](const ExportNode<T>) -> std::string { return ""; })
               << "\n";

SeQuant/core/expressions/expr.cpp

@@ -14,6 +14,7 @@
 #include <SeQuant/core/tensor_canonicalizer.hpp>
 #include <SeQuant/core/tensor_network.hpp>
 #include <SeQuant/core/tensor_network/typedefs.hpp>
+#include <SeQuant/core/tree_index.hpp>
 #include <SeQuant/core/utility/macros.hpp>
 
 #include <range/v3/algorithm/any_of.hpp>
@@ -65,6 +66,14 @@ const ExprPtr &Expr::operator[](std::size_t idx) const {
   return begin()[idx];
 }
 
+ExprPtr &ExprPtr::operator[](const TreeIndex &idx) {
+  return idx.select_from(*this);
+}
+
+const ExprPtr &ExprPtr::operator[](const TreeIndex &idx) const {
+  return idx.select_from(*this);
+}
+
 ExprPtr &Expr::at(std::size_t idx) { return (*this)[idx]; }
 
 const ExprPtr &Expr::at(std::size_t idx) const { return (*this)[idx]; }
@@ -183,6 +192,12 @@ Expr &Expr::operator+=(const Expr &) { throw not_implemented("operator+="); }
 
 Expr &Expr::operator-=(const Expr &) { throw not_implemented("operator-="); }
 
+Expr &Expr::operator[](const TreeIndex &idx) { return idx.select_from(*this); }
+
+const Expr &Expr::operator[](const TreeIndex &idx) const {
+  return idx.select_from(*this);
+}
+
 ExprPtr adjoint(const ExprPtr &expr) {
   auto result = expr->clone();
   result->adjoint();

SeQuant/core/expressions/expr.hpp

@@ -4,6 +4,7 @@
 #include <SeQuant/core/expressions/expr_iterator.hpp>
 #include <SeQuant/core/expressions/expr_ptr.hpp>
 #include <SeQuant/core/options.hpp>
+#include <SeQuant/core/tree_index.hpp>
 #include <SeQuant/core/utility/macros.hpp>
 
 #include <boost/core/demangle.hpp>
@@ -381,6 +382,10 @@ class Expr : public std::enable_shared_from_this<Expr> {
   ExprPtr &operator[](std::size_t idx);
   const ExprPtr &operator[](std::size_t idx) const;
 
+  /// @return The subexpression identified by the given index
+  Expr &operator[](const TreeIndex &idx);
+  const Expr &operator[](const TreeIndex &idx) const;
+
   ExprPtr &at(std::size_t idx);
   const ExprPtr &at(std::size_t idx) const;
 

SeQuant/core/expressions/expr_ptr.hpp

@@ -4,6 +4,7 @@
 #include <SeQuant/core/container.hpp>
 #include <SeQuant/core/expr_fwd.hpp>
 #include <SeQuant/core/expressions/traits.hpp>
+#include <SeQuant/core/tree_index.hpp>
 #include <SeQuant/core/utility/macros.hpp>
 
 #include <memory>
@@ -111,6 +112,10 @@ class ExprPtr : public std::shared_ptr<Expr> {
   /// @return reference to @c *this
   ExprPtr &operator*=(const ExprPtr &);
 
+  /// @return The subexpression identified by the given index
+  ExprPtr &operator[](const TreeIndex &idx);
+  const ExprPtr &operator[](const TreeIndex &idx) const;
+
   /// @tparam T an Expr type
   /// @return true if this object is of type @c T
   template <typename T>

SeQuant/core/io/latex/latex.cpp

@@ -5,12 +5,15 @@
 #include <SeQuant/core/container.hpp>
 #include <SeQuant/core/expressions/constant.hpp>
 #include <SeQuant/core/expressions/power.hpp>
+#include <SeQuant/core/expressions/result_expr.hpp>
 #include <SeQuant/core/io/latex/latex.hpp>
 #include <SeQuant/core/rational.hpp>
 #include <SeQuant/core/utility/string.hpp>
 
 #include <map>
 #include <optional>
+#include <sstream>
+#include <string>
 #include <vector>
 
 namespace sequant::io::latex {
@@ -64,6 +67,19 @@ std::wstring to_string(const Power& power) {
   return result;
 }
 
+std::wstring to_string(const ResultExpr& expr) {
+  std::wstringstream stream;
+  if (expr.produces_tensor()) {
+    stream << to_string(expr.result_as_tensor());
+  } else {
+    stream << to_string(expr.result_as_variable());
+  }
+
+  stream << " = " << to_string(expr.expression());
+
+  return stream.str();
+}
+
 namespace detail {
 
 template <typename Char, typename Traits, typename Alloc>

SeQuant/core/io/latex/latex.hpp

@@ -19,7 +19,8 @@
 
 namespace sequant {
 class Power;
-}
+class ResultExpr;
+}  // namespace sequant
 
 namespace sequant::io::latex {
 
@@ -105,6 +106,8 @@ std::wstring to_string(const rational& num);
 
 std::wstring to_string(const Power& power);
 
+std::wstring to_string(const ResultExpr& expr);
+
 namespace detail {
 
 template <typename Char, typename Traits, typename Alloc>

SeQuant/core/meta.hpp

@@ -9,6 +9,7 @@
 #include <memory>
 #include <ranges>
 #include <type_traits>
+#include <utility>
 
 #include <range/v3/range/access.hpp>
 #include <range/v3/range/traits.hpp>
@@ -570,6 +571,27 @@ static_assert(
     std::same_as<mimic_constness_t<const int &, float &>, const float &>);
 static_assert(std::same_as<mimic_constness_t<int &, const float &>, float &>);
 
+#if defined(__cpp_lib_forward_like) && __cpp_lib_forward_like >= 202207L
+using std::forward_like;
+#else
+template <class T, class U>
+// Taken from cppreference https://en.cppreference.com/cpp/utility/forward_like
+constexpr auto &&forward_like(U &&x) noexcept {
+  constexpr bool is_adding_const = std::is_const_v<std::remove_reference_t<T>>;
+  if constexpr (std::is_lvalue_reference_v<T &&>) {
+    if constexpr (is_adding_const)
+      return std::as_const(x);
+    else
+      return static_cast<U &>(x);
+  } else {
+    if constexpr (is_adding_const)
+      return std::move(std::as_const(x));
+    else
+      return std::move(x);
+  }
+}
+#endif
+
 ///
 /// True if @p T is a range of rank @p Rank whose value type is convertible to
 /// @p V.

SeQuant/core/tree_index.cpp

@@ -0,0 +1,71 @@
+#include <SeQuant/core/container.hpp>
+#include <SeQuant/core/expr.hpp>
+#include <SeQuant/core/tree_index.hpp>
+#include <SeQuant/core/utility/exception.hpp>
+#include <SeQuant/core/utility/macros.hpp>
+
+#include <concepts>
+#include <cstddef>
+#include <initializer_list>
+#include <ranges>
+#include <string>
+#include <type_traits>
+
+namespace sequant {
+
+TreeIndex::TreeIndex(container::svector<std::size_t> positions)
+    : positions_(std::move(positions)) {}
+
+TreeIndex::TreeIndex(std::initializer_list<std::size_t> positions)
+    : TreeIndex(container::svector<std::size_t>(std::move(positions))) {}
+
+template <typename ExprType>
+ExprType &select_impl(ExprType &expr,
+                      const container::svector<std::size_t> &positions) {
+  constexpr bool is_expr_ptr =
+      std::same_as<std::remove_cvref_t<ExprType>, ExprPtr>;
+
+  ExprType *selected = &expr;
+
+  for (std::size_t current : positions) {
+    if (current >= selected->size()) {
+      throw Exception("Position " + std::to_string(current) +
+                      " is out of bounds for an expression of dimension " +
+                      std::to_string(selected->size()));
+    }
+
+    using std::ranges::begin;
+
+    auto it = begin(*selected) + current;
+
+    if constexpr (is_expr_ptr) {
+      selected = &(*it);
+    } else {
+      selected = &(*(*it));
+    }
+  }
+
+  SEQUANT_ASSERT(selected);
+
+  return *selected;
+}
+
+ExprPtr &TreeIndex::select_from(ExprPtr &expr) const {
+  return select_impl(expr, positions_);
+}
+
+const ExprPtr &TreeIndex::select_from(const ExprPtr &expr) const {
+  return select_impl(expr, positions_);
+}
+
+Expr &TreeIndex::select_from(Expr &expr) const {
+  return select_impl(expr, positions_);
+}
+
+const Expr &TreeIndex::select_from(const Expr &expr) const {
+  return select_impl(expr, positions_);
+}
+
+std::size_t TreeIndex::depth() const { return positions_.size(); }
+
+}  // namespace sequant