Fix annotations

Author: vaalfreja

include/p4mlir/Dialect/P4HIR/P4HIR_Ops.td

@@ -540,7 +540,8 @@ def CmpOp : P4HIR_Op<"cmp",
 def ScopeOp : P4HIR_Op<"scope", [
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        RecursivelySpeculatable,
-       RecursiveMemoryEffects, AutomaticAllocationScope,
+       DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+       AutomaticAllocationScope,
        NoRegionArguments, Annotated]> {
   let summary = "Represents a P4 scope";
   let description = [{

lib/Dialect/P4HIR/P4HIR_Ops.cpp

@@ -926,6 +926,24 @@ LogicalResult P4HIR::ScopeOp::verify() {
     return success();
 }
 
+void P4HIR::ScopeOp::getEffects(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>> &effects) {
+    if (getAnnotations()) {
+        effects.emplace_back(MemoryEffects::Write::get(), ExternResource::get());
+        return;
+    }
+
+    for (Region &region : getOperation()->getRegions()) {
+        for (Block &block : region) {
+            for (Operation &op : block) {
+                if (auto opEffects = mlir::getEffectsRecursively(&op)) {
+                    effects.append(opEffects->begin(), opEffects->end());
+                }
+            }
+        }
+    }
+}
+
 LogicalResult P4HIR::ScopeOp::canonicalize(P4HIR::ScopeOp op, PatternRewriter &rewriter) {
     // Canonicalize scope: one without variables could be inlined
     if (op.getOps<VariableOp>().empty() && op.getScopeRegion().hasOneBlock() &&

lib/Transforms/FlattenCFG.cpp

@@ -86,9 +86,7 @@ class ScopeOpFlattening : public mlir::OpRewritePattern<P4HIR::ScopeOp> {
         mlir::OpBuilder::InsertionGuard guard(rewriter);
         auto loc = scopeOp.getLoc();
 
-        // Empty scope: just remove it.
-        // TODO: Decide if we'd need to do something with annotated scopes
-        if (scopeOp.isEmpty()) {
+        if (scopeOp.isEmpty() && !scopeOp.getAnnotations()) {
             rewriter.eraseOp(scopeOp);
             return mlir::success();
         }

test/Transforms/Folds/scope_simplification.mlir

@@ -68,7 +68,10 @@ module {
 
   // CHECK-LABEL: @empty_scope_with_annotations_preserved
   p4hir.func @empty_scope_with_annotations_preserved(%arg0: !ref_i32i) {
-    // TODO: Doesn't work(but should)
+    // Empty scopes with annotations should be preserved
+    // CHECK: p4hir.scope annotations {name = "test"} {
+    // CHECK-NEXT: }
+    // CHECK-NEXT: p4hir.return
     p4hir.scope annotations {name = "test"} {
     }
     p4hir.return