inlining.h

/*
 * SPDX-FileCopyrightText: 2013 Barefoot Networks, Inc.
 * Copyright 2013-present Barefoot Networks, Inc.
 *
 * SPDX-License-Identifier: Apache-2.0
 */
 
#ifndef FRONTENDS_P4_INLINING_H_
#define FRONTENDS_P4_INLINING_H_
 
#include "commonInlining.h"
#include "frontends/common/parser_options.h"
#include "frontends/common/resolveReferences/referenceMap.h"
#include "frontends/p4/evaluator/evaluator.h"
#include "frontends/p4/evaluator/substituteParameters.h"
#include "frontends/p4/typeChecking/typeChecker.h"
#include "frontends/p4/unusedDeclarations.h"
#include "ir/ir.h"
#include "lib/ordered_map.h"
 
// These are various data structures needed by the parser/parser and control/control inliners.
// This only works correctly after local variable initializers have been removed,
// and after the SideEffectOrdering pass has been executed.
 
namespace P4 {

struct CallInfo : public IHasDbPrint {
    const IR::IContainer *caller;
    const IR::IContainer *callee;
    const IR::Declaration_Instance *instantiation;  // callee instantiation
    // Each instantiation may be invoked multiple times.
    std::set<const IR::MethodCallStatement *> invocations;  // all invocations within the caller
 
    CallInfo(const IR::IContainer *caller, const IR::IContainer *callee,
             const IR::Declaration_Instance *instantiation)
        : caller(caller), callee(callee), instantiation(instantiation) {
        CHECK_NULL(caller);
        CHECK_NULL(callee);
        CHECK_NULL(instantiation);
    }
    void addInvocation(const IR::MethodCallStatement *statement) { invocations.emplace(statement); }
    void dbprint(std::ostream &out) const {
        out << "Inline " << callee << " into " << caller << " with " << invocations.size()
            << " invocations";
    }
};
 
class SymRenameMap {
    std::map<const IR::IDeclaration *, cstring> internalName;
    std::map<const IR::IDeclaration *, cstring> externalName;
 
 public:
    void setNewName(const IR::IDeclaration *decl, cstring name, cstring extName) {
        CHECK_NULL(decl);
        BUG_CHECK(!name.isNullOrEmpty() && !extName.isNullOrEmpty(), "Empty name");
        LOG3("setNewName " << dbp(decl) << " to " << name);
        if (internalName.find(decl) != internalName.end()) BUG("%1%: already renamed", decl);
        internalName.emplace(decl, name);
        externalName.emplace(decl, extName);
    }
    cstring getName(const IR::IDeclaration *decl) const {
        CHECK_NULL(decl);
        BUG_CHECK(internalName.find(decl) != internalName.end(), "%1%: no new name", decl);
        auto result = ::P4::get(internalName, decl);
        return result;
    }
    cstring getExtName(const IR::IDeclaration *decl) const {
        CHECK_NULL(decl);
        BUG_CHECK(externalName.find(decl) != externalName.end(), "%1%: no external name", decl);
        auto result = ::P4::get(externalName, decl);
        return result;
    }
    bool isRenamed(const IR::IDeclaration *decl) const {
        CHECK_NULL(decl);
        return internalName.find(decl) != internalName.end();
    }
};
 
struct PerInstanceSubstitutions {
    ParameterSubstitution paramSubst;
    TypeVariableSubstitution tvs;
    SymRenameMap renameMap;
    PerInstanceSubstitutions() = default;
    PerInstanceSubstitutions(const PerInstanceSubstitutions &other)
        : paramSubst(other.paramSubst), tvs(other.tvs), renameMap(other.renameMap) {}
    template <class T>
    const T *rename(ReferenceMap *refMap, const IR::Node *node);
};

struct InlineSummary : public IHasDbPrint {
    struct PerCaller {  // information for each caller
        typedef std::pair<const IR::MethodCallStatement *, const IR::PathExpression *>
            InlinedInvocationInfo;

        struct key_hash {
            std::size_t operator()(const InlinedInvocationInfo &k) const {
                std::ostringstream oss;
                std::get<0>(k)->dbprint(oss);
                std::get<1>(k)->dbprint(oss);
                return std::hash<std::string>{}(oss.str());
            }
        };

        struct key_equal {
            bool operator()(const InlinedInvocationInfo &v0,
                            const InlinedInvocationInfo &v1) const {
                return std::get<0>(v0)->equiv(*std::get<0>(v1)) &&
                       std::get<1>(v0)->equiv(*std::get<1>(v1));
            }
        };

        std::map<const IR::Declaration_Instance *, const IR::IContainer *> declToCallee;
        std::map<const IR::Declaration_Instance *, PerInstanceSubstitutions *> substitutions;
        std::map<const IR::MethodCallStatement *, const IR::Declaration_Instance *> callToInstance;

        std::unordered_map<const InlinedInvocationInfo, const IR::ID, key_hash, key_equal>
            invocationToState;

        const IR::MethodCallStatement *uniqueCaller(
            const IR::Declaration_Instance *instance) const {
            const IR::MethodCallStatement *call = nullptr;
            for (auto m : callToInstance) {
                if (m.second == instance) {
                    if (call == nullptr)
                        call = m.first;
                    else
                        return nullptr;
                }
            }
            return call;
        }
    };
    std::map<const IR::IContainer *, PerCaller> callerToWork;
 
    void add(const CallInfo *cci) {
        callerToWork[cci->caller].declToCallee[cci->instantiation] = cci->callee;
        for (auto mcs : cci->invocations)
            callerToWork[cci->caller].callToInstance[mcs] = cci->instantiation;
    }
    void dbprint(std::ostream &out) const {
        out << "Inline " << callerToWork.size() << " call sites";
    }
};
 
// Inling information constructed here.
class InlineList {
    // We use an ordered map to make the iterator deterministic
    ordered_map<const IR::Declaration_Instance *, CallInfo *> inlineMap;
    std::vector<CallInfo *> toInline;  // sorted in order of inlining
    const bool allowMultipleCalls = true;
 
 public:
    void addInstantiation(const IR::IContainer *caller, const IR::IContainer *callee,
                          const IR::Declaration_Instance *instantiation) {
        CHECK_NULL(caller);
        CHECK_NULL(callee);
        CHECK_NULL(instantiation);
        LOG3("Inline instantiation " << dbp(instantiation));
        auto inst = new CallInfo(caller, callee, instantiation);
        inlineMap[instantiation] = inst;
    }
 
    size_t size() const { return inlineMap.size(); }
 
    void addInvocation(const IR::Declaration_Instance *instance,
                       const IR::MethodCallStatement *statement) {
        CHECK_NULL(instance);
        CHECK_NULL(statement);
        LOG3("Inline invocation " << dbp(instance) << " at " << dbp(statement));
        auto info = inlineMap[instance];
        BUG_CHECK(info, "Could not locate instance %1% invoked by %2%", instance, statement);
        info->addInvocation(statement);
    }
 
    void replace(const IR::IContainer *container, const IR::IContainer *replacement) {
        CHECK_NULL(container);
        CHECK_NULL(replacement);
        LOG3("Replacing " << dbp(container) << " with " << dbp(replacement));
        for (auto e : toInline) {
            if (e->callee == container) e->callee = replacement;
            if (e->caller == container) e->caller = replacement;
        }
    }
 
    void analyze();
    InlineSummary *next();
};

class DiscoverInlining : public Inspector {
    InlineList *inlineList;  // output: result is here
    ReferenceMap *refMap;    // input
    TypeMap *typeMap;        // input
    IHasBlock *evaluator;    // used to obtain the toplevel block
    IR::ToplevelBlock *toplevel;
 
 public:
    bool allowParsers = true;
    bool allowControls = true;
 
    DiscoverInlining(InlineList *inlineList, ReferenceMap *refMap, TypeMap *typeMap,
                     IHasBlock *evaluator)
        : inlineList(inlineList),
          refMap(refMap),
          typeMap(typeMap),
          evaluator(evaluator),
          toplevel(nullptr) {
        CHECK_NULL(inlineList);
        CHECK_NULL(refMap);
        CHECK_NULL(typeMap);
        CHECK_NULL(evaluator);
        setName("DiscoverInlining");
        visitDagOnce = false;
    }
    Visitor::profile_t init_apply(const IR::Node *node) override {
        toplevel = evaluator->getToplevelBlock();
        CHECK_NULL(toplevel);
        return Inspector::init_apply(node);
    }
    void visit_all(const IR::Block *block);
    bool preorder(const IR::Block *block) override {
        visit_all(block);
        return false;
    }
    bool preorder(const IR::ControlBlock *block) override;
    bool preorder(const IR::ParserBlock *block) override;
    void postorder(const IR::MethodCallStatement *statement) override;
    // We don't care to visit the program, we just visit the blocks.
    bool preorder(const IR::P4Program *) override {
        visit_all(toplevel);
        return false;
    }
};

class GeneralInliner : public AbstractInliner<InlineList, InlineSummary> {
    ReferenceMap *refMap;
    TypeMap *typeMap;
    InlineSummary::PerCaller *workToDo;
    bool optimizeParserInlining;
 
 public:
    explicit GeneralInliner(ReferenceMap *refMap, bool _optimizeParserInlining)
        : refMap(refMap),
          typeMap(new TypeMap()),
          workToDo(nullptr),
          optimizeParserInlining(_optimizeParserInlining) {
        setName("GeneralInliner");
        visitDagOnce = false;
    }
    // controlled visiting order
    const IR::Node *preorder(IR::MethodCallStatement *statement) override;
    template <class P4Block, class P4BlockType>
    void inline_subst(P4Block *caller, IR::IndexedVector<IR::Declaration> P4Block::*blockLocals,
                      const P4BlockType *P4Block::*blockType);
    const IR::Node *preorder(IR::P4Control *caller) override;
    const IR::Node *preorder(IR::P4Parser *caller) override;
    const IR::Node *preorder(IR::ParserState *state) override;
    Visitor::profile_t init_apply(const IR::Node *node) override;
};

class InlinePass : public PassManager {
    InlineList toInline;
 
 public:
    InlinePass(ReferenceMap *refMap, TypeMap *typeMap, EvaluatorPass *evaluator,
               const RemoveUnusedPolicy &policy, bool optimizeParserInlining)
        : PassManager({new TypeChecking(refMap, typeMap),
                       new DiscoverInlining(&toInline, refMap, typeMap, evaluator),
                       new InlineDriver<InlineList, InlineSummary>(
                           &toInline, new GeneralInliner(refMap, optimizeParserInlining)),
                       new RemoveAllUnusedDeclarations(policy)}) {
        setName("InlinePass");
    }
};

class Inline : public PassManager {
    static std::set<cstring> noPropagateAnnotations;
    ReferenceMap refMap;
 
 public:
    Inline(TypeMap *typeMap, const RemoveUnusedPolicy &policy, bool optimizeParserInlining,
           EvaluatorPass *evaluator = nullptr) {
        refMap.setIsV1(P4CContext::get().options().isv1());
        auto *evInstance = evaluator ? evaluator : new EvaluatorPass(&refMap, typeMap);
        addPasses({
            evInstance,
            new PassRepeated(
                {new InlinePass(&refMap, typeMap, evInstance, policy, optimizeParserInlining),
                 // After inlining the output of the evaluator changes, so we have to
                 // run it again
                 evInstance}),
        });
        setName("Inline");
    }

    static void setAnnotationNoPropagate(cstring name) { noPropagateAnnotations.emplace(name); }

    static bool isAnnotationNoPropagate(cstring name) { return noPropagateAnnotations.count(name); }
};
 
}  // namespace P4
 
#endif /* FRONTENDS_P4_INLINING_H_ */