diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
index 0ba9931b2fd..e7cf63cb47a 100644
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -332,6 +332,57 @@ acquireLocksOnSubLinks(Node *node, acquireLocksOnSubLinks_context *context)
 	return expression_tree_walker(node, acquireLocksOnSubLinks, context);
 }
 
+static List *
+add_generated_columns_to_targetlist(List *targetList,
+									CmdType commandType,
+									Relation target_relation,
+									RangeTblEntry *target_rte,
+									int target_rt_index,
+									int result_relation,
+									int nomatch_varno)
+{
+	TupleDesc	tupdesc = RelationGetDescr(target_relation);
+	List	   *result;
+
+	if (!(tupdesc->constr &&
+		  (tupdesc->constr->has_generated_stored ||
+		   tupdesc->constr->has_generated_virtual)))
+		return targetList;
+
+	result = list_copy(targetList);
+
+	for (int i = 0; i < tupdesc->natts; i++)
+	{
+		Form_pg_attribute attr = TupleDescAttr(tupdesc, i);
+		Node	   *expr;
+		TargetEntry *te;
+
+		if (!attr->attgenerated)
+			continue;
+
+		expr = build_generation_expression(target_relation, i + 1);
+		ChangeVarNodes(expr, 1, target_rt_index, 0);
+		expr = ReplaceVarsFromTargetList(expr,
+										 target_rt_index,
+										 0,
+										 target_rte,
+										 targetList,
+										 result_relation,
+										 (commandType == CMD_UPDATE) ?
+										 REPLACEVARS_CHANGE_VARNO :
+										 REPLACEVARS_SUBSTITUTE_NULL,
+										 nomatch_varno,
+										 NULL);
+
+		te = makeTargetEntry((Expr *) expr,
+							 i + 1,
+							 pstrdup(NameStr(attr->attname)),
+							 false);
+		result = lappend(result, te);
+	}
+
+	return result;
+}
 
 /*
  * rewriteRuleAction -
@@ -645,24 +696,16 @@ rewriteRuleAction(Query *parsetree,
 	{
 		RangeTblEntry *new_rte = rt_fetch(new_varno, sub_action->rtable);
 		Relation	new_rel;
+		List	   *new_targetlist;
 
-		/*
-		 * First, expand any generated column references in the NEW relation.
-		 * This must happen before ReplaceVarsFromTargetList, because the
-		 * query's target list won't contain entries for generated columns
-		 * (they are removed by rewriteTargetListIU).  Without this step,
-		 * ReplaceVarsFromTargetList would fall through to
-		 * REPLACEVARS_CHANGE_VARNO (for UPDATE) or
-		 * REPLACEVARS_SUBSTITUTE_NULL (for INSERT), producing wrong results
-		 * when the generated column depends on columns being modified.
-		 */
 		new_rel = relation_open(new_rte->relid, NoLock);
-		sub_action = (Query *)
-			expand_generated_columns_internal((Node *) sub_action,
-											  new_rel, new_varno,
-											  new_rte,
-											  sub_action->resultRelation,
-											  true);
+		new_targetlist = add_generated_columns_to_targetlist(parsetree->targetList,
+															 event,
+															 new_rel,
+															 new_rte,
+															 new_varno,
+															 sub_action->resultRelation,
+															 current_varno);
 		relation_close(new_rel, NoLock);
 
 		sub_action = (Query *)
@@ -670,7 +713,7 @@ rewriteRuleAction(Query *parsetree,
 									  new_varno,
 									  0,
 									  new_rte,
-									  parsetree->targetList,
+									  new_targetlist,
 									  sub_action->resultRelation,
 									  (event == CMD_UPDATE) ?
 									  REPLACEVARS_CHANGE_VARNO :
diff --git a/src/test/regress/expected/generated_stored.out b/src/test/regress/expected/generated_stored.out
index 363bd7ccad5..be421b738bc 100644
--- a/src/test/regress/expected/generated_stored.out
+++ b/src/test/regress/expected/generated_stored.out
@@ -1600,20 +1600,22 @@ CREATE TABLE gtest28b (LIKE gtest28a INCLUDING GENERATED);
 
 -- rules
 -- NEW.b in a rule action should reflect the generated column's new value
-CREATE TABLE gtest_rule (a int, b int GENERATED ALWAYS AS (a * 2) STORED);
+CREATE TABLE gtest_rule (a int, c int, b int GENERATED ALWAYS AS (a + c) STORED);
 CREATE TABLE gtest_rule_log (op text, old_b int, new_b int);
 CREATE RULE gtest_rule_upd AS ON UPDATE TO gtest_rule
   DO ALSO INSERT INTO gtest_rule_log VALUES ('UPD', OLD.b, NEW.b);
 CREATE RULE gtest_rule_ins AS ON INSERT TO gtest_rule
   DO ALSO INSERT INTO gtest_rule_log VALUES ('INS', NULL, NEW.b);
-INSERT INTO gtest_rule (a) VALUES (1);
+INSERT INTO gtest_rule (a, c) VALUES (1, 5);
 UPDATE gtest_rule SET a = 10;
+UPDATE gtest_rule SET c = 7;
 SELECT * FROM gtest_rule_log;
  op  | old_b | new_b 
 -----+-------+-------
- INS |       |     2
- UPD |     2 |    20
-(2 rows)
+ INS |       |     6
+ UPD |     6 |    15
+ UPD |    15 |    17
+(3 rows)
 
 DROP TABLE gtest_rule, gtest_rule_log;
 -- sanity check of system catalog
diff --git a/src/test/regress/expected/generated_virtual.out b/src/test/regress/expected/generated_virtual.out
index 5bbd2f684ce..35b54f98ca1 100644
--- a/src/test/regress/expected/generated_virtual.out
+++ b/src/test/regress/expected/generated_virtual.out
@@ -1522,20 +1522,22 @@ CREATE TABLE gtest28b (LIKE gtest28a INCLUDING GENERATED);
 
 -- rules
 -- NEW.b in a rule action should reflect the generated column's new value
-CREATE TABLE gtest_rule (a int, b int GENERATED ALWAYS AS (a * 2) STORED);
+CREATE TABLE gtest_rule (a int, c int, b int GENERATED ALWAYS AS (a + c) VIRTUAL);
 CREATE TABLE gtest_rule_log (op text, old_b int, new_b int);
 CREATE RULE gtest_rule_upd AS ON UPDATE TO gtest_rule
   DO ALSO INSERT INTO gtest_rule_log VALUES ('UPD', OLD.b, NEW.b);
 CREATE RULE gtest_rule_ins AS ON INSERT TO gtest_rule
   DO ALSO INSERT INTO gtest_rule_log VALUES ('INS', NULL, NEW.b);
-INSERT INTO gtest_rule (a) VALUES (1);
+INSERT INTO gtest_rule (a, c) VALUES (1, 5);
 UPDATE gtest_rule SET a = 10;
+UPDATE gtest_rule SET c = 7;
 SELECT * FROM gtest_rule_log;
  op  | old_b | new_b 
 -----+-------+-------
- INS |       |     2
- UPD |     2 |    20
-(2 rows)
+ INS |       |     6
+ UPD |     6 |    15
+ UPD |    15 |    17
+(3 rows)
 
 DROP TABLE gtest_rule, gtest_rule_log;
 -- sanity check of system catalog
diff --git a/src/test/regress/sql/generated_stored.sql b/src/test/regress/sql/generated_stored.sql
index 3265a1e88c6..bda1fe760e8 100644
--- a/src/test/regress/sql/generated_stored.sql
+++ b/src/test/regress/sql/generated_stored.sql
@@ -800,14 +800,15 @@ CREATE TABLE gtest28b (LIKE gtest28a INCLUDING GENERATED);
 
 -- rules
 -- NEW.b in a rule action should reflect the generated column's new value
-CREATE TABLE gtest_rule (a int, b int GENERATED ALWAYS AS (a * 2) STORED);
+CREATE TABLE gtest_rule (a int, c int, b int GENERATED ALWAYS AS (a + c) STORED);
 CREATE TABLE gtest_rule_log (op text, old_b int, new_b int);
 CREATE RULE gtest_rule_upd AS ON UPDATE TO gtest_rule
   DO ALSO INSERT INTO gtest_rule_log VALUES ('UPD', OLD.b, NEW.b);
 CREATE RULE gtest_rule_ins AS ON INSERT TO gtest_rule
   DO ALSO INSERT INTO gtest_rule_log VALUES ('INS', NULL, NEW.b);
-INSERT INTO gtest_rule (a) VALUES (1);
+INSERT INTO gtest_rule (a, c) VALUES (1, 5);
 UPDATE gtest_rule SET a = 10;
+UPDATE gtest_rule SET c = 7;
 SELECT * FROM gtest_rule_log;
 DROP TABLE gtest_rule, gtest_rule_log;
 
diff --git a/src/test/regress/sql/generated_virtual.sql b/src/test/regress/sql/generated_virtual.sql
index cd1cd938975..b7c4c711e61 100644
--- a/src/test/regress/sql/generated_virtual.sql
+++ b/src/test/regress/sql/generated_virtual.sql
@@ -813,14 +813,15 @@ CREATE TABLE gtest28b (LIKE gtest28a INCLUDING GENERATED);
 
 -- rules
 -- NEW.b in a rule action should reflect the generated column's new value
-CREATE TABLE gtest_rule (a int, b int GENERATED ALWAYS AS (a * 2) STORED);
+CREATE TABLE gtest_rule (a int, c int, b int GENERATED ALWAYS AS (a + c) VIRTUAL);
 CREATE TABLE gtest_rule_log (op text, old_b int, new_b int);
 CREATE RULE gtest_rule_upd AS ON UPDATE TO gtest_rule
   DO ALSO INSERT INTO gtest_rule_log VALUES ('UPD', OLD.b, NEW.b);
 CREATE RULE gtest_rule_ins AS ON INSERT TO gtest_rule
   DO ALSO INSERT INTO gtest_rule_log VALUES ('INS', NULL, NEW.b);
-INSERT INTO gtest_rule (a) VALUES (1);
+INSERT INTO gtest_rule (a, c) VALUES (1, 5);
 UPDATE gtest_rule SET a = 10;
+UPDATE gtest_rule SET c = 7;
 SELECT * FROM gtest_rule_log;
 DROP TABLE gtest_rule, gtest_rule_log;
 
