Fix issue: annotations get optimized out.

include/p4mlir/Dialect/P4HIR/P4HIR_Ops.td

@@ -585,7 +585,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

@@ -1039,6 +1039,26 @@ LogicalResult P4HIR::ScopeOp::verify() {
     return success();
 }
 
+void P4HIR::ScopeOp::getEffects(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>> &effects) {
+    if (auto annotations = getAnnotations()) {
+        bool isEmpty =
+            getScopeRegion().hasOneBlock() && llvm::hasSingleElement(getScopeRegion().front());
+
+        if (llvm::any_of(annotations->getValue(),
+                         [&](auto attr) { return attr.getName() != "atomic" || !isEmpty; })) {
+            effects.emplace_back(MemoryEffects::Write::get(), SideEffects::DefaultResource::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() &&
@@ -1294,6 +1314,69 @@ Block *P4HIR::IfOp::elseBlock() {
 P4HIR::YieldOp P4HIR::IfOp::elseYield() { return cast<P4HIR::YieldOp>(&elseBlock()->back()); }
 
 namespace {
+
+static std::optional<bool> getConstantBool(mlir::Value value) {
+    auto cst = value.getDefiningOp<P4HIR::ConstOp>();
+    if (!cst) return std::nullopt;
+
+    auto boolAttr = mlir::dyn_cast<P4HIR::BoolAttr>(cst.getValue());
+    if (!boolAttr) return std::nullopt;
+
+    return boolAttr.getValue();
+}
+
+static bool isBranchHint(llvm::StringRef name) { return name == "likely" || name == "unlikely"; }
+
+static mlir::DictionaryAttr filterBranchHints(mlir::DictionaryAttr annotations,
+                                              mlir::MLIRContext *ctx) {
+    if (!annotations) return nullptr;
+
+    auto attrs = annotations.getValue();
+
+    if (llvm::none_of(attrs, [](auto attr) { return isBranchHint(attr.getName().getValue()); })) {
+        return annotations;
+    }
+
+    llvm::SmallVector<mlir::NamedAttribute> filtered;
+    llvm::copy_if(attrs, std::back_inserter(filtered),
+                  [](auto attr) { return !isBranchHint(attr.getName().getValue()); });
+    return filtered.empty() ? nullptr : mlir::DictionaryAttr::get(ctx, filtered);
+}
+
+static void inlineBlockAndReplaceOp(mlir::Block *block, mlir::Operation *op,
+                                    mlir::PatternRewriter &rewriter) {
+    mlir::Operation *terminator = block->getTerminator();
+    mlir::ValueRange results = terminator->getOperands();
+    rewriter.inlineBlockBefore(block, op, /*blockArgs=*/{});
+    op->getNumResults() == 0 ? rewriter.eraseOp(op) : rewriter.replaceOp(op, results);
+    rewriter.eraseOp(terminator);
+}
+
+static mlir::LogicalResult foldConstantIfBranch(P4HIR::IfOp ifOp, mlir::Region &branchRegion,
+                                                mlir::Block *branchBlock,
+                                                mlir::DictionaryAttr branchAnnotations,
+                                                mlir::PatternRewriter &rewriter) {
+    if (!branchBlock) {
+        rewriter.eraseOp(ifOp);
+        return mlir::success();
+    }
+
+    auto filteredAnnotations = filterBranchHints(branchAnnotations, rewriter.getContext());
+    if (!filteredAnnotations) {
+        inlineBlockAndReplaceOp(branchBlock, ifOp, rewriter);
+        return mlir::success();
+    }
+
+    auto scopeOp = rewriter.create<P4HIR::ScopeOp>(ifOp.getLoc(), filteredAnnotations,
+                                                   [](mlir::OpBuilder &, mlir::Location) {});
+    scopeOp.getScopeRegion().front().erase();
+    rewriter.inlineRegionBefore(branchRegion, scopeOp.getScopeRegion(),
+                                scopeOp.getScopeRegion().begin());
+    ifOp.getNumResults() == 0 ? rewriter.eraseOp(ifOp)
+                              : rewriter.replaceOp(ifOp, scopeOp.getResults());
+    return mlir::success();
+}
+
 struct RemoveEmptyElseBranch : public OpRewritePattern<P4HIR::IfOp> {
     using OpRewritePattern<P4HIR::IfOp>::OpRewritePattern;
 
@@ -1358,10 +1441,26 @@ struct InvertIfCondition : public OpRewritePattern<P4HIR::IfOp> {
         return success();
     }
 };
+
+struct FoldConstantIf : public OpRewritePattern<P4HIR::IfOp> {
+    using OpRewritePattern<P4HIR::IfOp>::OpRewritePattern;
+
+    LogicalResult matchAndRewrite(P4HIR::IfOp ifOp, PatternRewriter &rewriter) const override {
+        auto constValue = getConstantBool(ifOp.getCondition());
+        if (!constValue) return failure();
+
+        return *constValue ? foldConstantIfBranch(ifOp, ifOp.getThenRegion(), ifOp.thenBlock(),
+                                                  ifOp.getThenAnnotationsAttr(), rewriter)
+                           : foldConstantIfBranch(ifOp, ifOp.getElseRegion(), ifOp.elseBlock(),
+                                                  ifOp.getElseAnnotationsAttr(), rewriter);
+    }
+};
+
 }  // namespace
 
 void P4HIR::IfOp::getCanonicalizationPatterns(RewritePatternSet &results, MLIRContext *context) {
-    results.add<RemoveEmptyElseBranch, RemoveEmptyThenBranch, InvertIfCondition>(context);
+    results.add<RemoveEmptyElseBranch, RemoveEmptyThenBranch, InvertIfCondition, FoldConstantIf>(
+        context);
 }
 
 static mlir::LogicalResult verifyReturnLike(mlir::Operation *op, bool mayHaveNoOperands) {

test/Transforms/Folds/if.mlir

@@ -84,4 +84,63 @@ module {
     p4hir.return %result : !i32i
   }
 
+  // CHECK-LABEL: constant_fold_with_annotation
+  p4hir.func @constant_fold_with_annotation(%arg0: !p4hir.ref<!i32i>) {
+    // Constant true with annotation should create scope with annotation preserved
+    // CHECK:      %c1_i32i = p4hir.const #int1_i32i
+    // CHECK-NEXT: p4hir.scope annotations {name = "fast_path"} {
+    // CHECK-NEXT:   p4hir.assign %c1_i32i, %arg0 : <!i32i>
+    // CHECK-NEXT: }
+    // CHECK-NEXT: p4hir.return
+
+    %true = p4hir.const #p4hir.bool<true> : !p4hir.bool
+    %c1_i32i = p4hir.const #int1_i32i
+
+    p4hir.if %true annotations {name = "fast_path"} {
+      p4hir.assign %c1_i32i, %arg0 : <!i32i>
+    }
+
+    p4hir.return
+  }
+
+  // CHECK-LABEL: constant_fold_filters_branch_hints
+  p4hir.func @constant_fold_filters_branch_hints(%arg0: !p4hir.ref<!i32i>) {
+    // Constant true with likely annotation should create scope without likely (branch hint doesn't apply to scope)
+    // CHECK:      %c1_i32i = p4hir.const #int1_i32i
+    // CHECK-NEXT: p4hir.scope annotations {name = "debug"} {
+    // CHECK-NEXT:   p4hir.assign %c1_i32i, %arg0 : <!i32i>
+    // CHECK-NEXT: }
+    // CHECK-NEXT: p4hir.return
+
+    %true = p4hir.const #p4hir.bool<true> : !p4hir.bool
+    %c1_i32i = p4hir.const #int1_i32i
+
+    p4hir.if %true annotations {likely, name = "debug"} {
+      p4hir.assign %c1_i32i, %arg0 : <!i32i>
+    }
+
+    p4hir.return
+  }
+
+  // CHECK-LABEL: constant_fold_only_branch_hints
+  p4hir.func @constant_fold_only_branch_hints(%arg0: !p4hir.ref<!i32i>) {
+    // Constant false with only unlikely annotation - should inline directly (no annotations to preserve)
+    // CHECK:      %c3_i32i = p4hir.const #int3_i32i
+    // CHECK-NEXT: p4hir.assign %c3_i32i, %arg0 : <!i32i>
+    // CHECK-NEXT: p4hir.return
+    // CHECK-NOT:  p4hir.scope
+
+    %false = p4hir.const #p4hir.bool<false> : !p4hir.bool
+    %c2_i32i = p4hir.const #int3_i32i
+
+    p4hir.if %false {
+      %dummy = p4hir.const #int1_i32i
+      p4hir.assign %dummy, %arg0 : <!i32i>
+    } else annotations {unlikely} {
+      p4hir.assign %c2_i32i, %arg0 : <!i32i>
+    }
+
+    p4hir.return
+  }
+
 }

test/Transforms/Folds/scope_simplification.mlir

@@ -72,7 +72,9 @@ module {
 
   // CHECK-LABEL: @empty_scope_with_annotations_preserved
   p4hir.func @empty_scope_with_annotations_preserved(%arg0: !ref_i32i) {
-    // TODO: Doesn't work(but should)
+    // CHECK: p4hir.scope annotations {name = "test"} {
+    // CHECK-NEXT: }
+    // CHECK-NEXT: p4hir.return
     p4hir.scope annotations {name = "test"} {
     }
     p4hir.return
@@ -94,4 +96,27 @@ module {
     }
     p4hir.return
   }
+
+  // CHECK-LABEL: @empty_scope_with_atomic_removed
+  p4hir.func @empty_scope_with_atomic_removed(%arg0: !ref_i32i) {
+    // CHECK-NOT: p4hir.scope
+    // CHECK: p4hir.return
+    p4hir.scope annotations {atomic} {
+    }
+    p4hir.return
+  }
+
+  // CHECK-LABEL: @nonempty_scope_with_atomic_preserved
+  p4hir.func @nonempty_scope_with_atomic_preserved(%arg0: !ref_i32i) {
+    // CHECK: %[[C:.*]] = p4hir.const #int1_i32i
+    // CHECK-NEXT: p4hir.scope annotations {atomic} {
+    // CHECK-NEXT: p4hir.assign %[[C]], %arg0
+    // CHECK-NEXT: }
+    // CHECK-NEXT: p4hir.return
+    %c1 = p4hir.const #int1_i32i
+    p4hir.scope annotations {atomic} {
+      p4hir.assign %c1, %arg0 : <!i32i>
+    }
+    p4hir.return
+  }
 }