switch_expander.h

Go to the documentation of this file.

#pragma once
 
#include <memory>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
 
#include "sdfg/data_flow/library_node.h"
#include "sdfg/passes/expansion/lib_node_expander.h"
 
namespace sdfg::passes {

class SwitchExpander : public LibNodeExpander {
    std::vector<std::shared_ptr<const LibNodeExpander>> expanders_;
    std::unordered_map<std::string, const LibNodeExpander*> by_code_;
    std::vector<const LibNodeExpander*> fallback_;
 
public:
    template<typename... Expanders>
    explicit SwitchExpander(Expanders... expanders) {
        expanders_.reserve(sizeof...(Expanders));
        (register_expander(std::move(expanders)), ...);
    }

    explicit SwitchExpander(std::vector<std::shared_ptr<const LibNodeExpander>> expanders) {
        expanders_.reserve(expanders.size());
        for (auto& expander : expanders) {
            register_expander(std::move(expander));
        }
    }

    void register_expander(std::shared_ptr<const LibNodeExpander> expander) {
        if (!expander) {
            return;
        }
        if (const auto* code = expander->matched_code()) {
            register_for_code(*code, std::move(expander));
        } else {
            fallback_.push_back(expander.get());
            expanders_.push_back(std::move(expander));
        }
    }

    void register_for_code(const data_flow::LibraryNodeCode& code, std::shared_ptr<const LibNodeExpander> expander) {
        if (!expander) {
            return;
        }
        auto [it, inserted] = by_code_.emplace(code.value(), expander.get());
        if (!inserted) {
            throw std::invalid_argument(
                "SwitchExpander: more than one expander registered for library node code '" + code.value() + "'"
            );
        }
        expanders_.push_back(std::move(expander));
    }

    const LibNodeExpander* for_lib_node(const data_flow::LibraryNode& node) const override {
        if (auto it = by_code_.find(node.code().value()); it != by_code_.end()) {
            if (const auto* match = it->second->for_lib_node(node)) {
                return match;
            }
        }
        for (const auto* expander : fallback_) {
            if (const auto* match = expander->for_lib_node(node)) {
                return match;
            }
        }
        return nullptr;
    }

    ExpandOutcome handle_expand(
        ExpandContext& context, structured_control_flow::Block& block, data_flow::LibraryNode& node
    ) const override {
        if (const auto* expander = for_lib_node(node); expander && expander != this) {
            return expander->handle_expand(context, block, node);
        }
        return context.unapplicable();
    }
};
 
} // namespace sdfg::passes