canonical_trace.hpp

Go to the documentation of this file.

/*
 * Copyright © 2025-2026 Akreon LLC
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <cstddef>
#include <cstdint>
#include <limits>
#include <tuple>
#include <type_traits>
#include <utility>
 
#include "akreon/convex/cone_traits.hpp"
#include "akreon/convex/convex_traits.hpp"
 
namespace convex {
 
enum struct CanonicalConeKind : uint8_t {
  NONNEG_ORTHANT = 0,
  SECOND_ORDER = 1,
  POSITIVE_SEMIDEFINITE = 2,
  EXP = 3,
  POWER = 4,
  UNKNOWN = 255
};
 
inline constexpr std::size_t canonical_no_parent =
    std::numeric_limits<std::size_t>::max();
 
struct CanonicalNodeRecord {
  std::size_t node_id{}; 
 
  std::size_t parent_id{canonical_no_parent};
 
  AtomType atom_type{};
 
  std::size_t child_index{};
};
 
struct CanonicalConeRecord {
  std::size_t node_id{};          
  std::size_t local_cone_index{}; 
  std::size_t
      global_cone_index{}; 
  CanonicalConeKind cone_kind{CanonicalConeKind::UNKNOWN};
  std::size_t cone_offset{}; 
  std::size_t cone_dim{};    
};
 
struct CanonicalEqualityRecord {
  std::size_t node_id{}; 
  std::size_t
      local_eq_index{};     
  std::size_t row_offset{}; 
  std::size_t row_count{};  
};
 
template <typename NodeRecords = std::tuple<>,
          typename ConeRecords = std::tuple<>,
          typename EqualityRecords = std::tuple<>>
struct CanonicalTrace {
  NodeRecords nodes{};          
  ConeRecords cones{};          
  EqualityRecords equalities{}; 
};
 
using EmptyCanonicalTrace = CanonicalTrace<>;
 
template <typename Trace>
struct canonical_trace_traits;
 
template <typename NodeTuple, typename ConeTuple, typename EqualityTuple>
struct canonical_trace_traits<
    CanonicalTrace<NodeTuple, ConeTuple, EqualityTuple>> {
  using Nodes = NodeTuple;
  using Cones = ConeTuple;
  using Equalities = EqualityTuple;
 
  static constexpr std::size_t NumNodes = std::tuple_size_v<NodeTuple>;
  static constexpr std::size_t NumCones = std::tuple_size_v<ConeTuple>;
  static constexpr std::size_t NumEqualities = std::tuple_size_v<EqualityTuple>;
};
 
template <typename Trace>
using canonical_trace_nodes_t =
    typename canonical_trace_traits<linalg::remove_cvref_t<Trace>>::Nodes;
 
template <typename Trace>
using canonical_trace_cones_t =
    typename canonical_trace_traits<linalg::remove_cvref_t<Trace>>::Cones;
 
template <typename Trace>
using canonical_trace_equalities_t =
    typename canonical_trace_traits<linalg::remove_cvref_t<Trace>>::Equalities;
 
template <typename Trace>
inline constexpr std::size_t canonical_trace_num_nodes_v =
    canonical_trace_traits<linalg::remove_cvref_t<Trace>>::NumNodes;
 
template <typename Trace>
inline constexpr std::size_t canonical_trace_num_cones_v =
    canonical_trace_traits<linalg::remove_cvref_t<Trace>>::NumCones;
 
template <typename Trace>
inline constexpr std::size_t canonical_trace_num_equalities_v =
    canonical_trace_traits<linalg::remove_cvref_t<Trace>>::NumEqualities;
 
template <typename Cone>
struct canonical_cone_kind {
  static constexpr CanonicalConeKind value = CanonicalConeKind::UNKNOWN;
};
 
template <linalg::VectorContainer Vec>
struct canonical_cone_kind<conic::NonnegOrthantCone<Vec>> {
  static constexpr CanonicalConeKind value = CanonicalConeKind::NONNEG_ORTHANT;
};
 
template <linalg::VectorContainer Vec>
struct canonical_cone_kind<conic::SecondOrderCone<Vec>> {
  static constexpr CanonicalConeKind value = CanonicalConeKind::SECOND_ORDER;
};
 
template <linalg::VectorContainer Vec, std::size_t Dim>
struct canonical_cone_kind<conic::PositiveSemidefiniteCone<Vec, Dim>> {
  static constexpr CanonicalConeKind value =
      CanonicalConeKind::POSITIVE_SEMIDEFINITE;
};
 
template <linalg::VectorContainer Vec>
struct canonical_cone_kind<conic::ExpCone<Vec>> {
  static constexpr CanonicalConeKind value = CanonicalConeKind::EXP;
};
 
template <linalg::VectorContainer Vec>
struct canonical_cone_kind<conic::PowerCone<Vec>> {
  static constexpr CanonicalConeKind value = CanonicalConeKind::POWER;
};
 
template <typename Cone>
inline constexpr CanonicalConeKind canonical_cone_kind_v =
    canonical_cone_kind<linalg::remove_cvref_t<Cone>>::value;
 
template <typename ConeTuple, std::size_t Index>
consteval std::size_t cone_tuple_offset_before_v() {
  using Tuple = linalg::remove_cvref_t<ConeTuple>;
  if constexpr (Index == 0) {
    return 0;
  } else {
    using PrevCone = std::tuple_element_t<Index - 1, Tuple>;
    return cone_tuple_offset_before_v<ConeTuple, Index - 1>() +
           conic::cone_dim_t<PrevCone>;
  }
}
 
template <typename Trace, typename NodeRecords, typename ConeRecords,
          typename EqualityRecords>
constexpr auto extend_canonical_trace(const Trace& trace,
                                      const NodeRecords& nodes,
                                      const ConeRecords& cones,
                                      const EqualityRecords& equalities) {
  return CanonicalTrace<decltype(std::tuple_cat(trace.nodes, nodes)),
                        decltype(std::tuple_cat(trace.cones, cones)),
                        decltype(std::tuple_cat(trace.equalities, equalities))>{
      std::tuple_cat(trace.nodes, nodes), std::tuple_cat(trace.cones, cones),
      std::tuple_cat(trace.equalities, equalities)};
}
 
template <typename ConeTuple>
constexpr auto make_canonical_cone_records(std::size_t node_id,
                                           std::size_t global_cone_start,
                                           std::size_t cone_offset_start) {
  using Tuple = linalg::remove_cvref_t<ConeTuple>;
 
  auto make_records = [&]<std::size_t... Is>(std::index_sequence<Is...>) {
    return std::make_tuple(CanonicalConeRecord{
        node_id, Is, global_cone_start + Is,
        canonical_cone_kind_v<std::tuple_element_t<Is, Tuple>>,
        cone_offset_start + cone_tuple_offset_before_v<Tuple, Is>(),
        conic::cone_dim_t<std::tuple_element_t<Is, Tuple>>}...);
  };
 
  return make_records(std::make_index_sequence<std::tuple_size_v<Tuple>>{});
}
 
} // namespace convex