flattenInterfaceStructs.h

/*
Copyright 2018 VMware, Inc.
 
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
 
    http://www.apache.org/licenses/LICENSE-2.0
 
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
 
#ifndef MIDEND_FLATTENINTERFACESTRUCTS_H_
#define MIDEND_FLATTENINTERFACESTRUCTS_H_
 
#include "frontends/p4/typeChecking/typeChecker.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) {
        auto vec = new IR::IndexedVector<IR::StructField>();
        flatten(typeMap, cstring::empty, type, type->annotations, vec, policy);
        if (type->is<IR::Type_Struct>()) {
            replacementType =
                new IR::Type_Struct(type->srcInfo, type->name, IR::Annotations::empty, *vec);
        } else if (type->is<IR::Type_Header>()) {
            replacementType =
                new IR::Type_Header(type->srcInfo, type->name, IR::Annotations::empty, *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::Annotations *annotations, IR::IndexedVector<IR::StructField> *fields,
                 AnnotationSelectionPolicy *policy) {
        // Drop name annotations
        annotations = annotations->where(
            [](const IR::Annotation *a) { return a->name != IR::Annotation::nameAnnotation; });
        if (auto st = type->to<T>()) {
            auto sannotations = 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) {
                auto na = new IR::Annotations();
                na->append(sannotations);
                na->append(annotations);
                na->append(f->annotations);
                auto ft = typeMap->getType(f, true);
                flatten(typeMap, prefix + "." + f->name, ft, na, 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), annotations, type->getP4Type()));
        LOG3("Flatten: " << type << " | " << prefix);
    }

    const IR::StructExpression *explode(const IR::Expression *root, cstring prefix) {
        auto vec = new IR::IndexedVector<IR::NamedExpression>();
        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, *vec);
    }
};

struct NestedStructMap {
    P4::ReferenceMap *refMap;
    P4::TypeMap *typeMap;
 
    ordered_map<const IR::Type *, StructTypeReplacement<IR::Type_Struct> *> replacement;
 
    NestedStructMap(P4::ReferenceMap *refMap, P4::TypeMap *typeMap)
        : refMap(refMap), typeMap(typeMap) {
        CHECK_NULL(refMap);
        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 {
    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 {
 public:
    FlattenInterfaceStructs(ReferenceMap *refMap, TypeMap *typeMap) {
        auto sm = new NestedStructMap(refMap, typeMap);
        passes.push_back(new TypeChecking(refMap, 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_ */