specializeGenericFunctions.h

/*
 * Copyright 2020 VMware, Inc.
 * SPDX-FileCopyrightText: 2020 VMware, Inc.
 *
 * SPDX-License-Identifier: Apache-2.0
 */
 
#ifndef FRONTENDS_P4_SPECIALIZEGENERICFUNCTIONS_H_
#define FRONTENDS_P4_SPECIALIZEGENERICFUNCTIONS_H_
 
#include "frontends/common/resolveReferences/referenceMap.h"
#include "frontends/common/resolveReferences/resolveReferences.h"
#include "frontends/p4/typeChecking/typeChecker.h"
#include "ir/ir.h"
 
namespace P4 {

struct FunctionSpecialization {
    cstring name;
    const IR::Function *original;
    const IR::Function *specialized;
    const IR::MethodCallExpression *invocation;
    const IR::Node *insertBefore;
 
    FunctionSpecialization(cstring name, const IR::MethodCallExpression *invocation,
                           const IR::Function *function, const IR::Node *insert)
        : name(name),
          original(function),
          specialized(nullptr),
          invocation(invocation),
          insertBefore(insert) {
        CHECK_NULL(invocation);
        CHECK_NULL(invocation);
        CHECK_NULL(insertBefore);
    }
};
 
struct FunctionSpecializationMap {
    TypeMap *typeMap;
    ordered_map<const IR::MethodCallExpression *, FunctionSpecialization *> map;
    // Keep track of the values in the above map which are already
    // inserted in the program.
    std::set<FunctionSpecialization *> inserted;
 
    void add(const IR::MethodCallExpression *mce, const IR::Function *func, const IR::Node *insert,
             NameGenerator *nameGen) {
        cstring name = nameGen->newName(func->name.string_view());
        map.emplace(mce, new FunctionSpecialization(name, mce, func, insert));
    }
    FunctionSpecialization *get(const IR::MethodCallExpression *mce) const {
        return ::P4::get(map, mce);
    }
    IR::Vector<IR::Node> *getInsertions(const IR::Node *insertionPoint) {
        IR::Vector<IR::Node> *result = nullptr;
        for (auto s : map) {
            if (inserted.find(s.second) != inserted.end()) continue;
            if (s.second->insertBefore == insertionPoint) {
                if (result == nullptr) result = new IR::Vector<IR::Node>();
                LOG2("Will insert " << dbp(s.second->specialized) << " before "
                                    << dbp(insertionPoint));
                result->push_back(s.second->specialized);
                inserted.emplace(s.second);
            }
        }
        return result;
    }
};

class FindFunctionSpecializations : public Inspector, public ResolutionContext {
    FunctionSpecializationMap *specMap;
    MinimalNameGenerator nameGen;
 
 public:
    explicit FindFunctionSpecializations(FunctionSpecializationMap *specMap) : specMap(specMap) {
        CHECK_NULL(specMap);
        setName("FindFunctionSpecializations");
    }
 
    bool preorder(const IR::MethodCallExpression *call) override;
    profile_t init_apply(const IR::Node *node) override;
};

class SpecializeFunctions : public Transform {
    FunctionSpecializationMap *specMap;
 
 public:
    explicit SpecializeFunctions(FunctionSpecializationMap *specMap) : specMap(specMap) {
        CHECK_NULL(specMap);
        setName("SpecializeFunctions");
    }
    const IR::Node *postorder(IR::Function *function) override;
    const IR::Node *postorder(IR::MethodCallExpression *) override;
    const IR::Node *insert(const IR::Node *before);
    const IR::Node *preorder(IR::P4Control *control) override { return insert(control); }
    const IR::Node *preorder(IR::P4Parser *parser) override { return insert(parser); }
    const IR::Node *preorder(IR::Function *function) override { return insert(function); }
};
 
class SpecializeGenericFunctions : public PassManager {
    FunctionSpecializationMap specMap;
 
 public:
    explicit SpecializeGenericFunctions(TypeMap *typeMap) {
        passes.emplace_back(new TypeChecking(nullptr, typeMap));
        passes.emplace_back(new FindFunctionSpecializations(&specMap));
        passes.emplace_back(new SpecializeFunctions(&specMap));
        specMap.typeMap = typeMap;
        setName("SpecializeGenericFunctions");
    }
};
 
}  // namespace P4
 
#endif /* FRONTENDS_P4_SPECIALIZEGENERICFUNCTIONS_H_ */