node.hpp

Go to the documentation of this file.

//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
#ifndef node_h
#define node_h
 
#include <iostream>
#include <string>
#include <memory>
#include <iomanip>
#include <functional>
#include <type_traits>
 
#include "backend.hpp"
 
namespace graph {
//******************************************************************************
//  Base leaf node.
//******************************************************************************
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class leaf_node : public std::enable_shared_from_this<leaf_node<T, SAFE_MATH>> {
    protected:
        const size_t hash;
        const size_t complexity;
        std::unordered_map<size_t, std::shared_ptr<leaf_node<T, SAFE_MATH>>> df_cache;
        const bool contains_pseudo;
 
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        leaf_node(const std::string s,
                  const size_t count,
                  const bool pseudo) :
        hash(std::hash<std::string>{} (s)),
        complexity(count), contains_pseudo(pseudo) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual ~leaf_node() {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual backend::buffer<T> evaluate() = 0;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual std::shared_ptr<leaf_node<T, SAFE_MATH>> reduce() {
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual std::shared_ptr<leaf_node<T, SAFE_MATH>>
        df(std::shared_ptr<leaf_node<T, SAFE_MATH>> x) {
            return std::shared_ptr<leaf_node<T, SAFE_MATH>> ();
        };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void compile_preamble(std::ostringstream &stream,
                                      jit::register_map &registers,
                                      jit::visiter_map &visited,
                                      jit::register_usage &usage,
                                      jit::texture1d_list &textures1d,
                                      jit::texture2d_list &textures2d,
                                      jit::preamble_map &pre_funcs,
                                      int &avail_const_mem) {
#ifdef SHOW_USE_COUNT
            if (usage.find(this) == usage.end()) {
                usage[this] = 1;
            } else {
                ++usage[this];
            }
#endif
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual std::shared_ptr<leaf_node<T, SAFE_MATH>>
        compile(std::ostringstream &stream,
                jit::register_map &registers,
                const jit::register_map &thread_mem,
                const jit::register_usage &usage) = 0;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_match(std::shared_ptr<leaf_node<T, SAFE_MATH>> x) {
            return this == x.get();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        bool is_power_base_match(std::shared_ptr<leaf_node<T, SAFE_MATH>> x) {
            return this->get_power_base()->is_match(x->get_power_base());
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void set(const T d) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void set(const size_t index,
                         const T d) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void set(const std::vector<T> &d) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void set(const backend::buffer<T> &d) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void to_latex() const {};
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual std::shared_ptr<leaf_node<T, SAFE_MATH>> to_vizgraph(std::stringstream &stream,
                                                                     jit::register_map &registers) = 0;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_constant() const {
            return false;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool has_constant_zero() const {
            return false;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        bool is_normal() {
            return this->evaluate().is_normal();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_all_variables() const {
            return false;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_power_like() const {
            return false;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual std::shared_ptr<leaf_node<T, SAFE_MATH>> get_power_base() {
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual std::shared_ptr<leaf_node<T, SAFE_MATH>> get_power_exponent() const = 0;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        size_t get_hash() const {
            return hash;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        size_t get_complexity() const {
            return complexity;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool has_pseudo() const {
            return contains_pseudo;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual std::shared_ptr<leaf_node<T, SAFE_MATH>> remove_pseudo() {
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void endline(std::ostringstream &stream,
                             const jit::register_usage &usage,
                             const char end=';')
#ifndef SHOW_USE_COUNT
                             const
#endif
                             {
            stream << end
#ifdef SHOW_USE_COUNT
                   << " // used " << usage.at(this)
#endif
                   << std::endl;
        }
 
//  Create one struct that holds both caches: for constructed nodes and for the
//  backend buffers
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        struct caches_t {
            std::unordered_map<size_t, std::shared_ptr<leaf_node<T, SAFE_MATH>>> nodes;
            std::unordered_map<size_t, backend::buffer<T>> backends;
        };
 
        inline static thread_local caches_t caches;
 
        typedef T base;
    };
 
    template<jit::float_scalar T, bool SAFE_MATH=false>
    using shared_leaf = std::shared_ptr<leaf_node<T, SAFE_MATH>>;
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    constexpr shared_leaf<T, SAFE_MATH> null_leaf() {
        return shared_leaf<T, SAFE_MATH> ();
    }
    template<jit::float_scalar T, bool SAFE_MATH=false>
    using output_nodes = std::vector<shared_leaf<T, SAFE_MATH>>;
 
    template<jit::float_scalar T, bool SAFE_MATH=false>
    constexpr shared_leaf<T, SAFE_MATH> zero();
    template<jit::float_scalar T, bool SAFE_MATH=false>
    constexpr shared_leaf<T, SAFE_MATH> one();
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    void make_vizgraph(shared_leaf<T, SAFE_MATH> node) {
        std::stringstream stream;
        jit::register_map registers;
        stream << std::setprecision(jit::max_digits10<T> ());
 
        stream << "graph \"\" {" << std::endl;
        stream << "    node [fontname = \"Helvetica\", ordering = out]" << std::endl << std::endl;
        node->to_vizgraph(stream, registers);
        stream << "}" << std::endl;
 
        std::cout << stream.str() << std::endl;
    }
 
//******************************************************************************
//******************************************************************************
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class index_node final : public leaf_node<T, SAFE_MATH> {
    private:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static std::string to_string() {
            return "i";
        }
        
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        index_node() :
        leaf_node<T, SAFE_MATH> (index_node::to_string(), 1, false) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual backend::buffer<T> evaluate() {
            return backend::buffer<T> ();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> df(shared_leaf<T, SAFE_MATH> x) {
            return this->is_match(x) ? one<T, SAFE_MATH> () : zero<T, SAFE_MATH> ();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual std::shared_ptr<leaf_node<T, SAFE_MATH>>
        compile(std::ostringstream &stream,
                jit::register_map &registers,
                const jit::register_map &thread_mem,
                const jit::register_usage &usage) {
            if (registers.find(this) == registers.end()) {
                registers[this] = jit::to_string('i', this);
                stream << "        const ";
                if constexpr (jit::use_cuda()) {
                    stream << "int " << registers[this] << " = index";
                } else if constexpr (jit::use_metal<T> ()) {
                    stream << "uint " << registers[this] << " = index";
                } else {
                    stream << "size_t " << registers[this] << " = i";
                }
                this->endline(stream, usage);
            }
 
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void to_latex() const {
            std::cout << "i";
        };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> to_vizgraph(std::stringstream &stream,
                                                      jit::register_map &registers) {
            if (registers.find(this) == registers.end()) {
                const std::string name = jit::to_string('i', this);
                registers[this] = name;
                stream << "    " << name
                       << " [label = \"i\", shape = box, style = \"rounded,filled\", fillcolor = black, fontcolor = white];" << std::endl;
            }
 
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> get_power_exponent() const {
            return one<T, SAFE_MATH> ();
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    shared_leaf<T, SAFE_MATH> index() {
        auto temp = std::make_shared<index_node<T, SAFE_MATH>> ();
//  Test for hash collisions.
        for (size_t i = temp->get_hash(); i < std::numeric_limits<size_t>::max(); i++) {
            if (leaf_node<T, SAFE_MATH>::caches.nodes.find(i) ==
                leaf_node<T, SAFE_MATH>::caches.nodes.end()) {
                leaf_node<T, SAFE_MATH>::caches.nodes[i] = temp;
                return temp;
            } else if (temp->is_match(leaf_node<T, SAFE_MATH>::caches.nodes[i])) {
                return leaf_node<T, SAFE_MATH>::caches.nodes[i];
            }
        }
#if defined(__clang__) || defined(__GNUC__)
        __builtin_unreachable();
#else
        assert(false && "Should never reach.");
#endif
    }
 
    template<jit::float_scalar T, bool SAFE_MATH=false>
    using shared_index = std::shared_ptr<index_node<T, SAFE_MATH>>;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    shared_index<T, SAFE_MATH> index_cast(shared_leaf<T, SAFE_MATH> x) {
        return std::dynamic_pointer_cast<index_node<T, SAFE_MATH>> (x);
    }
 
//******************************************************************************
//  Constant node.
//******************************************************************************
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class constant_node final : public leaf_node<T, SAFE_MATH> {
    private:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static std::string to_string(const T d) {
            return jit::format_to_string<T> (d);
        }
 
        const backend::buffer<T> data;
 
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        constant_node(const backend::buffer<T> &d) :
        leaf_node<T, SAFE_MATH> (constant_node::to_string(d.at(0)), 1, false), data(d) {
            if constexpr (SAFE_MATH) {
                assert(d.is_normal() && "Denormal encountered");
            }
            assert(d.size() == 1 && "Constants need to be scalar functions.");
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual backend::buffer<T> evaluate() {
            return data;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> df(shared_leaf<T, SAFE_MATH> x) {
            return this->is_match(x) ? one<T, SAFE_MATH> () : zero<T, SAFE_MATH> ();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH>
        compile(std::ostringstream &stream,
                jit::register_map &registers,
                const jit::register_map &thread_mem,
                const jit::register_usage &usage) {
            if (registers.find(this) == registers.end()) {
#ifdef USE_CONSTANT_CACHE
                registers[this] = jit::to_string('r', this);
                stream << "        const ";
                jit::add_type<T> (stream);
                const T temp = this->evaluate().at(0);
 
                stream << " " << registers[this] << " = ";
                if constexpr (jit::complex_scalar<T>) {
                    jit::add_type<T> (stream);
                }
                stream << temp;
                this->endline(stream, usage);
#else
                if constexpr (jit::complex_scalar<T>) {
                    registers[this] = jit::get_type_string<T> () + "("
                                    + jit::format_to_string(this->evaluate().at(0))
                                    + ")";
                } else {
                    registers[this] = "(" + jit::get_type_string<T> () + ")"
                                    + jit::format_to_string(this->evaluate().at(0));
                }
#endif
            }
 
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_match(shared_leaf<T, SAFE_MATH> x) {
            if (this == x.get()) {
                return true;
            }
 
            auto x_cast = constant_cast(x);
            if (x_cast.get()) {
                return this->evaluate() == x_cast->evaluate();
            }
 
            return false;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        bool is(const T d) {
            return data.size() == 1 && data.at(0) == d;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        bool is_integer() {
            const auto temp = this->evaluate().at(0);
            return std::imag(temp) == 0 &&
                   fmod(std::real(temp), 1.0) == 0.0;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void to_latex() const {
            std::cout << data.at(0);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> to_vizgraph(std::stringstream &stream,
                                                      jit::register_map &registers) {
            if (registers.find(this) == registers.end()) {
                const std::string name = jit::to_string('r', this);
                registers[this] = name;
                stream << "    " << name
                       << " [label = \"" << this->evaluate().at(0)
                       << "\", shape = box, style = \"rounded,filled\", fillcolor = black, fontcolor = white];" << std::endl;
            }
 
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_constant() const {
            return true;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool has_constant_zero() const {
            return data.has_zero();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_power_like() const {
            return true;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> get_power_base() {
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> get_power_exponent() const {
            return one<T, SAFE_MATH> ();
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    shared_leaf<T, SAFE_MATH> constant(const backend::buffer<T> &d) {
        auto temp = std::make_shared<constant_node<T, SAFE_MATH>> (d);
//  Test for hash collisions.
        for (size_t i = temp->get_hash(); i < std::numeric_limits<size_t>::max(); i++) {
            if (leaf_node<T, SAFE_MATH>::caches.nodes.find(i) ==
                leaf_node<T, SAFE_MATH>::caches.nodes.end()) {
                leaf_node<T, SAFE_MATH>::caches.nodes[i] = temp;
                return temp;
            } else if (temp->is_match(leaf_node<T, SAFE_MATH>::caches.nodes[i])) {
                return leaf_node<T, SAFE_MATH>::caches.nodes[i];
            }
        }
#if defined(__clang__) || defined(__GNUC__)
        __builtin_unreachable();
#else
        assert(false && "Should never reach.");
#endif
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    shared_leaf<T, SAFE_MATH> constant(const T d) {
        return constant<T, SAFE_MATH> (backend::buffer<T> (1, d));
    }
 
//  Define some common constants.
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH>
    constexpr shared_leaf<T, SAFE_MATH> zero() {
        return constant<T, SAFE_MATH> (static_cast<T> (0.0));
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH>
    constexpr shared_leaf<T, SAFE_MATH> one() {
        return constant<T, SAFE_MATH> (static_cast<T> (1.0));
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    constexpr shared_leaf<T, SAFE_MATH> none() {
        return constant<T, SAFE_MATH> (static_cast<T> (-1.0));
    }
 
    template<jit::complex_scalar T>
    constexpr T i = T(0.0, 1.0);
 
    template<jit::float_scalar T, bool SAFE_MATH=false>
    using shared_constant = std::shared_ptr<constant_node<T, SAFE_MATH>>;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    shared_constant<T, SAFE_MATH> constant_cast(shared_leaf<T, SAFE_MATH> x) {
        return std::dynamic_pointer_cast<constant_node<T, SAFE_MATH>> (x);
    }
 
//******************************************************************************
//  Base straight node.
//******************************************************************************
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class straight_node : public leaf_node<T, SAFE_MATH> {
    protected:
        shared_leaf<T, SAFE_MATH> arg;
 
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        straight_node(shared_leaf<T, SAFE_MATH> a,
                      const std::string s) :
        leaf_node<T, SAFE_MATH> (s, a->get_complexity() + 1, a->has_pseudo()),
        arg(a) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual backend::buffer<T> evaluate() {
            return arg->evaluate();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void compile_preamble(std::ostringstream &stream,
                                      jit::register_map &registers,
                                      jit::visiter_map &visited,
                                      jit::register_usage &usage,
                                      jit::texture1d_list &textures1d,
                                      jit::texture2d_list &textures2d,
                                      jit::preamble_map &pre_funcs,
                                      int &avail_const_mem) {
            if (!visited.contains(this)) {
                arg->compile_preamble(stream, registers,
                                      visited, usage,
                                      textures1d, textures2d,
                                      pre_funcs, avail_const_mem);
                visited.insert(this);
#ifdef SHOW_USE_COUNT
                usage[this] = 1;
            } else {
                ++usage[this];
#endif
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH>
        compile(std::ostringstream &stream,
                jit::register_map &registers,
                const jit::register_map &thread_mem,
                const jit::register_usage &usage) {
            return arg->compile(stream, registers, thread_mem, usage);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        shared_leaf<T, SAFE_MATH> get_arg() const {
            return arg;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_all_variables() const {
            return arg->is_all_variables();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> get_power_exponent() const {
            return one<T, SAFE_MATH> ();
        }
    };
 
//******************************************************************************
//  Base branch node.
//******************************************************************************
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class branch_node : public leaf_node<T, SAFE_MATH> {
    protected:
        shared_leaf<T, SAFE_MATH> left;
        shared_leaf<T, SAFE_MATH> right;
 
    public:
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        branch_node(shared_leaf<T, SAFE_MATH> l,
                    shared_leaf<T, SAFE_MATH> r,
                    const std::string s) :
        leaf_node<T, SAFE_MATH> (s, l->get_complexity() + r->get_complexity() + 1,
                                 l->has_pseudo() || r->has_pseudo()),
        left(l), right(r) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        branch_node(shared_leaf<T, SAFE_MATH> l,
                    shared_leaf<T, SAFE_MATH> r,
                    const std::string s,
                    const size_t count,
                    const bool pseudo) :
        leaf_node<T, SAFE_MATH> (s, count, pseudo),
        left(l), right(r) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void compile_preamble(std::ostringstream &stream,
                                      jit::register_map &registers,
                                      jit::visiter_map &visited,
                                      jit::register_usage &usage,
                                      jit::texture1d_list &textures1d,
                                      jit::texture2d_list &textures2d,
                                      jit::preamble_map &pre_funcs,
                                      int &avail_const_mem) {
            if (!visited.contains(this)) {
                left->compile_preamble(stream, registers,
                                       visited, usage,
                                       textures1d, textures2d,
                                       pre_funcs, avail_const_mem);
                right->compile_preamble(stream, registers,
                                        visited, usage,
                                        textures1d, textures2d,
                                        pre_funcs, avail_const_mem);
                visited.insert(this);
#ifdef SHOW_USE_COUNT
                usage[this] = 1;
            } else {
                ++usage[this];
#endif
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        shared_leaf<T, SAFE_MATH> get_left() const {
            return left;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        shared_leaf<T, SAFE_MATH> get_right() const {
            return right;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_all_variables() const {
            return left->is_all_variables() &&
                   right->is_all_variables();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual std::shared_ptr<leaf_node<T, SAFE_MATH>>
        get_power_exponent() const {
            return one<T, SAFE_MATH> ();
        }
    };
 
//******************************************************************************
//  Base triple node.
//******************************************************************************
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class triple_node : public branch_node<T, SAFE_MATH> {
    protected:
        shared_leaf<T, SAFE_MATH> middle;
 
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        triple_node(shared_leaf<T, SAFE_MATH> l,
                    shared_leaf<T, SAFE_MATH> m,
                    shared_leaf<T, SAFE_MATH> r,
                    const std::string s) :
        branch_node<T, SAFE_MATH> (l, r, s,
                                   l->get_complexity() +
                                   m->get_complexity() +
                                   r->get_complexity(),
                                   l->has_pseudo() ||
                                   m->has_pseudo() ||
                                   r->has_pseudo()),
        middle(m) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void compile_preamble(std::ostringstream &stream,
                                      jit::register_map &registers,
                                      jit::visiter_map &visited,
                                      jit::register_usage &usage,
                                      jit::texture1d_list &textures1d,
                                      jit::texture2d_list &textures2d,
                                      jit::preamble_map &pre_funcs,
                                      int &avail_const_mem) {
            if (!visited.contains(this)) {
                this->left->compile_preamble(stream, registers,
                                             visited, usage,
                                             textures1d, textures2d,
                                             pre_funcs, avail_const_mem);
                this->middle->compile_preamble(stream, registers,
                                               visited, usage,
                                               textures1d, textures2d,
                                               pre_funcs, avail_const_mem);
                this->right->compile_preamble(stream, registers,
                                              visited, usage,
                                              textures1d, textures2d,
                                              pre_funcs, avail_const_mem);
                visited.insert(this);
#ifdef SHOW_USE_COUNT
                usage[this] = 1;
            } else {
                ++usage[this];
#endif
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        shared_leaf<T, SAFE_MATH> get_middle() const {
            return middle;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_all_variables() const {
            return this->left->is_all_variables()   &&
                   middle->is_all_variables() &&
                   this->right->is_all_variables();
        }
    };
 
//******************************************************************************
//  Base N arg node.
//******************************************************************************
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<size_t N, jit::float_scalar T, bool SAFE_MATH=false>
    class n_branch_node : public leaf_node<T, SAFE_MATH> {
    protected:
        std::array<shared_leaf<T, SAFE_MATH>, N> branches;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        bool any_has_pseudo(std::array<shared_leaf<T, SAFE_MATH>, N> &branches) {
            for (auto &b : branches) {
                const bool test = b->has_pseudo();
                if (test) {
                    return test;
                }
            }
            return false;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        size_t total_complexity(std::array<shared_leaf<T, SAFE_MATH>, N> &branches) {
            size_t complexity = 1;
            for (auto &b : branches) {
                complexity += b->get_complexity();
            }
            return complexity;
        }
 
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        n_branch_node(std::array<shared_leaf<T, SAFE_MATH>, N> branches,
                      const std::string s) :
        leaf_node<T, SAFE_MATH> (s, n_branch_node::total_complexity(branches),
                                 n_branch_node::any_has_pseudo(branches)),
        branches(branches) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void compile_preamble(std::ostringstream &stream,
                                      jit::register_map &registers,
                                      jit::visiter_map &visited,
                                      jit::register_usage &usage,
                                      jit::texture1d_list &textures1d,
                                      jit::texture2d_list &textures2d,
                                      jit::preamble_map &pre_funcs,
                                      int &avail_const_mem) {
            if (!visited.contains(this)) {
                for (auto &b : branches) {
                    b->compile_preamble(stream, registers,
                                        visited, usage,
                                        textures1d, textures2d,
                                        pre_funcs, avail_const_mem);
                }
 
            visited.insert(this);
#ifdef SHOW_USE_COUNT
                usage[this] = 1;
            } else {
                ++usage[this];
#endif
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        shared_leaf<T, SAFE_MATH> get_arg(const size_t index) const {
            return branches[index];
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_all_variables() const {
            for (auto b : branches) {
                const bool test = b->is_all_variables();
                if (!test) {
                    return test;
                }
            }
            return true;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual std::shared_ptr<leaf_node<T, SAFE_MATH>>
        get_power_exponent() const {
            return one<T, SAFE_MATH> ();
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<typename T>
    struct has_no_derivative : std::false_type {};
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false,
             class BASE_NODE=leaf_node<T, SAFE_MATH>, size_t N=1>
    class no_derivative : public BASE_NODE {
    public:
        template<typename S=bool>
        std::shared_ptr<leaf_node<T, SAFE_MATH>>
        df(std::shared_ptr<leaf_node<T, SAFE_MATH>> x) requires(has_no_derivative<S>::value);
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        no_derivative(const std::string s)
        requires(std::is_base_of_v<leaf_node<T, SAFE_MATH>,
                                   no_derivative<T, SAFE_MATH,
                                                 leaf_node<T, SAFE_MATH>>>) :
        leaf_node<T, SAFE_MATH> (s, 0, false) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        no_derivative(shared_leaf<T, SAFE_MATH> arg,
                      const std::string s)
        requires(std::is_base_of_v<straight_node<T, SAFE_MATH>,
                                   no_derivative<T, SAFE_MATH,
                                                 straight_node<T, SAFE_MATH>>>) :
        straight_node<T, SAFE_MATH> (arg, s) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        no_derivative(shared_leaf<T, SAFE_MATH> l,
                      shared_leaf<T, SAFE_MATH> r,
                      const std::string s)
        requires(std::is_base_of_v<branch_node<T, SAFE_MATH>,
                                   no_derivative<T, SAFE_MATH,
                                                 branch_node<T, SAFE_MATH>>>) :
        branch_node<T, SAFE_MATH> (l, r, s) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        no_derivative(shared_leaf<T, SAFE_MATH> l,
                      shared_leaf<T, SAFE_MATH> m,
                      shared_leaf<T, SAFE_MATH> r,
                      const std::string s)
        requires(std::is_base_of_v<triple_node<T, SAFE_MATH>,
                                   no_derivative<T, SAFE_MATH,
                                                 triple_node<T, SAFE_MATH>>>) :
        triple_node<T, SAFE_MATH> (l, m, r, s) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        no_derivative(std::array<shared_leaf<T, SAFE_MATH>, N> b,
                      const std::string s)
        requires(std::is_base_of_v<n_branch_node<N, T, SAFE_MATH>,
                                   no_derivative<T, SAFE_MATH,
                                                 n_branch_node<N, T,
                                                               SAFE_MATH>,
                                                 N>>) :
        n_branch_node<N, T, SAFE_MATH> (b, s) {}
    };
 
//******************************************************************************
//  Variable node.
//******************************************************************************
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class variable_node final : public leaf_node<T, SAFE_MATH> {
    private:
        backend::buffer<T> buffer;
        const std::string symbol;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static std::string to_string(variable_node<T, SAFE_MATH> *p) {
            return jit::format_to_string(reinterpret_cast<size_t> (p));
        }
 
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        variable_node(const size_t s,
                      const std::string &symbol) :
        leaf_node<T, SAFE_MATH> (variable_node::to_string(this), 1, false),
        buffer(s), symbol(symbol) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        variable_node(const size_t s, const T d,
                      const std::string &symbol) :
        leaf_node<T, SAFE_MATH> (variable_node::to_string(this), 1, false),
        buffer(s, d), symbol(symbol) {
            assert(buffer.is_normal() && "NaN or Inf value.");
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        variable_node(const std::vector<T> &d,
                      const std::string &symbol) :
        leaf_node<T, SAFE_MATH> (variable_node::to_string(this), 1, false),
        buffer(d), symbol(symbol) {
            assert(buffer.is_normal() && "NaN or Inf value.");
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        variable_node(const backend::buffer<T> &d,
                      const std::string &symbol) :
        leaf_node<T, SAFE_MATH> (variable_node::to_string(this), 1, false),
        buffer(d), symbol(symbol) {
            assert(buffer.is_normal() && "NaN or Inf value.");
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual backend::buffer<T> evaluate() {
            return buffer;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> df(shared_leaf<T, SAFE_MATH> x) {
            return constant<T, SAFE_MATH> (static_cast<T> (this->is_match(x)));
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void compile_preamble(std::ostringstream &stream,
                                      jit::register_map &registers,
                                      jit::visiter_map &visited,
                                      jit::register_usage &usage,
                                      jit::texture1d_list &textures1d,
                                      jit::texture2d_list &textures2d,
                                      jit::preamble_map &pre_funcs,
                                      int &avail_const_mem) {
            if (usage.find(this) == usage.end()) {
                usage[this] = 1;
#ifdef SHOW_USE_COUNT
            } else {
                ++usage[this];
#endif
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH>
        compile(std::ostringstream &stream,
                jit::register_map &registers,
                const jit::register_map &thread_mem,
                const jit::register_usage &usage) {
           return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void set(const T d) {
            buffer.set(d);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void set(const size_t index, const T d) {
            buffer[index] = d;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void set(const std::vector<T> &d) {
            buffer.set(d);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void set(const backend::buffer<T> &d) {
            buffer = d;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        std::string get_symbol() const {
            return symbol;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void to_latex() const {
            std::cout << get_symbol();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> to_vizgraph(std::stringstream &stream,
                                                      jit::register_map &registers) {
            if (registers.find(this) == registers.end()) {
                const std::string name = jit::to_string('r', this);
                registers[this] = name;
                stream << "    " << name
                       << " [label = \"" << this->get_symbol()
                       << "\", shape = box];" << std::endl;
            }
 
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        size_t size() {
            return buffer.size();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        T *data() {
            return buffer.data();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_all_variables() const {
            return true;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_power_like() const {
            return true;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> get_power_base() {
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> get_power_exponent() const {
            return one<T, SAFE_MATH> ();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void endline(std::ostringstream &stream,
                             const jit::register_usage &usage,
                             const char end=';')
        #ifndef SHOW_USE_COUNT
                             const
        #endif
                             {
                    stream << end << " // " << symbol
        #ifdef SHOW_USE_COUNT
                           << " used " << usage.at(this)
        #endif
                           << std::endl;
                }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    shared_leaf<T, SAFE_MATH> variable(const size_t s,
                                       const std::string &symbol) {
        return std::make_shared<variable_node<T, SAFE_MATH>> (s, symbol);
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    shared_leaf<T, SAFE_MATH> variable(const size_t s, const T d,
                                       const std::string &symbol) {
        return std::make_shared<variable_node<T, SAFE_MATH>> (s, d, symbol);
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    shared_leaf<T, SAFE_MATH> variable(const std::vector<T> &d,
                                       const std::string &symbol) {
        return std::make_shared<variable_node<T, SAFE_MATH>> (d, symbol);
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    shared_leaf<T, SAFE_MATH> variable(const backend::buffer<T> &d,
                                       const std::string &symbol) {
        return std::make_shared<variable_node<T, SAFE_MATH>> (d, symbol);
    }
 
    template<jit::float_scalar T, bool SAFE_MATH=false>
    using shared_variable = std::shared_ptr<variable_node<T, SAFE_MATH>>;
    template<jit::float_scalar T, bool SAFE_MATH=false>
    using input_nodes = std::vector<shared_variable<T, SAFE_MATH>>;
    template<jit::float_scalar T, bool SAFE_MATH=false>
    using map_nodes = std::vector<std::pair<shared_leaf<T, SAFE_MATH>,
                                            shared_variable<T, SAFE_MATH>>>;
    template<jit::float_scalar T, bool SAFE_MATH=false>
    using copy_nodes = std::vector<std::pair<shared_variable<T, SAFE_MATH>,
                                             shared_variable<T, SAFE_MATH>>>;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    shared_variable<T, SAFE_MATH> variable_cast(shared_leaf<T, SAFE_MATH> x) {
        return std::dynamic_pointer_cast<variable_node<T, SAFE_MATH>> (x);
    }
 
//******************************************************************************
//  Pseudo variable node.
//******************************************************************************
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class pseudo_variable_node final : public straight_node<T, SAFE_MATH> {
    private:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static std::string to_string(leaf_node<T, SAFE_MATH> *p) {
            return jit::format_to_string(reinterpret_cast<size_t> (p));
        }
 
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        pseudo_variable_node(shared_leaf<T, SAFE_MATH> a) :
        straight_node<T, SAFE_MATH> (a, pseudo_variable_node::to_string(a.get())) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> df(shared_leaf<T, SAFE_MATH> x) {
            return constant<T, SAFE_MATH> (static_cast<T> (this->is_match(x)));
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void to_latex() const {
            std::cout << "\\left(";
            this->arg->to_latex();
            std::cout << "\\right)";
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_all_variables() const {
            return true;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_power_like() const {
            return true;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> get_power_base() {
            return this->arg->get_power_base();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> get_power_exponent() const {
            return this->arg->get_power_exponent();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool has_pseudo() const {
            return true;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> remove_pseudo() {
            return this->arg;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T, SAFE_MATH> to_vizgraph(std::stringstream &stream,
                                                      jit::register_map &registers) {
            if (registers.find(this) == registers.end()) {
                const std::string name = jit::to_string('r', this);
                registers[this] = name;
                stream << "    " << name
                       << " [label = \"pseudo_variable\", shape = oval, style = filled, fillcolor = blue, fontcolor = white];" << std::endl;
 
                auto a = this->arg->to_vizgraph(stream, registers);
                stream << "    " << name << " -- " << registers[a.get()] << ";" << std::endl;
            }
 
            return this->shared_from_this();
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    shared_leaf<T, SAFE_MATH> pseudo_variable(shared_leaf<T, SAFE_MATH> x) {
        return std::make_shared<pseudo_variable_node<T, SAFE_MATH>> (x);
    }
 
    template<jit::float_scalar T, bool SAFE_MATH=false>
    using shared_pseudo_variable = std::shared_ptr<pseudo_variable_node<T, SAFE_MATH>>;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    shared_pseudo_variable<T, SAFE_MATH> pseudo_variable_cast(shared_leaf<T, SAFE_MATH> &x) {
        return std::dynamic_pointer_cast<pseudo_variable_node<T, SAFE_MATH>> (x);
    }
}
 
#endif /* node_h */