flattenUnions.h

/*
 * Copyright 2022 Intel Corp.
 * SPDX-FileCopyrightText: 2022 Intel Corp.
 *
 * SPDX-License-Identifier: Apache-2.0
 */
 
#ifndef MIDEND_FLATTENUNIONS_H_
#define MIDEND_FLATTENUNIONS_H_
 
#include "./frontends/p4/parserControlFlow.h"
#include "./frontends/p4/simplifyDefUse.h"
#include "./frontends/p4/unusedDeclarations.h"
#include "frontends/p4/typeChecking/typeChecker.h"
#include "ir/ir.h"
namespace P4 {

class DoFlattenHeaderUnion : public Transform {
 protected:
    P4::ReferenceMap *refMap;
    P4::TypeMap *typeMap;
    std::map<cstring, std::map<cstring, cstring>> replacementMap;
    // Replacement map needed to add element-wise header declaration in right context
    std::map<IR::Declaration_Variable *, IR::IndexedVector<IR::Declaration>> replaceDVMap;
 
 public:
    DoFlattenHeaderUnion(P4::ReferenceMap *refMap, P4::TypeMap *typeMap)
        : refMap(refMap), typeMap(typeMap) {}
    const IR::Node *postorder(IR::Type_Struct *sf) override;
    const IR::Node *postorder(IR::Declaration_Variable *dv) override;
    const IR::Node *postorder(IR::Member *m) override;
    const IR::Node *postorder(IR::P4Parser *parser) override;
    const IR::Node *postorder(IR::P4Control *control) override;
    const IR::Node *postorder(IR::P4Action *action) override;
    bool hasHeaderUnionField(IR::Type_Struct *s);
};

class DoFlattenHeaderUnionStack : public DoFlattenHeaderUnion {
    std::map<cstring, std::vector<cstring>> stackMap;
 
 public:
    DoFlattenHeaderUnionStack(P4::ReferenceMap *refMap, P4::TypeMap *typeMap)
        : DoFlattenHeaderUnion(refMap, typeMap) {
        setName("DoFlattenHeaderUnionStack");
    }
    const IR::Node *postorder(IR::Type_Struct *sf) override;
    const IR::Node *postorder(IR::ArrayIndex *e) override;
    const IR::Node *postorder(IR::Declaration_Variable *dv) override;
    bool hasHeaderUnionStackField(IR::Type_Struct *s);
};

class HandleValidityHeaderUnion : public Transform {
    P4::ReferenceMap *refMap;
    P4::TypeMap *typeMap;
    IR::IndexedVector<IR::Declaration> toInsert;  // temporaries
 
 public:
    HandleValidityHeaderUnion(P4::ReferenceMap *refMap, P4::TypeMap *typeMap)
        : refMap(refMap), typeMap(typeMap) {
        setName("HandleValidityHeaderUnion");
    }
    const IR::Node *postorder(IR::AssignmentStatement *assn) override;
    const IR::Node *postorder(IR::IfStatement *a) override;
    const IR::Node *postorder(IR::SwitchStatement *a) override;
    const IR::Node *postorder(IR::MethodCallStatement *mcs) override;
    const IR::Node *postorder(IR::P4Parser *parser) override;
    const IR::Node *postorder(IR::P4Control *control) override;
    const IR::Node *postorder(IR::P4Action *action) override;
    const IR::MethodCallStatement *processValidityForStr(const IR::Statement *s,
                                                         const IR::Member *m, cstring headerElement,
                                                         cstring setValid);
    const IR::Node *setInvalidforRest(const IR::Statement *s, const IR::Member *m,
                                      const IR::Type_HeaderUnion *hu, cstring exclude,
                                      bool setValidforCurrMem);
    const IR::Node *expandIsValid(const IR::Statement *a, const IR::MethodCallExpression *mce,
                                  IR::IndexedVector<IR::StatOrDecl> &code_block);
};

class DoExpandHeaderUnionStackPushPop : public Transform {
    P4::TypeMap *typeMap;
 
 public:
    explicit DoExpandHeaderUnionStackPushPop(P4::TypeMap *typeMap) : typeMap(typeMap) {
        setName("DoExpandHeaderUnionStackPushPop");
    }
    const IR::Node *postorder(IR::MethodCallStatement *mcs) override;
};
 
class RemoveUnusedHUDeclarations : public Transform {
    const UsedDeclSet &used;
 
 public:
    explicit RemoveUnusedHUDeclarations(const UsedDeclSet &used) : used(used) {}
 
    const IR::Node *preorder(IR::Type_HeaderUnion *type) override {
        if (!used.isUsed(getOriginal<IR::IDeclaration>())) return nullptr;
        return type;
    }
};
 
class RemoveAllUnusedHUDDeclarations : public PassManager {
    UsedDeclSet used;
 
 public:
    RemoveAllUnusedHUDDeclarations()
        : PassManager({
              new CollectUsedDeclarations(used),
              new RemoveUnusedHUDeclarations(used),
          }) {
        setName("RemoveAllUnusedHUDDeclarations");
        setStopOnError(true);
    }
};

class FlattenHeaderUnion : public PassManager {
 public:
    FlattenHeaderUnion(P4::ReferenceMap *refMap, P4::TypeMap *typeMap, bool loopsUnroll = true) {
        passes.push_back(new P4::TypeChecking(refMap, typeMap));
        // Lower push_front/pop_front on header union stacks first; the new nodes need
        // types before HandleValidityHeaderUnion.
        if (loopsUnroll) {
            passes.push_back(new DoExpandHeaderUnionStackPushPop(typeMap));
            passes.push_back(new P4::ClearTypeMap(typeMap));
            passes.push_back(new P4::TypeInference(typeMap, false));
            passes.push_back(new P4::TypeChecking(refMap, typeMap));
        }
        passes.push_back(new HandleValidityHeaderUnion(refMap, typeMap));
        // Stack flattening is only applicable if parser loops are unrolled and
        // header union stack elements are accessed using [] notation. This pass does not handle
        // .next .last etc accessors for stack elements.
        if (loopsUnroll) {
            passes.push_back(new DoFlattenHeaderUnionStack(refMap, typeMap));
            passes.push_back(new P4::RemoveAllUnusedDeclarations(RemoveUnusedPolicy()));
        }
        passes.push_back(new P4::ClearTypeMap(typeMap));
        passes.push_back(new P4::TypeInference(typeMap, false));
        passes.push_back(new P4::TypeChecking(refMap, typeMap));
        passes.push_back(new DoFlattenHeaderUnion(refMap, typeMap));
        passes.push_back(new P4::RemoveAllUnusedDeclarations(RemoveUnusedPolicy()));
        passes.push_back(new P4::RemoveAllUnusedHUDDeclarations());
        passes.push_back(new P4::ClearTypeMap(typeMap));
        passes.push_back(new P4::TypeChecking(nullptr, typeMap));
        passes.push_back(new P4::RemoveParserIfs(typeMap));
    }
};
}  // namespace P4
 
#endif /* MIDEND_FLATTENUNIONS_H_ */