flattenInterfaceStructs.h

/*
 * Copyright 2018 VMware, Inc.
 * SPDX-FileCopyrightText: 2018 VMware, Inc.
 *
 * SPDX-License-Identifier: Apache-2.0
 */
 
#ifndef MIDEND_FLATTENINTERFACESTRUCTS_H_
#define MIDEND_FLATTENINTERFACESTRUCTS_H_
 
#include "frontends/common/resolveReferences/resolveReferences.h"
#include "frontends/p4/typeChecking/typeChecker.h"
#include "ir/annotations.h"
#include "ir/ir.h"
#include "lib/cstring.h"
 
namespace P4 {

class AnnotationSelectionPolicy {
 public:
    virtual ~AnnotationSelectionPolicy() {}

    virtual bool keep(const IR::Annotation *) { return true; }
};

template <typename T>
// T is the type of objects that will be replaced.  E.g., IR::Type_Struct
struct StructTypeReplacement : public IHasDbPrint {
    StructTypeReplacement(const P4::TypeMap *typeMap, const IR::Type_StructLike *type,
                          AnnotationSelectionPolicy *policy) {
        IR::IndexedVector<IR::StructField> vec;
        flatten(typeMap, cstring::empty, type, type->annotations, vec, policy);
        if (type->is<IR::Type_Struct>()) {
            replacementType = new IR::Type_Struct(type->srcInfo, type->name, std::move(vec));
        } else if (type->is<IR::Type_Header>()) {
            replacementType = new IR::Type_Header(type->srcInfo, type->name, std::move(vec));
        } else {
            BUG("Unexpected type %1%", type);
        }
    }
 
    // Maps nested field names to final field names.
    // In our example this could be:
    // .t.s.a -> _t_s_a0;
    // .t.s.b -> _t_s_b1;
    // .t.y -> _t_y2;
    // .x -> _x3;
    std::map<cstring, cstring> fieldNameRemap;
    // Maps internal fields names to types.
    // .t -> T
    // .t.s -> S
    std::map<cstring, const IR::Type_StructLike *> structFieldMap;
    // Holds a new flat type
    // struct M {
    //    bit _t_s_a0;
    //    bool _t_s_b1;
    //    bit<6> _t_y2;
    //    bit<3> _x3;
    // }
    const IR::Type *replacementType;
    virtual void dbprint(std::ostream &out) const { out << replacementType; }
 
    // Helper for constructor
    void flatten(const P4::TypeMap *typeMap, cstring prefix, const IR::Type *type,
                 const IR::Vector<IR::Annotation> &annotations,
                 IR::IndexedVector<IR::StructField> &fields, AnnotationSelectionPolicy *policy) {
        // Drop name annotations
        auto ann = IR::Annotations::withoutNameAnnotation(annotations);
        if (auto st = type->to<T>()) {
            IR::Vector<IR::Annotation> sann =
                st->annotations.where([policy](const IR::Annotation *annot) {
                    if (!policy) return false;
                    return policy->keep(annot);
                });
            structFieldMap.emplace(prefix, st);
            for (auto f : st->fields) {
                IR::Vector<IR::Annotation> fann(sann);
                fann.append(ann);
                fann.append(f->annotations);
                auto ft = typeMap->getType(f, true);
                flatten(typeMap, prefix + "." + f->name, ft, fann, fields, policy);
            }
            return;
        }
        cstring fieldName = prefix.replace('.', '_') + std::to_string(fieldNameRemap.size());
        fieldNameRemap.emplace(prefix, fieldName);
        fields.push_back(new IR::StructField(IR::ID(fieldName), ann, type->getP4Type()));
        LOG3("Flatten: " << type << " | " << prefix);
    }

    const IR::StructExpression *explode(const IR::Expression *root, cstring prefix) {
        IR::IndexedVector<IR::NamedExpression> vec;
        auto fieldType = ::P4::get(structFieldMap, prefix);
        BUG_CHECK(fieldType, "No field for %1%", prefix);
        for (auto f : fieldType->fields) {
            cstring fieldName = prefix + "." + f->name.name;
            auto newFieldname = ::P4::get(fieldNameRemap, fieldName);
            const IR::Expression *expr;
            if (!newFieldname.isNullOrEmpty()) {
                expr = new IR::Member(root, newFieldname);
            } else {
                expr = explode(root, fieldName);
            }
            vec.push_back(new IR::NamedExpression(f->name, expr));
        }
        auto type = fieldType->getP4Type()->template to<IR::Type_Name>();
        return new IR::StructExpression(root->srcInfo, type, type, std::move(vec));
    }
};

struct NestedStructMap {
    P4::TypeMap *typeMap;
 
    ordered_map<const IR::Type *, StructTypeReplacement<IR::Type_Struct> *> replacement;
 
    explicit NestedStructMap(P4::TypeMap *typeMap) : typeMap(typeMap) { CHECK_NULL(typeMap); }
    void createReplacement(const IR::Type_Struct *type);
    StructTypeReplacement<IR::Type_Struct> *getReplacement(const IR::Type *type) const {
        return ::P4::get(replacement, type);
    }
    bool empty() const { return replacement.empty(); }
};

class FindTypesToReplace : public Inspector {
    NestedStructMap *map;
 
 public:
    explicit FindTypesToReplace(NestedStructMap *map) : map(map) {
        setName("FindTypesToReplace");
        CHECK_NULL(map);
    }
    bool preorder(const IR::Declaration_Instance *inst) override;
};

class ReplaceStructs : public Transform, P4WriteContext, ResolutionContext {
    NestedStructMap *replacementMap;
    std::map<const IR::Parameter *, StructTypeReplacement<IR::Type_Struct> *> toReplace;
 
 public:
    explicit ReplaceStructs(NestedStructMap *sm) : replacementMap(sm) {
        CHECK_NULL(sm);
        setName("ReplaceStructs");
    }
 
    const IR::Node *preorder(IR::P4Program *program) override;
    const IR::Node *postorder(IR::Member *expression) override;
    const IR::Node *preorder(IR::P4Parser *parser) override;
    const IR::Node *preorder(IR::P4Control *control) override;
    const IR::Node *postorder(IR::Type_Struct *type) override;
};
 
class FlattenInterfaceStructs final : public PassManager {
    NestedStructMap sm;
 
 public:
    explicit FlattenInterfaceStructs(TypeMap *typeMap) : sm(typeMap) {
        passes.push_back(new TypeChecking(nullptr, typeMap));
        passes.push_back(new FindTypesToReplace(&sm));
        passes.push_back(new ReplaceStructs(&sm));
        passes.push_back(new ClearTypeMap(typeMap));
        setName("FlattenInterfaceStructs");
    }
};
 
}  // namespace P4
 
#endif /* MIDEND_FLATTENINTERFACESTRUCTS_H_ */