Preserve annotations on functions properly

- when dumping P4 code - when cloning functions in removeReturns and specialization
p4lang · Mar 14, 2024 · 00f8c94 · 00f8c94
1 parent 98d166f
commit 00f8c94
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 2 deletions.
diff --git a/frontends/p4/removeReturns.cpp b/frontends/p4/removeReturns.cpp
@@ -85,7 +85,8 @@ const IR::Node *DoRemoveReturns::preorder(IR::Function *function) {
     if (retvalDecl != nullptr) body->push_back(retvalDecl);
     body->components.append(function->body->components);
     if (returnsVal) body->push_back(new IR::ReturnStatement(new IR::PathExpression(returnedValue)));
-    auto result = new IR::Function(function->srcInfo, function->name, function->type, body);
+    auto result = new IR::Function(function->srcInfo, function->name, function->annotations,
+                                   function->type, body);
     pop();
     BUG_CHECK(stack.empty(), "Non-empty stack");
     prune();

diff --git a/frontends/p4/specializeGenericFunctions.cpp b/frontends/p4/specializeGenericFunctions.cpp
@@ -60,7 +60,8 @@ const IR::Node *SpecializeFunctions::postorder(IR::Function *function) {
             LOG3("Substitution " << ts);
             auto specialized = function->apply(tsv)->to<IR::Function>();
             auto renamed = new IR::Function(specialized->srcInfo, it.second->name,
-                                            specialized->type, specialized->body);
+                                            specialized->annotations, specialized->type,
+                                            specialized->body);
             it.second->specialized = renamed;
             LOG3("Specializing " << function << " as " << renamed);
         }

diff --git a/frontends/p4/toP4/toP4.cpp b/frontends/p4/toP4/toP4.cpp
@@ -430,6 +430,10 @@ bool ToP4::preorder(const IR::Method *m) {
 
 bool ToP4::preorder(const IR::Function *function) {
     dump(1);
+    if (!function->annotations->annotations.empty()) {
+        visit(function->annotations);
+        builder.spc();
+    }
     auto t = function->type;
     BUG_CHECK(t != nullptr, "Function %1% has no type", function);
     if (t->returnType != nullptr) {