Program Listing for File expr.hpp

↰ Return to documentation for file (SeQuant/core/expressions/expr.hpp)

#ifndef SEQUANT_EXPRESSIONS_EXPR_HPP
#define SEQUANT_EXPRESSIONS_EXPR_HPP

#include <SeQuant/core/context.hpp>
#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>

#include <atomic>
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <optional>
#include <ranges>
#include <string>
#include <string_view>
#include <type_traits>
#include <typeindex>
#include <typeinfo>
#include <utility>

namespace sequant {

static const wchar_t adjoint_label = L'\u207A';

inline bool is_adjoint_label(std::wstring_view label) {
  return !label.empty() && label.back() == adjoint_label;
}

inline std::wstring_view strip_adjoint_label(std::wstring_view label) {
  return is_adjoint_label(label) ? label.substr(0, label.size() - 1) : label;
}

inline void toggle_adjoint_label(std::wstring &label) {
  if (is_adjoint_label(label))
    label.pop_back();
  else
    label.push_back(adjoint_label);
}

namespace detail {

constexpr std::uint64_t fnv1a_64(std::string_view str) {
  std::uint64_t hash = 0xcbf29ce484222325ull;
  for (const char c : str) {
    hash ^= static_cast<unsigned char>(c);
    hash *= 0x100000001b3ull;
  }
  return hash;
}

constexpr std::string uint_to_string(std::uint64_t n) {
  std::string result;
  do {
    result.insert(result.begin(), static_cast<char>('0' + n % 10));
    n /= 10;
  } while (n != 0);
  return result;
}

template <typename T>
constexpr std::string compiler_type_name() {
#if defined(_MSC_VER)
  return __FUNCSIG__;
#else
  return __PRETTY_FUNCTION__;
#endif
}

struct CanonicalMark {
  std::atomic<std::uint64_t> stamp{0};
  std::atomic<std::size_t> seal{0};

  CanonicalMark() = default;
  CanonicalMark(const CanonicalMark &other) noexcept { *this = other; }
  CanonicalMark(CanonicalMark &&other) noexcept { *this = std::move(other); }
  CanonicalMark &operator=(const CanonicalMark &other) noexcept {
    stamp.store(other.stamp.load(std::memory_order_relaxed),
                std::memory_order_relaxed);
    seal.store(other.seal.load(std::memory_order_relaxed),
               std::memory_order_relaxed);
    return *this;
  }
  CanonicalMark &operator=(CanonicalMark &&other) noexcept {
    if (this != &other) {
      *this = static_cast<const CanonicalMark &>(other);
      other.reset();
    }
    return *this;
  }

  void reset() noexcept {
    stamp.store(0, std::memory_order_relaxed);
    seal.store(0, std::memory_order_relaxed);
  }
};

}  // namespace detail

template <typename T>
constexpr std::string type_name_of() {
  if constexpr (requires {
                  {
                    T::static_type_name(std::type_identity<T>{})
                  } -> std::convertible_to<std::string>;
                })
    return std::string(T::static_type_name(std::type_identity<T>{}));
  else
    return detail::compiler_type_name<T>();
}


class Expr : public std::enable_shared_from_this<Expr> {
 public:
  using hash_type = std::size_t;
  using type_id_type = std::uint64_t;
  using type_rank_type = std::uint8_t;

  static constexpr type_rank_type default_type_rank = 128;

  Expr() = default;
  virtual ~Expr() = default;

  bool is_atom() const { return empty(); }

  virtual bool is_zero() const { return false; }

  virtual std::wstring to_latex() const;

  virtual ExprPtr clone() const = 0;

  ExprPtr exprptr_from_this() {
    if (weak_from_this().use_count() == 0)
      return this->clone();
    else
      return static_cast<ExprPtr>(this->shared_from_this());
  }

  ExprPtr exprptr_from_this() const {
    if (weak_from_this().use_count() == 0)
      return this->clone();
    else
      return static_cast<const ExprPtr>(
          std::const_pointer_cast<Expr>(this->shared_from_this()));
  }

  virtual ExprPtr canonicalize(
      CanonicalizeOptions = CanonicalizeOptions::default_options()) {
    return {};  // by default do nothing and return nullptr
  }

  virtual ExprPtr rapid_canonicalize(
      CanonicalizeOptions opts =
          CanonicalizeOptions::default_options().copy_and_set(
              CanonicalizationMethod::Rapid)) {
    return this->canonicalize(opts.copy_and_set(CanonicalizationMethod::Rapid));
  }

  bool is_canonical(const CanonicalizeOptions &opts =
                        CanonicalizeOptions::default_options()) const;

  void mark_canonical(
      const CanonicalizeOptions &opts,
      std::uint64_t contexts_version = current_contexts_version()) const;

  // clang-format off
  // clang-format on
  template <typename Visitor>
  bool visit(Visitor &&visitor, const bool atoms_only = false) {
    return visit_impl(*this, std::forward<Visitor>(visitor), atoms_only);
  }

  template <typename Visitor>
  bool visit(Visitor &&visitor, const bool atoms_only = false) const {
    return visit_impl(*this, std::forward<Visitor>(visitor), atoms_only);
  }

  Expr &expr() { return *this; }
  const Expr &expr() const { return *this; }

  template <typename T, typename Enabler = void>
  struct is_shared_ptr_of_expr : std::false_type {};
  template <typename T>
  struct is_shared_ptr_of_expr<std::shared_ptr<T>,
                               std::enable_if_t<is_expr_v<T>>>
      : std::true_type {};
  template <typename T, typename Enabler = void>
  struct is_shared_ptr_of_expr_or_derived : std::false_type {};
  template <typename T>
  struct is_shared_ptr_of_expr_or_derived<std::shared_ptr<T>,
                                          std::enable_if_t<is_an_expr_v<T>>>
      : std::true_type {};

  virtual bool is_scalar() const {
    if (is_atom()) return false;
    for (auto it = begin(); it != end(); ++it) {
      if (!(*it)->is_scalar()) return false;
    }
    return true;
  }

  virtual bool is_cnumber() const {
    if (is_atom())
      return true;
    else {
      bool result = true;
      for (auto it = begin(); result && it != end(); ++it) {
        result &= (*it)->is_cnumber();
      }
      return result;
    }
  }

  bool commutes_with(const Expr &that) const {
    auto this_is_atom = is_atom();
    auto that_is_atom = that.is_atom();
    bool result = true;
    if (this_is_atom && that_is_atom) {
      result =
          this->is_cnumber() || that.is_cnumber() || commutes_with_atom(that);
    } else if (this_is_atom) {
      if (!this->is_cnumber()) {
        for (auto it = that.begin(); result && it != that.end(); ++it) {
          result &= this->commutes_with(**it);
        }
      }
    } else {
      for (auto it = this->begin(); result && it != this->end(); ++it) {
        result &= (*it)->commutes_with(that);
      }
    }
    return result;
  }

  virtual void adjoint() = 0;

  hash_type hash_value(
      std::function<hash_type(const std::shared_ptr<const Expr> &)> hasher = {})
      const {
    return hasher ? hasher(shared_from_this()) : memoizing_hash();
  }

  virtual type_id_type type_id() const = 0;

  friend inline bool operator==(const Expr &a, const Expr &b);

  template <typename T>
    requires(is_an_expr_v<T>)
  bool operator<(const T &that) const {
    if (type_id() ==
        that.type_id()) {  // if same type, use generic (or type-specific, if
                           // available) comparison
      return static_less_than(static_cast<const Expr &>(that));
    } else {  // order types by type id
      return type_id() < that.type_id();
    }
  }

  static constexpr type_id_type make_type_id(type_rank_type rank,
                                             std::string_view name) {
    return (static_cast<type_id_type>(rank) << 56) |
           (detail::fnv1a_64(name) >> 8);
  }

  static constexpr type_rank_type type_rank_of(type_id_type id) {
    return static_cast<type_rank_type>(id >> 56);
  }

  template <typename T>
  static type_id_type get_type_id() {
    static const type_id_type id = [] {
      const std::string name = sequant::type_name_of<T>();
      type_rank_type rank = default_type_rank;
      if constexpr (requires { T::type_rank; }) {
        static_assert(std::in_range<type_rank_type>(T::type_rank),
                      "type_rank must fit in Expr::type_rank_type");
        rank = static_cast<type_rank_type>(T::type_rank);
      }
      const type_id_type result = make_type_id(rank, name);
      register_type_id(result, name, typeid(T));
      return result;
    }();
    return id;
  }

  template <typename T>
  bool is() const {
    if constexpr (is_expr_v<T>)
      return true;
    else if constexpr (std::is_base_of_v<Expr, T>)
      return this->type_id() == get_type_id<std::remove_cvref_t<T>>();
    else
      return dynamic_cast<const T *>(this) != nullptr;
  }

  template <typename T>
  const T &as() const {
    SEQUANT_ASSERT(this->is<T>());
    if constexpr (std::is_base_of_v<Expr, T>) {
      return static_cast<const T &>(*this);
    } else
      return dynamic_cast<const T &>(*this);
  }

  template <typename T>
  T &as() {
    SEQUANT_ASSERT(this->is<T>());
    if constexpr (std::is_base_of_v<Expr, T>) {
      return static_cast<T &>(*this);
    } else
      return dynamic_cast<T &>(*this);
  }

  std::string type_name() const {
    return boost::core::demangle(typeid(*this).name());
  }

  ExprIterator begin();
  ExprIterator end();
  ConstExprIterator begin() const;
  ConstExprIterator end() const;
  ConstExprIterator cbegin() const;
  ConstExprIterator cend() const;

  virtual ExprIterator begin_subexpr();
  virtual ExprIterator end_subexpr();
  virtual ConstExprIterator begin_subexpr() const;
  virtual ConstExprIterator end_subexpr() const;

  std::size_t size() const;

  bool empty() const;

  ExprPtr &operator[](std::size_t idx);
  const ExprPtr &operator[](std::size_t idx) const;

  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;

  ExprPtr &front();
  const ExprPtr &front() const;

  ExprPtr &back();
  const ExprPtr &back() const;

 private:
  [[noreturn]] void throw_out_of_range(std::size_t idx) const;

  template <
      typename E, typename Visitor,
      typename = std::enable_if_t<std::is_same_v<std::remove_cvref_t<E>, Expr>>>
  static bool visit_impl(E &&expr, Visitor &&visitor, const bool atoms_only) {
    if (expr.weak_from_this().use_count() == 0)
      throw Exception(
          "Expr::visit: cannot visit expressions not managed by shared_ptr");
    for (auto &subexpr_ptr : expr.expr()) {
      const auto subexpr_is_an_atom = subexpr_ptr->is_atom();
      const auto need_to_visit_subexpr = !atoms_only || subexpr_is_an_atom;
      bool visited = false;
      if (!subexpr_is_an_atom)  // if not a leaf, recur into it
        visited = visit_impl(*subexpr_ptr, std::forward<Visitor>(visitor),
                             atoms_only);
      // call on the subexpression itself, if not yet done so
      if (need_to_visit_subexpr && !visited) visitor(subexpr_ptr);
    }
    // N.B. can only visit itself if visitor is nonmutating!
    bool this_visited = false;
    if constexpr (std::is_invocable_r_v<void, std::remove_reference_t<Visitor>,
                                        const ExprPtr &>) {
      if (!atoms_only || expr.is_atom()) {
        const ExprPtr this_exprptr = expr.exprptr_from_this();
        visitor(this_exprptr);
        this_visited = true;
      }
    }
    return this_visited;
  }

 protected:
  Expr(Expr &&) = default;
  Expr(const Expr &) = default;
  Expr &operator=(Expr &&) = default;
  Expr &operator=(const Expr &) = default;

  mutable std::optional<hash_type> hash_value_;  // not initialized by default
  virtual hash_type memoizing_hash() const {
    static const hash_type default_hash_value = 0;
    if (hash_value_)
      return *hash_value_;
    else
      return default_hash_value;
  }
  virtual void reset_hash_value() const {
    hash_value_.reset();
    reset_canonical_mark();
  }

  void reset_canonical_mark() const { canonical_mark_.reset(); }

  void copy_canonical_mark(const Expr &other) const {
    canonical_mark_ = other.canonical_mark_;
  }

  virtual bool static_equal(const Expr &that) const = 0;

  virtual bool static_less_than(const Expr &that) const {
    return this->hash_value() < that.hash_value();
  }

  virtual bool commutes_with_atom([[maybe_unused]] const Expr &that) const {
    return true;
  }

 private:
  static void register_type_id(type_id_type id, const std::string &name,
                               std::type_index type);
  mutable detail::CanonicalMark canonical_mark_;

  void stamp_canonical_subtree() const;

  std::optional<std::size_t> canonical_subtree_digest() const;
};  // class Expr

namespace expr_type_rank {
inline constexpr Expr::type_rank_type tensor = 10;
inline constexpr Expr::type_rank_type product = 20;
inline constexpr Expr::type_rank_type constant = 30;
inline constexpr Expr::type_rank_type sum = 40;
inline constexpr Expr::type_rank_type variable = 50;
inline constexpr Expr::type_rank_type power = 60;
inline constexpr Expr::type_rank_type boperator = 250;
inline constexpr Expr::type_rank_type foperator = 251;
inline constexpr Expr::type_rank_type bnoperator = 252;
inline constexpr Expr::type_rank_type fnoperator = 253;
}  // namespace expr_type_rank

static_assert(std::ranges::sized_range<Expr>);
static_assert(std::ranges::bidirectional_range<Expr>);
static_assert(std::ranges::random_access_range<Expr>);

template <>
struct Expr::is_shared_ptr_of_expr<ExprPtr, void> : std::true_type {};
template <>
struct Expr::is_shared_ptr_of_expr_or_derived<ExprPtr, void> : std::true_type {
};

inline bool operator==(const Expr &a, const Expr &b) {
  if (a.type_id() != b.type_id())
    return false;
  else
    return a.static_equal(b);
}

struct proportional_to {
  bool operator()(const ExprPtr &expr1, const ExprPtr &expr2) const;
};

}  // namespace sequant

#endif  // SEQUANT_EXPRESSIONS_EXPR_HPP