Support for op= assignments in the frontend/midend

frontends/CMakeLists.txt

@@ -45,6 +45,7 @@ set (P4_FRONTEND_SRCS
   p4/parserControlFlow.cpp
   p4/reassociation.cpp
   p4/redundantParsers.cpp
+  p4/removeOpAssign.cpp
   p4/removeParameters.cpp
   p4/removeReturns.cpp
   p4/reservedWords.cpp
@@ -123,6 +124,7 @@ set (P4_FRONTEND_HDRS
   p4/parserControlFlow.h
   p4/reassociation.h
   p4/redundantParsers.h
+  p4/removeOpAssign.h
   p4/removeParameters.h
   p4/removeReturns.h
   p4/reservedWords.h

frontends/p4/frontend.cpp

@@ -50,6 +50,7 @@ limitations under the License.
 #include "parserControlFlow.h"
 #include "reassociation.h"
 #include "redundantParsers.h"
+#include "removeOpAssign.h"
 #include "removeParameters.h"
 #include "removeReturns.h"
 #include "resetHeaders.h"
@@ -216,6 +217,7 @@ const IR::P4Program *FrontEnd::run(const CompilerOptions &options, const IR::P4P
         new MoveDeclarations(),  // Move all local declarations to the beginning
         new MoveInitializers(),
         new SideEffectOrdering(&typeMap, policy->skipSideEffectOrdering()),
+        policy->removeOpAssign() ? new RemoveOpAssign : nullptr,
         new SimplifyControlFlow(&typeMap),
         new SimplifySwitch(&typeMap),
         new MoveDeclarations(),  // Move all local declarations to the beginning

frontends/p4/frontend.h

@@ -43,6 +43,11 @@ class FrontEndPolicy : public RemoveUnusedPolicy {
     /// @returns Defaults to nullptr.
     virtual ParseAnnotations *getParseAnnotations() const { return nullptr; }
 
+    /// Indicates whether OpAssignmentExpressions should be expanded and replaced
+    /// by simple assignments.  This depends on SideEffectOrdering to be correct, so
+    /// should probably be false if skipSideEffectOrdering is true.
+    virtual bool removeOpAssign() const { return true; }
+
     /// Indicates whether the side-effect-ordering pass should be skipped.
     /// @returns Defaults to false.
     // TODO: This should probably not be allowed to be skipped at all.

frontends/p4/removeOpAssign.cpp

@@ -0,0 +1,27 @@
+/*
+Copyright 2025 Nvidia, 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.
+*/
+
+#include "removeOpAssign.h"
+
+#include "cloner.h"
+
+namespace P4 {
+
+const IR::Node *RemoveOpAssign::finish(IR::AssignmentStatement *as, const IR::Expression *e) {
+    return new IR::AssignmentStatement(as->srcInfo, as->left->apply(CloneExpressions()), e);
+}
+
+}  // namespace P4

frontends/p4/removeOpAssign.h

@@ -0,0 +1,62 @@
+/*
+Copyright 2025 Nvidia, 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 FRONTENDS_P4_REMOVEOPASSIGN_H_
+#define FRONTENDS_P4_REMOVEOPASSIGN_H_
+
+#include "ir/ir.h"
+
+namespace P4 {
+
+class RemoveOpAssign : public Transform {
+    const IR::Node *finish(IR::AssignmentStatement *as, const IR::Expression *e);
+
+#define PREORDER(OP)                                        \
+    const IR::Node *preorder(IR::OP##Assign *as) override { \
+        prune();                                            \
+        return finish(as, new IR::OP(as->left, as->right)); \
+    }
+    PREORDER(Mul)
+    PREORDER(Div)
+    PREORDER(Mod)
+    PREORDER(Add)
+    PREORDER(Sub)
+    PREORDER(AddSat)
+    PREORDER(SubSat)
+    PREORDER(Shl)
+    PREORDER(Shr)
+    PREORDER(BAnd)
+    PREORDER(BOr)
+    PREORDER(BXor)
+#undef PREORDER
+
+    const IR::Node *preorder(IR::AssignmentStatement *s) override {
+        prune();
+        return s;
+    }
+    const IR::Node *preorder(IR::Expression *e) override {
+        prune();
+        return e;
+    }
+    const IR::Node *preorder(IR::Annotation *a) override {
+        prune();
+        return a;
+    }
+};
+
+}  // namespace P4
+
+#endif /* FRONTENDS_P4_REMOVEOPASSIGN_H_ */

frontends/p4/simplifyDefUse.cpp

@@ -794,6 +794,7 @@ class FindUninitialized : public Inspector {
         Log::TempIndent indent;
         LOG3("FU Visiting " << dbp(statement) << " " << statement << indent);
         if (!unreachable) {
+            if (statement->is<IR::OpAssignmentStatement>()) visit(statement->left);
             lhs = true;
             visit(statement->left);
             checkHeaderFieldWrite(statement->left, statement->left);

frontends/p4/toP4/toP4.cpp

@@ -1091,6 +1091,17 @@ bool ToP4::preorder(const IR::AssignmentStatement *a) {
     return false;
 }
 
+bool ToP4::preorder(const IR::OpAssignmentStatement *a) {
+    dump(2);
+    visit(a->left);
+    builder.append(" ");
+    builder.append(a->getStringOp());
+    builder.append("= ");
+    visit(a->right);
+    builder.endOfStatement();
+    return false;
+}
+
 bool ToP4::preorder(const IR::BlockStatement *s) {
     dump(1);
     if (printAnnotations(s)) builder.spc();

frontends/p4/toP4/toP4.h

@@ -220,6 +220,7 @@ class ToP4 : public Inspector, ResolutionContext {
 
     // statements
     bool preorder(const IR::AssignmentStatement *s) override;
+    bool preorder(const IR::OpAssignmentStatement *s) override;
     bool preorder(const IR::BlockStatement *s) override;
     bool preorder(const IR::MethodCallStatement *s) override;
     bool preorder(const IR::EmptyStatement *s) override;

frontends/p4/typeChecking/readOnlyTypeInference.cpp

@@ -140,6 +140,7 @@ DEFINE_POSTORDER(IR::ReturnStatement)
 DEFINE_POSTORDER(IR::IfStatement)
 DEFINE_POSTORDER(IR::SwitchStatement)
 DEFINE_POSTORDER(IR::AssignmentStatement)
+DEFINE_POSTORDER(IR::OpAssignmentStatement)
 DEFINE_POSTORDER(IR::ForInStatement)
 DEFINE_POSTORDER(IR::ActionListElement)
 DEFINE_POSTORDER(IR::KeyElement)

frontends/p4/typeChecking/typeCheckStmt.cpp

@@ -137,11 +137,25 @@ const IR::Node *TypeInferenceBase::postorder(const IR::AssignmentStatement *assi
     }
 
     auto newInit = assignment(assign, ltype, assign->right);
-    if (newInit != assign->right)
-        assign = new IR::AssignmentStatement(assign->srcInfo, assign->left, newInit);
+    if (newInit != assign->right) {
+        auto *clone = assign->clone();
+        clone->right = newInit;
+        assign = clone;
+    }
     return assign;
 }
 
+const IR::Node *TypeInferenceBase::postorder(const IR::OpAssignmentStatement *assign) {
+    auto ltype = getType(assign->left);
+    if (ltype == nullptr) return assign;
+    if (!ltype->is<IR::Type_Bits>()) {
+        typeError("%1%=: cannot be applied to '%2%' with type '%3%'", assign->getStringOp(),
+                  assign->left, ltype->toString());
+        return assign;
+    }
+    return TypeInferenceBase::postorder(static_cast<const IR::AssignmentStatement *>(assign));
+}
+
 const IR::Node *TypeInferenceBase::postorder(const IR::ForInStatement *forin) {
     LOG3("TI Visiting " << dbp(getOriginal()));
     auto ltype = getType(forin->ref);

frontends/p4/typeChecking/typeChecker.h

@@ -329,6 +329,7 @@ class TypeInferenceBase : public virtual Visitor, public ResolutionContext {
     const IR::Node *postorder(const IR::IfStatement *stat);
     const IR::Node *postorder(const IR::SwitchStatement *stat);
     const IR::Node *postorder(const IR::AssignmentStatement *stat);
+    const IR::Node *postorder(const IR::OpAssignmentStatement *stat);
     const IR::Node *postorder(const IR::ForInStatement *stat);
     const IR::Node *postorder(const IR::ActionListElement *elem);
     const IR::Node *postorder(const IR::KeyElement *elem);
@@ -472,6 +473,7 @@ class ReadOnlyTypeInference : public virtual Inspector, public TypeInferenceBase
     void postorder(const IR::IfStatement *stat) override;
     void postorder(const IR::SwitchStatement *stat) override;
     void postorder(const IR::AssignmentStatement *stat) override;
+    void postorder(const IR::OpAssignmentStatement *stat) override;
     void postorder(const IR::ForInStatement *stat) override;
     void postorder(const IR::ActionListElement *elem) override;
     void postorder(const IR::KeyElement *elem) override;
@@ -608,6 +610,7 @@ class TypeInference : public virtual Transform, public TypeInferenceBase {
     const IR::Node *postorder(IR::IfStatement *stat) override;
     const IR::Node *postorder(IR::SwitchStatement *stat) override;
     const IR::Node *postorder(IR::AssignmentStatement *stat) override;
+    const IR::Node *postorder(IR::OpAssignmentStatement *stat) override;
     const IR::Node *postorder(IR::ForInStatement *stat) override;
     const IR::Node *postorder(IR::ActionListElement *elem) override;
     const IR::Node *postorder(IR::KeyElement *elem) override;

frontends/p4/typeChecking/typeInference.cpp

@@ -130,6 +130,7 @@ DEFINE_POSTORDER(IR::ReturnStatement)
 DEFINE_POSTORDER(IR::IfStatement)
 DEFINE_POSTORDER(IR::SwitchStatement)
 DEFINE_POSTORDER(IR::AssignmentStatement)
+DEFINE_POSTORDER(IR::OpAssignmentStatement)
 DEFINE_POSTORDER(IR::ForInStatement)
 DEFINE_POSTORDER(IR::ActionListElement)
 DEFINE_POSTORDER(IR::KeyElement)

frontends/parsers/p4/p4lexer.ll

@@ -283,6 +283,19 @@ using Parser = P4::P4Parser;
 "<="    { BEGIN(driver.saveState); driver.template_args = false; return makeToken(LE); }
 "++"    { BEGIN(driver.saveState); driver.template_args = false; return makeToken(PP); }
 
+"+="    { BEGIN(driver.saveState); driver.template_args = false; return makeToken(ASSIGN_PLUS); }
+"-="    { BEGIN(driver.saveState); driver.template_args = false; return makeToken(ASSIGN_MINUS); }
+"*="    { BEGIN(driver.saveState); driver.template_args = false; return makeToken(ASSIGN_MUL); }
+"/="    { BEGIN(driver.saveState); driver.template_args = false; return makeToken(ASSIGN_DIV); }
+"%="    { BEGIN(driver.saveState); driver.template_args = false; return makeToken(ASSIGN_MOD); }
+"&="    { BEGIN(driver.saveState); driver.template_args = false; return makeToken(ASSIGN_BIT_AND); }
+"|="    { BEGIN(driver.saveState); driver.template_args = false; return makeToken(ASSIGN_BIT_OR); }
+"^="    { BEGIN(driver.saveState); driver.template_args = false; return makeToken(ASSIGN_XOR); }
+"<<="   { BEGIN(driver.saveState); driver.template_args = false; return makeToken(ASSIGN_SHL); }
+">>="   { BEGIN(driver.saveState); driver.template_args = false; return makeToken(ASSIGN_SHR); }
+"|+|="  { BEGIN(driver.saveState); driver.template_args = false; return makeToken(ASSIGN_PLUS_SAT); }
+"|-|="  { BEGIN(driver.saveState); driver.template_args = false; return makeToken(ASSIGN_MINUS_SAT); }
+
 "+"     { BEGIN(driver.saveState); driver.template_args = false; return makeToken(PLUS); }
 "{#}"   { BEGIN(driver.saveState); driver.template_args = false; return makeToken(INVALID); }
 "|+|"   { BEGIN(driver.saveState); driver.template_args = false; return makeToken(PLUS_SAT); }

frontends/parsers/p4/p4parser.ypp

@@ -231,6 +231,18 @@ using namespace P4;
 %token<Token>      QUESTION "?"
 %token<Token>      DOT "."
 %token<Token>      ASSIGN "="
+%token<Token>      ASSIGN_PLUS "+="
+%token<Token>      ASSIGN_MINUS "-="
+%token<Token>      ASSIGN_MUL "*="
+%token<Token>      ASSIGN_DIV "/="
+%token<Token>      ASSIGN_MOD "%="
+%token<Token>      ASSIGN_BIT_AND "&="
+%token<Token>      ASSIGN_BIT_OR "|="
+%token<Token>      ASSIGN_XOR "^="
+%token<Token>      ASSIGN_SHL "<<="
+%token<Token>      ASSIGN_SHR ">>="
+%token<Token>      ASSIGN_PLUS_SAT "|+|="
+%token<Token>      ASSIGN_MINUS_SAT "|-|="
 %token<Token>      SEMICOLON ";"
 %token<Token>      AT "@"
 %token<Token>      PP "++"
@@ -1270,6 +1282,32 @@ assignmentOrMethodCallStatementWithoutSemicolon
     | lvalue "=" expression
         { $$ = new IR::AssignmentStatement(@2, $1, $3); }
 
+    | lvalue "*=" expression
+        { $$ = new IR::MulAssign(@2, $1, $3); }
+    | lvalue "/=" expression
+        { $$ = new IR::DivAssign(@2, $1, $3); }
+    | lvalue "%=" expression
+        { $$ = new IR::ModAssign(@2, $1, $3); }
+    | lvalue "+=" expression
+        { $$ = new IR::AddAssign(@2, $1, $3); }
+    | lvalue "-=" expression
+        { $$ = new IR::SubAssign(@2, $1, $3); }
+    | lvalue "|+|=" expression
+        { $$ = new IR::AddSatAssign(@2, $1, $3); }
+    | lvalue "|-|=" expression
+        { $$ = new IR::SubSatAssign(@2, $1, $3); }
+    | lvalue "<<=" expression
+        { $$ = new IR::ShlAssign(@2, $1, $3); }
+    | lvalue ">>=" expression
+        { $$ = new IR::ShrAssign(@2, $1, $3); }
+    | lvalue R_ANGLE_SHIFT ">=" expression
+        { $$ = new IR::ShrAssign(@2+@3, $1, $4); }
+    | lvalue "&=" expression
+        { $$ = new IR::BAndAssign(@2, $1, $3); }
+    | lvalue "|=" expression
+        { $$ = new IR::BOrAssign(@2, $1, $3); }
+    | lvalue "^=" expression
+        { $$ = new IR::BXorAssign(@2, $1, $3); }
     ;
 
 emptyStatement

ir/dbprint-stmt.cpp

@@ -41,6 +41,12 @@ void IR::AssignmentStatement::dbprint(std::ostream &out) const {
     if (!prec) out << ';';
 }
 
+void IR::OpAssignmentStatement::dbprint(std::ostream &out) const {
+    int prec = getprec(out);
+    out << Prec_Low << left << " " << getStringOp() << "= " << right << setprec(prec);
+    if (!prec) out << ';';
+}
+
 void IR::IfStatement::dbprint(std::ostream &out) const {
     int prec = getprec(out);
     out << Prec_Low << "if (" << condition << ") {" << indent << setprec(0) << Log::endl << ifTrue;