codac2_analytic_flat_input_layout.h

Go to the documentation of this file.

 
#pragma once
 
#include <typeindex>
#include <unordered_map>
 
#include "codac2_AnalyticExprWrapper.h"
#include "codac2_AnalyticType.h"
#include "codac2_ExprBase.h"
#include "codac2_FunctionArgsList.h"
 
namespace codac2
{
  class ScalarVar;
  class VectorVar;
  class MatrixVar;

  struct FlatInputBinding
  {
    std::type_index type = typeid(ScalarType); 
    Index offset = 0;                          
    Index rows = 1;                            
    Index cols = 1;                            

    bool is_scalar() const;

    bool is_vector() const;

    bool is_matrix() const;
  };

  class FlatInputLayout
  {
    public:

      explicit FlatInputLayout(const FunctionArgsList& args);

      Index size() const;

      bool same_domain_as(const FlatInputLayout& other) const;

      bool same_domain_as(const FunctionArgsList& other_args) const;

      Index flat_index_of(const ScalarVar& x) const;

      Index flat_index_of(const VectorVar& x, Index i) const;

      Index flat_index_of(const MatrixVar& x, Index i, Index j) const;

      bool flat_index_of(const ScalarExpr& x, Index& flat_index) const;

      const FlatInputBinding& binding_of(const ExprID& id) const;

      const FlatInputBinding* find_binding(const ExprID& id) const;
 
    private:
 
      std::unordered_map<Index,FlatInputBinding> _bindings; 
      Index _size = 0;                                      
  };
 
  template<typename T>
  static std::shared_ptr<ExprBase> as_expr_base(const AnalyticExprWrapper<T>& e)
  {
    return std::static_pointer_cast<ExprBase>(
      std::shared_ptr<AnalyticExpr<T>>(e)
    );
  }
 
  template<typename T>
    requires std::is_base_of_v<AnalyticTypeBase,T>
  class AnalyticFunction;
 
  template<typename T>
  inline AnalyticFunction<T> unaryize_function(const AnalyticFunction<T>& f)
  {
    if(f.nb_args() == 0 || (f.nb_args() == 1 && std::dynamic_pointer_cast<VectorVar>(f.args()[0])))
      return f;
 
    FlatInputLayout layout(f.args());
    VectorVar flat_x(layout.size(), "x");
 
    auto y = std::dynamic_pointer_cast<AnalyticExpr<T>>(f.expr()->copy());
    assert(y && "unaryize_function: unable to copy analytic expression");
 
    for(const auto& arg : f.args())
    {
      const auto& b = layout.binding_of(arg->unique_id());
 
      if(std::dynamic_pointer_cast<ScalarVar>(arg))
        y->replace_arg(arg->unique_id(), as_expr_base(flat_x[b.offset]));
 
      else if(std::dynamic_pointer_cast<VectorVar>(arg))
      {
        assert(b.cols == 1 && "unaryize_function: invalid flat binding for vector input");
        y->replace_arg(
          arg->unique_id(),
          as_expr_base(flat_x.subvector(b.offset, b.offset + b.rows - 1))
        );
      }
 
      else
      {
        assert(false && "unaryize_function: only scalar/vector input arguments are currently supported");
      }
    }
 
    return AnalyticFunction<T>({flat_x}, y);
  }
}