When flattening a struct expression, use the underlying type, not the…

… components.

backends/p4tools/common/lib/gen_eq.cpp

@@ -24,8 +24,8 @@ const IR::Expression *GenEq::resolveSingletonList(const IR::Expression *expr) {
 
 const IR::Expression *GenEq::equateListTypes(const IR::Expression *left,
                                              const IR::Expression *right) {
-    std::vector<const IR::Expression *> leftElems = flattenListOrStructExpression(left);
-    std::vector<const IR::Expression *> rightElems = flattenListOrStructExpression(right);
+    std::vector<const IR::Expression *> leftElems = IR::flattenListOrStructExpression(left);
+    std::vector<const IR::Expression *> rightElems = IR::flattenListOrStructExpression(right);
 
     auto leftElemsSize = leftElems.size();
     auto rightElemsSize = rightElems.size();

backends/p4tools/common/lib/gen_eq.h

@@ -5,7 +5,7 @@
 
 namespace P4Tools {
 
-/// Generates an equality on two input expressions, recursing into lists and structs.
+/// Generates an semantic equality on two input expressions, recursing into lists and structs.
 /// This supports fuzzy matching on singleton lists: singleton lists are considered the same as
 /// their singleton elements. This is implemented by eagerly recursing into singleton lists before
 /// attempting to generate the equality.

ir/irutils.cpp

@@ -1,6 +1,5 @@
 #include "ir/irutils.h"
 
-#include <algorithm>
 #include <cmath>
 #include <map>
 #include <tuple>
@@ -176,18 +175,21 @@ const IR::Expression *getDefaultValue(const IR::Type *type, const Util::SourceIn
 
 std::vector<const Expression *> flattenStructExpression(const StructExpression *structExpr) {
     std::vector<const Expression *> exprList;
-    for (const auto *listElem : structExpr->components) {
+    // Ensure that the underlying type is a Type_StructLike.
+    // TODO: How do fail gracefully if we get a Type_Name?
+    const auto *structType = structExpr->type->to<IR::Type_StructLike>();
+    BUG_CHECK(structType != nullptr, "%1%: expected a struct type, received %2%", structExpr->type,
+              structExpr->node_type_name());
+
+    // We use the underlying struct type, which will gives us the right field ordering.
+    for (const auto *typeField : structType->fields) {
+        const auto *listElem = structExpr->getField(typeField->name);
         if (const auto *subStructExpr = listElem->expression->to<StructExpression>()) {
             auto subList = flattenStructExpression(subStructExpr);
             exprList.insert(exprList.end(), subList.begin(), subList.end());
-        } else if (const auto *headerStackExpr =
-                       listElem->expression->to<HeaderStackExpression>()) {
-            for (const auto *headerStackElem : headerStackExpr->components) {
-                // We assume there are no nested header stacks.
-                auto subList =
-                    flattenStructExpression(headerStackElem->checkedTo<IR::StructExpression>());
-                exprList.insert(exprList.end(), subList.begin(), subList.end());
-            }
+        } else if (const auto *subListExpr = listElem->to<BaseListExpression>()) {
+            auto subList = flattenListExpression(subListExpr);
+            exprList.insert(exprList.end(), subList.begin(), subList.end());
         } else {
             exprList.emplace_back(listElem->expression);
         }
@@ -198,9 +200,12 @@ std::vector<const Expression *> flattenStructExpression(const StructExpression *
 std::vector<const Expression *> flattenListExpression(const BaseListExpression *listExpr) {
     std::vector<const Expression *> exprList;
     for (const auto *listElem : listExpr->components) {
-        if (const auto *subListExpr = listElem->to<ListExpression>()) {
+        if (const auto *subListExpr = listElem->to<BaseListExpression>()) {
             auto subList = flattenListExpression(subListExpr);
             exprList.insert(exprList.end(), subList.begin(), subList.end());
+        } else if (const auto *subStructExpr = listElem->to<IR::StructExpression>()) {
+            auto subList = flattenStructExpression(subStructExpr);
+            exprList.insert(exprList.end(), subList.begin(), subList.end());
         } else {
             exprList.emplace_back(listElem);
         }

ir/irutils.h

@@ -74,12 +74,16 @@ const IR::Constant *convertBoolLiteral(const IR::BoolLiteral *lit);
 /// Given an StructExpression, returns a flat vector of the expressions contained in that
 /// struct. Unfortunately, list and struct expressions are similar but have no common ancestors.
 /// This is why we require two separate methods.
+/// Note that this function will fail if the type of @param structExpr is not a Type_Name.
 std::vector<const Expression *> flattenStructExpression(const StructExpression *structExpr);
 
 /// Given an BaseListExpression, returns a flat vector of the expressions contained in that
 /// list.
 std::vector<const Expression *> flattenListExpression(const BaseListExpression *listExpr);
 
+/// Given a StructExpression or BaseListExpression, returns a flat vector of the expressions
+/// contained in that list.
+/// Note that this function will fail if the type of any input struct expression is not a Type_Name.
 std::vector<const Expression *> flattenListOrStructExpression(const Expression *listLikeExpr);
 
 /* =========================================================================================